qwen-code/packages/vscode-ide-companion/src/acp/acpConnection.ts

/**
 * @license
 * Copyright 2025 Qwen Team
 * SPDX-License-Identifier: Apache-2.0
 */

import { JSONRPC_VERSION } from '../shared/acpTypes.js';
import type {
  AcpBackend,
  AcpMessage,
  AcpPermissionRequest,
  AcpResponse,
  AcpSessionUpdate,
} from '../shared/acpTypes.js';
import type { ChildProcess, SpawnOptions } from 'child_process';
import { spawn } from 'child_process';
import type {
  PendingRequest,
  AcpConnectionCallbacks,
} from './connectionTypes.js';
import { AcpMessageHandler } from './acpMessageHandler.js';
import { AcpSessionManager } from './acpSessionManager.js';

/**
 * ACP Connection Handler for VSCode Extension
 *
 * This class implements the client side of the ACP (Agent Communication Protocol).
 *
 * Implementation Status:
 *
 * Client Methods (Methods this class implements, called by CLI):
 * ✅ session/update - Handle session updates via onSessionUpdate callback
 * ✅ session/request_permission - Request user permission for tool execution
 * ✅ fs/read_text_file - Read file from workspace
 * ✅ fs/write_text_file - Write file to workspace
 *
 * Agent Methods (Methods CLI implements, called by this class):
 * ✅ initialize - Initialize ACP protocol connection
 * ✅ authenticate - Authenticate with selected auth method
 * ✅ session/new - Create new chat session
 * ✅ session/prompt - Send user message to agent
 * ✅ session/cancel - Cancel current generation
 * ✅ session/load - Load previous session
 * ✅ session/save - Save current session
 *
 * Custom Methods (Not in standard ACP):
 * ⚠️  session/list - List available sessions (custom extension)
 */
export class AcpConnection {
  private child: ChildProcess | null = null;
  private pendingRequests = new Map<number, PendingRequest<unknown>>();
  private nextRequestId = { value: 0 };
  private backend: AcpBackend | null = null;

  // 模块实例
  private messageHandler: AcpMessageHandler;
  private sessionManager: AcpSessionManager;

  // 回调函数
  onSessionUpdate: (data: AcpSessionUpdate) => void = () => {};
  onPermissionRequest: (data: AcpPermissionRequest) => Promise<{
    optionId: string;
  }> = () => Promise.resolve({ optionId: 'allow' });
  onEndTurn: () => void = () => {};

  constructor() {
    this.messageHandler = new AcpMessageHandler();
    this.sessionManager = new AcpSessionManager();
  }

  /**
   * Connect to ACP backend
   *
   * @param backend - Backend type
   * @param cliPath - CLI path
   * @param workingDir - Working directory
   * @param extraArgs - Extra command line arguments
   */
  async connect(
    backend: AcpBackend,
    cliPath: string,
    workingDir: string = process.cwd(),
    extraArgs: string[] = [],
  ): Promise<void> {
    if (this.child) {
      this.disconnect();
    }

    this.backend = backend;

    const isWindows = process.platform === 'win32';
    const env = { ...process.env };

    // If proxy is configured in extraArgs, also set it as environment variable
    // This ensures token refresh requests also use the proxy
    const proxyArg = extraArgs.find(
      (arg, i) => arg === '--proxy' && i + 1 < extraArgs.length,
    );
    if (proxyArg) {
      const proxyIndex = extraArgs.indexOf('--proxy');
      const proxyUrl = extraArgs[proxyIndex + 1];
      console.log('[ACP] Setting proxy environment variables:', proxyUrl);

      env.HTTP_PROXY = proxyUrl;
      env.HTTPS_PROXY = proxyUrl;
      env.http_proxy = proxyUrl;
      env.https_proxy = proxyUrl;
    }

    let spawnCommand: string;
    let spawnArgs: string[];

    if (cliPath.startsWith('npx ')) {
      const parts = cliPath.split(' ');
      spawnCommand = isWindows ? 'npx.cmd' : 'npx';
      spawnArgs = [...parts.slice(1), '--experimental-acp', ...extraArgs];
    } else {
      spawnCommand = cliPath;
      spawnArgs = ['--experimental-acp', ...extraArgs];
    }

    console.log('[ACP] Spawning command:', spawnCommand, spawnArgs.join(' '));

    const options: SpawnOptions = {
      cwd: workingDir,
      stdio: ['pipe', 'pipe', 'pipe'],
      env,
      shell: isWindows,
    };

    this.child = spawn(spawnCommand, spawnArgs, options);
    await this.setupChildProcessHandlers(backend);
  }

