package handler import ( "context" "encoding/json" "net/http" "net/http/httptest" "testing" "time" "github.com/jackc/pgx/v5/pgtype" db "github.com/multica-ai/multica/server/pkg/db/generated" ) func TestRuntimeHandlersRejectMalformedRuntimeID(t *testing.T) { tests := []struct { name string method string path string handle func(http.ResponseWriter, *http.Request) }{ { name: "usage", method: "GET", path: "/api/runtimes/not-a-uuid/usage", handle: testHandler.GetRuntimeUsage, }, { name: "task activity", method: "GET", path: "/api/runtimes/not-a-uuid/task-activity", handle: testHandler.GetRuntimeTaskActivity, }, { name: "delete", method: "DELETE", path: "/api/runtimes/not-a-uuid", handle: testHandler.DeleteAgentRuntime, }, { name: "models", method: "POST", path: "/api/runtimes/not-a-uuid/models", handle: testHandler.InitiateListModels, }, { name: "update", method: "POST", path: "/api/runtimes/not-a-uuid/update", handle: testHandler.InitiateUpdate, }, { name: "local skills", method: "POST", path: "/api/runtimes/not-a-uuid/local-skills", handle: testHandler.InitiateListLocalSkills, }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { w := httptest.NewRecorder() req := newRequest(tt.method, tt.path, nil) req = withURLParam(req, "runtimeId", "not-a-uuid") tt.handle(w, req) if w.Code != http.StatusBadRequest { t.Fatalf("%s: expected 400 for malformed runtimeId, got %d: %s", tt.name, w.Code, w.Body.String()) } }) } } // TestGetRuntimeUsage_BucketsByUsageTime ensures a task that was enqueued on // one calendar day but whose tokens were reported the next day (e.g. execution // crossed midnight, or the task sat in the queue) is attributed to the day // tokens were actually produced, not the enqueue day. It also verifies the // ?days=N cutoff covers the full earliest calendar day, not just "now minus N // days" which would clip the morning of that day. func TestGetRuntimeUsage_BucketsByUsageTime(t *testing.T) { if testHandler == nil { t.Skip("database not available") } ctx := context.Background() // Pick a runtime bound to the fixture workspace. var runtimeID string if err := testPool.QueryRow(ctx, ` SELECT id FROM agent_runtime WHERE workspace_id = $1 LIMIT 1 `, testWorkspaceID).Scan(&runtimeID); err != nil { t.Fatalf("fetch runtime: %v", err) } var agentID string if err := testPool.QueryRow(ctx, ` SELECT id FROM agent WHERE workspace_id = $1 LIMIT 1 `, testWorkspaceID).Scan(&agentID); err != nil { t.Fatalf("fetch agent: %v", err) } // Create an issue for the tasks to reference. var issueID string if err := testPool.QueryRow(ctx, ` INSERT INTO issue (workspace_id, title, creator_id, creator_type) VALUES ($1, 'runtime usage test', $2, 'member') RETURNING id `, testWorkspaceID, testUserID).Scan(&issueID); err != nil { t.Fatalf("create issue: %v", err) } t.Cleanup(func() { testPool.Exec(ctx, `DELETE FROM issue WHERE id = $1`, issueID) }) // enqueued yesterday 23:58 UTC, finished today 00:05 UTC — tokens belong to today. now := time.Now().UTC() today := time.Date(now.Year(), now.Month(), now.Day(), 0, 0, 0, 0, time.UTC) yesterdayLate := today.Add(-2 * time.Minute) todayEarly := today.Add(5 * time.Minute) // Task that ran entirely yesterday around 05:00 — used to verify the // ?days cutoff isn't clipping yesterday's morning. yesterdayMorning := today.Add(-19 * time.Hour) insertTaskWithUsage := func(enqueueAt, usageAt time.Time, inputTokens int64) string { var taskID string if err := testPool.QueryRow(ctx, ` INSERT INTO agent_task_queue (agent_id, issue_id, runtime_id, status, created_at) VALUES ($1, $2, $3, 'completed', $4) RETURNING id `, agentID, issueID, runtimeID, enqueueAt).Scan(&taskID); err != nil { t.Fatalf("insert task: %v", err) } if _, err := testPool.Exec(ctx, ` INSERT INTO task_usage (task_id, provider, model, input_tokens, output_tokens, created_at) VALUES ($1, 'claude', 'claude-3-5-sonnet', $2, 0, $3) `, taskID, inputTokens, usageAt); err != nil { t.Fatalf("insert task_usage: %v", err) } t.Cleanup(func() { testPool.Exec(ctx, `DELETE FROM agent_task_queue WHERE id = $1`, taskID) }) return taskID } insertTaskWithUsage(yesterdayLate, todayEarly, 1000) // cross-midnight insertTaskWithUsage(yesterdayMorning, yesterdayMorning, 2000) // full-day yesterday // ListRuntimeUsage now reads from the `task_usage_daily` rollup // table maintained by the cron-driven rollup_task_usage_daily() // function. In production the watermarked wrapper waits a 5 min // safety lag before consuming rows; here we drive the underlying // window function directly with a wide-open range so the freshly // inserted fixture rows are guaranteed to be aggregated before the // handler is called. Each test invocation gets its own isolated // daily buckets keyed by (date, runtime, provider, model), so // re-running the test is idempotent (the upsert just rewrites the // same totals). if _, err := testPool.Exec(ctx, ` SELECT rollup_task_usage_daily_window('-infinity'::timestamptz, 'infinity'::timestamptz) `); err != nil { t.Fatalf("rollup_task_usage_daily_window: %v", err) } t.Cleanup(func() { testPool.Exec(ctx, ` DELETE FROM task_usage_daily WHERE runtime_id = $1 AND bucket_date IN ($2::date, $3::date) `, runtimeID, today, today.Add(-24*time.Hour)) }) // Call the handler with ?days=1 at whatever "now" is. That should include // both today and yesterday in full. w := httptest.NewRecorder() req := newRequest("GET", "/api/runtimes/"+runtimeID+"/usage?days=1", nil) req = withURLParam(req, "runtimeId", runtimeID) testHandler.GetRuntimeUsage(w, req) if w.Code != http.StatusOK { t.Fatalf("GetRuntimeUsage: expected 200, got %d: %s", w.Code, w.Body.String()) } var resp []RuntimeUsageResponse if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { t.Fatalf("decode response: %v", err) } byDate := make(map[string]int64) for _, r := range resp { byDate[r.Date] += r.InputTokens } todayKey := today.Format("2006-01-02") yesterdayKey := today.Add(-24 * time.Hour).Format("2006-01-02") // Cross-midnight task must attribute to today (tu.created_at), not yesterday // (atq.created_at). Before the fix this was 0 on today / 1000 on yesterday. if byDate[todayKey] != 1000 { t.Errorf("cross-midnight task: today bucket expected 1000 input tokens, got %d (full map: %v)", byDate[todayKey], byDate) } // Yesterday's morning task must still be included — this is what breaks // when ?days=N is interpreted as a rolling window instead of calendar days. if byDate[yesterdayKey] != 2000 { t.Errorf("yesterday morning task: yesterday bucket expected 2000 input tokens, got %d (full map: %v)", byDate[yesterdayKey], byDate) } } func TestGetRuntimeUsageDailyRollupCutoffUsesRuntimeTimezone(t *testing.T) { if testHandler == nil { t.Skip("database not available") } ctx := context.Background() runtimeID := handlerTestRuntimeID(t) loc, err := time.LoadLocation("Asia/Shanghai") if err != nil { t.Fatalf("load location: %v", err) } cutoff := time.Date(2026, 5, 4, 0, 0, 0, 0, loc) cutoffDate := cutoff.Format("2006-01-02") extraDate := cutoff.AddDate(0, 0, -1).Format("2006-01-02") var originalTZ string if err := testPool.QueryRow(ctx, `SELECT timezone FROM agent_runtime WHERE id = $1`, runtimeID).Scan(&originalTZ); err != nil { t.Fatalf("read runtime timezone: %v", err) } t.Cleanup(func() { testPool.Exec(ctx, `UPDATE agent_runtime SET timezone = $1 WHERE id = $2`, originalTZ, runtimeID) testPool.Exec(ctx, `DELETE FROM task_usage_daily WHERE runtime_id = $1 AND provider = 'cutoff-test'`, runtimeID) }) if _, err := testPool.Exec(ctx, `UPDATE agent_runtime SET timezone = 'Asia/Shanghai' WHERE id = $1`, runtimeID); err != nil { t.Fatalf("set runtime timezone: %v", err) } if _, err := testPool.Exec(ctx, ` INSERT INTO task_usage_daily ( bucket_date, workspace_id, runtime_id, provider, model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, event_count ) VALUES ($1::date, $3, $4, 'cutoff-test', 'old-day', 111, 0, 0, 0, 1), ($2::date, $3, $4, 'cutoff-test', 'cutoff-day', 222, 0, 0, 0, 1) ON CONFLICT (bucket_date, workspace_id, runtime_id, provider, model) DO UPDATE SET input_tokens = EXCLUDED.input_tokens, output_tokens = EXCLUDED.output_tokens, cache_read_tokens = EXCLUDED.cache_read_tokens, cache_write_tokens = EXCLUDED.cache_write_tokens, event_count = EXCLUDED.event_count `, extraDate, cutoffDate, testWorkspaceID, runtimeID); err != nil { t.Fatalf("seed rollup rows: %v", err) } origRollup := testHandler.cfg.UseDailyRollupForRuntimeUsage testHandler.cfg.UseDailyRollupForRuntimeUsage = true t.Cleanup(func() { testHandler.cfg.UseDailyRollupForRuntimeUsage = origRollup }) resp, err := testHandler.listRuntimeUsage(ctx, parseUUID(runtimeID), "Asia/Shanghai", pgtype.Timestamptz{ Time: cutoff, Valid: true, }) if err != nil { t.Fatalf("listRuntimeUsage: %v", err) } byDate := make(map[string]int64) for _, row := range resp { if row.Provider == "cutoff-test" { byDate[row.Date] += row.InputTokens } } if byDate[cutoffDate] != 222 { t.Fatalf("expected cutoff date %s to be included with 222 tokens, got map %v", cutoffDate, byDate) } if byDate[extraDate] != 0 { t.Fatalf("expected extra date %s to be excluded, got map %v", extraDate, byDate) } } func TestUpdateAgentRuntimeTimezoneValidatesPermissionAndValue(t *testing.T) { if testHandler == nil { t.Skip("database not available") } ctx := context.Background() runtimeID := handlerTestRuntimeID(t) var originalTZ string if err := testPool.QueryRow(ctx, `SELECT timezone FROM agent_runtime WHERE id = $1`, runtimeID).Scan(&originalTZ); err != nil { t.Fatalf("read runtime timezone: %v", err) } t.Cleanup(func() { testPool.Exec(ctx, `UPDATE agent_runtime SET timezone = $1 WHERE id = $2`, originalTZ, runtimeID) testPool.Exec(ctx, `DELETE FROM task_usage_daily WHERE runtime_id = $1 AND provider = 'patch-tz-test'`, runtimeID) }) w := httptest.NewRecorder() req := newRequest("PATCH", "/api/runtimes/"+runtimeID, map[string]string{"timezone": "Asia/Shanghai"}) req = withURLParam(req, "runtimeId", runtimeID) testHandler.UpdateAgentRuntime(w, req) if w.Code != http.StatusOK { t.Fatalf("valid timezone: expected 200, got %d: %s", w.Code, w.Body.String()) } var resp AgentRuntimeResponse if err := json.NewDecoder(w.Body).Decode(&resp); err != nil { t.Fatalf("decode response: %v", err) } if resp.Timezone != "Asia/Shanghai" { t.Fatalf("expected timezone Asia/Shanghai, got %q", resp.Timezone) } w = httptest.NewRecorder() req = newRequest("PATCH", "/api/runtimes/"+runtimeID, map[string]string{"timezone": "Mars/Olympus"}) req = withURLParam(req, "runtimeId", runtimeID) testHandler.UpdateAgentRuntime(w, req) if w.Code != http.StatusBadRequest { t.Fatalf("invalid timezone: expected 400, got %d: %s", w.Code, w.Body.String()) } var otherUserID string testPool.Exec(ctx, `DELETE FROM "user" WHERE email = 'runtime-tz-member@multica.ai'`) if err := testPool.QueryRow(ctx, ` INSERT INTO "user" (name, email) VALUES ('Runtime TZ Member', 'runtime-tz-member@multica.ai') RETURNING id `).Scan(&otherUserID); err != nil { t.Fatalf("create member user: %v", err) } t.Cleanup(func() { testPool.Exec(ctx, `DELETE FROM "user" WHERE id = $1`, otherUserID) }) if _, err := testPool.Exec(ctx, ` INSERT INTO member (workspace_id, user_id, role) VALUES ($1, $2, 'member') `, testWorkspaceID, otherUserID); err != nil { t.Fatalf("create member: %v", err) } w = httptest.NewRecorder() req = newRequest("PATCH", "/api/runtimes/"+runtimeID, map[string]string{"timezone": "Asia/Tokyo"}) req.Header.Set("X-User-ID", otherUserID) req = withURLParam(req, "runtimeId", runtimeID) testHandler.UpdateAgentRuntime(w, req) if w.Code != http.StatusForbidden { t.Fatalf("non-owner member: expected 403, got %d: %s", w.Code, w.Body.String()) } } func TestUpsertAgentRuntimePreservesTimezoneOverride(t *testing.T) { if testHandler == nil { t.Skip("database not available") } ctx := context.Background() testPool.Exec(ctx, ` DELETE FROM agent_runtime WHERE workspace_id = $1 AND daemon_id = 'tz-upsert-daemon' AND provider = 'tz-upsert-provider' `, testWorkspaceID) row, err := testHandler.Queries.UpsertAgentRuntime(ctx, db.UpsertAgentRuntimeParams{ WorkspaceID: parseUUID(testWorkspaceID), DaemonID: strToText("tz-upsert-daemon"), Name: "Timezone Upsert Runtime", RuntimeMode: "local", Provider: "tz-upsert-provider", Status: "online", DeviceInfo: "tz-upsert-device", Metadata: []byte(`{}`), OwnerID: parseUUID(testUserID), Timezone: "Asia/Shanghai", }) if err != nil { t.Fatalf("initial upsert: %v", err) } t.Cleanup(func() { testPool.Exec(ctx, `DELETE FROM agent_runtime WHERE id = $1`, row.ID) }) updated, err := testHandler.Queries.UpdateAgentRuntimeTimezone(ctx, db.UpdateAgentRuntimeTimezoneParams{ ID: row.ID, Timezone: "America/New_York", }) if err != nil { t.Fatalf("set override: %v", err) } if updated.Timezone != "America/New_York" { t.Fatalf("expected override to be set, got %q", updated.Timezone) } row, err = testHandler.Queries.UpsertAgentRuntime(ctx, db.UpsertAgentRuntimeParams{ WorkspaceID: parseUUID(testWorkspaceID), DaemonID: strToText("tz-upsert-daemon"), Name: "Timezone Upsert Runtime", RuntimeMode: "local", Provider: "tz-upsert-provider", Status: "online", DeviceInfo: "tz-upsert-device reconnect", Metadata: []byte(`{}`), OwnerID: pgtype.UUID{}, Timezone: "Asia/Tokyo", }) if err != nil { t.Fatalf("reconnect upsert: %v", err) } if row.Timezone != "America/New_York" { t.Fatalf("daemon reconnect should preserve user override, got %q", row.Timezone) } }