Merge pull request #142 from BaaRai/fix/llm-stored-key-endpoint

fix(llm): send a saved API key only to the endpoint it was saved for
2026-11-13 11:40:53 +08:00 · 2026-11-03 23:49:33 +01:00
parent e94471609f 334f6fdddc
commit 0093781c8b
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
 	}