From 0dd671ea260272755ff3dd5dad4eb91507a022d3 Mon Sep 17 00:00:00 2001 From: MelekTaus Date: Sun, 18 Oct 2026 12:10:02 +0200 Subject: [PATCH] Count each connected agent once in the editor badge (#429) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The Agent badge showed 2 or 3 while one Claude Code session was connected. The relay log against the open TCP connections in the container showed two causes. - A client may drop its event stream while the relay still asks the API about its token. The close was over before the relay listened for it, so the stream stayed in the list and counted for good. The relay now listens before the token check and keeps no stream that is already gone. It listens on the response, not the request. - Claude Code opens three sessions whenever it connects and keeps a stream on one of them. It never ends the others with a DELETE. Any session counted for two minutes after its last request, so the badge stood at 3 after each start. Now a session counts while its stream is open or it called a tool in the last two minutes. - The relay logs session start, stream open and close and each change of the count, with session ids cut to 8 characters. One case is left: after a relay restart Claude Code sometimes repeats its first call on a spare session. The badge then shows 2 for two minutes and drops to 1 on its own. Tests: - relay.test.ts: 27/27, three new. A dropped stream counts out, a stream dropped during the token check never counts, and spare sessions without a tool call do not count. - Counter-checks: without the early close listener the drop test turns red, with the old count two tests turn red. - Live: the badge showed 2 with one open SSE connection before. After the fix it shows "Agent ยท 1" and stays there after a normal call. Co-Authored-By: Claude Opus 5.5 --- tests/mcp/relay.test.ts | 59 +++++++++++++++++++++++++++++++++++++---- tools/mcp/relay.mjs | 31 +++++++++++++++++----- 2 files changed, 78 insertions(+), 12 deletions(-) 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 */