package agi import ( "encoding/json" "io" "net/http" "net/http/httptest" "net/url" "path/filepath" "strings" "testing" "time" "github.com/robertkrimen/otto" "imuslab.com/arozos/mod/agi/static" llm "imuslab.com/arozos/mod/aiservers/llm" database "imuslab.com/arozos/mod/database" user "imuslab.com/arozos/mod/user" ) // dbGateway returns a Gateway backed by a throwaway bolt database so the // config / pricing / metrics persistence paths can be exercised in tests. // Shared by every *_test.go file in this package (llm, cnn, ...). func dbGateway(t *testing.T) *Gateway { t.Helper() dbfile := filepath.Join(t.TempDir(), "test.db") sysdb, err := database.NewDatabase(dbfile, false) if err != nil { t.Fatalf("failed to create test database: %v", err) } t.Cleanup(func() { sysdb.Close() }) uh, err := user.NewUserHandler(sysdb, nil, nil, nil, nil) if err != nil { t.Fatalf("failed to create user handler: %v", err) } g := minimalGateway() g.Option.UserHandler = uh sysdb.NewTable(llmDBTable) return g } // ─── pure helpers ───────────────────────────────────────────────────────────── func TestParseLLMCallOptions(t *testing.T) { if opt := parseLLMCallOptions(""); opt.Model != "" { t.Errorf("empty string should yield zero options") } if opt := parseLLMCallOptions("undefined"); opt.Model != "" { t.Errorf("'undefined' should yield zero options") } if opt := parseLLMCallOptions("null"); opt.Model != "" { t.Errorf("'null' should yield zero options") } opt := parseLLMCallOptions(`{"model":"gpt-4o","system":"be brief","temperature":0.5,"max_tokens":42}`) if opt.Model != "gpt-4o" || opt.System != "be brief" { t.Errorf("unexpected parse: %+v", opt) } if opt.Temperature == nil || *opt.Temperature != 0.5 { t.Errorf("temperature not parsed") } if opt.MaxTokens == nil || *opt.MaxTokens != 42 { t.Errorf("max_tokens not parsed") } } func TestLLMMaskKey(t *testing.T) { cases := map[string]string{ "": "", "abc": "•••", "sk-1234567890": "••••7890", } for in, want := range cases { if got := llmMaskKey(in); got != want { t.Errorf("maskKey(%q) = %q, want %q", in, got, want) } } } func TestLLMExtClassification(t *testing.T) { if !llmIsImageExt(".png") || !llmIsImageExt(".jpeg") { t.Error("expected image extensions to be detected") } if llmIsImageExt(".txt") { t.Error(".txt should not be an image") } if !llmIsTextExt(".md") || !llmIsTextExt(".go") { t.Error("expected text extensions to be detected") } if llmIsTextExt(".png") { t.Error(".png should not be classified as text") } } // ─── persistence ────────────────────────────────────────────────────────────── func TestRecordLLMUsageAccumulatesAndCosts(t *testing.T) { g := dbGateway(t) sysdb := g.Option.UserHandler.GetDatabase() //Pricing: $2.50 / 1M input, $10.00 / 1M output sysdb.Write(llmDBTable, "pricing", map[string]LLMPricing{ "test-model": {InputPrice: 2.5, OutputPrice: 10.0}, }) g.recordLLMUsage("test-model", 1000, 500) g.recordLLMUsage("test-model", 1000, 500) m := g.getLLMMetrics() if m.TotalRequests != 2 { t.Errorf("expected 2 requests, got %d", m.TotalRequests) } if m.TotalPromptTokens != 2000 || m.TotalCompletionTokens != 1000 || m.TotalTokens != 3000 { t.Errorf("unexpected token totals: %+v", m) } //Each call: 1000/1e6*2.5 + 500/1e6*10 = 0.0075 ; two calls => 0.015 if got := m.TotalCost; got < 0.01499 || got > 0.01501 { t.Errorf("expected total cost ~0.015, got %v", got) } rec := m.PerModel["test-model"] if rec == nil || rec.Requests != 2 || rec.TotalTokens != 3000 { t.Errorf("per-model record incorrect: %+v", rec) } } func TestGetLLMConfigDefaultsCurrency(t *testing.T) { g := dbGateway(t) cfg := g.getLLMConfig() if cfg.Currency != "USD" { t.Errorf("expected default currency USD, got %q", cfg.Currency) } } // ─── orchestration (config resolution + metrics recording) ────────────────── // Wire-protocol mechanics (request shape, auth headers, response decoding) // are covered by mod/aiservers/llm's own tests; these only verify that the // AGI-layer orchestrator wires the client and the persisted config/metrics // together correctly. func TestLLMDoRequestFlow(t *testing.T) { var gotModel string var sawUserMessage bool srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { body, _ := io.ReadAll(r.Body) var req struct { Model string `json:"model"` Messages []struct { Role string `json:"role"` } `json:"messages"` } json.Unmarshal(body, &req) gotModel = req.Model for _, msg := range req.Messages { if msg.Role == "user" { sawUserMessage = true } } w.Header().Set("Content-Type", "application/json") io.WriteString(w, `{"model":"test-model", "choices":[{"index":0,"message":{"role":"assistant","content":"Hello from mock"},"finish_reason":"stop"}], "usage":{"prompt_tokens":1000,"completion_tokens":500,"total_tokens":1500}}`) })) defer srv.Close() g := dbGateway(t) sysdb := g.Option.UserHandler.GetDatabase() sysdb.Write(llmDBTable, "config", LLMConfig{ Endpoint: srv.URL, APIKey: "test-key", DefaultModel: "test-model", Currency: "USD", }) sysdb.Write(llmDBTable, "pricing", map[string]LLMPricing{ "test-model": {InputPrice: 2.5, OutputPrice: 10.0}, }) resp, err := g.llmDoRequest("", []llm.Message{{Role: "user", Content: "hi"}}, llmCallOptions{}) if err != nil { t.Fatalf("llmDoRequest returned error: %v", err) } if content := llmExtractContent(resp); content != "Hello from mock" { t.Errorf("unexpected content: %q", content) } if gotModel != "test-model" { t.Errorf("expected default model to be used, got %q", gotModel) } if !sawUserMessage { t.Error("server did not receive a user message") } //Metrics should have been recorded from the usage block m := g.getLLMMetrics() if m.TotalRequests != 1 || m.TotalTokens != 1500 { t.Errorf("metrics not recorded after request: %+v", m) } } func TestLLMDoStreamRequestFlow(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "text/event-stream") io.WriteString(w, "data: {\"model\":\"m\",\"choices\":[{\"delta\":{\"reasoning_content\":\"hmm\"}}]}\n\n") io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"Hi\"}}]}\n\n") io.WriteString(w, "data: {\"choices\":[{\"delta\":{\"content\":\"!\"},\"finish_reason\":\"stop\"}]}\n\n") io.WriteString(w, "data: {\"choices\":[],\"usage\":{\"prompt_tokens\":4,\"completion_tokens\":2,\"total_tokens\":6}}\n\n") io.WriteString(w, "data: [DONE]\n\n") })) defer srv.Close() g := dbGateway(t) g.Option.UserHandler.GetDatabase().Write(llmDBTable, "config", LLMConfig{Endpoint: srv.URL, DefaultModel: "m", APIFormat: "openai", Currency: "USD"}) var content, reasoning strings.Builder resp, err := g.llmDoStreamRequest("", []llm.Message{{Role: "user", Content: "hi"}}, llmCallOptions{}, func(d llm.StreamDelta) { content.WriteString(d.Content) reasoning.WriteString(d.Reasoning) }) if err != nil { t.Fatalf("llmDoStreamRequest error: %v", err) } if content.String() != "Hi!" { t.Errorf("streamed content = %q, want Hi!", content.String()) } if reasoning.String() != "hmm" { t.Errorf("streamed reasoning = %q, want hmm", reasoning.String()) } if llmExtractContent(resp) != "Hi!" { t.Errorf("assembled content wrong: %q", llmExtractContent(resp)) } //Usage from the streamed final chunk must be recorded like a blocking call. m := g.getLLMMetrics() if m.TotalRequests != 1 || m.TotalTokens != 6 { t.Errorf("metrics not recorded after stream: %+v", m) } } func TestLLMDoStreamRequestNoEndpoint(t *testing.T) { g := dbGateway(t) _, err := g.llmDoStreamRequest("m", []llm.Message{{Role: "user", Content: "hi"}}, llmCallOptions{}, nil) if err == nil { t.Error("expected error when endpoint is not configured") } } func TestLLMDoRequestNoEndpoint(t *testing.T) { g := dbGateway(t) _, err := g.llmDoRequest("m", []llm.Message{{Role: "user", Content: "hi"}}, llmCallOptions{}) if err == nil { t.Error("expected error when endpoint is not configured") } } func TestLLMDoRequestAnthropicFlow(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { w.Header().Set("Content-Type", "application/json") io.WriteString(w, `{"model":"claude-x", "content":[{"type":"text","text":"Hi from Claude"}], "usage":{"input_tokens":30,"output_tokens":12}, "stop_reason":"end_turn"}`) })) defer srv.Close() g := dbGateway(t) sysdb := g.Option.UserHandler.GetDatabase() sysdb.Write(llmDBTable, "config", LLMConfig{ Endpoint: srv.URL, APIKey: "anthropic-key", DefaultModel: "claude-x", APIFormat: "anthropic", Currency: "USD", }) //A system message in the unified array must be lifted to the top-level field //(verified directly in mod/aiservers/llm); here we only check the result //that reaches the AGI layer and that usage gets recorded. msgs := []llm.Message{ {Role: "system", Content: "be brief"}, {Role: "user", Content: "hello"}, } resp, err := g.llmDoRequest("", msgs, llmCallOptions{}) if err != nil { t.Fatalf("anthropic request errored: %v", err) } if content := llmExtractContent(resp); content != "Hi from Claude" { t.Errorf("unexpected content: %q", content) } //Usage mapping: input->prompt, output->completion. m := g.getLLMMetrics() if m.TotalPromptTokens != 30 || m.TotalCompletionTokens != 12 || m.TotalTokens != 42 { t.Errorf("usage not mapped/recorded correctly: %+v", m) } } func TestLLMDoRequestRecordsTokensPerSecond(t *testing.T) { srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { time.Sleep(25 * time.Millisecond) //ensure a measurable generation time w.Header().Set("Content-Type", "application/json") io.WriteString(w, `{"model":"m","choices":[{"message":{"role":"assistant","content":"hello world"}}], "usage":{"prompt_tokens":5,"completion_tokens":20,"total_tokens":25}}`) })) defer srv.Close() g := dbGateway(t) g.Option.UserHandler.GetDatabase().Write(llmDBTable, "config", LLMConfig{Endpoint: srv.URL, DefaultModel: "m", APIFormat: "openai"}) resp, err := g.llmDoRequest("", []llm.Message{{Role: "user", Content: "hi"}}, llmCallOptions{}) if err != nil { t.Fatalf("request errored: %v", err) } if resp.Usage.GenerationMs <= 0 { t.Errorf("expected generation_ms > 0, got %d", resp.Usage.GenerationMs) } if resp.Usage.TokensPerSecond <= 0 { t.Errorf("expected tokens_per_second > 0, got %v", resp.Usage.TokensPerSecond) } m := g.getLLMMetrics() if m.TotalGenerationMs <= 0 { t.Errorf("expected total generation ms recorded, got %d", m.TotalGenerationMs) } if m.SpeedSamples != 1 || m.SpeedSum <= 0 { t.Errorf("expected one speed sample recorded, got samples=%d sum=%v", m.SpeedSamples, m.SpeedSum) } if rec := m.PerModel["m"]; rec == nil || rec.GenerationMs <= 0 || rec.SpeedSamples != 1 { t.Errorf("per-model speed sample not recorded: %+v", rec) } } // The average speed must be the mean of per-request speeds, not total tokens // over total time (which is token-weighted and skews toward large requests). func TestLLMAverageSpeedIsMeanOfRequests(t *testing.T) { g := dbGateway(t) //Request A: 10 tokens in 1000ms -> 10 tok/s g.recordLLMUsage("m", 0, 10, 1000) //Request B: 1000 tokens in 10000ms -> 100 tok/s g.recordLLMUsage("m", 0, 1000, 10000) m := g.getLLMMetrics() if m.SpeedSamples != 2 { t.Fatalf("expected 2 speed samples, got %d", m.SpeedSamples) } avg := m.SpeedSum / float64(m.SpeedSamples) //Mean of speeds = (10 + 100) / 2 = 55 (NOT throughput 1010/11 ≈ 91.8). if avg < 54.9 || avg > 55.1 { t.Errorf("expected average speed ~55 tok/s, got %v", avg) } } // ─── quota enforcement ────────────────────────────────────────────────────── func TestLLMQuotaEnforcement(t *testing.T) { g := dbGateway(t) sysdb := g.Option.UserHandler.GetDatabase() sysdb.Write(llmDBTable, "quota", LLMQuota{Enabled: true, MaxTokens: 100, Period: "total"}) //Under the cap -> allowed. if err := g.llmCheckQuota(); err != nil { t.Fatalf("expected no error under quota, got %v", err) } //Consume past the cap. g.recordLLMUsage("m", 80, 40) // 120 tokens > 100 if err := g.llmCheckQuota(); err == nil { t.Error("expected quota error after exceeding token cap") } else if !strings.Contains(err.Error(), "quota") { t.Errorf("expected a quota error, got %v", err) } //Disabling the quota lifts the block. sysdb.Write(llmDBTable, "quota", LLMQuota{Enabled: false, MaxTokens: 100, Period: "total"}) if err := g.llmCheckQuota(); err != nil { t.Errorf("disabled quota should not block, got %v", err) } } func TestLLMDoRequestBlockedByQuota(t *testing.T) { g := dbGateway(t) sysdb := g.Option.UserHandler.GetDatabase() sysdb.Write(llmDBTable, "config", LLMConfig{Endpoint: "http://127.0.0.1:0", DefaultModel: "m", APIFormat: "openai"}) sysdb.Write(llmDBTable, "quota", LLMQuota{Enabled: true, MaxTokens: 10, Period: "total"}) g.recordLLMUsage("m", 20, 0) // exceed _, err := g.llmDoRequest("m", []llm.Message{{Role: "user", Content: "hi"}}, llmCallOptions{}) if err == nil || !strings.Contains(err.Error(), "quota") { t.Errorf("expected request to be blocked by quota, got %v", err) } } func TestLLMWindowExpired(t *testing.T) { now := time.Date(2026, 6, 11, 12, 0, 0, 0, time.UTC) if !llmWindowExpired(0, "daily", now) { t.Error("zero start should be considered expired") } yesterday := now.AddDate(0, 0, -1).Unix() if !llmWindowExpired(yesterday, "daily", now) { t.Error("yesterday should be expired for daily period") } if llmWindowExpired(now.Add(-1*time.Hour).Unix(), "daily", now) { t.Error("same day should not be expired for daily period") } lastMonth := now.AddDate(0, -1, 0).Unix() if !llmWindowExpired(lastMonth, "monthly", now) { t.Error("last month should be expired for monthly period") } if llmWindowExpired(now.AddDate(0, -1, 0).Unix(), "total", now) { t.Error("total period should never expire") } } // ─── config handler masking ───────────────────────────────────────────────── // HandleAIModelConfig keeps its original name/route - only the requirelib // identifier exposed to AGI scripts changed. func TestHandleAIModelConfigMaskingAndKeyRetention(t *testing.T) { g := dbGateway(t) sysdb := g.Option.UserHandler.GetDatabase() sysdb.Write(llmDBTable, "config", LLMConfig{ Endpoint: "https://api.example.com/v1", APIKey: "sk-supersecret9999", DefaultModel: "m", Currency: "USD", }) //GET should mask the key rec := httptest.NewRecorder() g.HandleAIModelConfig(rec, httptest.NewRequest("GET", "/system/aimodel/config", nil)) var got map[string]interface{} json.Unmarshal(rec.Body.Bytes(), &got) if got["hasKey"] != true { t.Errorf("expected hasKey true, got %v", got["hasKey"]) } if hint, _ := got["keyHint"].(string); !strings.HasSuffix(hint, "9999") || strings.Contains(hint, "supersecret") { t.Errorf("key not properly masked: %v", got["keyHint"]) } //POST without apikey should retain the saved key, but update endpoint form := url.Values{} form.Set("endpoint", "https://new.example.com/v1") form.Set("defaultModel", "m2") form.Set("currency", "EUR") req := httptest.NewRequest("POST", "/system/aimodel/config", strings.NewReader(form.Encode())) req.Header.Set("Content-Type", "application/x-www-form-urlencoded") g.HandleAIModelConfig(httptest.NewRecorder(), req) cfg := g.getLLMConfig() if cfg.APIKey != "sk-supersecret9999" { t.Errorf("API key should have been retained, got %q", cfg.APIKey) } if cfg.Endpoint != "https://new.example.com/v1" || cfg.DefaultModel != "m2" || cfg.Currency != "EUR" { t.Errorf("config not updated correctly: %+v", cfg) } } // ─── JS object exposure ───────────────────────────────────────────────────── func TestInjectLLMLib_JSObjectExposed(t *testing.T) { g := minimalGateway() vm := otto.New() payload := &static.AgiLibInjectionPayload{VM: vm, User: &user.User{Username: "alice"}} g.injectLLMFunctions(payload) for _, method := range []string{"chat", "chatWithFile", "request", "streamRequest", "usage", "models"} { val, err := vm.Run(`typeof llm.` + method) if err != nil { t.Fatalf("evaluating llm.%s: %v", method, err) } s, _ := val.ToString() if s != "function" { t.Errorf("llm.%s should be a function, got %q", method, s) } } }