diff --git a/packages/a2a-server/src/agent/task.test.ts b/packages/a2a-server/src/agent/task.test.ts index 217b2105b4..513867f4e2 100644 --- a/packages/a2a-server/src/agent/task.test.ts +++ b/packages/a2a-server/src/agent/task.test.ts @@ -56,4 +56,42 @@ describe('Task', () => { expect(requests).toEqual(originalRequests); }); + + describe('acceptAgentMessage', () => { + it('should set currentTraceId when event has traceId', async () => { + const mockConfig = createMockConfig(); + const mockEventBus: ExecutionEventBus = { + publish: vi.fn(), + on: vi.fn(), + off: vi.fn(), + once: vi.fn(), + removeAllListeners: vi.fn(), + finished: vi.fn(), + }; + + // @ts-expect-error - Calling private constructor for test purposes. + const task = new Task( + 'task-id', + 'context-id', + mockConfig as Config, + mockEventBus, + ); + + const event = { + type: 'content', + value: 'test', + traceId: 'test-trace-id', + }; + + await task.acceptAgentMessage(event); + + expect(mockEventBus.publish).toHaveBeenCalledWith( + expect.objectContaining({ + metadata: expect.objectContaining({ + traceId: 'test-trace-id', + }), + }), + ); + }); + }); }); diff --git a/packages/a2a-server/src/agent/task.ts b/packages/a2a-server/src/agent/task.ts index 18391d793e..a7b0e288c9 100644 --- a/packages/a2a-server/src/agent/task.ts +++ b/packages/a2a-server/src/agent/task.ts @@ -220,12 +220,14 @@ export class Task { final = false, timestamp?: string, metadataError?: string, + traceId?: string, ): TaskStatusUpdateEvent { const metadata: { coderAgent: CoderAgentMessage; model: string; userTier?: UserTierId; error?: string; + traceId?: string; } = { coderAgent: coderAgentMessage, model: this.config.getModel(), @@ -236,6 +238,10 @@ export class Task { metadata.error = metadataError; } + if (traceId) { + metadata.traceId = traceId; + } + return { kind: 'status-update', taskId: this.id, @@ -257,6 +263,7 @@ export class Task { messageParts?: Part[], // For more complex messages final = false, metadataError?: string, + traceId?: string, ): void { this.taskState = newState; let message: Message | undefined; @@ -281,6 +288,7 @@ export class Task { final, undefined, metadataError, + traceId, ); this.eventBus?.publish(event); } @@ -582,10 +590,13 @@ export class Task { const stateChange: StateChange = { kind: CoderAgentEvent.StateChangeEvent, }; + const traceId = + 'traceId' in event && event.traceId ? event.traceId : undefined; + switch (event.type) { case GeminiEventType.Content: logger.info('[Task] Sending agent message content...'); - this._sendTextContent(event.value); + this._sendTextContent(event.value, traceId); break; case GeminiEventType.ToolCallRequest: // This is now handled by the agent loop, which collects all requests @@ -624,11 +635,13 @@ export class Task { 'Task cancelled by user', undefined, true, + undefined, + traceId, ); break; case GeminiEventType.Thought: logger.info('[Task] Sending agent thought...'); - this._sendThought(event.value); + this._sendThought(event.value, traceId); break; case GeminiEventType.ChatCompressed: break; @@ -658,6 +671,7 @@ export class Task { undefined, false, errMessage, + traceId, ); break; } @@ -915,7 +929,7 @@ export class Task { } } - _sendTextContent(content: string): void { + _sendTextContent(content: string, traceId?: string): void { if (content === '') { return; } @@ -930,11 +944,14 @@ export class Task { textContent, message, false, + undefined, + undefined, + traceId, ), ); } - _sendThought(content: ThoughtSummary): void { + _sendThought(content: ThoughtSummary, traceId?: string): void { if (!content.subject && !content.description) { return; } @@ -956,7 +973,15 @@ export class Task { kind: CoderAgentEvent.ThoughtEvent, }; this.eventBus?.publish( - this._createStatusUpdateEvent(this.taskState, thought, message, false), + this._createStatusUpdateEvent( + this.taskState, + thought, + message, + false, + undefined, + undefined, + traceId, + ), ); } } diff --git a/packages/a2a-server/src/http/app.test.ts b/packages/a2a-server/src/http/app.test.ts index 76a4873f0a..787fc86050 100644 --- a/packages/a2a-server/src/http/app.test.ts +++ b/packages/a2a-server/src/http/app.test.ts @@ -624,4 +624,32 @@ describe('E2E Tests', () => { assertUniqueFinalEventIsLast(events); expect(events.length).toBe(10); }); + + it('should include traceId in status updates when available', async () => { + const traceId = 'test-trace-id'; + sendMessageStreamSpy.mockImplementation(async function* () { + yield* [ + { type: 'content', value: 'Hello', traceId }, + { type: 'thought', value: { subject: 'Thinking...' }, traceId }, + ]; + }); + + const agent = request.agent(app); + const res = await agent + .post('/') + .send(createStreamMessageRequest('hello', 'a2a-trace-id-test')) + .set('Content-Type', 'application/json') + .expect(200); + + const events = streamToSSEEvents(res.text); + + // The first two events are task-creation and working status + const textContentEvent = events[2].result as TaskStatusUpdateEvent; + expect(textContentEvent.kind).toBe('status-update'); + expect(textContentEvent.metadata?.['traceId']).toBe(traceId); + + const thoughtEvent = events[3].result as TaskStatusUpdateEvent; + expect(thoughtEvent.kind).toBe('status-update'); + expect(thoughtEvent.metadata?.['traceId']).toBe(traceId); + }); }); diff --git a/packages/core/src/code_assist/converter.test.ts b/packages/core/src/code_assist/converter.test.ts index c36c8424af..17dba1e4da 100644 --- a/packages/core/src/code_assist/converter.test.ts +++ b/packages/core/src/code_assist/converter.test.ts @@ -310,6 +310,27 @@ describe('converter', () => { const genaiRes = fromGenerateContentResponse(codeAssistRes); expect(genaiRes.modelVersion).toEqual('gemini-2.5-pro'); }); + + it('should handle traceId', () => { + const codeAssistRes: CaGenerateContentResponse = { + response: { + candidates: [], + }, + traceId: 'my-trace-id', + }; + const genaiRes = fromGenerateContentResponse(codeAssistRes); + expect(genaiRes.responseId).toEqual('my-trace-id'); + }); + + it('should handle missing traceId', () => { + const codeAssistRes: CaGenerateContentResponse = { + response: { + candidates: [], + }, + }; + const genaiRes = fromGenerateContentResponse(codeAssistRes); + expect(genaiRes.responseId).toBeUndefined(); + }); }); describe('toContents', () => { diff --git a/packages/core/src/code_assist/converter.ts b/packages/core/src/code_assist/converter.ts index 78e743136f..2b8b0a3a33 100644 --- a/packages/core/src/code_assist/converter.ts +++ b/packages/core/src/code_assist/converter.ts @@ -73,6 +73,7 @@ interface VertexGenerationConfig { export interface CaGenerateContentResponse { response: VertexGenerateContentResponse; + traceId?: string; } interface VertexGenerateContentResponse { @@ -139,6 +140,7 @@ export function fromGenerateContentResponse( out.promptFeedback = inres.promptFeedback; out.usageMetadata = inres.usageMetadata; out.modelVersion = inres.modelVersion; + out.responseId = res.traceId; return out; } diff --git a/packages/core/src/core/turn.test.ts b/packages/core/src/core/turn.test.ts index 97e7195d30..efe4761ce5 100644 --- a/packages/core/src/core/turn.test.ts +++ b/packages/core/src/core/turn.test.ts @@ -784,6 +784,68 @@ describe('Turn', () => { { type: GeminiEventType.Content, value: 'Success' }, ]); }); + + it('should yield content events with traceId', async () => { + const mockResponseStream = (async function* () { + yield { + type: StreamEventType.CHUNK, + value: { + candidates: [{ content: { parts: [{ text: 'Hello' }] } }], + responseId: 'trace-123', + } as GenerateContentResponse, + }; + })(); + mockSendMessageStream.mockResolvedValue(mockResponseStream); + + const events = []; + for await (const event of turn.run( + 'test-model', + [{ text: 'Hi' }], + new AbortController().signal, + )) { + events.push(event); + } + + expect(events).toEqual([ + { type: GeminiEventType.Content, value: 'Hello', traceId: 'trace-123' }, + ]); + }); + + it('should yield thought events with traceId', async () => { + const mockResponseStream = (async function* () { + yield { + type: StreamEventType.CHUNK, + value: { + candidates: [ + { + content: { + parts: [{ text: '[Thought: thinking]', thought: 'thinking' }], + }, + }, + ], + responseId: 'trace-456', + } as unknown as GenerateContentResponse, + }; + })(); + mockSendMessageStream.mockResolvedValue(mockResponseStream); + + const events = []; + for await (const event of turn.run( + 'test-model', + [{ text: 'Hi' }], + new AbortController().signal, + )) { + events.push(event); + } + + expect(events).toEqual([ + { + type: GeminiEventType.Thought, + value: { subject: '', description: '[Thought: thinking]' }, + traceId: 'trace-456', + }, + ]); + }); }); describe('getDebugResponses', () => { diff --git a/packages/core/src/core/turn.ts b/packages/core/src/core/turn.ts index fc6772bf3e..681a7e9be2 100644 --- a/packages/core/src/core/turn.ts +++ b/packages/core/src/core/turn.ts @@ -120,11 +120,13 @@ export interface ServerToolCallConfirmationDetails { export type ServerGeminiContentEvent = { type: GeminiEventType.Content; value: string; + traceId?: string; }; export type ServerGeminiThoughtEvent = { type: GeminiEventType.Thought; value: ThoughtSummary; + traceId?: string; }; export type ServerGeminiToolCallRequestEvent = { @@ -261,19 +263,22 @@ export class Turn { this.debugResponses.push(resp); + const traceId = resp.responseId; + const thoughtPart = resp.candidates?.[0]?.content?.parts?.[0]; if (thoughtPart?.thought) { const thought = parseThought(thoughtPart.text ?? ''); yield { type: GeminiEventType.Thought, value: thought, + traceId, }; continue; } const text = getResponseText(resp); if (text) { - yield { type: GeminiEventType.Content, value: text }; + yield { type: GeminiEventType.Content, value: text, traceId }; } // Handle function calls (requesting tool execution)