mirror of
https://github.com/toeverything/AFFiNE.git
synced 2026-07-18 10:36:22 +08:00
fix(server): schema preflight for tool call (#15172)
This commit is contained in:
@@ -1,3 +1,5 @@
|
||||
import { z } from 'zod';
|
||||
|
||||
import {
|
||||
type LlmBackendConfig,
|
||||
llmDispatchToolLoopStream,
|
||||
@@ -63,7 +65,7 @@ export function createToolExecutionCallback(
|
||||
};
|
||||
}
|
||||
|
||||
async function executeToolCall(
|
||||
export async function executeToolCall(
|
||||
tools: CopilotToolSet,
|
||||
request: LlmToolCallbackRequest,
|
||||
options: CopilotToolExecuteOptions
|
||||
@@ -103,7 +105,11 @@ async function executeToolCall(
|
||||
}
|
||||
|
||||
try {
|
||||
const output = await tool.execute(request.args, options);
|
||||
const args =
|
||||
tool.inputSchema instanceof z.ZodType
|
||||
? tool.inputSchema.parse(request.args)
|
||||
: request.args;
|
||||
const output = await tool.execute(args, options);
|
||||
return {
|
||||
callId: request.callId,
|
||||
name: request.name,
|
||||
@@ -173,3 +179,7 @@ export function createToolLoopBridge(
|
||||
);
|
||||
};
|
||||
}
|
||||
|
||||
// re-export for test consumers
|
||||
export type { LlmToolCallbackRequest } from '../../../../native';
|
||||
export type { CopilotToolSet, CopilotToolExecuteOptions } from '../../tools';
|
||||
|
||||
@@ -0,0 +1,111 @@
|
||||
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