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

* fix: add missing trace info and cancellation events

* fix: re-organize tool/request cancellation logging

Author: Mingholy

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

@@ -443,6 +443,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

@@ -85,6 +85,7 @@ export interface ToolCallRequestInfo {
   args: Record<string, unknown>;
   isClientInitiated: boolean;
   prompt_id: string;
+  response_id?: string;
 }
 
 export interface ToolCallResponseInfo {
@@ -203,6 +204,7 @@ export class Turn {
   readonly pendingToolCalls: ToolCallRequestInfo[];
   private debugResponses: GenerateContentResponse[];
   finishReason: FinishReason | undefined;
+  private currentResponseId?: string;
 
   constructor(
     private readonly chat: GeminiChat,
@@ -248,6 +250,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
@@ -347,6 +354,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',