fix(api): include auth header in local model probe (#1896)

This commit is contained in:
xiwuqi 2026-03-23 00:41:40 -05:00 committed by GitHub
parent 40279c8dde
commit d014f3e989
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 60 additions and 20 deletions

View file

@ -169,7 +169,7 @@ func TestGatewayStartReady_LocalModelWithoutAPIKey(t *testing.T) {
defer cleanup() defer cleanup()
resetModelProbeHooks(t) resetModelProbeHooks(t)
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
return false return false
} }
@ -206,8 +206,8 @@ func TestGatewayStartReady_LocalModelWithRunningService(t *testing.T) {
defer cleanup() defer cleanup()
resetModelProbeHooks(t) resetModelProbeHooks(t)
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
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 == ""
} }
cfg, err := config.LoadConfig(configPath) cfg, err := config.LoadConfig(configPath)
@ -240,7 +240,7 @@ func TestGatewayStartReady_RemoteVLLMWithAPIKeyDoesNotProbe(t *testing.T) {
defer cleanup() defer cleanup()
resetModelProbeHooks(t) 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) t.Fatalf("unexpected OpenAI-compatible probe for %q (%q)", apiBase, modelID)
return false return false
} }

View file

@ -82,14 +82,14 @@ func probeLocalModelAvailability(m config.ModelConfig) bool {
case "ollama": case "ollama":
return probeOllamaModelFunc(apiBase, modelID) return probeOllamaModelFunc(apiBase, modelID)
case "vllm": case "vllm":
return probeOpenAICompatibleModelFunc(apiBase, modelID) return probeOpenAICompatibleModelFunc(apiBase, modelID, m.APIKey)
case "github-copilot", "copilot": case "github-copilot", "copilot":
return probeTCPServiceFunc(apiBase) return probeTCPServiceFunc(apiBase)
case "claude-cli", "claudecli", "codex-cli", "codexcli": case "claude-cli", "claudecli", "codex-cli", "codexcli":
return true return true
default: default:
if hasLocalAPIBase(apiBase) { if hasLocalAPIBase(apiBase) {
return probeOpenAICompatibleModelFunc(apiBase, modelID) return probeOpenAICompatibleModelFunc(apiBase, modelID, m.APIKey)
} }
return false return false
} }
@ -209,7 +209,7 @@ func probeOllamaModel(apiBase, modelID string) bool {
Model string `json:"model"` Model string `json:"model"`
} `json:"models"` } `json:"models"`
} }
if err := getJSON(root+"/api/tags", &resp); err != nil { if err := getJSON(root+"/api/tags", &resp, ""); err != nil {
return false return false
} }
@ -221,7 +221,7 @@ func probeOllamaModel(apiBase, modelID string) bool {
return false return false
} }
func probeOpenAICompatibleModel(apiBase, modelID string) bool { func probeOpenAICompatibleModel(apiBase, modelID, apiKey string) bool {
if strings.TrimSpace(apiBase) == "" { if strings.TrimSpace(apiBase) == "" {
return false return false
} }
@ -231,7 +231,7 @@ func probeOpenAICompatibleModel(apiBase, modelID string) bool {
ID string `json:"id"` ID string `json:"id"`
} `json:"data"` } `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 return false
} }
@ -243,11 +243,14 @@ func probeOpenAICompatibleModel(apiBase, modelID string) bool {
return false 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) req, err := http.NewRequest(http.MethodGet, rawURL, nil)
if err != nil { if err != nil {
return err return err
} }
if apiKey = strings.TrimSpace(apiKey); apiKey != "" {
req.Header.Set("Authorization", "Bearer "+apiKey)
}
client := &http.Client{Timeout: modelProbeTimeout} client := &http.Client{Timeout: modelProbeTimeout}
resp, err := client.Do(req) resp, err := client.Do(req)

View file

@ -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")
}
}

View file

@ -36,11 +36,11 @@ func TestHandleListModels_ConfiguredStatusUsesRuntimeProbesForLocalModels(t *tes
var ollamaProbes []string var ollamaProbes []string
var tcpProbes []string var tcpProbes []string
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
mu.Lock() mu.Lock()
openAIProbes = append(openAIProbes, apiBase+"|"+modelID) openAIProbes = append(openAIProbes, apiBase+"|"+modelID+"|"+apiKey)
mu.Unlock() 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 { probeOllamaModelFunc = func(apiBase, modelID string) bool {
mu.Lock() mu.Lock()
@ -131,7 +131,7 @@ func TestHandleListModels_ConfiguredStatusUsesRuntimeProbesForLocalModels(t *tes
if !got["copilot-gpt-5.4"] { if !got["copilot-gpt-5.4"] {
t.Fatalf("copilot model configured = false, want true when local bridge probe succeeds") 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) t.Fatalf("openAI probes = %#v, want only local vllm probe", openAIProbes)
} }
if len(ollamaProbes) != 1 || ollamaProbes[0] != "http://localhost:11434/v1|llama3" { 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) started := make(chan string, 2)
release := make(chan struct{}) release := make(chan struct{})
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
started <- apiBase + "|" + modelID started <- apiBase + "|" + modelID
<-release <-release
return true return true
@ -265,9 +265,9 @@ func TestHandleListModels_NormalizesWildcardLocalAPIBaseForProbe(t *testing.T) {
resetModelProbeHooks(t) resetModelProbeHooks(t)
var gotProbe string var gotProbe string
probeOpenAICompatibleModelFunc = func(apiBase, modelID string) bool { probeOpenAICompatibleModelFunc = func(apiBase, modelID, apiKey string) bool {
gotProbe = apiBase + "|" + modelID gotProbe = apiBase + "|" + modelID + "|" + apiKey
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 == ""
} }
cfg, err := config.LoadConfig(configPath) cfg, err := config.LoadConfig(configPath)
@ -307,7 +307,7 @@ func TestHandleListModels_NormalizesWildcardLocalAPIBaseForProbe(t *testing.T) {
if !resp.Models[0].Configured { if !resp.Models[0].Configured {
t.Fatal("wildcard-bound local model configured = false, want true after probe host normalization") 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" { 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") t.Fatalf("probe api base = %q, want %q", gotProbe, "http://127.0.0.1:8000/v1|custom-model|")
} }
} }