From 1d2bc7d62685ec6d0e6dbdce673aa95798b5fddf Mon Sep 17 00:00:00 2001 From: Mitch Shao Date: Mon, 5 Oct 2026 14:12:21 -0700 Subject: [PATCH 1/2] feat: stream abort handling, connection pooling, and assistant-prefill continuation - Abort upstream Copilot requests when the client disconnects during streaming - Configure undici connection pool (16 connections per origin) - Skip debug JSON serialization unless debug logging is enabled - Append a 'Continue.' user message when the conversation ends with an assistant turn Co-Authored-By: Claude --- src/lib/proxy.ts | 20 ++- src/routes/chat-completions/handler.ts | 54 ++++++-- src/routes/messages/handler.ts | 116 +++++++++++++----- src/routes/messages/non-stream-translation.ts | 8 ++ .../copilot/create-chat-completions.ts | 14 ++- src/start.ts | 4 +- tests/anthropic-request.test.ts | 26 ++++ tests/create-chat-completions.test.ts | 13 ++ 8 files changed, 208 insertions(+), 47 deletions(-) diff --git a/src/lib/proxy.ts b/src/lib/proxy.ts index 22beb5819c..d8f4647ddf 100644 --- a/src/lib/proxy.ts +++ b/src/lib/proxy.ts @@ -2,11 +2,29 @@ import consola from "consola" import { getProxyForUrl } from "proxy-from-env" import { Agent, ProxyAgent, setGlobalDispatcher, type Dispatcher } from "undici" +const MAX_CONNECTIONS_PER_ORIGIN = 16 + +function createDirectAgent(): Agent { + return new Agent({ + connections: MAX_CONNECTIONS_PER_ORIGIN, + pipelining: 1, + }) +} + +export function initDirectConnectionPool(): void { + if (typeof Bun !== "undefined") return + + setGlobalDispatcher(createDirectAgent()) + consola.debug( + `HTTP connection pool configured with ${MAX_CONNECTIONS_PER_ORIGIN} connections per origin`, + ) +} + export function initProxyFromEnv(): void { if (typeof Bun !== "undefined") return try { - const direct = new Agent() + const direct = createDirectAgent() const proxies = new Map() // We only need a minimal dispatcher that implements `dispatch` at runtime. diff --git a/src/routes/chat-completions/handler.ts b/src/routes/chat-completions/handler.ts index 04a5ae9edd..ba5ad37018 100644 --- a/src/routes/chat-completions/handler.ts +++ b/src/routes/chat-completions/handler.ts @@ -18,7 +18,9 @@ export async function handleCompletion(c: Context) { await checkRateLimit(state) let payload = await c.req.json() - consola.debug("Request payload:", JSON.stringify(payload).slice(-400)) + if (consola.level >= 4) { + consola.debug("Request payload:", JSON.stringify(payload).slice(-400)) + } // Find the selected model const selectedModel = state.models?.data.find( @@ -47,20 +49,50 @@ export async function handleCompletion(c: Context) { consola.debug("Set max_tokens to:", JSON.stringify(payload.max_tokens)) } + if (payload.stream) { + consola.debug("Streaming response") + return streamSSE( + c, + async (stream) => { + const controller = new AbortController() + stream.onAbort(() => { + controller.abort() + }) + + const response = await createChatCompletions(payload, controller.signal) + if (isNonStreaming(response)) { + throw new Error("Expected a streaming response from Copilot") + } + + try { + for await (const chunk of response) { + if (consola.level >= 4) { + consola.debug("Streaming chunk:", JSON.stringify(chunk)) + } + await stream.writeSSE(chunk as SSEMessage) + } + } finally { + controller.abort() + } + }, + (error) => { + if (error.name !== "AbortError") { + consola.error("Streaming request failed:", error) + } + return Promise.resolve() + }, + ) + } + const response = await createChatCompletions(payload) + if (!isNonStreaming(response)) { + throw new Error("Expected a non-streaming response from Copilot") + } - if (isNonStreaming(response)) { + if (consola.level >= 4) { consola.debug("Non-streaming response:", JSON.stringify(response)) - return c.json(response) } - - consola.debug("Streaming response") - return streamSSE(c, async (stream) => { - for await (const chunk of response) { - consola.debug("Streaming chunk:", JSON.stringify(chunk)) - await stream.writeSSE(chunk as SSEMessage) - } - }) + return c.json(response) } const isNonStreaming = ( diff --git a/src/routes/messages/handler.ts b/src/routes/messages/handler.ts index 85dbf62437..860f6a49be 100644 --- a/src/routes/messages/handler.ts +++ b/src/routes/messages/handler.ts @@ -26,64 +26,118 @@ export async function handleCompletion(c: Context) { await checkRateLimit(state) const anthropicPayload = await c.req.json() - consola.debug("Anthropic request payload:", JSON.stringify(anthropicPayload)) + if (consola.level >= 4) { + consola.debug( + "Anthropic request payload:", + JSON.stringify(anthropicPayload), + ) + } const openAIPayload = translateToOpenAI(anthropicPayload) - consola.debug( - "Translated OpenAI request payload:", - JSON.stringify(openAIPayload), - ) + if (consola.level >= 4) { + consola.debug( + "Translated OpenAI request payload:", + JSON.stringify(openAIPayload), + ) + } if (state.manualApprove) { await awaitApproval() } + if (openAIPayload.stream) { + consola.debug("Streaming response from Copilot") + return streamSSE( + c, + async (stream) => { + const controller = new AbortController() + stream.onAbort(() => { + controller.abort() + }) + + const response = await createChatCompletions( + openAIPayload, + controller.signal, + ) + if (isNonStreaming(response)) { + throw new Error("Expected a streaming response from Copilot") + } + + try { + await writeStreamingResponse(stream, response) + } finally { + controller.abort() + } + }, + (error) => { + if (error.name !== "AbortError") { + consola.error("Streaming request failed:", error) + } + return Promise.resolve() + }, + ) + } + const response = await createChatCompletions(openAIPayload) + if (!isNonStreaming(response)) { + throw new Error("Expected a non-streaming response from Copilot") + } - if (isNonStreaming(response)) { + if (consola.level >= 4) { consola.debug( "Non-streaming response from Copilot:", JSON.stringify(response).slice(-400), ) - const anthropicResponse = translateToAnthropic(response) + } + const anthropicResponse = translateToAnthropic(response) + if (consola.level >= 4) { consola.debug( "Translated Anthropic response:", JSON.stringify(anthropicResponse), ) - return c.json(anthropicResponse) } + return c.json(anthropicResponse) +} - consola.debug("Streaming response from Copilot") - return streamSSE(c, async (stream) => { - const streamState: AnthropicStreamState = { - messageStartSent: false, - contentBlockIndex: 0, - contentBlockOpen: false, - toolCalls: {}, - } +async function writeStreamingResponse( + stream: Parameters[1]>[0], + response: Exclude< + Awaited>, + ChatCompletionResponse + >, +): Promise { + const streamState: AnthropicStreamState = { + messageStartSent: false, + contentBlockIndex: 0, + contentBlockOpen: false, + toolCalls: {}, + } - for await (const rawEvent of response) { + for await (const rawEvent of response) { + if (consola.level >= 4) { consola.debug("Copilot raw stream event:", JSON.stringify(rawEvent)) - if (rawEvent.data === "[DONE]") { - break - } + } + if (rawEvent.data === "[DONE]") { + break + } - if (!rawEvent.data) { - continue - } + if (!rawEvent.data) { + continue + } - const chunk = JSON.parse(rawEvent.data) as ChatCompletionChunk - const events = translateChunkToAnthropicEvents(chunk, streamState) + const chunk = JSON.parse(rawEvent.data) as ChatCompletionChunk + const events = translateChunkToAnthropicEvents(chunk, streamState) - for (const event of events) { + for (const event of events) { + if (consola.level >= 4) { consola.debug("Translated Anthropic event:", JSON.stringify(event)) - await stream.writeSSE({ - event: event.type, - data: JSON.stringify(event), - }) } + await stream.writeSSE({ + event: event.type, + data: JSON.stringify(event), + }) } - }) + } } const isNonStreaming = ( diff --git a/src/routes/messages/non-stream-translation.ts b/src/routes/messages/non-stream-translation.ts index dc41e63825..afc6404090 100644 --- a/src/routes/messages/non-stream-translation.ts +++ b/src/routes/messages/non-stream-translation.ts @@ -68,6 +68,14 @@ function translateAnthropicMessagesToOpenAI( : handleAssistantMessage(message), ) + const lastMessage = otherMessages.at(-1) + if ( + lastMessage?.role === "assistant" + && (!lastMessage.tool_calls || lastMessage.tool_calls.length === 0) + ) { + otherMessages.push({ role: "user", content: "Continue." }) + } + return [...systemMessages, ...otherMessages] } diff --git a/src/services/copilot/create-chat-completions.ts b/src/services/copilot/create-chat-completions.ts index 8534151daa..e08c4aa0b1 100644 --- a/src/services/copilot/create-chat-completions.ts +++ b/src/services/copilot/create-chat-completions.ts @@ -1,5 +1,6 @@ import consola from "consola" import { events } from "fetch-event-stream" +import { getGlobalDispatcher, type Dispatcher } from "undici" import { copilotHeaders, copilotBaseUrl } from "~/lib/api-config" import { HTTPError } from "~/lib/error" @@ -7,6 +8,7 @@ import { state } from "~/lib/state" export const createChatCompletions = async ( payload: ChatCompletionsPayload, + signal?: AbortSignal, ) => { if (!state.copilotToken) throw new Error("Copilot token not found") @@ -28,11 +30,17 @@ export const createChatCompletions = async ( "X-Initiator": isAgentCall ? "agent" : "user", } - const response = await fetch(`${copilotBaseUrl(state)}/chat/completions`, { + const requestInit: RequestInit & { dispatcher: Dispatcher } = { method: "POST", headers, body: JSON.stringify(payload), - }) + signal, + dispatcher: getGlobalDispatcher(), + } + const response = await fetch( + `${copilotBaseUrl(state)}/chat/completions`, + requestInit, + ) if (!response.ok) { consola.error("Failed to create chat completions", response) @@ -40,7 +48,7 @@ export const createChatCompletions = async ( } if (payload.stream) { - return events(response) + return events(response, signal) } return (await response.json()) as ChatCompletionResponse diff --git a/src/start.ts b/src/start.ts index 14abbbdffb..e9c24c5fa8 100644 --- a/src/start.ts +++ b/src/start.ts @@ -7,7 +7,7 @@ import { serve, type ServerHandler } from "srvx" import invariant from "tiny-invariant" import { ensurePaths } from "./lib/paths" -import { initProxyFromEnv } from "./lib/proxy" +import { initDirectConnectionPool, initProxyFromEnv } from "./lib/proxy" import { generateEnvScript } from "./lib/shell" import { state } from "./lib/state" import { setupCopilotToken, setupGitHubToken } from "./lib/token" @@ -30,6 +30,8 @@ interface RunServerOptions { export async function runServer(options: RunServerOptions): Promise { if (options.proxyEnv) { initProxyFromEnv() + } else { + initDirectConnectionPool() } if (options.verbose) { diff --git a/tests/anthropic-request.test.ts b/tests/anthropic-request.test.ts index 06c6637788..b67b47d238 100644 --- a/tests/anthropic-request.test.ts +++ b/tests/anthropic-request.test.ts @@ -154,6 +154,32 @@ describe("Anthropic to OpenAI translation logic", () => { "Let me think about this simple math problem...", ) expect(assistantMessage?.content).toContain("2+2 equals 4.") + expect(openAIPayload.messages.at(-1)).toEqual({ + role: "user", + content: "Continue.", + }) + }) + + test("should convert assistant prefill into a user-ended conversation", () => { + const anthropicPayload: AnthropicMessagesPayload = { + model: "claude-opus-5.5", + messages: [ + { role: "user", content: "Complete this response." }, + { role: "assistant", content: "The answer is" }, + ], + max_tokens: 100, + } + + const openAIPayload = translateToOpenAI(anthropicPayload) + + expect(openAIPayload.messages.at(-2)).toEqual({ + role: "assistant", + content: "The answer is", + }) + expect(openAIPayload.messages.at(-1)).toEqual({ + role: "user", + content: "Continue.", + }) }) test("should handle thinking blocks with tool calls", () => { diff --git a/tests/create-chat-completions.test.ts b/tests/create-chat-completions.test.ts index d18e741aa6..5b825072f0 100644 --- a/tests/create-chat-completions.test.ts +++ b/tests/create-chat-completions.test.ts @@ -54,3 +54,16 @@ test("sets X-Initiator to user if only user present", async () => { ).headers expect(headers["X-Initiator"]).toBe("user") }) + +test("passes the request abort signal to the upstream fetch", async () => { + const payload: ChatCompletionsPayload = { + messages: [{ role: "user", content: "hi" }], + model: "gpt-test", + } + const controller = new AbortController() + + await createChatCompletions(payload, controller.signal) + + const options = fetchMock.mock.calls.at(-1)?.[1] as RequestInit + expect(options.signal).toBe(controller.signal) +}) From a04477c81cdfc0f88e6c2b816e9cdd1457b0efc4 Mon Sep 17 00:00:00 2001 From: Mitch Shao Date: Sat, 17 Oct 2026 10:17:05 -0700 Subject: [PATCH 2/2] fix: refresh Copilot token and retry once on 401/403 The Copilot token can be revoked server-side before its scheduled refresh, leaving the proxy returning 403 forbidden until restart. On 401/403 from chat/completions, fetch a fresh Copilot token (deduplicated across concurrent requests) and retry once. Also stop the periodic refresh from rethrowing inside setInterval. Co-Authored-By: Claude --- src/lib/token.ts | 25 ++++++++-- .../copilot/create-chat-completions.ts | 46 +++++++++++++------ 2 files changed, 53 insertions(+), 18 deletions(-) diff --git a/src/lib/token.ts b/src/lib/token.ts index fc8d2785f5..a6ead351f6 100644 --- a/src/lib/token.ts +++ b/src/lib/token.ts @@ -28,6 +28,23 @@ export const setupCopilotToken = async () => { const refreshInterval = (refresh_in - 60) * 1000 setInterval(async () => { consola.debug("Refreshing Copilot token") + try { + await refreshCopilotToken() + } catch (error) { + consola.error("Failed to refresh Copilot token:", error) + } + }, refreshInterval) +} + +let inflightRefresh: Promise | undefined + +/** + * Fetch a new Copilot token and store it in state. Concurrent callers share + * one in-flight request, so a burst of 401/403 responses triggers a single + * refresh. + */ +export const refreshCopilotToken = (): Promise => { + inflightRefresh ??= (async () => { try { const { token } = await getCopilotToken() state.copilotToken = token @@ -35,11 +52,11 @@ export const setupCopilotToken = async () => { if (state.showToken) { consola.info("Refreshed Copilot token:", token) } - } catch (error) { - consola.error("Failed to refresh Copilot token:", error) - throw error + } finally { + inflightRefresh = undefined } - }, refreshInterval) + })() + return inflightRefresh } interface SetupGitHubTokenOptions { diff --git a/src/services/copilot/create-chat-completions.ts b/src/services/copilot/create-chat-completions.ts index e08c4aa0b1..8b8b30e3e2 100644 --- a/src/services/copilot/create-chat-completions.ts +++ b/src/services/copilot/create-chat-completions.ts @@ -5,6 +5,7 @@ import { getGlobalDispatcher, type Dispatcher } from "undici" import { copilotHeaders, copilotBaseUrl } from "~/lib/api-config" import { HTTPError } from "~/lib/error" import { state } from "~/lib/state" +import { refreshCopilotToken } from "~/lib/token" export const createChatCompletions = async ( payload: ChatCompletionsPayload, @@ -24,23 +25,40 @@ export const createChatCompletions = async ( ["assistant", "tool"].includes(msg.role), ) - // Build headers and add X-Initiator - const headers: Record = { - ...copilotHeaders(state, enableVision), - "X-Initiator": isAgentCall ? "agent" : "user", + const body = JSON.stringify(payload) + const send = () => { + // Build headers (with the current Copilot token) and add X-Initiator + const headers: Record = { + ...copilotHeaders(state, enableVision), + "X-Initiator": isAgentCall ? "agent" : "user", + } + + const requestInit: RequestInit & { dispatcher: Dispatcher } = { + method: "POST", + headers, + body, + signal, + dispatcher: getGlobalDispatcher(), + } + return fetch(`${copilotBaseUrl(state)}/chat/completions`, requestInit) } - const requestInit: RequestInit & { dispatcher: Dispatcher } = { - method: "POST", - headers, - body: JSON.stringify(payload), - signal, - dispatcher: getGlobalDispatcher(), + let response = await send() + + // The Copilot token can be revoked server-side before its scheduled + // refresh. On 401/403, fetch a fresh token and retry once. + if (response.status === 401 || response.status === 403) { + consola.warn( + `Copilot returned ${response.status}, refreshing Copilot token and retrying`, + ) + try { + await refreshCopilotToken() + await response.body?.cancel() + response = await send() + } catch (error) { + consola.error("Failed to refresh Copilot token:", error) + } } - const response = await fetch( - `${copilotBaseUrl(state)}/chat/completions`, - requestInit, - ) if (!response.ok) { consola.error("Failed to create chat completions", response)