feat: refactor message handling to use centralized flow handlers

Author: caozhiyuan

src/routes/messages/handler.ts

@@ -27,6 +27,12 @@ import { parseSubagentMarkerFromFirstUser } from "./subagent-marker"
 
 const logger = createHandlerLogger("messages-handler")
 
+export const messagesFlowHandlers = {
+  handleWithChatCompletions,
+  handleWithMessagesApi,
+  handleWithResponsesApi,
+}
+
 export async function handleCompletion(c: Context) {
   await checkRateLimit(state)
 
@@ -81,35 +87,47 @@ export async function handleCompletion(c: Context) {
   anthropicPayload.model = selectedModel?.id ?? anthropicPayload.model
 
   if (shouldUseMessagesApi(selectedModel)) {
-    return await handleWithMessagesApi(c, anthropicPayload, {
-      anthropicBetaHeader: anthropicBeta,
-      subagentMarker,
-      selectedModel,
-      requestId,
-      sessionId,
-      compactType,
-      logger,
-    })
+    return await messagesFlowHandlers.handleWithMessagesApi(
+      c,
+      anthropicPayload,
+      {
+        anthropicBetaHeader: anthropicBeta,
+        subagentMarker,
+        selectedModel,
+        requestId,
+        sessionId,
+        compactType,
+        logger,
+      },
+    )
   }
 
   if (shouldUseResponsesApi(selectedModel)) {
-    return await handleWithResponsesApi(c, anthropicPayload, {
+    return await messagesFlowHandlers.handleWithResponsesApi(
+      c,
+      anthropicPayload,
+      {
+        subagentMarker,
+        selectedModel,
+        requestId,
+        sessionId,
+        compactType,
+        logger,
+      },
+    )
+  }
+
+  return await messagesFlowHandlers.handleWithChatCompletions(
+    c,
+    anthropicPayload,
+    {
       subagentMarker,
-      selectedModel,
       requestId,
       sessionId,
       compactType,
       logger,
-    })
-  }
-
-  return await handleWithChatCompletions(c, anthropicPayload, {
-    subagentMarker,
-    requestId,
-    sessionId,
-    compactType,
-    logger,
-  })
+    },
+  )
 }
 
 const RESPONSES_ENDPOINT = "/responses"

tests/messages-handler.test.ts

@@ -1,4 +1,4 @@
-import { beforeEach, describe, expect, mock, test } from "bun:test"
+import { afterEach, beforeEach, describe, expect, mock, test } from "bun:test"
 import { Hono } from "hono"
 
 import type { AnthropicMessagesPayload } from "../src/routes/messages/anthropic-types"
@@ -36,21 +36,21 @@ const handleWithMessagesApi = mock(
     _c: unknown,
     _payload: AnthropicMessagesPayload,
     _options: FlowCallOptions,
-  ) => new Response("messages"),
+  ) => Promise.resolve(new Response("messages")),
 )
 const handleWithResponsesApi = mock(
   (
     _c: unknown,
     _payload: AnthropicMessagesPayload,
     _options: FlowCallOptions,
-  ) => new Response("responses"),
+  ) => Promise.resolve(new Response("responses")),
 )
 const handleWithChatCompletions = mock(
   (
     _c: unknown,
     _payload: AnthropicMessagesPayload,
     _options: FlowCallOptions,
-  ) => new Response("chat"),
+  ) => Promise.resolve(new Response("chat")),
 )
 
 await mock.module("~/lib/state", () => ({
@@ -73,13 +73,11 @@ await mock.module("~/lib/models", () => ({
 await mock.module("~/lib/utils", () => ({
   ...actualUtilsModule,
 }))
-await mock.module("~/routes/messages/api-flows", () => ({
-  handleWithMessagesApi,
-  handleWithResponsesApi,
-  handleWithChatCompletions,
-}))
+const { handleCompletion, messagesFlowHandlers } = await import(
+  "../src/routes/messages/handler"
+)
 
-const { handleCompletion } = await import("../src/routes/messages/handler")
+const defaultMessagesFlowHandlers = { ...messagesFlowHandlers }
 
 const createApp = () => {
   const app = new Hono()
@@ -102,12 +100,25 @@ beforeEach(() => {
   messagesApiEnabled = true
   selectedModel = undefined
 
+  messagesFlowHandlers.handleWithMessagesApi = handleWithMessagesApi
+  messagesFlowHandlers.handleWithResponsesApi = handleWithResponsesApi
+  messagesFlowHandlers.handleWithChatCompletions = handleWithChatCompletions
+
   findEndpointModel.mockClear()
   handleWithMessagesApi.mockClear()
   handleWithResponsesApi.mockClear()
   handleWithChatCompletions.mockClear()
 })
 
+afterEach(() => {
+  messagesFlowHandlers.handleWithMessagesApi =
+    defaultMessagesFlowHandlers.handleWithMessagesApi
+  messagesFlowHandlers.handleWithResponsesApi =
+    defaultMessagesFlowHandlers.handleWithResponsesApi
+  messagesFlowHandlers.handleWithChatCompletions =
+    defaultMessagesFlowHandlers.handleWithChatCompletions
+})
+
 describe("messages handler orchestration", () => {
   test("removes executeCode and rewrites getDiagnostics before forwarding tools", async () => {
     selectedModel = {