diff --git a/tests/mcp/relay.test.ts b/tests/mcp/relay.test.ts index b6d861eb..bbdfba8c 100644 --- a/tests/mcp/relay.test.ts +++ b/tests/mcp/relay.test.ts @@ -19,17 +19,20 @@ let running = false; let tickets: Map; let addresses: string[]; let throttled: boolean; +/** How long the fake API takes to know a token. */ +let verifyMs: number; const open: WebSocket[] = []; async function start(extra: Partial[0]> = {}): Promise { if (running) await relay.close(); relay = await startRelay({ api: { - verifyToken: (token: string, address: string): Promise => { + verifyToken: async (token: string, address: string): Promise => { addresses.push(address); - if (throttled) return Promise.resolve({ status: 429, retryAfter: '600', error: 'Too many' }); + if (verifyMs > 0) await new Promise((resolve) => setTimeout(resolve, verifyMs)); + if (throttled) return { status: 429, retryAfter: '600', error: 'Too many' }; const user = TOKENS[token]; - return Promise.resolve(user ? { status: 200, user } : { status: 401 }); + return user ? { status: 200, user } : { status: 401 }; }, redeemTicket: (ticket: string): Promise => { const user = tickets.get(ticket); @@ -48,6 +51,7 @@ beforeEach(async () => { tickets = new Map(); addresses = []; throttled = false; + verifyMs = 0; await start(); }); @@ -414,14 +418,59 @@ describe('tabs', () => { it('learn how many agents their account has on, and an agent gone quiet counts out', async () => { await start({ agentRecentMs: 150, pingMs: 50 }); - await initialize('fwa_alice'); + await call('fwa_alice', await initialize('fwa_alice'), 'get_view'); await new Promise((resolve) => setTimeout(resolve, 200)); const tab = await onTab('alice'); expect((await tab.next('agents')).count).toBe(0); - await initialize('fwa_alice'); + await call('fwa_alice', await initialize('fwa_alice'), 'get_view'); expect((await tab.next('agents')).count).toBe(1); expect((await tab.next('agents')).count).toBe(0); }); + + it('do not count the spare sessions a client opens and never calls a tool in', async () => { + const tab = await onTab('alice'); + expect((await tab.next('agents')).count).toBe(0); + await initialize('fwa_alice'); + await initialize('fwa_alice'); + await call('fwa_alice', await initialize('fwa_alice'), 'get_view'); + expect((await tab.next('agents')).count).toBe(1); + expect(tab.received.filter((message) => message.type === 'agents')).toEqual([]); + }); + + it('count out when the event stream of an agent drops, without a DELETE', async () => { + await start({ agentRecentMs: 150, pingMs: 50 }); + const tab = await onTab('alice'); + expect((await tab.next('agents')).count).toBe(0); + const session = await initialize('fwa_alice'); + const gone = new AbortController(); + const stream = await fetch(url(), { + headers: { Authorization: 'Bearer fwa_alice', 'Mcp-Session-Id': session, Accept: 'text/event-stream' }, + signal: gone.signal, + }); + expect(stream.status).toBe(200); + expect((await tab.next('agents')).count).toBe(1); + await new Promise((resolve) => setTimeout(resolve, 400)); + expect(tab.received.filter((message) => message.type === 'agents')).toEqual([]); + gone.abort(); + expect((await tab.next('agents')).count).toBe(0); + }); + + it('do not count an agent that drops its event stream while its token is still being checked', async () => { + const tab = await onTab('alice'); + expect((await tab.next('agents')).count).toBe(0); + const session = await initialize('fwa_alice'); + verifyMs = 100; + const gone = new AbortController(); + const stream = fetch(url(), { + headers: { Authorization: 'Bearer fwa_alice', 'Mcp-Session-Id': session, Accept: 'text/event-stream' }, + signal: gone.signal, + }).catch(() => null); + await new Promise((resolve) => setTimeout(resolve, 30)); + gone.abort(); + await stream; + await new Promise((resolve) => setTimeout(resolve, 200)); + expect(tab.received.filter((message) => message.type === 'agents')).toEqual([]); + }); }); describe('the view', () => { diff --git a/tools/mcp/relay.mjs b/tools/mcp/relay.mjs index 6b3bdf31..cd96459a 100644 --- a/tools/mcp/relay.mjs +++ b/tools/mcp/relay.mjs @@ -223,10 +223,11 @@ export async function startRelay(options) { subscriptions: new Set(), streams: new Set(), seenAt: Date.now(), + calledAt: 0, lastSummary: null, }; sessions.set(session.id, session); - log(`agent session for ${user}`); + log(`agent session ${session.id.slice(0, 8)} for ${user}`); tellAgents(key); return session; } @@ -297,6 +298,8 @@ export async function startRelay(options) { * see gets told so with the answer: not every client follows a subscription. */ async function callTool(session, name, args) { + session.calledAt = Date.now(); + tellAgents(session.key); const tab = tabs.get(session.key); if (!tab) return { content: [{ type: 'text', text: NOT_CONNECTED }], isError: true }; const seen = session.lastSummary; @@ -317,8 +320,12 @@ export async function startRelay(options) { if (!String(req.headers.accept ?? '').includes('text/event-stream')) { return reply(res, 405, { error: 'Ask for text/event-stream' }, { Allow: 'GET, POST, DELETE' }); } + // A client may drop the stream while its token is being checked: its close is over by the time + // the stream would be kept, and a stream kept then would count an agent that is gone, for good. + let dropped = false; + res.once('close', () => (dropped = true)); const user = await agent(req, res); - if (user === null) return; + if (user === null || dropped) return; const session = sessionOf(req, res, user); if (!session) return; res.writeHead(200, { @@ -329,9 +336,11 @@ export async function startRelay(options) { }); res.write(': open\n\n'); session.streams.add(res); + log(`stream open on ${session.id.slice(0, 8)}`); tellAgents(session.key); - req.on('close', () => { + res.on('close', () => { session.streams.delete(res); + log(`stream closed on ${session.id.slice(0, 8)}`); tellAgents(session.key); }); } @@ -341,6 +350,7 @@ export async function startRelay(options) { if (user === null) return; const session = sessionOf(req, res, user); if (!session) return; + log(`agent session ${session.id.slice(0, 8)} ended by the agent`); endSession(session); reply(res, 204); } @@ -360,18 +370,24 @@ export async function startRelay(options) { /** * How many agents an account has on, for its tab to show: a session counts while its event stream - * is open or it asked something lately. A client that went away without a DELETE leaves its - * session behind, and so does a revoked token until it is used again. + * is open or it called a tool lately. A client that went away without a DELETE leaves its + * session behind, and so does a revoked token until it is used again. Claude Code opens two + * spare sessions next to the one it keeps whenever it connects: they never call a tool. */ function tellAgents(key) { const tab = tabs.get(key); if (!tab) return; const now = Date.now(); const count = sessionsOf(key).filter( - (session) => session.streams.size > 0 || now - session.seenAt < agentRecentMs, + (session) => session.streams.size > 0 || now - session.calledAt < agentRecentMs, ).length; if (count === tab.agents) return; tab.agents = count; + const seen = sessionsOf(key).map( + (session) => + `${session.id.slice(0, 8)} streams ${session.streams.size} idle ${Math.round((now - session.seenAt) / 1000)}s`, + ); + log(`agents for ${tab.user}: ${count} [${seen.join('; ')}]`); tab.conn.send(JSON.stringify({ type: 'agents', count })); } @@ -617,5 +633,6 @@ function passOn(req, res, backend) { /** * @typedef {{status: number, user?: string, role?: string, error?: string, retryAfter?: string | null}} ApiAnswer * @typedef {{id: string, user: string, key: string, protocolVersion: string, subscriptions: Set, - * streams: Set, seenAt: number, lastSummary: string | null}} Session + * streams: Set, seenAt: number, calledAt: number, + * lastSummary: string | null}} Session */