From 08ec877198b7e940ab2b07ff3abbc24f0cfb794b Mon Sep 17 00:00:00 2001 From: Abhiram Shaji Date: Tue, 27 Oct 2026 21:26:29 -0400 Subject: [PATCH] fix(export): keep usage metadata on inline function handoff When a model call ends by asking for an inline function, the exported Strands agent returned right after forwarding messageStop. Strands sends that call's metadata event, which carries the token usage, directly after messageStop, so the caller never received usage for the call that handed off the tool and could not account for it. The stream loop now marks the handoff and inspects only the next event: a metadata event is forwarded, anything else is dropped, and the stream ends. Both loop variants (with and without execution limits) share the change. Runs without an inline function are unaffected. Closes #2695 --- .../assets.snapshot.test.ts.snap | 20 +++- src/assets/python/http/strands/base/main.py | 20 +++- .../strands-inline-function-stream.test.ts | 101 ++++++++++++++++++ 3 files changed, 133 insertions(+), 8 deletions(-) create mode 100644 src/cli/templates/__tests__/strands-inline-function-stream.test.ts diff --git a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap index 344bf98f24..603b83ee57 100644 --- a/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap +++ b/src/assets/__tests__/__snapshots__/assets.snapshot.test.ts.snap @@ -7397,10 +7397,17 @@ async def invoke(payload, context): try: {{#if inlineFunctionTools}} hit_inline_function = False + handed_off = False {{/if}} async for event in agent.stream_async( prompt, ): + {{#if inlineFunctionTools}} + if handed_off: + if isinstance(event, dict) and "metadata" in event.get("event", {}): + yield event + return + {{/if}} if not isinstance(event, dict) or "event" not in event: continue cbs = event["event"].get("contentBlockStart") @@ -7412,8 +7419,7 @@ async def invoke(payload, context): {{/if}} yield event {{#if inlineFunctionTools}} - if hit_inline_function and "messageStop" in event["event"]: - return + handed_off = hit_inline_function and "messageStop" in event["event"] {{/if}} if timeout_fired: @@ -7433,10 +7439,17 @@ async def invoke(payload, context): {{else}} {{#if inlineFunctionTools}} hit_inline_function = False + handed_off = False {{/if}} async for event in agent.stream_async( prompt, ): + {{#if inlineFunctionTools}} + if handed_off: + if isinstance(event, dict) and "metadata" in event.get("event", {}): + yield event + return + {{/if}} if not isinstance(event, dict) or "event" not in event: continue cbs = event["event"].get("contentBlockStart") @@ -7448,8 +7461,7 @@ async def invoke(payload, context): {{/if}} yield event {{#if inlineFunctionTools}} - if hit_inline_function and "messageStop" in event["event"]: - return + handed_off = hit_inline_function and "messageStop" in event["event"] {{/if}} {{/if}} diff --git a/src/assets/python/http/strands/base/main.py b/src/assets/python/http/strands/base/main.py index 69dca7e8e2..c647e62ab1 100644 --- a/src/assets/python/http/strands/base/main.py +++ b/src/assets/python/http/strands/base/main.py @@ -653,10 +653,17 @@ async def _timeout_watchdog(): try: {{#if inlineFunctionTools}} hit_inline_function = False + handed_off = False {{/if}} async for event in agent.stream_async( prompt, ): + {{#if inlineFunctionTools}} + if handed_off: + if isinstance(event, dict) and "metadata" in event.get("event", {}): + yield event + return + {{/if}} if not isinstance(event, dict) or "event" not in event: continue cbs = event["event"].get("contentBlockStart") @@ -668,8 +675,7 @@ async def _timeout_watchdog(): {{/if}} yield event {{#if inlineFunctionTools}} - if hit_inline_function and "messageStop" in event["event"]: - return + handed_off = hit_inline_function and "messageStop" in event["event"] {{/if}} if timeout_fired: @@ -689,10 +695,17 @@ async def _timeout_watchdog(): {{else}} {{#if inlineFunctionTools}} hit_inline_function = False + handed_off = False {{/if}} async for event in agent.stream_async( prompt, ): + {{#if inlineFunctionTools}} + if handed_off: + if isinstance(event, dict) and "metadata" in event.get("event", {}): + yield event + return + {{/if}} if not isinstance(event, dict) or "event" not in event: continue cbs = event["event"].get("contentBlockStart") @@ -704,8 +717,7 @@ async def _timeout_watchdog(): {{/if}} yield event {{#if inlineFunctionTools}} - if hit_inline_function and "messageStop" in event["event"]: - return + handed_off = hit_inline_function and "messageStop" in event["event"] {{/if}} {{/if}} diff --git a/src/cli/templates/__tests__/strands-inline-function-stream.test.ts b/src/cli/templates/__tests__/strands-inline-function-stream.test.ts new file mode 100644 index 0000000000..682d47a274 --- /dev/null +++ b/src/cli/templates/__tests__/strands-inline-function-stream.test.ts @@ -0,0 +1,101 @@ +import { copyAndRenderDir } from '../render.js'; +import { TEMPLATE_ROOT } from '../templateRoot.js'; +import { spawnSync } from 'node:child_process'; +import { mkdtempSync, readFileSync, rmSync } from 'node:fs'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { afterAll, describe, expect, it } from 'vitest'; + +const TEMPLATE_DIR = join(TEMPLATE_ROOT, 'python', 'http', 'strands', 'base'); +const outDirs: string[] = []; + +const toolUseStart = { contentBlockStart: { start: { toolUse: { toolUseId: 't1', name: 'lookup' } } } }; +const metadata = { metadata: { usage: { inputTokens: 10, outputTokens: 2 }, metrics: { latencyMs: 5 } } }; +const handOffCall = [ + { messageStart: { role: 'assistant' } }, + toolUseStart, + { contentBlockStop: { contentBlockIndex: 0 } }, + { messageStop: { stopReason: 'tool_use' } }, +]; +const nextCall = [{ messageStart: { role: 'assistant' } }, { messageStop: { stopReason: 'end_turn' } }]; + +async function renderMain(hasExecutionLimits: boolean): Promise { + const outDir = mkdtempSync(join(tmpdir(), 'strands-inline-')); + outDirs.push(outDir); + await copyAndRenderDir(TEMPLATE_DIR, outDir, { + isExportHarness: true, + inlineFunctionTools: [{ name: 'lookup', description: 'Look up', inputSchema: { type: 'object' } }], + hasExecutionLimits, + timeoutSeconds: 60, + }); + return readFileSync(join(outDir, 'main.py'), 'utf-8'); +} + +function extractBlock(source: string, from: string, to = from): string { + const lines = source.split('\n'); + const indentOf = (line: string) => line.length - line.trimStart().length; + const head = lines.findIndex(line => line.trimStart().startsWith(to)); + const start = lines.findIndex(line => line.trimStart().startsWith(from)); + const colon = lines.findIndex((line, i) => i >= head && line.trimEnd().endsWith(':')); + const end = lines.findIndex((line, i) => i > colon && line.trim() !== '' && indentOf(line) <= indentOf(lines[head]!)); + return lines + .slice(start, end) + .map(line => line.slice(indentOf(lines[start]!))) + .join('\n'); +} + +function forwardedEvents(main: string, stream: object[]): string[] { + const loop = extractBlock(main, 'hit_inline_function = False', 'async for event in agent.stream_async('); + const script = [ + 'import asyncio, json, sys', + '_INLINE_FUNCTION_NAMES = {"lookup"}', + extractBlock(main, 'def _is_inline_function_call'), + 'class FakeAgent:', + ' async def stream_async(self, prompt):', + ' for chunk in json.loads(sys.argv[1]):', + ' yield {"event": chunk}', + 'async def run(agent, prompt):', + loop.replace(/^/gm, ' '), + 'async def main():', + ' print(json.dumps([next(iter(e["event"])) async for e in run(FakeAgent(), "hi")]))', + 'asyncio.run(main())', + ].join('\n'); + const result = spawnSync('python3', ['-c', script, JSON.stringify(stream)], { encoding: 'utf-8' }); + expect(result.stderr).toBe(''); + return JSON.parse(result.stdout) as string[]; +} + +afterAll(() => { + for (const dir of outDirs) rmSync(dir, { recursive: true, force: true }); +}); + +describe.each([false, true])('strands inline function stream (hasExecutionLimits: %s)', hasExecutionLimits => { + it('forwards the usage metadata of the call that hands off an inline function', async () => { + const main = await renderMain(hasExecutionLimits); + + expect(forwardedEvents(main, [...handOffCall, metadata, ...nextCall])).toEqual([ + 'messageStart', + 'contentBlockStart', + 'contentBlockStop', + 'messageStop', + 'metadata', + ]); + }); + + it('stops at messageStop when the model sends no usage metadata', async () => { + const main = await renderMain(hasExecutionLimits); + + expect(forwardedEvents(main, [...handOffCall, ...nextCall])).toEqual([ + 'messageStart', + 'contentBlockStart', + 'contentBlockStop', + 'messageStop', + ]); + }); + + it('forwards every event of a run that ends without an inline function', async () => { + const main = await renderMain(hasExecutionLimits); + + expect(forwardedEvents(main, [...nextCall, metadata])).toEqual(['messageStart', 'messageStop', 'metadata']); + }); +});