perf(streaming): inline tool call streaming + fix large-context hang

RvVeen · RvVeen · commit bfa1693ce206 · 2026-06-19T09:09:47.000+02:00
- Stream content_block_start/delta/stop inline as tool events arrive,
  so the tool name is visible immediately rather than after full stream
- Stop accumulating tool input in totalContent; only the tool name is
  needed for the bracket-format fallback parser
- Short-circuit parseBracketToolCalls with includes('[Called ') guard
  to prevent catastrophic regex backtracking on large JSON payloads
- Detect truncated tool calls (no stop event) and close the block
  cleanly with a user-visible error instead of forwarding invalid JSON
- Add stopped/blockIndex to ToolCallState; update stream-transformer.ts
  and types.ts accordingly

fix(auth): use profileArn region for OIDC token refresh

- IDC login and kiro-cli sync both defaulted oidcRegion to us-east-1
  when no explicit idc_region was set, causing token refreshes to hit
  oidc.us-east-1.amazonaws.com for eu-central-1 accounts
- Now derive oidcRegion from the profileArn region, which is authoritative
diff --git a/src/__tests__/stream-transformer.test.ts b/src/__tests__/stream-transformer.test.ts
@@ -173,11 +173,14 @@ describe('transformSdkStream: tool call streaming', () => {
     ]
     const sdk = makeSdkResponse(events)
     const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1'))
-    const toolBlock = result.find(
-      (e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.arguments
-    )
-    expect(toolBlock).toBeDefined()
-    const args = JSON.parse(toolBlock.choices[0].delta.tool_calls[0].function.arguments)
+
+    // With inline streaming, args arrive as multiple partial_json deltas.
+    // Collect all argument chunks for this tool and concatenate.
+    const argChunks = result
+      .filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.arguments !== undefined)
+      .map((e: any) => e.choices[0].delta.tool_calls[0].function.arguments)
+      .join('')
+    const args = JSON.parse(argChunks)
     expect(args.filePath).toBe('/tmp/x.txt')
     expect(args.content).toBe('hello')
   })
@@ -193,6 +196,73 @@ describe('transformSdkStream: tool call streaming', () => {
     const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name)
     expect(toolBlocks.length).toBe(2)
   })
+
+  test('truncated tool call (no stop event) closes the block and emits error text', async () => {
+    // Simulates Kiro cutting the stream mid-tool-call (large .frm JSON that
+    // exceeds the response size limit).
+    // With inline streaming: content_block_start IS emitted (tool name visible),
+    // then on truncation the block is closed and an error text is appended.
+    const events = [
+      { assistantResponseEvent: { content: 'I will edit the file.' } },
+      {
+        toolUseEvent: {
+          toolUseId: 't-trunc',
+          name: 'replaceFileContent',
+          input: '{"path":"/app.frm","content":{"forms":[{"name":"form1"'
+        }
+      }
+      // note: no stop event — stream cut off mid-JSON
+    ]
+    const sdk = makeSdkResponse(events)
+    const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1'))
+
+    // content_block_start IS emitted (inline streaming — tool name is immediately visible)
+    const toolNameBlocks = result.filter(
+      (e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name
+    )
+    expect(toolNameBlocks.length).toBeGreaterThan(0)
+    expect(toolNameBlocks[0]!.choices[0].delta.tool_calls[0].function.name).toBe(
+      'replaceFileContent'
+    )
+
+    // Must emit a text delta explaining the truncation
+    const allText = result
+      .filter((e) => e.choices?.[0]?.delta?.content)
+      .map((e: any) => e.choices[0].delta.content)
+      .join('')
+    expect(allText).toContain('truncated')
+    expect(allText).toContain('replaceFileContent')
+  })
+
+  test('tool input is NOT accumulated into totalContent — bracket parser never sees large JSON', async () => {
+    // Regression: previously `totalContent += tc.input` caused catastrophic
+    // regex backtracking when the tool input was hundreds of KB of JSON.
+    const largeInput = JSON.stringify(
+      Array.from({ length: 100 }, (_, i) => ({ id: i, value: 'x'.repeat(500) }))
+    )
+    const events = [
+      {
+        toolUseEvent: {
+          toolUseId: 't-big',
+          name: 'write',
+          input: largeInput.slice(0, largeInput.length / 2)
+        }
+      },
+      { toolUseEvent: { toolUseId: 't-big', input: largeInput.slice(largeInput.length / 2) } },
+      { toolUseEvent: { toolUseId: 't-big', stop: true } }
+    ]
+    const sdk = makeSdkResponse(events)
+    const start = performance.now()
+    const result = await collectSdk(transformSdkStream(sdk, 'auto', 'conv-1'))
+    const elapsed = performance.now() - start
+
+    // Must complete in well under 1s — would hang for seconds with catastrophic backtracking
+    expect(elapsed).toBeLessThan(500)
+
+    // Tool call must still be emitted correctly
+    const toolBlocks = result.filter((e) => e.choices?.[0]?.delta?.tool_calls?.[0]?.function?.name)
+    expect(toolBlocks.length).toBeGreaterThan(0)
+  })
 })
 
 // ── SDK stream: real token usage ─────────────────────────────────────────────
