mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-08-30 21:29:37 +08:00
fix(server): tool call test
This commit is contained in:
DarkSky
@@ -7,6 +7,7 @@ import type { TestFn } from 'ava';
|
||||
import ava from 'ava';
|
||||
import { nanoid } from 'nanoid';
|
||||
import Sinon from 'sinon';
|
||||
import { z } from 'zod';
|
||||
|
||||
import {
|
||||
EventBus,
|
||||
@@ -27,6 +28,7 @@ import {
|
||||
WorkspaceModel,
|
||||
WorkspaceRole,
|
||||
} from '../../models';
|
||||
import type { LlmToolCallbackRequest } from '../../native';
|
||||
import { CopilotModule } from '../../plugins/copilot';
|
||||
import { CopilotContextService } from '../../plugins/copilot/context';
|
||||
import { CopilotContextResolver } from '../../plugins/copilot/context/resolver';
|
||||
@@ -63,9 +65,14 @@ import { ImageResultHost } from '../../plugins/copilot/runtime/hosts/image-resul
|
||||
import { ModelSelectionPolicy } from '../../plugins/copilot/runtime/model-selection-policy';
|
||||
import { PromptRuntime } from '../../plugins/copilot/runtime/prompt-runtime';
|
||||
import { getProviderRuntimeHost } from '../../plugins/copilot/runtime/provider-runtime-context';
|
||||
import { executeToolCall } from '../../plugins/copilot/runtime/tool/bridge';
|
||||
import { TurnOrchestrator } from '../../plugins/copilot/runtime/turn-orchestrator';
|
||||
import { ChatSessionService } from '../../plugins/copilot/session';
|
||||
import { CopilotStorage } from '../../plugins/copilot/storage';
|
||||
import type {
|
||||
CopilotToolExecuteOptions,
|
||||
CopilotToolSet,
|
||||
} from '../../plugins/copilot/tools';
|
||||
import { CopilotTranscriptionService } from '../../plugins/copilot/transcript';
|
||||
import { CopilotWorkspaceService } from '../../plugins/copilot/workspace';
|
||||
import { PaymentModule } from '../../plugins/payment';
|
||||
@@ -1272,6 +1279,143 @@ const wrapAsyncIter = async <T>(iter: AsyncIterable<T>) => {
|
||||
return result;
|
||||
};
|
||||
|
||||
test('tool bridge should validate zod tool args before execution', async t => {
|
||||
const execute = Sinon.stub().resolves({ message: 'executed' });
|
||||
const tools: CopilotToolSet = {
|
||||
safeTool: {
|
||||
inputSchema: z.object({
|
||||
name: z.string(),
|
||||
}),
|
||||
execute,
|
||||
},
|
||||
};
|
||||
const options: CopilotToolExecuteOptions = {};
|
||||
const request: LlmToolCallbackRequest = {
|
||||
callId: 'call-1',
|
||||
name: 'safeTool',
|
||||
args: { name: 123 },
|
||||
rawArgumentsText: '{"name":123}',
|
||||
};
|
||||
|
||||
const response = await executeToolCall(tools, request, options);
|
||||
|
||||
t.true(response.isError);
|
||||
t.true(response.output instanceof Object);
|
||||
t.regex((response.output as { message: string }).message, /Expected string/);
|
||||
t.false(execute.called);
|
||||
});
|
||||
|
||||
test('tool bridge should execute with parsed zod args and preserve callback response metadata', async t => {
|
||||
const execute = Sinon.stub().resolves({ message: 'executed' });
|
||||
const tools: CopilotToolSet = {
|
||||
safeTool: {
|
||||
inputSchema: z.object({
|
||||
name: z.string().trim(),
|
||||
}),
|
||||
execute,
|
||||
},
|
||||
};
|
||||
const options: CopilotToolExecuteOptions = {};
|
||||
const request: LlmToolCallbackRequest = {
|
||||
callId: 'call-2',
|
||||
name: 'safeTool',
|
||||
args: { name: ' AFFiNE ' },
|
||||
rawArgumentsText: '{"name":" AFFiNE "}',
|
||||
};
|
||||
|
||||
const response = await executeToolCall(tools, request, options);
|
||||
|
||||
t.false(response.isError ?? false);
|
||||
t.deepEqual(response, {
|
||||
callId: 'call-2',
|
||||
name: 'safeTool',
|
||||
args: { name: ' AFFiNE ' },
|
||||
rawArgumentsText: '{"name":" AFFiNE "}',
|
||||
argumentParseError: undefined,
|
||||
output: { message: 'executed' },
|
||||
});
|
||||
t.deepEqual(execute.firstCall.args, [{ name: 'AFFiNE' }, options]);
|
||||
});
|
||||
|
||||
test('tool bridge should reject malformed or unknown tool calls without executing tools', async t => {
|
||||
const execute = Sinon.stub().resolves({ message: 'executed' });
|
||||
const tools: CopilotToolSet = {
|
||||
safeTool: {
|
||||
inputSchema: z.record(z.unknown()),
|
||||
execute,
|
||||
},
|
||||
};
|
||||
|
||||
const invalidJsonResponse = await executeToolCall(
|
||||
tools,
|
||||
{
|
||||
callId: 'call-3',
|
||||
name: 'safeTool',
|
||||
args: {},
|
||||
rawArgumentsText: '{"unterminated"',
|
||||
argumentParseError: 'Unexpected end of JSON input',
|
||||
},
|
||||
{}
|
||||
);
|
||||
const missingToolResponse = await executeToolCall(
|
||||
tools,
|
||||
{
|
||||
callId: 'call-4',
|
||||
name: 'missingTool',
|
||||
args: {},
|
||||
rawArgumentsText: '{}',
|
||||
},
|
||||
{}
|
||||
);
|
||||
|
||||
t.deepEqual(invalidJsonResponse, {
|
||||
callId: 'call-3',
|
||||
name: 'safeTool',
|
||||
args: {},
|
||||
rawArgumentsText: '{"unterminated"',
|
||||
argumentParseError: 'Unexpected end of JSON input',
|
||||
isError: true,
|
||||
output: {
|
||||
message: 'Invalid tool arguments JSON',
|
||||
rawArguments: '{"unterminated"',
|
||||
error: 'Unexpected end of JSON input',
|
||||
},
|
||||
});
|
||||
t.deepEqual(missingToolResponse, {
|
||||
callId: 'call-4',
|
||||
name: 'missingTool',
|
||||
args: {},
|
||||
rawArgumentsText: '{}',
|
||||
argumentParseError: undefined,
|
||||
isError: true,
|
||||
output: { message: 'Tool not found: missingTool' },
|
||||
});
|
||||
t.false(execute.called);
|
||||
});
|
||||
|
||||
test('tool bridge should not mutate global object prototype for adversarial args', async t => {
|
||||
const execute = Sinon.stub().resolves({ message: 'executed' });
|
||||
const tools: CopilotToolSet = {
|
||||
safeTool: {
|
||||
inputSchema: z.record(z.unknown()),
|
||||
execute,
|
||||
},
|
||||
};
|
||||
const request: LlmToolCallbackRequest = {
|
||||
callId: 'call-5',
|
||||
name: 'safeTool',
|
||||
args: JSON.parse('{"__proto__":{"polluted":"yes"}}'),
|
||||
rawArgumentsText: '{"__proto__":{"polluted":"yes"}}',
|
||||
};
|
||||
|
||||
const response = await executeToolCall(tools, request, {});
|
||||
|
||||
t.true(Object.prototype.hasOwnProperty.call(request.args, '__proto__'));
|
||||
t.false(response.isError ?? false);
|
||||
t.is((Object.prototype as Record<string, unknown>).polluted, undefined);
|
||||
t.deepEqual(execute.firstCall.args[0], {});
|
||||
});
|
||||
|
||||
test('action stream should expose successful text action result as message', t => {
|
||||
t.deepEqual(
|
||||
projectActionEventToChatEvent('message-1', {
|
||||
|
||||
@@ -1,111 +0,0 @@
|
||||
import { z } from 'zod';
|
||||
|
||||
import { executeToolCall } from './bridge';
|
||||
import type {
|
||||
CopilotToolSet,
|
||||
LlmToolCallbackRequest,
|
||||
CopilotToolExecuteOptions,
|
||||
} from './bridge';
|
||||
|
||||
describe('executeToolCall maintains security boundary under adversarial input', () => {
|
||||
// Give mockTool a real Zod inputSchema so the Zod-parse path is actually exercised.
|
||||
// z.record(z.unknown()) accepts any plain key-value record.
|
||||
const inputSchema = z.record(z.unknown());
|
||||
|
||||
const mockTool = {
|
||||
inputSchema,
|
||||
execute: jest.fn().mockResolvedValue({ message: 'executed' }),
|
||||
};
|
||||
|
||||
const mockTools: CopilotToolSet = {
|
||||
safeTool: mockTool,
|
||||
};
|
||||
|
||||
// CopilotToolExecuteOptions only supports { signal?, messages? } — no `context`
|
||||
const baseOptions: CopilotToolExecuteOptions = {};
|
||||
|
||||
beforeEach(() => {
|
||||
mockTool.execute.mockClear();
|
||||
});
|
||||
|
||||
const payloads: Array<{ description: string; request: LlmToolCallbackRequest }> = [
|
||||
{
|
||||
description: 'prototype pollution payload',
|
||||
request: {
|
||||
callId: '1',
|
||||
name: 'safeTool',
|
||||
args: JSON.parse('{"__proto__":{"polluted":"yes"}}'),
|
||||
rawArgumentsText: '{"__proto__":{"polluted":"yes"}}',
|
||||
argumentParseError: null,
|
||||
},
|
||||
},
|
||||
{
|
||||
description: 'excessive nested object',
|
||||
request: {
|
||||
callId: '2',
|
||||
name: 'safeTool',
|
||||
args: { a: { b: { c: { d: { e: { f: { g: 'deep' } } } } } } },
|
||||
rawArgumentsText: '{"a":{"b":{"c":{"d":{"e":{"f":{"g":"deep"}}}}}}}',
|
||||
argumentParseError: null,
|
||||
},
|
||||
},
|
||||
{
|
||||
description: 'valid minimal input',
|
||||
request: {
|
||||
callId: '3',
|
||||
name: 'safeTool',
|
||||
args: {},
|
||||
rawArgumentsText: '{}',
|
||||
argumentParseError: null,
|
||||
},
|
||||
},
|
||||
];
|
||||
|
||||
test.each(payloads)(
|
||||
'handles $description without corrupting execution context',
|
||||
async ({ request }) => {
|
||||
const originalObjectProto = Object.prototype;
|
||||
|
||||
const response = await executeToolCall(mockTools, request, baseOptions);
|
||||
|
||||
expect(response).toBeDefined();
|
||||
expect(Object.prototype).toBe(originalObjectProto);
|
||||
expect((Object.prototype as any).polluted).toBeUndefined();
|
||||
|
||||
// Capture the actual arguments passed to mockTool.execute
|
||||
const [[actualArgs, actualOptions]] = mockTool.execute.mock.calls;
|
||||
|
||||
// Identity assertion: args must NOT be the same reference as request.args
|
||||
// This catches regressions where executeToolCall forwards request.args directly
|
||||
expect(actualArgs).not.toBe(request.args);
|
||||
|
||||
// Correctness assertion: args should deeply equal the parsed result
|
||||
const parsedArgs = inputSchema.parse(request.args);
|
||||
expect(actualArgs).toEqual(parsedArgs);
|
||||
|
||||
// Identity assertion: options should be the same reference
|
||||
expect(actualOptions).toBe(baseOptions);
|
||||
}
|
||||
);
|
||||
|
||||
test('rejects tool calls whose args fail schema validation', async () => {
|
||||
const strictTool = {
|
||||
inputSchema: z.object({ name: z.string() }),
|
||||
execute: jest.fn().mockResolvedValue({ message: 'executed' }),
|
||||
};
|
||||
const strictTools: CopilotToolSet = { strictTool };
|
||||
|
||||
const request: LlmToolCallbackRequest = {
|
||||
callId: '4',
|
||||
name: 'strictTool',
|
||||
// invalid: `name` must be a string, not a number
|
||||
args: { name: 123 },
|
||||
rawArgumentsText: '{"name":123}',
|
||||
argumentParseError: null,
|
||||
};
|
||||
|
||||
const response = await executeToolCall(strictTools, request, baseOptions);
|
||||
expect(response.isError).toBe(true);
|
||||
expect(strictTool.execute).not.toHaveBeenCalled();
|
||||
});
|
||||
});
|
||||
Reference in New Issue
Block a user