diff --git a/libs/ag-ui/src/lib/reducer.subagent.spec.ts b/libs/ag-ui/src/lib/reducer.subagent.spec.ts index f35004eaf..ee44ed45e 100644 --- a/libs/ag-ui/src/lib/reducer.subagent.spec.ts +++ b/libs/ag-ui/src/lib/reducer.subagent.spec.ts @@ -94,6 +94,47 @@ describe('reduceEvent SUBAGENT_* lifecycle', () => { expect(store.toolCalls()).toHaveLength(0); }); + it('attributed TOOL_CALL_START links the toolCallId onto the child message via parentMessageId', () => { + const store = makeStore(); + reduceEvent(ev({ type: 'SUBAGENT_STARTED', subagentRunId: 'sa-1', name: 'researcher' }), store); + reduceEvent(ev({ type: 'TEXT_MESSAGE_START', messageId: 'sa-1-m1', role: 'assistant', subagentRunId: 'sa-1' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_START', toolCallId: 't-1', toolCallName: 'web_search', subagentRunId: 'sa-1', parentMessageId: 'sa-1-m1' }), store); + const content = store.activities().get('sa-1')!.content(); + const msgs = content['messages'] as Array>; + const calls = content['toolCalls'] as Array>; + expect(msgs.find((m) => m['id'] === 'sa-1-m1')).toMatchObject({ toolCallIds: ['t-1'] }); + expect(calls).toHaveLength(1); + expect(calls[0]).toMatchObject({ id: 't-1', name: 'web_search' }); + }); + + it('attributed TOOL_CALL_START with a parentMessageId for an unseen message creates a child message slot', () => { + const store = makeStore(); + reduceEvent(ev({ type: 'SUBAGENT_STARTED', subagentRunId: 'sa-1', name: 'researcher' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_START', toolCallId: 't-1', toolCallName: 'web_search', subagentRunId: 'sa-1', parentMessageId: 'sa-1-m1' }), store); + const msgs = store.activities().get('sa-1')!.content()['messages'] as Array>; + expect(msgs).toHaveLength(1); + expect(msgs[0]).toMatchObject({ id: 'sa-1-m1', role: 'assistant', content: '', toolCallIds: ['t-1'] }); + }); + + it('attributed TOOL_CALL_START without a parentMessageId attaches to the most recently opened child message', () => { + const store = makeStore(); + reduceEvent(ev({ type: 'SUBAGENT_STARTED', subagentRunId: 'sa-1', name: 'researcher' }), store); + reduceEvent(ev({ type: 'TEXT_MESSAGE_START', messageId: 'sa-1-m1', role: 'assistant', subagentRunId: 'sa-1' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_START', toolCallId: 't-1', toolCallName: 'web_search', subagentRunId: 'sa-1' }), store); + const msgs = store.activities().get('sa-1')!.content()['messages'] as Array>; + expect(msgs.find((m) => m['id'] === 'sa-1-m1')).toMatchObject({ toolCallIds: ['t-1'] }); + }); + + it('attributed TOOL_CALL_START without a parentMessageId and no open child message leaves messages untouched', () => { + const store = makeStore(); + reduceEvent(ev({ type: 'SUBAGENT_STARTED', subagentRunId: 'sa-1', name: 'researcher' }), store); + reduceEvent(ev({ type: 'TOOL_CALL_START', toolCallId: 't-1', toolCallName: 'web_search', subagentRunId: 'sa-1' }), store); + const content = store.activities().get('sa-1')!.content(); + expect(content['messages']).toEqual([]); + const calls = content['toolCalls'] as Array>; + expect(calls).toMatchObject([{ id: 't-1', name: 'web_search' }]); + }); + it('an attributed event before SUBAGENT_STARTED creates the entry instead of dropping (buffer-not-drop)', () => { const store = makeStore(); reduceEvent(ev({ type: 'TEXT_MESSAGE_START', messageId: 'm-1', role: 'assistant', subagentRunId: 'sa-late' }), store); diff --git a/libs/ag-ui/src/lib/reducer.ts b/libs/ag-ui/src/lib/reducer.ts index 1c96a9281..d64567a40 100644 --- a/libs/ag-ui/src/lib/reducer.ts +++ b/libs/ag-ui/src/lib/reducer.ts @@ -690,9 +690,33 @@ function routeSubagentContentEvent(subagentRunId: string, event: BaseEvent, stor } case 'TEXT_MESSAGE_END': return c; - case 'TOOL_CALL_START': + case 'TOOL_CALL_START': { toolCalls.push({ id: e['toolCallId'], name: e['toolCallName'], args: {}, status: 'running' }); - return { ...c, toolCalls }; + // Mirror the parent TOOL_CALL_START handler: link the call to its + // parent child message so chat-subagent-card can draw it from + // message.toolCallIds, same as the top-level transcript does. + const parentId = e['parentMessageId'] as string | undefined; + if (parentId) { + const idx = messages.findIndex((m) => m['id'] === parentId); + if (idx < 0) { + messages.push({ id: parentId, role: 'assistant', content: '', toolCallIds: [e['toolCallId']] }); + } else { + messages[idx] = { + ...messages[idx], + toolCallIds: [...((messages[idx]['toolCallIds'] as unknown[]) ?? []), e['toolCallId']], + }; + } + } else if (messages.length > 0) { + // No parentMessageId on the wire: attach to the most recently + // opened child message, mirroring TEXT_MESSAGE_START's append order. + const lastIdx = messages.length - 1; + messages[lastIdx] = { + ...messages[lastIdx], + toolCallIds: [...((messages[lastIdx]['toolCallIds'] as unknown[]) ?? []), e['toolCallId']], + }; + } + return { ...c, messages, toolCalls }; + } case 'TOOL_CALL_ARGS': return parsedArgs === undefined ? c diff --git a/libs/ag-ui/src/lib/to-agent.spec.ts b/libs/ag-ui/src/lib/to-agent.spec.ts index e35bd7058..7d2873b8b 100644 --- a/libs/ag-ui/src/lib/to-agent.spec.ts +++ b/libs/ag-ui/src/lib/to-agent.spec.ts @@ -1271,6 +1271,12 @@ describe('SUBAGENT_* lifecycle projection', () => { expect(sa?.messages()).toEqual([ { id: 'm-1', role: 'assistant', content: 'Checking flights', delivery: expect.objectContaining({ phase: 'streaming' }) }, ]); + + // Attributed TOOL_CALL_START with parentMessageId must link the call + // onto the child message's toolCallIds so chat-subagent-card can draw + // it — not just append to the entry's toolCalls[]. + source.emit({ type: 'TOOL_CALL_START', toolCallId: 't-1', toolCallName: 'web_search', subagentRunId: 'sa-1', parentMessageId: 'm-1' } as never); + expect(agent.subagents!().get('sa-1')?.messages()[0].toolCallIds).toEqual(['t-1']); }); it('a wrapper read before SUBAGENT_STARTED reflects the real identity once STARTED arrives (no stale name/toolCallId)', () => {