fix: abort upstream requests on client disconnect, pool connections, handle assistant prefill

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<string, ProxyAgent>()
 
     // We only need a minimal dispatcher that implements `dispatch` at runtime.

src/lib/token.ts

@@ -28,18 +28,35 @@ 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<void> | 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<void> => {
+  inflightRefresh ??= (async () => {
     try {
       const { token } = await getCopilotToken()
       state.copilotToken = token
       consola.debug("Copilot token refreshed")
       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 {

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<ChatCompletionsPayload>()
-  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 = (

src/routes/messages/handler.ts

@@ -26,64 +26,118 @@ export async function handleCompletion(c: Context) {
   await checkRateLimit(state)
 
   const anthropicPayload = await c.req.json<AnthropicMessagesPayload>()
-  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<Parameters<typeof streamSSE>[1]>[0],
+  response: Exclude<
+    Awaited<ReturnType<typeof createChatCompletions>>,
+    ChatCompletionResponse
+  >,
+): Promise<void> {
+  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 = (

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]
 }
 

src/services/copilot/create-chat-completions.ts

@@ -1,12 +1,15 @@
 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"
 import { state } from "~/lib/state"
+import { refreshCopilotToken } from "~/lib/token"
 
 export const createChatCompletions = async (
   payload: ChatCompletionsPayload,
+  signal?: AbortSignal,
 ) => {
   if (!state.copilotToken) throw new Error("Copilot token not found")
 
@@ -22,25 +25,48 @@ export const createChatCompletions = async (
     ["assistant", "tool"].includes(msg.role),
   )
 
-  // Build headers and add X-Initiator
-  const headers: Record<string, string> = {
-    ...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<string, string> = {
+      ...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 response = await fetch(`${copilotBaseUrl(state)}/chat/completions`, {
-    method: "POST",
-    headers,
-    body: JSON.stringify(payload),
-  })
+  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)
+    }
+  }
 
   if (!response.ok) {
     consola.error("Failed to create chat completions", response)
     throw new HTTPError("Failed to create chat completions", response)
   }
 
   if (payload.stream) {
-    return events(response)
+    return events(response, signal)
   }
 
   return (await response.json()) as ChatCompletionResponse