From d014f3e989f82f292419edacaee622ddf189d1b0 Mon Sep 17 00:00:00 2001 From: xiwuqi <64734786+xiwuqi@users.noreply.github.com> Date: Mon, 23 Mar 2026 00:41:40 -0500 Subject: [PATCH] fix(api): include auth header in local model probe (#1896) --- web/backend/api/gateway_test.go | 8 +++--- web/backend/api/model_status.go | 15 ++++++----- web/backend/api/model_status_test.go | 37 ++++++++++++++++++++++++++++ web/backend/api/models_test.go | 20 +++++++-------- 4 files changed, 60 insertions(+), 20 deletions(-) create mode 100644 web/backend/api/model_status_test.go diff --git a/web/backend/api/gateway_test.go b/web/backend/api/gateway_test.go index 504d091af..387c5ac53 100644 --- a/web/backend/api/gateway_test.go +++ b/web/backend/api/gateway_test.go @@ -169,7 +169,7 @@ func TestGatewayStartReady_LocalModelWithoutAPIKey(t *testing.T) { defer cleanup() resetModelProbeHooks(t) - probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { + probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool { return false } @@ -206,8 +206,8 @@ func TestGatewayStartReady_LocalModelWithRunningService(t *testing.T) { defer cleanup() resetModelProbeHooks(t) - probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { - return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model" + probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool { + return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model" && apiKey == "" } cfg, err := config.LoadConfig(configPath) @@ -240,7 +240,7 @@ func TestGatewayStartReady_RemoteVLLMWithAPIKeyDoesNotProbe(t *testing.T) { defer cleanup() resetModelProbeHooks(t) - probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { + probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool { t.Fatalf("unexpected OpenAI-compatible probe for %q (%q)", apiBase, modelID) return false } diff --git a/web/backend/api/model_status.go b/web/backend/api/model_status.go index 22bf5c15b..b56fe5f39 100644 --- a/web/backend/api/model_status.go +++ b/web/backend/api/model_status.go @@ -82,14 +82,14 @@ func probeLocalModelAvailability(m config.ModelConfig) bool { case "ollama": return probeOllamaModelFunc(apiBase, modelID) case "vllm": - return probeOpenAICompatibleModelFunc(apiBase, modelID) + return probeOpenAICompatibleModelFunc(apiBase, modelID, m.APIKey) case "github-copilot", "copilot": return probeTCPServiceFunc(apiBase) case "claude-cli", "claudecli", "codex-cli", "codexcli": return true default: if hasLocalAPIBase(apiBase) { - return probeOpenAICompatibleModelFunc(apiBase, modelID) + return probeOpenAICompatibleModelFunc(apiBase, modelID, m.APIKey) } return false } @@ -209,7 +209,7 @@ func probeOllamaModel(apiBase, modelID string) bool { Model string `json:"model"` } `json:"models"` } - if err := getJSON(root+"/api/tags", &resp); err != nil { + if err := getJSON(root+"/api/tags", &resp, ""); err != nil { return false } @@ -221,7 +221,7 @@ func probeOllamaModel(apiBase, modelID string) bool { return false } -func probeOpenAICompatibleModel(apiBase, modelID string) bool { +func probeOpenAICompatibleModel(apiBase, modelID, apiKey string) bool { if strings.TrimSpace(apiBase) == "" { return false } @@ -231,7 +231,7 @@ func probeOpenAICompatibleModel(apiBase, modelID string) bool { ID string `json:"id"` } `json:"data"` } - if err := getJSON(strings.TrimRight(strings.TrimSpace(apiBase), "/")+"/models", &resp); err != nil { + if err := getJSON(strings.TrimRight(strings.TrimSpace(apiBase), "/")+"/models", &resp, apiKey); err != nil { return false } @@ -243,11 +243,14 @@ func probeOpenAICompatibleModel(apiBase, modelID string) bool { return false } -func getJSON(rawURL string, out any) error { +func getJSON(rawURL string, out any, apiKey string) error { req, err := http.NewRequest(http.MethodGet, rawURL, nil) if err != nil { return err } + if apiKey = strings.TrimSpace(apiKey); apiKey != "" { + req.Header.Set("Authorization", "Bearer "+apiKey) + } client := &http.Client{Timeout: modelProbeTimeout} resp, err := client.Do(req) diff --git a/web/backend/api/model_status_test.go b/web/backend/api/model_status_test.go new file mode 100644 index 000000000..047af7a4d --- /dev/null +++ b/web/backend/api/model_status_test.go @@ -0,0 +1,37 @@ +package api + +import ( + "net/http" + "net/http/httptest" + "testing" + + "github.com/sipeed/picoclaw/pkg/config" +) + +func TestProbeLocalModelAvailability_OpenAICompatibleIncludesAPIKey(t *testing.T) { + const apiKey = "test-api-key" + + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/v1/models" { + t.Fatalf("path = %q, want %q", r.URL.Path, "/v1/models") + } + if got := r.Header.Get("Authorization"); got != "Bearer "+apiKey { + http.Error(w, "missing auth", http.StatusUnauthorized) + return + } + + w.Header().Set("Content-Type", "application/json") + _, _ = w.Write([]byte(`{"data":[{"id":"custom-model"}]}`)) + })) + defer srv.Close() + + model := config.ModelConfig{ + Model: "openai/custom-model", + APIBase: srv.URL + "/v1", + APIKey: apiKey, + } + + if !probeLocalModelAvailability(model) { + t.Fatal("probeLocalModelAvailability() = false, want true when api_key is configured") + } +} diff --git a/web/backend/api/models_test.go b/web/backend/api/models_test.go index 2377b5b66..1ec5fb8c9 100644 --- a/web/backend/api/models_test.go +++ b/web/backend/api/models_test.go @@ -36,11 +36,11 @@ func TestHandleListModels_ConfiguredStatusUsesRuntimeProbesForLocalModels(t *tes var ollamaProbes []string var tcpProbes []string - probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { + probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool { mu.Lock() - openAIProbes = append(openAIProbes, apiBase+"|"+modelID) + openAIProbes = append(openAIProbes, apiBase+"|"+modelID+"|"+apiKey) mu.Unlock() - return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model" + return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model" && apiKey == "" } probeOllamaModelFunc = func(apiBase, modelID string) bool { mu.Lock() @@ -131,7 +131,7 @@ func TestHandleListModels_ConfiguredStatusUsesRuntimeProbesForLocalModels(t *tes if !got["copilot-gpt-5.4"] { t.Fatalf("copilot model configured = false, want true when local bridge probe succeeds") } - if len(openAIProbes) != 1 || openAIProbes[0] != "http://127.0.0.1:8000/v1|custom-model" { + if len(openAIProbes) != 1 || openAIProbes[0] != "http://127.0.0.1:8000/v1|custom-model|" { t.Fatalf("openAI probes = %#v, want only local vllm probe", openAIProbes) } if len(ollamaProbes) != 1 || ollamaProbes[0] != "http://localhost:11434/v1|llama3" { @@ -205,7 +205,7 @@ func TestHandleListModels_ProbesLocalModelsConcurrently(t *testing.T) { started := make(chan string, 2) release := make(chan struct{}) - probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { + probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool { started <- apiBase + "|" + modelID <-release return true @@ -265,9 +265,9 @@ func TestHandleListModels_NormalizesWildcardLocalAPIBaseForProbe(t *testing.T) { resetModelProbeHooks(t) var gotProbe string - probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { - gotProbe = apiBase + "|" + modelID - return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model" + probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool { + gotProbe = apiBase + "|" + modelID + "|" + apiKey + return apiBase == "http://127.0.0.1:8000/v1" && modelID == "custom-model" && apiKey == "" } cfg, err := config.LoadConfig(configPath) @@ -307,7 +307,7 @@ func TestHandleListModels_NormalizesWildcardLocalAPIBaseForProbe(t *testing.T) { if !resp.Models[0].Configured { t.Fatal("wildcard-bound local model configured = false, want true after probe host normalization") } - if gotProbe != "http://127.0.0.1:8000/v1|custom-model" { - t.Fatalf("probe api base = %q, want %q", gotProbe, "http://127.0.0.1:8000/v1|custom-model") + if gotProbe != "http://127.0.0.1:8000/v1|custom-model|" { + t.Fatalf("probe api base = %q, want %q", gotProbe, "http://127.0.0.1:8000/v1|custom-model|") } }