From 0cffb5b537d63cb79d98297e8f044a883e870ffd Mon Sep 17 00:00:00 2001
From: Andrey Avtomonov <andreybavt@gmail.com>
Date: Sat, 23 May 2026 10:37:08 +0200
Subject: [PATCH] fix(llm): allow StructuredOutput tool and raise maxTurns for
 generateObject
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The Claude Code agent SDK announces an internal pseudo-tool named
StructuredOutput in the system/init message whenever outputFormat is set
to { type: 'json_schema' }. The runtime's isolation check built its
allowedToolIds set only from MCP tool ids and treated StructuredOutput
as an unexpected host-injected tool, so every generateObject call threw
"Claude Code runtime isolation failed: tools=StructuredOutput ..." and
the table-descriptions and relationship-LLM-proposal enrichment stages
recorded null output across the board.

Whitelist StructuredOutput specifically in generateObject's
allowedToolIds — the check also enforces missing_tools symmetry, so
generateText and runAgentLoop, which do not see StructuredOutput, must
not require it.

generateObject also ran with maxTurns: 1, which the model intermittently
breached when it emitted thinking text before the structured response.
Raised to 5 to give the schema-bound call enough headroom without
allowing unbounded loops. The existing tests now exercise the path with
an init message that announces StructuredOutput so the regression cannot
slip back in.
---
 .../src/context/llm/claude-code-runtime.test.ts  | 14 +++++++++++---
 .../cli/src/context/llm/claude-code-runtime.ts   | 16 ++++++++++++++--
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/packages/cli/src/context/llm/claude-code-runtime.test.ts b/packages/cli/src/context/llm/claude-code-runtime.test.ts
index 38959140..b1003b78 100644
--- a/packages/cli/src/context/llm/claude-code-runtime.test.ts
+++ b/packages/cli/src/context/llm/claude-code-runtime.test.ts
@@ -91,9 +91,14 @@ describe('ClaudeCodeKtxLlmRuntime', () => {
     });
   });
 
-  it('validates structured output with the caller schema', async () => {
+  it('validates structured output with the caller schema and whitelists the SDK StructuredOutput tool', async () => {
     const schema = z.object({ answer: z.string() });
-    const query = vi.fn((_input: any) => stream([initMessage(), resultMessage({ structured_output: { answer: 'yes' } })]));
+    const query = vi.fn((_input: any) =>
+      stream([
+        initMessage({ tools: ['StructuredOutput'] }),
+        resultMessage({ structured_output: { answer: 'yes' } }),
+      ]),
+    );
     const runtime = new ClaudeCodeKtxLlmRuntime({
       projectDir: '/tmp/project',
       modelSlots: { default: 'sonnet' },
@@ -341,7 +346,10 @@ describe('ClaudeCodeKtxLlmRuntime', () => {
   it('passes scrubbed env to object generation and agent loops', async () => {
     const schema = z.object({ answer: z.string() });
     const objectQuery = vi.fn((_input: any) =>
-      stream([initMessage(), resultMessage({ structured_output: { answer: 'yes' } })]),
+      stream([
+        initMessage({ tools: ['StructuredOutput'] }),
+        resultMessage({ structured_output: { answer: 'yes' } }),
+      ]),
     );
     const objectRuntime = new ClaudeCodeKtxLlmRuntime({
       projectDir: '/tmp/project',
diff --git a/packages/cli/src/context/llm/claude-code-runtime.ts b/packages/cli/src/context/llm/claude-code-runtime.ts
index c6783d71..0eb3eadb 100644
--- a/packages/cli/src/context/llm/claude-code-runtime.ts
+++ b/packages/cli/src/context/llm/claude-code-runtime.ts
@@ -47,6 +47,13 @@ const BUILTIN_TOOLS = [
 
 const KTX_MCP_SERVER_NAME = 'ktx';
 
+// SDK-internal pseudo-tool that the Claude Code CLI announces in its
+// system/init message whenever outputFormat: { type: 'json_schema' } is set.
+// Structured output is returned via result.structured_output (not through
+// canUseTool), so the tool only needs to be whitelisted for generateObject's
+// init isolation check; generateText / runAgentLoop never see it.
+const STRUCTURED_OUTPUT_TOOL_NAME = 'StructuredOutput';
+
 function isResult(message: SDKMessage): message is SDKResultMessage {
   return message.type === 'result';
 }
@@ -238,7 +245,12 @@ export class ClaudeCodeKtxLlmRuntime implements KtxLlmRuntimePort {
         projectDir: this.deps.projectDir,
         model: modelForRole(this.deps.modelSlots, input.role),
         env: this.deps.env,
-        maxTurns: 1,
+        // Structured output occasionally takes more than one assistant turn —
+        // the model may emit thinking/text before the StructuredOutput tool
+        // call, or the SDK may count assistant + tool-result as separate turns.
+        // 5 leaves headroom without enabling unbounded loops; the json_schema
+        // constraint still forces the final answer to be the schema.
+        maxTurns: 5,
         tools: input.tools,
       }),
       outputFormat: { type: 'json_schema' as const, schema: jsonSchema(input.schema as z.ZodType) },
@@ -247,7 +259,7 @@ export class ClaudeCodeKtxLlmRuntime implements KtxLlmRuntimePort {
       query: this.runQuery,
       prompt: [input.system, input.prompt].filter(Boolean).join('\n\n'),
       options,
-      allowedToolIds: new Set(mcpToolIds(input.tools ?? {})),
+      allowedToolIds: new Set([...mcpToolIds(input.tools ?? {}), STRUCTURED_OUTPUT_TOOL_NAME]),
       expectedMcpServerNames: expectedMcpServerNames(input.tools),
     });
     const error = resultError(result);