diff --git a/apps/daemon/src/byok-tools.ts b/apps/daemon/src/byok-tools.ts index 8832b69ecc..0ac1d220c9 100644 --- a/apps/daemon/src/byok-tools.ts +++ b/apps/daemon/src/byok-tools.ts @@ -446,9 +446,10 @@ export interface BYOKToolContext { * semantics. */ videoPollIntervalMs?: number; /** Optional per-request init copied from the live chat turn. Used to - * forward the current proxy dispatcher into every upstream/download - * fetch the BYOK tool executor performs. */ - requestInit?: Pick; + * forward the current proxy dispatcher AND the client-cancellation + * signal into every upstream/download fetch the BYOK tool executor + * performs, so a disconnected client stops the tool loop's paid work. */ + requestInit?: Pick; } export interface ImageToolResult { diff --git a/apps/daemon/src/routes/chat.ts b/apps/daemon/src/routes/chat.ts index 8c087100a3..716a907a06 100644 --- a/apps/daemon/src/routes/chat.ts +++ b/apps/daemon/src/routes/chat.ts @@ -653,6 +653,20 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { return payload; }; + // A BYOK proxy stream must unwind when the client disconnects (Stop or a + // closed tab); otherwise the upstream completion — and any tool loop that + // would fire further paid image/video/speech rounds — keeps streaming and + // billing after the user is gone. Every upstream fetch below carries this + // signal. Mirrors the AbortController wiring already used by + // `/api/provider/models` and `/api/test/connection`. + const clientDisconnectSignal = (res: any, req?: any): AbortSignal => { + const controller = new AbortController(); + const abort = () => controller.abort(); + res.on('close', abort); + if (req) req.on('close', abort); + return controller.signal; + }; + const runAnthropicChatStream = async ( res: any, opts: { url: string; headers: Record; payload: any; logTag: string }, @@ -661,9 +675,11 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { let proxyDispatcher: ReturnType | null = null; try { proxyDispatcher = proxyDispatcherRequestInit(); + const signal = clientDisconnectSignal(res); sse.send('start', { model: opts.payload?.model }); const response = await fetch(opts.url, { ...proxyDispatcher.requestInit, + signal, method: 'POST', headers: { 'Content-Type': 'application/json', ...opts.headers }, body: JSON.stringify(opts.payload), @@ -749,9 +765,11 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { let proxyDispatcher: ReturnType | null = null; try { proxyDispatcher = proxyDispatcherRequestInit(); + const signal = clientDisconnectSignal(res); sse.send('start', { model: opts.model }); const response = await fetch(opts.url, { ...proxyDispatcher.requestInit, + signal, method: 'POST', headers: { 'Content-Type': 'application/json', ...opts.headers }, body: JSON.stringify(opts.payload), @@ -998,9 +1016,11 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { let proxyDispatcher: ReturnType | null = null; try { proxyDispatcher = proxyDispatcherRequestInit(); + const signal = clientDisconnectSignal(res); sse.send('start', { model }); const response = await fetch(url, { ...proxyDispatcher.requestInit, + signal, method: 'POST', headers: { 'Content-Type': 'application/json', @@ -1143,9 +1163,11 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { let proxyDispatcher: ReturnType | null = null; try { proxyDispatcher = proxyDispatcherRequestInit(); + const signal = clientDisconnectSignal(res); sse.send('start', { model }); const requestInit = { ...proxyDispatcher.requestInit, + signal, method: 'POST', headers: { 'Content-Type': 'application/json', @@ -1318,9 +1340,11 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { let proxyDispatcher: ReturnType | null = null; try { proxyDispatcher = proxyDispatcherRequestInit(); + const signal = clientDisconnectSignal(res); sse.send('start', { model }); const response = await fetch(url, { ...proxyDispatcher.requestInit, + signal, method: 'POST', headers: { 'Content-Type': 'application/json', Authorization: `Bearer ${apiKey}` }, body: JSON.stringify(payload), @@ -1842,10 +1866,12 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { let proxyDispatcher: ReturnType | null = null; try { proxyDispatcher = proxyDispatcherRequestInit(); - toolCtx.requestInit = proxyDispatcher.requestInit; + const signal = clientDisconnectSignal(res); + toolCtx.requestInit = { ...proxyDispatcher.requestInit, signal }; sse.send('start', { model }); const convMessages: any[] = Array.isArray(messages) ? [...messages] : []; for (let loop = 0; loop < MAX_BYOK_TOOL_LOOPS; loop++) { + if (signal.aborted) return sse.end(); const turn = await runAnthropicToolTurn(sse, anthropicUrl, headers, convMessages); if (turn.kind === 'error') return sse.end(); if (turn.kind === 'text_end') { @@ -1857,6 +1883,7 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { convMessages.push({ role: 'assistant', content: turn.assistantBlocks }); const toolResults: any[] = []; for (const call of turn.toolCalls) { + if (signal.aborted) return sse.end(); const result = await executeOneTool(call); const toolName = call?.function?.name ?? 'unknown'; if (result.ok) { @@ -1997,13 +2024,15 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { let proxyDispatcher: ReturnType | null = null; try { proxyDispatcher = proxyDispatcherRequestInit(); - toolCtx.requestInit = proxyDispatcher.requestInit; + const signal = clientDisconnectSignal(res); + toolCtx.requestInit = { ...proxyDispatcher.requestInit, signal }; sse.send('start', { model }); const contents: any[] = (Array.isArray(messages) ? messages : []).map((m: any) => ({ role: m.role === 'assistant' ? 'model' : 'user', parts: [{ text: typeof m.content === 'string' ? m.content : '' }], })); for (let loop = 0; loop < MAX_BYOK_TOOL_LOOPS; loop++) { + if (signal.aborted) return sse.end(); const turn = await runGeminiToolTurn(sse, geminiUrl, headers, contents); if (turn.kind === 'error') return sse.end(); if (turn.kind === 'text_end') { @@ -2015,6 +2044,7 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { contents.push({ role: 'model', parts: turn.functionCallParts }); const responseParts: any[] = []; for (const call of turn.toolCalls) { + if (signal.aborted) return sse.end(); const result = await executeOneTool(call); const toolName = call?.function?.name ?? 'unknown'; if (result.ok) { @@ -2125,9 +2155,11 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { try { proxyDispatcher = proxyDispatcherRequestInit(); - toolCtx.requestInit = proxyDispatcher.requestInit; + const signal = clientDisconnectSignal(res, req); + toolCtx.requestInit = { ...proxyDispatcher.requestInit, signal }; sse.send('start', { model }); for (let loop = 0; loop < MAX_BYOK_TOOL_LOOPS; loop++) { + if (signal.aborted) return sse.end(); const turn = await runTurn(sse, workingMessages); if (turn.kind === 'error') return sse.end(); if (turn.kind === 'text_end') { @@ -2137,6 +2169,7 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) { // turn.kind === 'tool_calls' workingMessages.push(turn.assistantMessage); for (const call of turn.toolCalls) { + if (signal.aborted) return sse.end(); const result = await executeOneTool(call); // The tool result is delivered to the model as a `tool` role // message — a structured payload the model can interpret. We diff --git a/apps/daemon/tests/proxy-routes.test.ts b/apps/daemon/tests/proxy-routes.test.ts index bcea40b55c..c409db9c39 100644 --- a/apps/daemon/tests/proxy-routes.test.ts +++ b/apps/daemon/tests/proxy-routes.test.ts @@ -1922,6 +1922,62 @@ describe('API proxy routes', () => { }); } }); + + // A client that disconnects (Stop / closed tab) must not keep the upstream + // request billing. Every upstream proxy fetch has to carry an AbortSignal + // tied to the client connection so the in-flight completion — and any + // BYOK tool loop that would otherwise fire further paid rounds — unwinds + // when the client goes away. + it('passes a client-cancellation signal to the upstream on a simple proxy stream', async () => { + let upstreamInit: FetchInit | undefined; + const fetchMock = vi.fn((input: FetchInput, init?: FetchInit) => { + const url = String(input); + if (url.startsWith(baseUrl)) return realFetch(input, init); + upstreamInit = init; + return Promise.resolve(sseResponse('data: [DONE]\n\n')); + }); + vi.stubGlobal('fetch', fetchMock); + + const res = await realFetch(`${baseUrl}/api/proxy/openai/stream`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + baseUrl: 'https://api.example.com/v1', + apiKey: 'sk-test', + model: 'gpt-test', + messages: [{ role: 'user', content: 'hello' }], + }), + }); + await res.text(); + + expect(upstreamInit?.signal).toBeInstanceOf(AbortSignal); + }); + + it('passes a client-cancellation signal to the upstream on a BYOK tool-loop stream', async () => { + let upstreamInit: FetchInit | undefined; + const fetchMock = vi.fn((input: FetchInput, init?: FetchInit) => { + const url = String(input); + if (url.startsWith(baseUrl)) return realFetch(input, init); + upstreamInit = init; + return Promise.resolve(sseResponse('data: [DONE]\n\n')); + }); + vi.stubGlobal('fetch', fetchMock); + + const res = await realFetch(`${baseUrl}/api/proxy/senseaudio/stream`, { + method: 'POST', + headers: { 'content-type': 'application/json' }, + body: JSON.stringify({ + baseUrl: 'https://api.senseaudio.cn', + apiKey: 'sa-key', + model: 'senseaudio-s2', + projectId: 'test-project', + messages: [{ role: 'user', content: 'hello' }], + }), + }); + await res.text(); + + expect(upstreamInit?.signal).toBeInstanceOf(AbortSignal); + }); }); function sseResponse(text: string): Response {