fix: subagent hardening batch — eviction order, args-buffer scoping, field-based lookup, AIMOCK_MODE warning

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:

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"

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:

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"

libs/ag-ui/src/lib/reducer.subagent.spec.ts

@@ -145,6 +145,39 @@ describe('reduceEvent SUBAGENT_* lifecycle', () => {
     expect((content['state'] as Record<string, unknown>)['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<Record<string, unknown>>;
+    const calls2 = store.activities().get('sa-2')?.content()['toolCalls'] as Array<Record<string, unknown>>;
+    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);

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<string, string>());
-    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) => {

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', () => {

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.',

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 `<toolCallId>-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<Record<string, unknown>>({}),
+      messages: signal([]),
+    });
+    const subagents = new Map<string, SubagentStreamRef>([
+      ['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(() =>