File size: 3,859 Bytes
7a1ad33 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | /**
* @license
* Copyright 2025 Google LLC
* SPDX-License-Identifier: Apache-2.0
*/
import { describe, it, expect, vi, beforeEach } from 'vitest';
import { BaseToolInvocation, type ToolResult } from './tools.js';
import type { MessageBus } from '../confirmation-bus/message-bus.js';
import {
type Message,
MessageBusType,
type ToolConfirmationRequest,
type ToolConfirmationResponse,
} from '../confirmation-bus/types.js';
class TestBaseToolInvocation extends BaseToolInvocation<object, ToolResult> {
getDescription(): string {
return 'test description';
}
async execute(): Promise<ToolResult> {
return { llmContent: [], returnDisplay: '' };
}
}
describe('BaseToolInvocation', () => {
let messageBus: MessageBus;
let abortController: AbortController;
beforeEach(() => {
messageBus = {
publish: vi.fn(),
subscribe: vi.fn(),
unsubscribe: vi.fn(),
} as unknown as MessageBus;
abortController = new AbortController();
});
it('should propagate serverName to ToolConfirmationRequest', async () => {
const serverName = 'test-server';
const tool = new TestBaseToolInvocation(
{},
messageBus,
'test-tool',
'Test Tool',
serverName,
);
let capturedRequest: ToolConfirmationRequest | undefined;
vi.mocked(messageBus.publish).mockImplementation(
async (request: Message) => {
if (request.type === MessageBusType.TOOL_CONFIRMATION_REQUEST) {
capturedRequest = request;
}
},
);
let responseHandler:
| ((response: ToolConfirmationResponse) => void)
| undefined;
vi.mocked(messageBus.subscribe).mockImplementation(
(type: MessageBusType, handler: (message: Message) => void) => {
if (type === MessageBusType.TOOL_CONFIRMATION_RESPONSE) {
responseHandler = handler as (
response: ToolConfirmationResponse,
) => void;
}
},
);
const confirmationPromise = tool.shouldConfirmExecute(
abortController.signal,
);
// Wait for microtasks to ensure publish is called
await new Promise((resolve) => setTimeout(resolve, 0));
expect(messageBus.publish).toHaveBeenCalledTimes(1);
expect(capturedRequest).toBeDefined();
expect(capturedRequest?.type).toBe(
MessageBusType.TOOL_CONFIRMATION_REQUEST,
);
expect(capturedRequest?.serverName).toBe(serverName);
// Simulate response to finish the promise cleanly
if (responseHandler && capturedRequest) {
responseHandler({
type: MessageBusType.TOOL_CONFIRMATION_RESPONSE,
correlationId: capturedRequest.correlationId,
confirmed: true,
});
}
await confirmationPromise;
});
it('should NOT propagate serverName if not provided', async () => {
const tool = new TestBaseToolInvocation(
{},
messageBus,
'test-tool',
'Test Tool',
// no serverName
);
let capturedRequest: ToolConfirmationRequest | undefined;
vi.mocked(messageBus.publish).mockImplementation(
async (request: Message) => {
if (request.type === MessageBusType.TOOL_CONFIRMATION_REQUEST) {
capturedRequest = request;
}
},
);
// We need to mock subscribe to avoid hanging if we want to await the promise,
// but for this test we just need to check publish.
// We'll abort to clean up.
const confirmationPromise = tool.shouldConfirmExecute(
abortController.signal,
);
await new Promise((resolve) => setTimeout(resolve, 0));
expect(messageBus.publish).toHaveBeenCalledTimes(1);
expect(capturedRequest).toBeDefined();
expect(capturedRequest?.serverName).toBeUndefined();
abortController.abort();
try {
await confirmationPromise;
} catch {
// ignore abort error
}
});
});
|