Skip to content

Commit bd990cd

Browse files
ochafikochafik
committed
apply review suggestions
1 parent 92754c9 commit bd990cd

2 files changed

Lines changed: 152 additions & 1 deletion

File tree

‎src/server/index.test.ts‎

Lines changed: 144 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1771,6 +1771,150 @@ describe('createMessage validation', () => {
17711771
})
17721772
).resolves.toMatchObject({ model: 'test-model' });
17731773
});
1774+
1775+
test('should throw when user sends text instead of tool_result after tool_use', async () => {
1776+
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
1777+
1778+
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
1779+
1780+
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
1781+
model: 'test-model',
1782+
role: 'assistant',
1783+
content: { type: 'text', text: 'Response' }
1784+
}));
1785+
1786+
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
1787+
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
1788+
1789+
// User ignores tool_use and sends text instead
1790+
await expect(
1791+
server.createMessage({
1792+
messages: [
1793+
{ role: 'user', content: { type: 'text', text: 'hello' } },
1794+
{ role: 'assistant', content: { type: 'tool_use', id: 'call_1', name: 'test_tool', input: {} } },
1795+
{ role: 'user', content: { type: 'text', text: 'actually nevermind' } }
1796+
],
1797+
maxTokens: 100,
1798+
tools: [{ name: 'test_tool', inputSchema: { type: 'object' } }]
1799+
})
1800+
).rejects.toThrow('ids of tool_result blocks and tool_use blocks from previous message do not match');
1801+
});
1802+
1803+
test('should throw when only some tool_results are provided for parallel tool_use', async () => {
1804+
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
1805+
1806+
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
1807+
1808+
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
1809+
model: 'test-model',
1810+
role: 'assistant',
1811+
content: { type: 'text', text: 'Response' }
1812+
}));
1813+
1814+
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
1815+
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
1816+
1817+
// Parallel tool_use but only one tool_result provided
1818+
await expect(
1819+
server.createMessage({
1820+
messages: [
1821+
{ role: 'user', content: { type: 'text', text: 'hello' } },
1822+
{
1823+
role: 'assistant',
1824+
content: [
1825+
{ type: 'tool_use', id: 'call_1', name: 'tool_a', input: {} },
1826+
{ type: 'tool_use', id: 'call_2', name: 'tool_b', input: {} }
1827+
]
1828+
},
1829+
{ role: 'user', content: { type: 'tool_result', toolUseId: 'call_1', content: [] } }
1830+
],
1831+
maxTokens: 100,
1832+
tools: [
1833+
{ name: 'tool_a', inputSchema: { type: 'object' } },
1834+
{ name: 'tool_b', inputSchema: { type: 'object' } }
1835+
]
1836+
})
1837+
).rejects.toThrow('ids of tool_result blocks and tool_use blocks from previous message do not match');
1838+
});
1839+
1840+
test('should validate tool_use/tool_result even without tools in current request', async () => {
1841+
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
1842+
1843+
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
1844+
1845+
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
1846+
model: 'test-model',
1847+
role: 'assistant',
1848+
content: { type: 'text', text: 'Response' }
1849+
}));
1850+
1851+
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
1852+
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
1853+
1854+
// Previous request returned tool_use, now sending tool_result without tools param
1855+
await expect(
1856+
server.createMessage({
1857+
messages: [
1858+
{ role: 'user', content: { type: 'text', text: 'hello' } },
1859+
{ role: 'assistant', content: { type: 'tool_use', id: 'call_1', name: 'test_tool', input: {} } },
1860+
{ role: 'user', content: { type: 'tool_result', toolUseId: 'wrong_id', content: [] } }
1861+
],
1862+
maxTokens: 100
1863+
// Note: no tools param - this is a follow-up request after tool execution
1864+
})
1865+
).rejects.toThrow('ids of tool_result blocks and tool_use blocks from previous message do not match');
1866+
});
1867+
1868+
test('should allow valid tool_use/tool_result without tools in current request', async () => {
1869+
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
1870+
1871+
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: { tools: {} } } });
1872+
1873+
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
1874+
model: 'test-model',
1875+
role: 'assistant',
1876+
content: { type: 'text', text: 'Response' }
1877+
}));
1878+
1879+
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
1880+
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
1881+
1882+
// Previous request returned tool_use, now sending matching tool_result without tools param
1883+
await expect(
1884+
server.createMessage({
1885+
messages: [
1886+
{ role: 'user', content: { type: 'text', text: 'hello' } },
1887+
{ role: 'assistant', content: { type: 'tool_use', id: 'call_1', name: 'test_tool', input: {} } },
1888+
{ role: 'user', content: { type: 'tool_result', toolUseId: 'call_1', content: [] } }
1889+
],
1890+
maxTokens: 100
1891+
// Note: no tools param - this is a follow-up request after tool execution
1892+
})
1893+
).resolves.toMatchObject({ model: 'test-model' });
1894+
});
1895+
1896+
test('should handle empty messages array', async () => {
1897+
const server = new Server({ name: 'test server', version: '1.0' }, { capabilities: {} });
1898+
1899+
const client = new Client({ name: 'test client', version: '1.0' }, { capabilities: { sampling: {} } });
1900+
1901+
client.setRequestHandler(CreateMessageRequestSchema, async () => ({
1902+
model: 'test-model',
1903+
role: 'assistant',
1904+
content: { type: 'text', text: 'Response' }
1905+
}));
1906+
1907+
const [clientTransport, serverTransport] = InMemoryTransport.createLinkedPair();
1908+
await Promise.all([client.connect(clientTransport), server.connect(serverTransport)]);
1909+
1910+
// Empty messages array should not crash
1911+
await expect(
1912+
server.createMessage({
1913+
messages: [],
1914+
maxTokens: 100
1915+
})
1916+
).resolves.toMatchObject({ model: 'test-model' });
1917+
});
17741918
});
17751919

17761920
test('should respect log level for transport with sessionId', async () => {

‎src/server/index.ts‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -328,16 +328,22 @@ export class Server<
328328
}
329329

330330
async createMessage(params: CreateMessageRequest['params'], options?: RequestOptions) {
331+
// Capability check - only required when tools/toolChoice are provided
331332
if (params.tools || params.toolChoice) {
332333
if (!this._clientCapabilities?.sampling?.tools) {
333334
throw new Error('Client does not support sampling tools capability.');
334335
}
336+
}
335337

338+
// Message structure validation - always validate tool_use/tool_result pairs.
339+
// These may appear even without tools/toolChoice in the current request when
340+
// a previous sampling request returned tool_use and this is a follow-up with results.
341+
if (params.messages.length > 0) {
336342
const lastMessage = params.messages[params.messages.length - 1];
337343
const lastContent = Array.isArray(lastMessage.content) ? lastMessage.content : [lastMessage.content];
338344
const hasToolResults = lastContent.some(c => c.type === 'tool_result');
339345

340-
const previousMessage = params.messages[params.messages.length - 2];
346+
const previousMessage = params.messages.length > 1 ? params.messages[params.messages.length - 2] : undefined;
341347
const previousContent = previousMessage
342348
? Array.isArray(previousMessage.content)
343349
? previousMessage.content
@@ -363,6 +369,7 @@ export class Server<
363369
}
364370
}
365371
}
372+
366373
return this.request({ method: 'sampling/createMessage', params }, CreateMessageResultSchema, options);
367374
}
368375

0 commit comments

Comments
 (0)