Merge pull request #67 from BaaRai/feat/llm-test-stored-key

fix: llm/test stored-key fallback skips the runtime's key normalization

db/config.go

@@ -6,6 +6,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"strings"
 	"time"
 )
 
@@ -160,6 +161,10 @@ ORDER BY is_default DESC, priority DESC, id ASC`)
 
 // SaveProfile inserts (id==0) or updates a profile. Empty apiKey on update keeps existing.
 func (d *DB) SaveProfile(p *LLMProfile) (int64, error) {
+	// Pasted keys can carry surrounding whitespace; every reader that goes
+	// through ConfigFrom trims, so store the trimmed form or the raw value
+	// leaks into paths that assign it directly (the llm/test fallback).
+	p.APIKey = strings.TrimSpace(p.APIKey)
 	hint := p.APIKeyHint
 	if len(p.APIKey) >= 4 {
 		hint = "…" + p.APIKey[len(p.APIKey)-4:]

server/llm_test_fallback_test.go

@@ -0,0 +1,94 @@
+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"
+)
+
+// The llm/test stored-key fallback must send the same header bytes the runtime
+// sends: form keys are trimmed inside ConfigFrom, the DB fallback used to be
+// assigned raw, so a key saved with surrounding whitespace ran fine in
+// conversations but 401'd the connection test.
+func TestLLMTestStoredKeyFallbackIsTrimmed(t *testing.T) {
+	m, err := NewManager(t.TempDir(), "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { m.Close() })
+	if m.pg == nil {
+		t.Skip("no postgres for this run")
+	}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	s := New(ctx, m, t.TempDir(), t.TempDir(), t.TempDir())
+
+	var gotKey atomic.Value
+	fake := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		gotKey.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}}`))
+	}))
+	t.Cleanup(fake.Close)
+
+	pid, err := m.pg.SaveProfile(&db.LLMProfile{
+		Name: "ws-key", Format: "anthropic", Model: "claude-x",
+		BaseURL: fake.URL, APIKey: "k1 \t",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	t.Cleanup(func() { _, _ = m.pg.Exec("DELETE FROM llm_profiles WHERE id=$1", pid) })
+
+	// The handler trims the fallback key: the profile was saved through
+	// SaveProfile (which now trims), so also prove the storage trim itself.
+	stored, err := m.pg.ProfileByID(pid)
+	if err != nil || stored == nil {
+		t.Fatalf("ProfileByID: %v %v", stored, err)
+	}
+	if stored.APIKey != "k1" {
+		t.Fatalf("SaveProfile must trim surrounding whitespace, got %q", stored.APIKey)
+	}
+
+	body, _ := json.Marshal(map[string]any{
+		"provider": "anthropic", "model": "claude-x", "base_url": fake.URL,
+		"api_key": "", "profile_id": pid, "streaming": false,
+	})
+	req := httptest.NewRequest(http.MethodPost, "/api/llm/test", bytes.NewReader(body))
+	req.Header.Set("Content-Type", "application/json")
+	rec := httptest.NewRecorder()
+	s.testLLM(rec, req)
+
+	if k, _ := gotKey.Load().(string); k != "k1" {
+		t.Fatalf("test request must send the trimmed stored key, header was %q", k)
+	}
+	// The reply body only proves the call completed; the header is the assertion.
+	if rec.Code != http.StatusOK {
+		t.Fatalf("handler status %d", rec.Code)
+	}
+	var resp struct {
+		OK    bool   `json:"ok"`
+		Error string `json:"error"`
+	}
+	_ = json.Unmarshal(rec.Body.Bytes(), &resp)
+	if !resp.OK {
+		t.Logf("note: provider round-trip returned ok=false (%s); header assertion stands", resp.Error)
+	}
+}
+
+// Sanity: the runtime path and the test path must agree byte-for-byte on the
+// key they would send for the same stored profile.
+func TestLLMTestFallbackMatchesRuntimeNormalization(t *testing.T) {
+	const raw = " k1 \n"
+	runtime := agent.ConfigFrom("anthropic", "m", "https://x.example", raw, "")
+	if runtime.APIKey != "k1" {
+		t.Fatalf("ConfigFrom trims: %q", runtime.APIKey)
+	}
+}

server/server.go

@@ -1404,11 +1404,14 @@ func (s *Server) testLLM(w http.ResponseWriter, r *http.Request) {
 	// Same for the session header name: when the form omits it, the stored profile's value is the fallback.
 	if req.ProfileID != nil && (cfg.APIKey == "" || cfg.SessionHeaderKey == "") {
 		if p, err := s.m.pg.ProfileByID(*req.ProfileID); err == nil && p != nil {
+			// 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.
 			if cfg.APIKey == "" {
-				cfg.APIKey = p.APIKey
+				cfg.APIKey = strings.TrimSpace(p.APIKey)
 			}
 			if cfg.SessionHeaderKey == "" {
-				cfg.SessionHeaderKey = p.SessionHeaderKey
+				cfg.SessionHeaderKey = strings.TrimSpace(p.SessionHeaderKey)
 			}
 		}
 	}