Merge pull request #1327 from QwenLM/feat/vscode-ide-companion-context-left

context left on vscode ide companion

Author: tanzhenxin

packages/cli/src/acp-integration/acpAgent.ts

@@ -19,6 +19,7 @@ import {
   type Config,
   type ConversationRecord,
   type DeviceAuthorizationData,
+  tokenLimit,
 } from '@qwen-code/qwen-code-core';
 import type { ApprovalModeValue } from './schema.js';
 import * as acp from './acp.js';
@@ -165,9 +166,30 @@ class GeminiAgent {
     this.setupFileSystem(config);
 
     const session = await this.createAndStoreSession(config);
+    const configuredModel = (
+      config.getModel() ||
+      this.config.getModel() ||
+      ''
+    ).trim();
+    const modelId = configuredModel || 'default';
+    const modelName = configuredModel || modelId;
 
     return {
       sessionId: session.getId(),
+      models: {
+        currentModelId: modelId,
+        availableModels: [
+          {
+            modelId,
+            name: modelName,
+            description: null,
+            _meta: {
+              contextLimit: tokenLimit(modelId),
+            },
+          },
+        ],
+        _meta: null,
+      },
     };
   }
 

packages/cli/src/acp-integration/schema.ts

@@ -93,6 +93,7 @@ export type ModeInfo = z.infer<typeof modeInfoSchema>;
 export type ModesData = z.infer<typeof modesDataSchema>;
 
 export type AgentInfo = z.infer<typeof agentInfoSchema>;
+export type ModelInfo = z.infer<typeof modelInfoSchema>;
 
 export type PromptCapabilities = z.infer<typeof promptCapabilitiesSchema>;
 
