diff --git a/internal/router/scheduler/fifo_test.go b/internal/router/scheduler/fifo_test.go index 6d177bb2..b98de055 100644 --- a/internal/router/scheduler/fifo_test.go +++ b/internal/router/scheduler/fifo_test.go @@ -1,6 +1,7 @@ package scheduler import ( + "context" "errors" "io" "net/http" @@ -54,8 +55,9 @@ type stopRec struct { // fakeEffects is an in-memory scheduler.Effects. Tests program process states // and GrantServe outcomes, then assert on the recorded calls. type fakeEffects struct { - states map[string]process.ProcessState // model -> state; missing => not handled - serveResult map[string]bool // GrantServe return per model (default true) + states map[string]process.ProcessState // model -> state; missing => not handled + serveResult map[string]bool // GrantServe return per model (default true) + lastServeReq HandlerReq starts []startRec grants []grantRec @@ -98,6 +100,7 @@ func (f *fakeEffects) GrantServe(req HandlerReq, modelID string) bool { if v, set := f.serveResult[modelID]; set { ok = v } + f.lastServeReq = req f.grants = append(f.grants, grantRec{model: modelID, serve: ok}) return ok } @@ -169,6 +172,99 @@ func TestFIFO_FastPath(t *testing.T) { } } +func TestFIFO_GrantSetsPriorityMetadata(t *testing.T) { + eff := newFakeEffects() + eff.states["a"] = process.StateReady + cfg := config.FifoConfig{Priority: map[string]int{"a": 7}} + s := NewFIFO("test", logmon.NewWriter(io.Discard), &stubPlanner{}, cfg, nil, eff) + + ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "a", Metadata: make(map[string]string)}) + s.OnRequest(HandlerReq{Model: "a", Ctx: ctx}) + + if got := eff.served("a"); got != 1 { + t.Fatalf("served(a)=%d want 1", got) + } + data, ok := shared.ReadContext(eff.lastServeReq.Ctx) + if !ok { + t.Fatal("context data missing from granted request") + } + if data.Metadata["fifo_priority"] != "7" { + t.Errorf("fifo_priority = %q, want 7", data.Metadata["fifo_priority"]) + } +} + +func TestFIFO_GrantSetsPriorityMetadata_DefaultZero(t *testing.T) { + // A model that is not listed in the Priority map should get fifo_priority="0". + eff := newFakeEffects() + eff.states["unlisted"] = process.StateReady + cfg := config.FifoConfig{Priority: map[string]int{"other": 5}} // "unlisted" absent + s := NewFIFO("test", logmon.NewWriter(io.Discard), &stubPlanner{}, cfg, nil, eff) + + ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "unlisted", Metadata: make(map[string]string)}) + s.OnRequest(HandlerReq{Model: "unlisted", Ctx: ctx}) + + if got := eff.served("unlisted"); got != 1 { + t.Fatalf("served(unlisted)=%d want 1", got) + } + data, ok := shared.ReadContext(eff.lastServeReq.Ctx) + if !ok { + t.Fatal("context data missing from granted request") + } + if data.Metadata["fifo_priority"] != "0" { + t.Errorf("fifo_priority = %q, want %q", data.Metadata["fifo_priority"], "0") + } +} + +func TestFIFO_GrantSetsPriorityMetadata_NoMetadataMap(t *testing.T) { + // When the request context has no Metadata map, grantHandler must not crash. + // It should log a debug message and still grant the request. + eff := newFakeEffects() + eff.states["a"] = process.StateReady + cfg := config.FifoConfig{Priority: map[string]int{"a": 3}} + s := NewFIFO("test", logmon.NewWriter(io.Discard), &stubPlanner{}, cfg, nil, eff) + + // No Metadata map in the context data — SetReqData will return an error. + ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "a"}) + s.OnRequest(HandlerReq{Model: "a", Ctx: ctx}) + + // The grant must still succeed despite the missing metadata map. + if got := eff.served("a"); got != 1 { + t.Fatalf("served(a)=%d want 1 (metadata error must not prevent grant)", got) + } +} + +func TestFIFO_GrantSetsPriorityMetadata_AfterSwapCompletion(t *testing.T) { + // Priority metadata must be set for waiters granted via OnSwapDone, not just + // requests that hit the fast path. + eff := newFakeEffects() + eff.states["a"] = process.StateStopped // forces a swap + cfg := config.FifoConfig{Priority: map[string]int{"a": 9}} + s := NewFIFO("test", logmon.NewWriter(io.Discard), &stubPlanner{}, cfg, nil, eff) + + ctx := shared.SetContext(context.Background(), shared.ReqContextData{ModelID: "a", Metadata: make(map[string]string)}) + s.OnRequest(HandlerReq{Model: "a", Ctx: ctx}) + + // Swap is in flight; no grant yet. + if got := eff.served("a"); got != 0 { + t.Fatalf("served(a)=%d want 0 before swap done", got) + } + + // Complete the swap. + eff.states["a"] = process.StateReady + s.OnSwapDone(SwapDone{ModelID: "a"}) + + if got := eff.served("a"); got != 1 { + t.Fatalf("served(a)=%d want 1 after swap done", got) + } + data, ok := shared.ReadContext(eff.lastServeReq.Ctx) + if !ok { + t.Fatal("context data missing from granted request after swap") + } + if data.Metadata["fifo_priority"] != "9" { + t.Errorf("fifo_priority = %q, want %q", data.Metadata["fifo_priority"], "9") + } +} + func TestFIFO_ModelNotFound(t *testing.T) { eff := newFakeEffects() // no states => model unknown s := newFIFO(&stubPlanner{}, eff) diff --git a/internal/server/metrics_test.go b/internal/server/metrics_test.go index 04412cd5..e019dd6a 100644 --- a/internal/server/metrics_test.go +++ b/internal/server/metrics_test.go @@ -1,9 +1,13 @@ package server import ( + "net/http" + "net/http/httptest" + "strings" "testing" "time" + "github.com/mostlygeek/llama-swap/internal/shared" "github.com/tidwall/gjson" ) @@ -56,6 +60,109 @@ func TestServer_ProcessStreamingResponse_NoData(t *testing.T) { } } +func TestMetricsMonitor_RecordMetadata(t *testing.T) { + mm := newMetricsMonitor(nil, 10, 0) + r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"usage":{}}`)) + r = r.WithContext(shared.SetContext(r.Context(), shared.ReqContextData{ + ModelID: "m", + Metadata: map[string]string{"client": "web", "trace": "abc"}, + })) + + w := httptest.NewRecorder() + copier := newBodyCopier(w) + copier.WriteHeader(http.StatusOK) + copier.Write([]byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2}}`)) + + mm.record("m", r, copier, 0, nil, nil) + + entries := mm.getMetrics() + if len(entries) != 1 { + t.Fatalf("want 1 entry, got %d", len(entries)) + } + if entries[0].Metadata["client"] != "web" { + t.Errorf("client = %q, want web", entries[0].Metadata["client"]) + } + if entries[0].Metadata["trace"] != "abc" { + t.Errorf("trace = %q, want abc", entries[0].Metadata["trace"]) + } +} + +func TestMetricsMonitor_RecordMetadata_EmptyMap(t *testing.T) { + // An empty Metadata map in context must NOT set tm.Metadata (omitempty semantics). + mm := newMetricsMonitor(nil, 10, 0) + r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{}`)) + r = r.WithContext(shared.SetContext(r.Context(), shared.ReqContextData{ + ModelID: "m", + Metadata: map[string]string{}, // empty, not nil + })) + + w := httptest.NewRecorder() + copier := newBodyCopier(w) + copier.WriteHeader(http.StatusOK) + copier.Write([]byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2}}`)) + + mm.record("m", r, copier, 0, nil, nil) + + entries := mm.getMetrics() + if len(entries) != 1 { + t.Fatalf("want 1 entry, got %d", len(entries)) + } + if entries[0].Metadata != nil { + t.Errorf("Metadata should be nil for empty context metadata, got %v", entries[0].Metadata) + } +} + +func TestMetricsMonitor_RecordMetadata_NoContextData(t *testing.T) { + // A request with no ReqContextData in context should produce nil Metadata. + mm := newMetricsMonitor(nil, 10, 0) + r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{}`)) + // No shared.SetContext call — no ReqContextData in context. + + w := httptest.NewRecorder() + copier := newBodyCopier(w) + copier.WriteHeader(http.StatusOK) + copier.Write([]byte(`{"usage":{"prompt_tokens":3,"completion_tokens":4}}`)) + + mm.record("m", r, copier, 0, nil, nil) + + entries := mm.getMetrics() + if len(entries) != 1 { + t.Fatalf("want 1 entry, got %d", len(entries)) + } + if entries[0].Metadata != nil { + t.Errorf("Metadata should be nil when no context data, got %v", entries[0].Metadata) + } +} + +func TestMetricsMonitor_RecordMetadata_DeepCopy(t *testing.T) { + // Mutating the original context metadata after record() must not affect the stored entry. + mm := newMetricsMonitor(nil, 10, 0) + original := map[string]string{"key": "before"} + r := httptest.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{}`)) + r = r.WithContext(shared.SetContext(r.Context(), shared.ReqContextData{ + ModelID: "m", + Metadata: original, + })) + + w := httptest.NewRecorder() + copier := newBodyCopier(w) + copier.WriteHeader(http.StatusOK) + copier.Write([]byte(`{"usage":{"prompt_tokens":1,"completion_tokens":2}}`)) + + mm.record("m", r, copier, 0, nil, nil) + + // Mutate the original map after record. + original["key"] = "after" + + entries := mm.getMetrics() + if len(entries) != 1 { + t.Fatalf("want 1 entry, got %d", len(entries)) + } + if entries[0].Metadata["key"] != "before" { + t.Errorf("Metadata[key] = %q, want %q (deep copy expected)", entries[0].Metadata["key"], "before") + } +} + func TestServer_ParseMetrics_Infill(t *testing.T) { // /infill responses are arrays; timings live in the last element. body := `[{"content":"a"},{"content":"b","timings":{"prompt_n":5,"predicted_n":9,"prompt_ms":10,"predicted_ms":20}}]` diff --git a/internal/shared/http_test.go b/internal/shared/http_test.go index 8bbdd69c..c5fe20cb 100644 --- a/internal/shared/http_test.go +++ b/internal/shared/http_test.go @@ -387,6 +387,105 @@ func TestExtractContext_ApiKey(t *testing.T) { } } +func TestSetReqData(t *testing.T) { + ctx := SetContext(context.Background(), ReqContextData{Model: "llama3", ModelID: "llama3", Metadata: make(map[string]string)}) + + if err := SetReqData(ctx, "client", "web"); err != nil { + t.Fatalf("SetReqData: %v", err) + } + if err := SetReqData(ctx, "trace", "abc123"); err != nil { + t.Fatalf("SetReqData: %v", err) + } + + data, ok := ReadContext(ctx) + if !ok { + t.Fatal("context data missing") + } + if data.Metadata["client"] != "web" { + t.Errorf("client = %q, want %q", data.Metadata["client"], "web") + } + if data.Metadata["trace"] != "abc123" { + t.Errorf("trace = %q, want %q", data.Metadata["trace"], "abc123") + } +} + +func TestSetReqData_Errors(t *testing.T) { + if err := SetReqData(context.Background(), "k", "v"); err == nil { + t.Error("expected error when no request context data exists") + } + ctx := SetContext(context.Background(), ReqContextData{Model: "llama3", ModelID: "llama3"}) + if err := SetReqData(ctx, "k", "v"); err == nil { + t.Error("expected error when metadata map is missing") + } +} + +func TestSetReqData_NilContext(t *testing.T) { + // nil context must return an error without panicking. + err := SetReqData(nil, "k", "v") + if err == nil { + t.Error("expected error for nil context, got nil") + } +} + +func TestSetReqData_OverwritesExistingKey(t *testing.T) { + ctx := SetContext(context.Background(), ReqContextData{ + Model: "m", + Metadata: map[string]string{"key": "old"}, + }) + if err := SetReqData(ctx, "key", "new"); err != nil { + t.Fatalf("SetReqData: %v", err) + } + data, _ := ReadContext(ctx) + if data.Metadata["key"] != "new" { + t.Errorf("key = %q, want %q", data.Metadata["key"], "new") + } +} + +func TestExtractContext_MetadataInitialized_GET(t *testing.T) { + r, _ := http.NewRequest(http.MethodGet, "/?model=llama3", nil) + got, err := extractContext(r) + if err != nil { + t.Fatalf("extractContext: %v", err) + } + if got.Metadata == nil { + t.Error("Metadata should be initialized (not nil) for GET requests") + } +} + +func TestExtractContext_MetadataInitialized_JSON(t *testing.T) { + r, _ := http.NewRequest(http.MethodPost, "/v1/chat/completions", strings.NewReader(`{"model":"llama3"}`)) + r.Header.Set("Content-Type", "application/json") + got, err := extractContext(r) + if err != nil { + t.Fatalf("extractContext: %v", err) + } + if got.Metadata == nil { + t.Error("Metadata should be initialized (not nil) for JSON POST requests") + } +} + +func TestExtractContext_MetadataInitialized_Form(t *testing.T) { + r, _ := http.NewRequest(http.MethodPost, "/v1/audio/transcriptions", strings.NewReader("model=whisper-1")) + r.Header.Set("Content-Type", "application/x-www-form-urlencoded") + got, err := extractContext(r) + if err != nil { + t.Fatalf("extractContext: %v", err) + } + if got.Metadata == nil { + t.Error("Metadata should be initialized (not nil) for form POST requests") + } +} + +func TestExtractContext_MetadataIsWritable(t *testing.T) { + // Verify the initialized map is writable — i.e. SetReqData can use it. + r, _ := http.NewRequest(http.MethodGet, "/?model=llama3", nil) + got, _ := extractContext(r) + ctx := SetContext(context.Background(), got) + if err := SetReqData(ctx, "x", "y"); err != nil { + t.Fatalf("SetReqData on extractContext Metadata: %v", err) + } +} + func TestServer_ExtractAPIKey(t *testing.T) { basicHeader := func(user, pass string) string { return "Basic " + base64.StdEncoding.EncodeToString([]byte(user+":"+pass)) diff --git a/ui-svelte/src/lib/types.test.ts b/ui-svelte/src/lib/types.test.ts new file mode 100644 index 00000000..a0900ea1 --- /dev/null +++ b/ui-svelte/src/lib/types.test.ts @@ -0,0 +1,217 @@ +import { describe, it, expect } from "vitest"; +import type { ActivityLogEntry, TokenMetrics } from "./types"; + +// Baseline token metrics used across tests. +const baseTokens: TokenMetrics = { + cache_tokens: 0, + input_tokens: 10, + output_tokens: 5, + prompt_per_second: 100, + tokens_per_second: 50, +}; + +function makeEntry(overrides: Partial = {}): ActivityLogEntry { + return { + id: 0, + timestamp: "2024-01-01T00:00:00Z", + model: "llama3", + req_path: "/v1/chat/completions", + resp_content_type: "application/json", + resp_status_code: 200, + tokens: baseTokens, + duration_ms: 100, + has_capture: false, + ...overrides, + }; +} + +describe("ActivityLogEntry", () => { + describe("metadata field", () => { + it("accepts an entry without metadata (undefined)", () => { + const entry = makeEntry(); + expect(entry.metadata).toBeUndefined(); + }); + + it("accepts an entry with metadata populated", () => { + const entry = makeEntry({ metadata: { client: "web", trace: "abc123" } }); + expect(entry.metadata).toEqual({ client: "web", trace: "abc123" }); + }); + + it("accepts an empty metadata object", () => { + const entry = makeEntry({ metadata: {} }); + expect(entry.metadata).toEqual({}); + }); + + it("allows reading a key from metadata when present", () => { + const entry = makeEntry({ metadata: { fifo_priority: "7" } }); + expect(entry.metadata?.["fifo_priority"]).toBe("7"); + }); + + it("returns undefined when accessing a missing key via optional chaining", () => { + const entry = makeEntry(); + expect(entry.metadata?.["fifo_priority"]).toBeUndefined(); + }); + + it("returns undefined for a missing key when metadata is an empty object", () => { + const entry = makeEntry({ metadata: {} }); + expect(entry.metadata?.["missing"]).toBeUndefined(); + }); + + it("supports metadata with multiple entries", () => { + const meta: Record = { a: "1", b: "2", c: "3" }; + const entry = makeEntry({ metadata: meta }); + expect(Object.keys(entry.metadata!)).toHaveLength(3); + expect(entry.metadata!["b"]).toBe("2"); + }); + }); + + describe("ActivityLogEntry structure", () => { + it("round-trips through JSON with metadata", () => { + const entry = makeEntry({ metadata: { client: "web" } }); + const json = JSON.stringify(entry); + const parsed: ActivityLogEntry = JSON.parse(json); + expect(parsed.metadata).toEqual({ client: "web" }); + }); + + it("round-trips through JSON without metadata (field omitted)", () => { + const entry = makeEntry(); + const json = JSON.stringify(entry); + const parsed: ActivityLogEntry = JSON.parse(json); + // When metadata is undefined it is dropped by JSON.stringify. + expect(parsed.metadata).toBeUndefined(); + }); + + it("round-trips through JSON with null metadata preserved as-is", () => { + // Explicit null from a server that omits the field still satisfies + // the optional type when accessed with optional chaining. + const raw = { ...makeEntry(), metadata: null }; + const json = JSON.stringify(raw); + const parsed = JSON.parse(json) as ActivityLogEntry; + // null is falsy — optional chaining returns undefined for null too. + expect(parsed.metadata ?? undefined).toBeUndefined(); + }); + }); +}); + +// META_PREFIX helpers — mirror the logic in Activity.svelte so we can test it +// without importing the Svelte component. +const META_PREFIX = "meta:"; + +function isMetaKey(key: string): boolean { + return key.startsWith(META_PREFIX); +} + +function metaKey(name: string): string { + return META_PREFIX + name; +} + +function metaLabel(key: string): string { + return key.slice(META_PREFIX.length); +} + +describe("Activity.svelte META_PREFIX helpers", () => { + describe("isMetaKey", () => { + it("returns true for keys with meta: prefix", () => { + expect(isMetaKey("meta:fifo_priority")).toBe(true); + expect(isMetaKey("meta:client")).toBe(true); + expect(isMetaKey("meta:")).toBe(true); // empty suffix is still prefixed + }); + + it("returns false for standard column keys", () => { + expect(isMetaKey("id")).toBe(false); + expect(isMetaKey("model")).toBe(false); + expect(isMetaKey("capture")).toBe(false); + expect(isMetaKey("duration")).toBe(false); + }); + + it("returns false for partial or incorrect prefixes", () => { + expect(isMetaKey("meta")).toBe(false); + expect(isMetaKey("Meta:key")).toBe(false); // case-sensitive + expect(isMetaKey("")).toBe(false); + }); + }); + + describe("metaKey", () => { + it("prepends META_PREFIX to the name", () => { + expect(metaKey("fifo_priority")).toBe("meta:fifo_priority"); + expect(metaKey("client")).toBe("meta:client"); + }); + + it("handles empty name", () => { + expect(metaKey("")).toBe("meta:"); + }); + + it("is the inverse of metaLabel", () => { + const name = "some_key"; + expect(metaLabel(metaKey(name))).toBe(name); + }); + }); + + describe("metaLabel", () => { + it("strips META_PREFIX and returns the bare name", () => { + expect(metaLabel("meta:fifo_priority")).toBe("fifo_priority"); + expect(metaLabel("meta:client")).toBe("client"); + }); + + it("handles empty suffix", () => { + expect(metaLabel("meta:")).toBe(""); + }); + + it("is the inverse of metaKey", () => { + const key = "meta:trace_id"; + expect(metaKey(metaLabel(key))).toBe(key); + }); + }); + + describe("metadata column derivation", () => { + it("derives unique metadata keys from a list of entries", () => { + const entries: ActivityLogEntry[] = [ + makeEntry({ metadata: { client: "web", trace: "a" } }), + makeEntry({ metadata: { client: "mobile" } }), + makeEntry({ metadata: { fifo_priority: "3" } }), + makeEntry({}), // no metadata + ]; + + const keys = Array.from( + new Set(entries.flatMap((m) => Object.keys(m.metadata || {}))) + ).sort(); + + expect(keys).toEqual(["client", "fifo_priority", "trace"]); + }); + + it("returns empty array when no entries have metadata", () => { + const entries: ActivityLogEntry[] = [makeEntry(), makeEntry()]; + const keys = Array.from( + new Set(entries.flatMap((m) => Object.keys(m.metadata || {}))) + ).sort(); + expect(keys).toHaveLength(0); + }); + + it("maps metadata keys to meta:-prefixed column keys", () => { + const metaKeys = ["client", "fifo_priority"]; + const columnKeys = metaKeys.map(metaKey); + expect(columnKeys).toEqual(["meta:client", "meta:fifo_priority"]); + }); + + it("resolves metadata value for a column key", () => { + const entry = makeEntry({ metadata: { fifo_priority: "7", client: "web" } }); + const key = "meta:fifo_priority"; + const value = entry.metadata?.[metaLabel(key)] ?? "-"; + expect(value).toBe("7"); + }); + + it("falls back to '-' for a column key not present in entry metadata", () => { + const entry = makeEntry({ metadata: { client: "web" } }); + const key = "meta:fifo_priority"; + const value = entry.metadata?.[metaLabel(key)] ?? "-"; + expect(value).toBe("-"); + }); + + it("falls back to '-' when entry has no metadata at all", () => { + const entry = makeEntry(); // metadata is undefined + const key = "meta:anything"; + const value = entry.metadata?.[metaLabel(key)] ?? "-"; + expect(value).toBe("-"); + }); + }); +});