diff --git a/cockpit/runtimes/aws-strands/python/src/subagent_emitter.py b/cockpit/runtimes/aws-strands/python/src/subagent_emitter.py index e858babb1..bab1487fb 100644 --- a/cockpit/runtimes/aws-strands/python/src/subagent_emitter.py +++ b/cockpit/runtimes/aws-strands/python/src/subagent_emitter.py @@ -59,13 +59,33 @@ class _DelegationState: # Finished entries are kept (not popped) so the tool's trailing result-string # yield and any stragglers stay suppressed. Growth is capped at _MAX_SESSIONS -# (dict insertion order = age; each entry is a short key string plus a -# 4-field dataclass, ~414 bytes measured with sys.getsizeof — the cap bounds -# the dict at ~200 KiB). +# (dict insertion order = age; eviction prefers finished sessions, see +# _eviction_candidate; each entry is a short key string plus a 4-field +# dataclass, ~414 bytes measured with sys.getsizeof — the cap bounds the dict +# at ~200 KiB). _sessions: dict[str, _DelegationState] = {} _MAX_SESSIONS = 512 +def _eviction_candidate(current: str) -> str: + """Pick the session to drop when the cap is exceeded: the oldest FINISHED + session first (its only job is straggler suppression); only when every + other session is still in flight, the oldest in-flight one — evicting an + unfinished session makes its next event re-emit SUBAGENT_STARTED, so that + is the last resort that keeps the cap a hard memory bound. Never the + session being registered right now.""" + oldest_unfinished: str | None = None + for key, state in _sessions.items(): + if key == current: + continue + if state.finished: + return key + if oldest_unfinished is None: + oldest_unfinished = key + assert oldest_unfinished is not None # cap > 1, so another key exists + return oldest_unfinished + + def _subagent_run_id(tool_use_id: str) -> str: return f"{tool_use_id}-sub" @@ -82,7 +102,7 @@ async def emit_subagent_events(ctx: ToolStreamEventContext): state = _DelegationState() _sessions[ctx.tool_use_id] = state while len(_sessions) > _MAX_SESSIONS: - del _sessions[next(k for k in _sessions if k != ctx.tool_use_id)] + del _sessions[_eviction_candidate(ctx.tool_use_id)] run_id = _subagent_run_id(ctx.tool_use_id) data = ctx.stream_data if state.finished and isinstance(data, dict) and "init_event_loop" in data: diff --git a/cockpit/runtimes/aws-strands/python/tests/test_subagent_emitter.py b/cockpit/runtimes/aws-strands/python/tests/test_subagent_emitter.py index 0f2c7a073..e070d2cf9 100644 --- a/cockpit/runtimes/aws-strands/python/tests/test_subagent_emitter.py +++ b/cockpit/runtimes/aws-strands/python/tests/test_subagent_emitter.py @@ -209,6 +209,44 @@ async def test_sessions_growth_is_capped(): assert f"call_{subagent_emitter._MAX_SESSIONS + 4}" in subagent_emitter._sessions +async def test_eviction_prefers_finished_sessions_over_in_flight_ones(): + # Fill the cap with UNFINISHED (in-flight) delegations, then one finished + # one inserted LAST (the youngest entry — plain oldest-first eviction + # would keep it and drop an in-flight session, whose next event would + # then re-emit SUBAGENT_STARTED). The finished entry must go first. + for i in range(subagent_emitter._MAX_SESSIONS - 1): + await _drive([{"data": "x"}], tool_use_id=f"call_inflight_{i}") + await _drive([{"result": object()}], tool_use_id="call_finished") + assert len(subagent_emitter._sessions) == subagent_emitter._MAX_SESSIONS + + out = await _drive([{"data": "y"}], tool_use_id="call_one_more") + + assert len(subagent_emitter._sessions) == subagent_emitter._MAX_SESSIONS + assert "call_finished" not in subagent_emitter._sessions + assert "call_one_more" in subagent_emitter._sessions + for i in range(subagent_emitter._MAX_SESSIONS - 1): + assert f"call_inflight_{i}" in subagent_emitter._sessions + # The newcomer started normally (STARTED + message open + delta). + assert [ev.type for ev in out] == [ + EventType.SUBAGENT_STARTED, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + ] + + # A straggler on a surviving in-flight session does NOT re-emit STARTED. + again = await _drive([{"data": "z"}], tool_use_id="call_inflight_0") + assert [ev.type for ev in again] == [EventType.TEXT_MESSAGE_CONTENT] + + +async def test_eviction_falls_back_to_oldest_in_flight_when_none_finished(): + for i in range(subagent_emitter._MAX_SESSIONS): + await _drive([{"data": "x"}], tool_use_id=f"call_inflight_{i}") + await _drive([{"data": "y"}], tool_use_id="call_one_more") + assert len(subagent_emitter._sessions) == subagent_emitter._MAX_SESSIONS + assert "call_inflight_0" not in subagent_emitter._sessions + assert "call_one_more" in subagent_emitter._sessions + + async def test_ids_derive_from_tool_use_id(): out = await _drive([{"data": "x"}, {"result": object()}], tool_use_id="call_other") assert out[0].subagent_run_id == "call_other-sub" diff --git a/deployments/ag-ui-dev/deps/aws_strands/src/subagent_emitter.py b/deployments/ag-ui-dev/deps/aws_strands/src/subagent_emitter.py index e858babb1..bab1487fb 100644 --- a/deployments/ag-ui-dev/deps/aws_strands/src/subagent_emitter.py +++ b/deployments/ag-ui-dev/deps/aws_strands/src/subagent_emitter.py @@ -59,13 +59,33 @@ class _DelegationState: # Finished entries are kept (not popped) so the tool's trailing result-string # yield and any stragglers stay suppressed. Growth is capped at _MAX_SESSIONS -# (dict insertion order = age; each entry is a short key string plus a -# 4-field dataclass, ~414 bytes measured with sys.getsizeof — the cap bounds -# the dict at ~200 KiB). +# (dict insertion order = age; eviction prefers finished sessions, see +# _eviction_candidate; each entry is a short key string plus a 4-field +# dataclass, ~414 bytes measured with sys.getsizeof — the cap bounds the dict +# at ~200 KiB). _sessions: dict[str, _DelegationState] = {} _MAX_SESSIONS = 512 +def _eviction_candidate(current: str) -> str: + """Pick the session to drop when the cap is exceeded: the oldest FINISHED + session first (its only job is straggler suppression); only when every + other session is still in flight, the oldest in-flight one — evicting an + unfinished session makes its next event re-emit SUBAGENT_STARTED, so that + is the last resort that keeps the cap a hard memory bound. Never the + session being registered right now.""" + oldest_unfinished: str | None = None + for key, state in _sessions.items(): + if key == current: + continue + if state.finished: + return key + if oldest_unfinished is None: + oldest_unfinished = key + assert oldest_unfinished is not None # cap > 1, so another key exists + return oldest_unfinished + + def _subagent_run_id(tool_use_id: str) -> str: return f"{tool_use_id}-sub" @@ -82,7 +102,7 @@ async def emit_subagent_events(ctx: ToolStreamEventContext): state = _DelegationState() _sessions[ctx.tool_use_id] = state while len(_sessions) > _MAX_SESSIONS: - del _sessions[next(k for k in _sessions if k != ctx.tool_use_id)] + del _sessions[_eviction_candidate(ctx.tool_use_id)] run_id = _subagent_run_id(ctx.tool_use_id) data = ctx.stream_data if state.finished and isinstance(data, dict) and "init_event_loop" in data: diff --git a/deployments/ag-ui-dev/deps/aws_strands/tests/test_subagent_emitter.py b/deployments/ag-ui-dev/deps/aws_strands/tests/test_subagent_emitter.py index 0f2c7a073..e070d2cf9 100644 --- a/deployments/ag-ui-dev/deps/aws_strands/tests/test_subagent_emitter.py +++ b/deployments/ag-ui-dev/deps/aws_strands/tests/test_subagent_emitter.py @@ -209,6 +209,44 @@ async def test_sessions_growth_is_capped(): assert f"call_{subagent_emitter._MAX_SESSIONS + 4}" in subagent_emitter._sessions +async def test_eviction_prefers_finished_sessions_over_in_flight_ones(): + # Fill the cap with UNFINISHED (in-flight) delegations, then one finished + # one inserted LAST (the youngest entry — plain oldest-first eviction + # would keep it and drop an in-flight session, whose next event would + # then re-emit SUBAGENT_STARTED). The finished entry must go first. + for i in range(subagent_emitter._MAX_SESSIONS - 1): + await _drive([{"data": "x"}], tool_use_id=f"call_inflight_{i}") + await _drive([{"result": object()}], tool_use_id="call_finished") + assert len(subagent_emitter._sessions) == subagent_emitter._MAX_SESSIONS + + out = await _drive([{"data": "y"}], tool_use_id="call_one_more") + + assert len(subagent_emitter._sessions) == subagent_emitter._MAX_SESSIONS + assert "call_finished" not in subagent_emitter._sessions + assert "call_one_more" in subagent_emitter._sessions + for i in range(subagent_emitter._MAX_SESSIONS - 1): + assert f"call_inflight_{i}" in subagent_emitter._sessions + # The newcomer started normally (STARTED + message open + delta). + assert [ev.type for ev in out] == [ + EventType.SUBAGENT_STARTED, + EventType.TEXT_MESSAGE_START, + EventType.TEXT_MESSAGE_CONTENT, + ] + + # A straggler on a surviving in-flight session does NOT re-emit STARTED. + again = await _drive([{"data": "z"}], tool_use_id="call_inflight_0") + assert [ev.type for ev in again] == [EventType.TEXT_MESSAGE_CONTENT] + + +async def test_eviction_falls_back_to_oldest_in_flight_when_none_finished(): + for i in range(subagent_emitter._MAX_SESSIONS): + await _drive([{"data": "x"}], tool_use_id=f"call_inflight_{i}") + await _drive([{"data": "y"}], tool_use_id="call_one_more") + assert len(subagent_emitter._sessions) == subagent_emitter._MAX_SESSIONS + assert "call_inflight_0" not in subagent_emitter._sessions + assert "call_one_more" in subagent_emitter._sessions + + async def test_ids_derive_from_tool_use_id(): out = await _drive([{"data": "x"}, {"result": object()}], tool_use_id="call_other") assert out[0].subagent_run_id == "call_other-sub" diff --git a/libs/ag-ui/src/lib/reducer.subagent.spec.ts b/libs/ag-ui/src/lib/reducer.subagent.spec.ts index 2ba8f6235..f35004eaf 100644 --- a/libs/ag-ui/src/lib/reducer.subagent.spec.ts +++ b/libs/ag-ui/src/lib/reducer.subagent.spec.ts @@ -145,6 +145,39 @@ describe('reduceEvent SUBAGENT_* lifecycle', () => { expect((content['state'] as Record)['error']).toBe('rate limited'); }); + it('two subagents reusing a toolCallId keep separate args buffers', () => { + const store = makeStore(); + reduceEvent(ev({ type: 'SUBAGENT_STARTED', subagentRunId: 'sa-1', name: 'researcher' }), store); + reduceEvent(ev({ type: 'SUBAGENT_STARTED', subagentRunId: 'sa-2', name: 'forecaster' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_START', toolCallId: 't-1', toolCallName: 'web_search', subagentRunId: 'sa-1' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_START', toolCallId: 't-1', toolCallName: 'forecast', subagentRunId: 'sa-2' }), store); + // Interleaved fragments: a shared buffer would concatenate them into + // `{"q":"{"city":"x"}"paris"}` and neither child would ever parse. + reduceEvent(ev({ type: 'TOOL_CALL_ARGS', toolCallId: 't-1', delta: '{"q":', subagentRunId: 'sa-1' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_ARGS', toolCallId: 't-1', delta: '{"city":"x"}', subagentRunId: 'sa-2' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_ARGS', toolCallId: 't-1', delta: '"paris"}', subagentRunId: 'sa-1' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_END', toolCallId: 't-1', subagentRunId: 'sa-1' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_END', toolCallId: 't-1', subagentRunId: 'sa-2' }), store); + const calls1 = store.activities().get('sa-1')?.content()['toolCalls'] as Array>; + const calls2 = store.activities().get('sa-2')?.content()['toolCalls'] as Array>; + expect(calls1[0]).toMatchObject({ id: 't-1', status: 'complete', args: { q: 'paris' } }); + expect(calls2[0]).toMatchObject({ id: 't-1', status: 'complete', args: { city: 'x' } }); + }); + + it('RUN_STARTED drops a dangling parent args buffer so a same-id call in the next run parses cleanly', () => { + const store = makeStore(); + // A run that dies mid-stream: ARGS fragment arrives, TOOL_CALL_END never does. + reduceEvent(ev({ type: 'TOOL_CALL_START', toolCallId: 't-1', toolCallName: 'web_search' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_ARGS', toolCallId: 't-1', delta: '{"a":' }), store); + reduceEvent(ev({ type: 'RUN_STARTED' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_START', toolCallId: 't-1', toolCallName: 'web_search' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_ARGS', toolCallId: 't-1', delta: '{"b":1}' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_END', toolCallId: 't-1' }), store); + const call = store.toolCalls().find((t) => t.id === 't-1'); + expect(call?.args).toEqual({ b: 1 }); + expect(call?.status).toBe('complete'); + }); + it('unattributed events behave exactly as before (regression)', () => { const store = makeStore(); reduceEvent(ev({ type: 'TEXT_MESSAGE_START', messageId: 'm-1', role: 'assistant' }), store); diff --git a/libs/ag-ui/src/lib/reducer.ts b/libs/ag-ui/src/lib/reducer.ts index b9ff7f66b..1c96a9281 100644 --- a/libs/ag-ui/src/lib/reducer.ts +++ b/libs/ag-ui/src/lib/reducer.ts @@ -155,6 +155,10 @@ export function reduceEvent(event: BaseEvent, store: ReducerStore): void { store.interrupt.set(undefined); store.customEvents.set([]); store.activities.set(new Map()); + // A run boundary is the reset point for in-flight args accumulation: + // a TOOL_CALL_ARGS fragment whose TOOL_CALL_END never arrived (aborted + // or errored run) must not prefix a same-id call in the next run. + store.argsBuffers?.clear(); return; } case 'RUN_FINISHED': { @@ -609,6 +613,10 @@ const SUBAGENT_ROUTED_TYPES = new Set([ 'TOOL_CALL_START', 'TOOL_CALL_ARGS', 'TOOL_CALL_END', 'TOOL_CALL_RESULT', ]); +function subagentArgsBufferKey(subagentRunId: string, toolCallId: string): string { + return `subagent:${subagentRunId}:${toolCallId}`; +} + /** Get-or-create the ActivityEntry for a subagent run, keyed by * subagentRunId — mirrors ACTIVITY_SNAPSHOT's creation branch exactly * (same generation allocation, same activities-map replace idiom) so @@ -652,13 +660,16 @@ function routeSubagentContentEvent(subagentRunId: string, event: BaseEvent, stor if (event.type === 'TOOL_CALL_ARGS') { // Same accumulated-buffer rule as the parent handler: deltas are JSON // fragments; parse the accumulation, keep last-good args. + // Keyed by subagent run AND tool-call id: two children reusing the same + // toolCallId (real for providers that mint per-child ids) must never + // interleave fragments into one buffer. const buffers = (store.argsBuffers ??= new Map()); - const key = `subagent:${e['toolCallId']}`; + const key = subagentArgsBufferKey(subagentRunId, e['toolCallId'] as string); const buffer = (buffers.get(key) ?? '') + ((e['delta'] as string) ?? ''); buffers.set(key, buffer); parsedArgs = tryParseArgs(buffer); } else if (event.type === 'TOOL_CALL_END') { - store.argsBuffers?.delete(`subagent:${e['toolCallId']}`); + store.argsBuffers?.delete(subagentArgsBufferKey(subagentRunId, e['toolCallId'] as string)); } entry.content.update((c) => { diff --git a/libs/e2e-harness/src/aimock-mode.spec.ts b/libs/e2e-harness/src/aimock-mode.spec.ts index b5bf0b778..183b396a8 100644 --- a/libs/e2e-harness/src/aimock-mode.spec.ts +++ b/libs/e2e-harness/src/aimock-mode.spec.ts @@ -1,11 +1,33 @@ // SPDX-License-Identifier: MIT -import { describe, it, expect, afterEach } from 'vitest'; +import { describe, it, expect, afterEach, vi } from 'vitest'; import { resolveAimockLaunch } from './aimock-mode'; describe('resolveAimockLaunch', () => { const saved = { ...process.env }; afterEach(() => { process.env = { ...saved }; + vi.restoreAllMocks(); + }); + + it('unrecognized AIMOCK_MODE warns once and falls back to replay', () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + process.env['AIMOCK_MODE'] = 'Record'; // case-sensitive: not 'record' + process.env['OPENAI_API_KEY'] = 'sk-real'; + const launch = resolveAimockLaunch('/repo/x/fixtures'); + expect(launch.startOptions).toEqual({ mode: 'replay', fixturePath: '/repo/x/fixtures' }); + expect(launch.openaiApiKey).toBe('test-not-used'); + expect(warn).toHaveBeenCalledTimes(1); + expect(warn.mock.calls[0][0]).toContain('Record'); + expect(warn.mock.calls[0][0]).toContain('replay'); + }); + + it('recognized modes do not warn', () => { + const warn = vi.spyOn(console, 'warn').mockImplementation(() => undefined); + process.env['AIMOCK_MODE'] = 'replay'; + resolveAimockLaunch('/repo/x/fixtures'); + delete process.env['AIMOCK_MODE']; + resolveAimockLaunch('/repo/x/fixtures'); + expect(warn).not.toHaveBeenCalled(); }); it('defaults to replay against the fixtures dir with a placeholder key', () => { diff --git a/libs/e2e-harness/src/aimock-mode.ts b/libs/e2e-harness/src/aimock-mode.ts index 41a4c27e2..3497614de 100644 --- a/libs/e2e-harness/src/aimock-mode.ts +++ b/libs/e2e-harness/src/aimock-mode.ts @@ -16,7 +16,13 @@ export interface AimockLaunch { * `.aimock-recordings/` next to the fixtures dir). */ export function resolveAimockLaunch(fixturesDir: string): AimockLaunch { - if (process.env['AIMOCK_MODE'] === 'record') { + const mode = process.env['AIMOCK_MODE']; + if (mode !== undefined && mode !== 'record' && mode !== 'replay') { + // A typo (`Record`, `recording`, ...) would otherwise silently replay + // against stale fixtures while the operator believes they are recording. + console.warn(`[aimock-harness] unrecognized AIMOCK_MODE="${mode}" — falling back to replay`); + } + if (mode === 'record') { if (!process.env['OPENAI_API_KEY']) { throw new Error( '[aimock-harness] AIMOCK_MODE=record requires OPENAI_API_KEY — the record proxy forwards requests to the live provider.', diff --git a/libs/langgraph/src/lib/agent.fn.spec.ts b/libs/langgraph/src/lib/agent.fn.spec.ts index 9ee35d649..1dde835bf 100644 --- a/libs/langgraph/src/lib/agent.fn.spec.ts +++ b/libs/langgraph/src/lib/agent.fn.spec.ts @@ -2,7 +2,8 @@ import { describe, it, expect, vi, beforeEach } from 'vitest'; import { TestBed } from '@angular/core/testing'; import { signal } from '@angular/core'; import type { AIMessage as CoreAIMessage } from '@langchain/core/messages'; -import { agent } from './agent.fn'; +import { agent, resolveSubagentsByMessage } from './agent.fn'; +import type { SubagentStreamRef } from './agent.types'; import { MockAgentTransport } from './transport/mock-stream.transport'; import type { AgentTransport, StreamEvent } from './agent.types'; import type { ThreadState } from '@langchain/langgraph-sdk'; @@ -894,6 +895,36 @@ describe('agent', () => { expect(ref.getSubagent('missing')).toBeUndefined(); }); + it('getSubagentsByMessage resolves through the Subagent.toolCallId field, not the map key', () => { + // The neutral contract anchors a subagent on `Subagent.toolCallId`; the + // map KEY is an adapter detail (LangGraph keys by the id today, AG-UI + // keys activities by `-sub`). A key that diverges from the + // field must not hide the subagent from the message lookup. + const sub = (toolCallId: string, name: string): SubagentStreamRef => ({ + toolCallId, + name, + status: signal<'pending' | 'running' | 'complete' | 'error'>('running'), + values: signal>({}), + messages: signal([]), + }); + const subagents = new Map([ + ['call_x-sub', sub('call_x', 'researcher')], + ['call_y-sub', sub('call_y', 'reviewer')], + ['unrelated-sub', sub('call_z', 'other')], + ]); + const msg = { + id: 'ai-1', + type: 'ai', + content: '', + tool_calls: [ + { id: 'call_x', name: 'task', args: {} }, + { id: 'call_y', name: 'task', args: {} }, + ], + } as unknown as CoreAIMessage; + + expect(resolveSubagentsByMessage(msg, subagents).map(sa => sa.name)).toEqual(['researcher', 'reviewer']); + }); + it('events$ is an Observable-like with .subscribe', () => { const transport = new MockAgentTransport(); const ref = withInjectionContext(() => diff --git a/libs/langgraph/src/lib/agent.fn.ts b/libs/langgraph/src/lib/agent.fn.ts index 96208488e..0cfc06cbd 100644 --- a/libs/langgraph/src/lib/agent.fn.ts +++ b/libs/langgraph/src/lib/agent.fn.ts @@ -628,13 +628,7 @@ export function agent< getSubagent: (toolCallId) => subagentsSig().get(toolCallId), getSubagentsByType: (type) => [...subagentsSig().values()].filter(sa => sa.name === type), - getSubagentsByMessage: (msg) => { - const ids = getToolCallIds(msg); - const subagents = subagentsSig(); - return ids - .map(id => subagents.get(id)) - .filter((subagent): subagent is SubagentStreamRef => subagent != null); - }, + getSubagentsByMessage: (msg) => resolveSubagentsByMessage(msg, subagentsSig()), customEvents: customSig, branch: branchSig, setBranch: (b) => branch$.next(b), @@ -825,6 +819,28 @@ function toSubagent( }; } +/** + * Resolve the subagents spawned by an AI message's tool calls. + * + * Anchors on the contract field `SubagentStreamRef.toolCallId`, not the map + * KEY: the key is an adapter detail (LangGraph keys by the tool-call id, + * AG-UI keys activities by `-sub`) and only the field is + * guaranteed to equal the id on the message's `tool_calls` — the same rule + * `` applies when it re-indexes `agent.subagents()`. + * + * @internal Exported for unit tests only — not part of the public API. + */ +export function resolveSubagentsByMessage( + msg: CoreAIMessage, + subagents: ReadonlyMap, +): SubagentStreamRef[] { + const byToolCallId = new Map(); + subagents.forEach(sa => byToolCallId.set(sa.toolCallId, sa)); + return getToolCallIds(msg) + .map(id => byToolCallId.get(id)) + .filter((subagent): subagent is SubagentStreamRef => subagent != null); +} + function getToolCallIds(msg: CoreAIMessage): string[] { const raw = msg as unknown as Record; const toolCalls = raw['tool_calls'];