@@ -254,8 +255,26 @@ export const authenticateUpdateSchema = z.object({
 
 export type AuthenticateUpdate = z.infer<typeof authenticateUpdateSchema>;
 
+export const acpMetaSchema = z.record(z.unknown()).nullable().optional();
+
+export const modelIdSchema = z.string();
+
+export const modelInfoSchema = z.object({
+  _meta: acpMetaSchema,
+  description: z.string().nullable().optional(),
+  modelId: modelIdSchema,
+  name: z.string(),
+});
+
+export const sessionModelStateSchema = z.object({
+  _meta: acpMetaSchema,
+  availableModels: z.array(modelInfoSchema),
+  currentModelId: modelIdSchema,
+});
+
 export const newSessionResponseSchema = z.object({
   sessionId: z.string(),
+  models: sessionModelStateSchema,
 });
 
 export const loadSessionResponseSchema = z.null();
@@ -514,6 +533,13 @@ export const currentModeUpdateSchema = z.object({
 
 export type CurrentModeUpdate = z.infer<typeof currentModeUpdateSchema>;
 
+export const currentModelUpdateSchema = z.object({
+  sessionUpdate: z.literal('current_model_update'),
+  model: modelInfoSchema,
+});
+
+export type CurrentModelUpdate = z.infer<typeof currentModelUpdateSchema>;
+
 export const sessionUpdateSchema = z.union([
   z.object({
     content: contentBlockSchema,
@@ -555,6 +581,7 @@ export const sessionUpdateSchema = z.union([
     sessionUpdate: z.literal('plan'),
   }),
   currentModeUpdateSchema,
+  currentModelUpdateSchema,
   availableCommandsUpdateSchema,
 ]);
 

packages/vscode-ide-companion/src/services/acpConnection.ts

@@ -146,6 +146,8 @@ export class AcpConnection {
       console.error(
         `[ACP qwen] Process exited with code: ${code}, signal: ${signal}`,
       );
+      // Clear pending requests when process exits
+      this.pendingRequests.clear();
     });
 
     // Wait for process to start

packages/vscode-ide-companion/src/services/qwenAgentManager.ts

@@ -8,6 +8,7 @@ import type {
   AcpSessionUpdate,
   AcpPermissionRequest,
   AuthenticateUpdateNotification,
+  ModelInfo,
 } from '../types/acpTypes.js';
 import type { ApprovalModeValue } from '../types/approvalModeValueTypes.js';
 import { QwenSessionReader, type QwenSession } from './qwenSessionReader.js';
@@ -17,13 +18,15 @@ import type {
   PlanEntry,
   ToolCallUpdateData,
   QwenAgentCallbacks,
+  UsageStatsPayload,
 } from '../types/chatTypes.js';
 import {
   QwenConnectionHandler,
   type QwenConnectionResult,
 } from '../services/qwenConnectionHandler.js';
 import { QwenSessionUpdateHandler } from './qwenSessionUpdateHandler.js';
 import { authMethod } from '../types/acpTypes.js';
+import { extractModelInfoFromNewSessionResult } from '../utils/acpModelInfo.js';
 import { isAuthenticationRequiredError } from '../utils/authErrors.js';
 import { handleAuthenticateUpdate } from '../utils/authNotificationHandler.js';
 
@@ -195,12 +198,16 @@ export class QwenAgentManager {
     options?: AgentConnectOptions,
   ): Promise<QwenConnectionResult> {
     this.currentWorkingDir = workingDir;
-    return this.connectionHandler.connect(
+    const res = await this.connectionHandler.connect(
       this.connection,
       workingDir,
       cliEntryPath,
       options,
     );
+    if (res.modelInfo && this.callbacks.onModelInfo) {
+      this.callbacks.onModelInfo(res.modelInfo);
+    }
+    return res;
   }
 
   /**
@@ -1091,9 +1098,10 @@ export class QwenAgentManager {
 
     this.sessionCreateInFlight = (async () => {
       try {
+        let newSessionResult: unknown;
         // Try to create a new ACP session. If Qwen asks for auth, let it handle authentication.
         try {
-          await this.connection.newSession(workingDir);
+          newSessionResult = await this.connection.newSession(workingDir);
         } catch (err) {
           const requiresAuth = isAuthenticationRequiredError(err);
 
@@ -1115,7 +1123,7 @@ export class QwenAgentManager {
               );
               // Add a slight delay to ensure auth state is settled
               await new Promise((resolve) => setTimeout(resolve, 300));
-              await this.connection.newSession(workingDir);
+              newSessionResult = await this.connection.newSession(workingDir);
             } catch (reauthErr) {
               console.error(
                 '[QwenAgentManager] Re-authentication failed:',
@@ -1127,6 +1135,13 @@ export class QwenAgentManager {
             throw err;
           }
         }
+
+        const modelInfo =
+          extractModelInfoFromNewSessionResult(newSessionResult);
+        if (modelInfo && this.callbacks.onModelInfo) {
+          this.callbacks.onModelInfo(modelInfo);
+        }
+
         const newSessionId = this.connection.currentSessionId;
         console.log(
           '[QwenAgentManager] New session created with ID:',
@@ -1257,6 +1272,22 @@ export class QwenAgentManager {
     this.sessionUpdateHandler.updateCallbacks(this.callbacks);
   }
 
+  /**
+   * Register callback for usage metadata updates
+   */
+  onUsageUpdate(callback: (stats: UsageStatsPayload) => void): void {
+    this.callbacks.onUsageUpdate = callback;
+    this.sessionUpdateHandler.updateCallbacks(this.callbacks);
+  }
+
+  /**
+   * Register callback for model info updates
+   */
+  onModelInfo(callback: (info: ModelInfo) => void): void {
+    this.callbacks.onModelInfo = callback;
+    this.sessionUpdateHandler.updateCallbacks(this.callbacks);
+  }
+
   /**
    * Disconnect
    */

packages/vscode-ide-companion/src/services/qwenConnectionHandler.ts

@@ -13,10 +13,13 @@
 import type { AcpConnection } from './acpConnection.js';
 import { isAuthenticationRequiredError } from '../utils/authErrors.js';
 import { authMethod } from '../types/acpTypes.js';
+import { extractModelInfoFromNewSessionResult } from '../utils/acpModelInfo.js';
+import type { ModelInfo } from '../types/acpTypes.js';
 
 export interface QwenConnectionResult {
   sessionCreated: boolean;
   requiresAuth: boolean;
+  modelInfo?: ModelInfo;
 }
 
 /**
@@ -44,6 +47,7 @@ export class QwenConnectionHandler {
     const autoAuthenticate = options?.autoAuthenticate ?? true;
     let sessionCreated = false;
     let requiresAuth = false;
+    let modelInfo: ModelInfo | undefined;
 
     // Build extra CLI arguments (only essential parameters)
     const extraArgs: string[] = [];
@@ -66,13 +70,15 @@ export class QwenConnectionHandler {
         console.log(
           '[QwenAgentManager] Creating new session (letting CLI handle authentication)...',
         );
-        await this.newSessionWithRetry(
+        const newSessionResult = await this.newSessionWithRetry(
           connection,
           workingDir,
           3,
           authMethod,
           autoAuthenticate,
         );
+        modelInfo =
+          extractModelInfoFromNewSessionResult(newSessionResult) || undefined;
         console.log('[QwenAgentManager] New session created successfully');
         sessionCreated = true;
       } catch (sessionError) {
@@ -99,7 +105,7 @@ export class QwenConnectionHandler {
     console.log(`\n========================================`);
     console.log(`[QwenAgentManager] ✅ CONNECT() COMPLETED SUCCESSFULLY`);
     console.log(`========================================\n`);
-    return { sessionCreated, requiresAuth };
+    return { sessionCreated, requiresAuth, modelInfo };
   }
 
   /**
@@ -115,15 +121,15 @@ export class QwenConnectionHandler {
     maxRetries: number,
     authMethod: string,
     autoAuthenticate: boolean,
-  ): Promise<void> {
+  ): Promise<unknown> {
     for (let attempt = 1; attempt <= maxRetries; attempt++) {
       try {
         console.log(
           `[QwenAgentManager] Creating session (attempt ${attempt}/${maxRetries})...`,
         );
-        await connection.newSession(workingDir);
+        const res = await connection.newSession(workingDir);
         console.log('[QwenAgentManager] Session created successfully');
-        return;
+        return res;
       } catch (error) {
         const errorMessage =
           error instanceof Error ? error.message : String(error);
@@ -155,11 +161,11 @@ export class QwenConnectionHandler {
               '[QwenAgentManager] newSessionWithRetry Authentication successful',
             );
             // Retry immediately after successful auth
-            await connection.newSession(workingDir);
+            const res = await connection.newSession(workingDir);
             console.log(
               '[QwenAgentManager] Session created successfully after auth',
             );
-            return;
+            return res;
           } catch (authErr) {
             console.error(
               '[QwenAgentManager] Re-authentication failed:',
@@ -180,5 +186,7 @@ export class QwenConnectionHandler {
         await new Promise((resolve) => setTimeout(resolve, delay));
       }
     }
+
+    throw new Error('Session creation failed unexpectedly');
   }
 }

packages/vscode-ide-companion/src/services/qwenSessionUpdateHandler.ts

@@ -10,9 +10,12 @@
  * Handles session updates from ACP and dispatches them to appropriate callbacks
  */
 
-import type { AcpSessionUpdate } from '../types/acpTypes.js';
+import type { AcpSessionUpdate, SessionUpdateMeta } from '../types/acpTypes.js';
 import type { ApprovalModeValue } from '../types/approvalModeValueTypes.js';
-import type { QwenAgentCallbacks } from '../types/chatTypes.js';
+import type {
+  QwenAgentCallbacks,
+  UsageStatsPayload,
+} from '../types/chatTypes.js';
 
 /**
  * Qwen Session Update Handler class
@@ -57,6 +60,7 @@ export class QwenSessionUpdateHandler {
         if (update.content?.text && this.callbacks.onStreamChunk) {
           this.callbacks.onStreamChunk(update.content.text);
         }
+        this.emitUsageMeta(update._meta);
         break;
 
       case 'agent_thought_chunk':
@@ -71,6 +75,7 @@ export class QwenSessionUpdateHandler {
             this.callbacks.onStreamChunk(update.content.text);
           }
         }
+        this.emitUsageMeta(update._meta);
         break;
 
       case 'tool_call': {
@@ -160,4 +165,17 @@ export class QwenSessionUpdateHandler {
         break;
     }
   }
+
+  private emitUsageMeta(meta?: SessionUpdateMeta): void {
+    if (!meta || !this.callbacks.onUsageUpdate) {
+      return;
+    }
+
+    const payload: UsageStatsPayload = {
+      usage: meta.usage || undefined,
+      durationMs: meta.durationMs ?? undefined,
+    };
+
+    this.callbacks.onUsageUpdate(payload);
+  }
 }

packages/vscode-ide-companion/src/types/acpTypes.ts

@@ -48,6 +48,35 @@ export interface ContentBlock {
   uri?: string;
 }
 
+export interface UsageMetadata {
+  promptTokens?: number | null;
+  completionTokens?: number | null;
+  thoughtsTokens?: number | null;
+  totalTokens?: number | null;
+  cachedTokens?: number | null;
+}
+
+export interface SessionUpdateMeta {
+  usage?: UsageMetadata | null;
+  durationMs?: number | null;
+}
+
+export type AcpMeta = Record<string, unknown>;
+export type ModelId = string;
+
+export interface ModelInfo {
+  _meta?: AcpMeta | null;
+  description?: string | null;
+  modelId: ModelId;
+  name: string;
+}
+
+export interface SessionModelState {
+  _meta?: AcpMeta | null;
+  availableModels: ModelInfo[];
+  currentModelId: ModelId;
+}
+
 export interface UserMessageChunkUpdate extends BaseSessionUpdate {
   update: {
     sessionUpdate: 'user_message_chunk';
@@ -59,13 +88,15 @@ export interface AgentMessageChunkUpdate extends BaseSessionUpdate {
   update: {
     sessionUpdate: 'agent_message_chunk';
     content: ContentBlock;
+    _meta?: SessionUpdateMeta;
   };
 }
 
 export interface AgentThoughtChunkUpdate extends BaseSessionUpdate {
   update: {
     sessionUpdate: 'agent_thought_chunk';
     content: ContentBlock;
+    _meta?: SessionUpdateMeta;
   };
 }