orchard

  1package main
  2
  3import (
  4	"encoding/json"
  5	"fmt"
  6	"net/http"
  7	"net/http/httptest"
  8	"path/filepath"
  9	"strings"
 10	"testing"
 11	"time"
 12)
 13
 14func testStore(t *testing.T) *Store {
 15	t.Helper()
 16	s, err := OpenStore(filepath.Join(t.TempDir(), "llm.db"))
 17	if err != nil {
 18		t.Fatal(err)
 19	}
 20	t.Cleanup(func() { _ = s.Close() })
 21	return s
 22}
 23
 24func TestAKeyIsUsableOnceAndNeverRecoverable(t *testing.T) {
 25	s := testStore(t)
 26	secret, k, err := s.NewKey("chat")
 27	if err != nil {
 28		t.Fatal(err)
 29	}
 30	if !strings.HasPrefix(secret, keyPrefix) {
 31		t.Errorf("secret = %q, want the %s prefix", secret, keyPrefix)
 32	}
 33	got, ok := s.Authenticate(secret)
 34	if !ok {
 35		t.Fatal("a fresh key did not authenticate")
 36	}
 37	if got.Name != "chat" || got.ID != k.ID {
 38		t.Errorf("resolved to %+v", got)
 39	}
 40	// The plaintext must not be anywhere the listing can reach.
 41	keys, err := s.Keys()
 42	if err != nil {
 43		t.Fatal(err)
 44	}
 45	for _, listed := range keys {
 46		if strings.Contains(secret, listed.Prefix) && len(listed.Prefix) >= len(secret) {
 47			t.Error("the whole secret is in the listing")
 48		}
 49	}
 50}
 51
 52func TestAWrongKeyIsRefused(t *testing.T) {
 53	s := testStore(t)
 54	if _, _, err := s.NewKey("chat"); err != nil {
 55		t.Fatal(err)
 56	}
 57	for _, bad := range []string{"", "nonsense", keyPrefix + "wrong", "Bearer something"} {
 58		if _, ok := s.Authenticate(bad); ok {
 59			t.Errorf("%q authenticated", bad)
 60		}
 61	}
 62}
 63
 64// Revoking has to take effect on the next request, which is the whole reason
 65// the key is checked every time rather than cached.
 66func TestRevokingStopsAKeyImmediately(t *testing.T) {
 67	s := testStore(t)
 68	secret, k, _ := s.NewKey("search")
 69	if _, ok := s.Authenticate(secret); !ok {
 70		t.Fatal("not usable before revoking")
 71	}
 72	if err := s.Revoke(k.ID); err != nil {
 73		t.Fatal(err)
 74	}
 75	if _, ok := s.Authenticate(secret); ok {
 76		t.Error("a revoked key still authenticates")
 77	}
 78}
 79
 80func TestAKeyNeedsAName(t *testing.T) {
 81	s := testStore(t)
 82	if _, _, err := s.NewKey("   "); err == nil {
 83		t.Error("an unnamed key was accepted")
 84	}
 85}
 86
 87func TestCallsAreLoggedAndRolledUp(t *testing.T) {
 88	s := testStore(t)
 89	_, k, _ := s.NewKey("chat")
 90	s.LogCall(Call{KeyID: k.ID, Caller: "chat", Model: "local",
 91		Messages: `[{"role":"user","content":"hello"}]`, Completion: "hi",
 92		PromptTok: 10, OutputTok: 4, DecodeTPS: 60, Status: 200})
 93	s.LogCall(Call{KeyID: k.ID, Caller: "chat", Status: 500, Err: "upstream refused"})
 94
 95	calls, err := s.Calls("chat", 10)
 96	if err != nil {
 97		t.Fatal(err)
 98	}
 99	if len(calls) != 2 {
100		t.Fatalf("got %d calls", len(calls))
101	}
102	// The prompt and the completion are the point of the log, so assert they
103	// survived rather than only that a row exists.
104	if !strings.Contains(calls[1].Messages, "hello") || calls[1].Completion != "hi" {
105		t.Errorf("the text was not kept: %+v", calls[1])
106	}
107
108	usage, err := s.Usage(time.Now().Add(-time.Hour))
109	if err != nil {
110		t.Fatal(err)
111	}
112	if len(usage) != 1 {
113		t.Fatalf("got %d callers", len(usage))
114	}
115	u := usage[0]
116	if u.Calls != 2 || u.PromptTok != 10 || u.OutputTok != 4 || u.Errors != 1 {
117		t.Errorf("usage = %+v", u)
118	}
119}
120
121func TestPruneDropsOldCallsOnly(t *testing.T) {
122	s := testStore(t)
123	s.LogCall(Call{Caller: "chat", Completion: "recent"})
124	if _, err := s.db.Exec(`INSERT INTO calls(caller, completion, at) VALUES(?,?,?)`,
125		"chat", "ancient", time.Now().Add(-200*24*time.Hour).Unix()); err != nil {
126		t.Fatal(err)
127	}
128	n, err := s.Prune(90 * 24 * time.Hour)
129	if err != nil {
130		t.Fatal(err)
131	}
132	if n != 1 {
133		t.Errorf("pruned %d rows, want 1", n)
134	}
135	calls, _ := s.Calls("", 10)
136	if len(calls) != 1 || calls[0].Completion != "recent" {
137		t.Errorf("wrong row survived: %+v", calls)
138	}
139}
140
141// The gateway is the security boundary, so an unkeyed request must not reach
142// upstream at all rather than being refused after it.
143func TestAnUnkeyedRequestNeverReachesUpstream(t *testing.T) {
144	reached := false
145	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
146		reached = true
147	}))
148	defer up.Close()
149
150	s := &site{store: testStore(t), upstream: up.URL, client: up.Client()}
151	rec := httptest.NewRecorder()
152	req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"local"}`))
153	s.requireKey(s.completions)(rec, req)
154
155	if rec.Code != http.StatusUnauthorized {
156		t.Errorf("status = %d, want 401", rec.Code)
157	}
158	if reached {
159		t.Error("an unkeyed request was forwarded upstream")
160	}
161}
162
163func TestAKeyedCallIsForwardedAndWrittenDown(t *testing.T) {
164	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
165		w.Header().Set("Content-Type", "application/json")
166		fmt.Fprint(w, `{"choices":[{"message":{"content":"the answer"}}],
167			"usage":{"prompt_tokens":31,"completion_tokens":7},
168			"timings":{"predicted_per_second":58.5}}`)
169	}))
170	defer up.Close()
171
172	store := testStore(t)
173	secret, _, _ := store.NewKey("chat")
174	s := &site{store: store, upstream: up.URL, client: up.Client()}
175
176	rec := httptest.NewRecorder()
177	req := httptest.NewRequest("POST", "/v1/chat/completions",
178		strings.NewReader(`{"model":"local","messages":[{"role":"user","content":"ask"}]}`))
179	req.Header.Set("Authorization", "Bearer "+secret)
180	s.requireKey(s.completions)(rec, req)
181
182	if rec.Code != 200 {
183		t.Fatalf("status = %d", rec.Code)
184	}
185	if !strings.Contains(rec.Body.String(), "the answer") {
186		t.Errorf("the caller did not get the upstream body: %s", rec.Body.String())
187	}
188
189	calls, _ := store.Calls("chat", 5)
190	if len(calls) != 1 {
191		t.Fatalf("logged %d calls", len(calls))
192	}
193	c := calls[0]
194	if c.Completion != "the answer" {
195		t.Errorf("completion = %q", c.Completion)
196	}
197	if c.PromptTok != 31 || c.OutputTok != 7 {
198		t.Errorf("tokens = %d/%d, want 31/7", c.PromptTok, c.OutputTok)
199	}
200	if !strings.Contains(c.Messages, "ask") {
201		t.Errorf("the prompt was not kept: %q", c.Messages)
202	}
203	if c.Status != 200 {
204		t.Errorf("status = %d", c.Status)
205	}
206}
207
208// Incognito is the one case where a call runs and leaves nothing behind, so the
209// test is that the answer still arrives and the log is still empty.
210func TestAnIncognitoCallIsForwardedAndNotWrittenDown(t *testing.T) {
211	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
212		w.Header().Set("Content-Type", "application/json")
213		fmt.Fprint(w, `{"choices":[{"message":{"content":"the answer"}}]}`)
214	}))
215	defer up.Close()
216
217	store := testStore(t)
218	secret, _, _ := store.NewKey("chat")
219	s := &site{store: store, upstream: up.URL, client: up.Client()}
220
221	rec := httptest.NewRecorder()
222	req := httptest.NewRequest("POST", "/v1/chat/completions",
223		strings.NewReader(`{"model":"local","messages":[{"role":"user","content":"a secret"}]}`))
224	req.Header.Set("Authorization", "Bearer "+secret)
225	req.Header.Set(incognitoHeader, "1")
226	s.requireKey(s.completions)(rec, req)
227
228	if rec.Code != 200 {
229		t.Fatalf("status = %d", rec.Code)
230	}
231	if !strings.Contains(rec.Body.String(), "the answer") {
232		t.Errorf("the caller did not get the upstream body: %s", rec.Body.String())
233	}
234	calls, _ := store.Calls("", 5)
235	if len(calls) != 0 {
236		t.Fatalf("an incognito call was logged: %+v", calls)
237	}
238}
239
240// A streamed answer is the one worth recording and the one easiest to lose,
241// since the text only exists as deltas passing through.
242func TestAStreamedCallIsReassembledForTheLog(t *testing.T) {
243	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
244		w.Header().Set("Content-Type", "text/event-stream")
245		for _, chunk := range []string{
246			`{"choices":[{"delta":{"content":"one "}}]}`,
247			`{"choices":[{"delta":{"content":"two"}}]}`,
248			`{"choices":[],"usage":{"prompt_tokens":9,"completion_tokens":2},"timings":{"predicted_per_second":61.0}}`,
249		} {
250			fmt.Fprintf(w, "data: %s\n\n", chunk)
251		}
252		fmt.Fprint(w, "data: [DONE]\n\n")
253	}))
254	defer up.Close()
255
256	store := testStore(t)
257	secret, _, _ := store.NewKey("search")
258	s := &site{store: store, upstream: up.URL, client: up.Client()}
259
260	rec := httptest.NewRecorder()
261	req := httptest.NewRequest("POST", "/v1/chat/completions",
262		strings.NewReader(`{"model":"local","stream":true,"messages":[{"role":"user","content":"go"}]}`))
263	req.Header.Set("Authorization", "Bearer "+secret)
264	s.requireKey(s.completions)(rec, req)
265
266	// The caller gets the events unchanged.
267	if !strings.Contains(rec.Body.String(), "data: [DONE]") {
268		t.Errorf("the stream did not pass through:\n%s", rec.Body.String())
269	}
270	calls, _ := store.Calls("search", 5)
271	if len(calls) != 1 {
272		t.Fatalf("logged %d calls", len(calls))
273	}
274	if calls[0].Completion != "one two" {
275		t.Errorf("completion = %q, want %q", calls[0].Completion, "one two")
276	}
277	if calls[0].OutputTok != 2 || calls[0].DecodeTPS != 61.0 {
278		t.Errorf("stats lost: %d tokens, %.1f tok/s", calls[0].OutputTok, calls[0].DecodeTPS)
279	}
280}
281
282// A caller whose key is refused upstream still gets a row, since a log with the
283// failures missing is the one that cannot explain an outage.
284func TestAnUpstreamFailureIsStillLogged(t *testing.T) {
285	up := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
286		http.Error(w, "no model loaded", http.StatusServiceUnavailable)
287	}))
288	defer up.Close()
289
290	store := testStore(t)
291	secret, _, _ := store.NewKey("chat")
292	s := &site{store: store, upstream: up.URL, client: up.Client()}
293
294	rec := httptest.NewRecorder()
295	req := httptest.NewRequest("POST", "/v1/chat/completions", strings.NewReader(`{"model":"local"}`))
296	req.Header.Set("Authorization", "Bearer "+secret)
297	s.requireKey(s.completions)(rec, req)
298
299	calls, _ := store.Calls("chat", 5)
300	if len(calls) != 1 {
301		t.Fatalf("logged %d calls", len(calls))
302	}
303	if calls[0].Status != http.StatusServiceUnavailable {
304		t.Errorf("status = %d, want 503", calls[0].Status)
305	}
306}
307
308func TestTheApiKeyHeaderIsAcceptedToo(t *testing.T) {
309	store := testStore(t)
310	secret, _, _ := store.NewKey("aiagent")
311	s := &site{store: store}
312	r := httptest.NewRequest("POST", "/v1/chat/completions", nil)
313	r.Header.Set("X-Api-Key", secret)
314	if _, ok := s.authenticate(r); !ok {
315		t.Error("a key sent as X-Api-Key was refused")
316	}
317}
318
319func TestShortRendersMessagesReadably(t *testing.T) {
320	msgs, _ := json.Marshal([]map[string]string{
321		{"role": "system", "content": "be brief"},
322		{"role": "user", "content": "what is the time"},
323	})
324	got := short(string(msgs), 200)
325	if !strings.Contains(got, "system: be brief") || !strings.Contains(got, "user: what is the time") {
326		t.Errorf("got %q", got)
327	}
328	if len(short(string(msgs), 10)) > 13 {
329		t.Errorf("not truncated: %q", short(string(msgs), 10))
330	}
331}