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