diff --git a/src/__tests__/tool-call-parser.test.ts b/src/__tests__/tool-call-parser.test.ts
@@ -65,6 +65,33 @@ describe('deduplicateToolCalls', () => {
   })
 })
 
+describe('parseBracketToolCalls: large JSON content (catastrophic backtracking guard)', () => {
+  test('returns empty array fast when "[Called " is absent — even with huge JSON', () => {
+    // A 500 KB JSON array with many nested braces — previously caused catastrophic
+    // backtracking in the regex. The short-circuit guard must make this instant.
+    const largeJson = JSON.stringify(
+      Array.from({ length: 200 }, (_, i) => ({
+        id: i,
+        data: 'x'.repeat(1000),
+        nested: { a: 1, b: [1, 2, 3] }
+      }))
+    )
+    const start = performance.now()
+    const result = parseBracketToolCalls(largeJson)
+    const elapsed = performance.now() - start
+    expect(result).toHaveLength(0)
+    expect(elapsed).toBeLessThan(50) // must be near-instant, not seconds
+  })
+
+  test('still finds bracket tool calls when "[Called " is present in large text', () => {
+    const prefix = 'x'.repeat(10000)
+    const text = `${prefix}[Called bash with args: {"command":"ls"}]${prefix}`
+    const result = parseBracketToolCalls(text)
+    expect(result).toHaveLength(1)
+    expect(result[0]!.name).toBe('bash')
+  })
+})
+
 describe('cleanToolCallsFromText', () => {
   test('removes bracket tool call from text', () => {
     const text = 'Here is the result [Called bash with args: {"command":"ls"}] done.'
diff --git a/src/core/auth/idc-auth-method.ts b/src/core/auth/idc-auth-method.ts
@@ -61,8 +61,14 @@ export class IdcAuthMethod {
     const invokedWithoutPrompts = !inputs || Object.keys(inputs).length === 0
 
     const startUrl = normalizeStartUrl(inputs?.start_url || this.config.idc_start_url) || undefined
-    const oidcRegion: KiroRegion = normalizeRegion(inputs?.idc_region || this.config.idc_region)
+    // For the OIDC device-code flow, prefer explicit idc_region, then fall back to
+    // the region from a pre-configured profileArn, then default_region. This ensures
+    // accounts with a eu-central-1 profileArn don't hit oidc.us-east-1.amazonaws.com.
     const configuredProfileArn = this.config.idc_profile_arn
+    const arnRegion = extractRegionFromArn(configuredProfileArn)
+    const oidcRegion: KiroRegion = normalizeRegion(
+      inputs?.idc_region || this.config.idc_region || arnRegion || configuredServiceRegion
+    )
     logger.log('IDC authorize: resolved defaults', {
       hasInputs: !!inputs && Object.keys(inputs).length > 0,
       invokedWithoutPrompts,
@@ -160,7 +166,7 @@ export class IdcAuthMethod {
             email,
             authMethod: 'idc',
             region: serviceRegion,
-            oidcRegion,
+            oidcRegion: oidcRegion,
             clientId: token.clientId,
             clientSecret: token.clientSecret,
             profileArn,
diff --git a/src/plugin/streaming/sdk-stream-transformer.ts b/src/plugin/streaming/sdk-stream-transformer.ts
@@ -130,21 +130,82 @@ export async function* transformSdkStream(
         }
       } else if (event.toolUseEvent) {
         const tc = event.toolUseEvent
+        // Only accumulate the tool *name* into totalContent — the input JSON is
+        // never bracket-format and can be hundreds of KB; including it causes
+        // catastrophic backtracking in parseBracketToolCalls.
         if (tc.name) totalContent += tc.name
-        if (tc.input) totalContent += tc.input
 
         if (tc.toolUseId) {
           if (currentToolCall && currentToolCall.toolUseId === tc.toolUseId) {
             currentToolCall.input += tc.input || ''
+            if (tc.input) {
+              const _c = convertToOpenAI(
+                {
+                  type: 'content_block_delta',
+                  index: currentToolCall.blockIndex,
+                  delta: { type: 'input_json_delta', partial_json: tc.input }
+                },
+                conversationId,
+                model
+              )
+              if (_c !== null) yield _c
+            }
           } else if (tc.name) {
             if (currentToolCall) toolCalls.push(currentToolCall)
+            for (const ev of stopBlock(streamState.textBlockIndex, streamState)) {
+              const _c = convertToOpenAI(ev, conversationId, model)
+              if (_c !== null) yield _c
+            }
+            const blockIndex = streamState.nextBlockIndex++
             currentToolCall = {
               toolUseId: tc.toolUseId,
               name: toolNameMapper ? toolNameMapper(tc.name) : tc.name,
-              input: tc.input || ''
+              input: tc.input || '',
+              stopped: false,
+              blockIndex
+            }
+            {
+              const _c = convertToOpenAI(
+                {
+                  type: 'content_block_start',
+                  index: blockIndex,
+                  content_block: {
+                    type: 'tool_use',
+                    id: tc.toolUseId,
+                    name: currentToolCall.name,
+                    input: {}
+                  }
+                },
+                conversationId,
+                model
+              )
+              if (_c !== null) yield _c
+            }
+            if (tc.input) {
+              const _c = convertToOpenAI(
+                {
+                  type: 'content_block_delta',
+                  index: blockIndex,
+                  delta: { type: 'input_json_delta', partial_json: tc.input }
+                },
+                conversationId,
+                model
+              )
+              if (_c !== null) yield _c
             }
           }
           if (tc.stop && currentToolCall) {
+            currentToolCall.stopped = true
+            // The AI SDK dispatches tools on message_delta, not on block stop,
+            // so closing inline is safe and lets the client show progress immediately.
+            {
+              const _c = convertToOpenAI(
+                { type: 'content_block_stop', index: currentToolCall.blockIndex },
+                conversationId,
+                model
+              )
+              if (_c !== null) yield _c
+            }
             toolCalls.push(currentToolCall)
             currentToolCall = null
           }
@@ -172,7 +233,27 @@ export async function* transformSdkStream(
     }
 
     if (currentToolCall) {
+      // Stream cut off mid-tool-call. Close the open block so the event stream
+      // is well-formed, then append a visible error so the model can recover.
+      logger.debug(
+        `[STREAM] Truncated tool call: name=${currentToolCall.name} id=${currentToolCall.toolUseId} inputLen=${currentToolCall.input.length}`
+      )
+      if (currentToolCall.blockIndex !== undefined) {
+        const _c = convertToOpenAI(
+          { type: 'content_block_stop', index: currentToolCall.blockIndex },
+          conversationId,
+          model
+        )
+        if (_c !== null) yield _c
+      }
       toolCalls.push(currentToolCall)
+      for (const ev of createTextDeltaEvents(
+        `\n\n[Kiro: tool call "${currentToolCall.name}" was truncated mid-stream — the response was too large. Try a smaller operation or split the task.]`,
+        streamState
+      )) {
+        const _c = convertToOpenAI(ev, conversationId, model)
+        if (_c !== null) yield _c
+      }
       currentToolCall = null
     }
 
@@ -205,25 +286,32 @@ export async function* transformSdkStream(
       if (_c !== null) yield _c
     }
 
-    const bracketToolCalls = parseBracketToolCalls(totalContent)
+    const bracketToolCalls = totalContent.includes('[Called ')
+      ? parseBracketToolCalls(totalContent)
+      : []
     if (bracketToolCalls.length > 0) {
       for (const btc of bracketToolCalls) {
         toolCalls.push({
           toolUseId: btc.toolUseId,
           name: btc.name,
-          input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input)
+          input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input),
+          stopped: true
+          // no blockIndex — these are emitted post-stream below
         })
       }
     }
 
     const dedupedToolCalls = deduplicateToolCallsByContent(toolCalls)
 
-    if (dedupedToolCalls.length > 0) {
-      const baseIndex = streamState.nextBlockIndex
-      for (let i = 0; i < dedupedToolCalls.length; i++) {
-        const tc = dedupedToolCalls[i]
+    // SDK tool calls were already emitted inline (content_block_start/delta/stop
+    // during streaming). Only bracket-format tool calls (blockIndex undefined)
+    // need to be emitted here.
+    const postStreamCalls = dedupedToolCalls.filter((tc) => tc.blockIndex === undefined)
+    if (postStreamCalls.length > 0) {
+      for (let i = 0; i < postStreamCalls.length; i++) {
+        const tc = postStreamCalls[i]
         if (!tc) continue
-        const blockIndex = baseIndex + i
+        const blockIndex = streamState.nextBlockIndex++
 
         {
           const _c = convertToOpenAI(
@@ -259,10 +347,7 @@ export async function* transformSdkStream(
             {
               type: 'content_block_delta',
               index: blockIndex,
-              delta: {
-                type: 'input_json_delta',
-                partial_json: inputJson
-              }
+              delta: { type: 'input_json_delta', partial_json: inputJson }
             },
             conversationId,
             model
diff --git a/src/plugin/streaming/stream-transformer.ts b/src/plugin/streaming/stream-transformer.ts
@@ -168,11 +168,13 @@ export async function* transformKiroStream(
               currentToolCall = {
                 toolUseId: tc.toolUseId,
                 name: tc.name,
-                input: tc.input || ''
+                input: tc.input || '',
+                stopped: false
               }
             }
 
             if (tc.stop && currentToolCall) {
+              currentToolCall.stopped = true
               toolCalls.push(currentToolCall)
               currentToolCall = null
             }
@@ -227,13 +229,16 @@ export async function* transformKiroStream(
       if (_c !== null) yield _c
     }
 
-    const bracketToolCalls = parseBracketToolCalls(totalContent)
+    const bracketToolCalls = totalContent.includes('[Called ')
+      ? parseBracketToolCalls(totalContent)
+      : []
     if (bracketToolCalls.length > 0) {
       for (const btc of bracketToolCalls) {
         toolCalls.push({
           toolUseId: btc.toolUseId,
           name: btc.name,
-          input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input)
+          input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input),
+          stopped: true
         })
       }
     }
diff --git a/src/plugin/streaming/types.ts b/src/plugin/streaming/types.ts
@@ -22,6 +22,14 @@ export interface ToolCallState {
   toolUseId: string
   name: string
   input: string
+  /** true only when the SDK sent a stop event for this tool call */
+  stopped: boolean
+  /**
+   * Block index assigned when content_block_start was emitted inline during
+   * streaming. Undefined for bracket-format tool calls which are emitted
+   * post-stream. Used to avoid re-emitting SDK tool calls at the end.
+   */
+  blockIndex?: number
 }
 
 export const THINKING_START_TAG = '<thinking>'
diff --git a/src/plugin/sync/kiro-cli.ts b/src/plugin/sync/kiro-cli.ts
@@ -46,10 +46,12 @@ export async function syncFromKiroCli() {
 
         const isIdc = row.key.includes('odic')
         const authMethod = isIdc ? 'idc' : 'desktop'
-        const oidcRegion = normalizeRegion(data.region)
         let profileArn: string | undefined = data.profile_arn || data.profileArn
         if (!profileArn && isIdc) profileArn = activeProfileArn || readActiveProfileArnFromKiroCli()
-        const serviceRegion = extractRegionFromArn(profileArn) || oidcRegion
+        // serviceRegion wins over data.region: kiro-cli stores data.region as the
+        // OIDC region (often us-east-1) regardless of where the account actually lives.
+        const serviceRegion = extractRegionFromArn(profileArn) || normalizeRegion(data.region)
+        const oidcRegion = serviceRegion
         const startUrl: string | undefined =
           typeof data.start_url === 'string'
             ? data.start_url

Original file line number	Diff line number	Diff line change
`@@ -168,11 +168,13 @@ export async function* transformKiroStream(`
`168`	`168`	`currentToolCall = {`
`169`	`169`	`toolUseId: tc.toolUseId,`
`170`	`170`	`name: tc.name,`
`171`		`- input: tc.input \|\| ''`
	`171`	`+ input: tc.input \|\| '',`
	`172`	`+ stopped: false`
`172`	`173`	`}`
`173`	`174`	`}`
`174`	`175`
`175`	`176`	`if (tc.stop && currentToolCall) {`
	`177`	`+ currentToolCall.stopped = true`
`176`	`178`	`toolCalls.push(currentToolCall)`
`177`	`179`	`currentToolCall = null`
`178`	`180`	`}`
`@@ -227,13 +229,16 @@ export async function* transformKiroStream(`
`227`	`229`	`if (_c !== null) yield _c`
`228`	`230`	`}`
`229`	`231`
`230`		`- const bracketToolCalls = parseBracketToolCalls(totalContent)`
	`232`	`+ const bracketToolCalls = totalContent.includes('[Called ')`
	`233`	`+ ? parseBracketToolCalls(totalContent)`
	`234`	`+ : []`
`231`	`235`	`if (bracketToolCalls.length > 0) {`
`232`	`236`	`for (const btc of bracketToolCalls) {`
`233`	`237`	`toolCalls.push({`
`234`	`238`	`toolUseId: btc.toolUseId,`
`235`	`239`	`name: btc.name,`
`236`		`- input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input)`
	`240`	`+ input: typeof btc.input === 'string' ? btc.input : JSON.stringify(btc.input),`
	`241`	`+ stopped: true`
`237`	`242`	`})`
`238`	`243`	`}`
`239`	`244`	`}`