Skip to content
Merged
3 changes: 3 additions & 0 deletions packages/a2a-server/src/utils/testing_utils.ts
Original file line number Diff line number Diff line change
Expand Up @@ -50,6 +50,9 @@ export function createMockConfig(
getEmbeddingModel: vi.fn().mockReturnValue('text-embedding-004'),
getSessionId: vi.fn().mockReturnValue('test-session-id'),
getUserTier: vi.fn(),
getEnableMessageBusIntegration: vi.fn().mockReturnValue(false),
getMessageBus: vi.fn(),
getPolicyEngine: vi.fn(),
...overrides,
} as unknown as Config;

Expand Down
3 changes: 3 additions & 0 deletions packages/cli/src/ui/hooks/useToolScheduler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,9 @@ const mockConfig = {
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getShellExecutionConfig: () => ({ terminalWidth: 80, terminalHeight: 24 }),
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const mockTool = new MockTool({
Expand Down
5 changes: 5 additions & 0 deletions packages/core/src/confirmation-bus/types.ts
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,11 @@ export interface ToolConfirmationResponse {
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE;
correlationId: string;
confirmed: boolean;
/**
* When true, indicates that policy decision was ASK_USER and the tool should
* show its legacy confirmation UI instead of auto-proceeding.
*/
requiresUserConfirmation?: boolean;
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Does this need to be optional?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The way this is written right now is that this only comes into play when the ASK_USER policy decision is set. Otherwise it's null and falls through to the legacy behviour.

}

export interface ToolPolicyRejection {
Expand Down
40 changes: 37 additions & 3 deletions packages/core/src/core/coreToolScheduler.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -255,6 +255,9 @@ describe('CoreToolScheduler', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -332,6 +335,9 @@ describe('CoreToolScheduler', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null,
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -365,15 +371,18 @@ describe('CoreToolScheduler', () => {
describe('getToolSuggestion', () => {
it('should suggest the top N closest tool names for a typo', () => {
// Create mocked tool registry
const mockToolRegistry = {
getAllToolNames: () => ['list_files', 'read_file', 'write_file'],
} as unknown as ToolRegistry;
const mockConfig = {
getToolRegistry: () => mockToolRegistry,
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;
const mockToolRegistry = {
getAllToolNames: () => ['list_files', 'read_file', 'write_file'],
} as unknown as ToolRegistry;

// Create scheduler
const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -448,6 +457,9 @@ describe('CoreToolScheduler with payload', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -768,6 +780,9 @@ describe('CoreToolScheduler edit cancellation', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -874,6 +889,9 @@ describe('CoreToolScheduler YOLO mode', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -981,6 +999,9 @@ describe('CoreToolScheduler request queueing', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -1113,6 +1134,9 @@ describe('CoreToolScheduler request queueing', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -1215,6 +1239,9 @@ describe('CoreToolScheduler request queueing', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -1287,6 +1314,9 @@ describe('CoreToolScheduler request queueing', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

const testTool = new TestApprovalTool(mockConfig);
Expand Down Expand Up @@ -1475,6 +1505,8 @@ describe('CoreToolScheduler Sequential Execution', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null,
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down Expand Up @@ -1595,6 +1627,8 @@ describe('CoreToolScheduler Sequential Execution', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null,
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
} as unknown as Config;

const scheduler = new CoreToolScheduler({
Expand Down
31 changes: 31 additions & 0 deletions packages/core/src/core/coreToolScheduler.ts
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,8 @@ import * as path from 'node:path';
import { doesToolInvocationMatch } from '../utils/tool-utils.js';
import levenshtein from 'fast-levenshtein';
import { ShellToolInvocation } from '../tools/shell.js';
import type { ToolConfirmationRequest } from '../confirmation-bus/types.js';
import { MessageBusType } from '../confirmation-bus/types.js';

export type ValidatingToolCall = {
status: 'validating';
Expand Down Expand Up @@ -352,6 +354,15 @@ export class CoreToolScheduler {
this.onToolCallsUpdate = options.onToolCallsUpdate;
this.getPreferredEditor = options.getPreferredEditor;
this.onEditorClose = options.onEditorClose;

// Subscribe to message bus for ASK_USER policy decisions
if (this.config.getEnableMessageBusIntegration()) {
const messageBus = this.config.getMessageBus();
messageBus.subscribe(
MessageBusType.TOOL_CONFIRMATION_REQUEST,
this.handleToolConfirmationRequest.bind(this),
);
}
}

private setStatusInternal(
Expand Down Expand Up @@ -1160,6 +1171,26 @@ export class CoreToolScheduler {
});
}

/**
* Handle tool confirmation requests from the message bus when policy decision is ASK_USER.
* This publishes a response with requiresUserConfirmation=true to signal the tool
* that it should fall back to its legacy confirmation UI.
*/
private handleToolConfirmationRequest(
request: ToolConfirmationRequest,
): void {
// When ASK_USER policy decision is made, the message bus emits the request here.
// We respond with requiresUserConfirmation=true to tell the tool to use its
// legacy confirmation flow (which will show diffs, URLs, etc in the UI).
const messageBus = this.config.getMessageBus();
messageBus.publish({
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
correlationId: request.correlationId,
confirmed: false, // Not auto-approved
requiresUserConfirmation: true, // Use legacy UI confirmation
});
}

private async autoApproveCompatiblePendingTools(
signal: AbortSignal,
triggeringCallId: string,
Expand Down
3 changes: 3 additions & 0 deletions packages/core/src/core/nonInteractiveToolExecutor.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,9 @@ describe('executeToolCall', () => {
getUseSmartEdit: () => false,
getUseModelRouter: () => false,
getGeminiClient: () => null, // No client needed for these tests
getEnableMessageBusIntegration: () => false,
getMessageBus: () => null,
getPolicyEngine: () => null,
} as unknown as Config;

abortController = new AbortController();
Expand Down
18 changes: 17 additions & 1 deletion packages/core/src/tools/message-bus-integration.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -52,6 +52,22 @@ class TestToolInvocation extends BaseToolInvocation<TestParams, TestResult> {
testValue: this.params.testParam,
};
}

override async shouldConfirmExecute(
abortSignal: AbortSignal,
): Promise<false> {
// This conditional is here to allow testing of the case where there is no message bus.
if (this.messageBus) {
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should this conditional always evaluate to true? Can we assert on it instead of making it a conditional?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The conditional supports the case where there is no message bus in the test. It's flag guarded so I wanted to test both scenarios. See the test on line 222 for example.

const decision = await this.getMessageBusDecision(abortSignal);
if (decision === 'ALLOW') {
return false;
}
if (decision === 'DENY') {
throw new Error('Tool execution denied by policy');
}
}
return false;
}
}

class TestTool extends BaseDeclarativeTool<TestParams, TestResult> {
Expand Down Expand Up @@ -200,7 +216,7 @@ describe('Message Bus Integration', () => {
abortController.abort();

await expect(confirmationPromise).rejects.toThrow(
'Tool confirmation aborted',
'Tool execution denied by policy',
);
});

Expand Down
Loading