abbyswag/qwen-code
1
0
Fork 0
mirror of https://github.com/QwenLM/qwen-code.git synced 2025-12-19 09:33:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: implement SDK MCP server support and enhance control request handling

Browse Source 
- Added new `SdkMcpController` to manage communication between CLI MCP clients and SDK MCP servers.
- Introduced `createSdkMcpServer` function for creating SDK-embedded MCP servers.
- Updated configuration options to support both external and SDK MCP servers.
- Enhanced timeout settings for various SDK operations, including MCP requests.
- Refactored existing control request handling to accommodate new SDK MCP server functionality.
- Updated tests to cover new SDK MCP server features and ensure proper integration.
This commit is contained in:
mingholy.lmh
2025-12-04 17:01:13 +08:00
parent a58d3f7aaf
commit 322ce80e2c
29 changed files with 2473 additions and 837 deletions
Download Patch File Download Diff File Expand all files Collapse all files

65
packages/cli/src/nonInteractive/control/ControlDispatcher.ts
View File

@@ -16,9 +16,12 @@
* Controllers:
* - SystemController: initialize, interrupt, set_model, supported_commands
* - PermissionController: can_use_tool, set_permission_mode
* - MCPController: mcp_message, mcp_server_status
* - SdkMcpController: mcp_server_status (mcp_message handled via callback)
* - HookController: hook_callback
*
* Note: mcp_message requests are NOT routed through the dispatcher. CLI MCP
* clients send messages via SdkMcpController.createSendSdkMcpMessage() callback.
*
* Note: Control request types are centrally defined in the ControlRequestType
* enum in packages/sdk/typescript/src/types/controlRequests.ts
*/
@@ -27,7 +30,7 @@ import type { IControlContext } from './ControlContext.js';
import type { IPendingRequestRegistry } from './controllers/baseController.js';
import { SystemController } from './controllers/systemController.js';
import { PermissionController } from './controllers/permissionController.js';
// import { MCPController } from './controllers/mcpController.js';
import { SdkMcpController } from './controllers/sdkMcpController.js';
// import { HookController } from './controllers/hookController.js';
import type {
CLIControlRequest,
@@ -65,7 +68,7 @@ export class ControlDispatcher implements IPendingRequestRegistry {
// Make controllers publicly accessible
readonly systemController: SystemController;
readonly permissionController: PermissionController;
// readonly mcpController: MCPController;
readonly sdkMcpController: SdkMcpController;
// readonly hookController: HookController;
// Central pending request registries
@@ -88,7 +91,11 @@ export class ControlDispatcher implements IPendingRequestRegistry {
this,
'PermissionController',
);
// this.mcpController = new MCPController(context, this, 'MCPController');
this.sdkMcpController = new SdkMcpController(
context,
this,
'SdkMcpController',
);
// this.hookController = new HookController(context, this, 'HookController');
// Listen for main abort signal
@@ -228,10 +235,10 @@ export class ControlDispatcher implements IPendingRequestRegistry {
}
this.pendingOutgoingRequests.clear();
// Cleanup controllers (MCP controller will close all clients)
// Cleanup controllers
this.systemController.cleanup();
this.permissionController.cleanup();
// this.mcpController.cleanup();
this.sdkMcpController.cleanup();
// this.hookController.cleanup();
}
@@ -291,6 +298,47 @@ export class ControlDispatcher implements IPendingRequestRegistry {
}
}
/**
* Get count of pending incoming requests (for debugging)
*/
getPendingIncomingRequestCount(): number {
return this.pendingIncomingRequests.size;
}
/**
* Wait for all incoming request handlers to complete.
*
* Uses polling since we don't have direct Promise references to handlers.
* The pendingIncomingRequests map is managed by BaseController:
* - Registered when handler starts (in handleRequest)
* - Deregistered when handler completes (success or error)
*
* @param pollIntervalMs - How often to check (default 50ms)
* @param timeoutMs - Maximum wait time (default 30s)
*/
async waitForPendingIncomingRequests(
pollIntervalMs: number = 50,
timeoutMs: number = 30000,
): Promise<void> {
const startTime = Date.now();
while (this.pendingIncomingRequests.size > 0) {
if (Date.now() - startTime > timeoutMs) {
if (this.context.debugMode) {
console.error(
`[ControlDispatcher] Timeout waiting for ${this.pendingIncomingRequests.size} pending incoming requests`,
);
}
break;
}
await new Promise((resolve) => setTimeout(resolve, pollIntervalMs));
}
if (this.context.debugMode && this.pendingIncomingRequests.size === 0) {
console.error('[ControlDispatcher] All incoming requests completed');
}
}
/**
* Returns the controller that handles the given request subtype
*/
@@ -306,9 +354,8 @@ export class ControlDispatcher implements IPendingRequestRegistry {
case 'set_permission_mode':
return this.permissionController;
// case 'mcp_message':
// case 'mcp_server_status':
// return this.mcpController;
case 'mcp_server_status':
return this.sdkMcpController;
// case 'hook_callback':
// return this.hookController;

47
packages/cli/src/nonInteractive/control/controllers/baseController.ts
View File

@@ -117,16 +117,41 @@ export abstract class BaseController {
* Send an outgoing control request to SDK
*
* Manages lifecycle: register -> send -> wait for response -> deregister
* Respects the provided AbortSignal for cancellation.
*/
async sendControlRequest(
payload: ControlRequestPayload,
timeoutMs: number = DEFAULT_REQUEST_TIMEOUT_MS,
signal?: AbortSignal,
): Promise<ControlResponse> {
// Check if already aborted
if (signal?.aborted) {
throw new Error('Request aborted');
}
const requestId = randomUUID();
return new Promise<ControlResponse>((resolve, reject) => {
// Setup abort handler
const abortHandler = () => {
this.registry.deregisterOutgoingRequest(requestId);
reject(new Error('Request aborted'));
if (this.context.debugMode) {
console.error(
`[${this.controllerName}] Outgoing request aborted: ${requestId}`,
);
}
};
if (signal) {
signal.addEventListener('abort', abortHandler, { once: true });
}
// Setup timeout
const timeoutId = setTimeout(() => {
if (signal) {
signal.removeEventListener('abort', abortHandler);
}
this.registry.deregisterOutgoingRequest(requestId);
reject(new Error('Control request timeout'));
if (this.context.debugMode) {
@@ -136,12 +161,27 @@ export abstract class BaseController {
}
}, timeoutMs);
// Wrap resolve/reject to clean up abort listener
const wrappedResolve = (response: ControlResponse) => {
if (signal) {
signal.removeEventListener('abort', abortHandler);
}
resolve(response);
};
const wrappedReject = (error: Error) => {
if (signal) {
signal.removeEventListener('abort', abortHandler);
}
reject(error);
};
// Register with central registry
this.registry.registerOutgoingRequest(
requestId,
this.controllerName,
resolve,
reject,
wrappedResolve,
wrappedReject,
timeoutId,
);
@@ -155,6 +195,9 @@ export abstract class BaseController {
try {
this.context.streamJson.send(request);
} catch (error) {
if (signal) {
signal.removeEventListener('abort', abortHandler);
}
this.registry.deregisterOutgoingRequest(requestId);
reject(error);
}

287
packages/cli/src/nonInteractive/control/controllers/mcpController.ts
View File

@@ -1,287 +0,0 @@
/**
* @license
* Copyright 2025 Qwen Team
* SPDX-License-Identifier: Apache-2.0
*/
/**
* MCP Controller
*
* Handles MCP-related control requests:
* - mcp_message: Route MCP messages
* - mcp_server_status: Return MCP server status
*/
import { BaseController } from './baseController.js';
import type { Client } from '@modelcontextprotocol/sdk/client/index.js';
import { ResultSchema } from '@modelcontextprotocol/sdk/types.js';
import type {
ControlRequestPayload,
CLIControlMcpMessageRequest,
} from '../../types.js';
import type {
MCPServerConfig,
WorkspaceContext,
} from '@qwen-code/qwen-code-core';
import {
connectToMcpServer,
MCP_DEFAULT_TIMEOUT_MSEC,
} from '@qwen-code/qwen-code-core';
export class MCPController extends BaseController {
/**
* Handle MCP control requests
*/
protected async handleRequestPayload(
payload: ControlRequestPayload,
_signal: AbortSignal,
): Promise<Record<string, unknown>> {
switch (payload.subtype) {
case 'mcp_message':
return this.handleMcpMessage(payload as CLIControlMcpMessageRequest);
case 'mcp_server_status':
return this.handleMcpStatus();
default:
throw new Error(`Unsupported request subtype in MCPController`);
}
}
/**
* Handle mcp_message request
*
* Routes JSON-RPC messages to MCP servers
*/
private async handleMcpMessage(
payload: CLIControlMcpMessageRequest,
): Promise<Record<string, unknown>> {
const serverNameRaw = payload.server_name;
if (
typeof serverNameRaw !== 'string' ||
serverNameRaw.trim().length === 0
) {
throw new Error('Missing server_name in mcp_message request');
}
const message = payload.message;
if (!message || typeof message !== 'object') {
throw new Error(
'Missing or invalid message payload for mcp_message request',
);
}
// Get or create MCP client
let clientEntry: { client: Client; config: MCPServerConfig };
try {
clientEntry = await this.getOrCreateMcpClient(serverNameRaw.trim());
} catch (error) {
throw new Error(
error instanceof Error
? error.message
: 'Failed to connect to MCP server',
);
}
const method = message.method;
if (typeof method !== 'string' || method.trim().length === 0) {
throw new Error('Invalid MCP message: missing method');
}
const jsonrpcVersion =
typeof message.jsonrpc === 'string' ? message.jsonrpc : '2.0';
const messageId = message.id;
const params = message.params;
const timeout =
typeof clientEntry.config.timeout === 'number'
? clientEntry.config.timeout
: MCP_DEFAULT_TIMEOUT_MSEC;
try {
// Handle notification (no id)
if (messageId === undefined) {
await clientEntry.client.notification({
method,
params,
});
return {
subtype: 'mcp_message',
mcp_response: {
jsonrpc: jsonrpcVersion,
id: null,
result: { success: true, acknowledged: true },
},
};
}
// Handle request (with id)
const result = await clientEntry.client.request(
{
method,
params,
},
ResultSchema,
{ timeout },
);
return {
subtype: 'mcp_message',
mcp_response: {
jsonrpc: jsonrpcVersion,
id: messageId,
result,
},
};
} catch (error) {
// If connection closed, remove from cache
if (error instanceof Error && /closed/i.test(error.message)) {
this.context.mcpClients.delete(serverNameRaw.trim());
}
const errorCode =
typeof (error as { code?: unknown })?.code === 'number'
? ((error as { code: number }).code as number)
: -32603;
const errorMessage =
error instanceof Error
? error.message
: 'Failed to execute MCP request';
const errorData = (error as { data?: unknown })?.data;
const errorBody: Record<string, unknown> = {
code: errorCode,
message: errorMessage,
};
if (errorData !== undefined) {
errorBody['data'] = errorData;
}
return {
subtype: 'mcp_message',
mcp_response: {
jsonrpc: jsonrpcVersion,
id: messageId ?? null,
error: errorBody,
},
};
}
}
/**
* Handle mcp_server_status request
*
* Returns status of registered MCP servers
*/
private async handleMcpStatus(): Promise<Record<string, unknown>> {
const status: Record<string, string> = {};
// Include SDK MCP servers
for (const serverName of this.context.sdkMcpServers) {
status[serverName] = 'connected';
}
// Include CLI-managed MCP clients
for (const serverName of this.context.mcpClients.keys()) {
status[serverName] = 'connected';
}
if (this.context.debugMode) {
console.error(
`[MCPController] MCP status: ${Object.keys(status).length} servers`,
);
}
return status;
}
/**
* Get or create MCP client for a server
*
* Implements lazy connection and caching
*/
private async getOrCreateMcpClient(
serverName: string,
): Promise<{ client: Client; config: MCPServerConfig }> {
// Check cache first
const cached = this.context.mcpClients.get(serverName);
if (cached) {
return cached;
}
// Get server configuration
const provider = this.context.config as unknown as {
getMcpServers?: () => Record<string, MCPServerConfig> | undefined;
getDebugMode?: () => boolean;
getWorkspaceContext?: () => unknown;
};
if (typeof provider.getMcpServers !== 'function') {
throw new Error(`MCP server "${serverName}" is not configured`);
}
const servers = provider.getMcpServers() ?? {};
const serverConfig = servers[serverName];
if (!serverConfig) {
throw new Error(`MCP server "${serverName}" is not configured`);
}
const debugMode =
typeof provider.getDebugMode === 'function'
? provider.getDebugMode()
: false;
const workspaceContext =
typeof provider.getWorkspaceContext === 'function'
? provider.getWorkspaceContext()
: undefined;
if (!workspaceContext) {
throw new Error('Workspace context is not available for MCP connection');
}
// Connect to MCP server
const client = await connectToMcpServer(
serverName,
serverConfig,
debugMode,
workspaceContext as WorkspaceContext,
);
// Cache the client
const entry = { client, config: serverConfig };
this.context.mcpClients.set(serverName, entry);
if (this.context.debugMode) {
console.error(`[MCPController] Connected to MCP server: ${serverName}`);
}
return entry;
}
/**
* Cleanup MCP clients
*/
override cleanup(): void {
if (this.context.debugMode) {
console.error(
`[MCPController] Cleaning up ${this.context.mcpClients.size} MCP clients`,
);
}
// Close all MCP clients
for (const [serverName, { client }] of this.context.mcpClients.entries()) {
try {
client.close();
} catch (error) {
if (this.context.debugMode) {
console.error(
`[MCPController] Failed to close MCP client ${serverName}:`,
error,
);
}
}
}
this.context.mcpClients.clear();
}
}

50
packages/cli/src/nonInteractive/control/controllers/permissionController.ts
View File

@@ -44,15 +44,23 @@ export class PermissionController extends BaseController {
*/
protected async handleRequestPayload(
payload: ControlRequestPayload,
_signal: AbortSignal,
signal: AbortSignal,
): Promise<Record<string, unknown>> {
if (signal.aborted) {
throw new Error('Request aborted');
}
switch (payload.subtype) {
case 'can_use_tool':
return this.handleCanUseTool(payload as CLIControlPermissionRequest);
return this.handleCanUseTool(
payload as CLIControlPermissionRequest,
signal,
);
case 'set_permission_mode':
return this.handleSetPermissionMode(
payload as CLIControlSetPermissionModeRequest,
signal,
);
default:
@@ -70,7 +78,12 @@ export class PermissionController extends BaseController {
*/
private async handleCanUseTool(
payload: CLIControlPermissionRequest,
signal: AbortSignal,
): Promise<Record<string, unknown>> {
if (signal.aborted) {
throw new Error('Request aborted');
}
const toolName = payload.tool_name;
if (
!toolName ||
@@ -192,7 +205,12 @@ export class PermissionController extends BaseController {
*/
private async handleSetPermissionMode(
payload: CLIControlSetPermissionModeRequest,
signal: AbortSignal,
): Promise<Record<string, unknown>> {
if (signal.aborted) {
throw new Error('Request aborted');
}
const mode = payload.mode;
const validModes: PermissionMode[] = [
'default',
@@ -373,6 +391,14 @@ export class PermissionController extends BaseController {
toolCall: WaitingToolCall,
): Promise<void> {
try {
// Check if already aborted
if (this.context.abortSignal?.aborted) {
await toolCall.confirmationDetails.onConfirm(
ToolConfirmationOutcome.Cancel,
);
return;
}
const inputFormat = this.context.config.getInputFormat?.();
const isStreamJsonMode = inputFormat === InputFormat.STREAM_JSON;
@@ -392,14 +418,18 @@ export class PermissionController extends BaseController {
toolCall.confirmationDetails,
);
const response = await this.sendControlRequest({
subtype: 'can_use_tool',
tool_name: toolCall.request.name,
tool_use_id: toolCall.request.callId,
input: toolCall.request.args,
permission_suggestions: permissionSuggestions,
blocked_path: null,
} as CLIControlPermissionRequest);
const response = await this.sendControlRequest(
{
subtype: 'can_use_tool',
tool_name: toolCall.request.name,
tool_use_id: toolCall.request.callId,
input: toolCall.request.args,
permission_suggestions: permissionSuggestions,
blocked_path: null,
} as CLIControlPermissionRequest,
undefined, // use default timeout
this.context.abortSignal,
);
if (response.subtype !== 'success') {
await toolCall.confirmationDetails.onConfirm(

138
packages/cli/src/nonInteractive/control/controllers/sdkMcpController.ts Normal file
View File

@@ -0,0 +1,138 @@
/**
* @license
* Copyright 2025 Qwen Team
* SPDX-License-Identifier: Apache-2.0
*/
/**
* SDK MCP Controller
*
* Handles MCP communication between CLI MCP clients and SDK MCP servers:
* - Provides sendSdkMcpMessage callback for CLI → SDK MCP message routing
* - mcp_server_status: Returns status of SDK MCP servers
*
* Message Flow (CLI MCP Client → SDK MCP Server):
* CLI MCP Client → SdkControlClientTransport.send() →
* sendSdkMcpMessage callback → control_request (mcp_message) → SDK →
* SDK MCP Server processes → control_response → CLI MCP Client
*/
import type { JSONRPCMessage } from '@modelcontextprotocol/sdk/types.js';
import { BaseController } from './baseController.js';
import type {
ControlRequestPayload,
CLIControlMcpMessageRequest,
} from '../../types.js';
const MCP_REQUEST_TIMEOUT = 30_000; // 30 seconds
export class SdkMcpController extends BaseController {
/**
* Handle SDK MCP control requests from ControlDispatcher
*
* Note: mcp_message requests are NOT handled here. CLI MCP clients
* send messages via the sendSdkMcpMessage callback directly, not
* through the control dispatcher.
*/
protected async handleRequestPayload(
payload: ControlRequestPayload,
signal: AbortSignal,
): Promise<Record<string, unknown>> {
if (signal.aborted) {
throw new Error('Request aborted');
}
switch (payload.subtype) {
case 'mcp_server_status':
return this.handleMcpStatus();
default:
throw new Error(`Unsupported request subtype in SdkMcpController`);
}
}
/**
* Handle mcp_server_status request
*
* Returns status of all registered SDK MCP servers.
* SDK servers are considered "connected" if they are registered.
*/
private async handleMcpStatus(): Promise<Record<string, unknown>> {
const status: Record<string, string> = {};
for (const serverName of this.context.sdkMcpServers) {
// SDK MCP servers are "connected" once registered since they run in SDK process
status[serverName] = 'connected';
}
return {
subtype: 'mcp_server_status',
status,
};
}
/**
* Send MCP message to SDK server via control plane
*
* @param serverName - Name of the SDK MCP server
* @param message - MCP JSON-RPC message to send
* @returns MCP JSON-RPC response from SDK server
*/
private async sendMcpMessageToSdk(
serverName: string,
message: JSONRPCMessage,
): Promise<JSONRPCMessage> {
if (this.context.debugMode) {
console.error(
`[SdkMcpController] Sending MCP message to SDK server '${serverName}':`,
JSON.stringify(message),
);
}
// Send control request to SDK with the MCP message
const response = await this.sendControlRequest(
{
subtype: 'mcp_message',
server_name: serverName,
message: message as CLIControlMcpMessageRequest['message'],
},
MCP_REQUEST_TIMEOUT,
this.context.abortSignal,
);
// Extract MCP response from control response
const responsePayload = response.response as Record<string, unknown>;
const mcpResponse = responsePayload?.['mcp_response'] as JSONRPCMessage;
if (!mcpResponse) {
throw new Error(
`Invalid MCP response from SDK for server '${serverName}'`,
);
}
if (this.context.debugMode) {
console.error(
`[SdkMcpController] Received MCP response from SDK server '${serverName}':`,
JSON.stringify(mcpResponse),
);
}
return mcpResponse;
}
/**
* Create a callback function for sending MCP messages to SDK servers.
*
* This callback is used by McpClientManager/SdkControlClientTransport to send
* MCP messages from CLI MCP clients to SDK MCP servers via the control plane.
*
* @returns A function that sends MCP messages to SDK and returns the response
*/
createSendSdkMcpMessage(): (
serverName: string,
message: JSONRPCMessage,
) => Promise<JSONRPCMessage> {
return (serverName: string, message: JSONRPCMessage) =>
this.sendMcpMessageToSdk(serverName, message);
}
}

266
packages/cli/src/nonInteractive/control/controllers/systemController.ts
View File

@@ -18,9 +18,15 @@ import type {
ControlRequestPayload,
CLIControlInitializeRequest,
CLIControlSetModelRequest,
CLIMcpServerConfig,
} from '../../types.js';
import { CommandService } from '../../../services/CommandService.js';
import { BuiltinCommandLoader } from '../../../services/BuiltinCommandLoader.js';
import {
MCPServerConfig,
AuthProviderType,
type MCPOAuthConfig,
} from '@qwen-code/qwen-code-core';
export class SystemController extends BaseController {
/**
@@ -28,20 +34,30 @@ export class SystemController extends BaseController {
*/
protected async handleRequestPayload(
payload: ControlRequestPayload,
_signal: AbortSignal,
signal: AbortSignal,
): Promise<Record<string, unknown>> {
if (signal.aborted) {
throw new Error('Request aborted');
}
switch (payload.subtype) {
case 'initialize':
return this.handleInitialize(payload as CLIControlInitializeRequest);
return this.handleInitialize(
payload as CLIControlInitializeRequest,
signal,
);
case 'interrupt':
return this.handleInterrupt();
case 'set_model':
return this.handleSetModel(payload as CLIControlSetModelRequest);
return this.handleSetModel(
payload as CLIControlSetModelRequest,
signal,
);
case 'supported_commands':
return this.handleSupportedCommands();
return this.handleSupportedCommands(signal);
default:
throw new Error(`Unsupported request subtype in SystemController`);
@@ -51,46 +67,110 @@ export class SystemController extends BaseController {
/**
* Handle initialize request
*
* Registers SDK MCP servers and returns capabilities
* Processes SDK MCP servers config.
* SDK servers are registered in context.sdkMcpServers
* and added to config.mcpServers with the sdk type flag.
* External MCP servers are configured separately in settings.
*/
private async handleInitialize(
payload: CLIControlInitializeRequest,
signal: AbortSignal,
): Promise<Record<string, unknown>> {
if (signal.aborted) {
throw new Error('Request aborted');
}
this.context.config.setSdkMode(true);
if (payload.sdkMcpServers && typeof payload.sdkMcpServers === 'object') {
for (const serverName of Object.keys(payload.sdkMcpServers)) {
this.context.sdkMcpServers.add(serverName);
// Process SDK MCP servers
if (
payload.sdkMcpServers &&
typeof payload.sdkMcpServers === 'object' &&
payload.sdkMcpServers !== null
) {
const sdkServers: Record<string, MCPServerConfig> = {};
for (const [key, wireConfig] of Object.entries(payload.sdkMcpServers)) {
const name =
typeof wireConfig?.name === 'string' && wireConfig.name.trim().length
? wireConfig.name
: key;
this.context.sdkMcpServers.add(name);
sdkServers[name] = new MCPServerConfig(
undefined, // command
undefined, // args
undefined, // env
undefined, // cwd
undefined, // url
undefined, // httpUrl
undefined, // headers
undefined, // tcp
undefined, // timeout
true, // trust - SDK servers are trusted
undefined, // description
undefined, // includeTools
undefined, // excludeTools
undefined, // extensionName
undefined, // oauth
undefined, // authProviderType
undefined, // targetAudience
undefined, // targetServiceAccount
'sdk', // type
);
}
try {
this.context.config.addMcpServers(payload.sdkMcpServers);
if (this.context.debugMode) {
console.error(
`[SystemController] Added ${Object.keys(payload.sdkMcpServers).length} SDK MCP servers to config`,
);
}
} catch (error) {
if (this.context.debugMode) {
console.error(
'[SystemController] Failed to add SDK MCP servers:',
error,
);
const sdkServerCount = Object.keys(sdkServers).length;
if (sdkServerCount > 0) {
try {
this.context.config.addMcpServers(sdkServers);
if (this.context.debugMode) {
console.error(
`[SystemController] Added ${sdkServerCount} SDK MCP servers to config`,
);
}
} catch (error) {
if (this.context.debugMode) {
console.error(
'[SystemController] Failed to add SDK MCP servers:',
error,
);
}
}
}
}
if (payload.mcpServers && typeof payload.mcpServers === 'object') {
try {
this.context.config.addMcpServers(payload.mcpServers);
if (this.context.debugMode) {
console.error(
`[SystemController] Added ${Object.keys(payload.mcpServers).length} MCP servers to config`,
);
if (
payload.mcpServers &&
typeof payload.mcpServers === 'object' &&
payload.mcpServers !== null
) {
const externalServers: Record<string, MCPServerConfig> = {};
for (const [name, serverConfig] of Object.entries(payload.mcpServers)) {
const normalized = this.normalizeMcpServerConfig(
name,
serverConfig as CLIMcpServerConfig | undefined,
);
if (normalized) {
externalServers[name] = normalized;
}
} catch (error) {
if (this.context.debugMode) {
console.error('[SystemController] Failed to add MCP servers:', error);
}
const externalCount = Object.keys(externalServers).length;
if (externalCount > 0) {
try {
this.context.config.addMcpServers(externalServers);
if (this.context.debugMode) {
console.error(
`[SystemController] Added ${externalCount} external MCP servers to config`,
);
}
} catch (error) {
if (this.context.debugMode) {
console.error(
'[SystemController] Failed to add external MCP servers:',
error,
);
}
}
}
}
@@ -143,13 +223,96 @@ export class SystemController extends BaseController {
can_set_permission_mode:
typeof this.context.config.setApprovalMode === 'function',
can_set_model: typeof this.context.config.setModel === 'function',
/* TODO: sdkMcpServers support */
can_handle_mcp_message: false,
// SDK MCP servers are supported - messages routed through control plane
can_handle_mcp_message: true,
};
return capabilities;
}
private normalizeMcpServerConfig(
serverName: string,
config?: CLIMcpServerConfig,
): MCPServerConfig | null {
if (!config || typeof config !== 'object') {
if (this.context.debugMode) {
console.error(
`[SystemController] Ignoring invalid MCP server config for '${serverName}'`,
);
}
return null;
}
const authProvider = this.normalizeAuthProviderType(
config.authProviderType,
);
const oauthConfig = this.normalizeOAuthConfig(config.oauth);
return new MCPServerConfig(
config.command,
config.args,
config.env,
config.cwd,
config.url,
config.httpUrl,
config.headers,
config.tcp,
config.timeout,
config.trust,
config.description,
config.includeTools,
config.excludeTools,
config.extensionName,
oauthConfig,
authProvider,
config.targetAudience,
config.targetServiceAccount,
);
}
private normalizeAuthProviderType(
value?: string,
): AuthProviderType | undefined {
if (!value) {
return undefined;
}
switch (value) {
case AuthProviderType.DYNAMIC_DISCOVERY:
case AuthProviderType.GOOGLE_CREDENTIALS:
case AuthProviderType.SERVICE_ACCOUNT_IMPERSONATION:
return value;
default:
if (this.context.debugMode) {
console.error(
`[SystemController] Unsupported authProviderType '${value}', skipping`,
);
}
return undefined;
}
}
private normalizeOAuthConfig(
oauth?: CLIMcpServerConfig['oauth'],
): MCPOAuthConfig | undefined {
if (!oauth) {
return undefined;
}
return {
enabled: oauth.enabled,
clientId: oauth.clientId,
clientSecret: oauth.clientSecret,
authorizationUrl: oauth.authorizationUrl,
tokenUrl: oauth.tokenUrl,
scopes: oauth.scopes,
audiences: oauth.audiences,
redirectUri: oauth.redirectUri,
tokenParamName: oauth.tokenParamName,
registrationUrl: oauth.registrationUrl,
};
}
/**
* Handle interrupt request
*
@@ -183,7 +346,12 @@ export class SystemController extends BaseController {
*/
private async handleSetModel(
payload: CLIControlSetModelRequest,
signal: AbortSignal,
): Promise<Record<string, unknown>> {
if (signal.aborted) {
throw new Error('Request aborted');
}
const model = payload.model;
// Validate model parameter
@@ -223,8 +391,14 @@ export class SystemController extends BaseController {
*
* Returns list of supported slash commands loaded dynamically
*/
private async handleSupportedCommands(): Promise<Record<string, unknown>> {
const slashCommands = await this.loadSlashCommandNames();
private async handleSupportedCommands(
signal: AbortSignal,
): Promise<Record<string, unknown>> {
if (signal.aborted) {
throw new Error('Request aborted');
}
const slashCommands = await this.loadSlashCommandNames(signal);
return {
subtype: 'supported_commands',
@@ -235,15 +409,24 @@ export class SystemController extends BaseController {
/**
* Load slash command names using CommandService
*
* @param signal - AbortSignal to respect for cancellation
* @returns Promise resolving to array of slash command names
*/
private async loadSlashCommandNames(): Promise<string[]> {
const controller = new AbortController();
private async loadSlashCommandNames(signal: AbortSignal): Promise<string[]> {
if (signal.aborted) {
return [];
}
try {
const service = await CommandService.create(
[new BuiltinCommandLoader(this.context.config)],
controller.signal,
signal,
);
if (signal.aborted) {
return [];
}
const names = new Set<string>();
const commands = service.getCommands();
for (const command of commands) {
@@ -251,6 +434,11 @@ export class SystemController extends BaseController {
}
return Array.from(names).sort();
} catch (error) {
// Check if the error is due to abort
if (signal.aborted) {
return [];
}
if (this.context.debugMode) {
console.error(
'[SystemController] Failed to load slash commands:',
@@ -258,8 +446,6 @@ export class SystemController extends BaseController {
);
}
return [];
} finally {
controller.abort();
}
}
}

10
packages/cli/src/nonInteractive/session.test.ts
View File

@@ -153,6 +153,11 @@ describe('runNonInteractiveStreamJson', () => {
handleControlResponse: ReturnType<typeof vi.fn>;
handleCancel: ReturnType<typeof vi.fn>;
shutdown: ReturnType<typeof vi.fn>;
getPendingIncomingRequestCount: ReturnType<typeof vi.fn>;
waitForPendingIncomingRequests: ReturnType<typeof vi.fn>;
sdkMcpController: {
createSendSdkMcpMessage: ReturnType<typeof vi.fn>;
};
};
let mockConsolePatcher: {
patch: ReturnType<typeof vi.fn>;
@@ -187,6 +192,11 @@ describe('runNonInteractiveStreamJson', () => {
handleControlResponse: vi.fn(),
handleCancel: vi.fn(),
shutdown: vi.fn(),
getPendingIncomingRequestCount: vi.fn().mockReturnValue(0),
waitForPendingIncomingRequests: vi.fn().mockResolvedValue(undefined),
sdkMcpController: {
createSendSdkMcpMessage: vi.fn().mockReturnValue(vi.fn()),
},
};
(
ControlDispatcher as unknown as ReturnType<typeof vi.fn>

311
packages/cli/src/nonInteractive/session.ts
View File

@@ -4,7 +4,10 @@
* SPDX-License-Identifier: Apache-2.0
*/
import type { Config } from '@qwen-code/qwen-code-core';
import type {
Config,
ConfigInitializeOptions,
} from '@qwen-code/qwen-code-core';
import { StreamJsonInputReader } from './io/StreamJsonInputReader.js';
import { StreamJsonOutputAdapter } from './io/StreamJsonOutputAdapter.js';
import { ControlContext } from './control/ControlContext.js';
@@ -50,6 +53,12 @@ class Session {
private isShuttingDown: boolean = false;
private configInitialized: boolean = false;
// Single initialization promise that resolves when session is ready for user messages.
// Created lazily once initialization actually starts.
private initializationPromise: Promise<void> | null = null;
private initializationResolve: (() => void) | null = null;
private initializationReject: ((error: Error) => void) | null = null;
constructor(config: Config, initialPrompt?: CLIUserMessage) {
this.config = config;
this.sessionId = config.getSessionId();
@@ -66,12 +75,32 @@ class Session {
this.setupSignalHandlers();
}
private ensureInitializationPromise(): void {
if (this.initializationPromise) {
return;
}
this.initializationPromise = new Promise<void>((resolve, reject) => {
this.initializationResolve = () => {
resolve();
this.initializationResolve = null;
this.initializationReject = null;
};
this.initializationReject = (error: Error) => {
reject(error);
this.initializationResolve = null;
this.initializationReject = null;
};
});
}
private getNextPromptId(): string {
this.promptIdCounter++;
return `${this.sessionId}########${this.promptIdCounter}`;
}
private async ensureConfigInitialized(): Promise<void> {
private async ensureConfigInitialized(
options?: ConfigInitializeOptions,
): Promise<void> {
if (this.configInitialized) {
return;
}
@@ -81,7 +110,7 @@ class Session {
}
try {
await this.config.initialize();
await this.config.initialize(options);
this.configInitialized = true;
} catch (error) {
if (this.debugMode) {
@@ -91,6 +120,44 @@ class Session {
}
}
/**
* Mark initialization as complete
*/
private completeInitialization(): void {
if (this.initializationResolve) {
if (this.debugMode) {
console.error('[Session] Initialization complete');
}
this.initializationResolve();
this.initializationResolve = null;
this.initializationReject = null;
}
}
/**
* Mark initialization as failed
*/
private failInitialization(error: Error): void {
if (this.initializationReject) {
if (this.debugMode) {
console.error('[Session] Initialization failed:', error);
}
this.initializationReject(error);
this.initializationResolve = null;
this.initializationReject = null;
}
}
/**
* Wait for session to be ready for user messages
*/
private async waitForInitialization(): Promise<void> {
if (!this.initializationPromise) {
return;
}
await this.initializationPromise;
}
private ensureControlSystem(): void {
if (this.controlContext && this.dispatcher && this.controlService) {
return;
@@ -120,49 +187,114 @@ class Session {
return this.dispatcher;
}
private async handleFirstMessage(
/**
* Handle the first message to determine session mode (SDK vs direct).
* This is synchronous from the message loop's perspective - it starts
* async work but does not return a promise that the loop awaits.
*
* The initialization completes asynchronously and resolves initializationPromise
* when ready for user messages.
*/
private handleFirstMessage(
message:
| CLIMessage
| CLIControlRequest
| CLIControlResponse
| ControlCancelRequest,
): Promise<boolean> {
): void {
if (isControlRequest(message)) {
const request = message as CLIControlRequest;
this.controlSystemEnabled = true;
this.ensureControlSystem();
if (request.request.subtype === 'initialize') {
// Dispatch the initialize request first
await this.dispatcher?.dispatch(request);
// After handling initialize control request, initialize the config
// This is the SDK mode where config initialization is deferred
await this.ensureConfigInitialized();
return true;
if (request.request.subtype === 'initialize') {
// Start SDK mode initialization (fire-and-forget from loop perspective)
void this.initializeSdkMode(request);
return;
}
if (this.debugMode) {
console.error(
'[Session] Ignoring non-initialize control request during initialization',
);
}
return true;
return;
}
if (isCLIUserMessage(message)) {
this.controlSystemEnabled = false;
// For non-SDK mode (direct user message), initialize config if not already done
await this.ensureConfigInitialized();
this.enqueueUserMessage(message as CLIUserMessage);
return true;
// Start direct mode initialization (fire-and-forget from loop perspective)
void this.initializeDirectMode(message as CLIUserMessage);
return;
}
this.controlSystemEnabled = false;
return false;
}
private async handleControlRequest(
request: CLIControlRequest,
/**
* SDK mode initialization flow
* Dispatches initialize request and initializes config with MCP support
*/
private async initializeSdkMode(request: CLIControlRequest): Promise<void> {
this.ensureInitializationPromise();
try {
// Dispatch the initialize request first
// This registers SDK MCP servers in the control context
await this.dispatcher?.dispatch(request);
// Get sendSdkMcpMessage callback from SdkMcpController
// This callback is used by McpClientManager to send MCP messages
// from CLI MCP clients to SDK MCP servers via the control plane
const sendSdkMcpMessage =
this.dispatcher?.sdkMcpController.createSendSdkMcpMessage();
// Initialize config with SDK MCP message support
await this.ensureConfigInitialized({ sendSdkMcpMessage });
// Initialization complete!
this.completeInitialization();
} catch (error) {
if (this.debugMode) {
console.error('[Session] SDK mode initialization failed:', error);
}
this.failInitialization(
error instanceof Error ? error : new Error(String(error)),
);
}
}
/**
* Direct mode initialization flow
* Initializes config and enqueues the first user message
*/
private async initializeDirectMode(
userMessage: CLIUserMessage,
): Promise<void> {
this.ensureInitializationPromise();
try {
// Initialize config
await this.ensureConfigInitialized();
// Initialization complete!
this.completeInitialization();
// Enqueue the first user message for processing
this.enqueueUserMessage(userMessage);
} catch (error) {
if (this.debugMode) {
console.error('[Session] Direct mode initialization failed:', error);
}
this.failInitialization(
error instanceof Error ? error : new Error(String(error)),
);
}
}
/**
* Handle control request asynchronously (fire-and-forget from main loop).
* Errors are handled internally and responses sent by dispatcher.
*/
private handleControlRequestAsync(request: CLIControlRequest): void {
const dispatcher = this.getDispatcher();
if (!dispatcher) {
if (this.debugMode) {
@@ -171,9 +303,20 @@ class Session {
return;
}
await dispatcher.dispatch(request);
// Fire-and-forget: dispatch runs concurrently
// The dispatcher's pendingIncomingRequests tracks completion
void dispatcher.dispatch(request).catch((error) => {
if (this.debugMode) {
console.error('[Session] Control request dispatch error:', error);
}
// Error response is already sent by dispatcher.dispatch()
});
}
/**
* Handle control response - MUST be synchronous
* This resolves pending outgoing requests, breaking the deadlock cycle.
*/
private handleControlResponse(response: CLIControlResponse): void {
const dispatcher = this.getDispatcher();
if (!dispatcher) {
@@ -201,8 +344,8 @@ class Session {
return;
}
// Ensure config is initialized before processing user messages
await this.ensureConfigInitialized();
// Wait for initialization to complete before processing user messages
await this.waitForInitialization();
const promptId = this.getNextPromptId();
@@ -307,6 +450,45 @@ class Session {
process.on('SIGTERM', this.shutdownHandler);
}
/**
* Wait for all pending work to complete before shutdown
*/
private async waitForAllPendingWork(): Promise<void> {
// 1. Wait for initialization to complete (or fail)
try {
await this.waitForInitialization();
} catch (error) {
if (this.debugMode) {
console.error('[Session] Initialization error during shutdown:', error);
}
}
// 2. Wait for all control request handlers using dispatcher's tracking
if (this.dispatcher) {
const pendingCount = this.dispatcher.getPendingIncomingRequestCount();
if (pendingCount > 0 && this.debugMode) {
console.error(
`[Session] Waiting for ${pendingCount} pending control request handlers`,
);
}
await this.dispatcher.waitForPendingIncomingRequests();
}
// 3. Wait for user message processing queue
while (this.processingPromise) {
if (this.debugMode) {
console.error('[Session] Waiting for user message processing');
}
try {
await this.processingPromise;
} catch (error) {
if (this.debugMode) {
console.error('[Session] Error in user message processing:', error);
}
}
}
}
private async shutdown(): Promise<void> {
if (this.debugMode) {
console.error('[Session] Shutting down');
@@ -314,18 +496,8 @@ class Session {
this.isShuttingDown = true;
if (this.processingPromise) {
try {
await this.processingPromise;
} catch (error) {
if (this.debugMode) {
console.error(
'[Session] Error waiting for processing to complete:',
error,
);
}
}
}
// Wait for all pending work
await this.waitForAllPendingWork();
this.dispatcher?.shutdown();
this.cleanupSignalHandlers();
@@ -339,18 +511,30 @@ class Session {
}
}
/**
* Main message processing loop
*
* CRITICAL: This loop must NEVER await handlers that might need to
* send control requests and wait for responses. Such handlers must
* be started in fire-and-forget mode, allowing the loop to continue
* reading responses that resolve pending requests.
*
* Message handling order:
* 1. control_response - FIRST, synchronously resolves pending requests
* 2. First message - determines mode, starts async initialization
* 3. control_request - fire-and-forget, tracked by dispatcher
* 4. control_cancel - synchronous
* 5. user_message - enqueued for processing
*/
async run(): Promise<void> {
try {
if (this.debugMode) {
console.error('[Session] Starting session', this.sessionId);
}
// Handle initial prompt if provided (fire-and-forget)
if (this.initialPrompt !== null) {
const handled = await this.handleFirstMessage(this.initialPrompt);
if (handled && this.isShuttingDown) {
await this.shutdown();
return;
}
this.handleFirstMessage(this.initialPrompt);
}
try {
@@ -359,23 +543,33 @@ class Session {
break;
}
if (this.controlSystemEnabled === null) {
const handled = await this.handleFirstMessage(message);
if (handled) {
if (this.isShuttingDown) {
break;
}
continue;
}
// ============================================================
// CRITICAL: Handle control_response FIRST and SYNCHRONOUSLY
// This resolves pending outgoing requests, breaking deadlock.
// ============================================================
if (isControlResponse(message)) {
this.handleControlResponse(message as CLIControlResponse);
continue;
}
// Handle first message to determine session mode
if (this.controlSystemEnabled === null) {
this.handleFirstMessage(message);
continue;
}
// ============================================================
// CRITICAL: Handle control_request in FIRE-AND-FORGET mode
// DON'T await - let handler run concurrently while loop continues
// Dispatcher's pendingIncomingRequests tracks completion
// ============================================================
if (isControlRequest(message)) {
await this.handleControlRequest(message as CLIControlRequest);
} else if (isControlResponse(message)) {
this.handleControlResponse(message as CLIControlResponse);
this.handleControlRequestAsync(message as CLIControlRequest);
} else if (isControlCancel(message)) {
// Cancel is synchronous - OK to handle inline
this.handleControlCancel(message as ControlCancelRequest);
} else if (isCLIUserMessage(message)) {
// User messages are enqueued, processing runs separately
this.enqueueUserMessage(message as CLIUserMessage);
} else if (this.debugMode) {
if (
@@ -402,19 +596,8 @@ class Session {
throw streamError;
}
while (this.processingPromise) {
if (this.debugMode) {
console.error('[Session] Waiting for final processing to complete');
}
try {
await this.processingPromise;
} catch (error) {
if (this.debugMode) {
console.error('[Session] Error in final processing:', error);
}
}
}
// Stream ended - wait for all pending work before shutdown
await this.waitForAllPendingWork();
await this.shutdown();
} catch (error) {
if (this.debugMode) {

66
packages/cli/src/nonInteractive/types.ts
View File

@@ -1,8 +1,5 @@
/* eslint-disable @typescript-eslint/no-explicit-any */
import type {
MCPServerConfig,
SubagentConfig,
} from '@qwen-code/qwen-code-core';
import type { SubagentConfig } from '@qwen-code/qwen-code-core';
/**
* Annotation for attaching metadata to content blocks
@@ -298,11 +295,68 @@ export interface CLIControlPermissionRequest {
blocked_path: string | null;
}
/**
* Wire format for SDK MCP server config in initialization request.
* The actual Server instance stays in the SDK process.
*/
export interface SDKMcpServerConfig {
type: 'sdk';
name: string;
}
/**
* Wire format for external MCP server config in initialization request.
* Represents stdio/SSE/HTTP/TCP transports that must run in the CLI process.
*/
export interface CLIMcpServerConfig {
command?: string;
args?: string[];
env?: Record<string, string>;
cwd?: string;
url?: string;
httpUrl?: string;
headers?: Record<string, string>;
tcp?: string;
timeout?: number;
trust?: boolean;
description?: string;
includeTools?: string[];
excludeTools?: string[];
extensionName?: string;
oauth?: {
enabled?: boolean;
clientId?: string;
clientSecret?: string;
authorizationUrl?: string;
tokenUrl?: string;
scopes?: string[];
audiences?: string[];
redirectUri?: string;
tokenParamName?: string;
registrationUrl?: string;
};
authProviderType?:
| 'dynamic_discovery'
| 'google_credentials'
| 'service_account_impersonation';
targetAudience?: string;
targetServiceAccount?: string;
}
export interface CLIControlInitializeRequest {
subtype: 'initialize';
hooks?: HookRegistration[] | null;
sdkMcpServers?: Record<string, MCPServerConfig>;
mcpServers?: Record<string, MCPServerConfig>;
/**
* SDK MCP servers config
* These are MCP servers running in the SDK process, connected via control plane.
* External MCP servers are configured separately in settings, not via initialization.
*/
sdkMcpServers?: Record<string, Omit<SDKMcpServerConfig, 'instance'>>;
/**
* External MCP servers that the SDK wants the CLI to manage.
* These run outside the SDK process and require CLI-side transport setup.
*/
mcpServers?: Record<string, CLIMcpServerConfig>;
agents?: SubagentConfig[];
}