  /**
   * Set up child process handlers
   *
   * @param backend - Backend name
   */
  private async setupChildProcessHandlers(backend: string): Promise<void> {
    let spawnError: Error | null = null;

    this.child!.stderr?.on('data', (data) => {
      const message = data.toString();
      if (
        message.toLowerCase().includes('error') &&
        !message.includes('Loaded cached')
      ) {
        console.error(`[ACP ${backend}]:`, message);
      } else {
        console.log(`[ACP ${backend}]:`, message);
      }
    });

    this.child!.on('error', (error) => {
      spawnError = error;
    });

    this.child!.on('exit', (code, signal) => {
      console.error(
        `[ACP ${backend}] Process exited with code: ${code}, signal: ${signal}`,
      );
    });

    // Wait for process to start
    await new Promise((resolve) => setTimeout(resolve, 1000));

    if (spawnError) {
      throw spawnError;
    }

    if (!this.child || this.child.killed) {
      throw new Error(`${backend} ACP process failed to start`);
    }

    // Handle messages from ACP server
    let buffer = '';
    this.child.stdout?.on('data', (data) => {
      buffer += data.toString();
      const lines = buffer.split('\n');
      buffer = lines.pop() || '';

      for (const line of lines) {
        if (line.trim()) {
          try {
            const message = JSON.parse(line) as AcpMessage;
            console.log(
              '[ACP] <<< Received message:',
              JSON.stringify(message).substring(0, 500),
            );
            this.handleMessage(message);
          } catch (_error) {
            // Ignore non-JSON lines
            console.log(
              '[ACP] <<< Non-JSON line (ignored):',
              line.substring(0, 200),
            );
          }
        }
      }
    });

    // 初始化协议
    const res = await this.sessionManager.initialize(
      this.child,
      this.pendingRequests,
      this.nextRequestId,
    );

    console.log('[ACP] Initialization response:', res);
  }

  /**
   * Handle received messages
   *
   * @param message - ACP message
   */
  private handleMessage(message: AcpMessage): void {
    const callbacks: AcpConnectionCallbacks = {
      onSessionUpdate: this.onSessionUpdate,
      onPermissionRequest: this.onPermissionRequest,
      onEndTurn: this.onEndTurn,
    };

    // Handle message
    if ('method' in message) {
      // Request or notification
      this.messageHandler
        .handleIncomingRequest(message, callbacks)
        .then((result) => {
          if ('id' in message && typeof message.id === 'number') {
            this.messageHandler.sendResponseMessage(this.child, {
              jsonrpc: JSONRPC_VERSION,
              id: message.id,
              result,
            });
          }
        })
        .catch((error) => {
          if ('id' in message && typeof message.id === 'number') {
            this.messageHandler.sendResponseMessage(this.child, {
              jsonrpc: JSONRPC_VERSION,
              id: message.id,
              error: {
                code: -32603,
                message: error instanceof Error ? error.message : String(error),
              },
            });
          }
        });
    } else {
      // 响应
      this.messageHandler.handleMessage(
        message,
        this.pendingRequests,
        callbacks,
      );
    }
  }

  /**
   * Authenticate
   *
   * @param methodId - Authentication method ID
   * @returns Authentication response
   */
  async authenticate(methodId?: string): Promise<AcpResponse> {
    return this.sessionManager.authenticate(
      methodId,
      this.child,
      this.pendingRequests,
      this.nextRequestId,
    );
  }

  /**
   * Create new session
   *
   * @param cwd - Working directory
   * @returns New session response
   */
  async newSession(cwd: string = process.cwd()): Promise<AcpResponse> {
    return this.sessionManager.newSession(
      cwd,
      this.child,
      this.pendingRequests,
      this.nextRequestId,
    );
  }

  /**
   * Send prompt message
   *
   * @param prompt - Prompt content
   * @returns Response
   */
  async sendPrompt(prompt: string): Promise<AcpResponse> {
    return this.sessionManager.sendPrompt(
      prompt,
      this.child,
      this.pendingRequests,
      this.nextRequestId,
    );
  }

  /**
   * Load existing session
   *
   * @param sessionId - Session ID
   * @returns Load response
   */
  async loadSession(sessionId: string): Promise<AcpResponse> {
    return this.sessionManager.loadSession(
      sessionId,
      this.child,
      this.pendingRequests,
      this.nextRequestId,
    );
  }

  /**
   * Get session list
   *
   * @returns Session list response
   */
  async listSessions(): Promise<AcpResponse> {
    return this.sessionManager.listSessions(
      this.child,
      this.pendingRequests,
      this.nextRequestId,
    );
  }

  /**
   * Switch to specified session
   *
   * @param sessionId - Session ID
   * @returns Switch response
   */
  async switchSession(sessionId: string): Promise<AcpResponse> {
    return this.sessionManager.switchSession(sessionId, this.nextRequestId);
  }

  /**
   * Cancel current session prompt generation
   */
  async cancelSession(): Promise<void> {
    await this.sessionManager.cancelSession(this.child);
  }

  /**
   * Save current session
   *
   * @param tag - Save tag
   * @returns Save response
   */
  async saveSession(tag: string): Promise<AcpResponse> {
    return this.sessionManager.saveSession(
      tag,
      this.child,
      this.pendingRequests,
      this.nextRequestId,
    );
  }

  /**
   * Disconnect
   */
  disconnect(): void {
    if (this.child) {
      this.child.kill();
      this.child = null;
    }

    this.pendingRequests.clear();
    this.sessionManager.reset();
    this.backend = null;
  }

  /**
   * Check if connected
   */
  get isConnected(): boolean {
    return this.child !== null && !this.child.killed;
  }

  /**
   * Check if there is an active session
   */
  get hasActiveSession(): boolean {
    return this.sessionManager.getCurrentSessionId() !== null;
  }

  /**
   * Get current session ID
   */
  get currentSessionId(): string | null {
    return this.sessionManager.getCurrentSessionId();
  }
}