From 5722f21cddd0ce77c5bf861c887b9ce981b74d47 Mon Sep 17 00:00:00 2001 From: "CJACK." Date: Fri, 3 Apr 2026 01:49:33 +0800 Subject: [PATCH 1/2] Align docs and adapters with unified OpenAI proxy architecture --- README.MD | 26 +++-- README.en.md | 26 +++-- internal/adapter/claude/handler_messages.go | 88 +++------------ internal/adapter/claude/proxy_vercel_test.go | 42 +++++++ internal/adapter/claude/stream_status_test.go | 50 ++------- internal/adapter/gemini/handler_generate.go | 81 ++++---------- internal/adapter/gemini/handler_test.go | 104 ++++++++++++------ 7 files changed, 185 insertions(+), 232 deletions(-) diff --git a/README.MD b/README.MD index fe97880..7696bfc 100644 --- a/README.MD +++ b/README.MD @@ -31,7 +31,7 @@ flowchart LR Client["🖥️ 客户端 / SDK\n(OpenAI / Claude / Gemini)"] Upstream["☁️ DeepSeek API"] - subgraph DS2API["DS2API 3.x（统一 Go 路由内核）"] + subgraph DS2API["DS2API 3.x（统一 OpenAI 内核）"] Router["chi Router + 中间件\n(RequestID / RealIP / Logger / Recoverer / CORS)"] subgraph Adapters["协议适配层"] @@ -43,13 +43,13 @@ flowchart LR end subgraph Runtime["运行时核心能力"] + Bridge["CLIProxy 转换桥\n(多协议 <-> OpenAI)"] + OAEngine["OpenAI ChatCompletions\n(统一工具调用与流式语义)"] Auth["Auth Resolver\n(API key / bearer / x-goog-api-key)"] Pool["Account Pool + Queue\n(并发槽位 + 等待队列)"] DSClient["DeepSeek Client\n(Session / Auth / HTTP)"] Pow["PoW WASM\n(wazero 预加载)"] - SSE["SSE/Stream Engine\n(统一流式消费)"] Tool["Tool Sieve\n(Go/Node 语义对齐)"] - Render["Formatter\n(OpenAI/Claude/Gemini 输出)"] end end @@ -58,17 +58,18 @@ flowchart LR Router --> Admin Router --> WebUI - OA & CA & GA --> Auth - OA & CA & GA -.账号轮询.-> Pool - OA & CA & GA -.工具调用解析.-> Tool - OA & CA & GA -.流式处理.-> SSE - OA & CA & GA -.PoW 计算.-> Pow - + OA --> OAEngine + CA & GA --> Bridge + Bridge --> OAEngine + OAEngine --> Auth + OAEngine -.账号轮询.-> Pool + OAEngine -.工具调用解析.-> Tool + OAEngine -.PoW 计算.-> Pow Auth --> DSClient DSClient --> Upstream Upstream --> DSClient - DSClient --> Render - Render --> Client + OAEngine --> Bridge + Bridge --> Client ``` - **后端**：Go（`cmd/ds2api/`、`api/`、`internal/`），不依赖 Python 运行时 @@ -78,7 +79,8 @@ flowchart LR ### 3.0 底层架构调整（相较旧版本） - **统一路由内核**：所有协议入口统一汇聚到 `internal/server/router.go`，并在同一路由树中注册 OpenAI / Claude / Gemini / Admin / WebUI 路由，避免多入口行为漂移。 -- **适配器分层更清晰**：`internal/adapter/{openai,claude,gemini}` 只负责协议形态、错误格式和流式事件语义，DeepSeek 侧调用统一收敛到共享调用层。 +- **统一执行链路**：Claude / Gemini 入口先经 `internal/translatorcliproxy` 做协议转换，再进入 `openai.ChatCompletions` 统一处理工具调用与流式语义，最后再转换回原协议响应。 +- **适配器分层更清晰**：`internal/adapter/{claude,gemini}` 负责入口/出口协议封装，`internal/adapter/openai` 负责核心执行，DeepSeek 侧调用只保留在 OpenAI 内核中。 - **Tool Calling 双运行时对齐**：Go 侧（`internal/util`）与 Vercel Node 侧（`internal/js/helpers/stream-tool-sieve`）保持一致的解析/防泄漏语义，覆盖 JSON / XML / invoke / text-kv 多风格输入。 - **配置与运行时设置解耦**：静态配置（`config`）与运行时策略（`settings`）通过 Admin API 分离管理，支持热更新和密码轮换失效旧 JWT。 - **流式能力升级**：`/v1/responses` 与 `/v1/chat/completions` 共享更一致的工具调用增量输出策略，降低不同 SDK 下的行为差异。 diff --git a/README.en.md b/README.en.md index 347c426..65096ef 100644 --- a/README.en.md +++ b/README.en.md @@ -31,7 +31,7 @@ flowchart LR Client["🖥️ Clients / SDKs\n(OpenAI / Claude / Gemini)"] Upstream["☁️ DeepSeek API"] - subgraph DS2API["DS2API 3.x (Unified Go Routing Core)"] + subgraph DS2API["DS2API 3.x (Unified OpenAI Core)"] Router["chi Router + Middleware\n(RequestID / RealIP / Logger / Recoverer / CORS)"] subgraph Adapters["Protocol Adapters"] @@ -43,13 +43,13 @@ flowchart LR end subgraph Runtime["Runtime + Core Capabilities"] + Bridge["CLIProxy Bridge\n(multi-protocol <-> OpenAI)"] + OAEngine["OpenAI ChatCompletions\n(unified tools + stream semantics)"] Auth["Auth Resolver\n(API key / bearer / x-goog-api-key)"] Pool["Account Pool + Queue\n(in-flight slots + wait queue)"] DSClient["DeepSeek Client\n(session / auth / HTTP)"] Pow["PoW WASM\n(wazero preload)"] - SSE["SSE/Stream Engine\n(unified streaming consumption)"] Tool["Tool Sieve\n(Go/Node semantic parity)"] - Render["Formatter\n(OpenAI/Claude/Gemini output)"] end end @@ -58,17 +58,18 @@ flowchart LR Router --> Admin Router --> WebUI - OA & CA & GA --> Auth - OA & CA & GA -.account rotation.-> Pool - OA & CA & GA -.tool-call parsing.-> Tool - OA & CA & GA -.stream handling.-> SSE - OA & CA & GA -.PoW solving.-> Pow - + OA --> OAEngine + CA & GA --> Bridge + Bridge --> OAEngine + OAEngine --> Auth + OAEngine -.account rotation.-> Pool + OAEngine -.tool-call parsing.-> Tool + OAEngine -.PoW solving.-> Pow Auth --> DSClient DSClient --> Upstream Upstream --> DSClient - DSClient --> Render - Render --> Client + OAEngine --> Bridge + Bridge --> Client ``` - **Backend**: Go (`cmd/ds2api/`, `api/`, `internal/`), no Python runtime @@ -78,7 +79,8 @@ flowchart LR ### 3.0 Architecture Changes (vs older releases) - **Unified routing core**: all protocol entries are now centralized through `internal/server/router.go`, with OpenAI / Claude / Gemini / Admin / WebUI routes registered in one tree to avoid multi-entry drift. -- **Cleaner adapter boundaries**: `internal/adapter/{openai,claude,gemini}` focuses on protocol shapes, error contracts, and streaming semantics, while upstream DeepSeek invocation stays shared. +- **Unified execution chain**: Claude/Gemini entries are translated by `internal/translatorcliproxy`, then executed through `openai.ChatCompletions` for shared tool-calling and stream semantics, then translated back to the client protocol. +- **Cleaner adapter boundaries**: `internal/adapter/{claude,gemini}` handles protocol wrappers, while `internal/adapter/openai` remains the execution core; upstream DeepSeek calls are retained only in the OpenAI core. - **Tool-calling parity across runtimes**: Go (`internal/util`) and Vercel Node (`internal/js/helpers/stream-tool-sieve`) follow aligned parsing/anti-leak semantics across JSON / XML / invoke / text-kv inputs. - **Config/runtime separation**: static config (`config`) and runtime policy (`settings`) are managed independently via Admin APIs, enabling hot updates and password rotation with JWT invalidation. - **Streaming behavior upgrade**: `/v1/responses` and `/v1/chat/completions` now share a more consistent incremental tool-call emission strategy across SDK ecosystems. diff --git a/internal/adapter/claude/handler_messages.go b/internal/adapter/claude/handler_messages.go index ced0dc1..c066ac8 100644 --- a/internal/adapter/claude/handler_messages.go +++ b/internal/adapter/claude/handler_messages.go @@ -3,17 +3,12 @@ package claude import ( "bytes" "encoding/json" - "fmt" "io" "net/http" "net/http/httptest" "strings" - "time" - "ds2api/internal/auth" "ds2api/internal/config" - claudefmt "ds2api/internal/format/claude" - "ds2api/internal/sse" streamengine "ds2api/internal/stream" "ds2api/internal/translatorcliproxy" "ds2api/internal/util" @@ -25,80 +20,17 @@ func (h *Handler) Messages(w http.ResponseWriter, r *http.Request) { if strings.TrimSpace(r.Header.Get("anthropic-version")) == "" { r.Header.Set("anthropic-version", "2023-06-01") } - if h.OpenAI != nil { - if h.proxyViaOpenAI(w, r) { - return - } - } - a, err := h.Auth.Determine(r) - if err != nil { - status := http.StatusUnauthorized - detail := err.Error() - if err == auth.ErrNoAccount { - status = http.StatusTooManyRequests - } - writeClaudeError(w, status, detail) + if h.OpenAI == nil { + writeClaudeError(w, http.StatusInternalServerError, "OpenAI proxy backend unavailable.") return } - defer h.Auth.Release(a) - - var req map[string]any - if err := json.NewDecoder(r.Body).Decode(&req); err != nil { - writeClaudeError(w, http.StatusBadRequest, "invalid json") + if h.proxyViaOpenAI(w, r, h.Store) { return } - norm, err := normalizeClaudeRequest(h.Store, req) - if err != nil { - writeClaudeError(w, http.StatusBadRequest, err.Error()) - return - } - stdReq := norm.Standard - - sessionID, err := h.DS.CreateSession(r.Context(), a, 3) - if err != nil { - writeClaudeError(w, http.StatusUnauthorized, "invalid token.") - return - } - pow, err := h.DS.GetPow(r.Context(), a, 3) - if err != nil { - writeClaudeError(w, http.StatusUnauthorized, "Failed to get PoW") - return - } - requestPayload := stdReq.CompletionPayload(sessionID) - resp, err := h.DS.CallCompletion(r.Context(), a, requestPayload, pow, 3) - if err != nil { - writeClaudeError(w, http.StatusInternalServerError, "Failed to get Claude response.") - return - } - if resp.StatusCode != http.StatusOK { - defer resp.Body.Close() - body, _ := io.ReadAll(resp.Body) - writeClaudeError(w, http.StatusInternalServerError, string(body)) - return - } - - if stdReq.Stream { - h.handleClaudeStreamRealtime(w, r, resp, stdReq.ResponseModel, norm.NormalizedMessages, stdReq.Thinking, stdReq.Search, stdReq.ToolNames) - return - } - result := sse.CollectStream(resp, stdReq.Thinking, true) - respBody := claudefmt.BuildMessageResponse( - fmt.Sprintf("msg_%d", time.Now().UnixNano()), - stdReq.ResponseModel, - norm.NormalizedMessages, - result.Thinking, - result.Text, - stdReq.ToolNames, - ) - if result.OutputTokens > 0 { - if usage, ok := respBody["usage"].(map[string]any); ok { - usage["output_tokens"] = result.OutputTokens - } - } - writeJSON(w, http.StatusOK, respBody) + writeClaudeError(w, http.StatusBadGateway, "Failed to proxy Claude request.") } -func (h *Handler) proxyViaOpenAI(w http.ResponseWriter, r *http.Request) bool { +func (h *Handler) proxyViaOpenAI(w http.ResponseWriter, r *http.Request, store ConfigReader) bool { raw, err := io.ReadAll(r.Body) if err != nil { writeClaudeError(w, http.StatusBadRequest, "invalid body") @@ -111,7 +43,15 @@ func (h *Handler) proxyViaOpenAI(w http.ResponseWriter, r *http.Request) bool { } model, _ := req["model"].(string) stream := util.ToBool(req["stream"]) - translatedReq := translatorcliproxy.ToOpenAI(sdktranslator.FormatClaude, model, raw, stream) + + // Preserve claude_mapping (fast/slow/opus routing) while proxying via OpenAI. + translateModel := model + if store != nil { + if norm, normErr := normalizeClaudeRequest(store, cloneMap(req)); normErr == nil && strings.TrimSpace(norm.Standard.ResolvedModel) != "" { + translateModel = strings.TrimSpace(norm.Standard.ResolvedModel) + } + } + translatedReq := translatorcliproxy.ToOpenAI(sdktranslator.FormatClaude, translateModel, raw, stream) isVercelPrepare := strings.TrimSpace(r.URL.Query().Get("__stream_prepare")) == "1" isVercelRelease := strings.TrimSpace(r.URL.Query().Get("__stream_release")) == "1" diff --git a/internal/adapter/claude/proxy_vercel_test.go b/internal/adapter/claude/proxy_vercel_test.go index 0439641..4a7736b 100644 --- a/internal/adapter/claude/proxy_vercel_test.go +++ b/internal/adapter/claude/proxy_vercel_test.go @@ -8,6 +8,14 @@ import ( "testing" ) +type claudeProxyStoreStub struct { + mapping map[string]string +} + +func (s claudeProxyStoreStub) ClaudeMapping() map[string]string { + return s.mapping +} + type openAIProxyStub struct { status int body string @@ -22,6 +30,21 @@ func (s openAIProxyStub) ChatCompletions(w http.ResponseWriter, _ *http.Request) _, _ = w.Write([]byte(s.body)) } +type openAIProxyCaptureStub struct { + seenModel string +} + +func (s *openAIProxyCaptureStub) ChatCompletions(w http.ResponseWriter, r *http.Request) { + var req map[string]any + _ = json.NewDecoder(r.Body).Decode(&req) + if m, ok := req["model"].(string); ok { + s.seenModel = m + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(`{"id":"ok","choices":[{"message":{"role":"assistant","content":"ok"}}]}`)) +} + func TestClaudeProxyViaOpenAIVercelPreparePassthrough(t *testing.T) { h := &Handler{OpenAI: openAIProxyStub{status: 200, body: `{"lease_id":"lease_123","payload":{"a":1}}`}} req := httptest.NewRequest(http.MethodPost, "/anthropic/v1/messages?__stream_prepare=1", strings.NewReader(`{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hi"}],"stream":true}`)) @@ -40,3 +63,22 @@ func TestClaudeProxyViaOpenAIVercelPreparePassthrough(t *testing.T) { t.Fatalf("expected lease_id in prepare passthrough, got=%v", out) } } + +func TestClaudeProxyViaOpenAIPreservesClaudeMapping(t *testing.T) { + openAI := &openAIProxyCaptureStub{} + h := &Handler{ + Store: claudeProxyStoreStub{mapping: map[string]string{"fast": "deepseek-chat", "slow": "deepseek-reasoner"}}, + OpenAI: openAI, + } + req := httptest.NewRequest(http.MethodPost, "/anthropic/v1/messages", strings.NewReader(`{"model":"claude-3-opus","messages":[{"role":"user","content":"hi"}],"stream":false}`)) + rec := httptest.NewRecorder() + + h.Messages(rec, req) + + if rec.Code != http.StatusOK { + t.Fatalf("unexpected status: %d body=%s", rec.Code, rec.Body.String()) + } + if got := strings.TrimSpace(openAI.seenModel); got != "deepseek-reasoner" { + t.Fatalf("expected mapped proxy model deepseek-reasoner, got %q", got) + } +} diff --git a/internal/adapter/claude/stream_status_test.go b/internal/adapter/claude/stream_status_test.go index c3936de..8743c44 100644 --- a/internal/adapter/claude/stream_status_test.go +++ b/internal/adapter/claude/stream_status_test.go @@ -1,7 +1,6 @@ package claude import ( - "context" "net/http" "net/http/httptest" "strings" @@ -9,48 +8,17 @@ import ( "github.com/go-chi/chi/v5" chimw "github.com/go-chi/chi/v5/middleware" - - "ds2api/internal/auth" ) -type streamStatusClaudeAuthStub struct{} +type streamStatusClaudeOpenAIStub struct{} -func (streamStatusClaudeAuthStub) Determine(_ *http.Request) (*auth.RequestAuth, error) { - return &auth.RequestAuth{ - UseConfigToken: false, - DeepSeekToken: "direct-token", - CallerID: "caller:test", - TriedAccounts: map[string]bool{}, - }, nil +func (streamStatusClaudeOpenAIStub) ChatCompletions(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello\"},\"finish_reason\":null}]}\n\n")) + _, _ = w.Write([]byte("data: [DONE]\n\n")) } -func (streamStatusClaudeAuthStub) Release(_ *auth.RequestAuth) {} - -type streamStatusClaudeDSStub struct{} - -func (streamStatusClaudeDSStub) CreateSession(_ context.Context, _ *auth.RequestAuth, _ int) (string, error) { - return "session-id", nil -} - -func (streamStatusClaudeDSStub) GetPow(_ context.Context, _ *auth.RequestAuth, _ int) (string, error) { - return "pow", nil -} - -func (streamStatusClaudeDSStub) CallCompletion(_ context.Context, _ *auth.RequestAuth, _ map[string]any, _ string, _ int) (*http.Response, error) { - body := "data: {\"p\":\"response/content\",\"v\":\"hello\"}\n" + "data: [DONE]\n" - return &http.Response{ - StatusCode: http.StatusOK, - Header: make(http.Header), - Body: ioNopCloser{strings.NewReader(body)}, - }, nil -} - -type ioNopCloser struct { - *strings.Reader -} - -func (ioNopCloser) Close() error { return nil } - type streamStatusClaudeStoreStub struct{} func (streamStatusClaudeStoreStub) ClaudeMapping() map[string]string { @@ -73,9 +41,8 @@ func captureClaudeStatusMiddleware(statuses *[]int) func(http.Handler) http.Hand func TestClaudeMessagesStreamStatusCapturedAs200(t *testing.T) { statuses := make([]int, 0, 1) h := &Handler{ - Store: streamStatusClaudeStoreStub{}, - Auth: streamStatusClaudeAuthStub{}, - DS: streamStatusClaudeDSStub{}, + Store: streamStatusClaudeStoreStub{}, + OpenAI: streamStatusClaudeOpenAIStub{}, } r := chi.NewRouter() r.Use(captureClaudeStatusMiddleware(&statuses)) @@ -83,7 +50,6 @@ func TestClaudeMessagesStreamStatusCapturedAs200(t *testing.T) { reqBody := `{"model":"claude-sonnet-4-5","messages":[{"role":"user","content":"hi"}],"stream":true}` req := httptest.NewRequest(http.MethodPost, "/anthropic/v1/messages", strings.NewReader(reqBody)) - req.Header.Set("Authorization", "Bearer direct-token") req.Header.Set("Content-Type", "application/json") rec := httptest.NewRecorder() r.ServeHTTP(rec, req) diff --git a/internal/adapter/gemini/handler_generate.go b/internal/adapter/gemini/handler_generate.go index d2f33f1..5703d0b 100644 --- a/internal/adapter/gemini/handler_generate.go +++ b/internal/adapter/gemini/handler_generate.go @@ -10,7 +10,6 @@ import ( "github.com/go-chi/chi/v5" - "ds2api/internal/auth" "ds2api/internal/sse" "ds2api/internal/translatorcliproxy" "ds2api/internal/util" @@ -19,62 +18,14 @@ import ( ) func (h *Handler) handleGenerateContent(w http.ResponseWriter, r *http.Request, stream bool) { - if h.OpenAI != nil { - if h.proxyViaOpenAI(w, r, stream) { - return - } - } - a, err := h.Auth.Determine(r) - if err != nil { - status := http.StatusUnauthorized - detail := err.Error() - if err == auth.ErrNoAccount { - status = http.StatusTooManyRequests - } - writeGeminiError(w, status, detail) + if h.OpenAI == nil { + writeGeminiError(w, http.StatusInternalServerError, "OpenAI proxy backend unavailable.") return } - defer h.Auth.Release(a) - - var req map[string]any - if err := json.NewDecoder(r.Body).Decode(&req); err != nil { - writeGeminiError(w, http.StatusBadRequest, "invalid json") + if h.proxyViaOpenAI(w, r, stream) { return } - - routeModel := strings.TrimSpace(chi.URLParam(r, "model")) - stdReq, err := normalizeGeminiRequest(h.Store, routeModel, req, stream) - if err != nil { - writeGeminiError(w, http.StatusBadRequest, err.Error()) - return - } - - sessionID, err := h.DS.CreateSession(r.Context(), a, 3) - if err != nil { - if a.UseConfigToken { - writeGeminiError(w, http.StatusUnauthorized, "Account token is invalid. Please re-login the account in admin.") - } else { - writeGeminiError(w, http.StatusUnauthorized, "Invalid token.") - } - return - } - pow, err := h.DS.GetPow(r.Context(), a, 3) - if err != nil { - writeGeminiError(w, http.StatusUnauthorized, "Failed to get PoW (invalid token or unknown error).") - return - } - payload := stdReq.CompletionPayload(sessionID) - resp, err := h.DS.CallCompletion(r.Context(), a, payload, pow, 3) - if err != nil { - writeGeminiError(w, http.StatusInternalServerError, "Failed to get completion.") - return - } - - if stream { - h.handleStreamGenerateContent(w, r, resp, stdReq.ResponseModel, stdReq.FinalPrompt, stdReq.Thinking, stdReq.Search, stdReq.ToolNames) - return - } - h.handleNonStreamGenerateContent(w, resp, stdReq.ResponseModel, stdReq.FinalPrompt, stdReq.Thinking, stdReq.ToolNames) + writeGeminiError(w, http.StatusBadGateway, "Failed to proxy Gemini request.") } func (h *Handler) proxyViaOpenAI(w http.ResponseWriter, r *http.Request, stream bool) bool { @@ -139,13 +90,7 @@ func (h *Handler) proxyViaOpenAI(w http.ResponseWriter, r *http.Request, stream defer res.Body.Close() body, _ := io.ReadAll(res.Body) if res.StatusCode < 200 || res.StatusCode >= 300 { - for k, vv := range res.Header { - for _, v := range vv { - w.Header().Add(k, v) - } - } - w.WriteHeader(res.StatusCode) - _, _ = w.Write(body) + writeGeminiErrorFromOpenAI(w, res.StatusCode, body) return true } if isVercelPrepare { @@ -165,6 +110,22 @@ func (h *Handler) proxyViaOpenAI(w http.ResponseWriter, r *http.Request, stream return true } +func writeGeminiErrorFromOpenAI(w http.ResponseWriter, status int, raw []byte) { + message := strings.TrimSpace(string(raw)) + var parsed map[string]any + if err := json.Unmarshal(raw, &parsed); err == nil { + if errObj, ok := parsed["error"].(map[string]any); ok { + if msg, ok := errObj["message"].(string); ok && strings.TrimSpace(msg) != "" { + message = strings.TrimSpace(msg) + } + } + } + if message == "" { + message = http.StatusText(status) + } + writeGeminiError(w, status, message) +} + func (h *Handler) handleNonStreamGenerateContent(w http.ResponseWriter, resp *http.Response, model, finalPrompt string, thinkingEnabled bool, toolNames []string) { defer resp.Body.Close() if resp.StatusCode != http.StatusOK { diff --git a/internal/adapter/gemini/handler_test.go b/internal/adapter/gemini/handler_test.go index 9316833..20cc0e6 100644 --- a/internal/adapter/gemini/handler_test.go +++ b/internal/adapter/gemini/handler_test.go @@ -61,6 +61,40 @@ func (m testGeminiDS) CallCompletion(_ context.Context, _ *auth.RequestAuth, _ m return m.resp, nil } +type geminiOpenAIErrorStub struct { + status int + body string +} + +func (s geminiOpenAIErrorStub) ChatCompletions(w http.ResponseWriter, _ *http.Request) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(s.status) + _, _ = w.Write([]byte(s.body)) +} + +type geminiOpenAISuccessStub struct { + stream bool + body string +} + +func (s geminiOpenAISuccessStub) ChatCompletions(w http.ResponseWriter, _ *http.Request) { + if s.stream { + w.Header().Set("Content-Type", "text/event-stream") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"hello \"},\"finish_reason\":null}]}\n\n")) + _, _ = w.Write([]byte("data: {\"id\":\"chatcmpl-1\",\"object\":\"chat.completion.chunk\",\"choices\":[{\"index\":0,\"delta\":{\"content\":\"world\"},\"finish_reason\":\"stop\"}]}\n\n")) + _, _ = w.Write([]byte("data: [DONE]\n\n")) + return + } + out := s.body + if strings.TrimSpace(out) == "" { + out = `{"id":"chatcmpl-1","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"eval_javascript","arguments":"{\"code\":\"1+1\"}"}}]},"finish_reason":"tool_calls"}]}` + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write([]byte(out)) +} + func makeGeminiUpstreamResponse(lines ...string) *http.Response { body := strings.Join(lines, "\n") if !strings.HasSuffix(body, "\n") { @@ -98,14 +132,11 @@ func TestGeminiRoutesRegistered(t *testing.T) { } func TestGenerateContentReturnsFunctionCallParts(t *testing.T) { - upstream := makeGeminiUpstreamResponse( - `data: {"p":"response/content","v":"{\"tool_calls\":[{\"name\":\"eval_javascript\",\"input\":{\"code\":\"1+1\"}}]}"}`, - `data: [DONE]`, - ) h := &Handler{ Store: testGeminiConfig{}, - Auth: testGeminiAuth{}, - DS: testGeminiDS{resp: upstream}, + OpenAI: geminiOpenAISuccessStub{ + body: `{"id":"chatcmpl-1","object":"chat.completion","choices":[{"index":0,"message":{"role":"assistant","tool_calls":[{"id":"call_1","type":"function","function":{"name":"eval_javascript","arguments":"{\"code\":\"1+1\"}"}}]},"finish_reason":"tool_calls"}]}`, + }, } r := chi.NewRouter() RegisterRoutes(r, h) @@ -115,7 +146,6 @@ func TestGenerateContentReturnsFunctionCallParts(t *testing.T) { "tools":[{"functionDeclarations":[{"name":"eval_javascript","description":"eval","parameters":{"type":"object","properties":{"code":{"type":"string"}}}}]}] }` req := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:generateContent", strings.NewReader(body)) - req.Header.Set("Authorization", "Bearer direct-token") rec := httptest.NewRecorder() r.ServeHTTP(rec, req) if rec.Code != http.StatusOK { @@ -144,11 +174,7 @@ func TestGenerateContentReturnsFunctionCallParts(t *testing.T) { } func TestGenerateContentMixedToolSnippetAlsoTriggersFunctionCall(t *testing.T) { - upstream := makeGeminiUpstreamResponse( - `data: {"p":"response/content","v":"我来调用工具\n{\"tool_calls\":[{\"name\":\"eval_javascript\",\"input\":{\"code\":\"1+1\"}}]}"}`, - `data: [DONE]`, - ) - h := &Handler{Store: testGeminiConfig{}, Auth: testGeminiAuth{}, DS: testGeminiDS{resp: upstream}} + h := &Handler{Store: testGeminiConfig{}, OpenAI: geminiOpenAISuccessStub{}} r := chi.NewRouter() RegisterRoutes(r, h) @@ -157,7 +183,6 @@ func TestGenerateContentMixedToolSnippetAlsoTriggersFunctionCall(t *testing.T) { "tools":[{"functionDeclarations":[{"name":"eval_javascript","description":"eval","parameters":{"type":"object","properties":{"code":{"type":"string"}}}}]}] }` req := httptest.NewRequest(http.MethodPost, "/v1beta/models/gemini-2.5-pro:generateContent", strings.NewReader(body)) - req.Header.Set("Authorization", "Bearer direct-token") rec := httptest.NewRecorder() r.ServeHTTP(rec, req) @@ -180,38 +205,25 @@ func TestGenerateContentMixedToolSnippetAlsoTriggersFunctionCall(t *testing.T) { } func TestStreamGenerateContentEmitsSSE(t *testing.T) { - upstream := makeGeminiUpstreamResponse( - `data: {"p":"response/content","v":"hello "}`, - `data: {"p":"response/content","v":"world"}`, - `data: [DONE]`, - ) h := &Handler{ - Store: testGeminiConfig{}, - Auth: testGeminiAuth{}, - DS: testGeminiDS{resp: upstream}, + Store: testGeminiConfig{}, + OpenAI: geminiOpenAISuccessStub{stream: true}, } r := chi.NewRouter() RegisterRoutes(r, h) body := `{"contents":[{"role":"user","parts":[{"text":"hello"}]}]}` req := httptest.NewRequest(http.MethodPost, "/v1/models/gemini-2.5-pro:streamGenerateContent?alt=sse", strings.NewReader(body)) - req.Header.Set("Authorization", "Bearer direct-token") rec := httptest.NewRecorder() r.ServeHTTP(rec, req) if rec.Code != http.StatusOK { t.Fatalf("expected 200, got %d body=%s", rec.Code, rec.Body.String()) } - if !strings.Contains(rec.Body.String(), "data: ") { - t.Fatalf("expected SSE data frames, got body=%s", rec.Body.String()) - } - if !strings.Contains(rec.Body.String(), `"finishReason":"STOP"`) { - t.Fatalf("expected stream finish frame, got body=%s", rec.Body.String()) - } frames := extractGeminiSSEFrames(t, rec.Body.String()) if len(frames) == 0 { - t.Fatalf("expected non-empty sse frames, body=%s", rec.Body.String()) + t.Fatalf("expected non-empty stream frames, body=%s", rec.Body.String()) } last := frames[len(frames)-1] candidates, _ := last["candidates"].([]any) @@ -229,16 +241,44 @@ func TestStreamGenerateContentEmitsSSE(t *testing.T) { } } +func TestGenerateContentOpenAIProxyErrorUsesGeminiEnvelope(t *testing.T) { + h := &Handler{ + Store: testGeminiConfig{}, + OpenAI: geminiOpenAIErrorStub{status: http.StatusUnauthorized, body: `{"error":{"message":"invalid api key"}}`}, + } + r := chi.NewRouter() + RegisterRoutes(r, h) + + req := httptest.NewRequest(http.MethodPost, "/v1/models/gemini-2.5-pro:generateContent", strings.NewReader(`{"contents":[{"role":"user","parts":[{"text":"hi"}]}]}`)) + rec := httptest.NewRecorder() + r.ServeHTTP(rec, req) + + if rec.Code != http.StatusUnauthorized { + t.Fatalf("expected 401, got %d body=%s", rec.Code, rec.Body.String()) + } + var out map[string]any + if err := json.Unmarshal(rec.Body.Bytes(), &out); err != nil { + t.Fatalf("expected json body: %v", err) + } + errObj, _ := out["error"].(map[string]any) + if errObj["status"] != "UNAUTHENTICATED" { + t.Fatalf("expected Gemini status UNAUTHENTICATED, got=%v", errObj["status"]) + } + if errObj["message"] != "invalid api key" { + t.Fatalf("expected parsed error message, got=%v", errObj["message"]) + } +} + func extractGeminiSSEFrames(t *testing.T, body string) []map[string]any { t.Helper() scanner := bufio.NewScanner(strings.NewReader(body)) out := make([]map[string]any, 0, 4) for scanner.Scan() { line := strings.TrimSpace(scanner.Text()) - if !strings.HasPrefix(line, "data: ") { - continue + raw := line + if strings.HasPrefix(line, "data: ") { + raw = strings.TrimSpace(strings.TrimPrefix(line, "data: ")) } - raw := strings.TrimSpace(strings.TrimPrefix(line, "data: ")) if raw == "" { continue } From a6a9863fc3392368257c60c36a57dc495e57a964 Mon Sep 17 00:00:00 2001 From: "CJACK." Date: Fri, 3 Apr 2026 01:59:35 +0800 Subject: [PATCH 2/2] Preserve upstream headers on Gemini proxy error responses --- internal/adapter/gemini/handler_generate.go | 5 +++++ internal/adapter/gemini/handler_test.go | 23 ++++++++++++++++++++- 2 files changed, 27 insertions(+), 1 deletion(-) diff --git a/internal/adapter/gemini/handler_generate.go b/internal/adapter/gemini/handler_generate.go index 5703d0b..a7de92d 100644 --- a/internal/adapter/gemini/handler_generate.go +++ b/internal/adapter/gemini/handler_generate.go @@ -90,6 +90,11 @@ func (h *Handler) proxyViaOpenAI(w http.ResponseWriter, r *http.Request, stream defer res.Body.Close() body, _ := io.ReadAll(res.Body) if res.StatusCode < 200 || res.StatusCode >= 300 { + for k, vv := range res.Header { + for _, v := range vv { + w.Header().Add(k, v) + } + } writeGeminiErrorFromOpenAI(w, res.StatusCode, body) return true } diff --git a/internal/adapter/gemini/handler_test.go b/internal/adapter/gemini/handler_test.go index 20cc0e6..fdb4b79 100644 --- a/internal/adapter/gemini/handler_test.go +++ b/internal/adapter/gemini/handler_test.go @@ -64,9 +64,13 @@ func (m testGeminiDS) CallCompletion(_ context.Context, _ *auth.RequestAuth, _ m type geminiOpenAIErrorStub struct { status int body string + headers map[string]string } func (s geminiOpenAIErrorStub) ChatCompletions(w http.ResponseWriter, _ *http.Request) { + for k, v := range s.headers { + w.Header().Set(k, v) + } w.Header().Set("Content-Type", "application/json") w.WriteHeader(s.status) _, _ = w.Write([]byte(s.body)) @@ -244,7 +248,15 @@ func TestStreamGenerateContentEmitsSSE(t *testing.T) { func TestGenerateContentOpenAIProxyErrorUsesGeminiEnvelope(t *testing.T) { h := &Handler{ Store: testGeminiConfig{}, - OpenAI: geminiOpenAIErrorStub{status: http.StatusUnauthorized, body: `{"error":{"message":"invalid api key"}}`}, + OpenAI: geminiOpenAIErrorStub{ + status: http.StatusUnauthorized, + body: `{"error":{"message":"invalid api key"}}`, + headers: map[string]string{ + "WWW-Authenticate": `Bearer realm="example"`, + "Retry-After": "30", + "X-RateLimit-Remaining": "0", + }, + }, } r := chi.NewRouter() RegisterRoutes(r, h) @@ -267,6 +279,15 @@ func TestGenerateContentOpenAIProxyErrorUsesGeminiEnvelope(t *testing.T) { if errObj["message"] != "invalid api key" { t.Fatalf("expected parsed error message, got=%v", errObj["message"]) } + if got := rec.Header().Get("WWW-Authenticate"); got == "" { + t.Fatalf("expected WWW-Authenticate header to be preserved") + } + if got := rec.Header().Get("Retry-After"); got != "30" { + t.Fatalf("expected Retry-After header 30, got=%q", got) + } + if got := rec.Header().Get("X-RateLimit-Remaining"); got != "0" { + t.Fatalf("expected X-RateLimit-Remaining header 0, got=%q", got) + } } func extractGeminiSSEFrames(t *testing.T, body string) []map[string]any {