fix(llm): send a saved API key only to the endpoint it was saved for

POST /api/llm/test and POST /api/llm/models take provider and base_url from the form but fell back to the stored profile key (and, for the test, the global key) when the form had none. Any base_url typed into the form, or supplied by a crafted request, received the saved key. The saved key is now used only when the request's provider and base URL normalize to the ones it was saved with (sameLLMEndpoint, built on agent.ConfigFrom). Otherwise the key must be entered in the form, and the error says why.
2026-11-13 11:40:53 +08:00 · 2026-11-03 23:25:47 +01:00
parent 23cf14cf0a
commit 334f6fdddc
3 changed files with 175 additions and 9 deletions
@@ -0,0 +1,127 @@
+package server
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+
+	"github.com/Autumn-27/artex/agent"
+	"github.com/Autumn-27/artex/db"
+)
+
+func TestSameLLMEndpoint(t *testing.T) {
+	for _, tc := range []struct {
+		provA, baseA, provB, baseB string
+		want                       bool
+	}{
+		{"anthropic", "https://gw.example", "anthropic", "https://gw.example/", true},
+		{"anthropic", "https://gw.example/v1/messages", "anthropic", "https://gw.example", true},
+		{"", "", "anthropic", "", true}, // empty provider defaults to anthropic
+		{"openai", " https://gw.example/v1/chat/completions ", "openai", "https://gw.example/v1", true},
+		{"anthropic", "https://evil.example", "anthropic", "https://gw.example", false},
+		{"anthropic", "", "anthropic", "https://gw.example", false}, // canonical host != saved host
+		{"openai", "", "anthropic", "", false},                      // different canonical hosts
+		{"openai", "https://gw.example/v1", "openai-responses", "https://gw.example/v1", false},
+	} {
+		if got := sameLLMEndpoint(tc.provA, tc.baseA, tc.provB, tc.baseB); got != tc.want {
+			t.Errorf("sameLLMEndpoint(%q,%q,%q,%q)=%v want %v", tc.provA, tc.baseA, tc.provB, tc.baseB, got, tc.want)
+		}
+	}
+}
+
+// keyRecorder is a fake Anthropic endpoint that records the x-api-key it receives.
+func keyRecorder(t *testing.T) (*httptest.Server, *atomic.Int32, *atomic.Value) {
+	var hits atomic.Int32
+	var key atomic.Value
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		hits.Add(1)
+		key.Store(r.Header.Get("x-api-key"))
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"id":"m","type":"message","role":"assistant","model":"m","content":[{"type":"text","text":"OK"}],"stop_reason":"end_turn","usage":{"input_tokens":1,"output_tokens":1},"data":[]}`))
+	}))
+	t.Cleanup(srv.Close)
+	return srv, &hits, &key
+}
+
+func postJSON(t *testing.T, h http.HandlerFunc, path string, body map[string]any) map[string]any {
+	t.Helper()
+	b, _ := json.Marshal(body)
+	req := httptest.NewRequest(http.MethodPost, path, bytes.NewReader(b))
+	req.Header.Set("Content-Type", "application/json")
+	rec := httptest.NewRecorder()
+	h(rec, req)
+	var resp map[string]any
+	_ = json.Unmarshal(rec.Body.Bytes(), &resp)
+	return resp
+}
+
+// The global key must not follow a base_url typed into the test form.
+func TestLLMTestGlobalKeyOnlyForGlobalEndpoint(t *testing.T) {
+	saved, savedHits, savedKey := keyRecorder(t)
+	other, otherHits, _ := keyRecorder(t)
+	s := &Server{}
+	s.llmCfg = agent.ConfigFrom("anthropic", "m", saved.URL, "global-secret", "")
+
+	resp := postJSON(t, s.testLLM, "/api/llm/test", map[string]any{
+		"provider": "anthropic", "model": "m", "base_url": other.URL, "api_key": "", "streaming": false,
+	})
+	if otherHits.Load() != 0 || resp["ok"] != false {
+		t.Fatalf("global key went to a form-supplied host: hits=%d resp=%v", otherHits.Load(), resp)
+	}
+
+	postJSON(t, s.testLLM, "/api/llm/test", map[string]any{
+		"provider": "anthropic", "model": "m", "base_url": saved.URL + "/", "api_key": "", "streaming": false,
+	})
+	if k, _ := savedKey.Load().(string); savedHits.Load() != 1 || k != "global-secret" {
+		t.Fatalf("global key not used for its own endpoint: hits=%d key=%q", savedHits.Load(), k)
+	}
+}
+
+// A stored profile's key goes only to that profile's base URL, for both the
+// connection test and the model listing.
+func TestStoredProfileKeyOnlyForStoredEndpoint(t *testing.T) {
+	m, err := NewManager(t.TempDir(), "")
+	if err != nil {
+		t.Skipf("postgres unavailable (%v)", err)
+	}
+	t.Cleanup(func() { m.Close() })
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
+	s.cfgMu.Lock()
+	s.llmCfg = agent.Config{} // no global key to fall back on
+	s.cfgMu.Unlock()
+
+	saved, savedHits, savedKey := keyRecorder(t)
+	other, otherHits, _ := keyRecorder(t)
+	pid, err := m.pg.SaveProfile(&db.LLMProfile{
+		Name: "key-endpoint", Format: "anthropic", Model: "m", BaseURL: saved.URL, APIKey: "profile-secret",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _, _ = m.pg.Exec("DELETE FROM llm_profiles WHERE id=$1", pid) })
+
+	for _, ep := range []struct {
+		path string
+		h    http.HandlerFunc
+	}{{"/api/llm/test", s.testLLM}, {"/api/llm/models", s.pgListModels}} {
+		resp := postJSON(t, ep.h, ep.path, map[string]any{
+			"provider": "anthropic", "model": "m", "base_url": other.URL, "api_key": "", "profile_id": pid, "streaming": false,
+		})
+		if otherHits.Load() != 0 || resp["ok"] != false {
+			t.Fatalf("%s: stored key went to a form-supplied host: hits=%d resp=%v", ep.path, otherHits.Load(), resp)
+		}
+		before := savedHits.Load()
+		postJSON(t, ep.h, ep.path, map[string]any{
+			"provider": "anthropic", "model": "m", "base_url": saved.URL, "api_key": "", "profile_id": pid, "streaming": false,
+		})
+		if k, _ := savedKey.Load().(string); savedHits.Load() == before || k != "profile-secret" {
+			t.Fatalf("%s: stored key not used for its own endpoint: key=%q", ep.path, k)
+		}
+	}
+}
@@ -1466,19 +1466,38 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) {
 	if cfg.NeedsAPIKey() {
 		// API key resolution precedence: form input > the named profile's stored key > the global config's key.
 		// A stored profile's key is never sent back to the browser, so the form is empty when testing a stored profile and the key comes from the DB.
+		// A saved key is only sent to the endpoint it was saved for: base_url and
+		// provider come from the form, so without this check the test button sends
+		// the saved key to any host typed into the form.
+		withheld := false
 		if stored != nil && cfg.APIKey == "" {
-			// Trim like ConfigFrom does for form values: a key stored with
-			// surrounding whitespace runs fine in conversations (the runtime
-			// path re-trims) but would be sent raw here and fail auth.
-			cfg.APIKey = strings.TrimSpace(stored.APIKey)
+			if sameLLMEndpoint(req.Provider, req.BaseURL, stored.Format, stored.BaseURL) {
+				// Trim like ConfigFrom does for form values: a key stored with
+				// surrounding whitespace runs fine in conversations (the runtime
+				// path re-trims) but would be sent raw here and fail auth.
+				cfg.APIKey = strings.TrimSpace(stored.APIKey)
+			} else {
+				withheld = true
+			}
 		}
 		if cfg.APIKey == "" {
 			s.cfgMu.Lock()
-			cfg.APIKey = s.llmCfg.APIKey
+			global := s.llmCfg
 			s.cfgMu.Unlock()
+			if global.APIKey != "" {
+				if sameLLMEndpoint(req.Provider, req.BaseURL, global.Provider(), global.BaseURL) {
+					cfg.APIKey = global.APIKey
+				} else {
+					withheld = true
+				}
+			}
 		}
 		if cfg.APIKey == "" {
-			writeJSON(w, 200, map[string]any{"ok": false, "error": "no API key provided"})
+			msg := "no API key provided"
+			if withheld {
+				msg = "no API key provided: the saved key is only sent to the endpoint it was saved for; enter the key to test a different base URL or provider"
+			}
+			writeJSON(w, 200, map[string]any{"ok": false, "error": msg})
 			return
 		}
 	}
@@ -1496,6 +1515,16 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) {
 	})
 }

+// sameLLMEndpoint reports whether two provider/base URL pairs resolve to the same
+// endpoint after ConfigFrom's normalization (trimmed, endpoint suffix stripped;
+// an empty base URL means the provider's canonical host). A saved API key may only
+// be sent when the request's endpoint matches the one it was saved for.
+func sameLLMEndpoint(provA, baseA, provB, baseB string) bool {
+	a := agent.ConfigFrom(provA, "", baseA, "", "")
+	b := agent.ConfigFrom(provB, "", baseB, "", "")
+	return a.Format == b.Format && a.BaseURL == b.BaseURL
+}
+
 // truncateReply clips a connection-test reply for display. A model told to answer
 // "OK" can still ramble (or think out loud); the UI only needs enough to show it
 // really said something. Rune-based so multibyte text never splits mid-character.
@@ -2053,17 +2053,27 @@ func (s *Server) pgListModels(w http.ResponseWriter, r *http.Request) {
 		writeErr(w, 400, err.Error())
 		return
 	}
-	// Resolve API key: form input > profile stored key.
+	// Resolve API key: form input > profile stored key. The stored key is only
+	// sent to the endpoint it was saved for, since base_url comes from the form.
 	apiKey := strings.TrimSpace(req.APIKey)
+	withheld := false
 	if apiKey == "" && req.ProfileID != nil {
 		if p, err := s.m.pg.ProfileByID(*req.ProfileID); err == nil && p != nil {
-			apiKey = p.APIKey
+			if sameLLMEndpoint(req.Provider, req.BaseURL, p.Format, p.BaseURL) {
+				apiKey = p.APIKey
+			} else {
+				withheld = p.APIKey != ""
+			}
 		}
 	}
 	provider := strings.TrimSpace(req.Provider)
 	// Ollama serves its model list keyless on the local server: no key gate.
 	if apiKey == "" && provider != "ollama" {
-		writeJSON(w, 200, map[string]any{"ok": false, "error": "no API Key provided"})
+		msg := "no API Key provided"
+		if withheld {
+			msg = "no API Key provided: the saved key is only sent to the endpoint it was saved for; enter the key to list models from a different base URL or provider"
+		}
+		writeJSON(w, 200, map[string]any{"ok": false, "error": msg})
 		return
 	}