mirror of
https://github.com/sipeed/picoclaw.git
synced 2026-06-12 18:08:54 +00:00
fix(api): include auth header in local model probe (#1896)
This commit is contained in:
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
@@ -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|")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user