sampling: validate tools, tool_use, tool_result constraints (#1156)
This commit is contained in:
Olivier Chafik
GitHub
parent
41c6b35e07
commit
2a55dfd5f2
1 changed files with 380 additions and 1 deletions
|
|
@ -1582,6 +1582,341 @@ test('should respect log level for transport without sessionId', async () => {
|
|||
expect(clientTransport.onmessage).toHaveBeenCalled();
|
||||
});
|
||||
|
||||
describe('createMessage validation', () => {
|
||||
test('should throw when tools are provided without sampling.tools capability', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client(
|
||||
{ name: 'test client', version: '1.0' },
|
||||
{ capabilities: { sampling: {} } } // No tools capability
|
||||
);
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [{ role: 'user', content: { type: 'text', text: 'hello' } }],
|
||||
maxTokens: 100,
|
||||
tools: [{ name: 'test_tool', inputSchema: { type: 'object' } }]
|
||||
})
|
||||
).rejects.toThrow('Client does not support sampling tools capability.');
|
||||
});
|
||||
|
||||
test('should throw when toolChoice is provided without sampling.tools capability', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client(
|
||||
{ name: 'test client', version: '1.0' },
|
||||
{ capabilities: { sampling: {} } } // No tools capability
|
||||
);
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [{ role: 'user', content: { type: 'text', text: 'hello' } }],
|
||||
maxTokens: 100,
|
||||
toolChoice: { mode: 'auto' }
|
||||
})
|
||||
).rejects.toThrow('Client does not support sampling tools capability.');
|
||||
});
|
||||
|
||||
test('should throw when tool_result is mixed with other content', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [
|
||||
{ role: 'user', content: { type: 'text', text: 'hello' } },
|
||||
{ role: 'assistant', content: { type: 'tool_use', id: 'call_1', name: 'test_tool', input: {} } },
|
||||
{
|
||||
role: 'user',
|
||||
content: [
|
||||
{ type: 'tool_result', toolUseId: 'call_1', content: [] },
|
||||
{ type: 'text', text: 'mixed content' } // Mixed!
|
||||
]
|
||||
}
|
||||
],
|
||||
maxTokens: 100,
|
||||
tools: [{ name: 'test_tool', inputSchema: { type: 'object' } }]
|
||||
})
|
||||
).rejects.toThrow('The last message must contain only tool_result content if any is present');
|
||||
});
|
||||
|
||||
test('should throw when tool_result has no matching tool_use in previous message', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
// tool_result without previous tool_use
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [
|
||||
{ role: 'user', content: { type: 'text', text: 'hello' } },
|
||||
{ role: 'user', content: { type: 'tool_result', toolUseId: 'call_1', content: [] } }
|
||||
],
|
||||
maxTokens: 100,
|
||||
tools: [{ name: 'test_tool', inputSchema: { type: 'object' } }]
|
||||
})
|
||||
).rejects.toThrow('tool_result blocks are not matching any tool_use from the previous message');
|
||||
});
|
||||
|
||||
test('should throw when tool_result IDs do not match tool_use IDs', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [
|
||||
{ role: 'user', content: { type: 'text', text: 'hello' } },
|
||||
{ role: 'assistant', content: { type: 'tool_use', id: 'call_1', name: 'test_tool', input: {} } },
|
||||
{ role: 'user', content: { type: 'tool_result', toolUseId: 'wrong_id', content: [] } }
|
||||
],
|
||||
maxTokens: 100,
|
||||
tools: [{ name: 'test_tool', inputSchema: { type: 'object' } }]
|
||||
})
|
||||
).rejects.toThrow('ids of tool_result blocks and tool_use blocks from previous message do not match');
|
||||
});
|
||||
|
||||
test('should allow text-only messages with tools (no tool_results)', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [{ role: 'user', content: { type: 'text', text: 'hello' } }],
|
||||
maxTokens: 100,
|
||||
tools: [{ name: 'test_tool', inputSchema: { type: 'object' } }]
|
||||
})
|
||||
).resolves.toMatchObject({ model: 'test-model' });
|
||||
});
|
||||
|
||||
test('should allow valid matching tool_result/tool_use IDs', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [
|
||||
{ role: 'user', content: { type: 'text', text: 'hello' } },
|
||||
{ role: 'assistant', content: { type: 'tool_use', id: 'call_1', name: 'test_tool', input: {} } },
|
||||
{ role: 'user', content: { type: 'tool_result', toolUseId: 'call_1', content: [] } }
|
||||
],
|
||||
maxTokens: 100,
|
||||
tools: [{ name: 'test_tool', inputSchema: { type: 'object' } }]
|
||||
})
|
||||
).resolves.toMatchObject({ model: 'test-model' });
|
||||
});
|
||||
|
||||
test('should throw when user sends text instead of tool_result after tool_use', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
// User ignores tool_use and sends text instead
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [
|
||||
{ role: 'user', content: { type: 'text', text: 'hello' } },
|
||||
{ role: 'assistant', content: { type: 'tool_use', id: 'call_1', name: 'test_tool', input: {} } },
|
||||
{ role: 'user', content: { type: 'text', text: 'actually nevermind' } }
|
||||
],
|
||||
maxTokens: 100,
|
||||
tools: [{ name: 'test_tool', inputSchema: { type: 'object' } }]
|
||||
})
|
||||
).rejects.toThrow('ids of tool_result blocks and tool_use blocks from previous message do not match');
|
||||
});
|
||||
|
||||
test('should throw when only some tool_results are provided for parallel tool_use', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
// Parallel tool_use but only one tool_result provided
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [
|
||||
{ role: 'user', content: { type: 'text', text: 'hello' } },
|
||||
{
|
||||
role: 'assistant',
|
||||
content: [
|
||||
{ type: 'tool_use', id: 'call_1', name: 'tool_a', input: {} },
|
||||
{ type: 'tool_use', id: 'call_2', name: 'tool_b', input: {} }
|
||||
]
|
||||
},
|
||||
{ role: 'user', content: { type: 'tool_result', toolUseId: 'call_1', content: [] } }
|
||||
],
|
||||
maxTokens: 100,
|
||||
tools: [
|
||||
{ name: 'tool_a', inputSchema: { type: 'object' } },
|
||||
{ name: 'tool_b', inputSchema: { type: 'object' } }
|
||||
]
|
||||
})
|
||||
).rejects.toThrow('ids of tool_result blocks and tool_use blocks from previous message do not match');
|
||||
});
|
||||
|
||||
test('should validate tool_use/tool_result even without tools in current request', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
// Previous request returned tool_use, now sending tool_result without tools param
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [
|
||||
{ role: 'user', content: { type: 'text', text: 'hello' } },
|
||||
{ role: 'assistant', content: { type: 'tool_use', id: 'call_1', name: 'test_tool', input: {} } },
|
||||
{ role: 'user', content: { type: 'tool_result', toolUseId: 'wrong_id', content: [] } }
|
||||
],
|
||||
maxTokens: 100
|
||||
// Note: no tools param - this is a follow-up request after tool execution
|
||||
})
|
||||
).rejects.toThrow('ids of tool_result blocks and tool_use blocks from previous message do not match');
|
||||
});
|
||||
|
||||
test('should allow valid tool_use/tool_result without tools in current request', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
// Previous request returned tool_use, now sending matching tool_result without tools param
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [
|
||||
{ role: 'user', content: { type: 'text', text: 'hello' } },
|
||||
{ role: 'assistant', content: { type: 'tool_use', id: 'call_1', name: 'test_tool', input: {} } },
|
||||
{ role: 'user', content: { type: 'tool_result', toolUseId: 'call_1', content: [] } }
|
||||
],
|
||||
maxTokens: 100
|
||||
// Note: no tools param - this is a follow-up request after tool execution
|
||||
})
|
||||
).resolves.toMatchObject({ model: 'test-model' });
|
||||
});
|
||||
|
||||
test('should handle empty messages array', async () => {
|
||||
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
|
||||
|
||||
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: {} } });
|
||||
|
||||
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
|
||||
model: 'test-model',
|
||||
role: 'assistant',
|
||||
content: { type: 'text', text: 'Response' }
|
||||
}));
|
||||
|
||||
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
|
||||
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
|
||||
|
||||
// Empty messages array should not crash
|
||||
await expect(
|
||||
server.createMessage({
|
||||
messages: [],
|
||||
maxTokens: 100
|
||||
})
|
||||
).resolves.toMatchObject({ model: 'test-model' });
|
||||
});
|
||||
});
|
||||
|
||||
test('should respect log level for transport with sessionId', async () => {
|
||||
const server = new Server(
|
||||
{
|
||||
|
|
|
|||
Loading…
Reference in a new issue