External/open-design
1
0
Fork 0
mirror of https://github.com/nexu-io/open-design.git synced 2026-10-27 09:51:12 +01:00
Code Issues Packages Projects Releases Wiki Activity

fix(chat): abort BYOK proxy upstreams when the client disconnects (#5182)

Browse Source 
Co-authored-by: tomsen-ai <230283659+tomsen-ai@users.noreply.github.com>
This commit is contained in:
ShiZai
2026-07-08 11:35:11 +00:00
committed by GitHub
co-authored by tomsen-ai
parent d513f9fd78
commit 9abba14fc9
3 changed files with 96 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/daemon/src/byok-tools.ts
+4 -3
View File
@@ -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<RequestInit, 'dispatcher'>;
* 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<RequestInit, 'dispatcher' | 'signal'>;
}
export interface ImageToolResult {
apps/daemon/src/routes/chat.ts
+36 -3
View File
@@ -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<string, string>; payload: any; logTag: string },
@@ -661,9 +675,11 @@ export function registerChatRoutes(app: Express, ctx: RegisterChatRoutesDeps) {
let proxyDispatcher: ReturnType<typeof proxyDispatcherRequestInit> | 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<typeof proxyDispatcherRequestInit> | 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<typeof proxyDispatcherRequestInit> | 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<typeof proxyDispatcherRequestInit> | 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<typeof proxyDispatcherRequestInit> | 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<typeof proxyDispatcherRequestInit> | 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<typeof proxyDispatcherRequestInit> | 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
apps/daemon/tests/proxy-routes.test.ts
+56
View File
@@ -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 {