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
4 changed files with 60 additions and 20 deletions
+4 -4
View File
@@ -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
}
+9 -6
View File
@@ -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)
+37
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")
}
}
+10 -10
View File
@@ -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|")
}
}