diff --git a/db/config.go b/db/config.go index 43e5262..ad29fc9 100644 --- a/db/config.go +++ b/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:] diff --git a/server/llm_test_fallback_test.go b/server/llm_test_fallback_test.go new file mode 100644 index 0000000..052b76e --- /dev/null +++ b/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) + } +} diff --git a/server/server.go b/server/server.go index 5cbcf0b..a2c2e34 100644 --- a/server/server.go +++ b/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) } } }