feat: implement Copilot context cache control and prepare payload for chat completions

Author: caozhiyuan

src/routes/messages/api-flows.ts

@@ -36,6 +36,8 @@ import {
   createChatCompletions,
   type ChatCompletionChunk,
   type ChatCompletionResponse,
+  type ChatCompletionsPayload,
+  type Message,
 } from "~/services/copilot/create-chat-completions"
 import { createMessages } from "~/services/copilot/create-messages"
 import {
@@ -59,6 +61,12 @@ import {
   translateChunkToAnthropicEvents,
 } from "./stream-translation"
 
+const COPILOT_CONTEXT_CACHE_SYSTEM_MARKER_LIMIT = 2
+const COPILOT_CONTEXT_CACHE_NON_SYSTEM_MARKER_LIMIT = 2
+const COPILOT_CONTEXT_CACHE_CONTROL = {
+  type: "ephemeral",
+} as const
+
 export interface FlowBaseOptions {
   logger: ConsolaInstance
   subagentMarker?: SubagentMarker | null
@@ -83,6 +91,7 @@ export const handleWithChatCompletions = async (
 ) => {
   const { logger, subagentMarker, requestId, sessionId, compactType } = options
   const openAIPayload = translateToOpenAI(anthropicPayload)
+  prepareCopilotChatCompletionsPayload(openAIPayload)
   const recordUsage = createCopilotUsageRecorder({
     endpoint: "chat_completions",
     fallbackSessionId: sessionId,
@@ -339,6 +348,58 @@ export const handleWithMessagesApi = async (
   return c.json(response)
 }
 
+export const prepareCopilotChatCompletionsPayload = (
+  payload: ChatCompletionsPayload,
+): void => {
+  applyCopilotContextCache(payload)
+}
+
+const applyCopilotContextCache = (payload: ChatCompletionsPayload): void => {
+  const messageIndexes = selectCopilotContextCacheMessageIndexes(
+    payload.messages,
+  )
+  for (const messageIndex of messageIndexes) {
+    const message = payload.messages[messageIndex]
+    message.copilot_cache_control = { ...COPILOT_CONTEXT_CACHE_CONTROL }
+  }
+}
+
+const selectCopilotContextCacheMessageIndexes = (
+  messages: Array<Message>,
+): Array<number> => {
+  const systemIndexes = messages
+    .flatMap((message, index) =>
+      message.role === "system" && isCopilotContextCacheEligible(message) ?
+        [index]
+      : [],
+    )
+    .slice(0, COPILOT_CONTEXT_CACHE_SYSTEM_MARKER_LIMIT)
+  const reverseNonSystemIndexes = messages
+    .flatMap((message, index) =>
+      message.role !== "system" && isCopilotContextCacheEligible(message) ?
+        [index]
+      : [],
+    )
+    .reverse()
+    .slice(0, COPILOT_CONTEXT_CACHE_NON_SYSTEM_MARKER_LIMIT)
+
+  return uniqueIndexes([...systemIndexes, ...reverseNonSystemIndexes]).sort(
+    (a, b) => a - b,
+  )
+}
+
+const isCopilotContextCacheEligible = (message: Message): boolean => {
+  if (typeof message.content === "string") {
+    return message.content.length > 0
+  }
+
+  return Array.isArray(message.content) && message.content.length > 0
+}
+
+const uniqueIndexes = (indexes: Array<number>): Array<number> => [
+  ...new Set(indexes),
+]
+
 const isNonStreaming = (
   response: Awaited<ReturnType<typeof createChatCompletions>>,
 ): response is ChatCompletionResponse => Object.hasOwn(response, "choices")

src/routes/messages/non-stream-translation.ts

@@ -427,7 +427,7 @@ function mapContent(
     }
   }
   if (contentParts.length === 0) {
-    return null
+    return ""
   }
   return contentParts
 }

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

@@ -217,6 +217,7 @@ export interface Message {
   reasoning_content?: string | null
   reasoning_text?: string | null
   reasoning_opaque?: string | null
+  copilot_cache_control?: CopilotCacheControl
 }
 
 export interface ToolCall {
@@ -234,6 +235,10 @@ export interface CacheControl {
   type: "ephemeral"
 }
 
+export interface CopilotCacheControl {
+  type: "ephemeral"
+}
+
 export interface TextPart {
   type: "text"
   text: string

tests/messages-api-flows.test.ts

@@ -0,0 +1,196 @@
+import { beforeEach, expect, mock, test } from "bun:test"
+
+import type { AnthropicMessagesPayload } from "../src/routes/messages/anthropic-types"
+import type {
+  ChatCompletionResponse,
+  ChatCompletionsPayload,
+} from "../src/services/copilot/create-chat-completions"
+
+const actualChatCompletionsModule = await import(
+  "../src/services/copilot/create-chat-completions"
+)
+
+let capturedPayload: ChatCompletionsPayload | null = null
+
+const createChatCompletions = mock(
+  (payload: ChatCompletionsPayload): Promise<ChatCompletionResponse> => {
+    capturedPayload = payload
+    return Promise.resolve({
+      id: "chatcmpl-test",
+      object: "chat.completion",
+      created: 0,
+      model: payload.model,
+      choices: [
+        {
+          index: 0,
+          message: {
+            role: "assistant",
+            content: "ok",
+          },
+          logprobs: null,
+          finish_reason: "stop",
+        },
+      ],
+    })
+  },
+)
+
+await mock.module("~/services/copilot/create-chat-completions", () => ({
+  ...actualChatCompletionsModule,
+  createChatCompletions,
+}))
+
+const { handleWithChatCompletions, prepareCopilotChatCompletionsPayload } =
+  await import("../src/routes/messages/api-flows")
+
+const logger = {
+  debug: () => {},
+  warn: () => {},
+  error: () => {},
+} as unknown as Parameters<typeof handleWithChatCompletions>[2]["logger"]
+
+const createContext = () =>
+  ({
+    json: (body: unknown) => Response.json(body),
+  }) as Parameters<typeof handleWithChatCompletions>[0]
+
+beforeEach(() => {
+  capturedPayload = null
+  createChatCompletions.mockClear()
+})
+
+test("messages Chat Completions flow adds Copilot cache control to system and latest two non-system messages", async () => {
+  const payload: AnthropicMessagesPayload = {
+    model: "gpt-test",
+    max_tokens: 128,
+    system: [
+      {
+        type: "text",
+        text: "system prompt",
+      },
+    ],
+    messages: [
+      { role: "user", content: "first user" },
+      {
+        role: "user",
+        content: [
+          {
+            type: "text",
+            text: "second user",
+          },
+        ],
+      },
+      { role: "assistant", content: "older answer" },
+      { role: "user", content: "latest user" },
+      { role: "assistant", content: "latest answer" },
+    ],
+  }
+
+  const response = await handleWithChatCompletions(createContext(), payload, {
+    logger,
+    requestId: "request-1",
+  })
+
+  expect(response.status).toBe(200)
+  expect(createChatCompletions).toHaveBeenCalledTimes(1)
+  expect(capturedPayload?.messages).toEqual([
+    {
+      role: "system",
+      content: "system prompt",
+      copilot_cache_control: {
+        type: "ephemeral",
+      },
+    },
+    {
+      role: "user",
+      content: "first user",
+    },
+    {
+      role: "user",
+      content: [
+        {
+          type: "text",
+          text: "second user",
+        },
+      ],
+    },
+    {
+      role: "assistant",
+      content: "older answer",
+    },
+    {
+      role: "user",
+      content: "latest user",
+      copilot_cache_control: {
+        type: "ephemeral",
+      },
+    },
+    {
+      role: "assistant",
+      content: "latest answer",
+      copilot_cache_control: {
+        type: "ephemeral",
+      },
+    },
+  ])
+})
+
+test("Copilot Chat Completions payload preparation marks two system and latest two non-system messages", () => {
+  const payload: ChatCompletionsPayload = {
+    model: "gpt-test",
+    messages: [
+      { role: "system", content: "system one" },
+      { role: "system", content: "system two" },
+      { role: "system", content: "system three" },
+      { role: "user", content: "older user" },
+      { role: "assistant", content: "older assistant" },
+      { role: "user", content: "latest user" },
+      { role: "assistant", content: "latest assistant" },
+    ],
+  }
+
+  prepareCopilotChatCompletionsPayload(payload)
+
+  expect(payload.messages).toEqual([
+    {
+      role: "system",
+      content: "system one",
+      copilot_cache_control: {
+        type: "ephemeral",
+      },
+    },
+    {
+      role: "system",
+      content: "system two",
+      copilot_cache_control: {
+        type: "ephemeral",
+      },
+    },
+    {
+      role: "system",
+      content: "system three",
+    },
+    {
+      role: "user",
+      content: "older user",
+    },
+    {
+      role: "assistant",
+      content: "older assistant",
+    },
+    {
+      role: "user",
+      content: "latest user",
+      copilot_cache_control: {
+        type: "ephemeral",
+      },
+    },
+    {
+      role: "assistant",
+      content: "latest assistant",
+      copilot_cache_control: {
+        type: "ephemeral",
+      },
+    },
+  ])
+})