feat: Propagate traceId from code assist to response metadata (Fixes … (#11360)

Co-authored-by: owenofbrien <86964623+owenofbrien@users.noreply.github.com>

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',
+          }),
+        }),
+      );
+    });
+  });
 });

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,
+      ),
     );
   }
 }

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);
+  });
 });

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

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;
 }
 

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

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)