fix: add missing trace info and cancellation events (#791)

* fix: add missing trace info and cancellation events

* fix: re-organize tool/request cancellation logging

packages/cli/src/ui/hooks/useGeminiStream.test.tsx

@@ -54,7 +54,11 @@ const MockedGeminiClientClass = vi.hoisted(() =>
 const MockedUserPromptEvent = vi.hoisted(() =>
   vi.fn().mockImplementation(() => {}),
 );
+const MockedApiCancelEvent = vi.hoisted(() =>
+  vi.fn().mockImplementation(() => {}),
+);
 const mockParseAndFormatApiError = vi.hoisted(() => vi.fn());
+const mockLogApiCancel = vi.hoisted(() => vi.fn());
 
 // Vision auto-switch mocks (hoisted)
 const mockHandleVisionSwitch = vi.hoisted(() =>
@@ -71,7 +75,9 @@ vi.mock('@qwen-code/qwen-code-core', async (importOriginal) => {
     GitService: vi.fn(),
     GeminiClient: MockedGeminiClientClass,
     UserPromptEvent: MockedUserPromptEvent,
+    ApiCancelEvent: MockedApiCancelEvent,
     parseAndFormatApiError: mockParseAndFormatApiError,
+    logApiCancel: mockLogApiCancel,
   };
 });
 

packages/cli/src/ui/hooks/useGeminiStream.ts

@@ -31,6 +31,8 @@ import {
   ConversationFinishedEvent,
   ApprovalMode,
   parseAndFormatApiError,
+  logApiCancel,
+  ApiCancelEvent,
 } from '@qwen-code/qwen-code-core';
 import { type Part, type PartListUnion, FinishReason } from '@google/genai';
 import type {
@@ -223,6 +225,16 @@ export const useGeminiStream = (
     turnCancelledRef.current = true;
     isSubmittingQueryRef.current = false;
     abortControllerRef.current?.abort();
+
+    // Log API cancellation
+    const prompt_id = config.getSessionId() + '########' + getPromptCount();
+    const cancellationEvent = new ApiCancelEvent(
+      config.getModel(),
+      prompt_id,
+      config.getContentGeneratorConfig()?.authType,
+    );
+    logApiCancel(config, cancellationEvent);
+
     if (pendingHistoryItemRef.current) {
       addItem(pendingHistoryItemRef.current, Date.now());
     }
@@ -242,6 +254,8 @@ export const useGeminiStream = (
     setPendingHistoryItem,
     onCancelSubmit,
     pendingHistoryItemRef,
+    config,
+    getPromptCount,
   ]);
 
   useKeypress(
@@ -448,6 +462,7 @@ export const useGeminiStream = (
       if (turnCancelledRef.current) {
         return;
       }
+
       if (pendingHistoryItemRef.current) {
         if (pendingHistoryItemRef.current.type === 'tool_group') {
           const updatedTools = pendingHistoryItemRef.current.tools.map(

packages/cli/src/zed-integration/zedIntegration.ts

@@ -349,6 +349,7 @@ class Session {
         function_name: fc.name ?? '',
         function_args: args,
         duration_ms: durationMs,
+        status: 'error',
         success: false,
         error: error.message,
         tool_type:
@@ -467,6 +468,7 @@ class Session {
         function_name: fc.name,
         function_args: args,
         duration_ms: durationMs,
+        status: 'success',
         success: true,
         prompt_id: promptId,
         tool_type:

packages/core/src/core/coreToolScheduler.ts

@@ -401,7 +401,7 @@ export class CoreToolScheduler {
             }
           }
 
-          return {
+          const cancelledCall = {
             request: currentCall.request,
             tool: toolInstance,
             invocation,
@@ -426,6 +426,8 @@ export class CoreToolScheduler {
             durationMs,
             outcome,
           } as CancelledToolCall;
+
+          return cancelledCall;
         }
         case 'validating':
           return {

packages/core/src/core/openaiContentGenerator/pipeline.ts

@@ -227,6 +227,12 @@ export class ContentGenerationPipeline {
         mergedResponse.usageMetadata = lastResponse.usageMetadata;
       }
 
+      // Copy other essential properties from the current response
+      mergedResponse.responseId = response.responseId;
+      mergedResponse.createTime = response.createTime;
+      mergedResponse.modelVersion = response.modelVersion;
+      mergedResponse.promptFeedback = response.promptFeedback;
+
       // Update the collected responses with the merged response
       collectedGeminiResponses[collectedGeminiResponses.length - 1] =
         mergedResponse;

packages/core/src/core/turn.ts

@@ -84,6 +84,7 @@ export interface ToolCallRequestInfo {
   args: Record<string, unknown>;
   isClientInitiated: boolean;
   prompt_id: string;
+  response_id?: string;
 }
 
 export interface ToolCallResponseInfo {
@@ -202,6 +203,7 @@ export class Turn {
   readonly pendingToolCalls: ToolCallRequestInfo[];
   private debugResponses: GenerateContentResponse[];
   finishReason: FinishReason | undefined;
+  private currentResponseId?: string;
 
   constructor(
     private readonly chat: GeminiChat,
@@ -247,6 +249,11 @@ export class Turn {
 
         this.debugResponses.push(resp);
 
+        // Track the current response ID for tool call correlation
+        if (resp.responseId) {
+          this.currentResponseId = resp.responseId;
+        }
+
         const thoughtPart = resp.candidates?.[0]?.content?.parts?.[0];
         if (thoughtPart?.thought) {
           // Thought always has a bold "subject" part enclosed in double asterisks
@@ -346,6 +353,7 @@ export class Turn {
       args,
       isClientInitiated: false,
       prompt_id: this.prompt_id,
+      response_id: this.currentResponseId,
     };
 
     this.pendingToolCalls.push(toolCallRequest);

packages/core/src/subagents/subagent.ts

@@ -381,6 +381,7 @@ export class SubAgentScope {
         let roundText = '';
         let lastUsage: GenerateContentResponseUsageMetadata | undefined =
           undefined;
+        let currentResponseId: string | undefined = undefined;
         for await (const streamEvent of responseStream) {
           if (abortController.signal.aborted) {
             this.terminateMode = SubagentTerminateMode.CANCELLED;
@@ -395,6 +396,10 @@ export class SubAgentScope {
           // Handle chunk events
           if (streamEvent.type === 'chunk') {
             const resp = streamEvent.value;
+            // Track the response ID for tool call correlation
+            if (resp.responseId) {
+              currentResponseId = resp.responseId;
+            }
             if (resp.functionCalls) functionCalls.push(...resp.functionCalls);
             const content = resp.candidates?.[0]?.content;
             const parts = content?.parts || [];
@@ -455,6 +460,7 @@ export class SubAgentScope {
             abortController,
             promptId,
             turnCounter,
+            currentResponseId,
           );
         } else {
           // No tool calls — treat this as the model's final answer.
@@ -543,6 +549,7 @@ export class SubAgentScope {
    * @param {FunctionCall[]} functionCalls - An array of `FunctionCall` objects to process.
    * @param {ToolRegistry} toolRegistry - The tool registry to look up and execute tools.
    * @param {AbortController} abortController - An `AbortController` to signal cancellation of tool executions.
+   * @param {string} responseId - Optional API response ID for correlation with tool calls.
    * @returns {Promise<Content[]>} A promise that resolves to an array of `Content` parts representing the tool responses,
    *          which are then used to update the chat history.
    */
@@ -551,6 +558,7 @@ export class SubAgentScope {
     abortController: AbortController,
     promptId: string,
     currentRound: number,
+    responseId?: string,
   ): Promise<Content[]> {
     const toolResponseParts: Part[] = [];
 
@@ -704,6 +712,7 @@ export class SubAgentScope {
         args,
         isClientInitiated: true,
         prompt_id: promptId,
+        response_id: responseId,
       };
 
       const description = this.getToolDescription(toolName, args);

packages/core/src/telemetry/constants.ts

@@ -10,6 +10,7 @@ export const EVENT_USER_PROMPT = 'qwen-code.user_prompt';
 export const EVENT_TOOL_CALL = 'qwen-code.tool_call';
 export const EVENT_API_REQUEST = 'qwen-code.api_request';
 export const EVENT_API_ERROR = 'qwen-code.api_error';
+export const EVENT_API_CANCEL = 'qwen-code.api_cancel';
 export const EVENT_API_RESPONSE = 'qwen-code.api_response';
 export const EVENT_CLI_CONFIG = 'qwen-code.config';
 export const EVENT_FLASH_FALLBACK = 'qwen-code.flash_fallback';

packages/core/src/telemetry/index.ts

@@ -17,6 +17,7 @@ export { SpanStatusCode, ValueType } from '@opentelemetry/api';
 export { SemanticAttributes } from '@opentelemetry/semantic-conventions';
 export {
   logApiError,
+  logApiCancel,
   logApiRequest,
   logApiResponse,
   logChatCompression,
@@ -35,6 +36,7 @@ export {
 } from './sdk.js';
 export {
   ApiErrorEvent,
+  ApiCancelEvent,
   ApiRequestEvent,
   ApiResponseEvent,
   ConversationFinishedEvent,
@@ -54,4 +56,5 @@ export type {
   TelemetryEvent,
 } from './types.js';
 export * from './uiTelemetry.js';
+export { QwenLogger } from './qwen-logger/qwen-logger.js';
 export { DEFAULT_OTLP_ENDPOINT, DEFAULT_TELEMETRY_TARGET };

packages/core/src/telemetry/loggers.test.ts

@@ -550,6 +550,7 @@ describe('loggers', () => {
             2,
           ),
           duration_ms: 100,
+          status: 'success',
           success: true,
           decision: ToolCallDecision.ACCEPT,
           prompt_id: 'prompt-id-1',
@@ -619,6 +620,7 @@ describe('loggers', () => {
             2,
           ),
           duration_ms: 100,
+          status: 'error',
           success: false,
           decision: ToolCallDecision.REJECT,
           prompt_id: 'prompt-id-2',
@@ -691,6 +693,7 @@ describe('loggers', () => {
             2,
           ),
           duration_ms: 100,
+          status: 'success',
           success: true,
           decision: ToolCallDecision.MODIFY,
           prompt_id: 'prompt-id-3',
@@ -762,6 +765,7 @@ describe('loggers', () => {
             2,
           ),
           duration_ms: 100,
+          status: 'success',
           success: true,
           prompt_id: 'prompt-id-4',
           tool_type: 'native',
@@ -834,6 +838,7 @@ describe('loggers', () => {
             2,
           ),
           duration_ms: 100,
+          status: 'error',
           success: false,
           error: 'test-error',
           'error.message': 'test-error',

packages/core/src/telemetry/loggers.ts

@@ -12,6 +12,7 @@ import { safeJsonStringify } from '../utils/safeJsonStringify.js';
 import { UserAccountManager } from '../utils/userAccountManager.js';
 import {
   EVENT_API_ERROR,
+  EVENT_API_CANCEL,
   EVENT_API_REQUEST,
   EVENT_API_RESPONSE,
   EVENT_CHAT_COMPRESSION,
@@ -45,6 +46,7 @@ import { QwenLogger } from './qwen-logger/qwen-logger.js';
 import { isTelemetrySdkInitialized } from './sdk.js';
 import type {
   ApiErrorEvent,
+  ApiCancelEvent,
   ApiRequestEvent,
   ApiResponseEvent,
   ChatCompressionEvent,
@@ -282,6 +284,32 @@ export function logApiError(config: Config, event: ApiErrorEvent): void {
   );
 }
 
+export function logApiCancel(config: Config, event: ApiCancelEvent): void {
+  const uiEvent = {
+    ...event,
+    'event.name': EVENT_API_CANCEL,
+    'event.timestamp': new Date().toISOString(),
+  } as UiEvent;
+  uiTelemetryService.addEvent(uiEvent);
+  QwenLogger.getInstance(config)?.logApiCancelEvent(event);
+  if (!isTelemetrySdkInitialized()) return;
+
+  const attributes: LogAttributes = {
+    ...getCommonAttributes(config),
+    ...event,
+    'event.name': EVENT_API_CANCEL,
+    'event.timestamp': new Date().toISOString(),
+    model_name: event.model,
+  };
+
+  const logger = logs.getLogger(SERVICE_NAME);
+  const logRecord: LogRecord = {
+    body: `API request cancelled for ${event.model}.`,
+    attributes,
+  };
+  logger.emit(logRecord);
+}
+
 export function logApiResponse(config: Config, event: ApiResponseEvent): void {
   const uiEvent = {
     ...event,

packages/core/src/telemetry/qwen-logger/qwen-logger.ts

@@ -15,6 +15,7 @@ import type {
   ApiRequestEvent,
   ApiResponseEvent,
   ApiErrorEvent,
+  ApiCancelEvent,
   FileOperationEvent,
   FlashFallbackEvent,
   LoopDetectedEvent,
@@ -411,6 +412,7 @@ export class QwenLogger {
       {
         properties: {
           prompt_id: event.prompt_id,
+          response_id: event.response_id,
         },
         snapshots: JSON.stringify({
           function_name: event.function_name,
@@ -427,6 +429,19 @@ export class QwenLogger {
     this.flushIfNeeded();
   }
 
+  logApiCancelEvent(event: ApiCancelEvent): void {
+    const rumEvent = this.createActionEvent('api', 'api_cancel', {
+      properties: {
+        model: event.model,
+        prompt_id: event.prompt_id,
+        auth_type: event.auth_type,
+      },
+    });
+
+    this.enqueueLogEvent(rumEvent);
+    this.flushIfNeeded();
+  }
+
   logFileOperationEvent(event: FileOperationEvent): void {
     const rumEvent = this.createActionEvent(
       'file_operation',

packages/core/src/telemetry/types.ts

@@ -127,11 +127,13 @@ export class ToolCallEvent implements BaseTelemetryEvent {
   function_name: string;
   function_args: Record<string, unknown>;
   duration_ms: number;
-  success: boolean;
+  status: 'success' | 'error' | 'cancelled';
+  success: boolean; // Keep for backward compatibility
   decision?: ToolCallDecision;
   error?: string;
   error_type?: string;
   prompt_id: string;
+  response_id?: string;
   tool_type: 'native' | 'mcp';
   // eslint-disable-next-line @typescript-eslint/no-explicit-any
   metadata?: { [key: string]: any };
@@ -142,13 +144,15 @@ export class ToolCallEvent implements BaseTelemetryEvent {
     this.function_name = call.request.name;
     this.function_args = call.request.args;
     this.duration_ms = call.durationMs ?? 0;
-    this.success = call.status === 'success';
+    this.status = call.status;
+    this.success = call.status === 'success'; // Keep for backward compatibility
     this.decision = call.outcome
       ? getDecisionFromOutcome(call.outcome)
       : undefined;
     this.error = call.response.error?.message;
     this.error_type = call.response.errorType;
     this.prompt_id = call.request.prompt_id;
+    this.response_id = call.request.response_id;
     this.tool_type =
       typeof call.tool !== 'undefined' && call.tool instanceof DiscoveredMCPTool
         ? 'mcp'
@@ -224,6 +228,22 @@ export class ApiErrorEvent implements BaseTelemetryEvent {
   }
 }
 
+export class ApiCancelEvent implements BaseTelemetryEvent {
+  'event.name': 'api_cancel';
+  'event.timestamp': string;
+  model: string;
+  prompt_id: string;
+  auth_type?: string;
+
+  constructor(model: string, prompt_id: string, auth_type?: string) {
+    this['event.name'] = 'api_cancel';
+    this['event.timestamp'] = new Date().toISOString();
+    this.model = model;
+    this.prompt_id = prompt_id;
+    this.auth_type = auth_type;
+  }
+}
+
 export class ApiResponseEvent implements BaseTelemetryEvent {
   'event.name': 'api_response';
   'event.timestamp': string; // ISO 8601
@@ -542,6 +562,7 @@ export type TelemetryEvent =
   | ToolCallEvent
   | ApiRequestEvent
   | ApiErrorEvent
+  | ApiCancelEvent
   | ApiResponseEvent
   | FlashFallbackEvent
   | LoopDetectedEvent

packages/core/src/telemetry/uiTelemetry.test.ts

@@ -15,6 +15,7 @@ import {
   EVENT_TOOL_CALL,
 } from './constants.js';
 import type {
+  CancelledToolCall,
   CompletedToolCall,
   ErroredToolCall,
   SuccessfulToolCall,
@@ -25,7 +26,7 @@ import { MockTool } from '../test-utils/tools.js';
 
 const createFakeCompletedToolCall = (
   name: string,
-  success: boolean,
+  success: boolean | 'cancelled',
   duration = 100,
   outcome?: ToolConfirmationOutcome,
   error?: Error,
@@ -39,7 +40,7 @@ const createFakeCompletedToolCall = (
   };
   const tool = new MockTool(name);
 
-  if (success) {
+  if (success === true) {
     return {
       status: 'success',
       request,
@@ -63,6 +64,30 @@ const createFakeCompletedToolCall = (
       durationMs: duration,
       outcome,
     } as SuccessfulToolCall;
+  } else if (success === 'cancelled') {
+    return {
+      status: 'cancelled',
+      request,
+      tool,
+      invocation: tool.build({ param: 'test' }),
+      response: {
+        callId: request.callId,
+        responseParts: [
+          {
+            functionResponse: {
+              id: request.callId,
+              name,
+              response: { error: 'Tool cancelled' },
+            },
+          },
+        ],
+        error: new Error('Tool cancelled'),
+        errorType: ToolErrorType.UNKNOWN,
+        resultDisplay: 'Cancelled!',
+      },
+      durationMs: duration,
+      outcome,
+    } as CancelledToolCall;
   } else {
     return {
       status: 'error',
@@ -411,6 +436,40 @@ describe('UiTelemetryService', () => {
       });
     });
 
+    it('should process a single cancelled ToolCallEvent', () => {
+      const toolCall = createFakeCompletedToolCall(
+        'test_tool',
+        'cancelled',
+        180,
+        ToolConfirmationOutcome.Cancel,
+      );
+      service.addEvent({
+        ...structuredClone(new ToolCallEvent(toolCall)),
+        'event.name': EVENT_TOOL_CALL,
+      } as ToolCallEvent & { 'event.name': typeof EVENT_TOOL_CALL });
+
+      const metrics = service.getMetrics();
+      const { tools } = metrics;
+
+      expect(tools.totalCalls).toBe(1);
+      expect(tools.totalSuccess).toBe(0);
+      expect(tools.totalFail).toBe(1);
+      expect(tools.totalDurationMs).toBe(180);
+      expect(tools.totalDecisions[ToolCallDecision.REJECT]).toBe(1);
+      expect(tools.byName['test_tool']).toEqual({
+        count: 1,
+        success: 0,
+        fail: 1,
+        durationMs: 180,
+        decisions: {
+          [ToolCallDecision.ACCEPT]: 0,
+          [ToolCallDecision.REJECT]: 1,
+          [ToolCallDecision.MODIFY]: 0,
+          [ToolCallDecision.AUTO_ACCEPT]: 0,
+        },
+      });
+    });
+
     it('should process a ToolCallEvent with modify decision', () => {
       const toolCall = createFakeCompletedToolCall(
         'test_tool',
@@ -637,6 +696,34 @@ describe('UiTelemetryService', () => {
       expect(service.getLastPromptTokenCount()).toBe(0);
       expect(spy).toHaveBeenCalledOnce();
     });
+
+    it('should correctly set status field for success/error/cancelled calls', () => {
+      const successCall = createFakeCompletedToolCall(
+        'success_tool',
+        true,
+        100,
+      );
+      const errorCall = createFakeCompletedToolCall('error_tool', false, 150);
+      const cancelledCall = createFakeCompletedToolCall(
+        'cancelled_tool',
+        'cancelled',
+        200,
+      );
+
+      const successEvent = new ToolCallEvent(successCall);
+      const errorEvent = new ToolCallEvent(errorCall);
+      const cancelledEvent = new ToolCallEvent(cancelledCall);
+
+      // Verify status field is correctly set
+      expect(successEvent.status).toBe('success');
+      expect(errorEvent.status).toBe('error');
+      expect(cancelledEvent.status).toBe('cancelled');
+
+      // Verify backward compatibility with success field
+      expect(successEvent.success).toBe(true);
+      expect(errorEvent.success).toBe(false);
+      expect(cancelledEvent.success).toBe(false);
+    });
   });
 
   describe('Tool Call Event with Line Count Metadata', () => {