Skip to content
36 changes: 36 additions & 0 deletions packages/cli/src/nonInteractive/control/ControlDispatcher.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,7 @@ import type {
CLIControlInterruptRequest,
CLIControlSetModelRequest,
CLIControlSupportedCommandsRequest,
CLIControlGetContextUsageRequest,
} from '../types.js';

/**
Expand Down Expand Up @@ -242,6 +243,41 @@ describe('ControlDispatcher', () => {
});
});

it('should route get_context_usage request to system controller', async () => {
const request: CLIControlRequest = {
type: 'control_request',
request_id: 'req-ctx',
request: {
subtype: 'get_context_usage',
show_details: false,
} as CLIControlGetContextUsageRequest,
};

const mockResponse = {
subtype: 'get_context_usage',
totalTokens: 1000,
};

vi.mocked(mockSystemController.handleRequest).mockResolvedValue(
mockResponse,
);

await dispatcher.dispatch(request);

expect(mockSystemController.handleRequest).toHaveBeenCalledWith(
request.request,
'req-ctx',
);
expect(mockContext.streamJson.send).toHaveBeenCalledWith({
type: 'control_response',
response: {
subtype: 'success',
request_id: 'req-ctx',
response: mockResponse,
},
});
});

it('should send error response when controller throws error', async () => {
const request: CLIControlRequest = {
type: 'control_request',
Expand Down
3 changes: 2 additions & 1 deletion packages/cli/src/nonInteractive/control/ControlDispatcher.ts
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
* which wraps these controllers with a stable programmatic API.
*
* Controllers:
* - SystemController: initialize, interrupt, set_model, supported_commands
* - SystemController: initialize, interrupt, set_model, supported_commands, get_context_usage
* - PermissionController: can_use_tool, set_permission_mode
* - SdkMcpController: mcp_server_status (mcp_message handled via callback)
* - HookController: hook_callback
Expand Down Expand Up @@ -380,6 +380,7 @@ export class ControlDispatcher implements IPendingRequestRegistry {
case 'interrupt':
case 'set_model':
case 'supported_commands':
case 'get_context_usage':
return this.systemController;

case 'can_use_tool':
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@ import type {
CLIControlInitializeRequest,
CLIControlSetModelRequest,
CLIMcpServerConfig,
CLIControlGetContextUsageRequest,
} from '../../types.js';
import { getAvailableCommands } from '../../../nonInteractiveCliCommands.js';
import {
Expand Down Expand Up @@ -61,11 +62,51 @@ export class SystemController extends BaseController {
case 'supported_commands':
return this.handleSupportedCommands(signal);

case 'get_context_usage':
return this.handleGetContextUsage(
payload as CLIControlGetContextUsageRequest,
signal,
);

default:
throw new Error(`Unsupported request subtype in SystemController`);
}
}

private async handleGetContextUsage(
payload: CLIControlGetContextUsageRequest,
signal: AbortSignal,
): Promise<Record<string, unknown>> {
if (signal.aborted) {
throw new Error('Request aborted');
}

try {
const mod = await import('../../../ui/commands/contextCommand.js');
if (typeof mod.collectContextData !== 'function') {
throw new Error('collectContextData is not available');
}
const showDetails = payload.show_details ?? false;
const contextUsageItem = await mod.collectContextData(
this.context.config,
showDetails,
);

return {
subtype: 'get_context_usage',
...contextUsageItem,
};
} catch (error) {
const errorMessage =
error instanceof Error ? error.message : 'Failed to get context usage';
debugLogger.error(
'[SystemController] Failed to get context usage:',
error,
);
throw new Error(errorMessage);
}
}

/**
* Handle initialize request
*
Expand Down Expand Up @@ -212,6 +253,7 @@ export class SystemController extends BaseController {
can_set_permission_mode:
typeof this.context.config.setApprovalMode === 'function',
can_set_model: typeof this.context.config.setModel === 'function',
can_get_context_usage: true,
// SDK MCP servers are supported - messages routed through control plane
can_handle_mcp_message: true,
};
Expand Down
8 changes: 7 additions & 1 deletion packages/cli/src/nonInteractive/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -407,6 +407,11 @@ export interface CLIControlSupportedCommandsRequest {
subtype: 'supported_commands';
}

export interface CLIControlGetContextUsageRequest {
subtype: 'get_context_usage';
show_details?: boolean;
}

export type ControlRequestPayload =
| CLIControlInterruptRequest
| CLIControlPermissionRequest
Expand All @@ -416,7 +421,8 @@ export type ControlRequestPayload =
| CLIControlMcpMessageRequest
| CLIControlSetModelRequest
| CLIControlMcpStatusRequest
| CLIControlSupportedCommandsRequest;
| CLIControlSupportedCommandsRequest
| CLIControlGetContextUsageRequest;

export interface CLIControlRequest {
type: 'control_request';
Expand Down
2 changes: 2 additions & 0 deletions packages/cli/src/nonInteractiveCliCommands.ts
Original file line number Diff line number Diff line change
Expand Up @@ -37,13 +37,15 @@ const debugLogger = createDebugLogger('NON_INTERACTIVE_COMMANDS');
* - init: Initialize project configuration
* - summary: Generate session summary
* - compress: Compress conversation history
* - context: Show context window usage (read-only diagnostic)
*/
export const ALLOWED_BUILTIN_COMMANDS_NON_INTERACTIVE = [
'init',
'summary',
'compress',
'btw',
'bug',
'context',
] as const;

/**
Expand Down
Loading
Loading