mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-15 12:40:55 +08:00
BaaRai
2853f085e5
The connection test resolves its API key as: form input, else the stored profile's key, else the global key. Form keys pass through ConfigFrom, which trims surrounding whitespace; the stored-key fallback assigned p.APIKey raw. A key pasted with a trailing space or newline therefore behaved asymmetrically: conversations worked (the runtime path loads the profile through ConfigFrom and trims) while the connection test sent the malformed header and failed auth with 401. SaveProfile had the same gap one layer down, storing the raw pasted value. Changes: - the llm/test fallback now trims the stored key and session header name, matching ConfigFrom's normalization so the test sends byte-identical headers to what the runtime sends for the same profile - SaveProfile trims the key before storing (and before deriving the hint), so surrounding whitespace never reaches the database - tests: an integration test saves a profile with a whitespace-padded key, runs the real handler against a fake anthropic endpoint, and asserts the endpoint received the trimmed key; a second test pins the ConfigFrom trim the fallback must match
95 lines
3.1 KiB
Go
95 lines
3.1 KiB
Go
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)
|
|
}
|
|
}
|