diff --git a/packages/api/src/agent-side/tools.ts b/packages/api/src/agent-side/tools.ts index 42ce3166..a95a9921 100644 --- a/packages/api/src/agent-side/tools.ts +++ b/packages/api/src/agent-side/tools.ts @@ -27,7 +27,7 @@ import type {Connector, DataRecord, RecordPage, RecordTarget, RecordValues} from import {parseRecordValues} from '../records/record-values'; import {loadConnection} from '../records/target'; import {integrationOf} from '../utils/connection-actions'; -import {requireOpenTaskIfAny} from '../mcp/pre-task-clients'; +import {requireOpenTask} from './tasks'; import {isPlainObject, parseToolParams} from './tool-params'; /** @@ -114,8 +114,7 @@ export function listAgentTools(signedIn: SignedInAgent, connectionId: string): P export type AgentCaller = {agent: SignedInAgent; via: 'web' | 'mcp'}; /** Who runs a tool, and the task it is run for, by its id. */ -// PRE-TASK CLIENTS: null for a pre-task client's call (mcp/pre-task-clients.ts). -export type TaskCaller = AgentCaller & {taskId: string | null}; +export type TaskCaller = AgentCaller & {taskId: string}; // `actions` are what the agent can do with the connection, and `refuse` why it can't do another. type ToolRun = {caller: TaskCaller; target: RecordTarget; connector: Connector; tool: CollectionTool; actions: string[]; refuse: (actionId: string) => ApiError}; @@ -152,7 +151,7 @@ function log(run: ToolRun, outcome: 'allowed' | 'denied', title: string | null, export function runAgentTool(caller: TaskCaller, connectionId: string, toolName: string, params: unknown, options: {readOnly?: boolean} = {}): Promise> { return Do(async ($) => { // None of the three needs another's answer, so they are read at once; a refusal still comes in this order. - const [agentRead, taskRead, connectionRead] = await Promise.all([requireAgent(caller.agent), requireOpenTaskIfAny(caller.agent, caller.taskId), loadConnection(caller.agent.orgId, connectionId)]); + const [agentRead, taskRead, connectionRead] = await Promise.all([requireAgent(caller.agent), requireOpenTask(caller.agent, caller.taskId), loadConnection(caller.agent.orgId, connectionId)]); const agent = $(agentRead); $(taskRead); const {connection, integration, connector} = $(connectionRead); diff --git a/packages/api/src/mcp/pre-task-clients.ts b/packages/api/src/mcp/pre-task-clients.ts deleted file mode 100644 index ab0999d8..00000000 --- a/packages/api/src/mcp/pre-task-clients.ts +++ /dev/null @@ -1,36 +0,0 @@ -import type {IncomingHttpHeaders} from 'node:http'; - -import {Ok, type Result} from 'ts-results-es'; - -import type {ApiError} from '@proxy/utils'; -import type {SignedInAgent} from '@proxy/db/agent'; -import type {AgentTaskRow} from '@proxy/db/task'; - -import {requireOpenTask} from '../agent-side/tasks'; - -/** - * PRE-TASK CLIENTS: a failsafe for the move to tasks, to be removed once every client is on them. - * Everything that uses it is marked `PRE-TASK CLIENTS`. - * - * An MCP client whose session began before every call took a task holds the tool list from then, - * which has no `start_task`, so it can't make a call for a task. `initialize` hands out a session - * id now, so a request without one is from a session begun before. Its calls of list_connections, - * list_tools, read and run without a task run for none, logged as the entries from before tasks - * are, and their answer says first that the tool list changed, so the client fetches the one with - * start_task. - */ - -const TASK_TOOLS = ['list_connections', 'list_tools', 'read', 'run']; - -/** Whether a tool call is a pre-task client's: no session, and no task for a tool that takes one. */ -export function isPreTaskCall(headers: IncomingHttpHeaders, name: string, args: Record): boolean { - return headers['mcp-session-id'] === undefined && TASK_TOOLS.includes(name) && (typeof args.task !== 'string' || args.task === ''); -} - -/** The task a call is for, still under way; none for a pre-task client's call. */ -export function requireOpenTaskIfAny(agent: SignedInAgent, taskId: string | null): Promise> { - if (taskId === null) { - return Promise.resolve(Ok(null)); - } - return requireOpenTask(agent, taskId); -} diff --git a/packages/api/src/mcp/route.test.ts b/packages/api/src/mcp/route.test.ts index c76e1b0f..527cbb5b 100644 --- a/packages/api/src/mcp/route.test.ts +++ b/packages/api/src/mcp/route.test.ts @@ -67,12 +67,6 @@ async function post(body: unknown, headers: Record = {authorizat return res; } -// The headers of a client on a session `initialize` began, as every client is once on tasks. -async function onSession(): Promise> { - const initialized = await post(rpc('initialize', {protocolVersion: '2025-11-25'})); - return {authorization: 'Bearer token-1', 'mcp-session-id': initialized.headers['mcp-session-id'] ?? ''}; -} - function rpc(method: string, params: Record = {}, id = 1) { return {jsonrpc: '2.0', id, method, params}; } @@ -161,9 +155,7 @@ describe('handleMcpRoute', () => { }); test('a call without a task is refused, saying to start one; one for a task that is not the agent’s or is finished, as the agent side says', async () => { - const headers = await onSession(); - listAgentConnections.mockClear(); - const taskless = await post(rpc('tools/call', {name: 'list_connections', arguments: {}}), headers); + const taskless = await post(rpc('tools/call', {name: 'list_connections', arguments: {}})); requireOpenTask.mockResolvedValue(Err(ApiErr.notFound('task', 'task-9'))); const unknown = await post(rpc('tools/call', {name: 'list_tools', arguments: {task: 'task-9', connection: 'mail-1'}})); requireOpenTask.mockResolvedValue(Err(ApiErr.conflict('Task task-1 (Triage inbox) is finished. Start a new task for what you are doing now, and give its id as task.'))); @@ -182,7 +174,7 @@ describe('handleMcpRoute', () => { runAgentTool.mockResolvedValue(Ok({record: null})); await post(rpc('tools/call', {name: 'read', arguments: {task: 'task-1', connection: 'mail-1', tool: 'emails_list', params: {folder: 'Inbox'}}})); await post(rpc('tools/call', {name: 'run', arguments: {task: 'task-1', connection: 'mail-1', tool: 'emails_archive', params: {id: 'inbox-1'}}})); - const taskless = await post(rpc('tools/call', {name: 'run', arguments: {connection: 'mail-1', tool: 'emails_archive', params: {id: 'inbox-1'}}}), await onSession()); + const taskless = await post(rpc('tools/call', {name: 'run', arguments: {connection: 'mail-1', tool: 'emails_archive', params: {id: 'inbox-1'}}})); expect(runAgentTool.mock.calls).toEqual([ [{agent: signedIn, via: 'mcp', taskId: 'task-1'}, 'mail-1', 'emails_list', {folder: 'Inbox'}, {readOnly: true}], @@ -243,23 +235,6 @@ describe('handleMcpRoute', () => { expect(listing.body).toMatchObject({id: 1, result: {tools: expect.any(Array)}}); }); - // PRE-TASK CLIENTS (pre-task-clients.ts) - test('a call without a task from a session begun before tasks runs for none, told the tool list changed', async () => { - runAgentTool.mockResolvedValue(Ok({record: null})); - const streams = {authorization: 'Bearer token-1', accept: 'application/json, text/event-stream'}; - const listed = await post(rpc('tools/call', {name: 'list_connections', arguments: {}}), streams); - const ran = await post(rpc('tools/call', {name: 'run', arguments: {connection: 'mail-1', tool: 'emails_archive', params: {id: 'inbox-1'}}}), {authorization: 'Bearer token-1'}); - const withTask = await post(rpc('tools/call', {name: 'list_connections', arguments: {task: 'task-1'}}), streams); - - expect(listed.headers['content-type']).toBe('text/event-stream'); - expect(String(listed.body)).toStartWith('event: message\ndata: {"jsonrpc":"2.0","method":"notifications/tools/list_changed"}\n\n'); - expect(String(listed.body)).toContain('"connection":"mail-1"'); - expect(runAgentTool).toHaveBeenCalledWith({agent: signedIn, via: 'mcp', taskId: null}, 'mail-1', 'emails_archive', {id: 'inbox-1'}, {readOnly: false}); - expect(ran.body).toMatchObject({result: {structuredContent: {record: null}}}); - expect(requireOpenTask.mock.calls).toEqual([[signedIn, 'task-1']]); - expect(withTask.body).toMatchObject({result: {structuredContent: {connections: [{connection: 'mail-1'}]}}}); - }); - test('an MCP-Protocol-Version it does not know is refused', async () => { const res = await post(rpc('ping'), {authorization: 'Bearer token-1', 'mcp-protocol-version': '1999-01-01'}); diff --git a/packages/api/src/mcp/route.ts b/packages/api/src/mcp/route.ts index 5289cedd..4b04b35a 100644 --- a/packages/api/src/mcp/route.ts +++ b/packages/api/src/mcp/route.ts @@ -1,5 +1,4 @@ import {createHash, randomUUID} from 'node:crypto'; -import type {IncomingHttpHeaders} from 'node:http'; import type {Request, Response} from 'express'; import {z} from 'zod'; @@ -9,7 +8,6 @@ import type {SignedInAgent} from '@proxy/db/agent'; import {authenticateBearer, wwwAuthenticate} from '../oauth/authenticate'; import {requireAppUrl} from '../oauth/oauth-config'; import {log, serializeError} from '../observability/log'; -import {isPreTaskCall} from './pre-task-clients'; import {callMcpTool, MCP_TOOLS, serverInstructions} from './tools'; /** @@ -62,7 +60,7 @@ async function initialize(signedIn: SignedInAgent, params: Record { +async function answer(signedIn: SignedInAgent, message: unknown): Promise { const parsed = messageSchema.safeParse(message); if (!parsed.success) { return failure(null, -32600, 'Invalid request'); @@ -85,10 +83,7 @@ async function answer(signedIn: SignedInAgent, message: unknown, headers: Incomi if (!call.success || !MCP_TOOLS.some((tool) => tool.name === call.data.name)) { return failure(id, -32602, `Unknown tool. The tools are ${MCP_TOOLS.map((tool) => tool.name).join(', ')}.`); } - const args = call.data.arguments ?? {}; - // PRE-TASK CLIENTS: whether the call is a pre-task client's (pre-task-clients.ts). - const preTask = isPreTaskCall(headers, call.data.name, args); - return {jsonrpc: '2.0', id, result: await callMcpTool(signedIn, call.data.name, args, preTask)}; + return {jsonrpc: '2.0', id, result: await callMcpTool(signedIn, call.data.name, call.data.arguments ?? {})}; } return failure(id, -32601, `Method not found: ${method}`); } @@ -121,7 +116,7 @@ export async function handleMcpRoute(req: Request, res: Response): Promise const messages: unknown[] = Array.isArray(body) ? body : [body]; const answers: JsonRpcResponse[] = []; for (const message of messages) { - const reply = await answer(signedIn.value, message, req.headers); + const reply = await answer(signedIn.value, message); if (reply) { answers.push(reply); } @@ -133,7 +128,7 @@ export async function handleMcpRoute(req: Request, res: Response): Promise res.status(202).end(); return; } - if (holdsOldTools(req, messages)) { + if (holdsOldTools(req) && messages.some((message) => isMethod(message, 'tools/call'))) { res.setHeader('Content-Type', 'text/event-stream').setHeader('Cache-Control', 'no-cache'); res.end([LIST_CHANGED, ...answers].map((event) => `event: message\ndata: ${JSON.stringify(event)}\n\n`).join('')); return; @@ -145,32 +140,15 @@ function isMethod(message: unknown, method: string): boolean { return messageSchema.safeParse(message).data?.method === method; } -function toolCalls(messages: unknown[]): Array> { - return messages.flatMap((message) => { - const parsed = messageSchema.safeParse(message); - if (parsed.data?.method !== 'tools/call') { - return []; - } - const call = callParamsSchema.safeParse(parsed.data.params ?? {}); - return call.success ? [call.data] : []; - }); -} - -// Whether the request calls a tool from a session from before the tool list last changed, and -// the client can take a stream to be told so. A client with no session, or one that only takes -// JSON, can't be told. -function holdsOldTools(req: Request, messages: unknown[]): boolean { - const calls = toolCalls(messages); +// Whether the client's session is from before the tool list last changed, and it can take a +// stream to be told so. A client with no session, or one that only takes JSON, can't be told. +function holdsOldTools(req: Request): boolean { + const session = req.headers['mcp-session-id']; const accept = req.headers.accept ?? ''; - if (calls.length === 0 || !accept.includes('text/event-stream')) { + if (typeof session !== 'string' || !accept.includes('text/event-stream')) { return false; } - // PRE-TASK CLIENTS: a pre-task client's call is told too, so it fetches the list with start_task. - if (calls.some((call) => isPreTaskCall(req.headers, call.name, call.arguments ?? {}))) { - return true; - } - const session = req.headers['mcp-session-id']; - return typeof session === 'string' && !session.endsWith(`.${TOOLS_VERSION}`); + return !session.endsWith(`.${TOOLS_VERSION}`); } /** `GET` and `DELETE /mcp`: there is no stream to open and no session to end. */ diff --git a/packages/api/src/mcp/tools.ts b/packages/api/src/mcp/tools.ts index 2a0bd1c3..e2cc5a3a 100644 --- a/packages/api/src/mcp/tools.ts +++ b/packages/api/src/mcp/tools.ts @@ -6,11 +6,10 @@ import type {SignedInAgent} from '@proxy/db/agent'; import {findIntegration, INFO_INTEGRATION_ID, integrationFor, INTEGRATIONS, isOffered, type Tool} from '@proxy/integrations'; import {sendAgentFeedback} from '../agent-side/feedback'; -import {finishAgentTask, startAgentTask} from '../agent-side/tasks'; +import {finishAgentTask, requireOpenTask, startAgentTask} from '../agent-side/tasks'; import {listAgentConnections, listAgentTools, runAgentTool, type AgentConnection} from '../agent-side/tools'; import {requireAppUrl} from '../oauth/oauth-config'; import {log, serializeError} from '../observability/log'; -import {requireOpenTaskIfAny} from './pre-task-clients'; /** * The MCP server's tools: seven that stay the same whatever an agent may do, so a client's tool @@ -220,8 +219,7 @@ function errorMessage(error: ApiError, args: Record): string { return 'Something went wrong. Try again in a moment.'; } -// PRE-TASK CLIENTS: `preTask` is whether the call is a pre-task client's (pre-task-clients.ts). -function call(signedIn: SignedInAgent, name: string, args: Record, preTask: boolean): Promise, ApiError>> { +function call(signedIn: SignedInAgent, name: string, args: Record): Promise, ApiError>> { return Do(async ($) => { const caller = {agent: signedIn, via: 'mcp'} as const; if (name === 'start_task') { @@ -235,16 +233,15 @@ function call(signedIn: SignedInAgent, name: string, args: Record, preTask: boolean): Promise { - const result = await call(signedIn, name, args, preTask); +export async function callMcpTool(signedIn: SignedInAgent, name: string, args: Record): Promise { + const result = await call(signedIn, name, args); if (result.isErr()) { return {content: [{type: 'text', text: errorMessage(result.error, args)}], isError: true}; } diff --git a/packages/db/src/audit.ts b/packages/db/src/audit.ts index c334f608..b66a24f3 100644 --- a/packages/db/src/audit.ts +++ b/packages/db/src/audit.ts @@ -34,8 +34,7 @@ export async function logAgentRequest(data: { query: string | null; outcome: AuditOutcome; via: AuditVia; - // PRE-TASK CLIENTS: null for a call of an MCP client from before tasks (api's mcp/pre-task-clients.ts). - taskId: string | null; + taskId: string; }): Promise> { return wrapDb(async () => { await db.$transaction([