mirror of
https://github.com/BaaRai/ARTEX.git
synced 2026-11-15 04:30:55 +08:00
Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2853f085e5 |
@@ -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:]
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
+5
-2
@@ -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)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user