From 1fca2052313a72f3aaf627b7e695bbda74431e6e Mon Sep 17 00:00:00 2001 From: akarineren Date: Fri, 31 Jul 2026 16:01:32 +0900 Subject: [PATCH 01/11] Refactor project name handling: replace "unknown" with usage.UnknownProject for consistency across providers; update related tests --- internal/agentusage/pi.go | 2 +- internal/agentusage/providers_test.go | 2 +- internal/agentusage/qwen.go | 2 +- internal/claudeusage/loader.go | 108 ++++++++++------------- internal/claudeusage/loader_test.go | 122 +++++++++++++++++++------- internal/codexusage/loader.go | 14 +-- internal/codexusage/loader_test.go | 28 ++++++ internal/projectfile/projectfile.go | 4 +- internal/usage/usage.go | 14 +++ internal/usagedb/db.go | 118 +++++++++++++++++++++++++ internal/usagedb/db_test.go | 88 +++++++++++++++++++ 11 files changed, 400 insertions(+), 102 deletions(-) diff --git a/internal/agentusage/pi.go b/internal/agentusage/pi.go index 304eeac..c06812b 100644 --- a/internal/agentusage/pi.go +++ b/internal/agentusage/pi.go @@ -92,5 +92,5 @@ func piProject(path string) string { return parts[i+1] } } - return "unknown" + return usage.UnknownProject } diff --git a/internal/agentusage/providers_test.go b/internal/agentusage/providers_test.go index d906e79..9feb4a7 100644 --- a/internal/agentusage/providers_test.go +++ b/internal/agentusage/providers_test.go @@ -182,7 +182,7 @@ func TestProvidersLoadEntries(t *testing.T) { want: usage.ProviderPi, model: "[pi] gpt-5", sessionID: "session-a", - project: "unknown", + project: usage.UnknownProject, tokens: usage.TokenUsage{ OutputTokens: 333, TotalTokens: 333, diff --git a/internal/agentusage/qwen.go b/internal/agentusage/qwen.go index 36c2a40..34e01a5 100644 --- a/internal/agentusage/qwen.go +++ b/internal/agentusage/qwen.go @@ -100,5 +100,5 @@ func qwenProject(path string) string { return parts[i+1] } } - return "unknown" + return usage.UnknownProject } diff --git a/internal/claudeusage/loader.go b/internal/claudeusage/loader.go index 85520bb..973868f 100644 --- a/internal/claudeusage/loader.go +++ b/internal/claudeusage/loader.go @@ -26,6 +26,8 @@ type UsageEntry struct { Timestamp string `json:"timestamp"` Version *string `json:"version"` Entrypoint *string `json:"entrypoint"` + CWD *string `json:"cwd"` + GitBranch *string `json:"gitBranch"` Message UsageMessage `json:"message"` CostUSD *float64 `json:"costUSD"` RequestID *string `json:"requestId"` @@ -84,6 +86,7 @@ type LoadedEntry struct { Model string `json:"model,omitempty"` Language string `json:"language"` Client string `json:"client,omitempty"` + Branch string `json:"branch,omitempty"` UsageLimitResetTime *time.Time `json:"usage_limit_reset_time,omitempty"` } @@ -125,6 +128,7 @@ func ConvertEntries(entries []LoadedEntry) []usage.Entry { Language: usage.NormalizeLanguage(entry.Language), OS: usage.NormalizeOS(runtime.GOOS), Client: entry.Client, + Branch: entry.Branch, Usage: usage.TokenUsage{ InputTokens: tokens.InputTokens, OutputTokens: tokens.OutputTokens, @@ -217,8 +221,7 @@ func ReadUsageFile(path string) ([]LoadedEntry, error) { } defer file.Close() - project := ExtractProject(path) - sessionID, projectPath := ExtractSessionParts(path) + sessionID := ExtractSessionID(path) entries := make([]LoadedEntry, 0) reader := bufio.NewReader(file) lineNumber := 0 @@ -230,7 +233,7 @@ func ReadUsageFile(path string) ([]LoadedEntry, error) { start := offset offset += int64(len(line)) line = bytes.TrimRight(line, "\r\n") - entry, ok := parseUsageLine(line, project, sessionID, projectPath) + entry, ok := parseUsageLine(line, sessionID) if ok { entry.SourceFile = path entry.SourceLine = lineNumber @@ -251,26 +254,26 @@ func ReadUsageFile(path string) ([]LoadedEntry, error) { return entries, nil } -func ExtractProject(path string) string { - parts := pathParts(path) - for i, part := range parts { - if part != "projects" { - continue - } - if i+1 >= len(parts) || strings.TrimSpace(parts[i+1]) == "" { - return "unknown" - } - projectPath := normalizeClaudeProjectPathParts([]string{parts[i+1]}) - project := filepath.Base(filepath.Clean(projectPath)) - if strings.TrimSpace(project) == "" || project == "." || project == string(filepath.Separator) { - return "unknown" - } - return project +// projectFromCWD derives the project path and name from a transcript line's +// cwd field. It reports false for values that cannot name a project (empty, +// relative, or the filesystem root), so callers keep their fallback. +func projectFromCWD(cwd string) (string, string, bool) { + clean := strings.TrimSpace(cwd) + if clean == "" || !filepath.IsAbs(clean) { + return "", "", false } - return "unknown" + clean = filepath.Clean(clean) + name := filepath.Base(clean) + if name == "." || name == string(filepath.Separator) || strings.TrimSpace(name) == "" { + return "", "", false + } + return clean, name, true } -func ExtractSessionParts(path string) (string, string) { +// ExtractSessionID derives the session id from a usage file's location under +// the projects directory: projects//.jsonl for sessions and +// projects///subagents/.jsonl for subagents. +func ExtractSessionID(path string) string { parts := pathParts(path) relative := parts for i, part := range parts { @@ -288,55 +291,18 @@ func ExtractSessionParts(path string) (string, string) { } } if len(relative) == 2 && fileSessionID != "" { - return fileSessionID, normalizeClaudeProjectPathParts(relative[:1]) + return fileSessionID } if len(relative) >= 4 && relative[len(relative)-2] == "subagents" { - sessionID := relative[len(relative)-3] - return sessionID, normalizeClaudeProjectPathParts(relative[:len(relative)-3]) + return relative[len(relative)-3] } - - sessionID := "unknown" if len(relative) >= 2 { - sessionID = relative[len(relative)-2] - } - projectPath := "Unknown Project" - if len(relative) > 2 { - projectPath = normalizeClaudeProjectPathParts(relative[:len(relative)-2]) - } - return sessionID, projectPath -} - -func normalizeClaudeProjectPathParts(parts []string) string { - if len(parts) == 0 { - return "Unknown Project" + return relative[len(relative)-2] } - if decoded, ok := decodeClaudeProjectDir(parts[0]); ok { - if len(parts) == 1 { - return decoded - } - joined := append([]string{decoded}, parts[1:]...) - return filepath.Join(joined...) - } - projectPath := strings.Join(parts, string(filepath.Separator)) - if strings.TrimSpace(projectPath) == "" { - return "Unknown Project" - } - return projectPath -} - -func decodeClaudeProjectDir(segment string) (string, bool) { - segment = strings.TrimSpace(segment) - if !strings.HasPrefix(segment, "-") || len(segment) == 1 { - return "", false - } - decoded := string(filepath.Separator) + strings.ReplaceAll(strings.TrimPrefix(segment, "-"), "-", string(filepath.Separator)) - if filepath.Clean(decoded) == string(filepath.Separator) { - return "", false - } - return filepath.Clean(decoded), true + return "unknown" } -func parseUsageLine(line []byte, project, sessionID, projectPath string) (LoadedEntry, bool) { +func parseUsageLine(line []byte, sessionID string) (LoadedEntry, bool) { if !bytes.Contains(line, []byte(`"usage":{`)) { return LoadedEntry{}, false } @@ -356,6 +322,18 @@ func parseUsageLine(line []byte, project, sessionID, projectPath string) (Loaded return LoadedEntry{}, false } + // The transcript line's cwd is the only trustworthy project source: the + // directory name under ~/.claude/projects encodes "/", "-", "_" and "." + // identically, so decoding it is guesswork. No cwd means no project. + project := usage.UnknownProject + projectPath := "" + if data.CWD != nil { + if path, name, ok := projectFromCWD(*data.CWD); ok { + projectPath = path + project = name + } + } + model := "" if data.Message.Model != nil && *data.Message.Model != "" { model = *data.Message.Model @@ -369,6 +347,11 @@ func parseUsageLine(line []byte, project, sessionID, projectPath string) (Loaded client = usage.NormalizeClient(usage.ProviderClaude, *data.Entrypoint) } + branch := "" + if data.GitBranch != nil { + branch = strings.TrimSpace(*data.GitBranch) + } + return LoadedEntry{ Data: data, Timestamp: timestamp, @@ -379,6 +362,7 @@ func parseUsageLine(line []byte, project, sessionID, projectPath string) (Loaded Model: model, Language: languageFromContent(data.Message.Content), Client: client, + Branch: branch, UsageLimitResetTime: usageLimitResetTimeFromLine(line, data.IsAPIErrorMessage), }, true } diff --git a/internal/claudeusage/loader_test.go b/internal/claudeusage/loader_test.go index fa1387b..f9a852b 100644 --- a/internal/claudeusage/loader_test.go +++ b/internal/claudeusage/loader_test.go @@ -5,6 +5,8 @@ import ( "path/filepath" "testing" "time" + + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) func TestUsageFilesLimitsDiscoveryToProjectFilter(t *testing.T) { @@ -59,57 +61,115 @@ func TestProjectPathSegmentRejectsUnsafeValues(t *testing.T) { } } -func TestExtractSessionParts(t *testing.T) { +func TestExtractSessionID(t *testing.T) { tests := []struct { - name string - path string - wantSessionID string - wantProjectPath string + name string + path string + wantSessionID string }{ { - name: "modern", - path: "/home/me/.claude/projects/project-a/session-a.jsonl", - wantSessionID: "session-a", - wantProjectPath: "project-a", + name: "modern", + path: "/home/me/.claude/projects/project-a/session-a.jsonl", + wantSessionID: "session-a", }, { - name: "nested", - path: "/home/me/.claude/projects/project-a/session-a/chat.jsonl", - wantSessionID: "session-a", - wantProjectPath: "project-a", + name: "nested", + path: "/home/me/.claude/projects/project-a/session-a/chat.jsonl", + wantSessionID: "session-a", }, { - name: "subagent", - path: "/home/me/.claude/projects/project-a/session-a/subagents/worker.jsonl", - wantSessionID: "session-a", - wantProjectPath: "project-a", + name: "subagent", + path: "/home/me/.claude/projects/project-a/session-a/subagents/worker.jsonl", + wantSessionID: "session-a", }, { - name: "encoded absolute project path", - path: "/home/me/.claude/projects/-Users-eren-workspace-LABX-relink/session-a.jsonl", - wantSessionID: "session-a", - wantProjectPath: filepath.Join(string(filepath.Separator), "Users", "eren", "workspace", "LABX", "relink"), + name: "encoded absolute project path", + path: "/home/me/.claude/projects/-Users-eren-workspace-LABX-relink/session-a.jsonl", + wantSessionID: "session-a", }, } for _, tt := range tests { t.Run(tt.name, func(t *testing.T) { - sessionID, projectPath := ExtractSessionParts(tt.path) - if sessionID != tt.wantSessionID { + if sessionID := ExtractSessionID(tt.path); sessionID != tt.wantSessionID { t.Fatalf("sessionID = %q, want %q", sessionID, tt.wantSessionID) } - if projectPath != tt.wantProjectPath { - t.Fatalf("projectPath = %q, want %q", projectPath, tt.wantProjectPath) - } }) } } -func TestExtractProjectNormalizesEncodedClaudePath(t *testing.T) { - path := "/home/me/.claude/projects/-Users-eren-workspace-LABX-relink/session-a.jsonl" +func TestReadUsageFilePrefersLineCWDOverEncodedDir(t *testing.T) { + // The encoded directory name is ambiguous: decoding it yields + // "/Users/eren/workspace/tracklm/tracklm/nextjs". The line's cwd is the + // truth and must win. + path := filepath.Join(t.TempDir(), "projects", "-Users-eren-workspace-tracklm-tracklm-nextjs", "session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, `{"timestamp":"2026-05-21T01:02:03Z","cwd":"/Users/eren/workspace/tracklm/tracklm-nextjs","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}}`) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + if entries[0].Project != "tracklm-nextjs" { + t.Fatalf("project = %q, want tracklm-nextjs", entries[0].Project) + } + if entries[0].ProjectPath != "/Users/eren/workspace/tracklm/tracklm-nextjs" { + t.Fatalf("projectPath = %q, want /Users/eren/workspace/tracklm/tracklm-nextjs", entries[0].ProjectPath) + } +} + +func TestReadUsageFileCapturesGitBranch(t *testing.T) { + path := filepath.Join(t.TempDir(), "projects", "project-a", "session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, `{"timestamp":"2026-05-21T01:02:03Z","gitBranch":"feature/login","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}}`) - if got := ExtractProject(path); got != "relink" { - t.Fatalf("ExtractProject() = %q, want relink", got) + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + if entries[0].Branch != "feature/login" { + t.Fatalf("branch = %q, want feature/login", entries[0].Branch) + } + + converted := ConvertEntries(entries) + if converted[0].Branch != "feature/login" { + t.Fatalf("converted branch = %q, want feature/login", converted[0].Branch) + } +} + +func TestReadUsageFileReportsUnknownProjectWithoutCWD(t *testing.T) { + // The encoded directory name is not decodable ("/", "-", "_" and "." all + // become "-"), so a line without cwd has no project. + path := filepath.Join(t.TempDir(), "projects", "-Users-eren-workspace-LABX-relink", "session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, `{"timestamp":"2026-05-21T01:02:03Z","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}}`) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + if entries[0].Project != usage.UnknownProject { + t.Fatalf("project = %q, want %q", entries[0].Project, usage.UnknownProject) + } + if entries[0].ProjectPath != "" { + t.Fatalf("projectPath = %q, want empty", entries[0].ProjectPath) + } +} + +func TestProjectFromCWDRejectsUnusableValues(t *testing.T) { + for _, cwd := range []string{"", " ", "relative/path", "/"} { + if _, _, ok := projectFromCWD(cwd); ok { + t.Fatalf("projectFromCWD(%q) ok = true, want false", cwd) + } } } @@ -136,7 +196,7 @@ func TestReadUsageFileParsesUsageLines(t *testing.T) { mkdirAll(t, filepath.Dir(path)) writeFile(t, path, ` {"type":"user","message":{"content":"hello"}} -{"sessionId":"session-a","timestamp":"2026-05-21T01:02:03Z","version":"1.2.3","requestId":"req-1","message":{"id":"msg-1","model":"claude-sonnet-4-20250514","usage":{"input_tokens":10,"output_tokens":5,"cache_creation_input_tokens":2,"cache_read_input_tokens":3,"speed":"fast"}}} +{"sessionId":"session-a","timestamp":"2026-05-21T01:02:03Z","version":"1.2.3","requestId":"req-1","cwd":"/repo/project-a","message":{"id":"msg-1","model":"claude-sonnet-4-20250514","usage":{"input_tokens":10,"output_tokens":5,"cache_creation_input_tokens":2,"cache_read_input_tokens":3,"speed":"fast"}}} `) entries, err := ReadUsageFile(path) diff --git a/internal/codexusage/loader.go b/internal/codexusage/loader.go index 9b90bd7..fcf231a 100644 --- a/internal/codexusage/loader.go +++ b/internal/codexusage/loader.go @@ -112,7 +112,7 @@ func ReadUsageFile(path string) ([]usage.Entry, error) { entry.SourceLine = lineNumber entry.SourceStart = start entry.SourceEnd = offset - entry.ID = stableEntryID(entry) + entry.ID = StableEntryID(entry) entries = append(entries, entry) } } @@ -328,10 +328,14 @@ func candidatesFromText(text string, weight int) []langdetect.Candidate { return candidates } -func stableEntryID(entry usage.Entry) string { +// StableEntryID computes the deterministic id of a codex usage entry. The +// source file is keyed by its basename: archiving a session moves the file +// from sessions/ into archived_sessions/ without renaming it, and the id must +// survive that move or every archived session double-counts. +func StableEntryID(entry usage.Entry) string { return usage.StableID( string(usage.ProviderCodex), - entry.SourceFile, + filepath.Base(entry.SourceFile), strconv.Itoa(entry.SourceLine), entry.Timestamp.Format(time.RFC3339Nano), entry.Model, @@ -346,11 +350,11 @@ func stableEntryID(entry usage.Entry) string { func projectName(path string) string { clean := strings.TrimSpace(path) if clean == "" { - return "unknown" + return usage.UnknownProject } base := filepath.Base(filepath.Clean(clean)) if base == "." || base == string(filepath.Separator) || base == "" { - return "unknown" + return usage.UnknownProject } return base } diff --git a/internal/codexusage/loader_test.go b/internal/codexusage/loader_test.go index f51e3c6..9fa0584 100644 --- a/internal/codexusage/loader_test.go +++ b/internal/codexusage/loader_test.go @@ -146,3 +146,31 @@ func writeFile(t *testing.T, path, data string) { t.Fatal(err) } } + +func TestStableEntryIDSurvivesArchiveMove(t *testing.T) { + content := `{"timestamp":"2026-06-04T01:02:03Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo/app"}} +{"timestamp":"2026-06-04T01:02:04Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":4,"output_tokens":5,"reasoning_output_tokens":2,"total_tokens":15}}}} +` + dir := t.TempDir() + before := filepath.Join(dir, "sessions", "2026", "06", "04", "rollout-x.jsonl") + after := filepath.Join(dir, "archived_sessions", "rollout-x.jsonl") + mkdirAll(t, filepath.Dir(before)) + mkdirAll(t, filepath.Dir(after)) + writeFile(t, before, content) + writeFile(t, after, content) + + beforeEntries, err := ReadUsageFile(before) + if err != nil { + t.Fatal(err) + } + afterEntries, err := ReadUsageFile(after) + if err != nil { + t.Fatal(err) + } + if len(beforeEntries) != 1 || len(afterEntries) != 1 { + t.Fatalf("entries = %d/%d, want 1/1", len(beforeEntries), len(afterEntries)) + } + if beforeEntries[0].ID != afterEntries[0].ID { + t.Fatalf("id changed across archive move: %q vs %q", beforeEntries[0].ID, afterEntries[0].ID) + } +} diff --git a/internal/projectfile/projectfile.go b/internal/projectfile/projectfile.go index 27dcca1..da13911 100644 --- a/internal/projectfile/projectfile.go +++ b/internal/projectfile/projectfile.go @@ -10,6 +10,8 @@ import ( "path/filepath" "strings" "unicode/utf8" + + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) const ( @@ -193,7 +195,7 @@ func findVCSRoot(start searchStart) (string, bool) { func projectName(path string) string { name := filepath.Base(filepath.Clean(path)) if name == "." || name == string(filepath.Separator) || name == "" { - return "unknown" + return usage.UnknownProject } return name } diff --git a/internal/usage/usage.go b/internal/usage/usage.go index c7c0fcf..7b8de83 100644 --- a/internal/usage/usage.go +++ b/internal/usage/usage.go @@ -30,6 +30,20 @@ const ( const UnknownLanguage = "Unknown" +// UnknownProject is the single spelling every provider uses when a project +// name cannot be determined. +const UnknownProject = "Unknown" + +// NormalizeProject maps empty and legacy "unknown" spellings to +// UnknownProject so undetermined projects look the same everywhere. +func NormalizeProject(name string) string { + name = strings.TrimSpace(name) + if name == "" || strings.EqualFold(name, "unknown") { + return UnknownProject + } + return name +} + // FileFilter reports whether a source file must be parsed. Returning false // means the file's events are already ingested and parsing it again would be // wasted work. A nil FileFilter parses everything. diff --git a/internal/usagedb/db.go b/internal/usagedb/db.go index 4e18f99..9b6548f 100644 --- a/internal/usagedb/db.go +++ b/internal/usagedb/db.go @@ -12,6 +12,7 @@ import ( "path/filepath" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/codexusage" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" _ "modernc.org/sqlite" ) @@ -70,9 +71,125 @@ func Open(path string) (*DB, error) { _ = db.Close() return nil, fmt.Errorf("migrate usage db: %w", err) } + if err := migrate(db); err != nil { + _ = db.Close() + return nil, fmt.Errorf("migrate usage db: %w", err) + } return &DB{db: db}, nil } +// eventSchemaVersion tracks one-time data repairs, recorded in the SQLite +// user_version pragma. Version 1 rekeys codex events whose ids hashed the +// source file's full path: archiving a session moves its file from sessions/ +// into archived_sessions/, which changed every id and double-counted the +// whole session. +const eventSchemaVersion = 1 + +func migrate(db *sql.DB) error { + var version int + if err := db.QueryRow(`PRAGMA user_version`).Scan(&version); err != nil { + return err + } + if version >= eventSchemaVersion { + return nil + } + if version < 1 { + if err := rekeyCodexEvents(db); err != nil { + return err + } + } + _, err := db.Exec(fmt.Sprintf(`PRAGMA user_version = %d`, eventSchemaVersion)) + return err +} + +func rekeyCodexEvents(db *sql.DB) error { + rows, err := db.Query(`SELECT id, payload, status FROM usage_events`) + if err != nil { + return err + } + defer rows.Close() + + type keeper struct { + oldID string + payload string + rank int + } + keepers := make(map[string]keeper) + drop := make([]string, 0) + for rows.Next() { + var id, payload, status string + if err := rows.Scan(&id, &payload, &status); err != nil { + return err + } + var entry usage.Entry + if err := json.Unmarshal([]byte(payload), &entry); err != nil { + continue + } + if entry.Provider != usage.ProviderCodex { + continue + } + newID := codexusage.StableEntryID(entry) + candidate := keeper{oldID: id, payload: payload, rank: statusRank(status)} + current, exists := keepers[newID] + switch { + case !exists: + keepers[newID] = candidate + case candidate.rank > current.rank: + // Rows collapsing onto one id are the same event ingested twice + // from the file's pre- and post-archive paths; the uploaded copy + // wins so the duplicate is never re-uploaded. + drop = append(drop, current.oldID) + keepers[newID] = candidate + default: + drop = append(drop, id) + } + } + if err := rows.Err(); err != nil { + return err + } + + tx, err := db.Begin() + if err != nil { + return err + } + defer tx.Rollback() + + for _, id := range drop { + if _, err := tx.Exec(`DELETE FROM usage_events WHERE id = ?`, id); err != nil { + return err + } + } + for newID, kept := range keepers { + if kept.oldID == newID { + continue + } + var entry usage.Entry + if err := json.Unmarshal([]byte(kept.payload), &entry); err != nil { + continue + } + entry.ID = newID + payload, err := json.Marshal(entry) + if err != nil { + return err + } + if _, err := tx.Exec(`UPDATE usage_events SET id = ?, payload = ? WHERE id = ?`, newID, string(payload), kept.oldID); err != nil { + return err + } + } + return tx.Commit() +} + +func statusRank(status string) int { + switch status { + case "uploaded": + return 3 + case "rejected": + return 2 + default: // pending, failed + return 1 + } +} + func (s *DB) Close() error { return s.db.Close() } @@ -101,6 +218,7 @@ func (s *DB) InsertEvents(entries []usage.Entry) (int, error) { return 0, fmt.Errorf("usage event id is required") } entry.Language = usage.NormalizeLanguage(entry.Language) + entry.Project = usage.NormalizeProject(entry.Project) payload, err := json.Marshal(entry) if err != nil { return 0, fmt.Errorf("encode usage event %q: %w", entry.ID, err) diff --git a/internal/usagedb/db_test.go b/internal/usagedb/db_test.go index 79639e4..6143d1a 100644 --- a/internal/usagedb/db_test.go +++ b/internal/usagedb/db_test.go @@ -5,6 +5,7 @@ import ( "testing" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/codexusage" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -209,3 +210,90 @@ func testUsageEntry(id string) usage.Entry { }, } } + +func TestOpenRekeysCodexEventsToBasenameIDs(t *testing.T) { + path := filepath.Join(t.TempDir(), "usage.db") + db, err := Open(path) + if err != nil { + t.Fatal(err) + } + + legacyID := func(entry usage.Entry) string { + return usage.StableID( + string(usage.ProviderCodex), + entry.SourceFile, + "1", + entry.Timestamp.Format(time.RFC3339Nano), + entry.Model, + "1", "0", "2", "0", "3", + ) + } + + // The same event ingested twice: once from the live path, once after the + // file was archived. Full-path ids made them distinct rows. + live := testUsageEntry("") + live.SourceFile = "/home/me/.codex/sessions/2026/06/04/rollout-x.jsonl" + live.ID = legacyID(live) + archived := live + archived.SourceFile = "/home/me/.codex/archived_sessions/rollout-x.jsonl" + archived.ID = legacyID(archived) + if live.ID == archived.ID { + t.Fatal("test premise broken: legacy ids should differ across paths") + } + if _, err := db.InsertEvents([]usage.Entry{live, archived}); err != nil { + t.Fatal(err) + } + if err := db.MarkEventsUploaded([]string{live.ID}); err != nil { + t.Fatal(err) + } + // Pretend this database was written by a pre-migration binary. + if _, err := db.db.Exec(`PRAGMA user_version = 0`); err != nil { + t.Fatal(err) + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + + reopened, err := Open(path) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = reopened.Close() }) + + rows, err := reopened.db.Query(`SELECT id, status FROM usage_events`) + if err != nil { + t.Fatal(err) + } + defer rows.Close() + type row struct{ id, status string } + got := make([]row, 0) + for rows.Next() { + var r row + if err := rows.Scan(&r.id, &r.status); err != nil { + t.Fatal(err) + } + got = append(got, r) + } + if err := rows.Err(); err != nil { + t.Fatal(err) + } + + if len(got) != 1 { + t.Fatalf("rows = %d, want 1 (duplicates collapsed)", len(got)) + } + if got[0].id != codexusage.StableEntryID(live) { + t.Fatalf("id = %q, want basename-keyed id %q", got[0].id, codexusage.StableEntryID(live)) + } + if got[0].status != "uploaded" { + t.Fatalf("status = %q, want uploaded (uploaded copy must win)", got[0].status) + } + + // The rewritten payload must carry the new id so future uploads use it. + pruned, err := reopened.PruneUploaded(time.Now().Add(time.Hour)) + if err != nil { + t.Fatal(err) + } + if pruned != 1 { + t.Fatalf("pruned = %d, want 1", pruned) + } +} From 4df13ec3cd5434b1ca4c555e982e434955717f28 Mon Sep 17 00:00:00 2001 From: akarineren Date: Fri, 31 Jul 2026 19:32:28 +0900 Subject: [PATCH 02/11] feat: enhance codex usage tracking with patch confirmation and cumulative deltas - Introduced support for tracking file modifications through patches, including confirmation of successful patch applications via new event types. - Updated token usage calculations to utilize cumulative deltas, ensuring accurate reporting even in cases of counter resets. - Enhanced the `StableEntryID` function to derive IDs from event details rather than file paths, preventing double-counting during file moves. - Added tests to validate the new functionality, including handling of patch application success and failure scenarios. - Implemented relative path handling for uploaded entities and files to obscure local filesystem structure. --- cmd/tokitoki/main_test.go | 4 +- internal/claudeusage/loader.go | 201 ++++++++++++--- internal/claudeusage/loader_test.go | 142 ++++++++++ internal/codexusage/loader.go | 356 +++++++++++++++++++++++--- internal/codexusage/loader_test.go | 235 +++++++++++++++-- internal/usage/usage.go | 15 ++ internal/usagedb/db.go | 42 ++- internal/usagedb/db_test.go | 14 +- internal/usageupload/uploader.go | 95 ++++--- internal/usageupload/uploader_test.go | 35 +++ 10 files changed, 1009 insertions(+), 130 deletions(-) diff --git a/cmd/tokitoki/main_test.go b/cmd/tokitoki/main_test.go index db86be1..1a35447 100644 --- a/cmd/tokitoki/main_test.go +++ b/cmd/tokitoki/main_test.go @@ -175,7 +175,9 @@ func TestRunHeartbeatUploadsUnifiedIDEEvent(t *testing.T) { if event.SourceType != "ide" || event.SourceProvider != "eclipse" || event.EventKind != "heartbeat" { t.Fatalf("source fields = %+v, want Eclipse IDE heartbeat", event) } - if event.Entity != "/repo/src/App.java" || event.Language != "Java" { + // The entity is uploaded relative to the project folder: the absolute + // path would leak the machine layout that project_path_hash hides. + if event.Entity != "src/App.java" || event.Language != "Java" { t.Fatalf("entity/language = %q/%q, want Java file", event.Entity, event.Language) } if event.IsWrite == nil || !*event.IsWrite { diff --git a/internal/claudeusage/loader.go b/internal/claudeusage/loader.go index 973868f..77f7068 100644 --- a/internal/claudeusage/loader.go +++ b/internal/claudeusage/loader.go @@ -72,22 +72,27 @@ func (s *Speed) UnmarshalJSON(data []byte) error { } type LoadedEntry struct { - Data UsageEntry `json:"data"` - ID string `json:"id,omitempty"` - SourceFile string `json:"source_file,omitempty"` - SourceLine int `json:"source_line,omitempty"` - SourceStart int64 `json:"source_start,omitempty"` - SourceEnd int64 `json:"source_end,omitempty"` - Timestamp time.Time `json:"timestamp"` - Date string `json:"date"` - Project string `json:"project"` - SessionID string `json:"session_id"` - ProjectPath string `json:"project_path"` - Model string `json:"model,omitempty"` - Language string `json:"language"` - Client string `json:"client,omitempty"` - Branch string `json:"branch,omitempty"` - UsageLimitResetTime *time.Time `json:"usage_limit_reset_time,omitempty"` + Data UsageEntry `json:"data"` + ID string `json:"id,omitempty"` + SourceFile string `json:"source_file,omitempty"` + SourceLine int `json:"source_line,omitempty"` + SourceStart int64 `json:"source_start,omitempty"` + SourceEnd int64 `json:"source_end,omitempty"` + Timestamp time.Time `json:"timestamp"` + Date string `json:"date"` + Project string `json:"project"` + SessionID string `json:"session_id"` + ProjectPath string `json:"project_path"` + Model string `json:"model,omitempty"` + Language string `json:"language"` + Client string `json:"client,omitempty"` + Branch string `json:"branch,omitempty"` + Entity string `json:"entity,omitempty"` + IsWrite bool `json:"is_write,omitempty"` + LinesAdded uint64 `json:"lines_added,omitempty"` + LinesRemoved uint64 `json:"lines_removed,omitempty"` + Files []usage.FileChange `json:"files,omitempty"` + UsageLimitResetTime *time.Time `json:"usage_limit_reset_time,omitempty"` } type DailyProjectSummary struct { @@ -112,23 +117,38 @@ func ConvertEntries(entries []LoadedEntry) []usage.Entry { converted := make([]usage.Entry, 0, len(entries)) for _, entry := range entries { tokens := entry.Data.Message.Usage + var isWrite *bool + if entry.IsWrite { + t := true + isWrite = &t + } + entityType := "" + if entry.Entity != "" { + entityType = "file" + } converted = append(converted, usage.Entry{ - Provider: usage.ProviderClaude, - ID: entry.ID, - SourceFile: entry.SourceFile, - SourceLine: entry.SourceLine, - SourceStart: entry.SourceStart, - SourceEnd: entry.SourceEnd, - Timestamp: entry.Timestamp, - Date: entry.Date, - Project: entry.Project, - ProjectPath: entry.ProjectPath, - SessionID: entry.SessionID, - Model: entry.Model, - Language: usage.NormalizeLanguage(entry.Language), - OS: usage.NormalizeOS(runtime.GOOS), - Client: entry.Client, - Branch: entry.Branch, + Provider: usage.ProviderClaude, + ID: entry.ID, + SourceFile: entry.SourceFile, + SourceLine: entry.SourceLine, + SourceStart: entry.SourceStart, + SourceEnd: entry.SourceEnd, + Timestamp: entry.Timestamp, + Date: entry.Date, + Project: entry.Project, + ProjectPath: entry.ProjectPath, + SessionID: entry.SessionID, + Model: entry.Model, + Language: usage.NormalizeLanguage(entry.Language), + OS: usage.NormalizeOS(runtime.GOOS), + Client: entry.Client, + Branch: entry.Branch, + Entity: entry.Entity, + EntityType: entityType, + IsWrite: isWrite, + LinesAdded: entry.LinesAdded, + LinesRemoved: entry.LinesRemoved, + Files: entry.Files, Usage: usage.TokenUsage{ InputTokens: tokens.InputTokens, OutputTokens: tokens.OutputTokens, @@ -223,6 +243,10 @@ func ReadUsageFile(path string) ([]LoadedEntry, error) { sessionID := ExtractSessionID(path) entries := make([]LoadedEntry, 0) + // A file-modification diff belongs to the assistant message that issued + // the edit, which precedes its tool result in the transcript. Diffs seen + // before the first usage entry are held and attached to it. + pending := make([]patchStats, 0) reader := bufio.NewReader(file) lineNumber := 0 offset := int64(0) @@ -233,14 +257,23 @@ func ReadUsageFile(path string) ([]LoadedEntry, error) { start := offset offset += int64(len(line)) line = bytes.TrimRight(line, "\r\n") - entry, ok := parseUsageLine(line, sessionID) - if ok { + if entry, ok := parseUsageLine(line, sessionID); ok { entry.SourceFile = path entry.SourceLine = lineNumber entry.SourceStart = start entry.SourceEnd = offset entry.ID = stableEntryID(entry) entries = append(entries, entry) + for _, patch := range pending { + applyPatch(&entries[len(entries)-1], patch) + } + pending = pending[:0] + } else if patch, ok := parsePatchLine(line); ok { + if len(entries) == 0 { + pending = append(pending, patch) + } else { + applyPatch(&entries[len(entries)-1], patch) + } } } if readErr == nil { @@ -254,6 +287,105 @@ func ReadUsageFile(path string) ([]LoadedEntry, error) { return entries, nil } +type patchStats struct { + file string + added uint64 + removed uint64 +} + +// applyPatch accumulates one diff into the entry's totals and per-file +// breakdown. The entity is always the most-changed file so far. +func applyPatch(entry *LoadedEntry, patch patchStats) { + entry.LinesAdded += patch.added + entry.LinesRemoved += patch.removed + entry.IsWrite = true + if patch.file == "" { + return + } + + index := -1 + for i := range entry.Files { + if entry.Files[i].Path == patch.file { + index = i + break + } + } + if index < 0 { + entry.Files = append(entry.Files, usage.FileChange{Path: patch.file}) + index = len(entry.Files) - 1 + } + entry.Files[index].LinesAdded += patch.added + entry.Files[index].LinesRemoved += patch.removed + + best := index + for i := range entry.Files { + if entry.Files[i].LinesAdded+entry.Files[i].LinesRemoved > entry.Files[best].LinesAdded+entry.Files[best].LinesRemoved { + best = i + } + } + entry.Entity = entry.Files[best].Path +} + +func parsePatchLine(line []byte) (patchStats, bool) { + if !bytes.Contains(line, []byte(`"structuredPatch"`)) { + return patchStats{}, false + } + + var envelope struct { + ToolUseResult json.RawMessage `json:"toolUseResult"` + } + if err := json.Unmarshal(line, &envelope); err != nil || len(envelope.ToolUseResult) == 0 { + return patchStats{}, false + } + var result struct { + Type string `json:"type"` + FilePath string `json:"filePath"` + Content string `json:"content"` + StructuredPatch []struct { + Lines []string `json:"lines"` + } `json:"structuredPatch"` + } + if err := json.Unmarshal(envelope.ToolUseResult, &result); err != nil { + return patchStats{}, false + } + + // Creating a file records no diff hunks, only the full content: every + // content line is an added line. + if len(result.StructuredPatch) == 0 { + if result.Type != "create" || result.FilePath == "" { + return patchStats{}, false + } + return patchStats{file: result.FilePath, added: countLines(result.Content)}, true + } + + stats := patchStats{file: result.FilePath} + for _, hunk := range result.StructuredPatch { + for _, hunkLine := range hunk.Lines { + if len(hunkLine) == 0 { + continue + } + switch hunkLine[0] { + case '+': + stats.added++ + case '-': + stats.removed++ + } + } + } + return stats, true +} + +func countLines(content string) uint64 { + if content == "" { + return 0 + } + lines := uint64(strings.Count(content, "\n")) + if !strings.HasSuffix(content, "\n") { + lines++ + } + return lines +} + // projectFromCWD derives the project path and name from a transcript line's // cwd field. It reports false for values that cannot name a project (empty, // relative, or the filesystem root), so callers keep their fallback. @@ -740,4 +872,3 @@ func pathParts(path string) []string { }) return parts } - diff --git a/internal/claudeusage/loader_test.go b/internal/claudeusage/loader_test.go index f9a852b..b81a6f5 100644 --- a/internal/claudeusage/loader_test.go +++ b/internal/claudeusage/loader_test.go @@ -121,6 +121,111 @@ func TestReadUsageFilePrefersLineCWDOverEncodedDir(t *testing.T) { } } +func TestReadUsageFileAttributesPatchesToIssuingEntry(t *testing.T) { + path := filepath.Join(t.TempDir(), "projects", "project-a", "session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + // Entry msg-1 issues two edits; small.go changes 2 lines, big.go 3. The + // entity is the most-changed file, the line counts are the sum. + writeFile(t, path, ` +{"timestamp":"2026-05-21T01:02:03Z","cwd":"/repo/app","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}} +{"type":"user","timestamp":"2026-05-21T01:02:04Z","toolUseResult":{"filePath":"/repo/app/small.go","structuredPatch":[{"oldStart":1,"oldLines":1,"newStart":1,"newLines":1,"lines":["+added","-removed"," context"]}]}} +{"type":"user","timestamp":"2026-05-21T01:02:05Z","toolUseResult":{"filePath":"/repo/app/big.go","structuredPatch":[{"oldStart":1,"oldLines":0,"newStart":1,"newLines":3,"lines":["+a","+b","+c"]}]}} +{"timestamp":"2026-05-21T01:02:06Z","cwd":"/repo/app","message":{"id":"msg-2","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}} +`) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("len(entries) = %d, want 2", len(entries)) + } + + first := entries[0] + if first.LinesAdded != 4 || first.LinesRemoved != 1 { + t.Fatalf("lines = +%d/-%d, want +4/-1", first.LinesAdded, first.LinesRemoved) + } + if first.Entity != "/repo/app/big.go" { + t.Fatalf("entity = %q, want most-changed file big.go", first.Entity) + } + if !first.IsWrite { + t.Fatal("isWrite = false, want true") + } + + second := entries[1] + if second.IsWrite || second.LinesAdded != 0 || second.Entity != "" { + t.Fatalf("second entry inherited patch data: %+v", second) + } + + if len(first.Files) != 2 { + t.Fatalf("files = %+v, want small.go and big.go", first.Files) + } + for _, file := range first.Files { + switch file.Path { + case "/repo/app/small.go": + if file.LinesAdded != 1 || file.LinesRemoved != 1 { + t.Fatalf("small.go = +%d/-%d, want +1/-1", file.LinesAdded, file.LinesRemoved) + } + case "/repo/app/big.go": + if file.LinesAdded != 3 || file.LinesRemoved != 0 { + t.Fatalf("big.go = +%d/-%d, want +3/-0", file.LinesAdded, file.LinesRemoved) + } + default: + t.Fatalf("unexpected file %q", file.Path) + } + } + + converted := ConvertEntries(entries) + if converted[0].Entity != "/repo/app/big.go" || converted[0].EntityType != "file" { + t.Fatalf("converted entity = %q/%q, want big.go/file", converted[0].Entity, converted[0].EntityType) + } + if len(converted[0].Files) != 2 { + t.Fatalf("converted files = %+v, want 2", converted[0].Files) + } + if converted[0].IsWrite == nil || !*converted[0].IsWrite { + t.Fatal("converted isWrite not set") + } + if converted[0].LinesAdded != 4 || converted[0].LinesRemoved != 1 { + t.Fatalf("converted lines = +%d/-%d, want +4/-1", converted[0].LinesAdded, converted[0].LinesRemoved) + } + if converted[1].IsWrite != nil || converted[1].EntityType != "" { + t.Fatalf("converted second entry inherited patch data: %+v", converted[1]) + } +} + +func TestReadUsageFileHoldsPatchesBeforeFirstEntry(t *testing.T) { + path := filepath.Join(t.TempDir(), "projects", "project-a", "session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, ` +{"type":"user","timestamp":"2026-05-21T01:02:02Z","toolUseResult":{"filePath":"/repo/app/a.go","structuredPatch":[{"lines":["+x"]}]}} +{"timestamp":"2026-05-21T01:02:03Z","cwd":"/repo/app","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}} +`) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + if entries[0].LinesAdded != 1 || entries[0].Entity != "/repo/app/a.go" || !entries[0].IsWrite { + t.Fatalf("pending patch not attached: %+v", entries[0]) + } +} + +func TestParsePatchLineIgnoresNonPatchToolResults(t *testing.T) { + lines := [][]byte{ + []byte(`{"toolUseResult":"plain text mentioning structuredPatch"}`), + []byte(`{"toolUseResult":{"filePath":"/a.go","structuredPatch":[]}}`), + []byte(`{"message":{"content":"structuredPatch"}}`), + } + for _, line := range lines { + if _, ok := parsePatchLine(line); ok { + t.Fatalf("parsePatchLine(%s) ok = true, want false", line) + } + } +} + func TestReadUsageFileCapturesGitBranch(t *testing.T) { path := filepath.Join(t.TempDir(), "projects", "project-a", "session-a.jsonl") mkdirAll(t, filepath.Dir(path)) @@ -332,3 +437,40 @@ func containsPathSegment(path, segment string) bool { } return false } + +func TestApplyPatchAccumulatesSameFile(t *testing.T) { + entry := LoadedEntry{} + applyPatch(&entry, patchStats{file: "/repo/a.go", added: 2, removed: 1}) + applyPatch(&entry, patchStats{file: "/repo/b.go", added: 1}) + applyPatch(&entry, patchStats{file: "/repo/a.go", added: 1}) + + if len(entry.Files) != 2 { + t.Fatalf("files = %+v, want 2", entry.Files) + } + if entry.Files[0].Path != "/repo/a.go" || entry.Files[0].LinesAdded != 3 || entry.Files[0].LinesRemoved != 1 { + t.Fatalf("a.go = %+v, want +3/-1", entry.Files[0]) + } + if entry.Entity != "/repo/a.go" { + t.Fatalf("entity = %q, want cumulative most-changed a.go", entry.Entity) + } + if entry.LinesAdded != 4 || entry.LinesRemoved != 1 { + t.Fatalf("totals = +%d/-%d, want +4/-1", entry.LinesAdded, entry.LinesRemoved) + } +} + +func TestParsePatchLineCountsCreatedFileContent(t *testing.T) { + line := []byte(`{"toolUseResult":{"type":"create","filePath":"/repo/new.go","content":"package main\n\nfunc main() {}\n","structuredPatch":[]}}`) + stats, ok := parsePatchLine(line) + if !ok { + t.Fatal("parsePatchLine ok = false, want true") + } + if stats.file != "/repo/new.go" || stats.added != 3 || stats.removed != 0 { + t.Fatalf("stats = %+v, want new.go +3/-0", stats) + } + + empty := []byte(`{"toolUseResult":{"type":"create","filePath":"/repo/empty.go","content":"","structuredPatch":[]}}`) + stats, ok = parsePatchLine(empty) + if !ok || stats.added != 0 { + t.Fatalf("empty create = %+v ok=%v, want +0 ok=true", stats, ok) + } +} diff --git a/internal/codexusage/loader.go b/internal/codexusage/loader.go index fcf231a..9218054 100644 --- a/internal/codexusage/loader.go +++ b/internal/codexusage/loader.go @@ -38,9 +38,12 @@ type turnContextPayload struct { } type eventPayload struct { - Type string `json:"type"` - Info struct { - LastTokenUsage *tokenUsagePayload `json:"last_token_usage"` + Type string `json:"type"` + CallID string `json:"call_id"` + Success *bool `json:"success"` + Info struct { + LastTokenUsage *tokenUsagePayload `json:"last_token_usage"` + TotalTokenUsage *tokenUsagePayload `json:"total_token_usage"` } `json:"info"` } @@ -112,7 +115,6 @@ func ReadUsageFile(path string) ([]usage.Entry, error) { entry.SourceLine = lineNumber entry.SourceStart = start entry.SourceEnd = offset - entry.ID = StableEntryID(entry) entries = append(entries, entry) } } @@ -133,6 +135,20 @@ type fileState struct { model string language string client string + // awaiting holds parsed patches keyed by call_id until their tool output + // confirms the patch actually applied; pending holds confirmed patches + // waiting to be folded into the next token_count entry. + awaiting map[string][]patchFile + pending []patchFile + // prevTotal is the last seen cumulative token counter, the basis for + // per-event deltas. + prevTotal *tokenUsagePayload +} + +type patchFile struct { + path string + added uint64 + removed uint64 } func parseLine(line []byte, state *fileState) (usage.Entry, bool) { @@ -150,6 +166,9 @@ func parseLine(line []byte, state *fileState) (usage.Entry, bool) { } switch envelope.Type { + case "response_item": + handleResponseItem(envelope.Payload, state) + return usage.Entry{}, false case "session_meta": var payload sessionMetaPayload if err := json.Unmarshal(envelope.Payload, &payload); err != nil { @@ -182,6 +201,17 @@ func parseLine(line []byte, state *fileState) (usage.Entry, bool) { if err := json.Unmarshal(envelope.Payload, &payload); err != nil { return usage.Entry{}, false } + // Newer codex confirms patches with a dedicated event instead of a + // tool output; either resolves the same awaiting call_id. + if payload.Type == "patch_apply_end" { + if patches, ok := state.awaiting[payload.CallID]; ok { + delete(state.awaiting, payload.CallID) + if payload.Success != nil && *payload.Success { + state.pending = append(state.pending, patches...) + } + } + return usage.Entry{}, false + } if payload.Type != "token_count" || payload.Info.LastTokenUsage == nil { return usage.Entry{}, false } @@ -189,20 +219,29 @@ func parseLine(line []byte, state *fileState) (usage.Entry, bool) { if err != nil { return usage.Entry{}, false } - tokens := payload.Info.LastTokenUsage - - // Codex reports input_tokens as the FULL prompt (cached + non-cached), - // with cached_input_tokens being the cached portion. Match ccusage: - // real input = input_tokens - cached, and the cached part is cache read. - // Otherwise the cached prompt gets double-counted into input. - nonCachedInput := tokens.InputTokens - if tokens.CachedInputTokens <= tokens.InputTokens { - nonCachedInput = tokens.InputTokens - tokens.CachedInputTokens - } else { - nonCachedInput = 0 + last := *payload.Info.LastTokenUsage + + // The cumulative counter is authoritative: codex replays the same + // last_token_usage across duplicate emissions and retries, so summing + // it overcounts (up to +50% on real sessions). Each entry's usage is + // the counter delta; last_token_usage covers files without a counter + // and counter resets. The id stays derived from last_token_usage so + // this accounting change never shifts event identity. + event := last + if total := payload.Info.TotalTokenUsage; total != nil { + if state.prevTotal == nil { + event = *total + } else if delta, ok := diffTokenUsage(*total, *state.prevTotal); ok { + if delta == (tokenUsagePayload{}) { + state.prevTotal = total + return usage.Entry{}, false + } + event = delta + } + state.prevTotal = total } - return usage.Entry{ + entry := usage.Entry{ Provider: usage.ProviderCodex, Timestamp: timestamp, Date: timestamp.In(time.Local).Format("2006-01-02"), @@ -213,14 +252,14 @@ func parseLine(line []byte, state *fileState) (usage.Entry, bool) { Language: stateLanguage(state), OS: usage.NormalizeOS(runtime.GOOS), Client: state.client, - Usage: usage.TokenUsage{ - InputTokens: nonCachedInput, - CacheReadInputTokens: tokens.CachedInputTokens, - OutputTokens: tokens.OutputTokens, - ReasoningOutputTokens: tokens.ReasoningOutputTokens, - TotalTokens: tokens.TotalTokens, - }, - }, true + Usage: accountUsage(event), + } + idEntry := entry + idEntry.Usage = accountUsage(last) + entry.ID = StableEntryID(idEntry) + applyPatches(&entry, state.pending) + state.pending = nil + return entry, true default: return usage.Entry{}, false } @@ -230,6 +269,261 @@ func stateLanguage(state *fileState) string { return usage.NormalizeLanguage(state.language) } +// accountUsage maps a raw codex token payload to our accounting: input_tokens +// is the FULL prompt (cached + non-cached), so real input = input - cached +// and the cached portion is cache read. Matches ccusage. +func accountUsage(tokens tokenUsagePayload) usage.TokenUsage { + nonCached := uint64(0) + if tokens.CachedInputTokens <= tokens.InputTokens { + nonCached = tokens.InputTokens - tokens.CachedInputTokens + } + return usage.TokenUsage{ + InputTokens: nonCached, + CacheReadInputTokens: tokens.CachedInputTokens, + OutputTokens: tokens.OutputTokens, + ReasoningOutputTokens: tokens.ReasoningOutputTokens, + TotalTokens: tokens.TotalTokens, + } +} + +// diffTokenUsage reports the counter movement between two cumulative +// snapshots. ok is false when any field went backwards — a counter reset — +// and the caller falls back to last_token_usage. +func diffTokenUsage(current, previous tokenUsagePayload) (tokenUsagePayload, bool) { + if previous.InputTokens > current.InputTokens || + previous.CachedInputTokens > current.CachedInputTokens || + previous.OutputTokens > current.OutputTokens || + previous.ReasoningOutputTokens > current.ReasoningOutputTokens || + previous.TotalTokens > current.TotalTokens { + return tokenUsagePayload{}, false + } + return tokenUsagePayload{ + InputTokens: current.InputTokens - previous.InputTokens, + CachedInputTokens: current.CachedInputTokens - previous.CachedInputTokens, + OutputTokens: current.OutputTokens - previous.OutputTokens, + ReasoningOutputTokens: current.ReasoningOutputTokens - previous.ReasoningOutputTokens, + TotalTokens: current.TotalTokens - previous.TotalTokens, + }, true +} + +// handleResponseItem tracks file-modifying tool calls. A patch is parsed from +// the call, held until its output confirms success, then folded into the next +// token_count entry — the token event for the turn follows its tool calls. +func handleResponseItem(raw json.RawMessage, state *fileState) { + var item struct { + Type string `json:"type"` + CallID string `json:"call_id"` + Name string `json:"name"` + Input string `json:"input"` + Arguments string `json:"arguments"` + Output json.RawMessage `json:"output"` + } + if err := json.Unmarshal(raw, &item); err != nil || item.CallID == "" { + return + } + + switch item.Type { + case "custom_tool_call", "function_call": + patches := patchesFromCall(item.Input, item.Arguments, state.projectPath) + if len(patches) > 0 { + if state.awaiting == nil { + state.awaiting = make(map[string][]patchFile) + } + state.awaiting[item.CallID] = patches + } + case "custom_tool_call_output", "function_call_output": + patches, ok := state.awaiting[item.CallID] + if !ok { + return + } + delete(state.awaiting, item.CallID) + if outputSucceeded(item.Output) { + state.pending = append(state.pending, patches...) + } + } +} + +// patchesFromCall finds the apply_patch envelope in a tool call: directly in +// input, escaped inside a JSON-encoded input, or inside a shell heredoc in +// the JSON-encoded arguments. +func patchesFromCall(input, arguments, cwd string) []patchFile { + if strings.Contains(input, patchBegin) { + if patches := parsePatchEnvelope(input, cwd); len(patches) > 0 { + return patches + } + var decoded any + if err := json.Unmarshal([]byte(input), &decoded); err == nil { + if text := findPatchString(decoded); text != "" { + return parsePatchEnvelope(text, cwd) + } + } + return nil + } + if !strings.Contains(arguments, patchBegin) { + return nil + } + var decoded any + if err := json.Unmarshal([]byte(arguments), &decoded); err != nil { + return nil + } + if text := findPatchString(decoded); text != "" { + return parsePatchEnvelope(text, cwd) + } + return nil +} + +func findPatchString(value any) string { + switch typed := value.(type) { + case string: + if strings.Contains(typed, patchBegin) { + return typed + } + case []any: + for _, child := range typed { + if text := findPatchString(child); text != "" { + return text + } + } + case map[string]any: + for _, child := range typed { + if text := findPatchString(child); text != "" { + return text + } + } + } + return "" +} + +const patchBegin = "*** Begin Patch" + +// parsePatchEnvelope counts added and removed lines per file in codex's +// apply_patch format: "*** Add File: p", "*** Update File: p" and +// "*** Delete File: p" sections whose body lines carry +/- prefixes. +// Relative paths resolve against cwd; "*** Move to:" renames the section's +// file so the change lands on the destination path. +func parsePatchEnvelope(text, cwd string) []patchFile { + start := strings.Index(text, patchBegin) + if start < 0 { + return nil + } + files := make([]patchFile, 0) + var current *patchFile + for _, line := range strings.Split(text[start:], "\n") { + switch { + case strings.HasPrefix(line, "*** Add File: "), + strings.HasPrefix(line, "*** Update File: "), + strings.HasPrefix(line, "*** Delete File: "): + path := resolvePatchPath(cwd, line[strings.Index(line, ": ")+2:]) + if path == "" { + current = nil + continue + } + files = append(files, patchFile{path: path}) + current = &files[len(files)-1] + case strings.HasPrefix(line, "*** Move to: "): + if current != nil { + if moved := resolvePatchPath(cwd, strings.TrimPrefix(line, "*** Move to: ")); moved != "" { + current.path = moved + } + } + case strings.HasPrefix(line, "*** End Patch"): + current = nil + case strings.HasPrefix(line, "***"): + // Other directives keep the current file. + case current == nil: + case strings.HasPrefix(line, "+"): + current.added++ + case strings.HasPrefix(line, "-"): + current.removed++ + } + } + return files +} + +func resolvePatchPath(cwd, path string) string { + path = strings.TrimSpace(path) + if path == "" || filepath.IsAbs(path) || strings.TrimSpace(cwd) == "" { + return path + } + return filepath.Join(cwd, path) +} + +// outputSucceeded reports whether a tool output confirms the patch applied. +// The output is either a plain string or an object with output/exit_code; +// apply_patch prints "Success." and shell wrappers prepend "Exit code: 0". +func outputSucceeded(raw json.RawMessage) bool { + if len(raw) == 0 { + return false + } + var text string + if err := json.Unmarshal(raw, &text); err == nil { + return patchOutputOK(text) + } + var items []struct { + Text string `json:"text"` + } + if err := json.Unmarshal(raw, &items); err == nil { + var joined strings.Builder + for _, item := range items { + joined.WriteString(item.Text) + joined.WriteByte('\n') + } + return patchOutputOK(joined.String()) + } + var structured struct { + Output string `json:"output"` + Metadata struct { + ExitCode *int `json:"exit_code"` + } `json:"metadata"` + } + if err := json.Unmarshal(raw, &structured); err != nil { + return false + } + if structured.Metadata.ExitCode != nil { + return *structured.Metadata.ExitCode == 0 + } + return patchOutputOK(structured.Output) +} + +func patchOutputOK(text string) bool { + return strings.Contains(text, "Success.") || strings.HasPrefix(text, "Exit code: 0") +} + +// applyPatches folds confirmed patches into the entry: totals, per-file +// breakdown, and the entity as the most-changed file. +func applyPatches(entry *usage.Entry, patches []patchFile) { + if len(patches) == 0 { + return + } + isWrite := true + entry.IsWrite = &isWrite + for _, patch := range patches { + entry.LinesAdded += patch.added + entry.LinesRemoved += patch.removed + index := -1 + for i := range entry.Files { + if entry.Files[i].Path == patch.path { + index = i + break + } + } + if index < 0 { + entry.Files = append(entry.Files, usage.FileChange{Path: patch.path}) + index = len(entry.Files) - 1 + } + entry.Files[index].LinesAdded += patch.added + entry.Files[index].LinesRemoved += patch.removed + } + best := 0 + for i := range entry.Files { + if entry.Files[i].LinesAdded+entry.Files[i].LinesRemoved > entry.Files[best].LinesAdded+entry.Files[best].LinesRemoved { + best = i + } + } + entry.Entity = entry.Files[best].Path + entry.EntityType = "file" +} + func languageFromPayload(raw json.RawMessage) string { if len(raw) == 0 { return langdetect.Unknown @@ -328,15 +622,17 @@ func candidatesFromText(text string, weight int) []langdetect.Candidate { return candidates } -// StableEntryID computes the deterministic id of a codex usage entry. The -// source file is keyed by its basename: archiving a session moves the file -// from sessions/ into archived_sessions/ without renaming it, and the id must -// survive that move or every archived session double-counts. +// StableEntryID computes the deterministic id of a codex usage entry from the +// event itself — session, time, model and token counts — never from where the +// event is stored. Storage details (path, filename, line number) change when +// codex archives or rewrites a rollout file, and every id built on them +// eventually double-counts. Two token events in one session sharing the same +// millisecond and identical counts are the same event; collapsing them is +// correct. func StableEntryID(entry usage.Entry) string { return usage.StableID( string(usage.ProviderCodex), - filepath.Base(entry.SourceFile), - strconv.Itoa(entry.SourceLine), + entry.SessionID, entry.Timestamp.Format(time.RFC3339Nano), entry.Model, strconv.FormatUint(entry.Usage.InputTokens, 10), diff --git a/internal/codexusage/loader_test.go b/internal/codexusage/loader_test.go index 9fa0584..a441360 100644 --- a/internal/codexusage/loader_test.go +++ b/internal/codexusage/loader_test.go @@ -40,20 +40,74 @@ func TestReadUsageFileParsesTokenCountEvents(t *testing.T) { if entry.Language != "Unknown" { t.Fatalf("language = %q, want Unknown", entry.Language) } - // input_tokens (40) is the full prompt incl. cache; we report non-cached - // input (40 - 8 = 32) and move the cached portion to cache read, matching - // ccusage's codex token accounting. - if entry.Usage.InputTokens != 32 { - t.Fatalf("input tokens = %d, want non-cached input (40-8)", entry.Usage.InputTokens) + // Usage comes from the cumulative counter (first event: the counter + // itself), with input_tokens split into non-cached (100-20) and cache + // read (20). last_token_usage only feeds the id. + if entry.Usage.InputTokens != 80 { + t.Fatalf("input tokens = %d, want non-cached input (100-20)", entry.Usage.InputTokens) } - if entry.Usage.CacheReadInputTokens != 8 { - t.Fatalf("cache read tokens = %d, want 8 (cached portion)", entry.Usage.CacheReadInputTokens) + if entry.Usage.CacheReadInputTokens != 20 { + t.Fatalf("cache read tokens = %d, want 20 (cached portion)", entry.Usage.CacheReadInputTokens) } - if entry.Usage.ReasoningOutputTokens != 2 { - t.Fatalf("reasoning output tokens = %d, want 2", entry.Usage.ReasoningOutputTokens) + if entry.Usage.ReasoningOutputTokens != 3 { + t.Fatalf("reasoning output tokens = %d, want 3", entry.Usage.ReasoningOutputTokens) } - if entry.Usage.TotalTokens != 45 { - t.Fatalf("total tokens = %d, want 45", entry.Usage.TotalTokens) + if entry.Usage.TotalTokens != 110 { + t.Fatalf("total tokens = %d, want 110", entry.Usage.TotalTokens) + } +} + +func TestReadUsageFileUsesCumulativeDeltas(t *testing.T) { + content := `{"timestamp":"2026-06-04T01:02:03Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo/app"}} +{"timestamp":"2026-06-04T01:02:04Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":15},"total_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":15}}}} +{"timestamp":"2026-06-04T01:02:05Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":15},"total_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":15}}}} +{"timestamp":"2026-06-04T01:02:06Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":25,"cached_input_tokens":4,"output_tokens":12,"reasoning_output_tokens":1,"total_tokens":37},"total_token_usage":{"input_tokens":35,"cached_input_tokens":4,"output_tokens":17,"reasoning_output_tokens":1,"total_tokens":52}}}} +` + path := filepath.Join(t.TempDir(), "sessions", "rollout-x.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, content) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + // The second event is a duplicate emission (counter unchanged) and must + // vanish; the third is the counter delta, not its last_token_usage. + if len(entries) != 2 { + t.Fatalf("entries = %d, want 2 (duplicate skipped)", len(entries)) + } + if entries[0].Usage.TotalTokens != 15 { + t.Fatalf("first total = %d, want 15", entries[0].Usage.TotalTokens) + } + second := entries[1].Usage + if second.InputTokens != 21 || second.CacheReadInputTokens != 4 || second.OutputTokens != 12 || second.TotalTokens != 37 { + t.Fatalf("second usage = %+v, want delta 21/4/12/37", second) + } + if entries[0].ID == entries[1].ID { + t.Fatal("distinct events share an id") + } +} + +func TestReadUsageFileFallsBackToLastUsageOnCounterReset(t *testing.T) { + content := `{"timestamp":"2026-06-04T01:02:03Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo/app"}} +{"timestamp":"2026-06-04T01:02:04Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":100,"cached_input_tokens":0,"output_tokens":50,"reasoning_output_tokens":0,"total_tokens":150},"total_token_usage":{"input_tokens":100,"cached_input_tokens":0,"output_tokens":50,"reasoning_output_tokens":0,"total_tokens":150}}}} +{"timestamp":"2026-06-04T01:02:05Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":8,"cached_input_tokens":0,"output_tokens":3,"reasoning_output_tokens":0,"total_tokens":11},"total_token_usage":{"input_tokens":8,"cached_input_tokens":0,"output_tokens":3,"reasoning_output_tokens":0,"total_tokens":11}}}} +` + path := filepath.Join(t.TempDir(), "sessions", "rollout-x.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, content) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("entries = %d, want 2", len(entries)) + } + // The counter went backwards (reset): the event keeps its own + // last_token_usage instead of a bogus delta. + if entries[1].Usage.TotalTokens != 11 { + t.Fatalf("post-reset total = %d, want 11", entries[1].Usage.TotalTokens) } } @@ -147,30 +201,161 @@ func writeFile(t *testing.T, path, data string) { } } -func TestStableEntryIDSurvivesArchiveMove(t *testing.T) { +func TestStableEntryIDIndependentOfFileLocation(t *testing.T) { + // The id comes from the event (session + time + tokens), never from + // storage: archiving moves the file, and a rename must not matter either. content := `{"timestamp":"2026-06-04T01:02:03Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo/app"}} {"timestamp":"2026-06-04T01:02:04Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":4,"output_tokens":5,"reasoning_output_tokens":2,"total_tokens":15}}}} ` dir := t.TempDir() - before := filepath.Join(dir, "sessions", "2026", "06", "04", "rollout-x.jsonl") - after := filepath.Join(dir, "archived_sessions", "rollout-x.jsonl") - mkdirAll(t, filepath.Dir(before)) - mkdirAll(t, filepath.Dir(after)) - writeFile(t, before, content) - writeFile(t, after, content) - - beforeEntries, err := ReadUsageFile(before) + paths := []string{ + filepath.Join(dir, "sessions", "2026", "06", "04", "rollout-x.jsonl"), + filepath.Join(dir, "archived_sessions", "rollout-x.jsonl"), + filepath.Join(dir, "archived_sessions", "renamed-y.jsonl"), + } + ids := make([]string, 0, len(paths)) + for _, path := range paths { + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, content) + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries(%s) = %d, want 1", path, len(entries)) + } + ids = append(ids, entries[0].ID) + } + if ids[0] != ids[1] || ids[0] != ids[2] { + t.Fatalf("id depends on file location: %v", ids) + } +} + +func TestReadUsageFileAttributesConfirmedPatches(t *testing.T) { + content := `{"timestamp":"2026-06-04T01:02:03Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo/app"}} +{"timestamp":"2026-06-04T01:02:04Z","type":"response_item","payload":{"type":"custom_tool_call","call_id":"c1","name":"apply_patch","input":"*** Begin Patch\n*** Update File: /repo/app/main.go\n@@\n-old line\n+new line\n+extra line\n*** Add File: /repo/app/new.go\n+package app\n*** End Patch"}} +{"timestamp":"2026-06-04T01:02:05Z","type":"response_item","payload":{"type":"custom_tool_call_output","call_id":"c1","output":"Exit code: 0\nOutput:\nSuccess. Updated the following files:\nM /repo/app/main.go\nA /repo/app/new.go\n"}} +{"timestamp":"2026-06-04T01:02:06Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":15}}}} +{"timestamp":"2026-06-04T01:02:07Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":20,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":25}}}} +` + path := filepath.Join(t.TempDir(), "sessions", "rollout-x.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, content) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("entries = %d, want 2", len(entries)) + } + + first := entries[0] + if first.LinesAdded != 3 || first.LinesRemoved != 1 { + t.Fatalf("lines = +%d/-%d, want +3/-1", first.LinesAdded, first.LinesRemoved) + } + if first.Entity != "/repo/app/main.go" || first.EntityType != "file" { + t.Fatalf("entity = %q/%q, want most-changed main.go/file", first.Entity, first.EntityType) + } + if first.IsWrite == nil || !*first.IsWrite { + t.Fatal("isWrite not set") + } + if len(first.Files) != 2 { + t.Fatalf("files = %+v, want main.go and new.go", first.Files) + } + + second := entries[1] + if second.IsWrite != nil || len(second.Files) != 0 { + t.Fatalf("second entry inherited patches: %+v", second) + } +} + +func TestReadUsageFileIgnoresFailedPatches(t *testing.T) { + content := `{"timestamp":"2026-06-04T01:02:03Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo/app"}} +{"timestamp":"2026-06-04T01:02:04Z","type":"response_item","payload":{"type":"custom_tool_call","call_id":"c1","name":"apply_patch","input":"*** Begin Patch\n*** Update File: /repo/app/main.go\n+x\n*** End Patch"}} +{"timestamp":"2026-06-04T01:02:05Z","type":"response_item","payload":{"type":"custom_tool_call_output","call_id":"c1","output":"Exit code: 1\napply_patch: context mismatch"}} +{"timestamp":"2026-06-04T01:02:06Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":15}}}} +` + path := filepath.Join(t.TempDir(), "sessions", "rollout-x.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, content) + + entries, err := ReadUsageFile(path) if err != nil { t.Fatal(err) } - afterEntries, err := ReadUsageFile(after) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].IsWrite != nil || entries[0].LinesAdded != 0 { + t.Fatalf("failed patch was counted: %+v", entries[0]) + } +} + +func TestReadUsageFileParsesHeredocPatchInShellCall(t *testing.T) { + content := `{"timestamp":"2026-06-04T01:02:03Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo/app"}} +{"timestamp":"2026-06-04T01:02:04Z","type":"response_item","payload":{"type":"function_call","call_id":"c1","name":"exec_command","arguments":"{\"command\":[\"bash\",\"-lc\",\"apply_patch <<'EOF'\\n*** Begin Patch\\n*** Delete File: /repo/app/dead.go\\n*** End Patch\\nEOF\"]}"}} +{"timestamp":"2026-06-04T01:02:05Z","type":"response_item","payload":{"type":"function_call_output","call_id":"c1","output":"Exit code: 0\nSuccess. Updated the following files:\nD /repo/app/dead.go\n"}} +{"timestamp":"2026-06-04T01:02:06Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":15}}}} +` + path := filepath.Join(t.TempDir(), "sessions", "rollout-x.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, content) + + entries, err := ReadUsageFile(path) if err != nil { t.Fatal(err) } - if len(beforeEntries) != 1 || len(afterEntries) != 1 { - t.Fatalf("entries = %d/%d, want 1/1", len(beforeEntries), len(afterEntries)) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].IsWrite == nil || len(entries[0].Files) != 1 || entries[0].Files[0].Path != "/repo/app/dead.go" { + t.Fatalf("heredoc patch not captured: %+v", entries[0]) + } +} + +func TestReadUsageFileConfirmsPatchViaPatchApplyEnd(t *testing.T) { + content := `{"timestamp":"2026-06-04T01:02:03Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo/app"}} +{"timestamp":"2026-06-04T01:02:04Z","type":"response_item","payload":{"type":"custom_tool_call","call_id":"c1","name":"apply_patch","input":"*** Begin Patch\n*** Update File: src/main.go\n+x\n*** Move to: src/renamed.go\n*** End Patch"}} +{"timestamp":"2026-06-04T01:02:05Z","type":"event_msg","payload":{"type":"patch_apply_end","call_id":"c1","success":true,"stdout":"Success."}} +{"timestamp":"2026-06-04T01:02:06Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":15}}}} +` + path := filepath.Join(t.TempDir(), "sessions", "rollout-x.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, content) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + entry := entries[0] + if entry.IsWrite == nil || entry.LinesAdded != 1 { + t.Fatalf("patch_apply_end confirmation not applied: %+v", entry) + } + // Relative path resolved against cwd, and Move to wins as the final path. + if len(entry.Files) != 1 || entry.Files[0].Path != "/repo/app/src/renamed.go" { + t.Fatalf("files = %+v, want /repo/app/src/renamed.go", entry.Files) + } +} + +func TestReadUsageFileRejectsPatchApplyEndFailure(t *testing.T) { + content := `{"timestamp":"2026-06-04T01:02:03Z","type":"session_meta","payload":{"id":"session-1","cwd":"/repo/app"}} +{"timestamp":"2026-06-04T01:02:04Z","type":"response_item","payload":{"type":"custom_tool_call","call_id":"c1","name":"apply_patch","input":"*** Begin Patch\n*** Update File: src/main.go\n+x\n*** End Patch"}} +{"timestamp":"2026-06-04T01:02:05Z","type":"event_msg","payload":{"type":"patch_apply_end","call_id":"c1","success":false,"stderr":"invalid patch"}} +{"timestamp":"2026-06-04T01:02:06Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"cached_input_tokens":0,"output_tokens":5,"reasoning_output_tokens":0,"total_tokens":15}}}} +` + path := filepath.Join(t.TempDir(), "sessions", "rollout-x.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, content) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) } - if beforeEntries[0].ID != afterEntries[0].ID { - t.Fatalf("id changed across archive move: %q vs %q", beforeEntries[0].ID, afterEntries[0].ID) + if len(entries) != 1 || entries[0].IsWrite != nil { + t.Fatalf("failed patch counted: %+v", entries[0]) } } diff --git a/internal/usage/usage.go b/internal/usage/usage.go index 7b8de83..5696645 100644 --- a/internal/usage/usage.go +++ b/internal/usage/usage.go @@ -59,6 +59,13 @@ type TokenUsage struct { TotalTokens uint64 `json:"total_tokens"` } +// FileChange records the diff one event applied to a single file. +type FileChange struct { + Path string `json:"path"` + LinesAdded uint64 `json:"lines_added,omitempty"` + LinesRemoved uint64 `json:"lines_removed,omitempty"` +} + type Entry struct { Provider Provider `json:"provider"` ID string `json:"id,omitempty"` @@ -88,6 +95,14 @@ type Entry struct { Editor string `json:"editor,omitempty"` Category string `json:"category,omitempty"` IsWrite *bool `json:"is_write,omitempty"` + // LinesAdded/LinesRemoved count the source lines the agent added and + // removed in this event's file modifications, for providers that record + // diffs. Zero means "no diff recorded", not "no change". + LinesAdded uint64 `json:"lines_added,omitempty"` + LinesRemoved uint64 `json:"lines_removed,omitempty"` + // Files breaks the same modifications down per file. Entity is always + // the most-changed path in here; LinesAdded/LinesRemoved are the totals. + Files []FileChange `json:"files,omitempty"` Raw map[string]any `json:"raw,omitempty"` Usage TokenUsage `json:"usage"` } diff --git a/internal/usagedb/db.go b/internal/usagedb/db.go index 9b6548f..a25400e 100644 --- a/internal/usagedb/db.go +++ b/internal/usagedb/db.go @@ -67,23 +67,48 @@ func Open(path string) (*DB, error) { return nil, fmt.Errorf("open usage db: %w", err) } db.SetMaxOpenConns(1) + fresh, err := isFreshDatabase(db) + if err != nil { + _ = db.Close() + return nil, fmt.Errorf("inspect usage db: %w", err) + } if _, err := db.Exec(schema); err != nil { _ = db.Close() return nil, fmt.Errorf("migrate usage db: %w", err) } - if err := migrate(db); err != nil { + // A fresh database is born at the current version: the schema constant + // already describes the final shape, so the migration chain only ever + // runs against databases created by an older binary. + if fresh { + err = stampVersion(db) + } else { + err = migrate(db) + } + if err != nil { _ = db.Close() return nil, fmt.Errorf("migrate usage db: %w", err) } return &DB{db: db}, nil } +func isFreshDatabase(db *sql.DB) (bool, error) { + var count int + err := db.QueryRow(`SELECT count(*) FROM sqlite_master WHERE type = 'table' AND name = 'usage_events'`).Scan(&count) + return count == 0, err +} + +func stampVersion(db *sql.DB) error { + _, err := db.Exec(fmt.Sprintf(`PRAGMA user_version = %d`, eventSchemaVersion)) + return err +} + // eventSchemaVersion tracks one-time data repairs, recorded in the SQLite -// user_version pragma. Version 1 rekeys codex events whose ids hashed the -// source file's full path: archiving a session moves its file from sessions/ -// into archived_sessions/, which changed every id and double-counted the -// whole session. -const eventSchemaVersion = 1 +// user_version pragma. Versions 1 and 2 both rekey codex events: v1 dropped +// the source file's full path from the id (archiving moves the file), v2 +// dropped file position entirely in favor of session + timestamp + tokens. +// The rekey recomputes from payload, so any older version jumps straight to +// the current scheme in one pass. +const eventSchemaVersion = 2 func migrate(db *sql.DB) error { var version int @@ -93,13 +118,12 @@ func migrate(db *sql.DB) error { if version >= eventSchemaVersion { return nil } - if version < 1 { + if version < 2 { if err := rekeyCodexEvents(db); err != nil { return err } } - _, err := db.Exec(fmt.Sprintf(`PRAGMA user_version = %d`, eventSchemaVersion)) - return err + return stampVersion(db) } func rekeyCodexEvents(db *sql.DB) error { diff --git a/internal/usagedb/db_test.go b/internal/usagedb/db_test.go index 6143d1a..e54d31c 100644 --- a/internal/usagedb/db_test.go +++ b/internal/usagedb/db_test.go @@ -211,7 +211,7 @@ func testUsageEntry(id string) usage.Entry { } } -func TestOpenRekeysCodexEventsToBasenameIDs(t *testing.T) { +func TestOpenRekeysCodexEventsToSemanticIDs(t *testing.T) { path := filepath.Join(t.TempDir(), "usage.db") db, err := Open(path) if err != nil { @@ -297,3 +297,15 @@ func TestOpenRekeysCodexEventsToBasenameIDs(t *testing.T) { t.Fatalf("pruned = %d, want 1", pruned) } } + +func TestOpenStampsFreshDatabaseAtCurrentVersion(t *testing.T) { + db := openTestDB(t) + + var version int + if err := db.db.QueryRow(`PRAGMA user_version`).Scan(&version); err != nil { + t.Fatal(err) + } + if version != eventSchemaVersion { + t.Fatalf("user_version = %d, want %d", version, eventSchemaVersion) + } +} diff --git a/internal/usageupload/uploader.go b/internal/usageupload/uploader.go index f93b913..3fc40d2 100644 --- a/internal/usageupload/uploader.go +++ b/internal/usageupload/uploader.go @@ -12,6 +12,7 @@ import ( "io" "net/http" "os" + "path/filepath" "runtime" "strings" "time" @@ -54,12 +55,12 @@ type DevicePayload struct { } type Event struct { - ID string `json:"id"` - Provider string `json:"provider"` - SourceType string `json:"source_type,omitempty"` - SourceProvider string `json:"source_provider,omitempty"` - EventKind string `json:"event_kind,omitempty"` - Timestamp string `json:"timestamp"` + ID string `json:"id"` + Provider string `json:"provider"` + SourceType string `json:"source_type,omitempty"` + SourceProvider string `json:"source_provider,omitempty"` + EventKind string `json:"event_kind,omitempty"` + Timestamp string `json:"timestamp"` // The machine's IANA zone ("Asia/Tokyo"), omitted when it cannot be // resolved — see usage.MachineTimezone. Never a fixed abbreviation like // "JST": those are ambiguous across regions and cannot be re-expanded. @@ -70,28 +71,31 @@ type Event struct { // Australia/Lord_Howe — is a pure function of the two. Sending it as well // would be a second copy of a derived value, and the copy is what goes // stale when tzdata is corrected. - Timezone string `json:"timezone,omitempty"` - SessionID string `json:"session_id,omitempty"` - Project string `json:"project"` - ProjectPathHash string `json:"project_path_hash,omitempty"` - Model string `json:"model,omitempty"` - Language string `json:"language"` - OS string `json:"os,omitempty"` - Client string `json:"client,omitempty"` - Entity string `json:"entity,omitempty"` - EntityType string `json:"entity_type,omitempty"` - Branch string `json:"branch,omitempty"` - Editor string `json:"editor,omitempty"` - Category string `json:"category,omitempty"` - IsWrite *bool `json:"is_write,omitempty"` - Raw map[string]any `json:"raw,omitempty"` - InputTokens uint64 `json:"input_tokens,omitempty"` - OutputTokens uint64 `json:"output_tokens,omitempty"` - CachedInputTokens uint64 `json:"cached_input_tokens,omitempty"` - CacheCreationInputTokens uint64 `json:"cache_creation_input_tokens,omitempty"` - CacheReadInputTokens uint64 `json:"cache_read_input_tokens,omitempty"` - ReasoningOutputTokens uint64 `json:"reasoning_output_tokens,omitempty"` - TotalTokens uint64 `json:"total_tokens,omitempty"` + Timezone string `json:"timezone,omitempty"` + SessionID string `json:"session_id,omitempty"` + Project string `json:"project"` + ProjectPathHash string `json:"project_path_hash,omitempty"` + Model string `json:"model,omitempty"` + Language string `json:"language"` + OS string `json:"os,omitempty"` + Client string `json:"client,omitempty"` + Entity string `json:"entity,omitempty"` + EntityType string `json:"entity_type,omitempty"` + Branch string `json:"branch,omitempty"` + Editor string `json:"editor,omitempty"` + Category string `json:"category,omitempty"` + IsWrite *bool `json:"is_write,omitempty"` + LinesAdded uint64 `json:"lines_added,omitempty"` + LinesRemoved uint64 `json:"lines_removed,omitempty"` + Files []usage.FileChange `json:"files,omitempty"` + Raw map[string]any `json:"raw,omitempty"` + InputTokens uint64 `json:"input_tokens,omitempty"` + OutputTokens uint64 `json:"output_tokens,omitempty"` + CachedInputTokens uint64 `json:"cached_input_tokens,omitempty"` + CacheCreationInputTokens uint64 `json:"cache_creation_input_tokens,omitempty"` + CacheReadInputTokens uint64 `json:"cache_read_input_tokens,omitempty"` + ReasoningOutputTokens uint64 `json:"reasoning_output_tokens,omitempty"` + TotalTokens uint64 `json:"total_tokens,omitempty"` } type Response struct { @@ -295,12 +299,15 @@ func convertEvent(entry usage.Entry, zoneName string) Event { Language: usage.NormalizeLanguage(entry.Language), OS: entry.OS, Client: entry.Client, - Entity: entry.Entity, + Entity: relativeEntity(entry.ProjectPath, entry.Entity), EntityType: entry.EntityType, Branch: entry.Branch, Editor: entry.Editor, Category: entry.Category, IsWrite: entry.IsWrite, + LinesAdded: entry.LinesAdded, + LinesRemoved: entry.LinesRemoved, + Files: relativeFiles(entry.ProjectPath, entry.Files), Raw: entry.Raw, InputTokens: entry.Usage.InputTokens, OutputTokens: entry.Usage.OutputTokens, @@ -330,3 +337,33 @@ func hashProjectPath(path string) string { sum := sha256.Sum256([]byte(path)) return hex.EncodeToString(sum[:]) } + +// relativeEntity strips the local filesystem prefix from an entity path for +// the same reason the project path is uploaded as a hash: the server sees the +// file's place inside the project, never the machine's directory layout. +func relativeEntity(projectPath, entity string) string { + entity = strings.TrimSpace(entity) + if entity == "" { + return "" + } + if projectPath = strings.TrimSpace(projectPath); projectPath != "" { + if rel, err := filepath.Rel(projectPath, entity); err == nil && rel != ".." && !strings.HasPrefix(rel, ".."+string(filepath.Separator)) { + return filepath.ToSlash(rel) + } + } + return filepath.Base(entity) +} + +// relativeFiles applies the same machine-layout stripping to the per-file +// breakdown that relativeEntity applies to the entity. +func relativeFiles(projectPath string, files []usage.FileChange) []usage.FileChange { + if len(files) == 0 { + return nil + } + out := make([]usage.FileChange, len(files)) + for i, file := range files { + out[i] = file + out[i].Path = relativeEntity(projectPath, file.Path) + } + return out +} diff --git a/internal/usageupload/uploader_test.go b/internal/usageupload/uploader_test.go index e33b719..ea2a27c 100644 --- a/internal/usageupload/uploader_test.go +++ b/internal/usageupload/uploader_test.go @@ -73,3 +73,38 @@ func TestUploadUsesBaseURLEnvironment(t *testing.T) { t.Fatalf("response = %+v, want accepted event", resp) } } + +func TestRelativeEntityHidesMachineLayout(t *testing.T) { + tests := []struct { + projectPath string + entity string + want string + }{ + {"/Users/me/repo", "/Users/me/repo/pkg/a.go", "pkg/a.go"}, + {"/Users/me/repo", "/Users/me/repo/a.go", "a.go"}, + {"/Users/me/repo", "/Users/me/elsewhere/b.go", "b.go"}, + {"", "/Users/me/repo/c.go", "c.go"}, + {"/Users/me/repo", "", ""}, + } + for _, tt := range tests { + if got := relativeEntity(tt.projectPath, tt.entity); got != tt.want { + t.Fatalf("relativeEntity(%q, %q) = %q, want %q", tt.projectPath, tt.entity, got, tt.want) + } + } +} + +func TestRelativeFilesStripsMachineLayout(t *testing.T) { + files := relativeFiles("/Users/me/repo", []usage.FileChange{ + {Path: "/Users/me/repo/pkg/a.go", LinesAdded: 2}, + {Path: "/Users/me/other/b.go", LinesRemoved: 1}, + }) + if files[0].Path != "pkg/a.go" || files[0].LinesAdded != 2 { + t.Fatalf("files[0] = %+v, want pkg/a.go +2", files[0]) + } + if files[1].Path != "b.go" || files[1].LinesRemoved != 1 { + t.Fatalf("files[1] = %+v, want b.go -1", files[1]) + } + if relativeFiles("/p", nil) != nil { + t.Fatal("relativeFiles(nil) should be nil") + } +} From e30e6dd0a0ed8630de4478764ac490dbc3d801bb Mon Sep 17 00:00:00 2001 From: akarineren Date: Fri, 31 Jul 2026 20:36:34 +0900 Subject: [PATCH 03/11] Refactor OpenCode usage tracking and enhance project path resolution - Updated OpenCode message structure to include working directory path. - Refactored line counting logic to a dedicated function in the usage package. - Improved project path extraction from CWD, ensuring it handles various edge cases. - Enhanced OpenCode provider directory resolution for cross-platform compatibility. - Added comprehensive tests for OpenCode usage scenarios, including token billing and file change tracking. - Introduced new test cases to validate behavior with user messages and tool interactions. --- internal/agentusage/opencode.go | 527 +++++++++++++++++++++----- internal/agentusage/opencode_test.go | 388 +++++++++++++++++++ internal/agentusage/providers_test.go | 4 +- internal/claudeusage/loader.go | 31 +- internal/claudeusage/loader_test.go | 4 +- internal/usage/usage.go | 82 ++++ pkg/agentlib/agentlib.go | 7 +- 7 files changed, 922 insertions(+), 121 deletions(-) create mode 100644 internal/agentusage/opencode_test.go diff --git a/internal/agentusage/opencode.go b/internal/agentusage/opencode.go index a0c2b4b..00c6d4a 100644 --- a/internal/agentusage/opencode.go +++ b/internal/agentusage/opencode.go @@ -2,6 +2,7 @@ package agentusage import ( "bytes" + "database/sql" "encoding/json" "os" "path/filepath" @@ -12,6 +13,23 @@ import ( "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) +// OpenCode stores everything in one SQLite database: sessions own messages, +// messages own parts. Tokens live on the message, the working directory lives +// on both the session and the message, and file edits live in the parts. All +// three tables are read together because an entry needs all three. +type openCodeSession struct { + directory string + version string +} + +type openCodeMessage struct { + id string + sessionID string + created int64 + data map[string]any + parts []map[string]any +} + func loadOpenCodeEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { entriesByID := make(map[string]usage.Entry) for _, root := range paths { @@ -46,19 +64,14 @@ func loadOpenCodeRoot(root string, filter usage.FileFilter) ([]usage.Entry, erro case ".db": return loadOpenCodeDatabase(root) case ".json": - if entry, ok, err := parseOpenCodeMessageFile(root, "", ""); err != nil || ok { - if !ok { - return nil, err - } - return []usage.Entry{entry}, err - } + return loadOpenCodeLegacyFile(root) } return nil, nil } entries := make([]usage.Entry, 0) seenIDs := make(map[string]bool) - if dbPath := openCodeDBPath(root); dbPath != "" { + for _, dbPath := range openCodeDBPaths(root) { dbEntries, err := loadOpenCodeDatabase(dbPath) if err != nil { return nil, err @@ -71,6 +84,8 @@ func loadOpenCodeRoot(root string, filter usage.FileFilter) ([]usage.Entry, erro } } + // Older OpenCode releases wrote one JSON file per message instead of the + // database. Their users keep those files, so keep reading them. files := collectExt(filepath.Join(root, "storage", "message"), ".json") sort.Strings(files) files = filterFiles(files, filter) @@ -79,144 +94,482 @@ func loadOpenCodeRoot(root string, filter usage.FileFilter) ([]usage.Entry, erro if seenIDs[stableOpenCodeMessageID(stem)] { continue } - entry, ok, err := parseOpenCodeMessageFile(file, "", "") + legacy, err := loadOpenCodeLegacyFile(file) if err != nil { return nil, err } - if !ok { - continue - } - if entry.ID != "" && seenIDs[entry.ID] { - continue - } - if entry.ID != "" { - seenIDs[entry.ID] = true + for _, entry := range legacy { + if entry.ID != "" && seenIDs[entry.ID] { + continue + } + if entry.ID != "" { + seenIDs[entry.ID] = true + } + entries = append(entries, entry) } - entries = append(entries, entry) } sortEntries(entries) return entries, nil } -func openCodeDBPath(root string) string { - candidate := filepath.Join(root, "opencode.db") - if existingSQLiteFile(candidate) { - return candidate - } - entries, err := os.ReadDir(root) +// openCodeDBPaths lists the databases in a data root. OpenCode ships release +// channels side by side ("opencode.db", "opencode-nightly.db"), each a full +// database of its own. +func openCodeDBPaths(root string) []string { + matches, err := filepath.Glob(filepath.Join(root, "opencode*.db")) if err != nil { - return "" + return nil } - candidates := make([]string, 0) - for _, entry := range entries { - name := entry.Name() - if entry.IsDir() || !isOpenCodeChannelDB(name) { - continue + paths := make([]string, 0, len(matches)) + for _, candidate := range matches { + if existingSQLiteFile(candidate) { + paths = append(paths, candidate) } - candidates = append(candidates, filepath.Join(root, name)) - } - sort.Strings(candidates) - if len(candidates) == 0 { - return "" } - return candidates[0] + sort.Strings(paths) + return paths } -func isOpenCodeChannelDB(name string) bool { - if !strings.HasPrefix(name, "opencode-") || !strings.HasSuffix(name, ".db") { - return false +func loadOpenCodeDatabase(path string) ([]usage.Entry, error) { + db, err := openSQLite(path) + if err != nil { + return nil, err + } + defer db.Close() + + sessions := queryOpenCodeSessions(db) + messages := queryOpenCodeMessages(db) + if len(messages) == 0 { + return nil, nil } - channel := strings.TrimSuffix(strings.TrimPrefix(name, "opencode-"), ".db") - if channel == "" { - return false + attachOpenCodeParts(db, messages) + + // Token counts on a message are the running context total, not that turn's + // cost, so each entry bills the growth since the previous message of the + // same session. Messages are ordered by creation time for that subtraction + // to mean anything. + bySession := make(map[string][]*openCodeMessage) + for _, message := range messages { + bySession[message.sessionID] = append(bySession[message.sessionID], message) } - for _, ch := range channel { - if !(ch >= 'a' && ch <= 'z' || ch >= 'A' && ch <= 'Z' || ch >= '0' && ch <= '9' || ch == '_' || ch == '-') { - return false - } + + entries := make([]usage.Entry, 0, len(messages)) + for sessionID, sessionMessages := range bySession { + sort.SliceStable(sessionMessages, func(i, j int) bool { + if sessionMessages[i].created != sessionMessages[j].created { + return sessionMessages[i].created < sessionMessages[j].created + } + return sessionMessages[i].id < sessionMessages[j].id + }) + entries = append(entries, openCodeSessionEntries(path, sessions[sessionID], sessionMessages)...) } - return true + return entries, nil } -func loadOpenCodeDatabase(path string) ([]usage.Entry, error) { - db, err := openSQLite(path) +func queryOpenCodeSessions(db *sql.DB) map[string]openCodeSession { + sessions := make(map[string]openCodeSession) + rows, err := db.Query(`SELECT id, COALESCE(directory, ''), COALESCE(version, '') FROM session`) if err != nil { - return nil, err + return sessions } - defer db.Close() + defer rows.Close() + for rows.Next() { + var id, directory, version string + if err := rows.Scan(&id, &directory, &version); err != nil { + continue + } + sessions[id] = openCodeSession{directory: directory, version: version} + } + return sessions +} - rows, err := db.Query("SELECT id, session_id, data FROM message") +func queryOpenCodeMessages(db *sql.DB) []*openCodeMessage { + rows, err := db.Query(`SELECT id, session_id, time_created, CAST(data AS TEXT) FROM message`) if err != nil { - return nil, nil + return nil } defer rows.Close() - entries := make([]usage.Entry, 0) + messages := make([]*openCodeMessage, 0) for rows.Next() { - var rowID, rowSessionID, data string - if err := rows.Scan(&rowID, &rowSessionID, &data); err != nil { + var id, sessionID, data string + var created int64 + if err := rows.Scan(&id, &sessionID, &created, &data); err != nil { continue } - if entry, ok := openCodeMessageEntry(path, decodeJSONObjectString(data), rowID, rowSessionID, 0); ok { - entries = append(entries, entry) + record := decodeJSONObjectString(data) + if record == nil { + continue } + messages = append(messages, &openCodeMessage{ + id: firstNonEmpty(stringField(record, "id"), id), + sessionID: firstNonEmpty(stringField(record, "sessionID"), sessionID), + created: created, + data: record, + }) } - return entries, rows.Err() + return messages } -func parseOpenCodeMessageFile(path, rowID, rowSessionID string) (usage.Entry, bool, error) { - data, err := os.ReadFile(path) +func attachOpenCodeParts(db *sql.DB, messages []*openCodeMessage) { + byID := make(map[string][]*openCodeMessage, len(messages)) + for _, message := range messages { + byID[message.id] = append(byID[message.id], message) + } + + rows, err := db.Query(`SELECT message_id, CAST(data AS TEXT) FROM part ORDER BY time_created ASC, id ASC`) if err != nil { - return usage.Entry{}, false, err + return } - decoder := json.NewDecoder(bytes.NewReader(data)) - decoder.UseNumber() - var record map[string]any - if err := decoder.Decode(&record); err != nil { - return usage.Entry{}, false, nil + defer rows.Close() + for rows.Next() { + var messageID, data string + if err := rows.Scan(&messageID, &data); err != nil { + continue + } + targets := byID[messageID] + if len(targets) == 0 { + continue + } + record := decodeJSONObjectString(data) + if record == nil { + continue + } + for _, target := range targets { + target.parts = append(target.parts, record) + } } - entry, ok := openCodeMessageEntry(path, record, rowID, rowSessionID, 1) - return entry, ok, nil } -func openCodeMessageEntry(source string, record map[string]any, rowID, rowSessionID string, sourceLine int) (usage.Entry, bool) { - if record == nil { - return usage.Entry{}, false +func openCodeSessionEntries(source string, session openCodeSession, messages []*openCodeMessage) []usage.Entry { + entries := make([]usage.Entry, 0, len(messages)) + var previous usage.TokenUsage + for _, message := range messages { + // User messages carry an all-zero token block. Letting one reset the + // running counters would make the next assistant message look like a + // fresh session and bill it for the whole context again. + reported := openCodeTokens(message.data) + if !nonZero(reported) { + continue + } + billed, ok := billOpenCodeTokens(reported, previous) + previous = reported + if !ok { + continue + } + entry, ok := openCodeEntry(source, session, message, billed) + if !ok { + continue + } + entries = append(entries, entry) } - tokenBlock := objectAt(record["tokens"]) - if tokenBlock == nil { - return usage.Entry{}, false + return entries +} + +// openCodeTokens reads a message's token block exactly as OpenCode wrote it. +// The fields do not all mean the same thing: input and the cache counters are +// running context totals that grow with every turn, while output and reasoning +// are what the model produced on this turn alone. billOpenCodeTokens is what +// turns this mixture into one turn's cost. +func openCodeTokens(record map[string]any) usage.TokenUsage { + block := objectAt(record["tokens"]) + if block == nil { + return usage.TokenUsage{} } - cache := objectAt(tokenBlock["cache"]) - tokens := usage.TokenUsage{ - InputTokens: uintField(tokenBlock, "input"), - OutputTokens: uintField(tokenBlock, "output"), + cache := objectAt(block["cache"]) + return usage.TokenUsage{ + InputTokens: uintField(block, "input"), + OutputTokens: uintField(block, "output"), + ReasoningOutputTokens: uintField(block, "reasoning"), CacheCreationInputTokens: uintField(cache, "write"), CacheReadInputTokens: uintField(cache, "read"), } - tokens = applyTotalFallback(tokens, uintField(tokenBlock, "total")) - if !nonZero(tokens) { - return usage.Entry{}, false +} + +// billOpenCodeTokens charges a message for one turn: the growth of the running +// counters plus the per-turn counters as reported. The message's own "total" +// field is deliberately ignored — it sums the running counters, so it inherits +// their whole history and is not this turn's cost. +func billOpenCodeTokens(reported, previous usage.TokenUsage) (usage.TokenUsage, bool) { + billed := usage.TokenUsage{ + InputTokens: growth(reported.InputTokens, previous.InputTokens), + CacheCreationInputTokens: growth(reported.CacheCreationInputTokens, previous.CacheCreationInputTokens), + CacheReadInputTokens: growth(reported.CacheReadInputTokens, previous.CacheReadInputTokens), + OutputTokens: reported.OutputTokens, + ReasoningOutputTokens: reported.ReasoningOutputTokens, + } + if !nonZero(billed) { + return usage.TokenUsage{}, false + } + return applyTotalFallback(billed, 0), true +} + +// growth returns how much a running counter advanced. A counter that shrank +// means the session was compacted or restarted, so the new value stands alone. +func growth(current, previous uint64) uint64 { + if current <= previous { + return current } + return current - previous +} + +func openCodeEntry(source string, session openCodeSession, message *openCodeMessage, tokens usage.TokenUsage) (usage.Entry, bool) { + record := message.data model := stringField(record, "modelID") if model == "" { return usage.Entry{}, false } - timestamp := time.Unix(0, 0).UTC() + + timestamp := time.UnixMilli(message.created).UTC() if parsed, ok := parseTimestamp(objectAt(record["time"])["created"]); ok { timestamp = parsed } - sessionID := firstNonEmpty(rowSessionID, stringField(record, "sessionID"), "unknown") - messageID := firstNonEmpty(rowID, stringField(record, "id")) - entry := baseEntry(usage.ProviderOpenCode, timestamp, "opencode", "OpenCode", sessionID, model, "OpenCode", tokens) - setSource(&entry, source, sourceLine, 0, 0) - entry.ID = stableOpenCodeMessageID(messageID) + + // The message records the directory the agent actually ran in; the session + // directory covers messages written before that field existed. + cwd := firstNonEmpty( + stringField(objectAt(record["path"]), "cwd"), + stringField(objectAt(record["path"]), "root"), + session.directory, + ) + project, projectPath := usage.UnknownProject, "" + if path, name, ok := usage.ProjectFromCWD(cwd); ok { + project, projectPath = name, path + } + + sessionID := firstNonEmpty(message.sessionID, usage.UnknownProject) + entry := baseEntry(usage.ProviderOpenCode, timestamp, project, projectPath, sessionID, model, "OpenCode", tokens) + setSource(&entry, source, 0, 0, 0) + entry.ID = stableOpenCodeMessageID(message.id) if entry.ID == "" { entry.ID = stableEntryID(entry) } + + for _, part := range message.parts { + for _, change := range openCodePartChanges(part, cwd) { + entry.ApplyFileChange(change) + } + } return entry, true } +// openCodePartChanges extracts the file edits one message part performed. +// Only tool calls count: a "patch" part is a git snapshot of the whole +// worktree, which lists files nobody edited (.DS_Store, build output) and +// carries no line counts, so treating it as agent work invents writes. +func openCodePartChanges(part map[string]any, cwd string) []usage.FileChange { + if stringField(part, "type") != "tool" { + return nil + } + + state := objectAt(part["state"]) + if stringField(state, "status") != "completed" { + return nil + } + input := objectAt(state["input"]) + metadata := objectAt(state["metadata"]) + + switch stringField(part, "tool") { + case "edit": + return openCodeEditChanges(input, metadata, cwd) + case "write": + return openCodeWriteChanges(input, metadata, cwd) + case "patch", "apply_patch": + return openCodeApplyPatchChanges(input, metadata, cwd) + default: + return nil + } +} + +// openCodeEditChanges prefers the diff OpenCode computed. Without it the +// replaced and replacing strings still bound the change. +func openCodeEditChanges(input, metadata map[string]any, cwd string) []usage.FileChange { + path := usage.ResolvePath(cwd, stringField(input, "filePath")) + + if diff := objectAt(metadata["filediff"]); diff != nil { + added := uintField(diff, "additions") + removed := uintField(diff, "deletions") + if diffPath := firstNonEmpty(stringField(diff, "filePath"), stringField(diff, "file")); diffPath != "" { + path = usage.ResolvePath(cwd, diffPath) + } + if added != 0 || removed != 0 { + return []usage.FileChange{{Path: path, LinesAdded: added, LinesRemoved: removed}} + } + } + + added := usage.CountLines(stringField(input, "newString")) + removed := usage.CountLines(stringField(input, "oldString")) + if path == "" && added == 0 && removed == 0 { + return nil + } + return []usage.FileChange{{Path: path, LinesAdded: added, LinesRemoved: removed}} +} + +// openCodeWriteChanges counts a write as adding its whole content. Overwriting +// an existing file replaces lines this record does not describe, so only the +// added side is known. +func openCodeWriteChanges(input, metadata map[string]any, cwd string) []usage.FileChange { + path := usage.ResolvePath(cwd, firstNonEmpty( + stringField(input, "filePath"), + stringField(metadata, "filepath"), + )) + if path == "" { + return nil + } + return []usage.FileChange{{Path: path, LinesAdded: usage.CountLines(stringField(input, "content"))}} +} + +// openCodeApplyPatchChanges prefers the per-file summary OpenCode attaches and +// falls back to counting the raw patch envelope itself. +func openCodeApplyPatchChanges(input, metadata map[string]any, cwd string) []usage.FileChange { + if files, ok := metadata["files"].([]any); ok && len(files) > 0 { + changes := make([]usage.FileChange, 0, len(files)) + for _, raw := range files { + file := objectAt(raw) + if file == nil { + continue + } + path := usage.ResolvePath(cwd, firstNonEmpty(stringField(file, "filePath"), stringField(file, "file"))) + if path == "" { + continue + } + changes = append(changes, usage.FileChange{ + Path: path, + LinesAdded: uintField(file, "additions"), + LinesRemoved: uintField(file, "deletions"), + }) + } + if len(changes) > 0 { + return changes + } + } + + patchText := firstNonEmpty( + stringField(input, "patchText"), + stringField(input, "patch"), + stringField(metadata, "patch"), + ) + return parseOpenCodePatchEnvelope(patchText, cwd) +} + +// parseOpenCodePatchEnvelope counts the added and removed lines of each file in +// an apply_patch envelope. One envelope can carry any number of files. +func parseOpenCodePatchEnvelope(patchText, cwd string) []usage.FileChange { + if strings.TrimSpace(patchText) == "" { + return nil + } + + changes := make([]usage.FileChange, 0) + current := usage.FileChange{} + flush := func() { + if current.Path != "" { + changes = append(changes, current) + } + current = usage.FileChange{} + } + + for _, line := range strings.Split(patchText, "\n") { + switch { + case strings.HasPrefix(line, "*** Move to: "): + // A rename keeps the diff it already accumulated, under the new name. + if current.Path != "" { + current.Path = usage.ResolvePath(cwd, strings.TrimPrefix(line, "*** Move to: ")) + } + case strings.HasPrefix(line, "*** Add File: "), + strings.HasPrefix(line, "*** Update File: "), + strings.HasPrefix(line, "*** Delete File: "): + flush() + current.Path = usage.ResolvePath(cwd, openCodePatchHeaderPath(line)) + case strings.HasPrefix(line, "*** End Patch"): + flush() + case current.Path == "": + case strings.HasPrefix(line, "+"): + current.LinesAdded++ + case strings.HasPrefix(line, "-"): + current.LinesRemoved++ + } + } + flush() + if len(changes) == 0 { + return nil + } + return changes +} + +func openCodePatchHeaderPath(line string) string { + for _, prefix := range []string{"*** Add File: ", "*** Update File: ", "*** Delete File: "} { + if path, ok := strings.CutPrefix(line, prefix); ok { + return strings.TrimSpace(path) + } + } + return "" +} + +// loadOpenCodeLegacyFile reads one message from the pre-database storage layout, +// where the message and its parts each lived in their own JSON file. +func loadOpenCodeLegacyFile(path string) ([]usage.Entry, error) { + data, err := os.ReadFile(path) + if err != nil { + return nil, err + } + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + var record map[string]any + if err := decoder.Decode(&record); err != nil { + return nil, nil + } + + message := &openCodeMessage{ + id: stringField(record, "id"), + sessionID: stringField(record, "sessionID"), + data: record, + parts: readOpenCodeLegacyParts(path, stringField(record, "id")), + } + // One file holds one message, so there is no predecessor to subtract from: + // billing against a zero baseline charges it for everything it reports. + tokens, ok := billOpenCodeTokens(openCodeTokens(record), usage.TokenUsage{}) + if !ok { + return nil, nil + } + entry, ok := openCodeEntry(path, openCodeSession{}, message, tokens) + if !ok { + return nil, nil + } + entry.SourceLine = 1 + return []usage.Entry{entry}, nil +} + +// readOpenCodeLegacyParts loads the part files stored alongside a legacy +// message: storage/message//.json pairs with +// storage/part//*.json. +func readOpenCodeLegacyParts(messagePath, messageID string) []map[string]any { + if messageID == "" { + return nil + } + storage := filepath.Dir(filepath.Dir(filepath.Dir(messagePath))) + files := collectExt(filepath.Join(storage, "part", messageID), ".json") + sort.Strings(files) + + parts := make([]map[string]any, 0, len(files)) + for _, file := range files { + data, err := os.ReadFile(file) + if err != nil { + continue + } + decoder := json.NewDecoder(bytes.NewReader(data)) + decoder.UseNumber() + var record map[string]any + if err := decoder.Decode(&record); err != nil { + continue + } + parts = append(parts, record) + } + return parts +} + +// stableOpenCodeMessageID keys an entry to the message id OpenCode assigned. +// The id travels with the message no matter which database file or storage +// layout holds it, so re-ingesting the same message never double-counts. func stableOpenCodeMessageID(messageID string) string { if strings.TrimSpace(messageID) == "" { return "" diff --git a/internal/agentusage/opencode_test.go b/internal/agentusage/opencode_test.go new file mode 100644 index 0000000..90d90f2 --- /dev/null +++ b/internal/agentusage/opencode_test.go @@ -0,0 +1,388 @@ +package agentusage + +import ( + "database/sql" + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// openCodeDB builds a database shaped like the one OpenCode writes: a session +// row, message rows carrying the token block, and part rows carrying the tool +// calls that touched files. +type openCodeRow struct { + id string + sessionID string + created int64 + data string +} + +func writeOpenCodeDB(t *testing.T, sessions, messages, parts []openCodeRow) string { + t.Helper() + return writeOpenCodeDBAt(t, filepath.Join(t.TempDir(), "opencode.db"), sessions, messages, parts) +} + +func writeOpenCodeDBAt(t *testing.T, path string, sessions, messages, parts []openCodeRow) string { + t.Helper() + db, err := sql.Open("sqlite", path) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + schema := []string{ + `CREATE TABLE session (id TEXT PRIMARY KEY, directory TEXT, version TEXT)`, + `CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT, time_created INTEGER, data TEXT)`, + `CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, session_id TEXT, time_created INTEGER, data TEXT)`, + } + for _, statement := range schema { + if _, err := db.Exec(statement); err != nil { + t.Fatal(err) + } + } + for _, row := range sessions { + if _, err := db.Exec(`INSERT INTO session (id, directory, version) VALUES (?, ?, ?)`, row.id, row.sessionID, row.data); err != nil { + t.Fatal(err) + } + } + for _, row := range messages { + if _, err := db.Exec(`INSERT INTO message (id, session_id, time_created, data) VALUES (?, ?, ?, ?)`, + row.id, row.sessionID, row.created, row.data); err != nil { + t.Fatal(err) + } + } + for i, row := range parts { + if _, err := db.Exec(`INSERT INTO part (id, message_id, session_id, time_created, data) VALUES (?, ?, ?, ?, ?)`, + row.id, row.sessionID, "ses-1", int64(i), row.data); err != nil { + t.Fatal(err) + } + } + return path +} + +// TestOpenCodeBillsRunningInputAndPerTurnOutput pins down the token semantics. +// OpenCode's "input" is the whole context resent every turn, so only its growth +// is this turn's cost, while "output" and "reasoning" already describe one turn +// and must be taken as reported. Getting this backwards inflated totals ninefold. +func TestOpenCodeBillsRunningInputAndPerTurnOutput(t *testing.T) { + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []openCodeRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":1000,"output":50,"reasoning":20,"total":1070}}`}, + {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":1400,"output":30,"reasoning":10,"total":1440}}`}, + }, nil) + + entries, err := loadOpenCodeDatabase(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("len(entries) = %d, want 2", len(entries)) + } + sortEntries(entries) + + first, second := entries[0], entries[1] + if first.Usage.InputTokens != 1000 || first.Usage.OutputTokens != 50 || first.Usage.ReasoningOutputTokens != 20 { + t.Fatalf("first usage = %+v, want the full first turn", first.Usage) + } + // 1400 - 1000: the second turn only paid for the context it added. + if second.Usage.InputTokens != 400 { + t.Fatalf("second input = %d, want 400 (growth only)", second.Usage.InputTokens) + } + if second.Usage.OutputTokens != 30 || second.Usage.ReasoningOutputTokens != 10 { + t.Fatalf("second output/reasoning = %d/%d, want 30/10 as reported", + second.Usage.OutputTokens, second.Usage.ReasoningOutputTokens) + } + // The message's own "total" sums the running counters, so it must not leak + // into the billed total. + if second.Usage.TotalTokens != 440 { + t.Fatalf("second total = %d, want 440 (billed fields only)", second.Usage.TotalTokens) + } +} + +// TestOpenCodeIgnoresZeroTokenMessages guards the counter against user turns, +// which carry an all-zero token block. Treating one as a counter reset made the +// next assistant turn pay for the whole context a second time. +func TestOpenCodeIgnoresZeroTokenMessages(t *testing.T) { + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []openCodeRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":1000,"output":50}}`}, + {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"role":"user","tokens":{"input":0,"output":0,"cache":{"read":0,"write":0}}}`}, + {id: "msg-3", sessionID: "ses-1", created: 3000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":1200,"output":20}}`}, + }, nil) + + entries, err := loadOpenCodeDatabase(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("len(entries) = %d, want 2 (the user turn bills nothing)", len(entries)) + } + sortEntries(entries) + if got := entries[1].Usage.InputTokens; got != 200 { + t.Fatalf("input after a user turn = %d, want 200 (the user turn must not reset the counter)", got) + } +} + +// TestOpenCodeRestartsBillingAfterCompaction covers a shrinking counter: the +// context was compacted, so the smaller value is a fresh baseline, not a refund. +func TestOpenCodeRestartsBillingAfterCompaction(t *testing.T) { + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []openCodeRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":9000,"output":10}}`}, + {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":300,"output":10}}`}, + }, nil) + + entries, err := loadOpenCodeDatabase(path) + if err != nil { + t.Fatal(err) + } + sortEntries(entries) + if len(entries) != 2 { + t.Fatalf("len(entries) = %d, want 2", len(entries)) + } + if got := entries[1].Usage.InputTokens; got != 300 { + t.Fatalf("input after compaction = %d, want 300 (the whole new context)", got) + } +} + +// TestOpenCodeProjectComesFromMessageCwd pins the project to the directory the +// agent actually ran in, falling back to the session's directory. +func TestOpenCodeProjectComesFromMessageCwd(t *testing.T) { + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/session-dir"}}, + []openCodeRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","modelID":"gpt-5","path":{"cwd":"/repo/message-dir"},"tokens":{"input":10,"output":5}}`}, + {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":20,"output":5}}`}, + {id: "msg-3", sessionID: "ses-2", created: 3000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":30,"output":5}}`}, + }, nil) + + entries, err := loadOpenCodeDatabase(path) + if err != nil { + t.Fatal(err) + } + sortEntries(entries) + if len(entries) != 3 { + t.Fatalf("len(entries) = %d, want 3", len(entries)) + } + if entries[0].Project != "message-dir" || entries[0].ProjectPath != "/repo/message-dir" { + t.Fatalf("message cwd project = %q (%q), want message-dir", entries[0].Project, entries[0].ProjectPath) + } + if entries[1].Project != "session-dir" { + t.Fatalf("session fallback project = %q, want session-dir", entries[1].Project) + } + // A session row that does not exist leaves nothing to name the project with. + if entries[2].Project != usage.UnknownProject { + t.Fatalf("unknown session project = %q, want %q", entries[2].Project, usage.UnknownProject) + } +} + +// TestOpenCodeCollectsFileChangesFromTools covers all three editing tools in one +// turn, including the per-file breakdown and which file becomes the entity. +func TestOpenCodeCollectsFileChangesFromTools(t *testing.T) { + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []openCodeRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","modelID":"gpt-5","path":{"cwd":"/repo/demo"},"tokens":{"input":10,"output":5}}`}, + }, + []openCodeRow{ + // edit: the recorded diff wins over the raw strings. + {id: "prt-1", sessionID: "msg-1", data: `{"type":"tool","tool":"edit","state":{"status":"completed","input":{"filePath":"small.go","oldString":"a","newString":"b"},"metadata":{"filediff":{"filePath":"small.go","additions":3,"deletions":1}}}}`}, + // write: a relative path resolves against the cwd, content is all additions. + {id: "prt-2", sessionID: "msg-1", data: `{"type":"tool","tool":"write","state":{"status":"completed","input":{"filePath":"big.go","content":"1\n2\n3\n4\n5"},"metadata":{"exists":false}}}`}, + // apply_patch: the per-file summary is used verbatim. + {id: "prt-3", sessionID: "msg-1", data: `{"type":"tool","tool":"apply_patch","state":{"status":"completed","metadata":{"files":[{"filePath":"/abs/other.go","additions":2,"deletions":2}]}}}`}, + }) + + entries, err := loadOpenCodeDatabase(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + entry := entries[0] + + if entry.LinesAdded != 3+5+2 || entry.LinesRemoved != 1+0+2 { + t.Fatalf("totals = +%d/-%d, want +10/-3", entry.LinesAdded, entry.LinesRemoved) + } + if entry.IsWrite == nil || !*entry.IsWrite { + t.Fatal("IsWrite = false, want true when files changed") + } + // big.go changed 5 lines, more than small.go's 4 and other.go's 4. + if entry.Entity != "/repo/demo/big.go" { + t.Fatalf("entity = %q, want the most-changed file", entry.Entity) + } + if entry.EntityType != "file" { + t.Fatalf("entity type = %q, want file", entry.EntityType) + } + + want := map[string]usage.FileChange{ + "/repo/demo/small.go": {Path: "/repo/demo/small.go", LinesAdded: 3, LinesRemoved: 1}, + "/repo/demo/big.go": {Path: "/repo/demo/big.go", LinesAdded: 5}, + "/abs/other.go": {Path: "/abs/other.go", LinesAdded: 2, LinesRemoved: 2}, + } + if len(entry.Files) != len(want) { + t.Fatalf("files = %+v, want %d entries", entry.Files, len(want)) + } + for _, got := range entry.Files { + expected, ok := want[got.Path] + if !ok { + t.Fatalf("unexpected file %q", got.Path) + } + if got != expected { + t.Fatalf("file %q = %+v, want %+v", got.Path, got, expected) + } + } +} + +// TestOpenCodeIgnoresSnapshotPatchParts guards against counting git snapshots as +// agent edits. A "patch" part lists whatever the worktree snapshot covered — +// build output, .DS_Store — and records no line counts at all. +func TestOpenCodeIgnoresSnapshotPatchParts(t *testing.T) { + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []openCodeRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":10,"output":5}}`}, + }, + []openCodeRow{ + {id: "prt-1", sessionID: "msg-1", data: `{"type":"patch","hash":"deadbeef","files":["/repo/demo/.DS_Store"]}`}, + }) + + entries, err := loadOpenCodeDatabase(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + if len(entries[0].Files) != 0 || entries[0].IsWrite != nil { + t.Fatalf("snapshot part produced files %+v / IsWrite %v, want none", + entries[0].Files, entries[0].IsWrite) + } +} + +// TestOpenCodeSkipsIncompleteToolCalls keeps a tool that errored or is still +// running out of the line counts. +func TestOpenCodeSkipsIncompleteToolCalls(t *testing.T) { + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []openCodeRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":10,"output":5}}`}, + }, + []openCodeRow{ + {id: "prt-1", sessionID: "msg-1", data: `{"type":"tool","tool":"write","state":{"status":"error","input":{"filePath":"/repo/demo/a.go","content":"x"}}}`}, + {id: "prt-2", sessionID: "msg-1", data: `{"type":"tool","tool":"read","state":{"status":"completed","input":{"filePath":"/repo/demo/b.go"}}}`}, + }) + + entries, err := loadOpenCodeDatabase(path) + if err != nil { + t.Fatal(err) + } + if len(entries[0].Files) != 0 { + t.Fatalf("files = %+v, want none from a failed write and a read", entries[0].Files) + } +} + +// TestOpenCodeParsesPatchEnvelope covers apply_patch without a per-file summary, +// where the raw envelope is the only record of what changed. One envelope can +// carry several files, and a rename keeps the diff under the new name. +func TestOpenCodeParsesPatchEnvelope(t *testing.T) { + patch := "*** Begin Patch\\n" + + "*** Add File: new.go\\n" + + "+package main\\n" + + "+func main() {}\\n" + + "*** Update File: old.go\\n" + + "-was here\\n" + + "+is here\\n" + + "*** Move to: renamed.go\\n" + + "*** End Patch" + + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []openCodeRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","modelID":"gpt-5","path":{"cwd":"/repo/demo"},"tokens":{"input":10,"output":5}}`}, + }, + []openCodeRow{ + {id: "prt-1", sessionID: "msg-1", data: `{"type":"tool","tool":"apply_patch","state":{"status":"completed","input":{"patchText":"` + patch + `"}}}`}, + }) + + entries, err := loadOpenCodeDatabase(path) + if err != nil { + t.Fatal(err) + } + entry := entries[0] + if entry.LinesAdded != 3 || entry.LinesRemoved != 1 { + t.Fatalf("totals = +%d/-%d, want +3/-1", entry.LinesAdded, entry.LinesRemoved) + } + + got := make(map[string]usage.FileChange, len(entry.Files)) + for _, file := range entry.Files { + got[file.Path] = file + } + if change := got["/repo/demo/new.go"]; change.LinesAdded != 2 || change.LinesRemoved != 0 { + t.Fatalf("new.go = %+v, want +2/-0", change) + } + // The rename carries old.go's diff over to its new name. + if change := got["/repo/demo/renamed.go"]; change.LinesAdded != 1 || change.LinesRemoved != 1 { + t.Fatalf("renamed.go = %+v, want +1/-1", change) + } + if _, ok := got["/repo/demo/old.go"]; ok { + t.Fatal("old.go still present, want it recorded under the new name") + } +} + +// TestOpenCodeEntryIDSurvivesRelocation keys entries to the message id OpenCode +// assigned, so the same message read from another database file or from the +// legacy JSON layout is recognised as one event rather than counted twice. +func TestOpenCodeEntryIDSurvivesRelocation(t *testing.T) { + message := `{"role":"assistant","modelID":"gpt-5","path":{"cwd":"/repo/demo"},"tokens":{"input":10,"output":5}}` + first := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []openCodeRow{{id: "msg-1", sessionID: "ses-1", created: 1000, data: message}}, nil) + second := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/elsewhere/demo"}}, + []openCodeRow{{id: "msg-1", sessionID: "ses-1", created: 1000, data: message}}, nil) + + a, err := loadOpenCodeDatabase(first) + if err != nil { + t.Fatal(err) + } + b, err := loadOpenCodeDatabase(second) + if err != nil { + t.Fatal(err) + } + if a[0].ID != b[0].ID { + t.Fatalf("ids differ across databases: %q vs %q", a[0].ID, b[0].ID) + } + if a[0].ID != stableOpenCodeMessageID("msg-1") { + t.Fatalf("id = %q, want the message-keyed id", a[0].ID) + } +} + +// TestOpenCodeReadsEveryChannelDatabase covers OpenCode's release channels, +// which sit side by side in the data root as separate databases. +func TestOpenCodeReadsEveryChannelDatabase(t *testing.T) { + root := t.TempDir() + for _, channel := range []string{"opencode.db", "opencode-nightly.db"} { + writeOpenCodeDBAt(t, filepath.Join(root, channel), + []openCodeRow{{id: "ses-" + channel, sessionID: "/repo/demo"}}, + []openCodeRow{{ + id: "msg-" + channel, + sessionID: "ses-" + channel, + created: 1000, + data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":10,"output":5}}`, + }}, nil) + } + + entries, err := loadOpenCodeRoot(root, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("len(entries) = %d, want one per channel database", len(entries)) + } +} diff --git a/internal/agentusage/providers_test.go b/internal/agentusage/providers_test.go index 9feb4a7..2e06097 100644 --- a/internal/agentusage/providers_test.go +++ b/internal/agentusage/providers_test.go @@ -323,13 +323,13 @@ func TestProvidersLoadEntries(t *testing.T) { provider: func() ([]usage.Entry, error) { dir := t.TempDir() path := filepath.Join(dir, "storage", "message", "session-a", "msg-1.json") - writeFile(t, path, `{"id":"msg-1","sessionID":"session-a","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","time":{"created":1767312000000},"tokens":{"input":100,"output":50,"cache":{"read":10,"write":20}},"cost":0}`) + writeFile(t, path, `{"id":"msg-1","sessionID":"session-a","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","path":{"cwd":"/repo/demo"},"time":{"created":1767312000000},"tokens":{"input":100,"output":50,"cache":{"read":10,"write":20}},"cost":0}`) return OpenCodeProvider{}.WithPaths([]string{dir}).Entries() }, want: usage.ProviderOpenCode, model: "claude-sonnet-4-20250514", sessionID: "session-a", - project: "opencode", + project: "demo", tokens: usage.TokenUsage{ InputTokens: 100, OutputTokens: 50, diff --git a/internal/claudeusage/loader.go b/internal/claudeusage/loader.go index 77f7068..186c50a 100644 --- a/internal/claudeusage/loader.go +++ b/internal/claudeusage/loader.go @@ -355,7 +355,7 @@ func parsePatchLine(line []byte) (patchStats, bool) { if result.Type != "create" || result.FilePath == "" { return patchStats{}, false } - return patchStats{file: result.FilePath, added: countLines(result.Content)}, true + return patchStats{file: result.FilePath, added: usage.CountLines(result.Content)}, true } stats := patchStats{file: result.FilePath} @@ -375,33 +375,6 @@ func parsePatchLine(line []byte) (patchStats, bool) { return stats, true } -func countLines(content string) uint64 { - if content == "" { - return 0 - } - lines := uint64(strings.Count(content, "\n")) - if !strings.HasSuffix(content, "\n") { - lines++ - } - return lines -} - -// projectFromCWD derives the project path and name from a transcript line's -// cwd field. It reports false for values that cannot name a project (empty, -// relative, or the filesystem root), so callers keep their fallback. -func projectFromCWD(cwd string) (string, string, bool) { - clean := strings.TrimSpace(cwd) - if clean == "" || !filepath.IsAbs(clean) { - return "", "", false - } - clean = filepath.Clean(clean) - name := filepath.Base(clean) - if name == "." || name == string(filepath.Separator) || strings.TrimSpace(name) == "" { - return "", "", false - } - return clean, name, true -} - // ExtractSessionID derives the session id from a usage file's location under // the projects directory: projects//.jsonl for sessions and // projects///subagents/.jsonl for subagents. @@ -460,7 +433,7 @@ func parseUsageLine(line []byte, sessionID string) (LoadedEntry, bool) { project := usage.UnknownProject projectPath := "" if data.CWD != nil { - if path, name, ok := projectFromCWD(*data.CWD); ok { + if path, name, ok := usage.ProjectFromCWD(*data.CWD); ok { projectPath = path project = name } diff --git a/internal/claudeusage/loader_test.go b/internal/claudeusage/loader_test.go index b81a6f5..2b9f20b 100644 --- a/internal/claudeusage/loader_test.go +++ b/internal/claudeusage/loader_test.go @@ -272,8 +272,8 @@ func TestReadUsageFileReportsUnknownProjectWithoutCWD(t *testing.T) { func TestProjectFromCWDRejectsUnusableValues(t *testing.T) { for _, cwd := range []string{"", " ", "relative/path", "/"} { - if _, _, ok := projectFromCWD(cwd); ok { - t.Fatalf("projectFromCWD(%q) ok = true, want false", cwd) + if _, _, ok := usage.ProjectFromCWD(cwd); ok { + t.Fatalf("usage.ProjectFromCWD(%q) ok = true, want false", cwd) } } } diff --git a/internal/usage/usage.go b/internal/usage/usage.go index 5696645..585788a 100644 --- a/internal/usage/usage.go +++ b/internal/usage/usage.go @@ -3,6 +3,7 @@ package usage import ( "crypto/sha256" "encoding/hex" + "path/filepath" "sort" "strings" "time" @@ -66,6 +67,87 @@ type FileChange struct { LinesRemoved uint64 `json:"lines_removed,omitempty"` } +// ProjectFromCWD derives the project path and name from a working directory an +// agent recorded. It reports false for values that cannot name a project +// (empty, relative, or the filesystem root), so callers keep their fallback. +func ProjectFromCWD(cwd string) (string, string, bool) { + clean := strings.TrimSpace(cwd) + if clean == "" || !filepath.IsAbs(clean) { + return "", "", false + } + clean = filepath.Clean(clean) + name := filepath.Base(clean) + if name == "." || name == string(filepath.Separator) || strings.TrimSpace(name) == "" { + return "", "", false + } + return clean, name, true +} + +// CountLines counts the source lines in a blob of file content. A trailing +// newline does not add a line, so "a\nb\n" and "a\nb" both count as 2. +func CountLines(content string) uint64 { + if content == "" { + return 0 + } + lines := uint64(strings.Count(content, "\n")) + if !strings.HasSuffix(content, "\n") { + lines++ + } + return lines +} + +// ResolvePath makes a file path recorded by an agent absolute, interpreting a +// relative one against the working directory the agent ran in. +func ResolvePath(cwd, path string) string { + path = strings.TrimSpace(path) + switch { + case path == "": + return "" + case filepath.IsAbs(path): + return filepath.Clean(path) + case strings.TrimSpace(cwd) == "": + return path + default: + return filepath.Join(cwd, path) + } +} + +// ApplyFileChange folds one file's diff into an entry: it accumulates the +// per-file totals, keeps the entry totals in sync, and re-points Entity at the +// most-changed file. Every provider that records diffs funnels through here so +// "Entity is the biggest change" holds identically everywhere. +func (e *Entry) ApplyFileChange(change FileChange) { + e.LinesAdded += change.LinesAdded + e.LinesRemoved += change.LinesRemoved + write := true + e.IsWrite = &write + if change.Path == "" { + return + } + + found := false + for i := range e.Files { + if e.Files[i].Path == change.Path { + e.Files[i].LinesAdded += change.LinesAdded + e.Files[i].LinesRemoved += change.LinesRemoved + found = true + break + } + } + if !found { + e.Files = append(e.Files, change) + } + + best, bestWeight := "", uint64(0) + for _, file := range e.Files { + if weight := file.LinesAdded + file.LinesRemoved; weight >= bestWeight { + best, bestWeight = file.Path, weight + } + } + e.Entity = best + e.EntityType = "file" +} + type Entry struct { Provider Provider `json:"provider"` ID string `json:"id,omitempty"` diff --git a/pkg/agentlib/agentlib.go b/pkg/agentlib/agentlib.go index 99f741f..ef1d0aa 100644 --- a/pkg/agentlib/agentlib.go +++ b/pkg/agentlib/agentlib.go @@ -479,7 +479,12 @@ func DefaultProviderDirs() map[Provider][]string { ProviderKilo: {filepath.Join(home, ".local", "share", "kilo")}, ProviderHermes: {filepath.Join(home, ".hermes")}, ProviderCodebuff: {filepath.Join(home, ".config", "manicode"), filepath.Join(home, ".config", "manicode-dev"), filepath.Join(home, ".config", "manicode-staging")}, - ProviderOpenCode: {filepath.Join(home, ".local", "share", "opencode")}, + ProviderOpenCode: { + filepath.Join(home, ".local", "share", "opencode"), + filepath.Join(home, "Library", "Application Support", "opencode"), + filepath.Join(home, "AppData", "Local", "opencode"), + filepath.Join(home, "AppData", "Roaming", "opencode"), + }, ProviderGoose: { filepath.Join(home, ".local", "share", "goose", "sessions", "sessions.db"), filepath.Join(home, "Library", "Application Support", "goose", "sessions", "sessions.db"), From 22e56e51258013756354814ee370f955e55896c6 Mon Sep 17 00:00:00 2001 From: akarineren Date: Fri, 31 Jul 2026 20:42:38 +0900 Subject: [PATCH 04/11] feat: add OpenCode, pi-agent, and Qwen providers with entry loading functionality - Implemented OpenCode provider for loading usage entries from OpenCode data sources. - Created pi-agent provider to handle usage entries specific to the pi-agent. - Developed Qwen provider for loading usage entries from Qwen chat files. - Added shared base functionality for providers to manage paths and filters. - Implemented entry loading logic for each provider, including parsing and filtering. - Added unit tests for each provider to ensure correct entry loading and validation. - Introduced shared utilities for handling file operations and JSON parsing. --- internal/agentusage/amp.go | 182 -------- internal/agentusage/hermes.go | 103 ---- internal/agentusage/kilo.go | 146 ------ internal/agentusage/kimi.go | 198 -------- internal/agentusage/openclaw.go | 132 ------ internal/agentusage/pi.go | 96 ---- internal/agentusage/providers.go | 376 --------------- internal/agentusage/providers_test.go | 440 ------------------ internal/agentusage/qwen.go | 104 ----- internal/provider/amp/loader.go | 183 ++++++++ internal/provider/amp/provider.go | 37 ++ internal/provider/amp/provider_test.go | 32 ++ .../claude}/loader.go | 2 +- .../claude}/loader_test.go | 2 +- .../claude}/provider.go | 2 +- .../codebuff/loader.go} | 127 ++--- internal/provider/codebuff/provider.go | 37 ++ internal/provider/codebuff/provider_test.go | 34 ++ .../{codexusage => provider/codex}/loader.go | 2 +- .../codex}/loader_test.go | 2 +- .../codex}/provider.go | 2 +- internal/provider/copilot/filter_test.go | 33 ++ .../copilot.go => provider/copilot/loader.go} | 205 ++++---- internal/provider/copilot/provider.go | 37 ++ internal/provider/copilot/provider_test.go | 35 ++ .../droid.go => provider/droid/loader.go} | 63 +-- internal/provider/droid/provider.go | 37 ++ internal/provider/droid/provider_test.go | 35 ++ .../gemini.go => provider/gemini/loader.go} | 123 ++--- internal/provider/gemini/provider.go | 37 ++ internal/provider/gemini/provider_test.go | 36 ++ .../goose.go => provider/goose/loader.go} | 65 +-- internal/provider/goose/provider.go | 30 ++ internal/provider/goose/provider_test.go | 52 +++ internal/provider/hermes/loader.go | 98 ++++ internal/provider/hermes/provider.go | 30 ++ internal/provider/hermes/provider_test.go | 57 +++ internal/provider/kilo/loader.go | 89 ++++ internal/provider/kilo/provider.go | 30 ++ internal/provider/kilo/provider_test.go | 42 ++ internal/provider/kimi/loader.go | 199 ++++++++ internal/provider/kimi/provider.go | 37 ++ internal/provider/kimi/provider_test.go | 37 ++ internal/provider/openclaw/loader.go | 133 ++++++ internal/provider/openclaw/provider.go | 37 ++ internal/provider/openclaw/provider_test.go | 35 ++ .../opencode/loader.go} | 249 +++++----- .../opencode/loader_test.go} | 35 +- internal/provider/opencode/provider.go | 37 ++ internal/provider/opencode/provider_test.go | 34 ++ internal/provider/pi/loader.go | 97 ++++ internal/provider/pi/provider.go | 37 ++ internal/provider/pi/provider_test.go | 31 ++ internal/provider/qwen/loader.go | 105 +++++ internal/provider/qwen/provider.go | 37 ++ internal/provider/qwen/provider_test.go | 34 ++ internal/provider/shared/base.go | 37 ++ internal/provider/shared/entries.go | 16 + .../shared}/helpers.go | 104 ++--- .../{agentusage => provider/shared}/sqlite.go | 14 +- internal/provider/shared/sqlitefiles.go | 65 +++ internal/provider/shared/testing.go | 84 ++++ internal/usage/usage.go | 20 +- internal/usagedb/db.go | 4 +- internal/usagedb/db_test.go | 6 +- internal/usagescan/scanner.go | 48 +- 66 files changed, 2640 insertions(+), 2305 deletions(-) delete mode 100644 internal/agentusage/amp.go delete mode 100644 internal/agentusage/hermes.go delete mode 100644 internal/agentusage/kilo.go delete mode 100644 internal/agentusage/kimi.go delete mode 100644 internal/agentusage/openclaw.go delete mode 100644 internal/agentusage/pi.go delete mode 100644 internal/agentusage/providers.go delete mode 100644 internal/agentusage/providers_test.go delete mode 100644 internal/agentusage/qwen.go create mode 100644 internal/provider/amp/loader.go create mode 100644 internal/provider/amp/provider.go create mode 100644 internal/provider/amp/provider_test.go rename internal/{claudeusage => provider/claude}/loader.go (99%) rename internal/{claudeusage => provider/claude}/loader_test.go (99%) rename internal/{claudeusage => provider/claude}/provider.go (98%) rename internal/{agentusage/codebuff.go => provider/codebuff/loader.go} (59%) create mode 100644 internal/provider/codebuff/provider.go create mode 100644 internal/provider/codebuff/provider_test.go rename internal/{codexusage => provider/codex}/loader.go (99%) rename internal/{codexusage => provider/codex}/loader_test.go (99%) rename internal/{codexusage => provider/codex}/provider.go (98%) create mode 100644 internal/provider/copilot/filter_test.go rename internal/{agentusage/copilot.go => provider/copilot/loader.go} (53%) create mode 100644 internal/provider/copilot/provider.go create mode 100644 internal/provider/copilot/provider_test.go rename internal/{agentusage/droid.go => provider/droid/loader.go} (64%) create mode 100644 internal/provider/droid/provider.go create mode 100644 internal/provider/droid/provider_test.go rename internal/{agentusage/gemini.go => provider/gemini/loader.go} (51%) create mode 100644 internal/provider/gemini/provider.go create mode 100644 internal/provider/gemini/provider_test.go rename internal/{agentusage/goose.go => provider/goose/loader.go} (52%) create mode 100644 internal/provider/goose/provider.go create mode 100644 internal/provider/goose/provider_test.go create mode 100644 internal/provider/hermes/loader.go create mode 100644 internal/provider/hermes/provider.go create mode 100644 internal/provider/hermes/provider_test.go create mode 100644 internal/provider/kilo/loader.go create mode 100644 internal/provider/kilo/provider.go create mode 100644 internal/provider/kilo/provider_test.go create mode 100644 internal/provider/kimi/loader.go create mode 100644 internal/provider/kimi/provider.go create mode 100644 internal/provider/kimi/provider_test.go create mode 100644 internal/provider/openclaw/loader.go create mode 100644 internal/provider/openclaw/provider.go create mode 100644 internal/provider/openclaw/provider_test.go rename internal/{agentusage/opencode.go => provider/opencode/loader.go} (61%) rename internal/{agentusage/opencode_test.go => provider/opencode/loader_test.go} (95%) create mode 100644 internal/provider/opencode/provider.go create mode 100644 internal/provider/opencode/provider_test.go create mode 100644 internal/provider/pi/loader.go create mode 100644 internal/provider/pi/provider.go create mode 100644 internal/provider/pi/provider_test.go create mode 100644 internal/provider/qwen/loader.go create mode 100644 internal/provider/qwen/provider.go create mode 100644 internal/provider/qwen/provider_test.go create mode 100644 internal/provider/shared/base.go create mode 100644 internal/provider/shared/entries.go rename internal/{agentusage => provider/shared}/helpers.go (79%) rename internal/{agentusage => provider/shared}/sqlite.go (83%) create mode 100644 internal/provider/shared/sqlitefiles.go create mode 100644 internal/provider/shared/testing.go diff --git a/internal/agentusage/amp.go b/internal/agentusage/amp.go deleted file mode 100644 index 2377afd..0000000 --- a/internal/agentusage/amp.go +++ /dev/null @@ -1,182 +0,0 @@ -package agentusage - -import ( - "encoding/json" - "os" - "path/filepath" - "sort" - "strconv" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func loadAmpEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { - files := make([]string, 0) - for _, path := range paths { - if info, err := os.Stat(path); err == nil && !info.IsDir() { - if filepath.Ext(path) == ".json" { - files = append(files, path) - } - continue - } - files = append(files, collectExt(filepath.Join(path, "threads"), ".json")...) - if filepath.Base(path) == "threads" { - files = append(files, collectExt(path, ".json")...) - } - } - sort.Strings(files) - files = filterFiles(uniqueStrings(files), filter) - - entries := make([]usage.Entry, 0) - for _, file := range files { - fileEntries, err := parseAmpThreadFile(file) - if err != nil { - return nil, err - } - entries = append(entries, fileEntries...) - } - sortEntries(entries) - return entries, nil -} - -func parseAmpThreadFile(path string) ([]usage.Entry, error) { - thread, err := readJSONObject(path) - if err != nil || thread == nil { - return nil, err - } - threadID := stringField(thread, "id") - if threadID == "" { - return nil, nil - } - messages := arrayAt(thread["messages"]) - if ledger := objectAt(thread["usageLedger"]); ledger != nil { - if events := arrayAt(ledger["events"]); len(events) > 0 { - return ampLedgerEntries(path, threadID, messages, events), nil - } - } - return ampMessageEntries(path, threadID, messages), nil -} - -func ampLedgerEntries(path, threadID string, messages []any, events []any) []usage.Entry { - cacheTokens := ampCacheTokens(messages) - entries := make([]usage.Entry, 0) - for index, raw := range events { - event := objectAt(raw) - if event == nil { - continue - } - timestamp, ok := parseTimestamp(event["timestamp"]) - if !ok { - continue - } - model := stringField(event, "model") - if model == "" { - continue - } - tokenBlock := objectAt(event["tokens"]) - if tokenBlock == nil { - continue - } - cache := cacheTokens[int64Value(event["toMessageId"])] - tokens := usage.TokenUsage{ - InputTokens: uintField(tokenBlock, "input"), - OutputTokens: uintField(tokenBlock, "output"), - CacheCreationInputTokens: cache.cacheCreation, - CacheReadInputTokens: cache.cacheRead, - } - tokens = applyTotalFallback(tokens, uintField(tokenBlock, "total")) - if !nonZero(tokens) { - continue - } - messageID := stringValue(event["id"]) - entry := baseEntry(usage.ProviderAmp, timestamp, "amp", "Amp", threadID, model, "Amp", tokens) - setSource(&entry, path, index+1, 0, 0) - entry.ID = stableEntryID(entry, messageID) - entries = append(entries, entry) - } - return entries -} - -func ampMessageEntries(path, threadID string, messages []any) []usage.Entry { - entries := make([]usage.Entry, 0) - for index, raw := range messages { - message := objectAt(raw) - if message == nil || stringValue(message["role"]) != "assistant" { - continue - } - usageBlock := objectAt(message["usage"]) - if usageBlock == nil { - continue - } - timestamp, ok := parseTimestamp(usageBlock["timestamp"]) - if !ok { - timestamp, ok = parseTimestamp(message["timestamp"]) - } - if !ok { - continue - } - model := stringField(usageBlock, "model") - if model == "" { - model = stringValue(message["model"]) - } - if model == "" { - continue - } - tokens := usage.TokenUsage{ - InputTokens: uintField(usageBlock, "inputTokens"), - OutputTokens: uintField(usageBlock, "outputTokens"), - CacheCreationInputTokens: uintField(usageBlock, "cacheCreationInputTokens"), - CacheReadInputTokens: uintField(usageBlock, "cacheReadInputTokens"), - } - tokens = applyTotalFallback(tokens, uintField(usageBlock, "totalTokens")) - if !nonZero(tokens) { - continue - } - messageID := stringValue(message["messageId"]) - entry := baseEntry(usage.ProviderAmp, timestamp, "amp", "Amp", threadID, model, "Amp", tokens) - setSource(&entry, path, index+1, 0, 0) - entry.ID = stableEntryID(entry, messageID) - entries = append(entries, entry) - } - return entries -} - -type ampCache struct { - cacheCreation uint64 - cacheRead uint64 -} - -func ampCacheTokens(messages []any) map[int64]ampCache { - tokens := make(map[int64]ampCache) - for _, raw := range messages { - message := objectAt(raw) - if message == nil || stringValue(message["role"]) != "assistant" { - continue - } - id := int64Value(message["messageId"]) - if id == 0 { - continue - } - usageBlock := objectAt(message["usage"]) - tokens[id] = ampCache{ - cacheCreation: uintField(usageBlock, "cacheCreationInputTokens"), - cacheRead: uintField(usageBlock, "cacheReadInputTokens"), - } - } - return tokens -} - -func int64Value(value any) int64 { - switch typed := value.(type) { - case json.Number: - parsed, _ := strconv.ParseInt(typed.String(), 10, 64) - return parsed - case float64: - return int64(typed) - case string: - parsed, _ := strconv.ParseInt(typed, 10, 64) - return parsed - default: - return 0 - } -} diff --git a/internal/agentusage/hermes.go b/internal/agentusage/hermes.go deleted file mode 100644 index 57cfc5b..0000000 --- a/internal/agentusage/hermes.go +++ /dev/null @@ -1,103 +0,0 @@ -package agentusage - -import ( - "os" - "strings" - "time" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func loadHermesEntries(paths []string) ([]usage.Entry, error) { - dbPaths := sqliteDBPaths(paths, "state.db", nil) - entries := make([]usage.Entry, 0) - for _, dbPath := range dbPaths { - dbEntries, err := loadHermesDatabase(dbPath) - if err != nil { - return nil, err - } - entries = append(entries, dbEntries...) - } - sortEntries(entries) - return entries, nil -} - -func loadHermesDatabase(path string) ([]usage.Entry, error) { - db, err := openSQLite(path) - if err != nil { - return nil, err - } - defer db.Close() - - rows, err := db.Query(` - SELECT id, model, billing_provider, started_at, message_count, input_tokens, - output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, - estimated_cost_usd, actual_cost_usd - FROM sessions - WHERE model IS NOT NULL AND TRIM(model) != '' - `) - if err != nil { - return nil, nil - } - defer rows.Close() - - entries := make([]usage.Entry, 0) - for rows.Next() { - var sessionID, model, provider, startedAt, messageCount, input, output, cacheRead, cacheWrite, reasoning, estimatedCost, actualCost any - if !scanAny(rows, &sessionID, &model, &provider, &startedAt, &messageCount, &input, &output, &cacheRead, &cacheWrite, &reasoning, &estimatedCost, &actualCost) { - continue - } - entry, ok := hermesRowEntry(path, sessionID, model, startedAt, input, output, cacheRead, cacheWrite, reasoning) - if ok { - entries = append(entries, entry) - } - } - return entries, rows.Err() -} - -func hermesRowEntry(path string, sessionRaw, modelRaw, startedAt, input, output, cacheRead, cacheWrite, reasoning any) (usage.Entry, bool) { - sessionID := sqlString(sessionRaw) - model := strings.TrimSpace(sqlString(modelRaw)) - if sessionID == "" || model == "" { - return usage.Entry{}, false - } - timestamp, ok := hermesTimestamp(startedAt) - if !ok { - return usage.Entry{}, false - } - tokens := usage.TokenUsage{ - InputTokens: sqlUint(input), - OutputTokens: sqlUint(output), - CacheCreationInputTokens: sqlUint(cacheWrite), - CacheReadInputTokens: sqlUint(cacheRead), - ReasoningOutputTokens: sqlUint(reasoning), - } - if tokens.TotalTokens == 0 { - tokens.TotalTokens = totalUsage(tokens) - } - if !nonZero(tokens) { - return usage.Entry{}, false - } - entry := baseEntry(usage.ProviderHermes, timestamp, "hermes", "Hermes", sessionID, model, "Hermes Agent", tokens) - setSource(&entry, path, 0, 0, 0) - entry.ID = stableEntryID(entry, "hermes:"+sessionID) - return entry, true -} - -func hermesTimestamp(value any) (time.Time, bool) { - if parsed, ok := parseTimestamp(value); ok { - return parsed, true - } - if number, ok := sqlFloat(value); ok { - return timestampFromFloat(number) - } - if text := sqlString(value); text != "" { - return parseTimestampString(text) - } - return time.Time{}, false -} - -func existingSQLiteFile(path string) bool { - info, err := os.Stat(path) - return err == nil && !info.IsDir() -} diff --git a/internal/agentusage/kilo.go b/internal/agentusage/kilo.go deleted file mode 100644 index d06ff20..0000000 --- a/internal/agentusage/kilo.go +++ /dev/null @@ -1,146 +0,0 @@ -package agentusage - -import ( - "bytes" - "database/sql" - "encoding/json" - "errors" - "os" - "path/filepath" - "sort" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func loadKiloEntries(paths []string) ([]usage.Entry, error) { - dbPaths := sqliteDBPaths(paths, "kilo.db", nil) - entries := make([]usage.Entry, 0) - for _, dbPath := range dbPaths { - dbEntries, err := loadKiloDatabase(dbPath) - if err != nil { - return nil, err - } - entries = append(entries, dbEntries...) - } - sortEntries(entries) - return entries, nil -} - -func loadKiloDatabase(path string) ([]usage.Entry, error) { - db, err := openSQLite(path) - if errors.Is(err, os.ErrNotExist) { - return nil, nil - } - if err != nil { - return nil, err - } - defer db.Close() - - rows, err := db.Query("SELECT id, session_id, data FROM message") - if err != nil { - return nil, nil - } - defer rows.Close() - - entries := make([]usage.Entry, 0) - for rows.Next() { - var rowID, rowSessionID, data string - if err := rows.Scan(&rowID, &rowSessionID, &data); err != nil { - continue - } - if entry, ok := kiloMessageEntry(path, rowID, rowSessionID, data); ok { - entries = append(entries, entry) - } - } - return entries, rows.Err() -} - -func kiloMessageEntry(dbPath, rowID, rowSessionID, data string) (usage.Entry, bool) { - record := decodeJSONObjectString(data) - if record == nil || stringField(record, "role") != "assistant" { - return usage.Entry{}, false - } - tokenBlock := objectAt(record["tokens"]) - if tokenBlock == nil { - return usage.Entry{}, false - } - cache := objectAt(tokenBlock["cache"]) - tokens := usage.TokenUsage{ - InputTokens: uintField(tokenBlock, "input"), - OutputTokens: uintField(tokenBlock, "output"), - CacheCreationInputTokens: uintField(cache, "write"), - CacheReadInputTokens: uintField(cache, "read"), - ReasoningOutputTokens: uintField(tokenBlock, "reasoning"), - } - tokens = applyTotalFallback(tokens, uintField(tokenBlock, "total")) - if !nonZero(tokens) { - return usage.Entry{}, false - } - model := stringField(record, "modelID") - if model == "" { - return usage.Entry{}, false - } - timestamp, ok := parseTimestamp(objectAt(record["time"])["created"]) - if !ok { - return usage.Entry{}, false - } - sessionID := firstNonEmpty(stringField(record, "session_id"), rowSessionID, "unknown") - messageID := firstNonEmpty(stringField(record, "id"), rowID) - entry := baseEntry(usage.ProviderKilo, timestamp, "kilo", "Kilo", sessionID, model, "Kilo", tokens) - setSource(&entry, dbPath, 0, 0, 0) - entry.ID = stableEntryID(entry, messageID) - return entry, true -} - -func decodeJSONObjectString(data string) map[string]any { - decoder := json.NewDecoder(bytes.NewReader([]byte(data))) - decoder.UseNumber() - var record map[string]any - if err := decoder.Decode(&record); err != nil { - return nil - } - return record -} - -func sqliteDBPaths(paths []string, defaultFile string, extraNames func(string) bool) []string { - dbPaths := make([]string, 0) - for _, root := range paths { - info, err := os.Stat(root) - if err != nil { - continue - } - if !info.IsDir() { - if filepath.Base(root) == defaultFile || extraNames != nil && extraNames(filepath.Base(root)) { - dbPaths = append(dbPaths, root) - } - continue - } - candidate := filepath.Join(root, defaultFile) - if fileInfo, err := os.Stat(candidate); err == nil && !fileInfo.IsDir() { - dbPaths = append(dbPaths, candidate) - } - if extraNames == nil { - continue - } - entries, err := os.ReadDir(root) - if err != nil { - continue - } - for _, entry := range entries { - if entry.IsDir() || !extraNames(entry.Name()) { - continue - } - dbPaths = append(dbPaths, filepath.Join(root, entry.Name())) - } - } - sort.Strings(dbPaths) - return uniqueStrings(dbPaths) -} - -func scanAny(rows *sql.Rows, values ...*any) bool { - dest := make([]any, len(values)) - for i := range values { - dest[i] = values[i] - } - return rows.Scan(dest...) == nil -} diff --git a/internal/agentusage/kimi.go b/internal/agentusage/kimi.go deleted file mode 100644 index 4dc6a6e..0000000 --- a/internal/agentusage/kimi.go +++ /dev/null @@ -1,198 +0,0 @@ -package agentusage - -import ( - "path/filepath" - "sort" - "strings" - "time" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -const kimiDefaultModel = "kimi-for-coding" - -func loadKimiEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { - files := make([]string, 0) - for _, root := range paths { - files = append(files, collectFiles(filepath.Join(root, "sessions"), isKimiWireFile)...) - } - sort.Strings(files) - files = filterFiles(uniqueStrings(files), filter) - - entries := make([]usage.Entry, 0) - seen := make(map[string]bool) - for _, file := range files { - fileEntries, err := parseKimiWireFile(file) - if err != nil { - return nil, err - } - for _, entry := range fileEntries { - if seen[entry.ID] { - continue - } - seen[entry.ID] = true - entries = append(entries, entry) - } - } - sortEntries(entries) - return entries, nil -} - -func isKimiWireFile(path string) bool { - if filepath.Base(path) != "wire.jsonl" { - return false - } - parts := strings.Split(filepath.ToSlash(filepath.Clean(path)), "/") - for i := range parts { - if parts[i] != "sessions" { - continue - } - // Old layout: sessions///wire.jsonl - // New layout: sessions///agents//wire.jsonl - if i+3 == len(parts)-1 || i+5 == len(parts)-1 { - return true - } - } - return false -} - -// kimiRecord is one usable wire line, normalized across the old -// StatusUpdate format and the new Kimi Code usage.record format. -type kimiRecord struct { - tokens usage.TokenUsage - timestamp time.Time - hasTime bool - model string // empty means fall back to the config.json model - messageID string -} - -func parseKimiWireFile(path string) ([]usage.Entry, error) { - lines, err := readJSONLines(path, `usage`) - if err != nil { - return nil, err - } - configModel := kimiConfigModel(path) - sessionID := kimiSessionID(path) - fallback := fileModifiedTime(path) - entries := make([]usage.Entry, 0) - for _, line := range lines { - var record kimiRecord - var ok bool - if stringField(line.value, "type") == "usage.record" { - record, ok = parseKimiUsageRecord(line.value) - } else { - record, ok = parseKimiStatusUpdate(line.value) - } - if !ok { - continue - } - timestamp := record.timestamp - if !record.hasTime { - timestamp = fallback - } - model := record.model - if model == "" { - model = configModel - } - entry := baseEntry(usage.ProviderKimi, timestamp, "kimi", "Kimi", sessionID, model, "Kimi", record.tokens) - setSource(&entry, path, line.line, line.start, line.end) - entry.ID = stableEntryID(entry, record.messageID) - entries = append(entries, entry) - } - return entries, nil -} - -func parseKimiStatusUpdate(value map[string]any) (kimiRecord, bool) { - message := objectAt(value["message"]) - if stringField(message, "type") != "StatusUpdate" { - return kimiRecord{}, false - } - payload := objectAt(message["payload"]) - tokenUsage := objectAt(payload["token_usage"]) - if tokenUsage == nil { - return kimiRecord{}, false - } - tokens := usage.TokenUsage{ - InputTokens: uintField(tokenUsage, "input_other"), - OutputTokens: uintField(tokenUsage, "output"), - CacheCreationInputTokens: uintField(tokenUsage, "input_cache_creation"), - CacheReadInputTokens: uintField(tokenUsage, "input_cache_read"), - } - tokens = applyTotalFallback(tokens, uintField(tokenUsage, "total")) - if !nonZero(tokens) { - return kimiRecord{}, false - } - record := kimiRecord{tokens: tokens, messageID: stringField(payload, "message_id")} - record.timestamp, record.hasTime = parseTimestamp(value["timestamp"]) - return record, true -} - -func parseKimiUsageRecord(value map[string]any) (kimiRecord, bool) { - // Session-scoped records are cumulative totals; only turn records count. - if stringField(value, "usageScope") != "turn" { - return kimiRecord{}, false - } - tokenUsage := objectAt(value["usage"]) - if tokenUsage == nil { - return kimiRecord{}, false - } - tokens := usage.TokenUsage{ - InputTokens: uintField(tokenUsage, "inputOther"), - OutputTokens: uintField(tokenUsage, "output"), - CacheCreationInputTokens: uintField(tokenUsage, "inputCacheCreation"), - CacheReadInputTokens: uintField(tokenUsage, "inputCacheRead"), - } - tokens = applyTotalFallback(tokens, 0) - if !nonZero(tokens) { - return kimiRecord{}, false - } - record := kimiRecord{ - tokens: tokens, - model: strings.TrimPrefix(stringField(value, "model"), "kimi-code/"), - } - record.timestamp, record.hasTime = parseTimestamp(value["time"]) - return record, true -} - -// kimiSessionID returns the session directory name for either layout. -func kimiSessionID(path string) string { - dir := filepath.Dir(path) - if filepath.Base(filepath.Dir(dir)) == "agents" { - dir = filepath.Dir(filepath.Dir(dir)) - } - sessionID := filepath.Base(dir) - if sessionID == "" || sessionID == "." { - return "unknown" - } - return sessionID -} - -// kimiRoot walks up from a wire file to the directory containing "sessions", -// which is the Kimi data root regardless of layout depth. -func kimiRoot(path string) string { - for dir := filepath.Dir(path); ; { - parent := filepath.Dir(dir) - if filepath.Base(dir) == "sessions" { - return parent - } - if parent == dir { - return "" - } - dir = parent - } -} - -func kimiConfigModel(path string) string { - root := kimiRoot(path) - if root == "" { - return kimiDefaultModel - } - config, err := readJSONObject(filepath.Join(root, "config.json")) - if err != nil || config == nil { - return kimiDefaultModel - } - if model := stringField(config, "model"); model != "" { - return model - } - return kimiDefaultModel -} diff --git a/internal/agentusage/openclaw.go b/internal/agentusage/openclaw.go deleted file mode 100644 index 30bf4d3..0000000 --- a/internal/agentusage/openclaw.go +++ /dev/null @@ -1,132 +0,0 @@ -package agentusage - -import ( - "path/filepath" - "sort" - "strings" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func loadOpenClawEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { - files := make([]string, 0) - for _, path := range paths { - files = append(files, collectFiles(path, isOpenClawSessionFile)...) - } - sort.Strings(files) - files = filterFiles(files, filter) - - entries := make([]usage.Entry, 0) - for _, file := range files { - fileEntries, err := parseOpenClawSessionFile(file) - if err != nil { - return nil, err - } - entries = append(entries, fileEntries...) - } - sortEntries(entries) - return entries, nil -} - -func isOpenClawSessionFile(path string) bool { - name := filepath.Base(path) - index := strings.Index(name, ".jsonl") - if index < 0 { - return false - } - suffix := name[index:] - return suffix == ".jsonl" || strings.HasPrefix(suffix, ".jsonl.deleted.") || strings.HasPrefix(suffix, ".jsonl.reset.") -} - -func parseOpenClawSessionFile(path string) ([]usage.Entry, error) { - lines, err := readJSONLines(path) - if err != nil { - return nil, err - } - sessionID := openClawSessionID(path) - currentModel := "" - currentProvider := "" - entries := make([]usage.Entry, 0) - for _, line := range lines { - record := line.value - if isOpenClawModelChange(record) { - source := objectAt(record["data"]) - if source == nil { - source = record - } - if model := firstStringField(source, "modelId", "model"); model != "" { - currentModel = model - } - if provider := stringField(source, "provider"); provider != "" { - currentProvider = provider - } - continue - } - if stringField(record, "type") != "message" { - continue - } - message := objectAt(record["message"]) - if stringField(message, "role") != "assistant" { - continue - } - usageBlock := objectAt(message["usage"]) - if usageBlock == nil { - continue - } - timestamp, ok := parseTimestamp(message["timestamp"]) - if !ok { - timestamp, ok = parseTimestamp(record["timestamp"]) - } - if !ok { - timestamp = fileModifiedTime(path) - } - model := firstStringField(message, "modelId", "model") - if model == "" { - model = currentModel - } - if model == "" { - model = "unknown" - } - provider := stringField(message, "provider") - if provider == "" { - provider = currentProvider - } - tokens := usage.TokenUsage{ - InputTokens: uintField(usageBlock, "input"), - OutputTokens: uintField(usageBlock, "output"), - CacheCreationInputTokens: uintField(usageBlock, "cacheWrite"), - CacheReadInputTokens: uintField(usageBlock, "cacheRead"), - } - tokens = applyTotalFallback(tokens, uintField(usageBlock, "totalTokens")) - if !nonZero(tokens) { - continue - } - entry := baseEntry(usage.ProviderOpenClaw, timestamp, "openclaw", "OpenClaw", sessionID, "[openclaw] "+model, "OpenClaw", tokens) - setSource(&entry, path, line.line, line.start, line.end) - entry.ID = stableEntryID(entry, provider) - entries = append(entries, entry) - } - return entries, nil -} - -func isOpenClawModelChange(record map[string]any) bool { - if stringField(record, "type") == "model_change" { - return true - } - return stringField(record, "type") == "custom" && stringField(record, "customType") == "model-snapshot" -} - -func openClawSessionID(path string) string { - name := filepath.Base(path) - index := strings.Index(name, ".jsonl") - if index < 0 { - if name == "" { - return "unknown" - } - return name - } - if index == 0 { - return name - } - return name[:index] -} diff --git a/internal/agentusage/pi.go b/internal/agentusage/pi.go deleted file mode 100644 index c06812b..0000000 --- a/internal/agentusage/pi.go +++ /dev/null @@ -1,96 +0,0 @@ -package agentusage - -import ( - "path/filepath" - "sort" - "strings" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func loadPiEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { - files := make([]string, 0) - for _, path := range paths { - files = append(files, collectExt(path, ".jsonl")...) - } - sort.Strings(files) - files = filterFiles(files, filter) - - entries := make([]usage.Entry, 0) - for _, file := range files { - fileEntries, err := parsePiSessionFile(file) - if err != nil { - return nil, err - } - entries = append(entries, fileEntries...) - } - sortEntries(entries) - return entries, nil -} - -func parsePiSessionFile(path string) ([]usage.Entry, error) { - lines, err := readJSONLines(path, `"usage"`, `"message"`) - if err != nil { - return nil, err - } - project := piProject(path) - sessionID := piSessionID(path) - entries := make([]usage.Entry, 0) - for _, line := range lines { - if typ := stringField(line.value, "type"); typ != "" && typ != "message" { - continue - } - message := objectAt(line.value["message"]) - if stringField(message, "role") != "assistant" { - continue - } - usageBlock := objectAt(message["usage"]) - if usageBlock == nil { - continue - } - timestamp, ok := parseTimestamp(line.value["timestamp"]) - if !ok { - continue - } - tokens := usage.TokenUsage{ - InputTokens: uintField(usageBlock, "input"), - OutputTokens: uintField(usageBlock, "output"), - CacheCreationInputTokens: uintField(usageBlock, "cacheWrite"), - CacheReadInputTokens: uintField(usageBlock, "cacheRead"), - } - tokens = applyTotalFallback(tokens, uintField(usageBlock, "totalTokens")) - if !nonZero(tokens) { - continue - } - model := stringField(message, "model") - if model != "" { - model = "[pi] " + model - } - entry := baseEntry(usage.ProviderPi, timestamp, project, project, sessionID, model, "pi-agent", tokens) - setSource(&entry, path, line.line, line.start, line.end) - entry.ID = stableEntryID(entry) - entries = append(entries, entry) - } - return entries, nil -} - -func piSessionID(path string) string { - stem := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) - if before, after, ok := strings.Cut(stem, "_"); ok && before != "" && after != "" { - return after - } - if stem == "" { - return "unknown" - } - return stem -} - -func piProject(path string) string { - parts := strings.Split(filepath.ToSlash(filepath.Clean(path)), "/") - for i, part := range parts { - if part == "sessions" && i+1 < len(parts) && parts[i+1] != "" { - return parts[i+1] - } - } - return usage.UnknownProject -} diff --git a/internal/agentusage/providers.go b/internal/agentusage/providers.go deleted file mode 100644 index a290431..0000000 --- a/internal/agentusage/providers.go +++ /dev/null @@ -1,376 +0,0 @@ -package agentusage - -import ( - "sort" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" - "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" -) - -func sortEntriesByTimestampDesc(entries []usage.Entry) []usage.Entry { - sort.Slice(entries, func(i, j int) bool { - return entries[i].Timestamp.After(entries[j].Timestamp) - }) - return entries -} - -// providerBase carries the scan configuration shared by every agent -// provider. filter skips source files whose events are already ingested; -// the SQLite-backed providers (Kilo, Hermes, Goose) deliberately do not -// implement WithFileFilter because WAL keeps the main database file's stat -// unchanged while data grows in the -wal journal. -type providerBase struct { - paths []string - filter usage.FileFilter -} - -func newProviderBase(paths []string) providerBase { - return providerBase{paths: append([]string{}, paths...)} -} - -// CopilotProvider loads GitHub Copilot CLI usage entries. -type CopilotProvider struct{ providerBase } - -// GeminiProvider loads Gemini CLI usage entries. -type GeminiProvider struct{ providerBase } - -// KimiProvider loads Kimi usage entries. -type KimiProvider struct{ providerBase } - -// QwenProvider loads Qwen usage entries. -type QwenProvider struct{ providerBase } - -// OpenClawProvider loads OpenClaw usage entries. -type OpenClawProvider struct{ providerBase } - -// PiProvider loads pi-agent usage entries. -type PiProvider struct{ providerBase } - -// AmpProvider loads Amp usage entries. -type AmpProvider struct{ providerBase } - -// DroidProvider loads Droid usage entries. -type DroidProvider struct{ providerBase } - -// KiloProvider loads Kilo usage entries. -type KiloProvider struct{ providerBase } - -// HermesProvider loads Hermes Agent usage entries. -type HermesProvider struct{ providerBase } - -// CodebuffProvider loads Codebuff usage entries. -type CodebuffProvider struct{ providerBase } - -// OpenCodeProvider loads OpenCode usage entries. -type OpenCodeProvider struct{ providerBase } - -// GooseProvider loads Goose usage entries. -type GooseProvider struct{ providerBase } - -var ( - _ usageprovider.Provider = CopilotProvider{} - _ usageprovider.Provider = GeminiProvider{} - _ usageprovider.Provider = KimiProvider{} - _ usageprovider.Provider = QwenProvider{} - _ usageprovider.Provider = OpenClawProvider{} - _ usageprovider.Provider = PiProvider{} - _ usageprovider.Provider = AmpProvider{} - _ usageprovider.Provider = DroidProvider{} - _ usageprovider.Provider = KiloProvider{} - _ usageprovider.Provider = HermesProvider{} - _ usageprovider.Provider = CodebuffProvider{} - _ usageprovider.Provider = OpenCodeProvider{} - _ usageprovider.Provider = GooseProvider{} -) - -// WithPaths returns a GitHub Copilot CLI provider configured with data roots. -func (CopilotProvider) WithPaths(paths []string) usageprovider.Provider { - return CopilotProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the GitHub Copilot CLI provider id. -func (CopilotProvider) Provider() usage.Provider { return usage.ProviderCopilot } - -// WithFileFilter returns a GitHub Copilot CLI provider that skips source -// files the filter rejects. -func (p CopilotProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized GitHub Copilot CLI usage entries, newest first. -func (p CopilotProvider) Entries() ([]usage.Entry, error) { - entries, err := loadCopilotEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns a Gemini CLI provider configured with data roots. -func (GeminiProvider) WithPaths(paths []string) usageprovider.Provider { - return GeminiProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the Gemini CLI provider id. -func (GeminiProvider) Provider() usage.Provider { return usage.ProviderGemini } - -// WithFileFilter returns a Gemini CLI provider that skips source files the -// filter rejects. -func (p GeminiProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized Gemini CLI usage entries, newest first. -func (p GeminiProvider) Entries() ([]usage.Entry, error) { - entries, err := loadGeminiEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns a Kimi provider configured with data roots. -func (KimiProvider) WithPaths(paths []string) usageprovider.Provider { - return KimiProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the Kimi provider id. -func (KimiProvider) Provider() usage.Provider { return usage.ProviderKimi } - -// WithFileFilter returns a Kimi provider that skips source files the filter -// rejects. -func (p KimiProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized Kimi usage entries, newest first. -func (p KimiProvider) Entries() ([]usage.Entry, error) { - entries, err := loadKimiEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns a Qwen provider configured with data roots. -func (QwenProvider) WithPaths(paths []string) usageprovider.Provider { - return QwenProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the Qwen provider id. -func (QwenProvider) Provider() usage.Provider { return usage.ProviderQwen } - -// WithFileFilter returns a Qwen provider that skips source files the filter -// rejects. -func (p QwenProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized Qwen usage entries, newest first. -func (p QwenProvider) Entries() ([]usage.Entry, error) { - entries, err := loadQwenEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns an OpenClaw provider configured with data roots. -func (OpenClawProvider) WithPaths(paths []string) usageprovider.Provider { - return OpenClawProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the OpenClaw provider id. -func (OpenClawProvider) Provider() usage.Provider { return usage.ProviderOpenClaw } - -// WithFileFilter returns an OpenClaw provider that skips source files the -// filter rejects. -func (p OpenClawProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized OpenClaw usage entries, newest first. -func (p OpenClawProvider) Entries() ([]usage.Entry, error) { - entries, err := loadOpenClawEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns a pi-agent provider configured with data roots. -func (PiProvider) WithPaths(paths []string) usageprovider.Provider { - return PiProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the pi-agent provider id. -func (PiProvider) Provider() usage.Provider { return usage.ProviderPi } - -// WithFileFilter returns a pi-agent provider that skips source files the -// filter rejects. -func (p PiProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized pi-agent usage entries, newest first. -func (p PiProvider) Entries() ([]usage.Entry, error) { - entries, err := loadPiEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns an Amp provider configured with data roots. -func (AmpProvider) WithPaths(paths []string) usageprovider.Provider { - return AmpProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the Amp provider id. -func (AmpProvider) Provider() usage.Provider { return usage.ProviderAmp } - -// WithFileFilter returns an Amp provider that skips source files the filter -// rejects. -func (p AmpProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized Amp usage entries, newest first. -func (p AmpProvider) Entries() ([]usage.Entry, error) { - entries, err := loadAmpEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns a Droid provider configured with data roots. -func (DroidProvider) WithPaths(paths []string) usageprovider.Provider { - return DroidProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the Droid provider id. -func (DroidProvider) Provider() usage.Provider { return usage.ProviderDroid } - -// WithFileFilter returns a Droid provider that skips source files the filter -// rejects. -func (p DroidProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized Droid usage entries, newest first. -func (p DroidProvider) Entries() ([]usage.Entry, error) { - entries, err := loadDroidEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns a Kilo provider configured with data roots. -func (KiloProvider) WithPaths(paths []string) usageprovider.Provider { - return KiloProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the Kilo provider id. -func (KiloProvider) Provider() usage.Provider { return usage.ProviderKilo } - -// Entries loads normalized Kilo usage entries, newest first. -func (p KiloProvider) Entries() ([]usage.Entry, error) { - entries, err := loadKiloEntries(p.paths) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns a Hermes provider configured with data roots. -func (HermesProvider) WithPaths(paths []string) usageprovider.Provider { - return HermesProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the Hermes provider id. -func (HermesProvider) Provider() usage.Provider { return usage.ProviderHermes } - -// Entries loads normalized Hermes Agent usage entries, newest first. -func (p HermesProvider) Entries() ([]usage.Entry, error) { - entries, err := loadHermesEntries(p.paths) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns a Codebuff provider configured with data roots. -func (CodebuffProvider) WithPaths(paths []string) usageprovider.Provider { - return CodebuffProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the Codebuff provider id. -func (CodebuffProvider) Provider() usage.Provider { return usage.ProviderCodebuff } - -// WithFileFilter returns a Codebuff provider that skips source files the -// filter rejects. -func (p CodebuffProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized Codebuff usage entries, newest first. -func (p CodebuffProvider) Entries() ([]usage.Entry, error) { - entries, err := loadCodebuffEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns an OpenCode provider configured with data roots. -func (OpenCodeProvider) WithPaths(paths []string) usageprovider.Provider { - return OpenCodeProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the OpenCode provider id. -func (OpenCodeProvider) Provider() usage.Provider { return usage.ProviderOpenCode } - -// WithFileFilter returns an OpenCode provider that skips source files the -// filter rejects. The filter applies to message files only; the OpenCode -// database goes through the SQLite path and is always scanned. -func (p OpenCodeProvider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { - p.filter = filter - return p -} - -// Entries loads normalized OpenCode usage entries, newest first. -func (p OpenCodeProvider) Entries() ([]usage.Entry, error) { - entries, err := loadOpenCodeEntries(p.paths, p.filter) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} - -// WithPaths returns a Goose provider configured with data roots. -func (GooseProvider) WithPaths(paths []string) usageprovider.Provider { - return GooseProvider{providerBase: newProviderBase(paths)} -} - -// Provider returns the Goose provider id. -func (GooseProvider) Provider() usage.Provider { return usage.ProviderGoose } - -// Entries loads normalized Goose usage entries, newest first. -func (p GooseProvider) Entries() ([]usage.Entry, error) { - entries, err := loadGooseEntries(p.paths) - if err != nil { - return nil, err - } - return sortEntriesByTimestampDesc(entries), nil -} diff --git a/internal/agentusage/providers_test.go b/internal/agentusage/providers_test.go deleted file mode 100644 index 2e06097..0000000 --- a/internal/agentusage/providers_test.go +++ /dev/null @@ -1,440 +0,0 @@ -package agentusage - -import ( - "database/sql" - "os" - "path/filepath" - "testing" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func TestProviderFileFilterSkipsRejectedFiles(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "copilot.jsonl") - writeFile(t, path, `{"type":"span","traceId":"trace-1","spanId":"span-1","name":"chat claude-sonnet-4","endTime":[1775934264,967317833],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"claude-sonnet-4","gen_ai.conversation.id":"conv-1","gen_ai.usage.input_tokens":19452,"gen_ai.usage.output_tokens":281}}}`+"\n") - - rejected := make([]string, 0) - provider := CopilotProvider{}.WithPaths([]string{dir}).(CopilotProvider). - WithFileFilter(func(candidate string) bool { - rejected = append(rejected, candidate) - return false - }) - entries, err := provider.Entries() - if err != nil { - t.Fatal(err) - } - if len(entries) != 0 { - t.Fatalf("len(entries) = %d, want 0 when the filter rejects every file", len(entries)) - } - if len(rejected) != 1 || rejected[0] != path { - t.Fatalf("filter saw %#v, want the session file", rejected) - } -} - -func TestProvidersLoadEntries(t *testing.T) { - tests := []struct { - name string - provider func() ([]usage.Entry, error) - want usage.Provider - model string - sessionID string - project string - tokens usage.TokenUsage - }{ - { - name: "copilot", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - path := filepath.Join(dir, "copilot.jsonl") - writeFile(t, path, `{"type":"span","traceId":"trace-1","spanId":"span-1","name":"chat claude-sonnet-4","endTime":[1775934264,967317833],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"claude-sonnet-4","gen_ai.conversation.id":"conv-1","gen_ai.usage.input_tokens":19452,"gen_ai.usage.output_tokens":281,"gen_ai.usage.cache_read.input_tokens":123,"gen_ai.usage.cache_creation.input_tokens":25,"gen_ai.usage.reasoning.output_tokens":128}}}`+"\n") - return CopilotProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderCopilot, - model: "claude-sonnet-4", - sessionID: "conv-1", - project: "copilot", - tokens: usage.TokenUsage{ - InputTokens: 19329, - OutputTokens: 281, - CacheCreationInputTokens: 25, - CacheReadInputTokens: 123, - ReasoningOutputTokens: 128, - TotalTokens: 19886, - }, - }, - { - name: "gemini", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - path := filepath.Join(dir, "session-a.jsonl") - writeFile(t, path, - `{"sessionId":"session-a","projectHash":"project-a","startTime":"2026-05-17T11:07:00.000Z"}`+"\n"+ - `{"id":"msg-a","timestamp":"2026-05-17T11:07:32.000Z","type":"gemini","model":"gemini-3-flash-preview","tokens":{"input":15327,"output":23,"cached":11526,"thoughts":919,"tool":7,"total":16276}}`+"\n") - return GeminiProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderGemini, - model: "gemini-3-flash-preview", - sessionID: "session-a", - project: "gemini", - tokens: usage.TokenUsage{ - InputTokens: 3808, - OutputTokens: 23, - CacheReadInputTokens: 11526, - ReasoningOutputTokens: 919, - TotalTokens: 16276, - }, - }, - { - name: "kimi", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - writeFile(t, filepath.Join(dir, "config.json"), `{"model":"kimi-k2"}`) - path := filepath.Join(dir, "sessions", "group", "session-a", "wire.jsonl") - writeFile(t, path, - `{"type":"metadata","protocol_version":"1.3"}`+"\n"+ - `{"timestamp":1770983427.123,"message":{"type":"StatusUpdate","payload":{"token_usage":{"input_other":100,"output":50,"input_cache_read":10,"input_cache_creation":20},"message_id":"msg-1"}}}`+"\n") - return KimiProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderKimi, - model: "kimi-k2", - sessionID: "session-a", - project: "kimi", - tokens: usage.TokenUsage{ - InputTokens: 100, - OutputTokens: 50, - CacheCreationInputTokens: 20, - CacheReadInputTokens: 10, - TotalTokens: 180, - }, - }, - { - name: "kimi-code", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - path := filepath.Join(dir, "sessions", "wd_beeec_fe4684529abd", "session-b", "agents", "main", "wire.jsonl") - writeFile(t, path, - `{"type":"metadata","protocol_version":"1.4","created_at":1785336260355}`+"\n"+ - `{"type":"usage.record","model":"kimi-code/kimi-for-coding","usage":{"inputOther":3064,"output":76,"inputCacheRead":14848,"inputCacheCreation":0},"usageScope":"turn","time":1782113184943}`+"\n"+ - `{"type":"usage.record","model":"kimi-code/kimi-for-coding","usage":{"inputOther":5000,"output":200,"inputCacheRead":20000,"inputCacheCreation":100},"usageScope":"session","time":1782113185000}`+"\n") - return KimiProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderKimi, - model: "kimi-for-coding", - sessionID: "session-b", - project: "kimi", - tokens: usage.TokenUsage{ - InputTokens: 3064, - OutputTokens: 76, - CacheReadInputTokens: 14848, - TotalTokens: 17988, - }, - }, - { - name: "qwen", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - path := filepath.Join(dir, "projects", "project-a", "chats", "chat-a.jsonl") - writeFile(t, path, `{"type":"assistant","timestamp":"2026-01-02T00:00:00.000Z","sessionId":"session-a","model":"qwen3-coder","usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":20,"thoughtsTokenCount":5,"cachedContentTokenCount":3,"totalTokenCount":38}}`+"\n") - return QwenProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderQwen, - model: "qwen3-coder", - sessionID: "session-a", - project: "qwen", - tokens: usage.TokenUsage{ - InputTokens: 10, - OutputTokens: 20, - CacheReadInputTokens: 3, - ReasoningOutputTokens: 5, - TotalTokens: 38, - }, - }, - { - name: "openclaw", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - path := filepath.Join(dir, "agents", "main", "sessions", "abc.jsonl") - writeFile(t, path, - `{"type":"model_change","provider":"openai-codex","modelId":"gpt-5.2"}`+"\n"+ - `{"type":"message","message":{"role":"assistant","usage":{"input":1660,"output":55,"cacheRead":108928},"timestamp":1769753935279}}`+"\n") - return OpenClawProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderOpenClaw, - model: "[openclaw] gpt-5.2", - sessionID: "abc", - project: "openclaw", - tokens: usage.TokenUsage{ - InputTokens: 1660, - OutputTokens: 55, - CacheReadInputTokens: 108928, - TotalTokens: 110643, - }, - }, - { - name: "pi", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - path := filepath.Join(dir, "project-a", "agent_session-a.jsonl") - writeFile(t, path, `{"type":"message","timestamp":"2026-01-02T00:00:00.000Z","message":{"role":"assistant","model":"gpt-5","usage":{"totalTokens":333}}}`+"\n") - return PiProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderPi, - model: "[pi] gpt-5", - sessionID: "session-a", - project: usage.UnknownProject, - tokens: usage.TokenUsage{ - OutputTokens: 333, - TotalTokens: 333, - }, - }, - { - name: "amp", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - path := filepath.Join(dir, "threads", "thread.json") - writeFile(t, path, `{"id":"thread-a","usageLedger":{"events":[{"id":"event-a","timestamp":"2026-01-02T00:00:00.000Z","model":"gpt-5","tokens":{"input":1,"output":2}}]}}`) - return AmpProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderAmp, - model: "gpt-5", - sessionID: "thread-a", - project: "amp", - tokens: usage.TokenUsage{ - InputTokens: 1, - OutputTokens: 2, - TotalTokens: 3, - }, - }, - { - name: "droid", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - path := filepath.Join(dir, "session-a.settings.json") - writeFile(t, path, `{"model":"Claude-Sonnet-4-[Anthropic]","providerLock":"anthropic","providerLockTimestamp":"2026-01-02T00:00:00.000Z","tokenUsage":{"inputTokens":100,"outputTokens":50,"cacheCreationTokens":20,"cacheReadTokens":10,"thinkingTokens":5}}`) - return DroidProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderDroid, - model: "claude-sonnet-4", - sessionID: "session-a", - project: "droid", - tokens: usage.TokenUsage{ - InputTokens: 100, - OutputTokens: 50, - CacheCreationInputTokens: 20, - CacheReadInputTokens: 10, - ReasoningOutputTokens: 5, - TotalTokens: 185, - }, - }, - { - name: "kilo", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - dbPath := filepath.Join(dir, "kilo.db") - db := openTestSQLite(t, dbPath) - defer db.Close() - execSQL(t, db, "CREATE TABLE message (id TEXT, session_id TEXT, data TEXT)") - execSQL(t, db, `INSERT INTO message (id, session_id, data) VALUES (?, ?, ?)`, - "row-1", - "session-a", - `{"id":"msg-1","role":"assistant","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","time":{"created":1767312000000},"tokens":{"input":100,"output":50,"reasoning":5,"cache":{"read":10,"write":20}}}`, - ) - return KiloProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderKilo, - model: "claude-sonnet-4-20250514", - sessionID: "session-a", - project: "kilo", - tokens: usage.TokenUsage{ - InputTokens: 100, - OutputTokens: 50, - CacheCreationInputTokens: 20, - CacheReadInputTokens: 10, - ReasoningOutputTokens: 5, - TotalTokens: 185, - }, - }, - { - name: "hermes", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - dbPath := filepath.Join(dir, "state.db") - db := openTestSQLite(t, dbPath) - defer db.Close() - execSQL(t, db, `CREATE TABLE sessions ( - id TEXT PRIMARY KEY, - model TEXT, - billing_provider TEXT, - started_at REAL, - message_count INTEGER, - input_tokens INTEGER, - output_tokens INTEGER, - cache_read_tokens INTEGER, - cache_write_tokens INTEGER, - reasoning_tokens INTEGER, - estimated_cost_usd REAL, - actual_cost_usd REAL - )`) - execSQL(t, db, `INSERT INTO sessions ( - id, model, billing_provider, started_at, message_count, input_tokens, - output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, - estimated_cost_usd, actual_cost_usd - ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, - "session-a", "gpt-5.5", "openai", 1750000000.25, 42, 100, 50, 10, 20, 5, 0.12, 0.34, - ) - return HermesProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderHermes, - model: "gpt-5.5", - sessionID: "session-a", - project: "hermes", - tokens: usage.TokenUsage{ - InputTokens: 100, - OutputTokens: 50, - CacheCreationInputTokens: 20, - CacheReadInputTokens: 10, - ReasoningOutputTokens: 5, - TotalTokens: 185, - }, - }, - { - name: "codebuff", - provider: func() ([]usage.Entry, error) { - root := filepath.Join(t.TempDir(), "manicode") - path := filepath.Join(root, "projects", "project-a", "chats", "2026-01-02T03-04-05.000Z", "chat-messages.json") - writeFile(t, path, `[{"id":"assistant-message","role":"assistant","timestamp":"2026-01-02T03:04:06.000Z","metadata":{"model":"claude-sonnet-4-20250514","usage":{"inputTokens":100,"outputTokens":50,"cacheCreationInputTokens":20,"cacheReadInputTokens":10}}}]`) - return CodebuffProvider{}.WithPaths([]string{root}).Entries() - }, - want: usage.ProviderCodebuff, - model: "claude-sonnet-4-20250514", - sessionID: "manicode/project-a/2026-01-02T03-04-05.000Z", - project: "codebuff", - tokens: usage.TokenUsage{ - InputTokens: 100, - OutputTokens: 50, - CacheCreationInputTokens: 20, - CacheReadInputTokens: 10, - TotalTokens: 180, - }, - }, - { - name: "opencode", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - path := filepath.Join(dir, "storage", "message", "session-a", "msg-1.json") - writeFile(t, path, `{"id":"msg-1","sessionID":"session-a","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","path":{"cwd":"/repo/demo"},"time":{"created":1767312000000},"tokens":{"input":100,"output":50,"cache":{"read":10,"write":20}},"cost":0}`) - return OpenCodeProvider{}.WithPaths([]string{dir}).Entries() - }, - want: usage.ProviderOpenCode, - model: "claude-sonnet-4-20250514", - sessionID: "session-a", - project: "demo", - tokens: usage.TokenUsage{ - InputTokens: 100, - OutputTokens: 50, - CacheCreationInputTokens: 20, - CacheReadInputTokens: 10, - TotalTokens: 180, - }, - }, - { - name: "goose", - provider: func() ([]usage.Entry, error) { - dir := t.TempDir() - dbPath := filepath.Join(dir, "sessions.db") - db := openTestSQLite(t, dbPath) - defer db.Close() - execSQL(t, db, `CREATE TABLE sessions ( - id TEXT PRIMARY KEY, - model_config_json TEXT, - provider_name TEXT, - created_at TEXT, - total_tokens INTEGER, - input_tokens INTEGER, - output_tokens INTEGER, - accumulated_total_tokens INTEGER, - accumulated_input_tokens INTEGER, - accumulated_output_tokens INTEGER - )`) - execSQL(t, db, `INSERT INTO sessions ( - id, model_config_json, provider_name, created_at, - accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens - ) VALUES (?, ?, ?, ?, ?, ?, ?)`, - "session-a", `{"model_name":"claude-sonnet-4-20250514"}`, "anthropic", "2026-05-01 01:02:03", 180, 100, 50, - ) - return GooseProvider{}.WithPaths([]string{dbPath}).Entries() - }, - want: usage.ProviderGoose, - model: "claude-sonnet-4-20250514", - sessionID: "session-a", - project: "goose", - tokens: usage.TokenUsage{ - InputTokens: 100, - OutputTokens: 50, - ReasoningOutputTokens: 30, - TotalTokens: 180, - }, - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - entries, err := tt.provider() - if err != nil { - t.Fatal(err) - } - if len(entries) != 1 { - t.Fatalf("entries = %d, want 1: %#v", len(entries), entries) - } - entry := entries[0] - if entry.Provider != tt.want { - t.Fatalf("provider = %q, want %q", entry.Provider, tt.want) - } - if entry.Model != tt.model { - t.Fatalf("model = %q, want %q", entry.Model, tt.model) - } - if entry.SessionID != tt.sessionID { - t.Fatalf("session id = %q, want %q", entry.SessionID, tt.sessionID) - } - if entry.Project != tt.project { - t.Fatalf("project = %q, want %q", entry.Project, tt.project) - } - if entry.Usage != tt.tokens { - t.Fatalf("usage = %#v, want %#v", entry.Usage, tt.tokens) - } - if entry.ID == "" { - t.Fatal("ID is empty") - } - }) - } -} - -func writeFile(t *testing.T, path, data string) { - t.Helper() - if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { - t.Fatal(err) - } - if err := os.WriteFile(path, []byte(data), 0o600); err != nil { - t.Fatal(err) - } -} - -func openTestSQLite(t *testing.T, path string) *sql.DB { - t.Helper() - if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { - t.Fatal(err) - } - db, err := openSQLite(path) - if err != nil { - t.Fatal(err) - } - return db -} - -func execSQL(t *testing.T, db *sql.DB, statement string, args ...any) { - t.Helper() - if _, err := db.Exec(statement, args...); err != nil { - t.Fatal(err) - } -} diff --git a/internal/agentusage/qwen.go b/internal/agentusage/qwen.go deleted file mode 100644 index 34e01a5..0000000 --- a/internal/agentusage/qwen.go +++ /dev/null @@ -1,104 +0,0 @@ -package agentusage - -import ( - "path/filepath" - "sort" - "strings" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func loadQwenEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { - files := make([]string, 0) - for _, root := range paths { - files = append(files, collectFiles(filepath.Join(root, "projects"), isQwenChatFile)...) - if strings.Contains(filepath.ToSlash(root), "/projects/") { - files = append(files, collectFiles(root, isQwenChatFile)...) - } - } - sort.Strings(files) - files = filterFiles(uniqueStrings(files), filter) - - entries := make([]usage.Entry, 0) - for _, file := range files { - fileEntries, err := parseQwenChatFile(file) - if err != nil { - return nil, err - } - entries = append(entries, fileEntries...) - } - sortEntries(entries) - return entries, nil -} - -func isQwenChatFile(path string) bool { - if !strings.EqualFold(filepath.Ext(path), ".jsonl") { - return false - } - parts := strings.Split(filepath.ToSlash(filepath.Clean(path)), "/") - for i := 0; i+3 < len(parts); i++ { - if parts[i] == "projects" && parts[i+2] == "chats" { - return true - } - } - return false -} - -func parseQwenChatFile(path string) ([]usage.Entry, error) { - lines, err := readJSONLines(path, `"usageMetadata"`) - if err != nil { - return nil, err - } - project := qwenProject(path) - fallback := fileModifiedTime(path) - entries := make([]usage.Entry, 0) - for _, line := range lines { - record := line.value - if stringField(record, "type") != "assistant" { - continue - } - meta := objectAt(record["usageMetadata"]) - if meta == nil { - continue - } - timestamp, ok := parseTimestamp(record["timestamp"]) - if !ok { - timestamp = fallback - } - sessionID := stringField(record, "sessionId") - if sessionID == "" { - sessionID = project + "-" + strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) - } - model := stringField(record, "model") - if model == "" { - model = "unknown" - } - tokens := usage.TokenUsage{ - InputTokens: uintField(meta, "promptTokenCount"), - OutputTokens: uintField(meta, "candidatesTokenCount"), - CacheReadInputTokens: uintField(meta, "cachedContentTokenCount"), - ReasoningOutputTokens: uintField(meta, - "thoughtsTokenCount", - ), - } - tokens = applyTotalFallback(tokens, uintField(meta, "totalTokenCount")) - if !nonZero(tokens) { - continue - } - entry := baseEntry(usage.ProviderQwen, timestamp, "qwen", project, sessionID, model, "Qwen", tokens) - setSource(&entry, path, line.line, line.start, line.end) - entry.ID = stableEntryID(entry) - entries = append(entries, entry) - } - return entries, nil -} - -func qwenProject(path string) string { - parts := strings.Split(filepath.ToSlash(filepath.Clean(path)), "/") - for i := 0; i+3 < len(parts); i++ { - if parts[i] == "projects" && parts[i+2] == "chats" && parts[i+1] != "" { - return parts[i+1] - } - } - return usage.UnknownProject -} diff --git a/internal/provider/amp/loader.go b/internal/provider/amp/loader.go new file mode 100644 index 0000000..7005381 --- /dev/null +++ b/internal/provider/amp/loader.go @@ -0,0 +1,183 @@ +package amp + +import ( + "encoding/json" + "os" + "path/filepath" + "sort" + "strconv" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + files := make([]string, 0) + for _, path := range paths { + if info, err := os.Stat(path); err == nil && !info.IsDir() { + if filepath.Ext(path) == ".json" { + files = append(files, path) + } + continue + } + files = append(files, shared.CollectExt(filepath.Join(path, "threads"), ".json")...) + if filepath.Base(path) == "threads" { + files = append(files, shared.CollectExt(path, ".json")...) + } + } + sort.Strings(files) + files = shared.FilterFiles(shared.UniqueStrings(files), filter) + + entries := make([]usage.Entry, 0) + for _, file := range files { + fileEntries, err := parseThreadFile(file) + if err != nil { + return nil, err + } + entries = append(entries, fileEntries...) + } + shared.SortEntries(entries) + return entries, nil +} + +func parseThreadFile(path string) ([]usage.Entry, error) { + thread, err := shared.ReadJSONObject(path) + if err != nil || thread == nil { + return nil, err + } + threadID := shared.StringField(thread, "id") + if threadID == "" { + return nil, nil + } + messages := shared.ArrayAt(thread["messages"]) + if ledger := shared.ObjectAt(thread["usageLedger"]); ledger != nil { + if events := shared.ArrayAt(ledger["events"]); len(events) > 0 { + return ledgerEntries(path, threadID, messages, events), nil + } + } + return messageEntries(path, threadID, messages), nil +} + +func ledgerEntries(path, threadID string, messages []any, events []any) []usage.Entry { + cacheTokens := cacheTokens(messages) + entries := make([]usage.Entry, 0) + for index, raw := range events { + event := shared.ObjectAt(raw) + if event == nil { + continue + } + timestamp, ok := shared.ParseTimestamp(event["timestamp"]) + if !ok { + continue + } + model := shared.StringField(event, "model") + if model == "" { + continue + } + tokenBlock := shared.ObjectAt(event["tokens"]) + if tokenBlock == nil { + continue + } + cache := cacheTokens[int64Value(event["toMessageId"])] + tokens := usage.TokenUsage{ + InputTokens: shared.UintField(tokenBlock, "input"), + OutputTokens: shared.UintField(tokenBlock, "output"), + CacheCreationInputTokens: cache.cacheCreation, + CacheReadInputTokens: cache.cacheRead, + } + tokens = shared.ApplyTotalFallback(tokens, shared.UintField(tokenBlock, "total")) + if !shared.NonZero(tokens) { + continue + } + messageID := shared.StringValue(event["id"]) + entry := shared.BaseEntry(usage.ProviderAmp, timestamp, "amp", "Amp", threadID, model, "Amp", tokens) + shared.SetSource(&entry, path, index+1, 0, 0) + entry.ID = shared.StableEntryID(entry, messageID) + entries = append(entries, entry) + } + return entries +} + +func messageEntries(path, threadID string, messages []any) []usage.Entry { + entries := make([]usage.Entry, 0) + for index, raw := range messages { + message := shared.ObjectAt(raw) + if message == nil || shared.StringValue(message["role"]) != "assistant" { + continue + } + usageBlock := shared.ObjectAt(message["usage"]) + if usageBlock == nil { + continue + } + timestamp, ok := shared.ParseTimestamp(usageBlock["timestamp"]) + if !ok { + timestamp, ok = shared.ParseTimestamp(message["timestamp"]) + } + if !ok { + continue + } + model := shared.StringField(usageBlock, "model") + if model == "" { + model = shared.StringValue(message["model"]) + } + if model == "" { + continue + } + tokens := usage.TokenUsage{ + InputTokens: shared.UintField(usageBlock, "inputTokens"), + OutputTokens: shared.UintField(usageBlock, "outputTokens"), + CacheCreationInputTokens: shared.UintField(usageBlock, "cacheCreationInputTokens"), + CacheReadInputTokens: shared.UintField(usageBlock, "cacheReadInputTokens"), + } + tokens = shared.ApplyTotalFallback(tokens, shared.UintField(usageBlock, "totalTokens")) + if !shared.NonZero(tokens) { + continue + } + messageID := shared.StringValue(message["messageId"]) + entry := shared.BaseEntry(usage.ProviderAmp, timestamp, "amp", "Amp", threadID, model, "Amp", tokens) + shared.SetSource(&entry, path, index+1, 0, 0) + entry.ID = shared.StableEntryID(entry, messageID) + entries = append(entries, entry) + } + return entries +} + +type cache struct { + cacheCreation uint64 + cacheRead uint64 +} + +func cacheTokens(messages []any) map[int64]cache { + tokens := make(map[int64]cache) + for _, raw := range messages { + message := shared.ObjectAt(raw) + if message == nil || shared.StringValue(message["role"]) != "assistant" { + continue + } + id := int64Value(message["messageId"]) + if id == 0 { + continue + } + usageBlock := shared.ObjectAt(message["usage"]) + tokens[id] = cache{ + cacheCreation: shared.UintField(usageBlock, "cacheCreationInputTokens"), + cacheRead: shared.UintField(usageBlock, "cacheReadInputTokens"), + } + } + return tokens +} + +func int64Value(value any) int64 { + switch typed := value.(type) { + case json.Number: + parsed, _ := strconv.ParseInt(typed.String(), 10, 64) + return parsed + case float64: + return int64(typed) + case string: + parsed, _ := strconv.ParseInt(typed, 10, 64) + return parsed + default: + return 0 + } +} diff --git a/internal/provider/amp/provider.go b/internal/provider/amp/provider.go new file mode 100644 index 0000000..51f2ecf --- /dev/null +++ b/internal/provider/amp/provider.go @@ -0,0 +1,37 @@ +package amp + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Amp usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a Amp provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the Amp provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderAmp } + +// WithFileFilter returns a Amp provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized Amp usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/amp/provider_test.go b/internal/provider/amp/provider_test.go new file mode 100644 index 0000000..236f603 --- /dev/null +++ b/internal/provider/amp/provider_test.go @@ -0,0 +1,32 @@ +package amp + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the Amp smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + path := filepath.Join(dir, "threads", "thread.json") + shared.WriteFile(t, path, `{"id":"thread-a","usageLedger":{"events":[{"id":"event-a","timestamp":"2026-01-02T00:00:00.000Z","model":"gpt-5","tokens":{"input":1,"output":2}}]}}`) + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderAmp, + Model: "gpt-5", + SessionID: "thread-a", + Project: "amp", + Tokens: usage.TokenUsage{ + InputTokens: 1, + OutputTokens: 2, + TotalTokens: 3, + }, + }) +} diff --git a/internal/claudeusage/loader.go b/internal/provider/claude/loader.go similarity index 99% rename from internal/claudeusage/loader.go rename to internal/provider/claude/loader.go index 186c50a..8d1748a 100644 --- a/internal/claudeusage/loader.go +++ b/internal/provider/claude/loader.go @@ -1,4 +1,4 @@ -package claudeusage +package claude import ( "bufio" diff --git a/internal/claudeusage/loader_test.go b/internal/provider/claude/loader_test.go similarity index 99% rename from internal/claudeusage/loader_test.go rename to internal/provider/claude/loader_test.go index 2b9f20b..03410bc 100644 --- a/internal/claudeusage/loader_test.go +++ b/internal/provider/claude/loader_test.go @@ -1,4 +1,4 @@ -package claudeusage +package claude import ( "os" diff --git a/internal/claudeusage/provider.go b/internal/provider/claude/provider.go similarity index 98% rename from internal/claudeusage/provider.go rename to internal/provider/claude/provider.go index 6a96afc..33e5572 100644 --- a/internal/claudeusage/provider.go +++ b/internal/provider/claude/provider.go @@ -1,4 +1,4 @@ -package claudeusage +package claude import ( "sort" diff --git a/internal/agentusage/codebuff.go b/internal/provider/codebuff/loader.go similarity index 59% rename from internal/agentusage/codebuff.go rename to internal/provider/codebuff/loader.go index 6bac979..7f067bb 100644 --- a/internal/agentusage/codebuff.go +++ b/internal/provider/codebuff/loader.go @@ -1,4 +1,4 @@ -package agentusage +package codebuff import ( "bytes" @@ -10,10 +10,11 @@ import ( "strings" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) -type codebuffUsage struct { +type tokenUsage struct { model string inputTokens uint64 outputTokens uint64 @@ -22,17 +23,17 @@ type codebuffUsage struct { extraTotalTokens uint64 } -func loadCodebuffEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, root := range paths { - files = append(files, collectCodebuffChatFiles(root)...) + files = append(files, collectChatFiles(root)...) } sort.Strings(files) - files = filterFiles(uniqueStrings(files), filter) + files = shared.FilterFiles(shared.UniqueStrings(files), filter) entriesByID := make(map[string]usage.Entry) for _, file := range files { - fileEntries, err := parseCodebuffChatFile(file) + fileEntries, err := parseChatFile(file) if err != nil { return nil, err } @@ -44,11 +45,11 @@ func loadCodebuffEntries(paths []string, filter usage.FileFilter) ([]usage.Entry for _, entry := range entriesByID { entries = append(entries, entry) } - sortEntries(entries) + shared.SortEntries(entries) return entries, nil } -func collectCodebuffChatFiles(root string) []string { +func collectChatFiles(root string) []string { info, err := os.Stat(root) if err != nil { return nil @@ -63,12 +64,12 @@ func collectCodebuffChatFiles(root string) []string { if filepath.Base(root) != "projects" { projectRoot = filepath.Join(root, "projects") } - return collectFiles(projectRoot, func(path string) bool { + return shared.CollectFiles(projectRoot, func(path string) bool { return filepath.Base(path) == "chat-messages.json" }) } -func parseCodebuffChatFile(path string) ([]usage.Entry, error) { +func parseChatFile(path string) ([]usage.Entry, error) { data, err := os.ReadFile(path) if err != nil { return nil, err @@ -79,16 +80,16 @@ func parseCodebuffChatFile(path string) ([]usage.Entry, error) { if err := decoder.Decode(&messages); err != nil { return nil, nil } - sessionID, chatID := codebuffSessionContext(path) - chatTimestamp, hasChatTimestamp := parseCodebuffChatTimestamp(chatID) - fileTimestamp := fileModifiedTime(path) + sessionID, chatID := sessionContext(path) + chatTimestamp, hasChatTimestamp := parseChatTimestamp(chatID) + fileTimestamp := shared.FileModifiedTime(path) entries := make([]usage.Entry, 0) for index, raw := range messages { - message := objectAt(raw) - if !isCodebuffAssistant(message) { + message := shared.ObjectAt(raw) + if !isAssistant(message) { continue } - parsedUsage := extractCodebuffUsage(message) + parsedUsage := extractUsage(message) tokens := usage.TokenUsage{ InputTokens: parsedUsage.inputTokens, OutputTokens: parsedUsage.outputTokens, @@ -97,16 +98,16 @@ func parseCodebuffChatFile(path string) ([]usage.Entry, error) { ReasoningOutputTokens: parsedUsage.extraTotalTokens, } if tokens.TotalTokens == 0 { - tokens.TotalTokens = totalUsage(tokens) + tokens.TotalTokens = shared.TotalUsage(tokens) } - if !nonZero(tokens) { + if !shared.NonZero(tokens) { continue } model := parsedUsage.model if model == "" { model = "codebuff-unknown" } - timestamp, ok := codebuffMessageTimestamp(message) + timestamp, ok := messageTimestamp(message) if !ok && hasChatTimestamp { timestamp = chatTimestamp ok = true @@ -114,15 +115,15 @@ func parseCodebuffChatFile(path string) ([]usage.Entry, error) { if !ok { timestamp = fileTimestamp } - entry := baseEntry(usage.ProviderCodebuff, timestamp, "codebuff", "Codebuff", sessionID, model, "Codebuff", tokens) - setSource(&entry, path, index+1, 0, 0) - entry.ID = stableEntryID(entry, codebuffDedupKey(message, sessionID, timestamp, model, tokens, index)) + entry := shared.BaseEntry(usage.ProviderCodebuff, timestamp, "codebuff", "Codebuff", sessionID, model, "Codebuff", tokens) + shared.SetSource(&entry, path, index+1, 0, 0) + entry.ID = shared.StableEntryID(entry, dedupKey(message, sessionID, timestamp, model, tokens, index)) entries = append(entries, entry) } return entries, nil } -func codebuffSessionContext(path string) (string, string) { +func sessionContext(path string) (string, string) { chatID := filepath.Base(filepath.Dir(path)) chatsDir := filepath.Dir(filepath.Dir(path)) projectDir := filepath.Dir(chatsDir) @@ -140,50 +141,50 @@ func codebuffSessionContext(path string) (string, string) { return channel + "/" + project + "/" + chatID, chatID } -func isCodebuffAssistant(message map[string]any) bool { - role := firstStringField(message, "variant", "role") +func isAssistant(message map[string]any) bool { + role := shared.FirstStringField(message, "variant", "role") return role == "ai" || role == "agent" || role == "assistant" } -func extractCodebuffUsage(message map[string]any) codebuffUsage { - var usage codebuffUsage - metadata := objectAt(message["metadata"]) +func extractUsage(message map[string]any) tokenUsage { + var usage tokenUsage + metadata := shared.ObjectAt(message["metadata"]) if metadata != nil { - usage.model = stringField(metadata, "model") - mergeCodebuffUsage(&usage, parseCodebuffUsageObject(metadata["usage"])) - mergeCodebuffUsage(&usage, parseCodebuffUsageObject(objectAt(metadata["codebuff"])["usage"])) - if runState := codebuffRunStateUsage(metadata); runState != nil { + usage.model = shared.StringField(metadata, "model") + mergeCodebuffUsage(&usage, parseUsageObject(metadata["usage"])) + mergeCodebuffUsage(&usage, parseUsageObject(shared.ObjectAt(metadata["codebuff"])["usage"])) + if runState := runStateUsage(metadata); runState != nil { mergeCodebuffUsage(&usage, *runState) } } return usage } -func codebuffRunStateUsage(metadata map[string]any) *codebuffUsage { - history := arrayAt(objectAt(objectAt(objectAt(metadata["runState"])["sessionState"])["mainAgentState"])["messageHistory"]) +func runStateUsage(metadata map[string]any) *tokenUsage { + history := shared.ArrayAt(shared.ObjectAt(shared.ObjectAt(shared.ObjectAt(metadata["runState"])["sessionState"])["mainAgentState"])["messageHistory"]) if len(history) == 0 { return nil } - var usage codebuffUsage + var usage tokenUsage found := false for i := len(history) - 1; i >= 0; i-- { - entry := objectAt(history[i]) - if stringField(entry, "role") != "assistant" { + entry := shared.ObjectAt(history[i]) + if shared.StringField(entry, "role") != "assistant" { continue } - providerOptions := objectAt(entry["providerOptions"]) + providerOptions := shared.ObjectAt(entry["providerOptions"]) if providerOptions == nil { continue } - entryUsage := parseCodebuffUsageObject(providerOptions["usage"]) - codebuff := objectAt(providerOptions["codebuff"]) + entryUsage := parseUsageObject(providerOptions["usage"]) + codebuff := shared.ObjectAt(providerOptions["codebuff"]) if codebuff != nil { - mergeCodebuffUsage(&entryUsage, parseCodebuffUsageObject(codebuff["usage"])) + mergeCodebuffUsage(&entryUsage, parseUsageObject(codebuff["usage"])) if entryUsage.model == "" { - entryUsage.model = stringField(codebuff, "model") + entryUsage.model = shared.StringField(codebuff, "model") } } - if codebuffUsageHasTokens(entryUsage) || entryUsage.model != "" { + if usageHasTokens(entryUsage) || entryUsage.model != "" { found = true } mergeCodebuffUsage(&usage, entryUsage) @@ -194,27 +195,27 @@ func codebuffRunStateUsage(metadata map[string]any) *codebuffUsage { return &usage } -func parseCodebuffUsageObject(value any) codebuffUsage { - record := objectAt(value) +func parseUsageObject(value any) tokenUsage { + record := shared.ObjectAt(value) if record == nil { - return codebuffUsage{} + return tokenUsage{} } - parsed := codebuffUsage{ - model: stringField(record, "model"), + parsed := tokenUsage{ + model: shared.StringField(record, "model"), inputTokens: firstUint(record, "inputTokens", "input_tokens", "promptTokens", "prompt_tokens"), outputTokens: firstUint(record, "outputTokens", "output_tokens", "completionTokens", "completion_tokens"), cacheReadInputTokens: firstUint(record, "cacheReadInputTokens", "cache_read_input_tokens"), cacheCreationInputTokens: firstUint(record, "cacheCreationInputTokens", "cache_creation_input_tokens", "cacheCreationTokens", "cache_creation_tokens", "cachedTokensCreated", "cached_tokens_created"), } - parsed.cacheReadInputTokens = maxUint64(parsed.cacheReadInputTokens, firstUint(objectAt(record["promptTokensDetails"]), "cachedTokens")) - parsed.cacheReadInputTokens = maxUint64(parsed.cacheReadInputTokens, firstUint(objectAt(record["prompt_tokens_details"]), "cached_tokens")) + parsed.cacheReadInputTokens = maxUint64(parsed.cacheReadInputTokens, firstUint(shared.ObjectAt(record["promptTokensDetails"]), "cachedTokens")) + parsed.cacheReadInputTokens = maxUint64(parsed.cacheReadInputTokens, firstUint(shared.ObjectAt(record["prompt_tokens_details"]), "cached_tokens")) tokens := usage.TokenUsage{ InputTokens: parsed.inputTokens, OutputTokens: parsed.outputTokens, CacheCreationInputTokens: parsed.cacheCreationInputTokens, CacheReadInputTokens: parsed.cacheReadInputTokens, } - tokens = applyTotalFallback(tokens, firstUint(record, "totalTokens", "total_tokens", "total")) + tokens = shared.ApplyTotalFallback(tokens, firstUint(record, "totalTokens", "total_tokens", "total")) parsed.inputTokens = tokens.InputTokens parsed.outputTokens = tokens.OutputTokens parsed.cacheCreationInputTokens = tokens.CacheCreationInputTokens @@ -223,7 +224,7 @@ func parseCodebuffUsageObject(value any) codebuffUsage { return parsed } -func mergeCodebuffUsage(target *codebuffUsage, fallback codebuffUsage) { +func mergeCodebuffUsage(target *tokenUsage, fallback tokenUsage) { if target.inputTokens == 0 { target.inputTokens = fallback.inputTokens } @@ -244,21 +245,21 @@ func mergeCodebuffUsage(target *codebuffUsage, fallback codebuffUsage) { } } -func codebuffUsageHasTokens(value codebuffUsage) bool { +func usageHasTokens(value tokenUsage) bool { return value.inputTokens > 0 || value.outputTokens > 0 || value.cacheCreationInputTokens > 0 || value.cacheReadInputTokens > 0 || value.extraTotalTokens > 0 } -func codebuffMessageTimestamp(message map[string]any) (time.Time, bool) { - if timestamp, ok := parseTimestamp(message["timestamp"]); ok { +func messageTimestamp(message map[string]any) (time.Time, bool) { + if timestamp, ok := shared.ParseTimestamp(message["timestamp"]); ok { return timestamp, true } - if timestamp, ok := parseTimestamp(message["createdAt"]); ok { + if timestamp, ok := shared.ParseTimestamp(message["createdAt"]); ok { return timestamp, true } - return parseTimestamp(objectAt(message["metadata"])["timestamp"]) + return shared.ParseTimestamp(shared.ObjectAt(message["metadata"])["timestamp"]) } -func parseCodebuffChatTimestamp(chatID string) (time.Time, bool) { +func parseChatTimestamp(chatID string) (time.Time, bool) { date, clock, ok := strings.Cut(chatID, "T") if !ok { return time.Time{}, false @@ -268,21 +269,21 @@ func parseCodebuffChatTimestamp(chatID string) (time.Time, bool) { clock = clock[:index] + ":" + clock[index+1:] } } - return parseTimestampString(date + "T" + clock) + return shared.ParseTimestampString(date + "T" + clock) } -func codebuffDedupKey(message map[string]any, sessionID string, timestamp time.Time, model string, tokens usage.TokenUsage, index int) string { - if id := stringField(message, "id"); id != "" { +func dedupKey(message map[string]any, sessionID string, timestamp time.Time, model string, tokens usage.TokenUsage, index int) string { + if id := shared.StringField(message, "id"); id != "" { return "codebuff:" + sessionID + ":" + id } - return stableEntryID(baseEntry(usage.ProviderCodebuff, timestamp, "codebuff", "Codebuff", sessionID, model, "Codebuff", tokens), strconv.Itoa(index)) + return shared.StableEntryID(shared.BaseEntry(usage.ProviderCodebuff, timestamp, "codebuff", "Codebuff", sessionID, model, "Codebuff", tokens), strconv.Itoa(index)) } func firstUint(record map[string]any, keys ...string) uint64 { if record == nil { return 0 } - return uintField(record, keys...) + return shared.UintField(record, keys...) } func maxUint64(a, b uint64) uint64 { diff --git a/internal/provider/codebuff/provider.go b/internal/provider/codebuff/provider.go new file mode 100644 index 0000000..4da56db --- /dev/null +++ b/internal/provider/codebuff/provider.go @@ -0,0 +1,37 @@ +package codebuff + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Codebuff usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a Codebuff provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the Codebuff provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderCodebuff } + +// WithFileFilter returns a Codebuff provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized Codebuff usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/codebuff/provider_test.go b/internal/provider/codebuff/provider_test.go new file mode 100644 index 0000000..2a6a80a --- /dev/null +++ b/internal/provider/codebuff/provider_test.go @@ -0,0 +1,34 @@ +package codebuff + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the Codebuff smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + root := filepath.Join(t.TempDir(), "manicode") + path := filepath.Join(root, "projects", "project-a", "chats", "2026-01-02T03-04-05.000Z", "chat-messages.json") + shared.WriteFile(t, path, `[{"id":"assistant-message","role":"assistant","timestamp":"2026-01-02T03:04:06.000Z","metadata":{"model":"claude-sonnet-4-20250514","usage":{"inputTokens":100,"outputTokens":50,"cacheCreationInputTokens":20,"cacheReadInputTokens":10}}}]`) + return Provider{}.WithPaths([]string{root}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderCodebuff, + Model: "claude-sonnet-4-20250514", + SessionID: "manicode/project-a/2026-01-02T03-04-05.000Z", + Project: "codebuff", + Tokens: usage.TokenUsage{ + InputTokens: 100, + OutputTokens: 50, + CacheCreationInputTokens: 20, + CacheReadInputTokens: 10, + TotalTokens: 180, + }, + }) +} diff --git a/internal/codexusage/loader.go b/internal/provider/codex/loader.go similarity index 99% rename from internal/codexusage/loader.go rename to internal/provider/codex/loader.go index 9218054..c9d2c7f 100644 --- a/internal/codexusage/loader.go +++ b/internal/provider/codex/loader.go @@ -1,4 +1,4 @@ -package codexusage +package codex import ( "bufio" diff --git a/internal/codexusage/loader_test.go b/internal/provider/codex/loader_test.go similarity index 99% rename from internal/codexusage/loader_test.go rename to internal/provider/codex/loader_test.go index a441360..9ace368 100644 --- a/internal/codexusage/loader_test.go +++ b/internal/provider/codex/loader_test.go @@ -1,4 +1,4 @@ -package codexusage +package codex import ( "os" diff --git a/internal/codexusage/provider.go b/internal/provider/codex/provider.go similarity index 98% rename from internal/codexusage/provider.go rename to internal/provider/codex/provider.go index 02debf7..ff8f99a 100644 --- a/internal/codexusage/provider.go +++ b/internal/provider/codex/provider.go @@ -1,4 +1,4 @@ -package codexusage +package codex import ( "sort" diff --git a/internal/provider/copilot/filter_test.go b/internal/provider/copilot/filter_test.go new file mode 100644 index 0000000..1ca258c --- /dev/null +++ b/internal/provider/copilot/filter_test.go @@ -0,0 +1,33 @@ +package copilot + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" +) + +// TestFileFilterSkipsRejectedFiles proves the scanner can skip source files +// whose events are already ingested. +func TestFileFilterSkipsRejectedFiles(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "copilot.jsonl") + shared.WriteFile(t, path, `{"type":"span","traceId":"trace-1","spanId":"span-1","name":"chat claude-sonnet-4","endTime":[1775934264,967317833],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"claude-sonnet-4","gen_ai.conversation.id":"conv-1","gen_ai.usage.input_tokens":19452,"gen_ai.usage.output_tokens":281}}}`+"\n") + + rejected := make([]string, 0) + provider := Provider{}.WithPaths([]string{dir}).(Provider). + WithFileFilter(func(candidate string) bool { + rejected = append(rejected, candidate) + return false + }) + entries, err := provider.Entries() + if err != nil { + t.Fatal(err) + } + if len(entries) != 0 { + t.Fatalf("len(entries) = %d, want 0 when the filter rejects every file", len(entries)) + } + if len(rejected) != 1 || rejected[0] != path { + t.Fatalf("filter saw %#v, want the session file", rejected) + } +} diff --git a/internal/agentusage/copilot.go b/internal/provider/copilot/loader.go similarity index 53% rename from internal/agentusage/copilot.go rename to internal/provider/copilot/loader.go index e2aaefe..de8f164 100644 --- a/internal/agentusage/copilot.go +++ b/internal/provider/copilot/loader.go @@ -1,4 +1,4 @@ -package agentusage +package copilot import ( "fmt" @@ -7,40 +7,41 @@ import ( "strings" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) -func loadCopilotEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, root := range paths { - files = append(files, collectExt(root, ".jsonl")...) + files = append(files, shared.CollectExt(root, ".jsonl")...) } sort.Strings(files) - files = filterFiles(uniqueStrings(files), filter) + files = shared.FilterFiles(shared.UniqueStrings(files), filter) entries := make([]usage.Entry, 0) for _, file := range files { - fileEntries, err := parseCopilotOTELFile(file) + fileEntries, err := parseOTELFile(file) if err != nil { return nil, err } entries = append(entries, fileEntries...) } - sortEntries(entries) + shared.SortEntries(entries) return entries, nil } -type copilotSource int +type sourceKind int const ( - copilotChatSpan copilotSource = iota + copilotChatSpan sourceKind = iota copilotInferenceLog copilotAgentTurnLog copilotAgentSummarySpan ) -type copilotCandidate struct { - source copilotSource +type candidate struct { + source sourceKind traceID string responseID string sessionID string @@ -53,32 +54,32 @@ type copilotCandidate struct { sourceStart, sourceEnd int64 } -type copilotTraceContext struct { +type traceContext struct { model string sessionID string sessionIDPriority int } -func parseCopilotOTELFile(path string) ([]usage.Entry, error) { - lines, err := readJSONLines(path, `"attributes"`) +func parseOTELFile(path string) ([]usage.Entry, error) { + lines, err := shared.ReadJSONLines(path, `"attributes"`) if err != nil { return nil, err } - contexts := copilotTraceContexts(lines) - fallback := fileModifiedTime(path) - candidates := make([]copilotCandidate, 0) + contexts := traceContexts(lines) + fallback := shared.FileModifiedTime(path) + candidates := make([]candidate, 0) for index, line := range lines { - if candidate, ok := copilotRecordCandidate(path, line, index, fallback, contexts); ok { + if candidate, ok := recordCandidate(path, line, index, fallback, contexts); ok { candidates = append(candidates, candidate) } } - sets := copilotCandidateSets(candidates) + sets := candidateSets(candidates) entries := make([]usage.Entry, 0) for _, candidate := range candidates { if !shouldEmitCopilot(candidate, sets) { continue } - entry := baseEntry( + entry := shared.BaseEntry( usage.ProviderCopilot, candidate.timestamp, "copilot", @@ -88,30 +89,30 @@ func parseCopilotOTELFile(path string) ([]usage.Entry, error) { "GitHub Copilot CLI", candidate.tokens, ) - setSource(&entry, candidate.sourceFile, candidate.sourceLine, candidate.sourceStart, candidate.sourceEnd) - entry.ID = stableEntryID(entry, candidate.dedupKey) + shared.SetSource(&entry, candidate.sourceFile, candidate.sourceLine, candidate.sourceStart, candidate.sourceEnd) + entry.ID = shared.StableEntryID(entry, candidate.dedupKey) entries = append(entries, entry) } return entries, nil } -func copilotTraceContexts(lines []lineJSON) map[string]copilotTraceContext { - contexts := make(map[string]copilotTraceContext) +func traceContexts(lines []shared.LineJSON) map[string]traceContext { + contexts := make(map[string]traceContext) for _, line := range lines { - record := line.value - traceID := copilotTraceID(record) + record := line.Value + traceID := traceID(record) if traceID == "" { continue } - attrs := objectAt(record["attributes"]) + attrs := shared.ObjectAt(record["attributes"]) if attrs == nil { continue } context := contexts[traceID] if context.model == "" { - context.model = firstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") + context.model = shared.FirstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") } - if sessionID, priority := copilotBestSession(attrs); sessionID != "" && priority > context.sessionIDPriority { + if sessionID, priority := bestSession(attrs); sessionID != "" && priority > context.sessionIDPriority { context.sessionID = sessionID context.sessionIDPriority = priority } @@ -120,18 +121,18 @@ func copilotTraceContexts(lines []lineJSON) map[string]copilotTraceContext { return contexts } -func copilotRecordCandidate(path string, line lineJSON, index int, fallback time.Time, contexts map[string]copilotTraceContext) (copilotCandidate, bool) { - record := line.value - attrs := objectAt(record["attributes"]) +func recordCandidate(path string, line shared.LineJSON, index int, fallback time.Time, contexts map[string]traceContext) (candidate, bool) { + record := line.Value + attrs := shared.ObjectAt(record["attributes"]) if attrs == nil { - return copilotCandidate{}, false + return candidate{}, false } - source, ok := copilotRecordSource(record, attrs) + source, ok := recordSource(record, attrs) if !ok { - return copilotCandidate{}, false + return candidate{}, false } - input := uintField(attrs, "gen_ai.usage.input_tokens") - cacheRead := uintField(attrs, "gen_ai.usage.cache_read.input_tokens") + input := shared.UintField(attrs, "gen_ai.usage.input_tokens") + cacheRead := shared.UintField(attrs, "gen_ai.usage.cache_read.input_tokens") if cacheRead <= input { input -= cacheRead } else { @@ -139,25 +140,25 @@ func copilotRecordCandidate(path string, line lineJSON, index int, fallback time } tokens := usage.TokenUsage{ InputTokens: input, - OutputTokens: uintField(attrs, "gen_ai.usage.output_tokens"), - CacheCreationInputTokens: uintField(attrs, "gen_ai.usage.cache_write.input_tokens", "gen_ai.usage.cache_creation.input_tokens"), + OutputTokens: shared.UintField(attrs, "gen_ai.usage.output_tokens"), + CacheCreationInputTokens: shared.UintField(attrs, "gen_ai.usage.cache_write.input_tokens", "gen_ai.usage.cache_creation.input_tokens"), CacheReadInputTokens: cacheRead, - ReasoningOutputTokens: uintField(attrs, "gen_ai.usage.reasoning.output_tokens", "gen_ai.usage.reasoning_tokens"), + ReasoningOutputTokens: shared.UintField(attrs, "gen_ai.usage.reasoning.output_tokens", "gen_ai.usage.reasoning_tokens"), } - tokens = applyTotalFallback(tokens, uintField(attrs, "gen_ai.usage.total_tokens", "gen_ai.usage.total.token_count")) - if !nonZero(tokens) { - return copilotCandidate{}, false + tokens = shared.ApplyTotalFallback(tokens, shared.UintField(attrs, "gen_ai.usage.total_tokens", "gen_ai.usage.total.token_count")) + if !shared.NonZero(tokens) { + return candidate{}, false } - traceID := copilotTraceID(record) + traceID := traceID(record) context := contexts[traceID] - model := firstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") + model := shared.FirstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") if model == "" { model = context.model } if model == "" { model = "unknown" } - sessionID, _ := copilotBestSession(attrs) + sessionID, _ := bestSession(attrs) if sessionID == "" { sessionID = context.sessionID } @@ -167,12 +168,12 @@ func copilotRecordCandidate(path string, line lineJSON, index int, fallback time if sessionID == "" { sessionID = "unknown-session" } - timestamp, ok := copilotTimestamp(record) + timestamp, ok := timestamp(record) if !ok { timestamp = fallback } - responseID := stringField(attrs, "gen_ai.response.id") - return copilotCandidate{ + responseID := shared.StringField(attrs, "gen_ai.response.id") + return candidate{ source: source, traceID: traceID, responseID: responseID, @@ -180,90 +181,90 @@ func copilotRecordCandidate(path string, line lineJSON, index int, fallback time model: model, timestamp: timestamp, tokens: tokens, - dedupKey: copilotDedupKey(source, record, attrs, traceID, sessionID, timestamp, index), + dedupKey: dedupKey(source, record, attrs, traceID, sessionID, timestamp, index), sourceFile: path, - sourceLine: line.line, - sourceStart: line.start, - sourceEnd: line.end, + sourceLine: line.Line, + sourceStart: line.Start, + sourceEnd: line.End, }, true } -func copilotRecordSource(record, attrs map[string]any) (copilotSource, bool) { +func recordSource(record, attrs map[string]any) (sourceKind, bool) { switch { - case copilotIsChatSpan(record, attrs): + case isChatSpan(record, attrs): return copilotChatSpan, true - case copilotIsInferenceLog(record, attrs): + case isInferenceLog(record, attrs): return copilotInferenceLog, true - case copilotIsAgentTurnLog(record, attrs): + case isAgentTurnLog(record, attrs): return copilotAgentTurnLog, true - case copilotIsAgentSummarySpan(record, attrs): + case isAgentSummarySpan(record, attrs): return copilotAgentSummarySpan, true default: return 0, false } } -func copilotIsSpan(record map[string]any) bool { - if stringField(record, "type") == "span" { +func isSpan(record map[string]any) bool { + if shared.StringField(record, "type") == "span" { return true } - if stringField(record, "name") == "" { + if shared.StringField(record, "name") == "" { return false } - return stringField(record, "spanId") != "" || - stringField(record, "traceId") != "" || + return shared.StringField(record, "spanId") != "" || + shared.StringField(record, "traceId") != "" || record["startTime"] != nil || record["endTime"] != nil || record["duration"] != nil || record["kind"] != nil } -func copilotIsChatSpan(record, attrs map[string]any) bool { - return copilotIsSpan(record) && - (stringField(attrs, "gen_ai.operation.name") == "chat" || - strings.HasPrefix(stringField(record, "name"), "chat ")) +func isChatSpan(record, attrs map[string]any) bool { + return isSpan(record) && + (shared.StringField(attrs, "gen_ai.operation.name") == "chat" || + strings.HasPrefix(shared.StringField(record, "name"), "chat ")) } -func copilotIsAgentSummarySpan(record, attrs map[string]any) bool { - return copilotIsSpan(record) && - (stringField(attrs, "gen_ai.operation.name") == "invoke_agent" || - strings.HasPrefix(stringField(record, "name"), "invoke_agent ")) +func isAgentSummarySpan(record, attrs map[string]any) bool { + return isSpan(record) && + (shared.StringField(attrs, "gen_ai.operation.name") == "invoke_agent" || + strings.HasPrefix(shared.StringField(record, "name"), "invoke_agent ")) } -func copilotIsInferenceLog(record, attrs map[string]any) bool { - return !copilotIsSpan(record) && - (stringField(attrs, "event.name") == "gen_ai.client.inference.operation.details" || - strings.HasPrefix(copilotBody(record), "GenAI inference:")) +func isInferenceLog(record, attrs map[string]any) bool { + return !isSpan(record) && + (shared.StringField(attrs, "event.name") == "gen_ai.client.inference.operation.details" || + strings.HasPrefix(spanBody(record), "GenAI inference:")) } -func copilotIsAgentTurnLog(record, attrs map[string]any) bool { - return !copilotIsSpan(record) && - (stringField(attrs, "event.name") == "copilot_chat.agent.turn" || - strings.HasPrefix(copilotBody(record), "copilot_chat.agent.turn")) +func isAgentTurnLog(record, attrs map[string]any) bool { + return !isSpan(record) && + (shared.StringField(attrs, "event.name") == "copilot_chat.agent.turn" || + strings.HasPrefix(spanBody(record), "copilot_chat.agent.turn")) } -func copilotBody(record map[string]any) string { - if body := stringField(record, "body"); body != "" { +func spanBody(record map[string]any) string { + if body := shared.StringField(record, "body"); body != "" { return body } - return stringField(record, "_body") + return shared.StringField(record, "_body") } -func copilotTraceID(record map[string]any) string { - if traceID := stringField(record, "traceId"); traceID != "" { +func traceID(record map[string]any) string { + if traceID := shared.StringField(record, "traceId"); traceID != "" { return traceID } - return stringField(objectAt(record["spanContext"]), "traceId") + return shared.StringField(shared.ObjectAt(record["spanContext"]), "traceId") } -func copilotSpanID(record map[string]any) string { - if spanID := stringField(record, "spanId"); spanID != "" { +func spanID(record map[string]any) string { + if spanID := shared.StringField(record, "spanId"); spanID != "" { return spanID } - return stringField(objectAt(record["spanContext"]), "spanId") + return shared.StringField(shared.ObjectAt(record["spanContext"]), "spanId") } -func copilotBestSession(attrs map[string]any) (string, int) { +func bestSession(attrs map[string]any) (string, int) { candidates := []struct { key string priority int @@ -278,7 +279,7 @@ func copilotBestSession(attrs map[string]any) (string, int) { bestValue := "" bestPriority := 0 for _, candidate := range candidates { - if value := stringField(attrs, candidate.key); value != "" && candidate.priority > bestPriority { + if value := shared.StringField(attrs, candidate.key); value != "" && candidate.priority > bestPriority { bestValue = value bestPriority = candidate.priority } @@ -286,25 +287,25 @@ func copilotBestSession(attrs map[string]any) (string, int) { return bestValue, bestPriority } -func copilotTimestamp(record map[string]any) (time.Time, bool) { +func timestamp(record map[string]any) (time.Time, bool) { for _, key := range []string{"endTime", "startTime", "hrTime", "_hrTime", "time"} { - if timestamp, ok := timestampFromParts(record[key]); ok { + if timestamp, ok := shared.TimestampFromParts(record[key]); ok { return timestamp, true } } for _, key := range []string{"timestamp", "observedTimestamp"} { - if timestamp, ok := parseTimestamp(record[key]); ok { + if timestamp, ok := shared.ParseTimestamp(record[key]); ok { return timestamp, true } } - if raw := uintValue(record["timeUnixNano"]); raw > 0 { + if raw := shared.UintValue(record["timeUnixNano"]); raw > 0 { return time.UnixMilli(int64(raw / 1_000_000)), true } return time.Time{}, false } -func copilotDedupKey(source copilotSource, record, attrs map[string]any, traceID, sessionID string, timestamp time.Time, index int) string { - spanID := copilotSpanID(record) +func dedupKey(source sourceKind, record, attrs map[string]any, traceID, sessionID string, timestamp time.Time, index int) string { + spanID := spanID(record) switch source { case copilotChatSpan, copilotAgentSummarySpan: if traceID != "" && spanID != "" { @@ -317,7 +318,7 @@ func copilotDedupKey(source copilotSource, record, attrs map[string]any, traceID } return fmt.Sprintf("log:%s:%d:%d", sessionID, timestamp.UnixMilli(), index) case copilotAgentTurnLog: - turnIndex := uintField(attrs, "turn.index", "copilot_chat.turn.index") + turnIndex := shared.UintField(attrs, "turn.index", "copilot_chat.turn.index") turn := fmt.Sprintf("idx-%d", index) if turnIndex > 0 { turn = fmt.Sprintf("%d", turnIndex) @@ -327,11 +328,11 @@ func copilotDedupKey(source copilotSource, record, attrs map[string]any, traceID } return "agent-turn:" + sessionID + ":" + turn + fmt.Sprintf(":%d", index) default: - return fmt.Sprintf("%s:%d", filepath.Base(copilotSourceName(source)), index) + return fmt.Sprintf("%s:%d", filepath.Base(sourceName(source)), index) } } -func copilotSourceName(source copilotSource) string { +func sourceName(source sourceKind) string { switch source { case copilotChatSpan: return "chat" @@ -346,7 +347,7 @@ func copilotSourceName(source copilotSource) string { } } -type copilotSets struct { +type candidateSetMap struct { chatTraces map[string]bool inferenceTraces map[string]bool agentTurnTraces map[string]bool @@ -355,8 +356,8 @@ type copilotSets struct { agentTurnResponses map[string]bool } -func copilotCandidateSets(candidates []copilotCandidate) copilotSets { - sets := copilotSets{ +func candidateSets(candidates []candidate) candidateSetMap { + sets := candidateSetMap{ chatTraces: make(map[string]bool), inferenceTraces: make(map[string]bool), agentTurnTraces: make(map[string]bool), @@ -389,7 +390,7 @@ func copilotCandidateSets(candidates []copilotCandidate) copilotSets { return sets } -func shouldEmitCopilot(candidate copilotCandidate, sets copilotSets) bool { +func shouldEmitCopilot(candidate candidate, sets candidateSetMap) bool { traceMatch := func(values map[string]bool) bool { return candidate.traceID != "" && values[candidate.traceID] } diff --git a/internal/provider/copilot/provider.go b/internal/provider/copilot/provider.go new file mode 100644 index 0000000..8030c54 --- /dev/null +++ b/internal/provider/copilot/provider.go @@ -0,0 +1,37 @@ +package copilot + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads GitHub Copilot CLI usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a GitHub Copilot CLI provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the GitHub Copilot CLI provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderCopilot } + +// WithFileFilter returns a GitHub Copilot CLI provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized GitHub Copilot CLI usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/copilot/provider_test.go b/internal/provider/copilot/provider_test.go new file mode 100644 index 0000000..8e10127 --- /dev/null +++ b/internal/provider/copilot/provider_test.go @@ -0,0 +1,35 @@ +package copilot + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the GitHub Copilot CLI smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + path := filepath.Join(dir, "copilot.jsonl") + shared.WriteFile(t, path, `{"type":"span","traceId":"trace-1","spanId":"span-1","name":"chat claude-sonnet-4","endTime":[1775934264,967317833],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"claude-sonnet-4","gen_ai.conversation.id":"conv-1","gen_ai.usage.input_tokens":19452,"gen_ai.usage.output_tokens":281,"gen_ai.usage.cache_read.input_tokens":123,"gen_ai.usage.cache_creation.input_tokens":25,"gen_ai.usage.reasoning.output_tokens":128}}}`+"\n") + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderCopilot, + Model: "claude-sonnet-4", + SessionID: "conv-1", + Project: "copilot", + Tokens: usage.TokenUsage{ + InputTokens: 19329, + OutputTokens: 281, + CacheCreationInputTokens: 25, + CacheReadInputTokens: 123, + ReasoningOutputTokens: 128, + TotalTokens: 19886, + }, + }) +} diff --git a/internal/agentusage/droid.go b/internal/provider/droid/loader.go similarity index 64% rename from internal/agentusage/droid.go rename to internal/provider/droid/loader.go index fe6a4a7..63304b0 100644 --- a/internal/agentusage/droid.go +++ b/internal/provider/droid/loader.go @@ -1,4 +1,4 @@ -package agentusage +package droid import ( "bufio" @@ -7,20 +7,21 @@ import ( "sort" "strings" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) -func loadDroidEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, root := range paths { - files = append(files, collectFiles(root, isDroidSettingsFile)...) + files = append(files, shared.CollectFiles(root, isSettingsFile)...) } sort.Strings(files) - files = filterFiles(uniqueStrings(files), filter) + files = shared.FilterFiles(shared.UniqueStrings(files), filter) entries := make([]usage.Entry, 0) for _, file := range files { - entry, ok, err := parseDroidSettingsFile(file) + entry, ok, err := parseSettingsFile(file) if err != nil { return nil, err } @@ -28,58 +29,58 @@ func loadDroidEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, e entries = append(entries, entry) } } - sortEntries(entries) + shared.SortEntries(entries) return entries, nil } -func isDroidSettingsFile(path string) bool { +func isSettingsFile(path string) bool { return strings.HasSuffix(filepath.Base(path), ".settings.json") } -func parseDroidSettingsFile(path string) (usage.Entry, bool, error) { - settings, err := readJSONObject(path) +func parseSettingsFile(path string) (usage.Entry, bool, error) { + settings, err := shared.ReadJSONObject(path) if err != nil || settings == nil { return usage.Entry{}, false, err } - usageBlock := objectAt(settings["tokenUsage"]) + usageBlock := shared.ObjectAt(settings["tokenUsage"]) if usageBlock == nil { return usage.Entry{}, false, nil } tokens := usage.TokenUsage{ - InputTokens: uintField(usageBlock, "inputTokens"), - OutputTokens: uintField(usageBlock, "outputTokens"), - CacheCreationInputTokens: uintField(usageBlock, "cacheCreationTokens"), - CacheReadInputTokens: uintField(usageBlock, "cacheReadTokens"), - ReasoningOutputTokens: uintField(usageBlock, "thinkingTokens"), - } - tokens = applyTotalFallback(tokens, uintField(usageBlock, "totalTokens")) - if !nonZero(tokens) { + InputTokens: shared.UintField(usageBlock, "inputTokens"), + OutputTokens: shared.UintField(usageBlock, "outputTokens"), + CacheCreationInputTokens: shared.UintField(usageBlock, "cacheCreationTokens"), + CacheReadInputTokens: shared.UintField(usageBlock, "cacheReadTokens"), + ReasoningOutputTokens: shared.UintField(usageBlock, "thinkingTokens"), + } + tokens = shared.ApplyTotalFallback(tokens, shared.UintField(usageBlock, "totalTokens")) + if !shared.NonZero(tokens) { return usage.Entry{}, false, nil } - provider := normalizeDroidProvider(stringField(settings, "providerLock")) - model := normalizeDroidModel(stringField(settings, "model")) + provider := normalizeDroidProvider(shared.StringField(settings, "providerLock")) + model := normalizeDroidModel(shared.StringField(settings, "model")) if model == "" { - model, _ = droidSidecarModel(path) + model, _ = sidecarModel(path) } if model == "" { - model = droidDefaultModel(provider) + model = defaultModel(provider) } if model == "" { model = "unknown" } - timestamp, ok := parseTimestamp(settings["providerLockTimestamp"]) + timestamp, ok := shared.ParseTimestamp(settings["providerLockTimestamp"]) if !ok { - timestamp = fileModifiedTime(path) + timestamp = shared.FileModifiedTime(path) } sessionID := strings.TrimSuffix(filepath.Base(path), ".settings.json") if sessionID == "" { sessionID = "unknown" } - entry := baseEntry(usage.ProviderDroid, timestamp, "droid", "Droid", sessionID, model, "Droid", tokens) - setSource(&entry, path, 1, 0, 0) - entry.ID = stableEntryID(entry, sessionID) + entry := shared.BaseEntry(usage.ProviderDroid, timestamp, "droid", "Droid", sessionID, model, "Droid", tokens) + shared.SetSource(&entry, path, 1, 0, 0) + entry.ID = shared.StableEntryID(entry, sessionID) return entry, true, nil } @@ -141,7 +142,7 @@ func normalizeDroidProvider(provider string) string { } } -func droidDefaultModel(provider string) string { +func defaultModel(provider string) string { switch provider { case "anthropic": return "claude-unknown" @@ -156,7 +157,7 @@ func droidDefaultModel(provider string) string { } } -func droidSidecarModel(settingsPath string) (string, error) { +func sidecarModel(settingsPath string) (string, error) { name := filepath.Base(settingsPath) prefix := strings.TrimSuffix(name, ".settings.json") if prefix == name || prefix == "" { @@ -173,14 +174,14 @@ func droidSidecarModel(settingsPath string) (string, error) { scanner := bufio.NewScanner(file) for i := 0; i < 500 && scanner.Scan(); i++ { - if model := droidModelFromLine(scanner.Text()); model != "" { + if model := modelFromLine(scanner.Text()); model != "" { return model, nil } } return "", scanner.Err() } -func droidModelFromLine(line string) string { +func modelFromLine(line string) string { _, tail, ok := strings.Cut(line, "Model:") if !ok { return "" diff --git a/internal/provider/droid/provider.go b/internal/provider/droid/provider.go new file mode 100644 index 0000000..0277d7d --- /dev/null +++ b/internal/provider/droid/provider.go @@ -0,0 +1,37 @@ +package droid + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Droid usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a Droid provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the Droid provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderDroid } + +// WithFileFilter returns a Droid provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized Droid usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/droid/provider_test.go b/internal/provider/droid/provider_test.go new file mode 100644 index 0000000..cd386a5 --- /dev/null +++ b/internal/provider/droid/provider_test.go @@ -0,0 +1,35 @@ +package droid + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the Droid smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + path := filepath.Join(dir, "session-a.settings.json") + shared.WriteFile(t, path, `{"model":"Claude-Sonnet-4-[Anthropic]","providerLock":"anthropic","providerLockTimestamp":"2026-01-02T00:00:00.000Z","tokenUsage":{"inputTokens":100,"outputTokens":50,"cacheCreationTokens":20,"cacheReadTokens":10,"thinkingTokens":5}}`) + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderDroid, + Model: "claude-sonnet-4", + SessionID: "session-a", + Project: "droid", + Tokens: usage.TokenUsage{ + InputTokens: 100, + OutputTokens: 50, + CacheCreationInputTokens: 20, + CacheReadInputTokens: 10, + ReasoningOutputTokens: 5, + TotalTokens: 185, + }, + }) +} diff --git a/internal/agentusage/gemini.go b/internal/provider/gemini/loader.go similarity index 51% rename from internal/agentusage/gemini.go rename to internal/provider/gemini/loader.go index 8fa999a..df10108 100644 --- a/internal/agentusage/gemini.go +++ b/internal/provider/gemini/loader.go @@ -1,4 +1,4 @@ -package agentusage +package gemini import ( "path/filepath" @@ -6,37 +6,38 @@ import ( "strings" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) -func loadGeminiEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, root := range paths { - files = append(files, collectExt(root, ".json")...) - files = append(files, collectExt(root, ".jsonl")...) + files = append(files, shared.CollectExt(root, ".json")...) + files = append(files, shared.CollectExt(root, ".jsonl")...) } sort.Strings(files) - files = filterFiles(uniqueStrings(files), filter) + files = shared.FilterFiles(shared.UniqueStrings(files), filter) entries := make([]usage.Entry, 0) for _, file := range files { var fileEntries []usage.Entry var err error if strings.EqualFold(filepath.Ext(file), ".jsonl") { - fileEntries, err = parseGeminiJSONLFile(file) + fileEntries, err = parseJSONLFile(file) } else { - fileEntries, err = parseGeminiJSONFile(file) + fileEntries, err = parseJSONFile(file) } if err != nil { return nil, err } entries = append(entries, fileEntries...) } - sortEntries(entries) + shared.SortEntries(entries) return entries, nil } -type geminiTokens struct { +type tokens struct { input uint64 output uint64 cached uint64 @@ -46,63 +47,63 @@ type geminiTokens struct { hasTotal bool } -func parseGeminiJSONFile(path string) ([]usage.Entry, error) { - record, err := readJSONObject(path) +func parseJSONFile(path string) ([]usage.Entry, error) { + record, err := shared.ReadJSONObject(path) if err != nil || record == nil { return nil, err } - fallback := fileModifiedTime(path) - sessionID := firstStringField(record, "sessionId", "session_id") + fallback := shared.FileModifiedTime(path) + sessionID := shared.FirstStringField(record, "sessionId", "session_id") if sessionID == "" { sessionID = strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) } sessionTimestamp := firstTimestamp(record, fallback, "startTime", "lastUpdated", "timestamp", "created_at") - if messages := arrayAt(record["messages"]); len(messages) > 0 { + if messages := shared.ArrayAt(record["messages"]); len(messages) > 0 { entries := make([]usage.Entry, 0) for index, raw := range messages { - message := objectAt(raw) - if stringField(message, "type") != "gemini" { + message := shared.ObjectAt(raw) + if shared.StringField(message, "type") != "gemini" { continue } - if entry, ok := geminiDirectEntry(message, path, index+1, "", sessionID, sessionTimestamp); ok { + if entry, ok := directEntry(message, path, index+1, "", sessionID, sessionTimestamp); ok { entries = append(entries, entry) } } return entries, nil } - if stringField(record, "type") == "gemini" { - if entry, ok := geminiDirectEntry(record, path, 1, "", sessionID, fallback); ok { + if shared.StringField(record, "type") == "gemini" { + if entry, ok := directEntry(record, path, 1, "", sessionID, fallback); ok { return []usage.Entry{entry}, nil } return nil, nil } - return geminiStatsEntries(recordStats(record), path, 1, stringField(record, "model"), sessionID, sessionTimestamp), nil + return statsEntries(recordStats(record), path, 1, shared.StringField(record, "model"), sessionID, sessionTimestamp), nil } -func parseGeminiJSONLFile(path string) ([]usage.Entry, error) { - lines, err := readJSONLines(path) +func parseJSONLFile(path string) ([]usage.Entry, error) { + lines, err := shared.ReadJSONLines(path) if err != nil { return nil, err } - fallback := fileModifiedTime(path) + fallback := shared.FileModifiedTime(path) sessionID := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) currentModel := "" entries := make([]usage.Entry, 0) indexByMessageID := make(map[string]int) for _, line := range lines { - record := line.value - if value := firstStringField(record, "sessionId", "session_id"); value != "" { + record := line.Value + if value := shared.FirstStringField(record, "sessionId", "session_id"); value != "" { sessionID = value } - if model := stringField(record, "model"); model != "" { + if model := shared.StringField(record, "model"); model != "" { currentModel = model } - if stringField(record, "type") == "gemini" { - entry, ok := geminiDirectEntry(record, path, line.line, currentModel, sessionID, fallback) + if shared.StringField(record, "type") == "gemini" { + entry, ok := directEntry(record, path, line.Line, currentModel, sessionID, fallback) if !ok { continue } - messageID := stringField(record, "id") + messageID := shared.StringField(record, "id") if messageID != "" { if index, exists := indexByMessageID[messageID]; exists { entries[index] = entry @@ -115,34 +116,34 @@ func parseGeminiJSONLFile(path string) ([]usage.Entry, error) { } if stats := recordStats(record); stats != nil { timestamp := firstTimestamp(record, fallback, "timestamp") - entries = append(entries, geminiStatsEntries(stats, path, line.line, currentModel, sessionID, timestamp)...) + entries = append(entries, statsEntries(stats, path, line.Line, currentModel, sessionID, timestamp)...) } } return entries, nil } -func geminiDirectEntry(record map[string]any, path string, line int, modelHint, sessionID string, fallback time.Time) (usage.Entry, bool) { - tokens, ok := parseGeminiTokens(record["tokens"]) +func directEntry(record map[string]any, path string, line int, modelHint, sessionID string, fallback time.Time) (usage.Entry, bool) { + tokens, ok := parseTokens(record["tokens"]) if !ok { return usage.Entry{}, false } - model := stringField(record, "model") + model := shared.StringField(record, "model") if model == "" { model = modelHint } timestamp := firstTimestamp(record, fallback, "timestamp", "created_at") - return buildGeminiEntry(path, line, model, sessionID, timestamp, tokens, true, stringField(record, "id")) + return buildGeminiEntry(path, line, model, sessionID, timestamp, tokens, true, shared.StringField(record, "id")) } -func geminiStatsEntries(stats map[string]any, path string, line int, modelHint, sessionID string, timestamp time.Time) []usage.Entry { +func statsEntries(stats map[string]any, path string, line int, modelHint, sessionID string, timestamp time.Time) []usage.Entry { if stats == nil { return nil } - if models := objectAt(stats["models"]); models != nil { + if models := shared.ObjectAt(stats["models"]); models != nil { entries := make([]usage.Entry, 0) for model, raw := range models { - data := objectAt(raw) - tokens, ok := parseGeminiTokens(data["tokens"]) + data := shared.ObjectAt(raw) + tokens, ok := parseTokens(data["tokens"]) if !ok { continue } @@ -154,7 +155,7 @@ func geminiStatsEntries(stats map[string]any, path string, line int, modelHint, return entries } } - tokens, ok := parseGeminiTokens(stats) + tokens, ok := parseTokens(stats) if !ok { return nil } @@ -169,7 +170,7 @@ func geminiStatsEntries(stats map[string]any, path string, line int, modelHint, return []usage.Entry{entry} } -func buildGeminiEntry(path string, line int, model, sessionID string, timestamp time.Time, tokens geminiTokens, direct bool, messageID string) (usage.Entry, bool) { +func buildGeminiEntry(path string, line int, model, sessionID string, timestamp time.Time, tokens tokens, direct bool, messageID string) (usage.Entry, bool) { model = strings.TrimSpace(model) if model == "" { return usage.Entry{}, false @@ -182,37 +183,37 @@ func buildGeminiEntry(path string, line int, model, sessionID string, timestamp ReasoningOutputTokens: tokens.thoughts, } if tokens.hasTotal { - tokenUsage = applyTotalFallback(tokenUsage, tokens.total) + tokenUsage = shared.ApplyTotalFallback(tokenUsage, tokens.total) } else if tokenUsage.TotalTokens == 0 { - tokenUsage.TotalTokens = totalUsage(tokenUsage) + tokenUsage.TotalTokens = shared.TotalUsage(tokenUsage) } - if !nonZero(tokenUsage) { + if !shared.NonZero(tokenUsage) { return usage.Entry{}, false } - entry := baseEntry(usage.ProviderGemini, timestamp, "gemini", "Gemini", sessionID, model, "Gemini CLI", tokenUsage) - setSource(&entry, path, line, 0, 0) - entry.ID = stableEntryID(entry, messageID) + entry := shared.BaseEntry(usage.ProviderGemini, timestamp, "gemini", "Gemini", sessionID, model, "Gemini CLI", tokenUsage) + shared.SetSource(&entry, path, line, 0, 0) + entry.ID = shared.StableEntryID(entry, messageID) return entry, true } -func parseGeminiTokens(raw any) (geminiTokens, bool) { - record := objectAt(raw) +func parseTokens(raw any) (tokens, bool) { + record := shared.ObjectAt(raw) if record == nil { - return geminiTokens{}, false + return tokens{}, false } - tokens := geminiTokens{ - input: uintField(record, "input", "prompt", "input_tokens", "prompt_tokens"), - output: uintField(record, "output", "candidates", "output_tokens", "candidates_tokens"), - cached: uintField(record, "cached", "cached_tokens"), - thoughts: uintField(record, "thoughts", "reasoning", "thoughts_tokens", "reasoning_tokens"), - tool: uintField(record, "tool", "tool_tokens"), - total: uintField(record, "total", "total_tokens"), + tokens := tokens{ + input: shared.UintField(record, "input", "prompt", "input_tokens", "prompt_tokens"), + output: shared.UintField(record, "output", "candidates", "output_tokens", "candidates_tokens"), + cached: shared.UintField(record, "cached", "cached_tokens"), + thoughts: shared.UintField(record, "thoughts", "reasoning", "thoughts_tokens", "reasoning_tokens"), + tool: shared.UintField(record, "tool", "tool_tokens"), + total: shared.UintField(record, "total", "total_tokens"), } tokens.hasTotal = tokens.total > 0 return tokens, true } -func normalizeGeminiInput(tokens geminiTokens, direct bool) (uint64, uint64) { +func normalizeGeminiInput(tokens tokens, direct bool) (uint64, uint64) { if !direct { cachedPortion := tokens.input if tokens.cached < cachedPortion { @@ -233,16 +234,16 @@ func normalizeGeminiInput(tokens geminiTokens, direct bool) (uint64, uint64) { } func recordStats(record map[string]any) map[string]any { - if stats := objectAt(record["stats"]); stats != nil { + if stats := shared.ObjectAt(record["stats"]); stats != nil { return stats } - result := objectAt(record["result"]) - return objectAt(result["stats"]) + result := shared.ObjectAt(record["result"]) + return shared.ObjectAt(result["stats"]) } func firstTimestamp(record map[string]any, fallback time.Time, keys ...string) time.Time { for _, key := range keys { - if timestamp, ok := parseTimestamp(record[key]); ok { + if timestamp, ok := shared.ParseTimestamp(record[key]); ok { return timestamp } } diff --git a/internal/provider/gemini/provider.go b/internal/provider/gemini/provider.go new file mode 100644 index 0000000..d4ca2ad --- /dev/null +++ b/internal/provider/gemini/provider.go @@ -0,0 +1,37 @@ +package gemini + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Gemini CLI usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a Gemini CLI provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the Gemini CLI provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderGemini } + +// WithFileFilter returns a Gemini CLI provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized Gemini CLI usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/gemini/provider_test.go b/internal/provider/gemini/provider_test.go new file mode 100644 index 0000000..000cbdd --- /dev/null +++ b/internal/provider/gemini/provider_test.go @@ -0,0 +1,36 @@ +package gemini + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the Gemini CLI smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + path := filepath.Join(dir, "session-a.jsonl") + shared.WriteFile(t, path, + `{"sessionId":"session-a","projectHash":"project-a","startTime":"2026-05-17T11:07:00.000Z"}`+"\n"+ + `{"id":"msg-a","timestamp":"2026-05-17T11:07:32.000Z","type":"gemini","model":"gemini-3-flash-preview","tokens":{"input":15327,"output":23,"cached":11526,"thoughts":919,"tool":7,"total":16276}}`+"\n") + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderGemini, + Model: "gemini-3-flash-preview", + SessionID: "session-a", + Project: "gemini", + Tokens: usage.TokenUsage{ + InputTokens: 3808, + OutputTokens: 23, + CacheReadInputTokens: 11526, + ReasoningOutputTokens: 919, + TotalTokens: 16276, + }, + }) +} diff --git a/internal/agentusage/goose.go b/internal/provider/goose/loader.go similarity index 52% rename from internal/agentusage/goose.go rename to internal/provider/goose/loader.go index b36e054..6c24295 100644 --- a/internal/agentusage/goose.go +++ b/internal/provider/goose/loader.go @@ -1,4 +1,4 @@ -package agentusage +package goose import ( "path/filepath" @@ -6,15 +6,16 @@ import ( "strings" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) -func loadGooseEntries(paths []string) ([]usage.Entry, error) { - dbPaths := gooseDBPaths(paths) +func loadEntries(paths []string) ([]usage.Entry, error) { + dbPaths := dbPaths(paths) entries := make([]usage.Entry, 0) seen := make(map[string]bool) for _, dbPath := range dbPaths { - dbEntries, err := loadGooseDatabase(dbPath) + dbEntries, err := loadDatabase(dbPath) if err != nil { return nil, err } @@ -27,14 +28,14 @@ func loadGooseEntries(paths []string) ([]usage.Entry, error) { entries = append(entries, entry) } } - sortEntries(entries) + shared.SortEntries(entries) return entries, nil } -func gooseDBPaths(paths []string) []string { +func dbPaths(paths []string) []string { dbPaths := make([]string, 0) for _, root := range paths { - if existingSQLiteFile(root) { + if shared.ExistingSQLiteFile(root) { dbPaths = append(dbPaths, root) continue } @@ -43,17 +44,17 @@ func gooseDBPaths(paths []string) []string { filepath.Join(root, "sessions", "sessions.db"), filepath.Join(root, "data", "sessions", "sessions.db"), } { - if existingSQLiteFile(candidate) { + if shared.ExistingSQLiteFile(candidate) { dbPaths = append(dbPaths, candidate) } } } sort.Strings(dbPaths) - return uniqueStrings(dbPaths) + return shared.UniqueStrings(dbPaths) } -func loadGooseDatabase(path string) ([]usage.Entry, error) { - db, err := openSQLite(path) +func loadDatabase(path string) ([]usage.Entry, error) { + db, err := shared.OpenSQLite(path) if err != nil { return nil, err } @@ -74,29 +75,29 @@ func loadGooseDatabase(path string) ([]usage.Entry, error) { entries := make([]usage.Entry, 0) for rows.Next() { var id, modelConfig, providerName, createdAt, total, input, output, accumulatedTotal, accumulatedInput, accumulatedOutput any - if !scanAny(rows, &id, &modelConfig, &providerName, &createdAt, &total, &input, &output, &accumulatedTotal, &accumulatedInput, &accumulatedOutput) { + if !shared.ScanAny(rows, &id, &modelConfig, &providerName, &createdAt, &total, &input, &output, &accumulatedTotal, &accumulatedInput, &accumulatedOutput) { continue } - if entry, ok := gooseRowEntry(path, id, modelConfig, providerName, createdAt, total, input, output, accumulatedTotal, accumulatedInput, accumulatedOutput); ok { + if entry, ok := rowEntry(path, id, modelConfig, providerName, createdAt, total, input, output, accumulatedTotal, accumulatedInput, accumulatedOutput); ok { entries = append(entries, entry) } } return entries, rows.Err() } -func gooseRowEntry(path string, id, modelConfig, providerName, createdAt, total, input, output, accumulatedTotal, accumulatedInput, accumulatedOutput any) (usage.Entry, bool) { - sessionID := sqlString(id) - model := gooseModelName(sqlString(modelConfig)) +func rowEntry(path string, id, modelConfig, providerName, createdAt, total, input, output, accumulatedTotal, accumulatedInput, accumulatedOutput any) (usage.Entry, bool) { + sessionID := shared.SqlString(id) + model := modelName(shared.SqlString(modelConfig)) if sessionID == "" || model == "" { return usage.Entry{}, false } - timestamp, ok := gooseTimestamp(sqlString(createdAt)) + timestamp, ok := timestamp(shared.SqlString(createdAt)) if !ok { return usage.Entry{}, false } - inputTokens := firstPositive(sqlUint(accumulatedInput), sqlUint(input)) - outputTokens := firstPositive(sqlUint(accumulatedOutput), sqlUint(output)) - totalTokens := firstPositive(sqlUint(accumulatedTotal), sqlUint(total), inputTokens+outputTokens) + inputTokens := firstPositive(shared.SqlUint(accumulatedInput), shared.SqlUint(input)) + outputTokens := firstPositive(shared.SqlUint(accumulatedOutput), shared.SqlUint(output)) + totalTokens := firstPositive(shared.SqlUint(accumulatedTotal), shared.SqlUint(total), inputTokens+outputTokens) tokens := usage.TokenUsage{ InputTokens: inputTokens, OutputTokens: outputTokens, @@ -104,34 +105,34 @@ func gooseRowEntry(path string, id, modelConfig, providerName, createdAt, total, if totalTokens > inputTokens+outputTokens { tokens.ReasoningOutputTokens = totalTokens - inputTokens - outputTokens } - tokens.TotalTokens = totalUsage(tokens) - if !nonZero(tokens) { + tokens.TotalTokens = shared.TotalUsage(tokens) + if !shared.NonZero(tokens) { return usage.Entry{}, false } - entry := baseEntry(usage.ProviderGoose, timestamp, "goose", "Goose", sessionID, model, "Goose", tokens) - setSource(&entry, path, 0, 0, 0) - entry.ID = stableEntryID(entry, "goose:"+sessionID+":"+sqlString(providerName)) + entry := shared.BaseEntry(usage.ProviderGoose, timestamp, "goose", "Goose", sessionID, model, "Goose", tokens) + shared.SetSource(&entry, path, 0, 0, 0) + entry.ID = shared.StableEntryID(entry, "goose:"+sessionID+":"+shared.SqlString(providerName)) return entry, true } -func gooseModelName(config string) string { - record := decodeJSONObjectString(config) - return stringField(record, "model_name") +func modelName(config string) string { + record := shared.DecodeJSONObjectString(config) + return shared.StringField(record, "model_name") } -func gooseTimestamp(value string) (time.Time, bool) { +func timestamp(value string) (time.Time, bool) { value = strings.TrimSpace(value) if value == "" { return time.Time{}, false } - if timestamp, ok := parseTimestampString(value); ok { + if timestamp, ok := shared.ParseTimestampString(value); ok { return timestamp, true } if len(value) == 19 && value[4] == '-' && value[7] == '-' && (value[10] == ' ' || value[10] == 'T') { - return parseTimestampString(value[:10] + "T" + value[11:] + "Z") + return shared.ParseTimestampString(value[:10] + "T" + value[11:] + "Z") } if len(value) == 10 && value[4] == '-' && value[7] == '-' { - return parseTimestampString(value + "T00:00:00Z") + return shared.ParseTimestampString(value + "T00:00:00Z") } return time.Time{}, false } diff --git a/internal/provider/goose/provider.go b/internal/provider/goose/provider.go new file mode 100644 index 0000000..e504f8e --- /dev/null +++ b/internal/provider/goose/provider.go @@ -0,0 +1,30 @@ +package goose + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Goose usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a Goose provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the Goose provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderGoose } + +// Entries loads normalized Goose usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/goose/provider_test.go b/internal/provider/goose/provider_test.go new file mode 100644 index 0000000..796cedc --- /dev/null +++ b/internal/provider/goose/provider_test.go @@ -0,0 +1,52 @@ +package goose + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the Goose smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + dbPath := filepath.Join(dir, "sessions.db") + db := shared.OpenTestSQLite(t, dbPath) + defer db.Close() + shared.ExecSQL(t, db, `CREATE TABLE sessions ( + id TEXT PRIMARY KEY, + model_config_json TEXT, + provider_name TEXT, + created_at TEXT, + total_tokens INTEGER, + input_tokens INTEGER, + output_tokens INTEGER, + accumulated_total_tokens INTEGER, + accumulated_input_tokens INTEGER, + accumulated_output_tokens INTEGER + )`) + shared.ExecSQL(t, db, `INSERT INTO sessions ( + id, model_config_json, provider_name, created_at, + accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens + ) VALUES (?, ?, ?, ?, ?, ?, ?)`, + "session-a", `{"model_name":"claude-sonnet-4-20250514"}`, "anthropic", "2026-05-01 01:02:03", 180, 100, 50, + ) + return Provider{}.WithPaths([]string{dbPath}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderGoose, + Model: "claude-sonnet-4-20250514", + SessionID: "session-a", + Project: "goose", + Tokens: usage.TokenUsage{ + InputTokens: 100, + OutputTokens: 50, + ReasoningOutputTokens: 30, + TotalTokens: 180, + }, + }) +} diff --git a/internal/provider/hermes/loader.go b/internal/provider/hermes/loader.go new file mode 100644 index 0000000..95d99bc --- /dev/null +++ b/internal/provider/hermes/loader.go @@ -0,0 +1,98 @@ +package hermes + +import ( + "strings" + "time" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func loadEntries(paths []string) ([]usage.Entry, error) { + dbPaths := shared.SqliteDBPaths(paths, "state.db", nil) + entries := make([]usage.Entry, 0) + for _, dbPath := range dbPaths { + dbEntries, err := loadDatabase(dbPath) + if err != nil { + return nil, err + } + entries = append(entries, dbEntries...) + } + shared.SortEntries(entries) + return entries, nil +} + +func loadDatabase(path string) ([]usage.Entry, error) { + db, err := shared.OpenSQLite(path) + if err != nil { + return nil, err + } + defer db.Close() + + rows, err := db.Query(` + SELECT id, model, billing_provider, started_at, message_count, input_tokens, + output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, + estimated_cost_usd, actual_cost_usd + FROM sessions + WHERE model IS NOT NULL AND TRIM(model) != '' + `) + if err != nil { + return nil, nil + } + defer rows.Close() + + entries := make([]usage.Entry, 0) + for rows.Next() { + var sessionID, model, provider, startedAt, messageCount, input, output, cacheRead, cacheWrite, reasoning, estimatedCost, actualCost any + if !shared.ScanAny(rows, &sessionID, &model, &provider, &startedAt, &messageCount, &input, &output, &cacheRead, &cacheWrite, &reasoning, &estimatedCost, &actualCost) { + continue + } + entry, ok := rowEntry(path, sessionID, model, startedAt, input, output, cacheRead, cacheWrite, reasoning) + if ok { + entries = append(entries, entry) + } + } + return entries, rows.Err() +} + +func rowEntry(path string, sessionRaw, modelRaw, startedAt, input, output, cacheRead, cacheWrite, reasoning any) (usage.Entry, bool) { + sessionID := shared.SqlString(sessionRaw) + model := strings.TrimSpace(shared.SqlString(modelRaw)) + if sessionID == "" || model == "" { + return usage.Entry{}, false + } + timestamp, ok := timestamp(startedAt) + if !ok { + return usage.Entry{}, false + } + tokens := usage.TokenUsage{ + InputTokens: shared.SqlUint(input), + OutputTokens: shared.SqlUint(output), + CacheCreationInputTokens: shared.SqlUint(cacheWrite), + CacheReadInputTokens: shared.SqlUint(cacheRead), + ReasoningOutputTokens: shared.SqlUint(reasoning), + } + if tokens.TotalTokens == 0 { + tokens.TotalTokens = shared.TotalUsage(tokens) + } + if !shared.NonZero(tokens) { + return usage.Entry{}, false + } + entry := shared.BaseEntry(usage.ProviderHermes, timestamp, "hermes", "Hermes", sessionID, model, "Hermes Agent", tokens) + shared.SetSource(&entry, path, 0, 0, 0) + entry.ID = shared.StableEntryID(entry, "hermes:"+sessionID) + return entry, true +} + +func timestamp(value any) (time.Time, bool) { + if parsed, ok := shared.ParseTimestamp(value); ok { + return parsed, true + } + if number, ok := shared.SqlFloat(value); ok { + return shared.TimestampFromFloat(number) + } + if text := shared.SqlString(value); text != "" { + return shared.ParseTimestampString(text) + } + return time.Time{}, false +} diff --git a/internal/provider/hermes/provider.go b/internal/provider/hermes/provider.go new file mode 100644 index 0000000..01e06f0 --- /dev/null +++ b/internal/provider/hermes/provider.go @@ -0,0 +1,30 @@ +package hermes + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Hermes Agent usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a Hermes Agent provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the Hermes Agent provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderHermes } + +// Entries loads normalized Hermes Agent usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/hermes/provider_test.go b/internal/provider/hermes/provider_test.go new file mode 100644 index 0000000..e2113dc --- /dev/null +++ b/internal/provider/hermes/provider_test.go @@ -0,0 +1,57 @@ +package hermes + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the Hermes Agent smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + dbPath := filepath.Join(dir, "state.db") + db := shared.OpenTestSQLite(t, dbPath) + defer db.Close() + shared.ExecSQL(t, db, `CREATE TABLE sessions ( + id TEXT PRIMARY KEY, + model TEXT, + billing_provider TEXT, + started_at REAL, + message_count INTEGER, + input_tokens INTEGER, + output_tokens INTEGER, + cache_read_tokens INTEGER, + cache_write_tokens INTEGER, + reasoning_tokens INTEGER, + estimated_cost_usd REAL, + actual_cost_usd REAL + )`) + shared.ExecSQL(t, db, `INSERT INTO sessions ( + id, model, billing_provider, started_at, message_count, input_tokens, + output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, + estimated_cost_usd, actual_cost_usd + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + "session-a", "gpt-5.5", "openai", 1750000000.25, 42, 100, 50, 10, 20, 5, 0.12, 0.34, + ) + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderHermes, + Model: "gpt-5.5", + SessionID: "session-a", + Project: "hermes", + Tokens: usage.TokenUsage{ + InputTokens: 100, + OutputTokens: 50, + CacheCreationInputTokens: 20, + CacheReadInputTokens: 10, + ReasoningOutputTokens: 5, + TotalTokens: 185, + }, + }) +} diff --git a/internal/provider/kilo/loader.go b/internal/provider/kilo/loader.go new file mode 100644 index 0000000..1b85f74 --- /dev/null +++ b/internal/provider/kilo/loader.go @@ -0,0 +1,89 @@ +package kilo + +import ( + "errors" + "os" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func loadEntries(paths []string) ([]usage.Entry, error) { + dbPaths := shared.SqliteDBPaths(paths, "kilo.db", nil) + entries := make([]usage.Entry, 0) + for _, dbPath := range dbPaths { + dbEntries, err := loadDatabase(dbPath) + if err != nil { + return nil, err + } + entries = append(entries, dbEntries...) + } + shared.SortEntries(entries) + return entries, nil +} + +func loadDatabase(path string) ([]usage.Entry, error) { + db, err := shared.OpenSQLite(path) + if errors.Is(err, os.ErrNotExist) { + return nil, nil + } + if err != nil { + return nil, err + } + defer db.Close() + + rows, err := db.Query("SELECT id, session_id, data FROM message") + if err != nil { + return nil, nil + } + defer rows.Close() + + entries := make([]usage.Entry, 0) + for rows.Next() { + var rowID, rowSessionID, data string + if err := rows.Scan(&rowID, &rowSessionID, &data); err != nil { + continue + } + if entry, ok := messageEntry(path, rowID, rowSessionID, data); ok { + entries = append(entries, entry) + } + } + return entries, rows.Err() +} + +func messageEntry(dbPath, rowID, rowSessionID, data string) (usage.Entry, bool) { + record := shared.DecodeJSONObjectString(data) + if record == nil || shared.StringField(record, "role") != "assistant" { + return usage.Entry{}, false + } + tokenBlock := shared.ObjectAt(record["tokens"]) + if tokenBlock == nil { + return usage.Entry{}, false + } + cache := shared.ObjectAt(tokenBlock["cache"]) + tokens := usage.TokenUsage{ + InputTokens: shared.UintField(tokenBlock, "input"), + OutputTokens: shared.UintField(tokenBlock, "output"), + CacheCreationInputTokens: shared.UintField(cache, "write"), + CacheReadInputTokens: shared.UintField(cache, "read"), + ReasoningOutputTokens: shared.UintField(tokenBlock, "reasoning"), + } + tokens = shared.ApplyTotalFallback(tokens, shared.UintField(tokenBlock, "total")) + if !shared.NonZero(tokens) { + return usage.Entry{}, false + } + model := shared.StringField(record, "modelID") + if model == "" { + return usage.Entry{}, false + } + timestamp, ok := shared.ParseTimestamp(shared.ObjectAt(record["time"])["created"]) + if !ok { + return usage.Entry{}, false + } + sessionID := shared.FirstNonEmpty(shared.StringField(record, "session_id"), rowSessionID, "unknown") + messageID := shared.FirstNonEmpty(shared.StringField(record, "id"), rowID) + entry := shared.BaseEntry(usage.ProviderKilo, timestamp, "kilo", "Kilo", sessionID, model, "Kilo", tokens) + shared.SetSource(&entry, dbPath, 0, 0, 0) + entry.ID = shared.StableEntryID(entry, messageID) + return entry, true +} diff --git a/internal/provider/kilo/provider.go b/internal/provider/kilo/provider.go new file mode 100644 index 0000000..ec57af1 --- /dev/null +++ b/internal/provider/kilo/provider.go @@ -0,0 +1,30 @@ +package kilo + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Kilo usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a Kilo provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the Kilo provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderKilo } + +// Entries loads normalized Kilo usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/kilo/provider_test.go b/internal/provider/kilo/provider_test.go new file mode 100644 index 0000000..688926a --- /dev/null +++ b/internal/provider/kilo/provider_test.go @@ -0,0 +1,42 @@ +package kilo + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the Kilo smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + dbPath := filepath.Join(dir, "kilo.db") + db := shared.OpenTestSQLite(t, dbPath) + defer db.Close() + shared.ExecSQL(t, db, "CREATE TABLE message (id TEXT, session_id TEXT, data TEXT)") + shared.ExecSQL(t, db, `INSERT INTO message (id, session_id, data) VALUES (?, ?, ?)`, + "row-1", + "session-a", + `{"id":"msg-1","role":"assistant","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","time":{"created":1767312000000},"tokens":{"input":100,"output":50,"reasoning":5,"cache":{"read":10,"write":20}}}`, + ) + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderKilo, + Model: "claude-sonnet-4-20250514", + SessionID: "session-a", + Project: "kilo", + Tokens: usage.TokenUsage{ + InputTokens: 100, + OutputTokens: 50, + CacheCreationInputTokens: 20, + CacheReadInputTokens: 10, + ReasoningOutputTokens: 5, + TotalTokens: 185, + }, + }) +} diff --git a/internal/provider/kimi/loader.go b/internal/provider/kimi/loader.go new file mode 100644 index 0000000..beb12a5 --- /dev/null +++ b/internal/provider/kimi/loader.go @@ -0,0 +1,199 @@ +package kimi + +import ( + "path/filepath" + "sort" + "strings" + "time" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +const defaultModel = "kimi-for-coding" + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + files := make([]string, 0) + for _, root := range paths { + files = append(files, shared.CollectFiles(filepath.Join(root, "sessions"), isWireFile)...) + } + sort.Strings(files) + files = shared.FilterFiles(shared.UniqueStrings(files), filter) + + entries := make([]usage.Entry, 0) + seen := make(map[string]bool) + for _, file := range files { + fileEntries, err := parseWireFile(file) + if err != nil { + return nil, err + } + for _, entry := range fileEntries { + if seen[entry.ID] { + continue + } + seen[entry.ID] = true + entries = append(entries, entry) + } + } + shared.SortEntries(entries) + return entries, nil +} + +func isWireFile(path string) bool { + if filepath.Base(path) != "wire.jsonl" { + return false + } + parts := strings.Split(filepath.ToSlash(filepath.Clean(path)), "/") + for i := range parts { + if parts[i] != "sessions" { + continue + } + // Old layout: sessions///wire.jsonl + // New layout: sessions///agents//wire.jsonl + if i+3 == len(parts)-1 || i+5 == len(parts)-1 { + return true + } + } + return false +} + +// record is one usable wire line, normalized across the old +// StatusUpdate format and the new Kimi Code usage.record format. +type record struct { + tokens usage.TokenUsage + timestamp time.Time + hasTime bool + model string // empty means fall back to the config.json model + messageID string +} + +func parseWireFile(path string) ([]usage.Entry, error) { + lines, err := shared.ReadJSONLines(path, `usage`) + if err != nil { + return nil, err + } + configModel := configModel(path) + sessionID := sessionID(path) + fallback := shared.FileModifiedTime(path) + entries := make([]usage.Entry, 0) + for _, line := range lines { + var record record + var ok bool + if shared.StringField(line.Value, "type") == "usage.record" { + record, ok = parseUsageRecord(line.Value) + } else { + record, ok = parseStatusUpdate(line.Value) + } + if !ok { + continue + } + timestamp := record.timestamp + if !record.hasTime { + timestamp = fallback + } + model := record.model + if model == "" { + model = configModel + } + entry := shared.BaseEntry(usage.ProviderKimi, timestamp, "kimi", "Kimi", sessionID, model, "Kimi", record.tokens) + shared.SetSource(&entry, path, line.Line, line.Start, line.End) + entry.ID = shared.StableEntryID(entry, record.messageID) + entries = append(entries, entry) + } + return entries, nil +} + +func parseStatusUpdate(value map[string]any) (record, bool) { + message := shared.ObjectAt(value["message"]) + if shared.StringField(message, "type") != "StatusUpdate" { + return record{}, false + } + payload := shared.ObjectAt(message["payload"]) + tokenUsage := shared.ObjectAt(payload["token_usage"]) + if tokenUsage == nil { + return record{}, false + } + tokens := usage.TokenUsage{ + InputTokens: shared.UintField(tokenUsage, "input_other"), + OutputTokens: shared.UintField(tokenUsage, "output"), + CacheCreationInputTokens: shared.UintField(tokenUsage, "input_cache_creation"), + CacheReadInputTokens: shared.UintField(tokenUsage, "input_cache_read"), + } + tokens = shared.ApplyTotalFallback(tokens, shared.UintField(tokenUsage, "total")) + if !shared.NonZero(tokens) { + return record{}, false + } + record := record{tokens: tokens, messageID: shared.StringField(payload, "message_id")} + record.timestamp, record.hasTime = shared.ParseTimestamp(value["timestamp"]) + return record, true +} + +func parseUsageRecord(value map[string]any) (record, bool) { + // Session-scoped records are cumulative totals; only turn records count. + if shared.StringField(value, "usageScope") != "turn" { + return record{}, false + } + tokenUsage := shared.ObjectAt(value["usage"]) + if tokenUsage == nil { + return record{}, false + } + tokens := usage.TokenUsage{ + InputTokens: shared.UintField(tokenUsage, "inputOther"), + OutputTokens: shared.UintField(tokenUsage, "output"), + CacheCreationInputTokens: shared.UintField(tokenUsage, "inputCacheCreation"), + CacheReadInputTokens: shared.UintField(tokenUsage, "inputCacheRead"), + } + tokens = shared.ApplyTotalFallback(tokens, 0) + if !shared.NonZero(tokens) { + return record{}, false + } + record := record{ + tokens: tokens, + model: strings.TrimPrefix(shared.StringField(value, "model"), "kimi-code/"), + } + record.timestamp, record.hasTime = shared.ParseTimestamp(value["time"]) + return record, true +} + +// sessionID returns the session directory name for either layout. +func sessionID(path string) string { + dir := filepath.Dir(path) + if filepath.Base(filepath.Dir(dir)) == "agents" { + dir = filepath.Dir(filepath.Dir(dir)) + } + sessionID := filepath.Base(dir) + if sessionID == "" || sessionID == "." { + return "unknown" + } + return sessionID +} + +// kimiRoot walks up from a wire file to the directory containing "sessions", +// which is the Kimi data root regardless of layout depth. +func kimiRoot(path string) string { + for dir := filepath.Dir(path); ; { + parent := filepath.Dir(dir) + if filepath.Base(dir) == "sessions" { + return parent + } + if parent == dir { + return "" + } + dir = parent + } +} + +func configModel(path string) string { + root := kimiRoot(path) + if root == "" { + return defaultModel + } + config, err := shared.ReadJSONObject(filepath.Join(root, "config.json")) + if err != nil || config == nil { + return defaultModel + } + if model := shared.StringField(config, "model"); model != "" { + return model + } + return defaultModel +} diff --git a/internal/provider/kimi/provider.go b/internal/provider/kimi/provider.go new file mode 100644 index 0000000..60da2e7 --- /dev/null +++ b/internal/provider/kimi/provider.go @@ -0,0 +1,37 @@ +package kimi + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Kimi usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a Kimi provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the Kimi provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderKimi } + +// WithFileFilter returns a Kimi provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized Kimi usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/kimi/provider_test.go b/internal/provider/kimi/provider_test.go new file mode 100644 index 0000000..cd93a12 --- /dev/null +++ b/internal/provider/kimi/provider_test.go @@ -0,0 +1,37 @@ +package kimi + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the Kimi smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + shared.WriteFile(t, filepath.Join(dir, "config.json"), `{"model":"kimi-k2"}`) + path := filepath.Join(dir, "sessions", "group", "session-a", "wire.jsonl") + shared.WriteFile(t, path, + `{"type":"metadata","protocol_version":"1.3"}`+"\n"+ + `{"timestamp":1770983427.123,"message":{"type":"StatusUpdate","payload":{"token_usage":{"input_other":100,"output":50,"input_cache_read":10,"input_cache_creation":20},"message_id":"msg-1"}}}`+"\n") + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderKimi, + Model: "kimi-k2", + SessionID: "session-a", + Project: "kimi", + Tokens: usage.TokenUsage{ + InputTokens: 100, + OutputTokens: 50, + CacheCreationInputTokens: 20, + CacheReadInputTokens: 10, + TotalTokens: 180, + }, + }) +} diff --git a/internal/provider/openclaw/loader.go b/internal/provider/openclaw/loader.go new file mode 100644 index 0000000..50352d6 --- /dev/null +++ b/internal/provider/openclaw/loader.go @@ -0,0 +1,133 @@ +package openclaw + +import ( + "path/filepath" + "sort" + "strings" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + files := make([]string, 0) + for _, path := range paths { + files = append(files, shared.CollectFiles(path, isSessionFile)...) + } + sort.Strings(files) + files = shared.FilterFiles(files, filter) + + entries := make([]usage.Entry, 0) + for _, file := range files { + fileEntries, err := parseSessionFile(file) + if err != nil { + return nil, err + } + entries = append(entries, fileEntries...) + } + shared.SortEntries(entries) + return entries, nil +} + +func isSessionFile(path string) bool { + name := filepath.Base(path) + index := strings.Index(name, ".jsonl") + if index < 0 { + return false + } + suffix := name[index:] + return suffix == ".jsonl" || strings.HasPrefix(suffix, ".jsonl.deleted.") || strings.HasPrefix(suffix, ".jsonl.reset.") +} + +func parseSessionFile(path string) ([]usage.Entry, error) { + lines, err := shared.ReadJSONLines(path) + if err != nil { + return nil, err + } + sessionID := openClawSessionID(path) + currentModel := "" + currentProvider := "" + entries := make([]usage.Entry, 0) + for _, line := range lines { + record := line.Value + if isModelChange(record) { + source := shared.ObjectAt(record["data"]) + if source == nil { + source = record + } + if model := shared.FirstStringField(source, "modelId", "model"); model != "" { + currentModel = model + } + if provider := shared.StringField(source, "provider"); provider != "" { + currentProvider = provider + } + continue + } + if shared.StringField(record, "type") != "message" { + continue + } + message := shared.ObjectAt(record["message"]) + if shared.StringField(message, "role") != "assistant" { + continue + } + usageBlock := shared.ObjectAt(message["usage"]) + if usageBlock == nil { + continue + } + timestamp, ok := shared.ParseTimestamp(message["timestamp"]) + if !ok { + timestamp, ok = shared.ParseTimestamp(record["timestamp"]) + } + if !ok { + timestamp = shared.FileModifiedTime(path) + } + model := shared.FirstStringField(message, "modelId", "model") + if model == "" { + model = currentModel + } + if model == "" { + model = "unknown" + } + provider := shared.StringField(message, "provider") + if provider == "" { + provider = currentProvider + } + tokens := usage.TokenUsage{ + InputTokens: shared.UintField(usageBlock, "input"), + OutputTokens: shared.UintField(usageBlock, "output"), + CacheCreationInputTokens: shared.UintField(usageBlock, "cacheWrite"), + CacheReadInputTokens: shared.UintField(usageBlock, "cacheRead"), + } + tokens = shared.ApplyTotalFallback(tokens, shared.UintField(usageBlock, "totalTokens")) + if !shared.NonZero(tokens) { + continue + } + entry := shared.BaseEntry(usage.ProviderOpenClaw, timestamp, "openclaw", "OpenClaw", sessionID, "[openclaw] "+model, "OpenClaw", tokens) + shared.SetSource(&entry, path, line.Line, line.Start, line.End) + entry.ID = shared.StableEntryID(entry, provider) + entries = append(entries, entry) + } + return entries, nil +} + +func isModelChange(record map[string]any) bool { + if shared.StringField(record, "type") == "model_change" { + return true + } + return shared.StringField(record, "type") == "custom" && shared.StringField(record, "customType") == "model-snapshot" +} + +func openClawSessionID(path string) string { + name := filepath.Base(path) + index := strings.Index(name, ".jsonl") + if index < 0 { + if name == "" { + return "unknown" + } + return name + } + if index == 0 { + return name + } + return name[:index] +} diff --git a/internal/provider/openclaw/provider.go b/internal/provider/openclaw/provider.go new file mode 100644 index 0000000..5cbec48 --- /dev/null +++ b/internal/provider/openclaw/provider.go @@ -0,0 +1,37 @@ +package openclaw + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads OpenClaw usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a OpenClaw provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the OpenClaw provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderOpenClaw } + +// WithFileFilter returns a OpenClaw provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized OpenClaw usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/openclaw/provider_test.go b/internal/provider/openclaw/provider_test.go new file mode 100644 index 0000000..96a33a5 --- /dev/null +++ b/internal/provider/openclaw/provider_test.go @@ -0,0 +1,35 @@ +package openclaw + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the OpenClaw smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + path := filepath.Join(dir, "agents", "main", "sessions", "abc.jsonl") + shared.WriteFile(t, path, + `{"type":"model_change","provider":"openai-codex","modelId":"gpt-5.2"}`+"\n"+ + `{"type":"message","message":{"role":"assistant","usage":{"input":1660,"output":55,"cacheRead":108928},"timestamp":1769753935279}}`+"\n") + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderOpenClaw, + Model: "[openclaw] gpt-5.2", + SessionID: "abc", + Project: "openclaw", + Tokens: usage.TokenUsage{ + InputTokens: 1660, + OutputTokens: 55, + CacheReadInputTokens: 108928, + TotalTokens: 110643, + }, + }) +} diff --git a/internal/agentusage/opencode.go b/internal/provider/opencode/loader.go similarity index 61% rename from internal/agentusage/opencode.go rename to internal/provider/opencode/loader.go index 00c6d4a..54eec95 100644 --- a/internal/agentusage/opencode.go +++ b/internal/provider/opencode/loader.go @@ -1,4 +1,4 @@ -package agentusage +package opencode import ( "bytes" @@ -10,6 +10,7 @@ import ( "strings" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -17,12 +18,12 @@ import ( // messages own parts. Tokens live on the message, the working directory lives // on both the session and the message, and file edits live in the parts. All // three tables are read together because an entry needs all three. -type openCodeSession struct { +type session struct { directory string version string } -type openCodeMessage struct { +type message struct { id string sessionID string created int64 @@ -30,16 +31,16 @@ type openCodeMessage struct { parts []map[string]any } -func loadOpenCodeEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { entriesByID := make(map[string]usage.Entry) for _, root := range paths { - rootEntries, err := loadOpenCodeRoot(root, filter) + rootEntries, err := loadRoot(root, filter) if err != nil { return nil, err } for _, entry := range rootEntries { if entry.ID == "" { - entry.ID = stableEntryID(entry) + entry.ID = shared.StableEntryID(entry) } if _, exists := entriesByID[entry.ID]; !exists { entriesByID[entry.ID] = entry @@ -50,11 +51,11 @@ func loadOpenCodeEntries(paths []string, filter usage.FileFilter) ([]usage.Entry for _, entry := range entriesByID { entries = append(entries, entry) } - sortEntries(entries) + shared.SortEntries(entries) return entries, nil } -func loadOpenCodeRoot(root string, filter usage.FileFilter) ([]usage.Entry, error) { +func loadRoot(root string, filter usage.FileFilter) ([]usage.Entry, error) { info, err := os.Stat(root) if err != nil { return nil, nil @@ -62,17 +63,17 @@ func loadOpenCodeRoot(root string, filter usage.FileFilter) ([]usage.Entry, erro if !info.IsDir() { switch strings.ToLower(filepath.Ext(root)) { case ".db": - return loadOpenCodeDatabase(root) + return loadDatabase(root) case ".json": - return loadOpenCodeLegacyFile(root) + return loadLegacyFile(root) } return nil, nil } entries := make([]usage.Entry, 0) seenIDs := make(map[string]bool) - for _, dbPath := range openCodeDBPaths(root) { - dbEntries, err := loadOpenCodeDatabase(dbPath) + for _, dbPath := range dbPaths(root) { + dbEntries, err := loadDatabase(dbPath) if err != nil { return nil, err } @@ -86,15 +87,15 @@ func loadOpenCodeRoot(root string, filter usage.FileFilter) ([]usage.Entry, erro // Older OpenCode releases wrote one JSON file per message instead of the // database. Their users keep those files, so keep reading them. - files := collectExt(filepath.Join(root, "storage", "message"), ".json") + files := shared.CollectExt(filepath.Join(root, "storage", "message"), ".json") sort.Strings(files) - files = filterFiles(files, filter) + files = shared.FilterFiles(files, filter) for _, file := range files { stem := strings.TrimSuffix(filepath.Base(file), filepath.Ext(file)) - if seenIDs[stableOpenCodeMessageID(stem)] { + if seenIDs[stableMessageID(stem)] { continue } - legacy, err := loadOpenCodeLegacyFile(file) + legacy, err := loadLegacyFile(file) if err != nil { return nil, err } @@ -108,21 +109,21 @@ func loadOpenCodeRoot(root string, filter usage.FileFilter) ([]usage.Entry, erro entries = append(entries, entry) } } - sortEntries(entries) + shared.SortEntries(entries) return entries, nil } -// openCodeDBPaths lists the databases in a data root. OpenCode ships release +// dbPaths lists the databases in a data root. OpenCode ships release // channels side by side ("opencode.db", "opencode-nightly.db"), each a full // database of its own. -func openCodeDBPaths(root string) []string { +func dbPaths(root string) []string { matches, err := filepath.Glob(filepath.Join(root, "opencode*.db")) if err != nil { return nil } paths := make([]string, 0, len(matches)) for _, candidate := range matches { - if existingSQLiteFile(candidate) { + if shared.ExistingSQLiteFile(candidate) { paths = append(paths, candidate) } } @@ -130,25 +131,25 @@ func openCodeDBPaths(root string) []string { return paths } -func loadOpenCodeDatabase(path string) ([]usage.Entry, error) { - db, err := openSQLite(path) +func loadDatabase(path string) ([]usage.Entry, error) { + db, err := shared.OpenSQLite(path) if err != nil { return nil, err } defer db.Close() - sessions := queryOpenCodeSessions(db) - messages := queryOpenCodeMessages(db) + sessions := querySessions(db) + messages := queryMessages(db) if len(messages) == 0 { return nil, nil } - attachOpenCodeParts(db, messages) + attachParts(db, messages) // Token counts on a message are the running context total, not that turn's // cost, so each entry bills the growth since the previous message of the // same session. Messages are ordered by creation time for that subtraction // to mean anything. - bySession := make(map[string][]*openCodeMessage) + bySession := make(map[string][]*message) for _, message := range messages { bySession[message.sessionID] = append(bySession[message.sessionID], message) } @@ -161,13 +162,13 @@ func loadOpenCodeDatabase(path string) ([]usage.Entry, error) { } return sessionMessages[i].id < sessionMessages[j].id }) - entries = append(entries, openCodeSessionEntries(path, sessions[sessionID], sessionMessages)...) + entries = append(entries, sessionEntries(path, sessions[sessionID], sessionMessages)...) } return entries, nil } -func queryOpenCodeSessions(db *sql.DB) map[string]openCodeSession { - sessions := make(map[string]openCodeSession) +func querySessions(db *sql.DB) map[string]session { + sessions := make(map[string]session) rows, err := db.Query(`SELECT id, COALESCE(directory, ''), COALESCE(version, '') FROM session`) if err != nil { return sessions @@ -178,32 +179,32 @@ func queryOpenCodeSessions(db *sql.DB) map[string]openCodeSession { if err := rows.Scan(&id, &directory, &version); err != nil { continue } - sessions[id] = openCodeSession{directory: directory, version: version} + sessions[id] = session{directory: directory, version: version} } return sessions } -func queryOpenCodeMessages(db *sql.DB) []*openCodeMessage { +func queryMessages(db *sql.DB) []*message { rows, err := db.Query(`SELECT id, session_id, time_created, CAST(data AS TEXT) FROM message`) if err != nil { return nil } defer rows.Close() - messages := make([]*openCodeMessage, 0) + messages := make([]*message, 0) for rows.Next() { var id, sessionID, data string var created int64 if err := rows.Scan(&id, &sessionID, &created, &data); err != nil { continue } - record := decodeJSONObjectString(data) + record := shared.DecodeJSONObjectString(data) if record == nil { continue } - messages = append(messages, &openCodeMessage{ - id: firstNonEmpty(stringField(record, "id"), id), - sessionID: firstNonEmpty(stringField(record, "sessionID"), sessionID), + messages = append(messages, &message{ + id: shared.FirstNonEmpty(shared.StringField(record, "id"), id), + sessionID: shared.FirstNonEmpty(shared.StringField(record, "sessionID"), sessionID), created: created, data: record, }) @@ -211,8 +212,8 @@ func queryOpenCodeMessages(db *sql.DB) []*openCodeMessage { return messages } -func attachOpenCodeParts(db *sql.DB, messages []*openCodeMessage) { - byID := make(map[string][]*openCodeMessage, len(messages)) +func attachParts(db *sql.DB, messages []*message) { + byID := make(map[string][]*message, len(messages)) for _, message := range messages { byID[message.id] = append(byID[message.id], message) } @@ -231,7 +232,7 @@ func attachOpenCodeParts(db *sql.DB, messages []*openCodeMessage) { if len(targets) == 0 { continue } - record := decodeJSONObjectString(data) + record := shared.DecodeJSONObjectString(data) if record == nil { continue } @@ -241,23 +242,23 @@ func attachOpenCodeParts(db *sql.DB, messages []*openCodeMessage) { } } -func openCodeSessionEntries(source string, session openCodeSession, messages []*openCodeMessage) []usage.Entry { +func sessionEntries(source string, session session, messages []*message) []usage.Entry { entries := make([]usage.Entry, 0, len(messages)) var previous usage.TokenUsage for _, message := range messages { // User messages carry an all-zero token block. Letting one reset the // running counters would make the next assistant message look like a // fresh session and bill it for the whole context again. - reported := openCodeTokens(message.data) - if !nonZero(reported) { + reported := tokens(message.data) + if !shared.NonZero(reported) { continue } - billed, ok := billOpenCodeTokens(reported, previous) + billed, ok := billTokens(reported, previous) previous = reported if !ok { continue } - entry, ok := openCodeEntry(source, session, message, billed) + entry, ok := newEntry(source, session, message, billed) if !ok { continue } @@ -266,31 +267,31 @@ func openCodeSessionEntries(source string, session openCodeSession, messages []* return entries } -// openCodeTokens reads a message's token block exactly as OpenCode wrote it. +// tokens reads a message's token block exactly as OpenCode wrote it. // The fields do not all mean the same thing: input and the cache counters are // running context totals that grow with every turn, while output and reasoning -// are what the model produced on this turn alone. billOpenCodeTokens is what +// are what the model produced on this turn alone. billTokens is what // turns this mixture into one turn's cost. -func openCodeTokens(record map[string]any) usage.TokenUsage { - block := objectAt(record["tokens"]) +func tokens(record map[string]any) usage.TokenUsage { + block := shared.ObjectAt(record["tokens"]) if block == nil { return usage.TokenUsage{} } - cache := objectAt(block["cache"]) + cache := shared.ObjectAt(block["cache"]) return usage.TokenUsage{ - InputTokens: uintField(block, "input"), - OutputTokens: uintField(block, "output"), - ReasoningOutputTokens: uintField(block, "reasoning"), - CacheCreationInputTokens: uintField(cache, "write"), - CacheReadInputTokens: uintField(cache, "read"), + InputTokens: shared.UintField(block, "input"), + OutputTokens: shared.UintField(block, "output"), + ReasoningOutputTokens: shared.UintField(block, "reasoning"), + CacheCreationInputTokens: shared.UintField(cache, "write"), + CacheReadInputTokens: shared.UintField(cache, "read"), } } -// billOpenCodeTokens charges a message for one turn: the growth of the running +// billTokens charges a message for one turn: the growth of the running // counters plus the per-turn counters as reported. The message's own "total" // field is deliberately ignored — it sums the running counters, so it inherits // their whole history and is not this turn's cost. -func billOpenCodeTokens(reported, previous usage.TokenUsage) (usage.TokenUsage, bool) { +func billTokens(reported, previous usage.TokenUsage) (usage.TokenUsage, bool) { billed := usage.TokenUsage{ InputTokens: growth(reported.InputTokens, previous.InputTokens), CacheCreationInputTokens: growth(reported.CacheCreationInputTokens, previous.CacheCreationInputTokens), @@ -298,10 +299,10 @@ func billOpenCodeTokens(reported, previous usage.TokenUsage) (usage.TokenUsage, OutputTokens: reported.OutputTokens, ReasoningOutputTokens: reported.ReasoningOutputTokens, } - if !nonZero(billed) { + if !shared.NonZero(billed) { return usage.TokenUsage{}, false } - return applyTotalFallback(billed, 0), true + return shared.ApplyTotalFallback(billed, 0), true } // growth returns how much a running counter advanced. A counter that shrank @@ -313,23 +314,23 @@ func growth(current, previous uint64) uint64 { return current - previous } -func openCodeEntry(source string, session openCodeSession, message *openCodeMessage, tokens usage.TokenUsage) (usage.Entry, bool) { +func newEntry(source string, session session, message *message, tokens usage.TokenUsage) (usage.Entry, bool) { record := message.data - model := stringField(record, "modelID") + model := shared.StringField(record, "modelID") if model == "" { return usage.Entry{}, false } timestamp := time.UnixMilli(message.created).UTC() - if parsed, ok := parseTimestamp(objectAt(record["time"])["created"]); ok { + if parsed, ok := shared.ParseTimestamp(shared.ObjectAt(record["time"])["created"]); ok { timestamp = parsed } // The message records the directory the agent actually ran in; the session // directory covers messages written before that field existed. - cwd := firstNonEmpty( - stringField(objectAt(record["path"]), "cwd"), - stringField(objectAt(record["path"]), "root"), + cwd := shared.FirstNonEmpty( + shared.StringField(shared.ObjectAt(record["path"]), "cwd"), + shared.StringField(shared.ObjectAt(record["path"]), "root"), session.directory, ) project, projectPath := usage.UnknownProject, "" @@ -337,59 +338,59 @@ func openCodeEntry(source string, session openCodeSession, message *openCodeMess project, projectPath = name, path } - sessionID := firstNonEmpty(message.sessionID, usage.UnknownProject) - entry := baseEntry(usage.ProviderOpenCode, timestamp, project, projectPath, sessionID, model, "OpenCode", tokens) - setSource(&entry, source, 0, 0, 0) - entry.ID = stableOpenCodeMessageID(message.id) + sessionID := shared.FirstNonEmpty(message.sessionID, usage.UnknownProject) + entry := shared.BaseEntry(usage.ProviderOpenCode, timestamp, project, projectPath, sessionID, model, "OpenCode", tokens) + shared.SetSource(&entry, source, 0, 0, 0) + entry.ID = stableMessageID(message.id) if entry.ID == "" { - entry.ID = stableEntryID(entry) + entry.ID = shared.StableEntryID(entry) } for _, part := range message.parts { - for _, change := range openCodePartChanges(part, cwd) { + for _, change := range partChanges(part, cwd) { entry.ApplyFileChange(change) } } return entry, true } -// openCodePartChanges extracts the file edits one message part performed. +// partChanges extracts the file edits one message part performed. // Only tool calls count: a "patch" part is a git snapshot of the whole // worktree, which lists files nobody edited (.DS_Store, build output) and // carries no line counts, so treating it as agent work invents writes. -func openCodePartChanges(part map[string]any, cwd string) []usage.FileChange { - if stringField(part, "type") != "tool" { +func partChanges(part map[string]any, cwd string) []usage.FileChange { + if shared.StringField(part, "type") != "tool" { return nil } - state := objectAt(part["state"]) - if stringField(state, "status") != "completed" { + state := shared.ObjectAt(part["state"]) + if shared.StringField(state, "status") != "completed" { return nil } - input := objectAt(state["input"]) - metadata := objectAt(state["metadata"]) + input := shared.ObjectAt(state["input"]) + metadata := shared.ObjectAt(state["metadata"]) - switch stringField(part, "tool") { + switch shared.StringField(part, "tool") { case "edit": - return openCodeEditChanges(input, metadata, cwd) + return editChanges(input, metadata, cwd) case "write": - return openCodeWriteChanges(input, metadata, cwd) + return writeChanges(input, metadata, cwd) case "patch", "apply_patch": - return openCodeApplyPatchChanges(input, metadata, cwd) + return applyPatchChanges(input, metadata, cwd) default: return nil } } -// openCodeEditChanges prefers the diff OpenCode computed. Without it the +// editChanges prefers the diff OpenCode computed. Without it the // replaced and replacing strings still bound the change. -func openCodeEditChanges(input, metadata map[string]any, cwd string) []usage.FileChange { - path := usage.ResolvePath(cwd, stringField(input, "filePath")) +func editChanges(input, metadata map[string]any, cwd string) []usage.FileChange { + path := usage.ResolvePath(cwd, shared.StringField(input, "filePath")) - if diff := objectAt(metadata["filediff"]); diff != nil { - added := uintField(diff, "additions") - removed := uintField(diff, "deletions") - if diffPath := firstNonEmpty(stringField(diff, "filePath"), stringField(diff, "file")); diffPath != "" { + if diff := shared.ObjectAt(metadata["filediff"]); diff != nil { + added := shared.UintField(diff, "additions") + removed := shared.UintField(diff, "deletions") + if diffPath := shared.FirstNonEmpty(shared.StringField(diff, "filePath"), shared.StringField(diff, "file")); diffPath != "" { path = usage.ResolvePath(cwd, diffPath) } if added != 0 || removed != 0 { @@ -397,46 +398,46 @@ func openCodeEditChanges(input, metadata map[string]any, cwd string) []usage.Fil } } - added := usage.CountLines(stringField(input, "newString")) - removed := usage.CountLines(stringField(input, "oldString")) + added := usage.CountLines(shared.StringField(input, "newString")) + removed := usage.CountLines(shared.StringField(input, "oldString")) if path == "" && added == 0 && removed == 0 { return nil } return []usage.FileChange{{Path: path, LinesAdded: added, LinesRemoved: removed}} } -// openCodeWriteChanges counts a write as adding its whole content. Overwriting +// writeChanges counts a write as adding its whole content. Overwriting // an existing file replaces lines this record does not describe, so only the // added side is known. -func openCodeWriteChanges(input, metadata map[string]any, cwd string) []usage.FileChange { - path := usage.ResolvePath(cwd, firstNonEmpty( - stringField(input, "filePath"), - stringField(metadata, "filepath"), +func writeChanges(input, metadata map[string]any, cwd string) []usage.FileChange { + path := usage.ResolvePath(cwd, shared.FirstNonEmpty( + shared.StringField(input, "filePath"), + shared.StringField(metadata, "filepath"), )) if path == "" { return nil } - return []usage.FileChange{{Path: path, LinesAdded: usage.CountLines(stringField(input, "content"))}} + return []usage.FileChange{{Path: path, LinesAdded: usage.CountLines(shared.StringField(input, "content"))}} } -// openCodeApplyPatchChanges prefers the per-file summary OpenCode attaches and +// applyPatchChanges prefers the per-file summary OpenCode attaches and // falls back to counting the raw patch envelope itself. -func openCodeApplyPatchChanges(input, metadata map[string]any, cwd string) []usage.FileChange { +func applyPatchChanges(input, metadata map[string]any, cwd string) []usage.FileChange { if files, ok := metadata["files"].([]any); ok && len(files) > 0 { changes := make([]usage.FileChange, 0, len(files)) for _, raw := range files { - file := objectAt(raw) + file := shared.ObjectAt(raw) if file == nil { continue } - path := usage.ResolvePath(cwd, firstNonEmpty(stringField(file, "filePath"), stringField(file, "file"))) + path := usage.ResolvePath(cwd, shared.FirstNonEmpty(shared.StringField(file, "filePath"), shared.StringField(file, "file"))) if path == "" { continue } changes = append(changes, usage.FileChange{ Path: path, - LinesAdded: uintField(file, "additions"), - LinesRemoved: uintField(file, "deletions"), + LinesAdded: shared.UintField(file, "additions"), + LinesRemoved: shared.UintField(file, "deletions"), }) } if len(changes) > 0 { @@ -444,17 +445,17 @@ func openCodeApplyPatchChanges(input, metadata map[string]any, cwd string) []usa } } - patchText := firstNonEmpty( - stringField(input, "patchText"), - stringField(input, "patch"), - stringField(metadata, "patch"), + patchText := shared.FirstNonEmpty( + shared.StringField(input, "patchText"), + shared.StringField(input, "patch"), + shared.StringField(metadata, "patch"), ) - return parseOpenCodePatchEnvelope(patchText, cwd) + return parsePatchEnvelope(patchText, cwd) } -// parseOpenCodePatchEnvelope counts the added and removed lines of each file in +// parsePatchEnvelope counts the added and removed lines of each file in // an apply_patch envelope. One envelope can carry any number of files. -func parseOpenCodePatchEnvelope(patchText, cwd string) []usage.FileChange { +func parsePatchEnvelope(patchText, cwd string) []usage.FileChange { if strings.TrimSpace(patchText) == "" { return nil } @@ -479,7 +480,7 @@ func parseOpenCodePatchEnvelope(patchText, cwd string) []usage.FileChange { strings.HasPrefix(line, "*** Update File: "), strings.HasPrefix(line, "*** Delete File: "): flush() - current.Path = usage.ResolvePath(cwd, openCodePatchHeaderPath(line)) + current.Path = usage.ResolvePath(cwd, patchHeaderPath(line)) case strings.HasPrefix(line, "*** End Patch"): flush() case current.Path == "": @@ -496,7 +497,7 @@ func parseOpenCodePatchEnvelope(patchText, cwd string) []usage.FileChange { return changes } -func openCodePatchHeaderPath(line string) string { +func patchHeaderPath(line string) string { for _, prefix := range []string{"*** Add File: ", "*** Update File: ", "*** Delete File: "} { if path, ok := strings.CutPrefix(line, prefix); ok { return strings.TrimSpace(path) @@ -505,9 +506,9 @@ func openCodePatchHeaderPath(line string) string { return "" } -// loadOpenCodeLegacyFile reads one message from the pre-database storage layout, +// loadLegacyFile reads one message from the pre-database storage layout, // where the message and its parts each lived in their own JSON file. -func loadOpenCodeLegacyFile(path string) ([]usage.Entry, error) { +func loadLegacyFile(path string) ([]usage.Entry, error) { data, err := os.ReadFile(path) if err != nil { return nil, err @@ -519,19 +520,19 @@ func loadOpenCodeLegacyFile(path string) ([]usage.Entry, error) { return nil, nil } - message := &openCodeMessage{ - id: stringField(record, "id"), - sessionID: stringField(record, "sessionID"), + message := &message{ + id: shared.StringField(record, "id"), + sessionID: shared.StringField(record, "sessionID"), data: record, - parts: readOpenCodeLegacyParts(path, stringField(record, "id")), + parts: readLegacyParts(path, shared.StringField(record, "id")), } // One file holds one message, so there is no predecessor to subtract from: // billing against a zero baseline charges it for everything it reports. - tokens, ok := billOpenCodeTokens(openCodeTokens(record), usage.TokenUsage{}) + tokens, ok := billTokens(tokens(record), usage.TokenUsage{}) if !ok { return nil, nil } - entry, ok := openCodeEntry(path, openCodeSession{}, message, tokens) + entry, ok := newEntry(path, session{}, message, tokens) if !ok { return nil, nil } @@ -539,15 +540,15 @@ func loadOpenCodeLegacyFile(path string) ([]usage.Entry, error) { return []usage.Entry{entry}, nil } -// readOpenCodeLegacyParts loads the part files stored alongside a legacy +// readLegacyParts loads the part files stored alongside a legacy // message: storage/message//.json pairs with // storage/part//*.json. -func readOpenCodeLegacyParts(messagePath, messageID string) []map[string]any { +func readLegacyParts(messagePath, messageID string) []map[string]any { if messageID == "" { return nil } storage := filepath.Dir(filepath.Dir(filepath.Dir(messagePath))) - files := collectExt(filepath.Join(storage, "part", messageID), ".json") + files := shared.CollectExt(filepath.Join(storage, "part", messageID), ".json") sort.Strings(files) parts := make([]map[string]any, 0, len(files)) @@ -567,10 +568,10 @@ func readOpenCodeLegacyParts(messagePath, messageID string) []map[string]any { return parts } -// stableOpenCodeMessageID keys an entry to the message id OpenCode assigned. +// stableMessageID keys an entry to the message id OpenCode assigned. // The id travels with the message no matter which database file or storage // layout holds it, so re-ingesting the same message never double-counts. -func stableOpenCodeMessageID(messageID string) string { +func stableMessageID(messageID string) string { if strings.TrimSpace(messageID) == "" { return "" } diff --git a/internal/agentusage/opencode_test.go b/internal/provider/opencode/loader_test.go similarity index 95% rename from internal/agentusage/opencode_test.go rename to internal/provider/opencode/loader_test.go index 90d90f2..e676082 100644 --- a/internal/agentusage/opencode_test.go +++ b/internal/provider/opencode/loader_test.go @@ -1,10 +1,11 @@ -package agentusage +package opencode import ( "database/sql" "path/filepath" "testing" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -73,14 +74,14 @@ func TestOpenCodeBillsRunningInputAndPerTurnOutput(t *testing.T) { {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":1400,"output":30,"reasoning":10,"total":1440}}`}, }, nil) - entries, err := loadOpenCodeDatabase(path) + entries, err := loadDatabase(path) if err != nil { t.Fatal(err) } if len(entries) != 2 { t.Fatalf("len(entries) = %d, want 2", len(entries)) } - sortEntries(entries) + shared.SortEntries(entries) first, second := entries[0], entries[1] if first.Usage.InputTokens != 1000 || first.Usage.OutputTokens != 50 || first.Usage.ReasoningOutputTokens != 20 { @@ -113,14 +114,14 @@ func TestOpenCodeIgnoresZeroTokenMessages(t *testing.T) { {id: "msg-3", sessionID: "ses-1", created: 3000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":1200,"output":20}}`}, }, nil) - entries, err := loadOpenCodeDatabase(path) + entries, err := loadDatabase(path) if err != nil { t.Fatal(err) } if len(entries) != 2 { t.Fatalf("len(entries) = %d, want 2 (the user turn bills nothing)", len(entries)) } - sortEntries(entries) + shared.SortEntries(entries) if got := entries[1].Usage.InputTokens; got != 200 { t.Fatalf("input after a user turn = %d, want 200 (the user turn must not reset the counter)", got) } @@ -136,11 +137,11 @@ func TestOpenCodeRestartsBillingAfterCompaction(t *testing.T) { {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":300,"output":10}}`}, }, nil) - entries, err := loadOpenCodeDatabase(path) + entries, err := loadDatabase(path) if err != nil { t.Fatal(err) } - sortEntries(entries) + shared.SortEntries(entries) if len(entries) != 2 { t.Fatalf("len(entries) = %d, want 2", len(entries)) } @@ -160,11 +161,11 @@ func TestOpenCodeProjectComesFromMessageCwd(t *testing.T) { {id: "msg-3", sessionID: "ses-2", created: 3000, data: `{"role":"assistant","modelID":"gpt-5","tokens":{"input":30,"output":5}}`}, }, nil) - entries, err := loadOpenCodeDatabase(path) + entries, err := loadDatabase(path) if err != nil { t.Fatal(err) } - sortEntries(entries) + shared.SortEntries(entries) if len(entries) != 3 { t.Fatalf("len(entries) = %d, want 3", len(entries)) } @@ -197,7 +198,7 @@ func TestOpenCodeCollectsFileChangesFromTools(t *testing.T) { {id: "prt-3", sessionID: "msg-1", data: `{"type":"tool","tool":"apply_patch","state":{"status":"completed","metadata":{"files":[{"filePath":"/abs/other.go","additions":2,"deletions":2}]}}}`}, }) - entries, err := loadOpenCodeDatabase(path) + entries, err := loadDatabase(path) if err != nil { t.Fatal(err) } @@ -252,7 +253,7 @@ func TestOpenCodeIgnoresSnapshotPatchParts(t *testing.T) { {id: "prt-1", sessionID: "msg-1", data: `{"type":"patch","hash":"deadbeef","files":["/repo/demo/.DS_Store"]}`}, }) - entries, err := loadOpenCodeDatabase(path) + entries, err := loadDatabase(path) if err != nil { t.Fatal(err) } @@ -278,7 +279,7 @@ func TestOpenCodeSkipsIncompleteToolCalls(t *testing.T) { {id: "prt-2", sessionID: "msg-1", data: `{"type":"tool","tool":"read","state":{"status":"completed","input":{"filePath":"/repo/demo/b.go"}}}`}, }) - entries, err := loadOpenCodeDatabase(path) + entries, err := loadDatabase(path) if err != nil { t.Fatal(err) } @@ -310,7 +311,7 @@ func TestOpenCodeParsesPatchEnvelope(t *testing.T) { {id: "prt-1", sessionID: "msg-1", data: `{"type":"tool","tool":"apply_patch","state":{"status":"completed","input":{"patchText":"` + patch + `"}}}`}, }) - entries, err := loadOpenCodeDatabase(path) + entries, err := loadDatabase(path) if err != nil { t.Fatal(err) } @@ -347,18 +348,18 @@ func TestOpenCodeEntryIDSurvivesRelocation(t *testing.T) { []openCodeRow{{id: "ses-1", sessionID: "/elsewhere/demo"}}, []openCodeRow{{id: "msg-1", sessionID: "ses-1", created: 1000, data: message}}, nil) - a, err := loadOpenCodeDatabase(first) + a, err := loadDatabase(first) if err != nil { t.Fatal(err) } - b, err := loadOpenCodeDatabase(second) + b, err := loadDatabase(second) if err != nil { t.Fatal(err) } if a[0].ID != b[0].ID { t.Fatalf("ids differ across databases: %q vs %q", a[0].ID, b[0].ID) } - if a[0].ID != stableOpenCodeMessageID("msg-1") { + if a[0].ID != stableMessageID("msg-1") { t.Fatalf("id = %q, want the message-keyed id", a[0].ID) } } @@ -378,7 +379,7 @@ func TestOpenCodeReadsEveryChannelDatabase(t *testing.T) { }}, nil) } - entries, err := loadOpenCodeRoot(root, nil) + entries, err := loadRoot(root, nil) if err != nil { t.Fatal(err) } diff --git a/internal/provider/opencode/provider.go b/internal/provider/opencode/provider.go new file mode 100644 index 0000000..3455d71 --- /dev/null +++ b/internal/provider/opencode/provider.go @@ -0,0 +1,37 @@ +package opencode + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads OpenCode usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a OpenCode provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the OpenCode provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderOpenCode } + +// WithFileFilter returns a OpenCode provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized OpenCode usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/opencode/provider_test.go b/internal/provider/opencode/provider_test.go new file mode 100644 index 0000000..e04c80a --- /dev/null +++ b/internal/provider/opencode/provider_test.go @@ -0,0 +1,34 @@ +package opencode + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the OpenCode smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + path := filepath.Join(dir, "storage", "message", "session-a", "msg-1.json") + shared.WriteFile(t, path, `{"id":"msg-1","sessionID":"session-a","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","path":{"cwd":"/repo/demo"},"time":{"created":1767312000000},"tokens":{"input":100,"output":50,"cache":{"read":10,"write":20}},"cost":0}`) + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderOpenCode, + Model: "claude-sonnet-4-20250514", + SessionID: "session-a", + Project: "demo", + Tokens: usage.TokenUsage{ + InputTokens: 100, + OutputTokens: 50, + CacheCreationInputTokens: 20, + CacheReadInputTokens: 10, + TotalTokens: 180, + }, + }) +} diff --git a/internal/provider/pi/loader.go b/internal/provider/pi/loader.go new file mode 100644 index 0000000..d5a55b2 --- /dev/null +++ b/internal/provider/pi/loader.go @@ -0,0 +1,97 @@ +package pi + +import ( + "path/filepath" + "sort" + "strings" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + files := make([]string, 0) + for _, path := range paths { + files = append(files, shared.CollectExt(path, ".jsonl")...) + } + sort.Strings(files) + files = shared.FilterFiles(files, filter) + + entries := make([]usage.Entry, 0) + for _, file := range files { + fileEntries, err := parseSessionFile(file) + if err != nil { + return nil, err + } + entries = append(entries, fileEntries...) + } + shared.SortEntries(entries) + return entries, nil +} + +func parseSessionFile(path string) ([]usage.Entry, error) { + lines, err := shared.ReadJSONLines(path, `"usage"`, `"message"`) + if err != nil { + return nil, err + } + project := project(path) + sessionID := sessionID(path) + entries := make([]usage.Entry, 0) + for _, line := range lines { + if typ := shared.StringField(line.Value, "type"); typ != "" && typ != "message" { + continue + } + message := shared.ObjectAt(line.Value["message"]) + if shared.StringField(message, "role") != "assistant" { + continue + } + usageBlock := shared.ObjectAt(message["usage"]) + if usageBlock == nil { + continue + } + timestamp, ok := shared.ParseTimestamp(line.Value["timestamp"]) + if !ok { + continue + } + tokens := usage.TokenUsage{ + InputTokens: shared.UintField(usageBlock, "input"), + OutputTokens: shared.UintField(usageBlock, "output"), + CacheCreationInputTokens: shared.UintField(usageBlock, "cacheWrite"), + CacheReadInputTokens: shared.UintField(usageBlock, "cacheRead"), + } + tokens = shared.ApplyTotalFallback(tokens, shared.UintField(usageBlock, "totalTokens")) + if !shared.NonZero(tokens) { + continue + } + model := shared.StringField(message, "model") + if model != "" { + model = "[pi] " + model + } + entry := shared.BaseEntry(usage.ProviderPi, timestamp, project, project, sessionID, model, "pi-agent", tokens) + shared.SetSource(&entry, path, line.Line, line.Start, line.End) + entry.ID = shared.StableEntryID(entry) + entries = append(entries, entry) + } + return entries, nil +} + +func sessionID(path string) string { + stem := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) + if before, after, ok := strings.Cut(stem, "_"); ok && before != "" && after != "" { + return after + } + if stem == "" { + return "unknown" + } + return stem +} + +func project(path string) string { + parts := strings.Split(filepath.ToSlash(filepath.Clean(path)), "/") + for i, part := range parts { + if part == "sessions" && i+1 < len(parts) && parts[i+1] != "" { + return parts[i+1] + } + } + return usage.UnknownProject +} diff --git a/internal/provider/pi/provider.go b/internal/provider/pi/provider.go new file mode 100644 index 0000000..88fcd92 --- /dev/null +++ b/internal/provider/pi/provider.go @@ -0,0 +1,37 @@ +package pi + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads pi-agent usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a pi-agent provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the pi-agent provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderPi } + +// WithFileFilter returns a pi-agent provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized pi-agent usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/pi/provider_test.go b/internal/provider/pi/provider_test.go new file mode 100644 index 0000000..dc2585b --- /dev/null +++ b/internal/provider/pi/provider_test.go @@ -0,0 +1,31 @@ +package pi + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the pi-agent smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + path := filepath.Join(dir, "project-a", "agent_session-a.jsonl") + shared.WriteFile(t, path, `{"type":"message","timestamp":"2026-01-02T00:00:00.000Z","message":{"role":"assistant","model":"gpt-5","usage":{"totalTokens":333}}}`+"\n") + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderPi, + Model: "[pi] gpt-5", + SessionID: "session-a", + Project: usage.UnknownProject, + Tokens: usage.TokenUsage{ + OutputTokens: 333, + TotalTokens: 333, + }, + }) +} diff --git a/internal/provider/qwen/loader.go b/internal/provider/qwen/loader.go new file mode 100644 index 0000000..8122d40 --- /dev/null +++ b/internal/provider/qwen/loader.go @@ -0,0 +1,105 @@ +package qwen + +import ( + "path/filepath" + "sort" + "strings" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + files := make([]string, 0) + for _, root := range paths { + files = append(files, shared.CollectFiles(filepath.Join(root, "projects"), isChatFile)...) + if strings.Contains(filepath.ToSlash(root), "/projects/") { + files = append(files, shared.CollectFiles(root, isChatFile)...) + } + } + sort.Strings(files) + files = shared.FilterFiles(shared.UniqueStrings(files), filter) + + entries := make([]usage.Entry, 0) + for _, file := range files { + fileEntries, err := parseChatFile(file) + if err != nil { + return nil, err + } + entries = append(entries, fileEntries...) + } + shared.SortEntries(entries) + return entries, nil +} + +func isChatFile(path string) bool { + if !strings.EqualFold(filepath.Ext(path), ".jsonl") { + return false + } + parts := strings.Split(filepath.ToSlash(filepath.Clean(path)), "/") + for i := 0; i+3 < len(parts); i++ { + if parts[i] == "projects" && parts[i+2] == "chats" { + return true + } + } + return false +} + +func parseChatFile(path string) ([]usage.Entry, error) { + lines, err := shared.ReadJSONLines(path, `"usageMetadata"`) + if err != nil { + return nil, err + } + project := project(path) + fallback := shared.FileModifiedTime(path) + entries := make([]usage.Entry, 0) + for _, line := range lines { + record := line.Value + if shared.StringField(record, "type") != "assistant" { + continue + } + meta := shared.ObjectAt(record["usageMetadata"]) + if meta == nil { + continue + } + timestamp, ok := shared.ParseTimestamp(record["timestamp"]) + if !ok { + timestamp = fallback + } + sessionID := shared.StringField(record, "sessionId") + if sessionID == "" { + sessionID = project + "-" + strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) + } + model := shared.StringField(record, "model") + if model == "" { + model = "unknown" + } + tokens := usage.TokenUsage{ + InputTokens: shared.UintField(meta, "promptTokenCount"), + OutputTokens: shared.UintField(meta, "candidatesTokenCount"), + CacheReadInputTokens: shared.UintField(meta, "cachedContentTokenCount"), + ReasoningOutputTokens: shared.UintField(meta, + "thoughtsTokenCount", + ), + } + tokens = shared.ApplyTotalFallback(tokens, shared.UintField(meta, "totalTokenCount")) + if !shared.NonZero(tokens) { + continue + } + entry := shared.BaseEntry(usage.ProviderQwen, timestamp, "qwen", project, sessionID, model, "Qwen", tokens) + shared.SetSource(&entry, path, line.Line, line.Start, line.End) + entry.ID = shared.StableEntryID(entry) + entries = append(entries, entry) + } + return entries, nil +} + +func project(path string) string { + parts := strings.Split(filepath.ToSlash(filepath.Clean(path)), "/") + for i := 0; i+3 < len(parts); i++ { + if parts[i] == "projects" && parts[i+2] == "chats" && parts[i+1] != "" { + return parts[i+1] + } + } + return usage.UnknownProject +} diff --git a/internal/provider/qwen/provider.go b/internal/provider/qwen/provider.go new file mode 100644 index 0000000..c133531 --- /dev/null +++ b/internal/provider/qwen/provider.go @@ -0,0 +1,37 @@ +package qwen + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Qwen usage entries. +type Provider struct{ shared.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a Qwen provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = shared.NewBase(paths) + return p +} + +// Provider returns the Qwen provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderQwen } + +// WithFileFilter returns a Qwen provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized Qwen usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return shared.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/qwen/provider_test.go b/internal/provider/qwen/provider_test.go new file mode 100644 index 0000000..1b2ddb4 --- /dev/null +++ b/internal/provider/qwen/provider_test.go @@ -0,0 +1,34 @@ +package qwen + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the Qwen smoke test: a minimal fixture must produce +// exactly one entry with the expected identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "project-a", "chats", "chat-a.jsonl") + shared.WriteFile(t, path, `{"type":"assistant","timestamp":"2026-01-02T00:00:00.000Z","sessionId":"session-a","model":"qwen3-coder","usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":20,"thoughtsTokenCount":5,"cachedContentTokenCount":3,"totalTokenCount":38}}`+"\n") + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + Provider: usage.ProviderQwen, + Model: "qwen3-coder", + SessionID: "session-a", + Project: "qwen", + Tokens: usage.TokenUsage{ + InputTokens: 10, + OutputTokens: 20, + CacheReadInputTokens: 3, + ReasoningOutputTokens: 5, + TotalTokens: 38, + }, + }) +} diff --git a/internal/provider/shared/base.go b/internal/provider/shared/base.go new file mode 100644 index 0000000..70fd2bc --- /dev/null +++ b/internal/provider/shared/base.go @@ -0,0 +1,37 @@ +package shared + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// Base carries the scan configuration every provider needs: where to look and +// which files are already ingested. Providers embed it so the accessors and +// the WithPaths/WithFileFilter plumbing exist in one place instead of once per +// provider. +type Base struct { + paths []string + filter usage.FileFilter +} + +// NewBase returns a Base scanning the given data roots. +func NewBase(paths []string) Base { + return Base{paths: append([]string{}, paths...)} +} + +// Paths returns the data roots to scan. +func (b Base) Paths() []string { return b.paths } + +// Filter returns the file filter, or nil when every file must be parsed. +func (b Base) Filter() usage.FileFilter { return b.filter } + +// WithPathsSet returns a copy scanning the given data roots. +func (b Base) WithPathsSet(paths []string) Base { + b.paths = append([]string{}, paths...) + return b +} + +// WithFilterSet returns a copy that skips the source files filter rejects. +func (b Base) WithFilterSet(filter usage.FileFilter) Base { + b.filter = filter + return b +} diff --git a/internal/provider/shared/entries.go b/internal/provider/shared/entries.go new file mode 100644 index 0000000..942d357 --- /dev/null +++ b/internal/provider/shared/entries.go @@ -0,0 +1,16 @@ +package shared + +import ( + "sort" + + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// SortEntriesByTimestampDesc orders entries newest first, the order every +// provider returns them in. +func SortEntriesByTimestampDesc(entries []usage.Entry) []usage.Entry { + sort.Slice(entries, func(i, j int) bool { + return entries[i].Timestamp.After(entries[j].Timestamp) + }) + return entries +} diff --git a/internal/agentusage/helpers.go b/internal/provider/shared/helpers.go similarity index 79% rename from internal/agentusage/helpers.go rename to internal/provider/shared/helpers.go index 9a26e53..405397b 100644 --- a/internal/agentusage/helpers.go +++ b/internal/provider/shared/helpers.go @@ -1,4 +1,4 @@ -package agentusage +package shared import ( "bufio" @@ -20,7 +20,7 @@ import ( const maxInt64Uint = uint64(1<<63 - 1) -func collectFiles(root string, match func(string) bool) []string { +func CollectFiles(root string, match func(string) bool) []string { info, err := os.Stat(root) if err != nil { return nil @@ -55,8 +55,8 @@ func collectFiles(root string, match func(string) bool) []string { return files } -// filterFiles drops files the filter rejects. A nil filter keeps everything. -func filterFiles(files []string, filter usage.FileFilter) []string { +// FilterFiles drops files the filter rejects. A nil filter keeps everything. +func FilterFiles(files []string, filter usage.FileFilter) []string { if filter == nil { return files } @@ -69,14 +69,14 @@ func filterFiles(files []string, filter usage.FileFilter) []string { return kept } -func collectExt(root, ext string) []string { +func CollectExt(root, ext string) []string { ext = strings.ToLower(ext) - return collectFiles(root, func(path string) bool { + return CollectFiles(root, func(path string) bool { return strings.EqualFold(filepath.Ext(path), ext) }) } -func readJSONLines(path string, prefilter ...string) ([]lineJSON, error) { +func ReadJSONLines(path string, prefilter ...string) ([]LineJSON, error) { file, err := os.Open(path) if errors.Is(err, os.ErrNotExist) { return nil, nil @@ -86,7 +86,7 @@ func readJSONLines(path string, prefilter ...string) ([]lineJSON, error) { } defer file.Close() - lines := make([]lineJSON, 0) + lines := make([]LineJSON, 0) reader := bufio.NewReader(file) lineNumber := 0 offset := int64(0) @@ -102,11 +102,11 @@ func readJSONLines(path string, prefilter ...string) ([]lineJSON, error) { decoder := json.NewDecoder(bytes.NewReader(line)) decoder.UseNumber() if err := decoder.Decode(&value); err == nil { - lines = append(lines, lineJSON{ - value: value, - line: lineNumber, - start: start, - end: offset, + lines = append(lines, LineJSON{ + Value: value, + Line: lineNumber, + Start: start, + End: offset, }) } } @@ -122,11 +122,11 @@ func readJSONLines(path string, prefilter ...string) ([]lineJSON, error) { return lines, nil } -type lineJSON struct { - value map[string]any - line int - start int64 - end int64 +type LineJSON struct { + Value map[string]any + Line int + Start int64 + End int64 } func matchesPrefilter(line []byte, filters []string) bool { @@ -138,7 +138,7 @@ func matchesPrefilter(line []byte, filters []string) bool { return true } -func readJSONObject(path string) (map[string]any, error) { +func ReadJSONObject(path string) (map[string]any, error) { data, err := os.ReadFile(path) if errors.Is(err, os.ErrNotExist) { return nil, nil @@ -155,17 +155,17 @@ func readJSONObject(path string) (map[string]any, error) { return value, nil } -func objectAt(value any) map[string]any { +func ObjectAt(value any) map[string]any { object, _ := value.(map[string]any) return object } -func arrayAt(value any) []any { +func ArrayAt(value any) []any { array, _ := value.([]any) return array } -func stringValue(value any) string { +func StringValue(value any) string { switch typed := value.(type) { case string: return strings.TrimSpace(typed) @@ -176,23 +176,23 @@ func stringValue(value any) string { } } -func stringField(object map[string]any, key string) string { +func StringField(object map[string]any, key string) string { if object == nil { return "" } - return stringValue(object[key]) + return StringValue(object[key]) } -func firstStringField(object map[string]any, keys ...string) string { +func FirstStringField(object map[string]any, keys ...string) string { for _, key := range keys { - if value := stringField(object, key); value != "" { + if value := StringField(object, key); value != "" { return value } } return "" } -func firstNonEmpty(values ...string) string { +func FirstNonEmpty(values ...string) string { for _, value := range values { if strings.TrimSpace(value) != "" { return strings.TrimSpace(value) @@ -201,7 +201,7 @@ func firstNonEmpty(values ...string) string { return "" } -func uintValue(value any) uint64 { +func UintValue(value any) uint64 { switch typed := value.(type) { case json.Number: if parsed, err := strconv.ParseUint(typed.String(), 10, 64); err == nil { @@ -243,9 +243,9 @@ func floatToUint(value float64) uint64 { return uint64(math.Trunc(value)) } -func uintField(object map[string]any, keys ...string) uint64 { +func UintField(object map[string]any, keys ...string) uint64 { for _, key := range keys { - if value := uintValue(object[key]); value > 0 { + if value := UintValue(object[key]); value > 0 { return value } } @@ -267,24 +267,24 @@ func floatValue(value any) (float64, bool) { } } -func parseTimestamp(value any) (time.Time, bool) { +func ParseTimestamp(value any) (time.Time, bool) { switch typed := value.(type) { case string: - return parseTimestampString(typed) + return ParseTimestampString(typed) case json.Number: if integer, err := strconv.ParseUint(typed.String(), 10, 64); err == nil { return timestampFromScalar(integer) } if parsed, err := strconv.ParseFloat(typed.String(), 64); err == nil { - return timestampFromFloat(parsed) + return TimestampFromFloat(parsed) } case float64: - return timestampFromFloat(typed) + return TimestampFromFloat(typed) } return time.Time{}, false } -func parseTimestampString(raw string) (time.Time, bool) { +func ParseTimestampString(raw string) (time.Time, bool) { raw = strings.TrimSpace(raw) if raw == "" { return time.Time{}, false @@ -296,12 +296,12 @@ func parseTimestampString(raw string) (time.Time, bool) { return timestampFromScalar(parsed) } if parsed, err := strconv.ParseFloat(raw, 64); err == nil { - return timestampFromFloat(parsed) + return TimestampFromFloat(parsed) } return time.Time{}, false } -func timestampFromFloat(value float64) (time.Time, bool) { +func TimestampFromFloat(value float64) (time.Time, bool) { if !isFinite(value) || value <= 0 { return time.Time{}, false } @@ -336,13 +336,13 @@ func timestampFromScalar(raw uint64) (time.Time, bool) { return time.UnixMilli(int64(millis)), true } -func timestampFromParts(value any) (time.Time, bool) { - parts := arrayAt(value) +func TimestampFromParts(value any) (time.Time, bool) { + parts := ArrayAt(value) if len(parts) < 2 { return time.Time{}, false } - seconds := uintValue(parts[0]) - nanos := uintValue(parts[1]) + seconds := UintValue(parts[0]) + nanos := UintValue(parts[1]) if seconds == 0 { return time.Time{}, false } @@ -353,7 +353,7 @@ func timestampFromParts(value any) (time.Time, bool) { return time.UnixMilli(int64(millis)), true } -func fileModifiedTime(path string) time.Time { +func FileModifiedTime(path string) time.Time { info, err := os.Stat(path) if err != nil { return time.Unix(0, 0).UTC() @@ -369,7 +369,7 @@ func formatDate(timestamp time.Time) string { return timestamp.In(time.Local).Format("2006-01-02") } -func totalUsage(tokens usage.TokenUsage) uint64 { +func TotalUsage(tokens usage.TokenUsage) uint64 { return tokens.InputTokens + tokens.OutputTokens + tokens.CacheCreationInputTokens + @@ -378,8 +378,8 @@ func totalUsage(tokens usage.TokenUsage) uint64 { tokens.ReasoningOutputTokens } -func applyTotalFallback(tokens usage.TokenUsage, total uint64) usage.TokenUsage { - sum := totalUsage(tokens) +func ApplyTotalFallback(tokens usage.TokenUsage, total uint64) usage.TokenUsage { + sum := TotalUsage(tokens) if sum == 0 && total > 0 { tokens.OutputTokens = total tokens.TotalTokens = total @@ -396,11 +396,11 @@ func applyTotalFallback(tokens usage.TokenUsage, total uint64) usage.TokenUsage return tokens } -func nonZero(tokens usage.TokenUsage) bool { - return totalUsage(tokens) > 0 || tokens.TotalTokens > 0 +func NonZero(tokens usage.TokenUsage) bool { + return TotalUsage(tokens) > 0 || tokens.TotalTokens > 0 } -func baseEntry(provider usage.Provider, timestamp time.Time, project, projectPath, sessionID, model, client string, tokens usage.TokenUsage) usage.Entry { +func BaseEntry(provider usage.Provider, timestamp time.Time, project, projectPath, sessionID, model, client string, tokens usage.TokenUsage) usage.Entry { return usage.Entry{ Provider: provider, Timestamp: timestamp, @@ -416,14 +416,14 @@ func baseEntry(provider usage.Provider, timestamp time.Time, project, projectPat } } -func setSource(entry *usage.Entry, source string, line int, start, end int64) { +func SetSource(entry *usage.Entry, source string, line int, start, end int64) { entry.SourceFile = source entry.SourceLine = line entry.SourceStart = start entry.SourceEnd = end } -func stableEntryID(entry usage.Entry, extra ...string) string { +func StableEntryID(entry usage.Entry, extra ...string) string { parts := []string{ string(entry.Provider), entry.SourceFile, @@ -445,7 +445,7 @@ func stableEntryID(entry usage.Entry, extra ...string) string { return usage.StableID(parts...) } -func sortEntries(entries []usage.Entry) { +func SortEntries(entries []usage.Entry) { sort.Slice(entries, func(i, j int) bool { if !entries[i].Timestamp.Equal(entries[j].Timestamp) { return entries[i].Timestamp.Before(entries[j].Timestamp) @@ -454,7 +454,7 @@ func sortEntries(entries []usage.Entry) { }) } -func uniqueStrings(values []string) []string { +func UniqueStrings(values []string) []string { if len(values) == 0 { return nil } diff --git a/internal/agentusage/sqlite.go b/internal/provider/shared/sqlite.go similarity index 83% rename from internal/agentusage/sqlite.go rename to internal/provider/shared/sqlite.go index 96d1871..391aae4 100644 --- a/internal/agentusage/sqlite.go +++ b/internal/provider/shared/sqlite.go @@ -1,4 +1,4 @@ -package agentusage +package shared import ( "database/sql" @@ -10,7 +10,7 @@ import ( _ "modernc.org/sqlite" ) -func openSQLite(path string) (*sql.DB, error) { +func OpenSQLite(path string) (*sql.DB, error) { db, err := sql.Open("sqlite", path) if err != nil { return nil, err @@ -24,7 +24,7 @@ func openSQLite(path string) (*sql.DB, error) { return db, nil } -func sqlString(value any) string { +func SqlString(value any) string { switch typed := value.(type) { case nil: return "" @@ -44,7 +44,7 @@ func sqlString(value any) string { } } -func sqlUint(value any) uint64 { +func SqlUint(value any) uint64 { switch typed := value.(type) { case nil: return 0 @@ -59,14 +59,14 @@ func sqlUint(value any) uint64 { case float64: return floatToUint(typed) case []byte: - return uintValue(string(typed)) + return UintValue(string(typed)) case string: - return uintValue(typed) + return UintValue(typed) } return 0 } -func sqlFloat(value any) (float64, bool) { +func SqlFloat(value any) (float64, bool) { switch typed := value.(type) { case nil: return 0, false diff --git a/internal/provider/shared/sqlitefiles.go b/internal/provider/shared/sqlitefiles.go new file mode 100644 index 0000000..0344a4f --- /dev/null +++ b/internal/provider/shared/sqlitefiles.go @@ -0,0 +1,65 @@ +package shared + +import ( + "bytes" + "database/sql" + "encoding/json" + "os" + "path/filepath" + "sort" +) + +func ExistingSQLiteFile(path string) bool { + info, err := os.Stat(path) + return err == nil && !info.IsDir() +} +func DecodeJSONObjectString(data string) map[string]any { + decoder := json.NewDecoder(bytes.NewReader([]byte(data))) + decoder.UseNumber() + var record map[string]any + if err := decoder.Decode(&record); err != nil { + return nil + } + return record +} +func SqliteDBPaths(paths []string, defaultFile string, extraNames func(string) bool) []string { + dbPaths := make([]string, 0) + for _, root := range paths { + info, err := os.Stat(root) + if err != nil { + continue + } + if !info.IsDir() { + if filepath.Base(root) == defaultFile || extraNames != nil && extraNames(filepath.Base(root)) { + dbPaths = append(dbPaths, root) + } + continue + } + candidate := filepath.Join(root, defaultFile) + if fileInfo, err := os.Stat(candidate); err == nil && !fileInfo.IsDir() { + dbPaths = append(dbPaths, candidate) + } + if extraNames == nil { + continue + } + entries, err := os.ReadDir(root) + if err != nil { + continue + } + for _, entry := range entries { + if entry.IsDir() || !extraNames(entry.Name()) { + continue + } + dbPaths = append(dbPaths, filepath.Join(root, entry.Name())) + } + } + sort.Strings(dbPaths) + return UniqueStrings(dbPaths) +} +func ScanAny(rows *sql.Rows, values ...*any) bool { + dest := make([]any, len(values)) + for i := range values { + dest[i] = values[i] + } + return rows.Scan(dest...) == nil +} diff --git a/internal/provider/shared/testing.go b/internal/provider/shared/testing.go new file mode 100644 index 0000000..9e0541b --- /dev/null +++ b/internal/provider/shared/testing.go @@ -0,0 +1,84 @@ +package shared + +import ( + "database/sql" + "os" + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// WantEntry describes the single entry a provider must produce from a minimal +// fixture. Every provider asserts the same shape, so the check lives here +// rather than once per package. +type WantEntry struct { + Provider usage.Provider + Model string + SessionID string + Project string + Tokens usage.TokenUsage +} + +// AssertSingleEntry checks that a provider loaded exactly one entry matching +// want. It is the shared body of every provider's smoke test. +func AssertSingleEntry(t *testing.T, entries []usage.Entry, err error, want WantEntry) { + t.Helper() + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1: %#v", len(entries), entries) + } + entry := entries[0] + if entry.Provider != want.Provider { + t.Fatalf("provider = %q, want %q", entry.Provider, want.Provider) + } + if entry.Model != want.Model { + t.Fatalf("model = %q, want %q", entry.Model, want.Model) + } + if entry.SessionID != want.SessionID { + t.Fatalf("session id = %q, want %q", entry.SessionID, want.SessionID) + } + if entry.Project != want.Project { + t.Fatalf("project = %q, want %q", entry.Project, want.Project) + } + if entry.Usage != want.Tokens { + t.Fatalf("usage = %#v, want %#v", entry.Usage, want.Tokens) + } + if entry.ID == "" { + t.Fatal("ID is empty") + } +} + +// WriteFile writes a fixture file, creating its directory. +func WriteFile(t *testing.T, path, data string) { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(path, []byte(data), 0o600); err != nil { + t.Fatal(err) + } +} + +// OpenTestSQLite opens a fixture database, creating its directory. +func OpenTestSQLite(t *testing.T, path string) *sql.DB { + t.Helper() + if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { + t.Fatal(err) + } + db, err := OpenSQLite(path) + if err != nil { + t.Fatal(err) + } + return db +} + +// ExecSQL runs a fixture statement and fails the test if it errors. +func ExecSQL(t *testing.T, db *sql.DB, statement string, args ...any) { + t.Helper() + if _, err := db.Exec(statement, args...); err != nil { + t.Fatal(err) + } +} diff --git a/internal/usage/usage.go b/internal/usage/usage.go index 585788a..dc226d4 100644 --- a/internal/usage/usage.go +++ b/internal/usage/usage.go @@ -170,13 +170,13 @@ type Entry struct { // Client is the human-readable IDE or app source the request came from. // VS Code plugins are normalized across providers, but standalone apps // remain product-specific, e.g. "VS Code", "Codex Desktop", "Claude CLI". - Client string `json:"client,omitempty"` - Entity string `json:"entity,omitempty"` - EntityType string `json:"entity_type,omitempty"` - Branch string `json:"branch,omitempty"` - Editor string `json:"editor,omitempty"` - Category string `json:"category,omitempty"` - IsWrite *bool `json:"is_write,omitempty"` + Client string `json:"client,omitempty"` + Entity string `json:"entity,omitempty"` + EntityType string `json:"entity_type,omitempty"` + Branch string `json:"branch,omitempty"` + Editor string `json:"editor,omitempty"` + Category string `json:"category,omitempty"` + IsWrite *bool `json:"is_write,omitempty"` // LinesAdded/LinesRemoved count the source lines the agent added and // removed in this event's file modifications, for providers that record // diffs. Zero means "no diff recorded", not "no change". @@ -184,9 +184,9 @@ type Entry struct { LinesRemoved uint64 `json:"lines_removed,omitempty"` // Files breaks the same modifications down per file. Entity is always // the most-changed path in here; LinesAdded/LinesRemoved are the totals. - Files []FileChange `json:"files,omitempty"` - Raw map[string]any `json:"raw,omitempty"` - Usage TokenUsage `json:"usage"` + Files []FileChange `json:"files,omitempty"` + Raw map[string]any `json:"raw,omitempty"` + Usage TokenUsage `json:"usage"` } // NormalizeOS maps a Go runtime.GOOS value to a human-readable name. diff --git a/internal/usagedb/db.go b/internal/usagedb/db.go index a25400e..e107dfb 100644 --- a/internal/usagedb/db.go +++ b/internal/usagedb/db.go @@ -12,7 +12,7 @@ import ( "path/filepath" "time" - "github.com/tokitoki-dev/tokitoki-cli/internal/codexusage" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/codex" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" _ "modernc.org/sqlite" ) @@ -152,7 +152,7 @@ func rekeyCodexEvents(db *sql.DB) error { if entry.Provider != usage.ProviderCodex { continue } - newID := codexusage.StableEntryID(entry) + newID := codex.StableEntryID(entry) candidate := keeper{oldID: id, payload: payload, rank: statusRank(status)} current, exists := keepers[newID] switch { diff --git a/internal/usagedb/db_test.go b/internal/usagedb/db_test.go index e54d31c..9e0205a 100644 --- a/internal/usagedb/db_test.go +++ b/internal/usagedb/db_test.go @@ -5,7 +5,7 @@ import ( "testing" "time" - "github.com/tokitoki-dev/tokitoki-cli/internal/codexusage" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/codex" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -281,8 +281,8 @@ func TestOpenRekeysCodexEventsToSemanticIDs(t *testing.T) { if len(got) != 1 { t.Fatalf("rows = %d, want 1 (duplicates collapsed)", len(got)) } - if got[0].id != codexusage.StableEntryID(live) { - t.Fatalf("id = %q, want basename-keyed id %q", got[0].id, codexusage.StableEntryID(live)) + if got[0].id != codex.StableEntryID(live) { + t.Fatalf("id = %q, want basename-keyed id %q", got[0].id, codex.StableEntryID(live)) } if got[0].status != "uploaded" { t.Fatalf("status = %q, want uploaded (uploaded copy must win)", got[0].status) diff --git a/internal/usagescan/scanner.go b/internal/usagescan/scanner.go index 8f70bc5..9f9dd80 100644 --- a/internal/usagescan/scanner.go +++ b/internal/usagescan/scanner.go @@ -9,10 +9,22 @@ import ( "sort" "strings" - "github.com/tokitoki-dev/tokitoki-cli/internal/agentusage" - "github.com/tokitoki-dev/tokitoki-cli/internal/claudeusage" - "github.com/tokitoki-dev/tokitoki-cli/internal/codexusage" "github.com/tokitoki-dev/tokitoki-cli/internal/projectfile" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/amp" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/claude" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/codebuff" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/codex" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/copilot" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/droid" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/gemini" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/goose" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/hermes" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/kilo" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/kimi" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/openclaw" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/opencode" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/pi" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/qwen" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usagedb" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" @@ -42,21 +54,21 @@ type ProviderResult struct { // DefaultProviders returns the built-in usage providers. func DefaultProviders() []usageprovider.Provider { return []usageprovider.Provider{ - claudeusage.Provider{}, - codexusage.Provider{}, - agentusage.CopilotProvider{}, - agentusage.GeminiProvider{}, - agentusage.KimiProvider{}, - agentusage.QwenProvider{}, - agentusage.OpenClawProvider{}, - agentusage.PiProvider{}, - agentusage.AmpProvider{}, - agentusage.DroidProvider{}, - agentusage.KiloProvider{}, - agentusage.HermesProvider{}, - agentusage.CodebuffProvider{}, - agentusage.OpenCodeProvider{}, - agentusage.GooseProvider{}, + claude.Provider{}, + codex.Provider{}, + copilot.Provider{}, + gemini.Provider{}, + kimi.Provider{}, + qwen.Provider{}, + openclaw.Provider{}, + pi.Provider{}, + amp.Provider{}, + droid.Provider{}, + kilo.Provider{}, + hermes.Provider{}, + codebuff.Provider{}, + opencode.Provider{}, + goose.Provider{}, } } From 865d2085f5695dd1049cd53478b8b7a7a1b3d429 Mon Sep 17 00:00:00 2001 From: akarineren Date: Fri, 31 Jul 2026 20:49:04 +0900 Subject: [PATCH 05/11] Refactor: Remove shared package and migrate functionality to agentdata and agentdb - Deleted the shared package which contained various utility functions and types. - Moved relevant functions to the new agentdata package for handling JSON and file operations. - Introduced agentdb package for SQLite database interactions, including opening connections and querying. - Added providertest package for shared testing utilities across provider tests. - Updated usageprovider to utilize new base functionality for entry management. - Ensured all previous functionalities are preserved and properly integrated into the new structure. --- .../helpers.go => agentdata/agentdata.go} | 133 +++------------ internal/agentdb/agentdb.go | 157 ++++++++++++++++++ internal/provider/amp/loader.go | 90 +++++----- internal/provider/amp/provider.go | 7 +- internal/provider/amp/provider_test.go | 6 +- internal/provider/codebuff/loader.go | 68 ++++---- internal/provider/codebuff/provider.go | 7 +- internal/provider/codebuff/provider_test.go | 6 +- internal/provider/copilot/filter_test.go | 4 +- internal/provider/copilot/loader.go | 90 +++++----- internal/provider/copilot/provider.go | 7 +- internal/provider/copilot/provider_test.go | 6 +- internal/provider/droid/loader.go | 44 ++--- internal/provider/droid/provider.go | 7 +- internal/provider/droid/provider_test.go | 6 +- internal/provider/gemini/loader.go | 82 ++++----- internal/provider/gemini/provider.go | 7 +- internal/provider/gemini/provider_test.go | 6 +- internal/provider/goose/loader.go | 49 +++--- internal/provider/goose/provider.go | 7 +- internal/provider/goose/provider_test.go | 10 +- internal/provider/hermes/loader.go | 47 +++--- internal/provider/hermes/provider.go | 7 +- internal/provider/hermes/provider_test.go | 10 +- internal/provider/kilo/loader.go | 47 +++--- internal/provider/kilo/provider.go | 7 +- internal/provider/kilo/provider_test.go | 10 +- internal/provider/kimi/loader.go | 70 ++++---- internal/provider/kimi/provider.go | 7 +- internal/provider/kimi/provider_test.go | 8 +- internal/provider/openclaw/loader.go | 58 +++---- internal/provider/openclaw/provider.go | 7 +- internal/provider/openclaw/provider_test.go | 6 +- internal/provider/opencode/loader.go | 123 +++++++------- internal/provider/opencode/loader_test.go | 14 +- internal/provider/opencode/provider.go | 7 +- internal/provider/opencode/provider_test.go | 6 +- internal/provider/pi/loader.go | 42 ++--- internal/provider/pi/provider.go | 7 +- internal/provider/pi/provider_test.go | 6 +- internal/provider/qwen/loader.go | 44 ++--- internal/provider/qwen/provider.go | 7 +- internal/provider/qwen/provider_test.go | 6 +- internal/provider/shared/base.go | 37 ----- internal/provider/shared/entries.go | 16 -- internal/provider/shared/sqlite.go | 84 ---------- internal/provider/shared/sqlitefiles.go | 65 -------- .../providertest.go} | 8 +- internal/usageprovider/provider.go | 139 +++++++++++++++- 49 files changed, 857 insertions(+), 837 deletions(-) rename internal/{provider/shared/helpers.go => agentdata/agentdata.go} (69%) create mode 100644 internal/agentdb/agentdb.go delete mode 100644 internal/provider/shared/base.go delete mode 100644 internal/provider/shared/entries.go delete mode 100644 internal/provider/shared/sqlite.go delete mode 100644 internal/provider/shared/sqlitefiles.go rename internal/{provider/shared/testing.go => providertest/providertest.go} (87%) diff --git a/internal/provider/shared/helpers.go b/internal/agentdata/agentdata.go similarity index 69% rename from internal/provider/shared/helpers.go rename to internal/agentdata/agentdata.go index 405397b..04507e2 100644 --- a/internal/provider/shared/helpers.go +++ b/internal/agentdata/agentdata.go @@ -1,4 +1,6 @@ -package shared +// Package agentdata reads the JSON, JSONL and on-disk layouts that local AI +// agents write their logs in. +package agentdata import ( "bufio" @@ -9,7 +11,6 @@ import ( "math" "os" "path/filepath" - "runtime" "sort" "strconv" "strings" @@ -208,10 +209,10 @@ func UintValue(value any) uint64 { return parsed } if parsed, err := strconv.ParseFloat(typed.String(), 64); err == nil { - return floatToUint(parsed) + return FloatToUint(parsed) } case float64: - return floatToUint(typed) + return FloatToUint(typed) case int: if typed > 0 { return uint64(typed) @@ -227,14 +228,14 @@ func UintValue(value any) uint64 { return parsed } if parsed, err := strconv.ParseFloat(strings.TrimSpace(typed), 64); err == nil { - return floatToUint(parsed) + return FloatToUint(parsed) } } return 0 } -func floatToUint(value float64) uint64 { - if !isFinite(value) || value <= 0 { +func FloatToUint(value float64) uint64 { + if !IsFinite(value) || value <= 0 { return 0 } if value > float64(^uint64(0)) { @@ -252,21 +253,6 @@ func UintField(object map[string]any, keys ...string) uint64 { return 0 } -func floatValue(value any) (float64, bool) { - switch typed := value.(type) { - case json.Number: - parsed, err := strconv.ParseFloat(typed.String(), 64) - return parsed, err == nil && isFinite(parsed) - case float64: - return typed, isFinite(typed) - case string: - parsed, err := strconv.ParseFloat(strings.TrimSpace(typed), 64) - return parsed, err == nil && isFinite(parsed) - default: - return 0, false - } -} - func ParseTimestamp(value any) (time.Time, bool) { switch typed := value.(type) { case string: @@ -302,7 +288,7 @@ func ParseTimestampString(raw string) (time.Time, bool) { } func TimestampFromFloat(value float64) (time.Time, bool) { - if !isFinite(value) || value <= 0 { + if !IsFinite(value) || value <= 0 { return time.Time{}, false } if value < 100_000_000_000 { @@ -311,7 +297,7 @@ func TimestampFromFloat(value float64) (time.Time, bool) { return time.UnixMilli(int64(value)), true } -func isFinite(value float64) bool { +func IsFinite(value float64) bool { return !math.IsNaN(value) && !math.IsInf(value, 0) } @@ -365,95 +351,6 @@ func FileModifiedTime(path string) time.Time { return modified } -func formatDate(timestamp time.Time) string { - return timestamp.In(time.Local).Format("2006-01-02") -} - -func TotalUsage(tokens usage.TokenUsage) uint64 { - return tokens.InputTokens + - tokens.OutputTokens + - tokens.CacheCreationInputTokens + - tokens.CacheReadInputTokens + - tokens.CachedInputTokens + - tokens.ReasoningOutputTokens -} - -func ApplyTotalFallback(tokens usage.TokenUsage, total uint64) usage.TokenUsage { - sum := TotalUsage(tokens) - if sum == 0 && total > 0 { - tokens.OutputTokens = total - tokens.TotalTokens = total - return tokens - } - if total > sum { - tokens.ReasoningOutputTokens += total - sum - tokens.TotalTokens = total - return tokens - } - if tokens.TotalTokens == 0 { - tokens.TotalTokens = sum - } - return tokens -} - -func NonZero(tokens usage.TokenUsage) bool { - return TotalUsage(tokens) > 0 || tokens.TotalTokens > 0 -} - -func BaseEntry(provider usage.Provider, timestamp time.Time, project, projectPath, sessionID, model, client string, tokens usage.TokenUsage) usage.Entry { - return usage.Entry{ - Provider: provider, - Timestamp: timestamp, - Date: formatDate(timestamp), - Project: project, - ProjectPath: projectPath, - SessionID: sessionID, - Model: model, - Language: usage.UnknownLanguage, - OS: usage.NormalizeOS(runtime.GOOS), - Client: client, - Usage: tokens, - } -} - -func SetSource(entry *usage.Entry, source string, line int, start, end int64) { - entry.SourceFile = source - entry.SourceLine = line - entry.SourceStart = start - entry.SourceEnd = end -} - -func StableEntryID(entry usage.Entry, extra ...string) string { - parts := []string{ - string(entry.Provider), - entry.SourceFile, - strconv.Itoa(entry.SourceLine), - entry.Timestamp.Format(time.RFC3339Nano), - entry.Project, - entry.ProjectPath, - entry.SessionID, - entry.Model, - strconv.FormatUint(entry.Usage.InputTokens, 10), - strconv.FormatUint(entry.Usage.OutputTokens, 10), - strconv.FormatUint(entry.Usage.CacheCreationInputTokens, 10), - strconv.FormatUint(entry.Usage.CacheReadInputTokens, 10), - strconv.FormatUint(entry.Usage.CachedInputTokens, 10), - strconv.FormatUint(entry.Usage.ReasoningOutputTokens, 10), - strconv.FormatUint(entry.Usage.TotalTokens, 10), - } - parts = append(parts, extra...) - return usage.StableID(parts...) -} - -func SortEntries(entries []usage.Entry) { - sort.Slice(entries, func(i, j int) bool { - if !entries[i].Timestamp.Equal(entries[j].Timestamp) { - return entries[i].Timestamp.Before(entries[j].Timestamp) - } - return entries[i].ID < entries[j].ID - }) -} - func UniqueStrings(values []string) []string { if len(values) == 0 { return nil @@ -470,3 +367,13 @@ func UniqueStrings(values []string) []string { } return out } + +func DecodeJSONObjectString(data string) map[string]any { + decoder := json.NewDecoder(bytes.NewReader([]byte(data))) + decoder.UseNumber() + var record map[string]any + if err := decoder.Decode(&record); err != nil { + return nil + } + return record +} diff --git a/internal/agentdb/agentdb.go b/internal/agentdb/agentdb.go new file mode 100644 index 0000000..4bde412 --- /dev/null +++ b/internal/agentdb/agentdb.go @@ -0,0 +1,157 @@ +// Package agentdb reads the SQLite databases that local AI agents store their +// sessions in. +package agentdb + +import ( + "database/sql" + "encoding/json" + "fmt" + "math" + "os" + "path/filepath" + "sort" + "strconv" + "strings" + + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + + // agentdb is the only package that opens these databases, so the driver + // registers here. + _ "modernc.org/sqlite" +) + +func floatValue(value any) (float64, bool) { + switch typed := value.(type) { + case json.Number: + parsed, err := strconv.ParseFloat(typed.String(), 64) + return parsed, err == nil && agentdata.IsFinite(parsed) + case float64: + return typed, agentdata.IsFinite(typed) + case string: + parsed, err := strconv.ParseFloat(strings.TrimSpace(typed), 64) + return parsed, err == nil && agentdata.IsFinite(parsed) + default: + return 0, false + } +} + +func OpenSQLite(path string) (*sql.DB, error) { + db, err := sql.Open("sqlite", path) + if err != nil { + return nil, err + } + db.SetMaxOpenConns(1) + db.SetMaxIdleConns(1) + if err := db.Ping(); err != nil { + _ = db.Close() + return nil, err + } + return db, nil +} + +func SqlString(value any) string { + switch typed := value.(type) { + case nil: + return "" + case string: + return strings.TrimSpace(typed) + case []byte: + return strings.TrimSpace(string(typed)) + case int64: + return strconv.FormatInt(typed, 10) + case float64: + if !agentdata.IsFinite(typed) { + return "" + } + return strconv.FormatFloat(typed, 'f', -1, 64) + default: + return strings.TrimSpace(fmt.Sprint(typed)) + } +} + +func SqlUint(value any) uint64 { + switch typed := value.(type) { + case nil: + return 0 + case int64: + if typed > 0 { + return uint64(typed) + } + case int: + if typed > 0 { + return uint64(typed) + } + case float64: + return agentdata.FloatToUint(typed) + case []byte: + return agentdata.UintValue(string(typed)) + case string: + return agentdata.UintValue(typed) + } + return 0 +} + +func SqlFloat(value any) (float64, bool) { + switch typed := value.(type) { + case nil: + return 0, false + case float64: + return typed, !math.IsNaN(typed) && !math.IsInf(typed, 0) + case int64: + return float64(typed), true + case []byte: + return floatValue(string(typed)) + case string: + return floatValue(typed) + default: + return 0, false + } +} + +func ExistingSQLiteFile(path string) bool { + info, err := os.Stat(path) + return err == nil && !info.IsDir() +} + +func SqliteDBPaths(paths []string, defaultFile string, extraNames func(string) bool) []string { + dbPaths := make([]string, 0) + for _, root := range paths { + info, err := os.Stat(root) + if err != nil { + continue + } + if !info.IsDir() { + if filepath.Base(root) == defaultFile || extraNames != nil && extraNames(filepath.Base(root)) { + dbPaths = append(dbPaths, root) + } + continue + } + candidate := filepath.Join(root, defaultFile) + if fileInfo, err := os.Stat(candidate); err == nil && !fileInfo.IsDir() { + dbPaths = append(dbPaths, candidate) + } + if extraNames == nil { + continue + } + entries, err := os.ReadDir(root) + if err != nil { + continue + } + for _, entry := range entries { + if entry.IsDir() || !extraNames(entry.Name()) { + continue + } + dbPaths = append(dbPaths, filepath.Join(root, entry.Name())) + } + } + sort.Strings(dbPaths) + return agentdata.UniqueStrings(dbPaths) +} + +func ScanAny(rows *sql.Rows, values ...*any) bool { + dest := make([]any, len(values)) + for i := range values { + dest[i] = values[i] + } + return rows.Scan(dest...) == nil +} diff --git a/internal/provider/amp/loader.go b/internal/provider/amp/loader.go index 7005381..f36057b 100644 --- a/internal/provider/amp/loader.go +++ b/internal/provider/amp/loader.go @@ -7,7 +7,9 @@ import ( "sort" "strconv" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -20,13 +22,13 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) } continue } - files = append(files, shared.CollectExt(filepath.Join(path, "threads"), ".json")...) + files = append(files, agentdata.CollectExt(filepath.Join(path, "threads"), ".json")...) if filepath.Base(path) == "threads" { - files = append(files, shared.CollectExt(path, ".json")...) + files = append(files, agentdata.CollectExt(path, ".json")...) } } sort.Strings(files) - files = shared.FilterFiles(shared.UniqueStrings(files), filter) + files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) entries := make([]usage.Entry, 0) for _, file := range files { @@ -36,22 +38,22 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) } entries = append(entries, fileEntries...) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } func parseThreadFile(path string) ([]usage.Entry, error) { - thread, err := shared.ReadJSONObject(path) + thread, err := agentdata.ReadJSONObject(path) if err != nil || thread == nil { return nil, err } - threadID := shared.StringField(thread, "id") + threadID := agentdata.StringField(thread, "id") if threadID == "" { return nil, nil } - messages := shared.ArrayAt(thread["messages"]) - if ledger := shared.ObjectAt(thread["usageLedger"]); ledger != nil { - if events := shared.ArrayAt(ledger["events"]); len(events) > 0 { + messages := agentdata.ArrayAt(thread["messages"]) + if ledger := agentdata.ObjectAt(thread["usageLedger"]); ledger != nil { + if events := agentdata.ArrayAt(ledger["events"]); len(events) > 0 { return ledgerEntries(path, threadID, messages, events), nil } } @@ -62,37 +64,37 @@ func ledgerEntries(path, threadID string, messages []any, events []any) []usage. cacheTokens := cacheTokens(messages) entries := make([]usage.Entry, 0) for index, raw := range events { - event := shared.ObjectAt(raw) + event := agentdata.ObjectAt(raw) if event == nil { continue } - timestamp, ok := shared.ParseTimestamp(event["timestamp"]) + timestamp, ok := agentdata.ParseTimestamp(event["timestamp"]) if !ok { continue } - model := shared.StringField(event, "model") + model := agentdata.StringField(event, "model") if model == "" { continue } - tokenBlock := shared.ObjectAt(event["tokens"]) + tokenBlock := agentdata.ObjectAt(event["tokens"]) if tokenBlock == nil { continue } cache := cacheTokens[int64Value(event["toMessageId"])] tokens := usage.TokenUsage{ - InputTokens: shared.UintField(tokenBlock, "input"), - OutputTokens: shared.UintField(tokenBlock, "output"), + InputTokens: agentdata.UintField(tokenBlock, "input"), + OutputTokens: agentdata.UintField(tokenBlock, "output"), CacheCreationInputTokens: cache.cacheCreation, CacheReadInputTokens: cache.cacheRead, } - tokens = shared.ApplyTotalFallback(tokens, shared.UintField(tokenBlock, "total")) - if !shared.NonZero(tokens) { + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(tokenBlock, "total")) + if !usageprovider.NonZero(tokens) { continue } - messageID := shared.StringValue(event["id"]) - entry := shared.BaseEntry(usage.ProviderAmp, timestamp, "amp", "Amp", threadID, model, "Amp", tokens) - shared.SetSource(&entry, path, index+1, 0, 0) - entry.ID = shared.StableEntryID(entry, messageID) + messageID := agentdata.StringValue(event["id"]) + entry := usageprovider.BaseEntry(usage.ProviderAmp, timestamp, "amp", "Amp", threadID, model, "Amp", tokens) + usageprovider.SetSource(&entry, path, index+1, 0, 0) + entry.ID = usageprovider.StableEntryID(entry, messageID) entries = append(entries, entry) } return entries @@ -101,42 +103,42 @@ func ledgerEntries(path, threadID string, messages []any, events []any) []usage. func messageEntries(path, threadID string, messages []any) []usage.Entry { entries := make([]usage.Entry, 0) for index, raw := range messages { - message := shared.ObjectAt(raw) - if message == nil || shared.StringValue(message["role"]) != "assistant" { + message := agentdata.ObjectAt(raw) + if message == nil || agentdata.StringValue(message["role"]) != "assistant" { continue } - usageBlock := shared.ObjectAt(message["usage"]) + usageBlock := agentdata.ObjectAt(message["usage"]) if usageBlock == nil { continue } - timestamp, ok := shared.ParseTimestamp(usageBlock["timestamp"]) + timestamp, ok := agentdata.ParseTimestamp(usageBlock["timestamp"]) if !ok { - timestamp, ok = shared.ParseTimestamp(message["timestamp"]) + timestamp, ok = agentdata.ParseTimestamp(message["timestamp"]) } if !ok { continue } - model := shared.StringField(usageBlock, "model") + model := agentdata.StringField(usageBlock, "model") if model == "" { - model = shared.StringValue(message["model"]) + model = agentdata.StringValue(message["model"]) } if model == "" { continue } tokens := usage.TokenUsage{ - InputTokens: shared.UintField(usageBlock, "inputTokens"), - OutputTokens: shared.UintField(usageBlock, "outputTokens"), - CacheCreationInputTokens: shared.UintField(usageBlock, "cacheCreationInputTokens"), - CacheReadInputTokens: shared.UintField(usageBlock, "cacheReadInputTokens"), + InputTokens: agentdata.UintField(usageBlock, "inputTokens"), + OutputTokens: agentdata.UintField(usageBlock, "outputTokens"), + CacheCreationInputTokens: agentdata.UintField(usageBlock, "cacheCreationInputTokens"), + CacheReadInputTokens: agentdata.UintField(usageBlock, "cacheReadInputTokens"), } - tokens = shared.ApplyTotalFallback(tokens, shared.UintField(usageBlock, "totalTokens")) - if !shared.NonZero(tokens) { + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(usageBlock, "totalTokens")) + if !usageprovider.NonZero(tokens) { continue } - messageID := shared.StringValue(message["messageId"]) - entry := shared.BaseEntry(usage.ProviderAmp, timestamp, "amp", "Amp", threadID, model, "Amp", tokens) - shared.SetSource(&entry, path, index+1, 0, 0) - entry.ID = shared.StableEntryID(entry, messageID) + messageID := agentdata.StringValue(message["messageId"]) + entry := usageprovider.BaseEntry(usage.ProviderAmp, timestamp, "amp", "Amp", threadID, model, "Amp", tokens) + usageprovider.SetSource(&entry, path, index+1, 0, 0) + entry.ID = usageprovider.StableEntryID(entry, messageID) entries = append(entries, entry) } return entries @@ -150,18 +152,18 @@ type cache struct { func cacheTokens(messages []any) map[int64]cache { tokens := make(map[int64]cache) for _, raw := range messages { - message := shared.ObjectAt(raw) - if message == nil || shared.StringValue(message["role"]) != "assistant" { + message := agentdata.ObjectAt(raw) + if message == nil || agentdata.StringValue(message["role"]) != "assistant" { continue } id := int64Value(message["messageId"]) if id == 0 { continue } - usageBlock := shared.ObjectAt(message["usage"]) + usageBlock := agentdata.ObjectAt(message["usage"]) tokens[id] = cache{ - cacheCreation: shared.UintField(usageBlock, "cacheCreationInputTokens"), - cacheRead: shared.UintField(usageBlock, "cacheReadInputTokens"), + cacheCreation: agentdata.UintField(usageBlock, "cacheCreationInputTokens"), + cacheRead: agentdata.UintField(usageBlock, "cacheReadInputTokens"), } } return tokens diff --git a/internal/provider/amp/provider.go b/internal/provider/amp/provider.go index 51f2ecf..79772e5 100644 --- a/internal/provider/amp/provider.go +++ b/internal/provider/amp/provider.go @@ -1,19 +1,18 @@ package amp import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads Amp usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a Amp provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/amp/provider_test.go b/internal/provider/amp/provider_test.go index 236f603..e16318e 100644 --- a/internal/provider/amp/provider_test.go +++ b/internal/provider/amp/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,11 +14,11 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() path := filepath.Join(dir, "threads", "thread.json") - shared.WriteFile(t, path, `{"id":"thread-a","usageLedger":{"events":[{"id":"event-a","timestamp":"2026-01-02T00:00:00.000Z","model":"gpt-5","tokens":{"input":1,"output":2}}]}}`) + providertest.WriteFile(t, path, `{"id":"thread-a","usageLedger":{"events":[{"id":"event-a","timestamp":"2026-01-02T00:00:00.000Z","model":"gpt-5","tokens":{"input":1,"output":2}}]}}`) return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderAmp, Model: "gpt-5", SessionID: "thread-a", diff --git a/internal/provider/codebuff/loader.go b/internal/provider/codebuff/loader.go index 7f067bb..a88c36e 100644 --- a/internal/provider/codebuff/loader.go +++ b/internal/provider/codebuff/loader.go @@ -10,7 +10,9 @@ import ( "strings" "time" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -29,7 +31,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) files = append(files, collectChatFiles(root)...) } sort.Strings(files) - files = shared.FilterFiles(shared.UniqueStrings(files), filter) + files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) entriesByID := make(map[string]usage.Entry) for _, file := range files { @@ -45,7 +47,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) for _, entry := range entriesByID { entries = append(entries, entry) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } @@ -64,7 +66,7 @@ func collectChatFiles(root string) []string { if filepath.Base(root) != "projects" { projectRoot = filepath.Join(root, "projects") } - return shared.CollectFiles(projectRoot, func(path string) bool { + return agentdata.CollectFiles(projectRoot, func(path string) bool { return filepath.Base(path) == "chat-messages.json" }) } @@ -82,10 +84,10 @@ func parseChatFile(path string) ([]usage.Entry, error) { } sessionID, chatID := sessionContext(path) chatTimestamp, hasChatTimestamp := parseChatTimestamp(chatID) - fileTimestamp := shared.FileModifiedTime(path) + fileTimestamp := agentdata.FileModifiedTime(path) entries := make([]usage.Entry, 0) for index, raw := range messages { - message := shared.ObjectAt(raw) + message := agentdata.ObjectAt(raw) if !isAssistant(message) { continue } @@ -98,9 +100,9 @@ func parseChatFile(path string) ([]usage.Entry, error) { ReasoningOutputTokens: parsedUsage.extraTotalTokens, } if tokens.TotalTokens == 0 { - tokens.TotalTokens = shared.TotalUsage(tokens) + tokens.TotalTokens = usageprovider.TotalUsage(tokens) } - if !shared.NonZero(tokens) { + if !usageprovider.NonZero(tokens) { continue } model := parsedUsage.model @@ -115,9 +117,9 @@ func parseChatFile(path string) ([]usage.Entry, error) { if !ok { timestamp = fileTimestamp } - entry := shared.BaseEntry(usage.ProviderCodebuff, timestamp, "codebuff", "Codebuff", sessionID, model, "Codebuff", tokens) - shared.SetSource(&entry, path, index+1, 0, 0) - entry.ID = shared.StableEntryID(entry, dedupKey(message, sessionID, timestamp, model, tokens, index)) + entry := usageprovider.BaseEntry(usage.ProviderCodebuff, timestamp, "codebuff", "Codebuff", sessionID, model, "Codebuff", tokens) + usageprovider.SetSource(&entry, path, index+1, 0, 0) + entry.ID = usageprovider.StableEntryID(entry, dedupKey(message, sessionID, timestamp, model, tokens, index)) entries = append(entries, entry) } return entries, nil @@ -142,17 +144,17 @@ func sessionContext(path string) (string, string) { } func isAssistant(message map[string]any) bool { - role := shared.FirstStringField(message, "variant", "role") + role := agentdata.FirstStringField(message, "variant", "role") return role == "ai" || role == "agent" || role == "assistant" } func extractUsage(message map[string]any) tokenUsage { var usage tokenUsage - metadata := shared.ObjectAt(message["metadata"]) + metadata := agentdata.ObjectAt(message["metadata"]) if metadata != nil { - usage.model = shared.StringField(metadata, "model") + usage.model = agentdata.StringField(metadata, "model") mergeCodebuffUsage(&usage, parseUsageObject(metadata["usage"])) - mergeCodebuffUsage(&usage, parseUsageObject(shared.ObjectAt(metadata["codebuff"])["usage"])) + mergeCodebuffUsage(&usage, parseUsageObject(agentdata.ObjectAt(metadata["codebuff"])["usage"])) if runState := runStateUsage(metadata); runState != nil { mergeCodebuffUsage(&usage, *runState) } @@ -161,27 +163,27 @@ func extractUsage(message map[string]any) tokenUsage { } func runStateUsage(metadata map[string]any) *tokenUsage { - history := shared.ArrayAt(shared.ObjectAt(shared.ObjectAt(shared.ObjectAt(metadata["runState"])["sessionState"])["mainAgentState"])["messageHistory"]) + history := agentdata.ArrayAt(agentdata.ObjectAt(agentdata.ObjectAt(agentdata.ObjectAt(metadata["runState"])["sessionState"])["mainAgentState"])["messageHistory"]) if len(history) == 0 { return nil } var usage tokenUsage found := false for i := len(history) - 1; i >= 0; i-- { - entry := shared.ObjectAt(history[i]) - if shared.StringField(entry, "role") != "assistant" { + entry := agentdata.ObjectAt(history[i]) + if agentdata.StringField(entry, "role") != "assistant" { continue } - providerOptions := shared.ObjectAt(entry["providerOptions"]) + providerOptions := agentdata.ObjectAt(entry["providerOptions"]) if providerOptions == nil { continue } entryUsage := parseUsageObject(providerOptions["usage"]) - codebuff := shared.ObjectAt(providerOptions["codebuff"]) + codebuff := agentdata.ObjectAt(providerOptions["codebuff"]) if codebuff != nil { mergeCodebuffUsage(&entryUsage, parseUsageObject(codebuff["usage"])) if entryUsage.model == "" { - entryUsage.model = shared.StringField(codebuff, "model") + entryUsage.model = agentdata.StringField(codebuff, "model") } } if usageHasTokens(entryUsage) || entryUsage.model != "" { @@ -196,26 +198,26 @@ func runStateUsage(metadata map[string]any) *tokenUsage { } func parseUsageObject(value any) tokenUsage { - record := shared.ObjectAt(value) + record := agentdata.ObjectAt(value) if record == nil { return tokenUsage{} } parsed := tokenUsage{ - model: shared.StringField(record, "model"), + model: agentdata.StringField(record, "model"), inputTokens: firstUint(record, "inputTokens", "input_tokens", "promptTokens", "prompt_tokens"), outputTokens: firstUint(record, "outputTokens", "output_tokens", "completionTokens", "completion_tokens"), cacheReadInputTokens: firstUint(record, "cacheReadInputTokens", "cache_read_input_tokens"), cacheCreationInputTokens: firstUint(record, "cacheCreationInputTokens", "cache_creation_input_tokens", "cacheCreationTokens", "cache_creation_tokens", "cachedTokensCreated", "cached_tokens_created"), } - parsed.cacheReadInputTokens = maxUint64(parsed.cacheReadInputTokens, firstUint(shared.ObjectAt(record["promptTokensDetails"]), "cachedTokens")) - parsed.cacheReadInputTokens = maxUint64(parsed.cacheReadInputTokens, firstUint(shared.ObjectAt(record["prompt_tokens_details"]), "cached_tokens")) + parsed.cacheReadInputTokens = maxUint64(parsed.cacheReadInputTokens, firstUint(agentdata.ObjectAt(record["promptTokensDetails"]), "cachedTokens")) + parsed.cacheReadInputTokens = maxUint64(parsed.cacheReadInputTokens, firstUint(agentdata.ObjectAt(record["prompt_tokens_details"]), "cached_tokens")) tokens := usage.TokenUsage{ InputTokens: parsed.inputTokens, OutputTokens: parsed.outputTokens, CacheCreationInputTokens: parsed.cacheCreationInputTokens, CacheReadInputTokens: parsed.cacheReadInputTokens, } - tokens = shared.ApplyTotalFallback(tokens, firstUint(record, "totalTokens", "total_tokens", "total")) + tokens = usageprovider.ApplyTotalFallback(tokens, firstUint(record, "totalTokens", "total_tokens", "total")) parsed.inputTokens = tokens.InputTokens parsed.outputTokens = tokens.OutputTokens parsed.cacheCreationInputTokens = tokens.CacheCreationInputTokens @@ -250,13 +252,13 @@ func usageHasTokens(value tokenUsage) bool { } func messageTimestamp(message map[string]any) (time.Time, bool) { - if timestamp, ok := shared.ParseTimestamp(message["timestamp"]); ok { + if timestamp, ok := agentdata.ParseTimestamp(message["timestamp"]); ok { return timestamp, true } - if timestamp, ok := shared.ParseTimestamp(message["createdAt"]); ok { + if timestamp, ok := agentdata.ParseTimestamp(message["createdAt"]); ok { return timestamp, true } - return shared.ParseTimestamp(shared.ObjectAt(message["metadata"])["timestamp"]) + return agentdata.ParseTimestamp(agentdata.ObjectAt(message["metadata"])["timestamp"]) } func parseChatTimestamp(chatID string) (time.Time, bool) { @@ -269,21 +271,21 @@ func parseChatTimestamp(chatID string) (time.Time, bool) { clock = clock[:index] + ":" + clock[index+1:] } } - return shared.ParseTimestampString(date + "T" + clock) + return agentdata.ParseTimestampString(date + "T" + clock) } func dedupKey(message map[string]any, sessionID string, timestamp time.Time, model string, tokens usage.TokenUsage, index int) string { - if id := shared.StringField(message, "id"); id != "" { + if id := agentdata.StringField(message, "id"); id != "" { return "codebuff:" + sessionID + ":" + id } - return shared.StableEntryID(shared.BaseEntry(usage.ProviderCodebuff, timestamp, "codebuff", "Codebuff", sessionID, model, "Codebuff", tokens), strconv.Itoa(index)) + return usageprovider.StableEntryID(usageprovider.BaseEntry(usage.ProviderCodebuff, timestamp, "codebuff", "Codebuff", sessionID, model, "Codebuff", tokens), strconv.Itoa(index)) } func firstUint(record map[string]any, keys ...string) uint64 { if record == nil { return 0 } - return shared.UintField(record, keys...) + return agentdata.UintField(record, keys...) } func maxUint64(a, b uint64) uint64 { diff --git a/internal/provider/codebuff/provider.go b/internal/provider/codebuff/provider.go index 4da56db..433054f 100644 --- a/internal/provider/codebuff/provider.go +++ b/internal/provider/codebuff/provider.go @@ -1,19 +1,18 @@ package codebuff import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads Codebuff usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a Codebuff provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/codebuff/provider_test.go b/internal/provider/codebuff/provider_test.go index 2a6a80a..0d5a516 100644 --- a/internal/provider/codebuff/provider_test.go +++ b/internal/provider/codebuff/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,11 +14,11 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { root := filepath.Join(t.TempDir(), "manicode") path := filepath.Join(root, "projects", "project-a", "chats", "2026-01-02T03-04-05.000Z", "chat-messages.json") - shared.WriteFile(t, path, `[{"id":"assistant-message","role":"assistant","timestamp":"2026-01-02T03:04:06.000Z","metadata":{"model":"claude-sonnet-4-20250514","usage":{"inputTokens":100,"outputTokens":50,"cacheCreationInputTokens":20,"cacheReadInputTokens":10}}}]`) + providertest.WriteFile(t, path, `[{"id":"assistant-message","role":"assistant","timestamp":"2026-01-02T03:04:06.000Z","metadata":{"model":"claude-sonnet-4-20250514","usage":{"inputTokens":100,"outputTokens":50,"cacheCreationInputTokens":20,"cacheReadInputTokens":10}}}]`) return Provider{}.WithPaths([]string{root}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderCodebuff, Model: "claude-sonnet-4-20250514", SessionID: "manicode/project-a/2026-01-02T03-04-05.000Z", diff --git a/internal/provider/copilot/filter_test.go b/internal/provider/copilot/filter_test.go index 1ca258c..ce0b240 100644 --- a/internal/provider/copilot/filter_test.go +++ b/internal/provider/copilot/filter_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" ) // TestFileFilterSkipsRejectedFiles proves the scanner can skip source files @@ -12,7 +12,7 @@ import ( func TestFileFilterSkipsRejectedFiles(t *testing.T) { dir := t.TempDir() path := filepath.Join(dir, "copilot.jsonl") - shared.WriteFile(t, path, `{"type":"span","traceId":"trace-1","spanId":"span-1","name":"chat claude-sonnet-4","endTime":[1775934264,967317833],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"claude-sonnet-4","gen_ai.conversation.id":"conv-1","gen_ai.usage.input_tokens":19452,"gen_ai.usage.output_tokens":281}}}`+"\n") + providertest.WriteFile(t, path, `{"type":"span","traceId":"trace-1","spanId":"span-1","name":"chat claude-sonnet-4","endTime":[1775934264,967317833],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"claude-sonnet-4","gen_ai.conversation.id":"conv-1","gen_ai.usage.input_tokens":19452,"gen_ai.usage.output_tokens":281}}}`+"\n") rejected := make([]string, 0) provider := Provider{}.WithPaths([]string{dir}).(Provider). diff --git a/internal/provider/copilot/loader.go b/internal/provider/copilot/loader.go index de8f164..42c0b84 100644 --- a/internal/provider/copilot/loader.go +++ b/internal/provider/copilot/loader.go @@ -7,17 +7,19 @@ import ( "strings" "time" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, root := range paths { - files = append(files, shared.CollectExt(root, ".jsonl")...) + files = append(files, agentdata.CollectExt(root, ".jsonl")...) } sort.Strings(files) - files = shared.FilterFiles(shared.UniqueStrings(files), filter) + files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) entries := make([]usage.Entry, 0) for _, file := range files { @@ -27,7 +29,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) } entries = append(entries, fileEntries...) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } @@ -61,12 +63,12 @@ type traceContext struct { } func parseOTELFile(path string) ([]usage.Entry, error) { - lines, err := shared.ReadJSONLines(path, `"attributes"`) + lines, err := agentdata.ReadJSONLines(path, `"attributes"`) if err != nil { return nil, err } contexts := traceContexts(lines) - fallback := shared.FileModifiedTime(path) + fallback := agentdata.FileModifiedTime(path) candidates := make([]candidate, 0) for index, line := range lines { if candidate, ok := recordCandidate(path, line, index, fallback, contexts); ok { @@ -79,7 +81,7 @@ func parseOTELFile(path string) ([]usage.Entry, error) { if !shouldEmitCopilot(candidate, sets) { continue } - entry := shared.BaseEntry( + entry := usageprovider.BaseEntry( usage.ProviderCopilot, candidate.timestamp, "copilot", @@ -89,14 +91,14 @@ func parseOTELFile(path string) ([]usage.Entry, error) { "GitHub Copilot CLI", candidate.tokens, ) - shared.SetSource(&entry, candidate.sourceFile, candidate.sourceLine, candidate.sourceStart, candidate.sourceEnd) - entry.ID = shared.StableEntryID(entry, candidate.dedupKey) + usageprovider.SetSource(&entry, candidate.sourceFile, candidate.sourceLine, candidate.sourceStart, candidate.sourceEnd) + entry.ID = usageprovider.StableEntryID(entry, candidate.dedupKey) entries = append(entries, entry) } return entries, nil } -func traceContexts(lines []shared.LineJSON) map[string]traceContext { +func traceContexts(lines []agentdata.LineJSON) map[string]traceContext { contexts := make(map[string]traceContext) for _, line := range lines { record := line.Value @@ -104,13 +106,13 @@ func traceContexts(lines []shared.LineJSON) map[string]traceContext { if traceID == "" { continue } - attrs := shared.ObjectAt(record["attributes"]) + attrs := agentdata.ObjectAt(record["attributes"]) if attrs == nil { continue } context := contexts[traceID] if context.model == "" { - context.model = shared.FirstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") + context.model = agentdata.FirstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") } if sessionID, priority := bestSession(attrs); sessionID != "" && priority > context.sessionIDPriority { context.sessionID = sessionID @@ -121,9 +123,9 @@ func traceContexts(lines []shared.LineJSON) map[string]traceContext { return contexts } -func recordCandidate(path string, line shared.LineJSON, index int, fallback time.Time, contexts map[string]traceContext) (candidate, bool) { +func recordCandidate(path string, line agentdata.LineJSON, index int, fallback time.Time, contexts map[string]traceContext) (candidate, bool) { record := line.Value - attrs := shared.ObjectAt(record["attributes"]) + attrs := agentdata.ObjectAt(record["attributes"]) if attrs == nil { return candidate{}, false } @@ -131,8 +133,8 @@ func recordCandidate(path string, line shared.LineJSON, index int, fallback time if !ok { return candidate{}, false } - input := shared.UintField(attrs, "gen_ai.usage.input_tokens") - cacheRead := shared.UintField(attrs, "gen_ai.usage.cache_read.input_tokens") + input := agentdata.UintField(attrs, "gen_ai.usage.input_tokens") + cacheRead := agentdata.UintField(attrs, "gen_ai.usage.cache_read.input_tokens") if cacheRead <= input { input -= cacheRead } else { @@ -140,18 +142,18 @@ func recordCandidate(path string, line shared.LineJSON, index int, fallback time } tokens := usage.TokenUsage{ InputTokens: input, - OutputTokens: shared.UintField(attrs, "gen_ai.usage.output_tokens"), - CacheCreationInputTokens: shared.UintField(attrs, "gen_ai.usage.cache_write.input_tokens", "gen_ai.usage.cache_creation.input_tokens"), + OutputTokens: agentdata.UintField(attrs, "gen_ai.usage.output_tokens"), + CacheCreationInputTokens: agentdata.UintField(attrs, "gen_ai.usage.cache_write.input_tokens", "gen_ai.usage.cache_creation.input_tokens"), CacheReadInputTokens: cacheRead, - ReasoningOutputTokens: shared.UintField(attrs, "gen_ai.usage.reasoning.output_tokens", "gen_ai.usage.reasoning_tokens"), + ReasoningOutputTokens: agentdata.UintField(attrs, "gen_ai.usage.reasoning.output_tokens", "gen_ai.usage.reasoning_tokens"), } - tokens = shared.ApplyTotalFallback(tokens, shared.UintField(attrs, "gen_ai.usage.total_tokens", "gen_ai.usage.total.token_count")) - if !shared.NonZero(tokens) { + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(attrs, "gen_ai.usage.total_tokens", "gen_ai.usage.total.token_count")) + if !usageprovider.NonZero(tokens) { return candidate{}, false } traceID := traceID(record) context := contexts[traceID] - model := shared.FirstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") + model := agentdata.FirstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") if model == "" { model = context.model } @@ -172,7 +174,7 @@ func recordCandidate(path string, line shared.LineJSON, index int, fallback time if !ok { timestamp = fallback } - responseID := shared.StringField(attrs, "gen_ai.response.id") + responseID := agentdata.StringField(attrs, "gen_ai.response.id") return candidate{ source: source, traceID: traceID, @@ -205,14 +207,14 @@ func recordSource(record, attrs map[string]any) (sourceKind, bool) { } func isSpan(record map[string]any) bool { - if shared.StringField(record, "type") == "span" { + if agentdata.StringField(record, "type") == "span" { return true } - if shared.StringField(record, "name") == "" { + if agentdata.StringField(record, "name") == "" { return false } - return shared.StringField(record, "spanId") != "" || - shared.StringField(record, "traceId") != "" || + return agentdata.StringField(record, "spanId") != "" || + agentdata.StringField(record, "traceId") != "" || record["startTime"] != nil || record["endTime"] != nil || record["duration"] != nil || @@ -221,47 +223,47 @@ func isSpan(record map[string]any) bool { func isChatSpan(record, attrs map[string]any) bool { return isSpan(record) && - (shared.StringField(attrs, "gen_ai.operation.name") == "chat" || - strings.HasPrefix(shared.StringField(record, "name"), "chat ")) + (agentdata.StringField(attrs, "gen_ai.operation.name") == "chat" || + strings.HasPrefix(agentdata.StringField(record, "name"), "chat ")) } func isAgentSummarySpan(record, attrs map[string]any) bool { return isSpan(record) && - (shared.StringField(attrs, "gen_ai.operation.name") == "invoke_agent" || - strings.HasPrefix(shared.StringField(record, "name"), "invoke_agent ")) + (agentdata.StringField(attrs, "gen_ai.operation.name") == "invoke_agent" || + strings.HasPrefix(agentdata.StringField(record, "name"), "invoke_agent ")) } func isInferenceLog(record, attrs map[string]any) bool { return !isSpan(record) && - (shared.StringField(attrs, "event.name") == "gen_ai.client.inference.operation.details" || + (agentdata.StringField(attrs, "event.name") == "gen_ai.client.inference.operation.details" || strings.HasPrefix(spanBody(record), "GenAI inference:")) } func isAgentTurnLog(record, attrs map[string]any) bool { return !isSpan(record) && - (shared.StringField(attrs, "event.name") == "copilot_chat.agent.turn" || + (agentdata.StringField(attrs, "event.name") == "copilot_chat.agent.turn" || strings.HasPrefix(spanBody(record), "copilot_chat.agent.turn")) } func spanBody(record map[string]any) string { - if body := shared.StringField(record, "body"); body != "" { + if body := agentdata.StringField(record, "body"); body != "" { return body } - return shared.StringField(record, "_body") + return agentdata.StringField(record, "_body") } func traceID(record map[string]any) string { - if traceID := shared.StringField(record, "traceId"); traceID != "" { + if traceID := agentdata.StringField(record, "traceId"); traceID != "" { return traceID } - return shared.StringField(shared.ObjectAt(record["spanContext"]), "traceId") + return agentdata.StringField(agentdata.ObjectAt(record["spanContext"]), "traceId") } func spanID(record map[string]any) string { - if spanID := shared.StringField(record, "spanId"); spanID != "" { + if spanID := agentdata.StringField(record, "spanId"); spanID != "" { return spanID } - return shared.StringField(shared.ObjectAt(record["spanContext"]), "spanId") + return agentdata.StringField(agentdata.ObjectAt(record["spanContext"]), "spanId") } func bestSession(attrs map[string]any) (string, int) { @@ -279,7 +281,7 @@ func bestSession(attrs map[string]any) (string, int) { bestValue := "" bestPriority := 0 for _, candidate := range candidates { - if value := shared.StringField(attrs, candidate.key); value != "" && candidate.priority > bestPriority { + if value := agentdata.StringField(attrs, candidate.key); value != "" && candidate.priority > bestPriority { bestValue = value bestPriority = candidate.priority } @@ -289,16 +291,16 @@ func bestSession(attrs map[string]any) (string, int) { func timestamp(record map[string]any) (time.Time, bool) { for _, key := range []string{"endTime", "startTime", "hrTime", "_hrTime", "time"} { - if timestamp, ok := shared.TimestampFromParts(record[key]); ok { + if timestamp, ok := agentdata.TimestampFromParts(record[key]); ok { return timestamp, true } } for _, key := range []string{"timestamp", "observedTimestamp"} { - if timestamp, ok := shared.ParseTimestamp(record[key]); ok { + if timestamp, ok := agentdata.ParseTimestamp(record[key]); ok { return timestamp, true } } - if raw := shared.UintValue(record["timeUnixNano"]); raw > 0 { + if raw := agentdata.UintValue(record["timeUnixNano"]); raw > 0 { return time.UnixMilli(int64(raw / 1_000_000)), true } return time.Time{}, false @@ -318,7 +320,7 @@ func dedupKey(source sourceKind, record, attrs map[string]any, traceID, sessionI } return fmt.Sprintf("log:%s:%d:%d", sessionID, timestamp.UnixMilli(), index) case copilotAgentTurnLog: - turnIndex := shared.UintField(attrs, "turn.index", "copilot_chat.turn.index") + turnIndex := agentdata.UintField(attrs, "turn.index", "copilot_chat.turn.index") turn := fmt.Sprintf("idx-%d", index) if turnIndex > 0 { turn = fmt.Sprintf("%d", turnIndex) diff --git a/internal/provider/copilot/provider.go b/internal/provider/copilot/provider.go index 8030c54..8b1d6bf 100644 --- a/internal/provider/copilot/provider.go +++ b/internal/provider/copilot/provider.go @@ -1,19 +1,18 @@ package copilot import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads GitHub Copilot CLI usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a GitHub Copilot CLI provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/copilot/provider_test.go b/internal/provider/copilot/provider_test.go index 8e10127..71eacc3 100644 --- a/internal/provider/copilot/provider_test.go +++ b/internal/provider/copilot/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,11 +14,11 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() path := filepath.Join(dir, "copilot.jsonl") - shared.WriteFile(t, path, `{"type":"span","traceId":"trace-1","spanId":"span-1","name":"chat claude-sonnet-4","endTime":[1775934264,967317833],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"claude-sonnet-4","gen_ai.conversation.id":"conv-1","gen_ai.usage.input_tokens":19452,"gen_ai.usage.output_tokens":281,"gen_ai.usage.cache_read.input_tokens":123,"gen_ai.usage.cache_creation.input_tokens":25,"gen_ai.usage.reasoning.output_tokens":128}}}`+"\n") + providertest.WriteFile(t, path, `{"type":"span","traceId":"trace-1","spanId":"span-1","name":"chat claude-sonnet-4","endTime":[1775934264,967317833],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"claude-sonnet-4","gen_ai.conversation.id":"conv-1","gen_ai.usage.input_tokens":19452,"gen_ai.usage.output_tokens":281,"gen_ai.usage.cache_read.input_tokens":123,"gen_ai.usage.cache_creation.input_tokens":25,"gen_ai.usage.reasoning.output_tokens":128}}}`+"\n") return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderCopilot, Model: "claude-sonnet-4", SessionID: "conv-1", diff --git a/internal/provider/droid/loader.go b/internal/provider/droid/loader.go index 63304b0..99d8a57 100644 --- a/internal/provider/droid/loader.go +++ b/internal/provider/droid/loader.go @@ -7,17 +7,19 @@ import ( "sort" "strings" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, root := range paths { - files = append(files, shared.CollectFiles(root, isSettingsFile)...) + files = append(files, agentdata.CollectFiles(root, isSettingsFile)...) } sort.Strings(files) - files = shared.FilterFiles(shared.UniqueStrings(files), filter) + files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) entries := make([]usage.Entry, 0) for _, file := range files { @@ -29,7 +31,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) entries = append(entries, entry) } } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } @@ -38,28 +40,28 @@ func isSettingsFile(path string) bool { } func parseSettingsFile(path string) (usage.Entry, bool, error) { - settings, err := shared.ReadJSONObject(path) + settings, err := agentdata.ReadJSONObject(path) if err != nil || settings == nil { return usage.Entry{}, false, err } - usageBlock := shared.ObjectAt(settings["tokenUsage"]) + usageBlock := agentdata.ObjectAt(settings["tokenUsage"]) if usageBlock == nil { return usage.Entry{}, false, nil } tokens := usage.TokenUsage{ - InputTokens: shared.UintField(usageBlock, "inputTokens"), - OutputTokens: shared.UintField(usageBlock, "outputTokens"), - CacheCreationInputTokens: shared.UintField(usageBlock, "cacheCreationTokens"), - CacheReadInputTokens: shared.UintField(usageBlock, "cacheReadTokens"), - ReasoningOutputTokens: shared.UintField(usageBlock, "thinkingTokens"), - } - tokens = shared.ApplyTotalFallback(tokens, shared.UintField(usageBlock, "totalTokens")) - if !shared.NonZero(tokens) { + InputTokens: agentdata.UintField(usageBlock, "inputTokens"), + OutputTokens: agentdata.UintField(usageBlock, "outputTokens"), + CacheCreationInputTokens: agentdata.UintField(usageBlock, "cacheCreationTokens"), + CacheReadInputTokens: agentdata.UintField(usageBlock, "cacheReadTokens"), + ReasoningOutputTokens: agentdata.UintField(usageBlock, "thinkingTokens"), + } + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(usageBlock, "totalTokens")) + if !usageprovider.NonZero(tokens) { return usage.Entry{}, false, nil } - provider := normalizeDroidProvider(shared.StringField(settings, "providerLock")) - model := normalizeDroidModel(shared.StringField(settings, "model")) + provider := normalizeDroidProvider(agentdata.StringField(settings, "providerLock")) + model := normalizeDroidModel(agentdata.StringField(settings, "model")) if model == "" { model, _ = sidecarModel(path) } @@ -70,17 +72,17 @@ func parseSettingsFile(path string) (usage.Entry, bool, error) { model = "unknown" } - timestamp, ok := shared.ParseTimestamp(settings["providerLockTimestamp"]) + timestamp, ok := agentdata.ParseTimestamp(settings["providerLockTimestamp"]) if !ok { - timestamp = shared.FileModifiedTime(path) + timestamp = agentdata.FileModifiedTime(path) } sessionID := strings.TrimSuffix(filepath.Base(path), ".settings.json") if sessionID == "" { sessionID = "unknown" } - entry := shared.BaseEntry(usage.ProviderDroid, timestamp, "droid", "Droid", sessionID, model, "Droid", tokens) - shared.SetSource(&entry, path, 1, 0, 0) - entry.ID = shared.StableEntryID(entry, sessionID) + entry := usageprovider.BaseEntry(usage.ProviderDroid, timestamp, "droid", "Droid", sessionID, model, "Droid", tokens) + usageprovider.SetSource(&entry, path, 1, 0, 0) + entry.ID = usageprovider.StableEntryID(entry, sessionID) return entry, true, nil } diff --git a/internal/provider/droid/provider.go b/internal/provider/droid/provider.go index 0277d7d..19bd22a 100644 --- a/internal/provider/droid/provider.go +++ b/internal/provider/droid/provider.go @@ -1,19 +1,18 @@ package droid import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads Droid usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a Droid provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/droid/provider_test.go b/internal/provider/droid/provider_test.go index cd386a5..1a1615d 100644 --- a/internal/provider/droid/provider_test.go +++ b/internal/provider/droid/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,11 +14,11 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() path := filepath.Join(dir, "session-a.settings.json") - shared.WriteFile(t, path, `{"model":"Claude-Sonnet-4-[Anthropic]","providerLock":"anthropic","providerLockTimestamp":"2026-01-02T00:00:00.000Z","tokenUsage":{"inputTokens":100,"outputTokens":50,"cacheCreationTokens":20,"cacheReadTokens":10,"thinkingTokens":5}}`) + providertest.WriteFile(t, path, `{"model":"Claude-Sonnet-4-[Anthropic]","providerLock":"anthropic","providerLockTimestamp":"2026-01-02T00:00:00.000Z","tokenUsage":{"inputTokens":100,"outputTokens":50,"cacheCreationTokens":20,"cacheReadTokens":10,"thinkingTokens":5}}`) return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderDroid, Model: "claude-sonnet-4", SessionID: "session-a", diff --git a/internal/provider/gemini/loader.go b/internal/provider/gemini/loader.go index df10108..7cd0c0c 100644 --- a/internal/provider/gemini/loader.go +++ b/internal/provider/gemini/loader.go @@ -6,18 +6,20 @@ import ( "strings" "time" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, root := range paths { - files = append(files, shared.CollectExt(root, ".json")...) - files = append(files, shared.CollectExt(root, ".jsonl")...) + files = append(files, agentdata.CollectExt(root, ".json")...) + files = append(files, agentdata.CollectExt(root, ".jsonl")...) } sort.Strings(files) - files = shared.FilterFiles(shared.UniqueStrings(files), filter) + files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) entries := make([]usage.Entry, 0) for _, file := range files { @@ -33,7 +35,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) } entries = append(entries, fileEntries...) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } @@ -48,21 +50,21 @@ type tokens struct { } func parseJSONFile(path string) ([]usage.Entry, error) { - record, err := shared.ReadJSONObject(path) + record, err := agentdata.ReadJSONObject(path) if err != nil || record == nil { return nil, err } - fallback := shared.FileModifiedTime(path) - sessionID := shared.FirstStringField(record, "sessionId", "session_id") + fallback := agentdata.FileModifiedTime(path) + sessionID := agentdata.FirstStringField(record, "sessionId", "session_id") if sessionID == "" { sessionID = strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) } sessionTimestamp := firstTimestamp(record, fallback, "startTime", "lastUpdated", "timestamp", "created_at") - if messages := shared.ArrayAt(record["messages"]); len(messages) > 0 { + if messages := agentdata.ArrayAt(record["messages"]); len(messages) > 0 { entries := make([]usage.Entry, 0) for index, raw := range messages { - message := shared.ObjectAt(raw) - if shared.StringField(message, "type") != "gemini" { + message := agentdata.ObjectAt(raw) + if agentdata.StringField(message, "type") != "gemini" { continue } if entry, ok := directEntry(message, path, index+1, "", sessionID, sessionTimestamp); ok { @@ -71,39 +73,39 @@ func parseJSONFile(path string) ([]usage.Entry, error) { } return entries, nil } - if shared.StringField(record, "type") == "gemini" { + if agentdata.StringField(record, "type") == "gemini" { if entry, ok := directEntry(record, path, 1, "", sessionID, fallback); ok { return []usage.Entry{entry}, nil } return nil, nil } - return statsEntries(recordStats(record), path, 1, shared.StringField(record, "model"), sessionID, sessionTimestamp), nil + return statsEntries(recordStats(record), path, 1, agentdata.StringField(record, "model"), sessionID, sessionTimestamp), nil } func parseJSONLFile(path string) ([]usage.Entry, error) { - lines, err := shared.ReadJSONLines(path) + lines, err := agentdata.ReadJSONLines(path) if err != nil { return nil, err } - fallback := shared.FileModifiedTime(path) + fallback := agentdata.FileModifiedTime(path) sessionID := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) currentModel := "" entries := make([]usage.Entry, 0) indexByMessageID := make(map[string]int) for _, line := range lines { record := line.Value - if value := shared.FirstStringField(record, "sessionId", "session_id"); value != "" { + if value := agentdata.FirstStringField(record, "sessionId", "session_id"); value != "" { sessionID = value } - if model := shared.StringField(record, "model"); model != "" { + if model := agentdata.StringField(record, "model"); model != "" { currentModel = model } - if shared.StringField(record, "type") == "gemini" { + if agentdata.StringField(record, "type") == "gemini" { entry, ok := directEntry(record, path, line.Line, currentModel, sessionID, fallback) if !ok { continue } - messageID := shared.StringField(record, "id") + messageID := agentdata.StringField(record, "id") if messageID != "" { if index, exists := indexByMessageID[messageID]; exists { entries[index] = entry @@ -127,22 +129,22 @@ func directEntry(record map[string]any, path string, line int, modelHint, sessio if !ok { return usage.Entry{}, false } - model := shared.StringField(record, "model") + model := agentdata.StringField(record, "model") if model == "" { model = modelHint } timestamp := firstTimestamp(record, fallback, "timestamp", "created_at") - return buildGeminiEntry(path, line, model, sessionID, timestamp, tokens, true, shared.StringField(record, "id")) + return buildGeminiEntry(path, line, model, sessionID, timestamp, tokens, true, agentdata.StringField(record, "id")) } func statsEntries(stats map[string]any, path string, line int, modelHint, sessionID string, timestamp time.Time) []usage.Entry { if stats == nil { return nil } - if models := shared.ObjectAt(stats["models"]); models != nil { + if models := agentdata.ObjectAt(stats["models"]); models != nil { entries := make([]usage.Entry, 0) for model, raw := range models { - data := shared.ObjectAt(raw) + data := agentdata.ObjectAt(raw) tokens, ok := parseTokens(data["tokens"]) if !ok { continue @@ -183,31 +185,31 @@ func buildGeminiEntry(path string, line int, model, sessionID string, timestamp ReasoningOutputTokens: tokens.thoughts, } if tokens.hasTotal { - tokenUsage = shared.ApplyTotalFallback(tokenUsage, tokens.total) + tokenUsage = usageprovider.ApplyTotalFallback(tokenUsage, tokens.total) } else if tokenUsage.TotalTokens == 0 { - tokenUsage.TotalTokens = shared.TotalUsage(tokenUsage) + tokenUsage.TotalTokens = usageprovider.TotalUsage(tokenUsage) } - if !shared.NonZero(tokenUsage) { + if !usageprovider.NonZero(tokenUsage) { return usage.Entry{}, false } - entry := shared.BaseEntry(usage.ProviderGemini, timestamp, "gemini", "Gemini", sessionID, model, "Gemini CLI", tokenUsage) - shared.SetSource(&entry, path, line, 0, 0) - entry.ID = shared.StableEntryID(entry, messageID) + entry := usageprovider.BaseEntry(usage.ProviderGemini, timestamp, "gemini", "Gemini", sessionID, model, "Gemini CLI", tokenUsage) + usageprovider.SetSource(&entry, path, line, 0, 0) + entry.ID = usageprovider.StableEntryID(entry, messageID) return entry, true } func parseTokens(raw any) (tokens, bool) { - record := shared.ObjectAt(raw) + record := agentdata.ObjectAt(raw) if record == nil { return tokens{}, false } tokens := tokens{ - input: shared.UintField(record, "input", "prompt", "input_tokens", "prompt_tokens"), - output: shared.UintField(record, "output", "candidates", "output_tokens", "candidates_tokens"), - cached: shared.UintField(record, "cached", "cached_tokens"), - thoughts: shared.UintField(record, "thoughts", "reasoning", "thoughts_tokens", "reasoning_tokens"), - tool: shared.UintField(record, "tool", "tool_tokens"), - total: shared.UintField(record, "total", "total_tokens"), + input: agentdata.UintField(record, "input", "prompt", "input_tokens", "prompt_tokens"), + output: agentdata.UintField(record, "output", "candidates", "output_tokens", "candidates_tokens"), + cached: agentdata.UintField(record, "cached", "cached_tokens"), + thoughts: agentdata.UintField(record, "thoughts", "reasoning", "thoughts_tokens", "reasoning_tokens"), + tool: agentdata.UintField(record, "tool", "tool_tokens"), + total: agentdata.UintField(record, "total", "total_tokens"), } tokens.hasTotal = tokens.total > 0 return tokens, true @@ -234,16 +236,16 @@ func normalizeGeminiInput(tokens tokens, direct bool) (uint64, uint64) { } func recordStats(record map[string]any) map[string]any { - if stats := shared.ObjectAt(record["stats"]); stats != nil { + if stats := agentdata.ObjectAt(record["stats"]); stats != nil { return stats } - result := shared.ObjectAt(record["result"]) - return shared.ObjectAt(result["stats"]) + result := agentdata.ObjectAt(record["result"]) + return agentdata.ObjectAt(result["stats"]) } func firstTimestamp(record map[string]any, fallback time.Time, keys ...string) time.Time { for _, key := range keys { - if timestamp, ok := shared.ParseTimestamp(record[key]); ok { + if timestamp, ok := agentdata.ParseTimestamp(record[key]); ok { return timestamp } } diff --git a/internal/provider/gemini/provider.go b/internal/provider/gemini/provider.go index d4ca2ad..c5ecb3c 100644 --- a/internal/provider/gemini/provider.go +++ b/internal/provider/gemini/provider.go @@ -1,19 +1,18 @@ package gemini import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads Gemini CLI usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a Gemini CLI provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/gemini/provider_test.go b/internal/provider/gemini/provider_test.go index 000cbdd..93ef8a4 100644 --- a/internal/provider/gemini/provider_test.go +++ b/internal/provider/gemini/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,13 +14,13 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() path := filepath.Join(dir, "session-a.jsonl") - shared.WriteFile(t, path, + providertest.WriteFile(t, path, `{"sessionId":"session-a","projectHash":"project-a","startTime":"2026-05-17T11:07:00.000Z"}`+"\n"+ `{"id":"msg-a","timestamp":"2026-05-17T11:07:32.000Z","type":"gemini","model":"gemini-3-flash-preview","tokens":{"input":15327,"output":23,"cached":11526,"thoughts":919,"tool":7,"total":16276}}`+"\n") return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderGemini, Model: "gemini-3-flash-preview", SessionID: "session-a", diff --git a/internal/provider/goose/loader.go b/internal/provider/goose/loader.go index 6c24295..10fef6a 100644 --- a/internal/provider/goose/loader.go +++ b/internal/provider/goose/loader.go @@ -6,7 +6,10 @@ import ( "strings" "time" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -28,14 +31,14 @@ func loadEntries(paths []string) ([]usage.Entry, error) { entries = append(entries, entry) } } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } func dbPaths(paths []string) []string { dbPaths := make([]string, 0) for _, root := range paths { - if shared.ExistingSQLiteFile(root) { + if agentdb.ExistingSQLiteFile(root) { dbPaths = append(dbPaths, root) continue } @@ -44,17 +47,17 @@ func dbPaths(paths []string) []string { filepath.Join(root, "sessions", "sessions.db"), filepath.Join(root, "data", "sessions", "sessions.db"), } { - if shared.ExistingSQLiteFile(candidate) { + if agentdb.ExistingSQLiteFile(candidate) { dbPaths = append(dbPaths, candidate) } } } sort.Strings(dbPaths) - return shared.UniqueStrings(dbPaths) + return agentdata.UniqueStrings(dbPaths) } func loadDatabase(path string) ([]usage.Entry, error) { - db, err := shared.OpenSQLite(path) + db, err := agentdb.OpenSQLite(path) if err != nil { return nil, err } @@ -75,7 +78,7 @@ func loadDatabase(path string) ([]usage.Entry, error) { entries := make([]usage.Entry, 0) for rows.Next() { var id, modelConfig, providerName, createdAt, total, input, output, accumulatedTotal, accumulatedInput, accumulatedOutput any - if !shared.ScanAny(rows, &id, &modelConfig, &providerName, &createdAt, &total, &input, &output, &accumulatedTotal, &accumulatedInput, &accumulatedOutput) { + if !agentdb.ScanAny(rows, &id, &modelConfig, &providerName, &createdAt, &total, &input, &output, &accumulatedTotal, &accumulatedInput, &accumulatedOutput) { continue } if entry, ok := rowEntry(path, id, modelConfig, providerName, createdAt, total, input, output, accumulatedTotal, accumulatedInput, accumulatedOutput); ok { @@ -86,18 +89,18 @@ func loadDatabase(path string) ([]usage.Entry, error) { } func rowEntry(path string, id, modelConfig, providerName, createdAt, total, input, output, accumulatedTotal, accumulatedInput, accumulatedOutput any) (usage.Entry, bool) { - sessionID := shared.SqlString(id) - model := modelName(shared.SqlString(modelConfig)) + sessionID := agentdb.SqlString(id) + model := modelName(agentdb.SqlString(modelConfig)) if sessionID == "" || model == "" { return usage.Entry{}, false } - timestamp, ok := timestamp(shared.SqlString(createdAt)) + timestamp, ok := timestamp(agentdb.SqlString(createdAt)) if !ok { return usage.Entry{}, false } - inputTokens := firstPositive(shared.SqlUint(accumulatedInput), shared.SqlUint(input)) - outputTokens := firstPositive(shared.SqlUint(accumulatedOutput), shared.SqlUint(output)) - totalTokens := firstPositive(shared.SqlUint(accumulatedTotal), shared.SqlUint(total), inputTokens+outputTokens) + inputTokens := firstPositive(agentdb.SqlUint(accumulatedInput), agentdb.SqlUint(input)) + outputTokens := firstPositive(agentdb.SqlUint(accumulatedOutput), agentdb.SqlUint(output)) + totalTokens := firstPositive(agentdb.SqlUint(accumulatedTotal), agentdb.SqlUint(total), inputTokens+outputTokens) tokens := usage.TokenUsage{ InputTokens: inputTokens, OutputTokens: outputTokens, @@ -105,19 +108,19 @@ func rowEntry(path string, id, modelConfig, providerName, createdAt, total, inpu if totalTokens > inputTokens+outputTokens { tokens.ReasoningOutputTokens = totalTokens - inputTokens - outputTokens } - tokens.TotalTokens = shared.TotalUsage(tokens) - if !shared.NonZero(tokens) { + tokens.TotalTokens = usageprovider.TotalUsage(tokens) + if !usageprovider.NonZero(tokens) { return usage.Entry{}, false } - entry := shared.BaseEntry(usage.ProviderGoose, timestamp, "goose", "Goose", sessionID, model, "Goose", tokens) - shared.SetSource(&entry, path, 0, 0, 0) - entry.ID = shared.StableEntryID(entry, "goose:"+sessionID+":"+shared.SqlString(providerName)) + entry := usageprovider.BaseEntry(usage.ProviderGoose, timestamp, "goose", "Goose", sessionID, model, "Goose", tokens) + usageprovider.SetSource(&entry, path, 0, 0, 0) + entry.ID = usageprovider.StableEntryID(entry, "goose:"+sessionID+":"+agentdb.SqlString(providerName)) return entry, true } func modelName(config string) string { - record := shared.DecodeJSONObjectString(config) - return shared.StringField(record, "model_name") + record := agentdata.DecodeJSONObjectString(config) + return agentdata.StringField(record, "model_name") } func timestamp(value string) (time.Time, bool) { @@ -125,14 +128,14 @@ func timestamp(value string) (time.Time, bool) { if value == "" { return time.Time{}, false } - if timestamp, ok := shared.ParseTimestampString(value); ok { + if timestamp, ok := agentdata.ParseTimestampString(value); ok { return timestamp, true } if len(value) == 19 && value[4] == '-' && value[7] == '-' && (value[10] == ' ' || value[10] == 'T') { - return shared.ParseTimestampString(value[:10] + "T" + value[11:] + "Z") + return agentdata.ParseTimestampString(value[:10] + "T" + value[11:] + "Z") } if len(value) == 10 && value[4] == '-' && value[7] == '-' { - return shared.ParseTimestampString(value + "T00:00:00Z") + return agentdata.ParseTimestampString(value + "T00:00:00Z") } return time.Time{}, false } diff --git a/internal/provider/goose/provider.go b/internal/provider/goose/provider.go index e504f8e..4faab0f 100644 --- a/internal/provider/goose/provider.go +++ b/internal/provider/goose/provider.go @@ -1,19 +1,18 @@ package goose import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads Goose usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a Goose provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -26,5 +25,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/goose/provider_test.go b/internal/provider/goose/provider_test.go index 796cedc..7be40f0 100644 --- a/internal/provider/goose/provider_test.go +++ b/internal/provider/goose/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,9 +14,9 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() dbPath := filepath.Join(dir, "sessions.db") - db := shared.OpenTestSQLite(t, dbPath) + db := providertest.OpenTestSQLite(t, dbPath) defer db.Close() - shared.ExecSQL(t, db, `CREATE TABLE sessions ( + providertest.ExecSQL(t, db, `CREATE TABLE sessions ( id TEXT PRIMARY KEY, model_config_json TEXT, provider_name TEXT, @@ -28,7 +28,7 @@ func TestLoadsEntry(t *testing.T) { accumulated_input_tokens INTEGER, accumulated_output_tokens INTEGER )`) - shared.ExecSQL(t, db, `INSERT INTO sessions ( + providertest.ExecSQL(t, db, `INSERT INTO sessions ( id, model_config_json, provider_name, created_at, accumulated_total_tokens, accumulated_input_tokens, accumulated_output_tokens ) VALUES (?, ?, ?, ?, ?, ?, ?)`, @@ -37,7 +37,7 @@ func TestLoadsEntry(t *testing.T) { return Provider{}.WithPaths([]string{dbPath}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderGoose, Model: "claude-sonnet-4-20250514", SessionID: "session-a", diff --git a/internal/provider/hermes/loader.go b/internal/provider/hermes/loader.go index 95d99bc..a80725e 100644 --- a/internal/provider/hermes/loader.go +++ b/internal/provider/hermes/loader.go @@ -4,12 +4,15 @@ import ( "strings" "time" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) func loadEntries(paths []string) ([]usage.Entry, error) { - dbPaths := shared.SqliteDBPaths(paths, "state.db", nil) + dbPaths := agentdb.SqliteDBPaths(paths, "state.db", nil) entries := make([]usage.Entry, 0) for _, dbPath := range dbPaths { dbEntries, err := loadDatabase(dbPath) @@ -18,12 +21,12 @@ func loadEntries(paths []string) ([]usage.Entry, error) { } entries = append(entries, dbEntries...) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } func loadDatabase(path string) ([]usage.Entry, error) { - db, err := shared.OpenSQLite(path) + db, err := agentdb.OpenSQLite(path) if err != nil { return nil, err } @@ -44,7 +47,7 @@ func loadDatabase(path string) ([]usage.Entry, error) { entries := make([]usage.Entry, 0) for rows.Next() { var sessionID, model, provider, startedAt, messageCount, input, output, cacheRead, cacheWrite, reasoning, estimatedCost, actualCost any - if !shared.ScanAny(rows, &sessionID, &model, &provider, &startedAt, &messageCount, &input, &output, &cacheRead, &cacheWrite, &reasoning, &estimatedCost, &actualCost) { + if !agentdb.ScanAny(rows, &sessionID, &model, &provider, &startedAt, &messageCount, &input, &output, &cacheRead, &cacheWrite, &reasoning, &estimatedCost, &actualCost) { continue } entry, ok := rowEntry(path, sessionID, model, startedAt, input, output, cacheRead, cacheWrite, reasoning) @@ -56,8 +59,8 @@ func loadDatabase(path string) ([]usage.Entry, error) { } func rowEntry(path string, sessionRaw, modelRaw, startedAt, input, output, cacheRead, cacheWrite, reasoning any) (usage.Entry, bool) { - sessionID := shared.SqlString(sessionRaw) - model := strings.TrimSpace(shared.SqlString(modelRaw)) + sessionID := agentdb.SqlString(sessionRaw) + model := strings.TrimSpace(agentdb.SqlString(modelRaw)) if sessionID == "" || model == "" { return usage.Entry{}, false } @@ -66,33 +69,33 @@ func rowEntry(path string, sessionRaw, modelRaw, startedAt, input, output, cache return usage.Entry{}, false } tokens := usage.TokenUsage{ - InputTokens: shared.SqlUint(input), - OutputTokens: shared.SqlUint(output), - CacheCreationInputTokens: shared.SqlUint(cacheWrite), - CacheReadInputTokens: shared.SqlUint(cacheRead), - ReasoningOutputTokens: shared.SqlUint(reasoning), + InputTokens: agentdb.SqlUint(input), + OutputTokens: agentdb.SqlUint(output), + CacheCreationInputTokens: agentdb.SqlUint(cacheWrite), + CacheReadInputTokens: agentdb.SqlUint(cacheRead), + ReasoningOutputTokens: agentdb.SqlUint(reasoning), } if tokens.TotalTokens == 0 { - tokens.TotalTokens = shared.TotalUsage(tokens) + tokens.TotalTokens = usageprovider.TotalUsage(tokens) } - if !shared.NonZero(tokens) { + if !usageprovider.NonZero(tokens) { return usage.Entry{}, false } - entry := shared.BaseEntry(usage.ProviderHermes, timestamp, "hermes", "Hermes", sessionID, model, "Hermes Agent", tokens) - shared.SetSource(&entry, path, 0, 0, 0) - entry.ID = shared.StableEntryID(entry, "hermes:"+sessionID) + entry := usageprovider.BaseEntry(usage.ProviderHermes, timestamp, "hermes", "Hermes", sessionID, model, "Hermes Agent", tokens) + usageprovider.SetSource(&entry, path, 0, 0, 0) + entry.ID = usageprovider.StableEntryID(entry, "hermes:"+sessionID) return entry, true } func timestamp(value any) (time.Time, bool) { - if parsed, ok := shared.ParseTimestamp(value); ok { + if parsed, ok := agentdata.ParseTimestamp(value); ok { return parsed, true } - if number, ok := shared.SqlFloat(value); ok { - return shared.TimestampFromFloat(number) + if number, ok := agentdb.SqlFloat(value); ok { + return agentdata.TimestampFromFloat(number) } - if text := shared.SqlString(value); text != "" { - return shared.ParseTimestampString(text) + if text := agentdb.SqlString(value); text != "" { + return agentdata.ParseTimestampString(text) } return time.Time{}, false } diff --git a/internal/provider/hermes/provider.go b/internal/provider/hermes/provider.go index 01e06f0..4686392 100644 --- a/internal/provider/hermes/provider.go +++ b/internal/provider/hermes/provider.go @@ -1,19 +1,18 @@ package hermes import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads Hermes Agent usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a Hermes Agent provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -26,5 +25,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/hermes/provider_test.go b/internal/provider/hermes/provider_test.go index e2113dc..9e33013 100644 --- a/internal/provider/hermes/provider_test.go +++ b/internal/provider/hermes/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,9 +14,9 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() dbPath := filepath.Join(dir, "state.db") - db := shared.OpenTestSQLite(t, dbPath) + db := providertest.OpenTestSQLite(t, dbPath) defer db.Close() - shared.ExecSQL(t, db, `CREATE TABLE sessions ( + providertest.ExecSQL(t, db, `CREATE TABLE sessions ( id TEXT PRIMARY KEY, model TEXT, billing_provider TEXT, @@ -30,7 +30,7 @@ func TestLoadsEntry(t *testing.T) { estimated_cost_usd REAL, actual_cost_usd REAL )`) - shared.ExecSQL(t, db, `INSERT INTO sessions ( + providertest.ExecSQL(t, db, `INSERT INTO sessions ( id, model, billing_provider, started_at, message_count, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, estimated_cost_usd, actual_cost_usd @@ -40,7 +40,7 @@ func TestLoadsEntry(t *testing.T) { return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderHermes, Model: "gpt-5.5", SessionID: "session-a", diff --git a/internal/provider/kilo/loader.go b/internal/provider/kilo/loader.go index 1b85f74..e53f20e 100644 --- a/internal/provider/kilo/loader.go +++ b/internal/provider/kilo/loader.go @@ -4,12 +4,15 @@ import ( "errors" "os" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) func loadEntries(paths []string) ([]usage.Entry, error) { - dbPaths := shared.SqliteDBPaths(paths, "kilo.db", nil) + dbPaths := agentdb.SqliteDBPaths(paths, "kilo.db", nil) entries := make([]usage.Entry, 0) for _, dbPath := range dbPaths { dbEntries, err := loadDatabase(dbPath) @@ -18,12 +21,12 @@ func loadEntries(paths []string) ([]usage.Entry, error) { } entries = append(entries, dbEntries...) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } func loadDatabase(path string) ([]usage.Entry, error) { - db, err := shared.OpenSQLite(path) + db, err := agentdb.OpenSQLite(path) if errors.Is(err, os.ErrNotExist) { return nil, nil } @@ -52,38 +55,38 @@ func loadDatabase(path string) ([]usage.Entry, error) { } func messageEntry(dbPath, rowID, rowSessionID, data string) (usage.Entry, bool) { - record := shared.DecodeJSONObjectString(data) - if record == nil || shared.StringField(record, "role") != "assistant" { + record := agentdata.DecodeJSONObjectString(data) + if record == nil || agentdata.StringField(record, "role") != "assistant" { return usage.Entry{}, false } - tokenBlock := shared.ObjectAt(record["tokens"]) + tokenBlock := agentdata.ObjectAt(record["tokens"]) if tokenBlock == nil { return usage.Entry{}, false } - cache := shared.ObjectAt(tokenBlock["cache"]) + cache := agentdata.ObjectAt(tokenBlock["cache"]) tokens := usage.TokenUsage{ - InputTokens: shared.UintField(tokenBlock, "input"), - OutputTokens: shared.UintField(tokenBlock, "output"), - CacheCreationInputTokens: shared.UintField(cache, "write"), - CacheReadInputTokens: shared.UintField(cache, "read"), - ReasoningOutputTokens: shared.UintField(tokenBlock, "reasoning"), + InputTokens: agentdata.UintField(tokenBlock, "input"), + OutputTokens: agentdata.UintField(tokenBlock, "output"), + CacheCreationInputTokens: agentdata.UintField(cache, "write"), + CacheReadInputTokens: agentdata.UintField(cache, "read"), + ReasoningOutputTokens: agentdata.UintField(tokenBlock, "reasoning"), } - tokens = shared.ApplyTotalFallback(tokens, shared.UintField(tokenBlock, "total")) - if !shared.NonZero(tokens) { + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(tokenBlock, "total")) + if !usageprovider.NonZero(tokens) { return usage.Entry{}, false } - model := shared.StringField(record, "modelID") + model := agentdata.StringField(record, "modelID") if model == "" { return usage.Entry{}, false } - timestamp, ok := shared.ParseTimestamp(shared.ObjectAt(record["time"])["created"]) + timestamp, ok := agentdata.ParseTimestamp(agentdata.ObjectAt(record["time"])["created"]) if !ok { return usage.Entry{}, false } - sessionID := shared.FirstNonEmpty(shared.StringField(record, "session_id"), rowSessionID, "unknown") - messageID := shared.FirstNonEmpty(shared.StringField(record, "id"), rowID) - entry := shared.BaseEntry(usage.ProviderKilo, timestamp, "kilo", "Kilo", sessionID, model, "Kilo", tokens) - shared.SetSource(&entry, dbPath, 0, 0, 0) - entry.ID = shared.StableEntryID(entry, messageID) + sessionID := agentdata.FirstNonEmpty(agentdata.StringField(record, "session_id"), rowSessionID, "unknown") + messageID := agentdata.FirstNonEmpty(agentdata.StringField(record, "id"), rowID) + entry := usageprovider.BaseEntry(usage.ProviderKilo, timestamp, "kilo", "Kilo", sessionID, model, "Kilo", tokens) + usageprovider.SetSource(&entry, dbPath, 0, 0, 0) + entry.ID = usageprovider.StableEntryID(entry, messageID) return entry, true } diff --git a/internal/provider/kilo/provider.go b/internal/provider/kilo/provider.go index ec57af1..c7c3c4b 100644 --- a/internal/provider/kilo/provider.go +++ b/internal/provider/kilo/provider.go @@ -1,19 +1,18 @@ package kilo import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads Kilo usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a Kilo provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -26,5 +25,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/kilo/provider_test.go b/internal/provider/kilo/provider_test.go index 688926a..d0fe3b4 100644 --- a/internal/provider/kilo/provider_test.go +++ b/internal/provider/kilo/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,10 +14,10 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() dbPath := filepath.Join(dir, "kilo.db") - db := shared.OpenTestSQLite(t, dbPath) + db := providertest.OpenTestSQLite(t, dbPath) defer db.Close() - shared.ExecSQL(t, db, "CREATE TABLE message (id TEXT, session_id TEXT, data TEXT)") - shared.ExecSQL(t, db, `INSERT INTO message (id, session_id, data) VALUES (?, ?, ?)`, + providertest.ExecSQL(t, db, "CREATE TABLE message (id TEXT, session_id TEXT, data TEXT)") + providertest.ExecSQL(t, db, `INSERT INTO message (id, session_id, data) VALUES (?, ?, ?)`, "row-1", "session-a", `{"id":"msg-1","role":"assistant","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","time":{"created":1767312000000},"tokens":{"input":100,"output":50,"reasoning":5,"cache":{"read":10,"write":20}}}`, @@ -25,7 +25,7 @@ func TestLoadsEntry(t *testing.T) { return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderKilo, Model: "claude-sonnet-4-20250514", SessionID: "session-a", diff --git a/internal/provider/kimi/loader.go b/internal/provider/kimi/loader.go index beb12a5..65e3ddf 100644 --- a/internal/provider/kimi/loader.go +++ b/internal/provider/kimi/loader.go @@ -6,7 +6,9 @@ import ( "strings" "time" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -15,10 +17,10 @@ const defaultModel = "kimi-for-coding" func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, root := range paths { - files = append(files, shared.CollectFiles(filepath.Join(root, "sessions"), isWireFile)...) + files = append(files, agentdata.CollectFiles(filepath.Join(root, "sessions"), isWireFile)...) } sort.Strings(files) - files = shared.FilterFiles(shared.UniqueStrings(files), filter) + files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) entries := make([]usage.Entry, 0) seen := make(map[string]bool) @@ -35,7 +37,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) entries = append(entries, entry) } } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } @@ -68,18 +70,18 @@ type record struct { } func parseWireFile(path string) ([]usage.Entry, error) { - lines, err := shared.ReadJSONLines(path, `usage`) + lines, err := agentdata.ReadJSONLines(path, `usage`) if err != nil { return nil, err } configModel := configModel(path) sessionID := sessionID(path) - fallback := shared.FileModifiedTime(path) + fallback := agentdata.FileModifiedTime(path) entries := make([]usage.Entry, 0) for _, line := range lines { var record record var ok bool - if shared.StringField(line.Value, "type") == "usage.record" { + if agentdata.StringField(line.Value, "type") == "usage.record" { record, ok = parseUsageRecord(line.Value) } else { record, ok = parseStatusUpdate(line.Value) @@ -95,63 +97,63 @@ func parseWireFile(path string) ([]usage.Entry, error) { if model == "" { model = configModel } - entry := shared.BaseEntry(usage.ProviderKimi, timestamp, "kimi", "Kimi", sessionID, model, "Kimi", record.tokens) - shared.SetSource(&entry, path, line.Line, line.Start, line.End) - entry.ID = shared.StableEntryID(entry, record.messageID) + entry := usageprovider.BaseEntry(usage.ProviderKimi, timestamp, "kimi", "Kimi", sessionID, model, "Kimi", record.tokens) + usageprovider.SetSource(&entry, path, line.Line, line.Start, line.End) + entry.ID = usageprovider.StableEntryID(entry, record.messageID) entries = append(entries, entry) } return entries, nil } func parseStatusUpdate(value map[string]any) (record, bool) { - message := shared.ObjectAt(value["message"]) - if shared.StringField(message, "type") != "StatusUpdate" { + message := agentdata.ObjectAt(value["message"]) + if agentdata.StringField(message, "type") != "StatusUpdate" { return record{}, false } - payload := shared.ObjectAt(message["payload"]) - tokenUsage := shared.ObjectAt(payload["token_usage"]) + payload := agentdata.ObjectAt(message["payload"]) + tokenUsage := agentdata.ObjectAt(payload["token_usage"]) if tokenUsage == nil { return record{}, false } tokens := usage.TokenUsage{ - InputTokens: shared.UintField(tokenUsage, "input_other"), - OutputTokens: shared.UintField(tokenUsage, "output"), - CacheCreationInputTokens: shared.UintField(tokenUsage, "input_cache_creation"), - CacheReadInputTokens: shared.UintField(tokenUsage, "input_cache_read"), + InputTokens: agentdata.UintField(tokenUsage, "input_other"), + OutputTokens: agentdata.UintField(tokenUsage, "output"), + CacheCreationInputTokens: agentdata.UintField(tokenUsage, "input_cache_creation"), + CacheReadInputTokens: agentdata.UintField(tokenUsage, "input_cache_read"), } - tokens = shared.ApplyTotalFallback(tokens, shared.UintField(tokenUsage, "total")) - if !shared.NonZero(tokens) { + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(tokenUsage, "total")) + if !usageprovider.NonZero(tokens) { return record{}, false } - record := record{tokens: tokens, messageID: shared.StringField(payload, "message_id")} - record.timestamp, record.hasTime = shared.ParseTimestamp(value["timestamp"]) + record := record{tokens: tokens, messageID: agentdata.StringField(payload, "message_id")} + record.timestamp, record.hasTime = agentdata.ParseTimestamp(value["timestamp"]) return record, true } func parseUsageRecord(value map[string]any) (record, bool) { // Session-scoped records are cumulative totals; only turn records count. - if shared.StringField(value, "usageScope") != "turn" { + if agentdata.StringField(value, "usageScope") != "turn" { return record{}, false } - tokenUsage := shared.ObjectAt(value["usage"]) + tokenUsage := agentdata.ObjectAt(value["usage"]) if tokenUsage == nil { return record{}, false } tokens := usage.TokenUsage{ - InputTokens: shared.UintField(tokenUsage, "inputOther"), - OutputTokens: shared.UintField(tokenUsage, "output"), - CacheCreationInputTokens: shared.UintField(tokenUsage, "inputCacheCreation"), - CacheReadInputTokens: shared.UintField(tokenUsage, "inputCacheRead"), + InputTokens: agentdata.UintField(tokenUsage, "inputOther"), + OutputTokens: agentdata.UintField(tokenUsage, "output"), + CacheCreationInputTokens: agentdata.UintField(tokenUsage, "inputCacheCreation"), + CacheReadInputTokens: agentdata.UintField(tokenUsage, "inputCacheRead"), } - tokens = shared.ApplyTotalFallback(tokens, 0) - if !shared.NonZero(tokens) { + tokens = usageprovider.ApplyTotalFallback(tokens, 0) + if !usageprovider.NonZero(tokens) { return record{}, false } record := record{ tokens: tokens, - model: strings.TrimPrefix(shared.StringField(value, "model"), "kimi-code/"), + model: strings.TrimPrefix(agentdata.StringField(value, "model"), "kimi-code/"), } - record.timestamp, record.hasTime = shared.ParseTimestamp(value["time"]) + record.timestamp, record.hasTime = agentdata.ParseTimestamp(value["time"]) return record, true } @@ -188,11 +190,11 @@ func configModel(path string) string { if root == "" { return defaultModel } - config, err := shared.ReadJSONObject(filepath.Join(root, "config.json")) + config, err := agentdata.ReadJSONObject(filepath.Join(root, "config.json")) if err != nil || config == nil { return defaultModel } - if model := shared.StringField(config, "model"); model != "" { + if model := agentdata.StringField(config, "model"); model != "" { return model } return defaultModel diff --git a/internal/provider/kimi/provider.go b/internal/provider/kimi/provider.go index 60da2e7..5d41be4 100644 --- a/internal/provider/kimi/provider.go +++ b/internal/provider/kimi/provider.go @@ -1,19 +1,18 @@ package kimi import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads Kimi usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a Kimi provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/kimi/provider_test.go b/internal/provider/kimi/provider_test.go index cd93a12..017c923 100644 --- a/internal/provider/kimi/provider_test.go +++ b/internal/provider/kimi/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -13,15 +13,15 @@ import ( func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() - shared.WriteFile(t, filepath.Join(dir, "config.json"), `{"model":"kimi-k2"}`) + providertest.WriteFile(t, filepath.Join(dir, "config.json"), `{"model":"kimi-k2"}`) path := filepath.Join(dir, "sessions", "group", "session-a", "wire.jsonl") - shared.WriteFile(t, path, + providertest.WriteFile(t, path, `{"type":"metadata","protocol_version":"1.3"}`+"\n"+ `{"timestamp":1770983427.123,"message":{"type":"StatusUpdate","payload":{"token_usage":{"input_other":100,"output":50,"input_cache_read":10,"input_cache_creation":20},"message_id":"msg-1"}}}`+"\n") return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderKimi, Model: "kimi-k2", SessionID: "session-a", diff --git a/internal/provider/openclaw/loader.go b/internal/provider/openclaw/loader.go index 50352d6..3537fdf 100644 --- a/internal/provider/openclaw/loader.go +++ b/internal/provider/openclaw/loader.go @@ -5,17 +5,19 @@ import ( "sort" "strings" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, path := range paths { - files = append(files, shared.CollectFiles(path, isSessionFile)...) + files = append(files, agentdata.CollectFiles(path, isSessionFile)...) } sort.Strings(files) - files = shared.FilterFiles(files, filter) + files = agentdata.FilterFiles(files, filter) entries := make([]usage.Entry, 0) for _, file := range files { @@ -25,7 +27,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) } entries = append(entries, fileEntries...) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } @@ -40,7 +42,7 @@ func isSessionFile(path string) bool { } func parseSessionFile(path string) ([]usage.Entry, error) { - lines, err := shared.ReadJSONLines(path) + lines, err := agentdata.ReadJSONLines(path) if err != nil { return nil, err } @@ -51,70 +53,70 @@ func parseSessionFile(path string) ([]usage.Entry, error) { for _, line := range lines { record := line.Value if isModelChange(record) { - source := shared.ObjectAt(record["data"]) + source := agentdata.ObjectAt(record["data"]) if source == nil { source = record } - if model := shared.FirstStringField(source, "modelId", "model"); model != "" { + if model := agentdata.FirstStringField(source, "modelId", "model"); model != "" { currentModel = model } - if provider := shared.StringField(source, "provider"); provider != "" { + if provider := agentdata.StringField(source, "provider"); provider != "" { currentProvider = provider } continue } - if shared.StringField(record, "type") != "message" { + if agentdata.StringField(record, "type") != "message" { continue } - message := shared.ObjectAt(record["message"]) - if shared.StringField(message, "role") != "assistant" { + message := agentdata.ObjectAt(record["message"]) + if agentdata.StringField(message, "role") != "assistant" { continue } - usageBlock := shared.ObjectAt(message["usage"]) + usageBlock := agentdata.ObjectAt(message["usage"]) if usageBlock == nil { continue } - timestamp, ok := shared.ParseTimestamp(message["timestamp"]) + timestamp, ok := agentdata.ParseTimestamp(message["timestamp"]) if !ok { - timestamp, ok = shared.ParseTimestamp(record["timestamp"]) + timestamp, ok = agentdata.ParseTimestamp(record["timestamp"]) } if !ok { - timestamp = shared.FileModifiedTime(path) + timestamp = agentdata.FileModifiedTime(path) } - model := shared.FirstStringField(message, "modelId", "model") + model := agentdata.FirstStringField(message, "modelId", "model") if model == "" { model = currentModel } if model == "" { model = "unknown" } - provider := shared.StringField(message, "provider") + provider := agentdata.StringField(message, "provider") if provider == "" { provider = currentProvider } tokens := usage.TokenUsage{ - InputTokens: shared.UintField(usageBlock, "input"), - OutputTokens: shared.UintField(usageBlock, "output"), - CacheCreationInputTokens: shared.UintField(usageBlock, "cacheWrite"), - CacheReadInputTokens: shared.UintField(usageBlock, "cacheRead"), + InputTokens: agentdata.UintField(usageBlock, "input"), + OutputTokens: agentdata.UintField(usageBlock, "output"), + CacheCreationInputTokens: agentdata.UintField(usageBlock, "cacheWrite"), + CacheReadInputTokens: agentdata.UintField(usageBlock, "cacheRead"), } - tokens = shared.ApplyTotalFallback(tokens, shared.UintField(usageBlock, "totalTokens")) - if !shared.NonZero(tokens) { + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(usageBlock, "totalTokens")) + if !usageprovider.NonZero(tokens) { continue } - entry := shared.BaseEntry(usage.ProviderOpenClaw, timestamp, "openclaw", "OpenClaw", sessionID, "[openclaw] "+model, "OpenClaw", tokens) - shared.SetSource(&entry, path, line.Line, line.Start, line.End) - entry.ID = shared.StableEntryID(entry, provider) + entry := usageprovider.BaseEntry(usage.ProviderOpenClaw, timestamp, "openclaw", "OpenClaw", sessionID, "[openclaw] "+model, "OpenClaw", tokens) + usageprovider.SetSource(&entry, path, line.Line, line.Start, line.End) + entry.ID = usageprovider.StableEntryID(entry, provider) entries = append(entries, entry) } return entries, nil } func isModelChange(record map[string]any) bool { - if shared.StringField(record, "type") == "model_change" { + if agentdata.StringField(record, "type") == "model_change" { return true } - return shared.StringField(record, "type") == "custom" && shared.StringField(record, "customType") == "model-snapshot" + return agentdata.StringField(record, "type") == "custom" && agentdata.StringField(record, "customType") == "model-snapshot" } func openClawSessionID(path string) string { diff --git a/internal/provider/openclaw/provider.go b/internal/provider/openclaw/provider.go index 5cbec48..f64fe06 100644 --- a/internal/provider/openclaw/provider.go +++ b/internal/provider/openclaw/provider.go @@ -1,19 +1,18 @@ package openclaw import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads OpenClaw usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a OpenClaw provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/openclaw/provider_test.go b/internal/provider/openclaw/provider_test.go index 96a33a5..4890459 100644 --- a/internal/provider/openclaw/provider_test.go +++ b/internal/provider/openclaw/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,13 +14,13 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() path := filepath.Join(dir, "agents", "main", "sessions", "abc.jsonl") - shared.WriteFile(t, path, + providertest.WriteFile(t, path, `{"type":"model_change","provider":"openai-codex","modelId":"gpt-5.2"}`+"\n"+ `{"type":"message","message":{"role":"assistant","usage":{"input":1660,"output":55,"cacheRead":108928},"timestamp":1769753935279}}`+"\n") return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderOpenClaw, Model: "[openclaw] gpt-5.2", SessionID: "abc", diff --git a/internal/provider/opencode/loader.go b/internal/provider/opencode/loader.go index 54eec95..68cacfb 100644 --- a/internal/provider/opencode/loader.go +++ b/internal/provider/opencode/loader.go @@ -10,7 +10,10 @@ import ( "strings" "time" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -40,7 +43,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) } for _, entry := range rootEntries { if entry.ID == "" { - entry.ID = shared.StableEntryID(entry) + entry.ID = usageprovider.StableEntryID(entry) } if _, exists := entriesByID[entry.ID]; !exists { entriesByID[entry.ID] = entry @@ -51,7 +54,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) for _, entry := range entriesByID { entries = append(entries, entry) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } @@ -87,9 +90,9 @@ func loadRoot(root string, filter usage.FileFilter) ([]usage.Entry, error) { // Older OpenCode releases wrote one JSON file per message instead of the // database. Their users keep those files, so keep reading them. - files := shared.CollectExt(filepath.Join(root, "storage", "message"), ".json") + files := agentdata.CollectExt(filepath.Join(root, "storage", "message"), ".json") sort.Strings(files) - files = shared.FilterFiles(files, filter) + files = agentdata.FilterFiles(files, filter) for _, file := range files { stem := strings.TrimSuffix(filepath.Base(file), filepath.Ext(file)) if seenIDs[stableMessageID(stem)] { @@ -109,7 +112,7 @@ func loadRoot(root string, filter usage.FileFilter) ([]usage.Entry, error) { entries = append(entries, entry) } } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } @@ -123,7 +126,7 @@ func dbPaths(root string) []string { } paths := make([]string, 0, len(matches)) for _, candidate := range matches { - if shared.ExistingSQLiteFile(candidate) { + if agentdb.ExistingSQLiteFile(candidate) { paths = append(paths, candidate) } } @@ -132,7 +135,7 @@ func dbPaths(root string) []string { } func loadDatabase(path string) ([]usage.Entry, error) { - db, err := shared.OpenSQLite(path) + db, err := agentdb.OpenSQLite(path) if err != nil { return nil, err } @@ -198,13 +201,13 @@ func queryMessages(db *sql.DB) []*message { if err := rows.Scan(&id, &sessionID, &created, &data); err != nil { continue } - record := shared.DecodeJSONObjectString(data) + record := agentdata.DecodeJSONObjectString(data) if record == nil { continue } messages = append(messages, &message{ - id: shared.FirstNonEmpty(shared.StringField(record, "id"), id), - sessionID: shared.FirstNonEmpty(shared.StringField(record, "sessionID"), sessionID), + id: agentdata.FirstNonEmpty(agentdata.StringField(record, "id"), id), + sessionID: agentdata.FirstNonEmpty(agentdata.StringField(record, "sessionID"), sessionID), created: created, data: record, }) @@ -232,7 +235,7 @@ func attachParts(db *sql.DB, messages []*message) { if len(targets) == 0 { continue } - record := shared.DecodeJSONObjectString(data) + record := agentdata.DecodeJSONObjectString(data) if record == nil { continue } @@ -250,7 +253,7 @@ func sessionEntries(source string, session session, messages []*message) []usage // running counters would make the next assistant message look like a // fresh session and bill it for the whole context again. reported := tokens(message.data) - if !shared.NonZero(reported) { + if !usageprovider.NonZero(reported) { continue } billed, ok := billTokens(reported, previous) @@ -273,17 +276,17 @@ func sessionEntries(source string, session session, messages []*message) []usage // are what the model produced on this turn alone. billTokens is what // turns this mixture into one turn's cost. func tokens(record map[string]any) usage.TokenUsage { - block := shared.ObjectAt(record["tokens"]) + block := agentdata.ObjectAt(record["tokens"]) if block == nil { return usage.TokenUsage{} } - cache := shared.ObjectAt(block["cache"]) + cache := agentdata.ObjectAt(block["cache"]) return usage.TokenUsage{ - InputTokens: shared.UintField(block, "input"), - OutputTokens: shared.UintField(block, "output"), - ReasoningOutputTokens: shared.UintField(block, "reasoning"), - CacheCreationInputTokens: shared.UintField(cache, "write"), - CacheReadInputTokens: shared.UintField(cache, "read"), + InputTokens: agentdata.UintField(block, "input"), + OutputTokens: agentdata.UintField(block, "output"), + ReasoningOutputTokens: agentdata.UintField(block, "reasoning"), + CacheCreationInputTokens: agentdata.UintField(cache, "write"), + CacheReadInputTokens: agentdata.UintField(cache, "read"), } } @@ -299,10 +302,10 @@ func billTokens(reported, previous usage.TokenUsage) (usage.TokenUsage, bool) { OutputTokens: reported.OutputTokens, ReasoningOutputTokens: reported.ReasoningOutputTokens, } - if !shared.NonZero(billed) { + if !usageprovider.NonZero(billed) { return usage.TokenUsage{}, false } - return shared.ApplyTotalFallback(billed, 0), true + return usageprovider.ApplyTotalFallback(billed, 0), true } // growth returns how much a running counter advanced. A counter that shrank @@ -316,21 +319,21 @@ func growth(current, previous uint64) uint64 { func newEntry(source string, session session, message *message, tokens usage.TokenUsage) (usage.Entry, bool) { record := message.data - model := shared.StringField(record, "modelID") + model := agentdata.StringField(record, "modelID") if model == "" { return usage.Entry{}, false } timestamp := time.UnixMilli(message.created).UTC() - if parsed, ok := shared.ParseTimestamp(shared.ObjectAt(record["time"])["created"]); ok { + if parsed, ok := agentdata.ParseTimestamp(agentdata.ObjectAt(record["time"])["created"]); ok { timestamp = parsed } // The message records the directory the agent actually ran in; the session // directory covers messages written before that field existed. - cwd := shared.FirstNonEmpty( - shared.StringField(shared.ObjectAt(record["path"]), "cwd"), - shared.StringField(shared.ObjectAt(record["path"]), "root"), + cwd := agentdata.FirstNonEmpty( + agentdata.StringField(agentdata.ObjectAt(record["path"]), "cwd"), + agentdata.StringField(agentdata.ObjectAt(record["path"]), "root"), session.directory, ) project, projectPath := usage.UnknownProject, "" @@ -338,12 +341,12 @@ func newEntry(source string, session session, message *message, tokens usage.Tok project, projectPath = name, path } - sessionID := shared.FirstNonEmpty(message.sessionID, usage.UnknownProject) - entry := shared.BaseEntry(usage.ProviderOpenCode, timestamp, project, projectPath, sessionID, model, "OpenCode", tokens) - shared.SetSource(&entry, source, 0, 0, 0) + sessionID := agentdata.FirstNonEmpty(message.sessionID, usage.UnknownProject) + entry := usageprovider.BaseEntry(usage.ProviderOpenCode, timestamp, project, projectPath, sessionID, model, "OpenCode", tokens) + usageprovider.SetSource(&entry, source, 0, 0, 0) entry.ID = stableMessageID(message.id) if entry.ID == "" { - entry.ID = shared.StableEntryID(entry) + entry.ID = usageprovider.StableEntryID(entry) } for _, part := range message.parts { @@ -359,18 +362,18 @@ func newEntry(source string, session session, message *message, tokens usage.Tok // worktree, which lists files nobody edited (.DS_Store, build output) and // carries no line counts, so treating it as agent work invents writes. func partChanges(part map[string]any, cwd string) []usage.FileChange { - if shared.StringField(part, "type") != "tool" { + if agentdata.StringField(part, "type") != "tool" { return nil } - state := shared.ObjectAt(part["state"]) - if shared.StringField(state, "status") != "completed" { + state := agentdata.ObjectAt(part["state"]) + if agentdata.StringField(state, "status") != "completed" { return nil } - input := shared.ObjectAt(state["input"]) - metadata := shared.ObjectAt(state["metadata"]) + input := agentdata.ObjectAt(state["input"]) + metadata := agentdata.ObjectAt(state["metadata"]) - switch shared.StringField(part, "tool") { + switch agentdata.StringField(part, "tool") { case "edit": return editChanges(input, metadata, cwd) case "write": @@ -385,12 +388,12 @@ func partChanges(part map[string]any, cwd string) []usage.FileChange { // editChanges prefers the diff OpenCode computed. Without it the // replaced and replacing strings still bound the change. func editChanges(input, metadata map[string]any, cwd string) []usage.FileChange { - path := usage.ResolvePath(cwd, shared.StringField(input, "filePath")) + path := usage.ResolvePath(cwd, agentdata.StringField(input, "filePath")) - if diff := shared.ObjectAt(metadata["filediff"]); diff != nil { - added := shared.UintField(diff, "additions") - removed := shared.UintField(diff, "deletions") - if diffPath := shared.FirstNonEmpty(shared.StringField(diff, "filePath"), shared.StringField(diff, "file")); diffPath != "" { + if diff := agentdata.ObjectAt(metadata["filediff"]); diff != nil { + added := agentdata.UintField(diff, "additions") + removed := agentdata.UintField(diff, "deletions") + if diffPath := agentdata.FirstNonEmpty(agentdata.StringField(diff, "filePath"), agentdata.StringField(diff, "file")); diffPath != "" { path = usage.ResolvePath(cwd, diffPath) } if added != 0 || removed != 0 { @@ -398,8 +401,8 @@ func editChanges(input, metadata map[string]any, cwd string) []usage.FileChange } } - added := usage.CountLines(shared.StringField(input, "newString")) - removed := usage.CountLines(shared.StringField(input, "oldString")) + added := usage.CountLines(agentdata.StringField(input, "newString")) + removed := usage.CountLines(agentdata.StringField(input, "oldString")) if path == "" && added == 0 && removed == 0 { return nil } @@ -410,14 +413,14 @@ func editChanges(input, metadata map[string]any, cwd string) []usage.FileChange // an existing file replaces lines this record does not describe, so only the // added side is known. func writeChanges(input, metadata map[string]any, cwd string) []usage.FileChange { - path := usage.ResolvePath(cwd, shared.FirstNonEmpty( - shared.StringField(input, "filePath"), - shared.StringField(metadata, "filepath"), + path := usage.ResolvePath(cwd, agentdata.FirstNonEmpty( + agentdata.StringField(input, "filePath"), + agentdata.StringField(metadata, "filepath"), )) if path == "" { return nil } - return []usage.FileChange{{Path: path, LinesAdded: usage.CountLines(shared.StringField(input, "content"))}} + return []usage.FileChange{{Path: path, LinesAdded: usage.CountLines(agentdata.StringField(input, "content"))}} } // applyPatchChanges prefers the per-file summary OpenCode attaches and @@ -426,18 +429,18 @@ func applyPatchChanges(input, metadata map[string]any, cwd string) []usage.FileC if files, ok := metadata["files"].([]any); ok && len(files) > 0 { changes := make([]usage.FileChange, 0, len(files)) for _, raw := range files { - file := shared.ObjectAt(raw) + file := agentdata.ObjectAt(raw) if file == nil { continue } - path := usage.ResolvePath(cwd, shared.FirstNonEmpty(shared.StringField(file, "filePath"), shared.StringField(file, "file"))) + path := usage.ResolvePath(cwd, agentdata.FirstNonEmpty(agentdata.StringField(file, "filePath"), agentdata.StringField(file, "file"))) if path == "" { continue } changes = append(changes, usage.FileChange{ Path: path, - LinesAdded: shared.UintField(file, "additions"), - LinesRemoved: shared.UintField(file, "deletions"), + LinesAdded: agentdata.UintField(file, "additions"), + LinesRemoved: agentdata.UintField(file, "deletions"), }) } if len(changes) > 0 { @@ -445,10 +448,10 @@ func applyPatchChanges(input, metadata map[string]any, cwd string) []usage.FileC } } - patchText := shared.FirstNonEmpty( - shared.StringField(input, "patchText"), - shared.StringField(input, "patch"), - shared.StringField(metadata, "patch"), + patchText := agentdata.FirstNonEmpty( + agentdata.StringField(input, "patchText"), + agentdata.StringField(input, "patch"), + agentdata.StringField(metadata, "patch"), ) return parsePatchEnvelope(patchText, cwd) } @@ -521,10 +524,10 @@ func loadLegacyFile(path string) ([]usage.Entry, error) { } message := &message{ - id: shared.StringField(record, "id"), - sessionID: shared.StringField(record, "sessionID"), + id: agentdata.StringField(record, "id"), + sessionID: agentdata.StringField(record, "sessionID"), data: record, - parts: readLegacyParts(path, shared.StringField(record, "id")), + parts: readLegacyParts(path, agentdata.StringField(record, "id")), } // One file holds one message, so there is no predecessor to subtract from: // billing against a zero baseline charges it for everything it reports. @@ -548,7 +551,7 @@ func readLegacyParts(messagePath, messageID string) []map[string]any { return nil } storage := filepath.Dir(filepath.Dir(filepath.Dir(messagePath))) - files := shared.CollectExt(filepath.Join(storage, "part", messageID), ".json") + files := agentdata.CollectExt(filepath.Join(storage, "part", messageID), ".json") sort.Strings(files) parts := make([]map[string]any, 0, len(files)) diff --git a/internal/provider/opencode/loader_test.go b/internal/provider/opencode/loader_test.go index e676082..0cb35db 100644 --- a/internal/provider/opencode/loader_test.go +++ b/internal/provider/opencode/loader_test.go @@ -1,12 +1,12 @@ package opencode import ( - "database/sql" "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // openCodeDB builds a database shaped like the one OpenCode writes: a session @@ -26,7 +26,7 @@ func writeOpenCodeDB(t *testing.T, sessions, messages, parts []openCodeRow) stri func writeOpenCodeDBAt(t *testing.T, path string, sessions, messages, parts []openCodeRow) string { t.Helper() - db, err := sql.Open("sqlite", path) + db, err := agentdb.OpenSQLite(path) if err != nil { t.Fatal(err) } @@ -81,7 +81,7 @@ func TestOpenCodeBillsRunningInputAndPerTurnOutput(t *testing.T) { if len(entries) != 2 { t.Fatalf("len(entries) = %d, want 2", len(entries)) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) first, second := entries[0], entries[1] if first.Usage.InputTokens != 1000 || first.Usage.OutputTokens != 50 || first.Usage.ReasoningOutputTokens != 20 { @@ -121,7 +121,7 @@ func TestOpenCodeIgnoresZeroTokenMessages(t *testing.T) { if len(entries) != 2 { t.Fatalf("len(entries) = %d, want 2 (the user turn bills nothing)", len(entries)) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) if got := entries[1].Usage.InputTokens; got != 200 { t.Fatalf("input after a user turn = %d, want 200 (the user turn must not reset the counter)", got) } @@ -141,7 +141,7 @@ func TestOpenCodeRestartsBillingAfterCompaction(t *testing.T) { if err != nil { t.Fatal(err) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) if len(entries) != 2 { t.Fatalf("len(entries) = %d, want 2", len(entries)) } @@ -165,7 +165,7 @@ func TestOpenCodeProjectComesFromMessageCwd(t *testing.T) { if err != nil { t.Fatal(err) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) if len(entries) != 3 { t.Fatalf("len(entries) = %d, want 3", len(entries)) } diff --git a/internal/provider/opencode/provider.go b/internal/provider/opencode/provider.go index 3455d71..31dd4fc 100644 --- a/internal/provider/opencode/provider.go +++ b/internal/provider/opencode/provider.go @@ -1,19 +1,18 @@ package opencode import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads OpenCode usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a OpenCode provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/opencode/provider_test.go b/internal/provider/opencode/provider_test.go index e04c80a..7716a5c 100644 --- a/internal/provider/opencode/provider_test.go +++ b/internal/provider/opencode/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,11 +14,11 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() path := filepath.Join(dir, "storage", "message", "session-a", "msg-1.json") - shared.WriteFile(t, path, `{"id":"msg-1","sessionID":"session-a","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","path":{"cwd":"/repo/demo"},"time":{"created":1767312000000},"tokens":{"input":100,"output":50,"cache":{"read":10,"write":20}},"cost":0}`) + providertest.WriteFile(t, path, `{"id":"msg-1","sessionID":"session-a","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","path":{"cwd":"/repo/demo"},"time":{"created":1767312000000},"tokens":{"input":100,"output":50,"cache":{"read":10,"write":20}},"cost":0}`) return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderOpenCode, Model: "claude-sonnet-4-20250514", SessionID: "session-a", diff --git a/internal/provider/pi/loader.go b/internal/provider/pi/loader.go index d5a55b2..a9b40a6 100644 --- a/internal/provider/pi/loader.go +++ b/internal/provider/pi/loader.go @@ -5,17 +5,19 @@ import ( "sort" "strings" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, path := range paths { - files = append(files, shared.CollectExt(path, ".jsonl")...) + files = append(files, agentdata.CollectExt(path, ".jsonl")...) } sort.Strings(files) - files = shared.FilterFiles(files, filter) + files = agentdata.FilterFiles(files, filter) entries := make([]usage.Entry, 0) for _, file := range files { @@ -25,12 +27,12 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) } entries = append(entries, fileEntries...) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } func parseSessionFile(path string) ([]usage.Entry, error) { - lines, err := shared.ReadJSONLines(path, `"usage"`, `"message"`) + lines, err := agentdata.ReadJSONLines(path, `"usage"`, `"message"`) if err != nil { return nil, err } @@ -38,38 +40,38 @@ func parseSessionFile(path string) ([]usage.Entry, error) { sessionID := sessionID(path) entries := make([]usage.Entry, 0) for _, line := range lines { - if typ := shared.StringField(line.Value, "type"); typ != "" && typ != "message" { + if typ := agentdata.StringField(line.Value, "type"); typ != "" && typ != "message" { continue } - message := shared.ObjectAt(line.Value["message"]) - if shared.StringField(message, "role") != "assistant" { + message := agentdata.ObjectAt(line.Value["message"]) + if agentdata.StringField(message, "role") != "assistant" { continue } - usageBlock := shared.ObjectAt(message["usage"]) + usageBlock := agentdata.ObjectAt(message["usage"]) if usageBlock == nil { continue } - timestamp, ok := shared.ParseTimestamp(line.Value["timestamp"]) + timestamp, ok := agentdata.ParseTimestamp(line.Value["timestamp"]) if !ok { continue } tokens := usage.TokenUsage{ - InputTokens: shared.UintField(usageBlock, "input"), - OutputTokens: shared.UintField(usageBlock, "output"), - CacheCreationInputTokens: shared.UintField(usageBlock, "cacheWrite"), - CacheReadInputTokens: shared.UintField(usageBlock, "cacheRead"), + InputTokens: agentdata.UintField(usageBlock, "input"), + OutputTokens: agentdata.UintField(usageBlock, "output"), + CacheCreationInputTokens: agentdata.UintField(usageBlock, "cacheWrite"), + CacheReadInputTokens: agentdata.UintField(usageBlock, "cacheRead"), } - tokens = shared.ApplyTotalFallback(tokens, shared.UintField(usageBlock, "totalTokens")) - if !shared.NonZero(tokens) { + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(usageBlock, "totalTokens")) + if !usageprovider.NonZero(tokens) { continue } - model := shared.StringField(message, "model") + model := agentdata.StringField(message, "model") if model != "" { model = "[pi] " + model } - entry := shared.BaseEntry(usage.ProviderPi, timestamp, project, project, sessionID, model, "pi-agent", tokens) - shared.SetSource(&entry, path, line.Line, line.Start, line.End) - entry.ID = shared.StableEntryID(entry) + entry := usageprovider.BaseEntry(usage.ProviderPi, timestamp, project, project, sessionID, model, "pi-agent", tokens) + usageprovider.SetSource(&entry, path, line.Line, line.Start, line.End) + entry.ID = usageprovider.StableEntryID(entry) entries = append(entries, entry) } return entries, nil diff --git a/internal/provider/pi/provider.go b/internal/provider/pi/provider.go index 88fcd92..b16f39a 100644 --- a/internal/provider/pi/provider.go +++ b/internal/provider/pi/provider.go @@ -1,19 +1,18 @@ package pi import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads pi-agent usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a pi-agent provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/pi/provider_test.go b/internal/provider/pi/provider_test.go index dc2585b..72993e1 100644 --- a/internal/provider/pi/provider_test.go +++ b/internal/provider/pi/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,11 +14,11 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() path := filepath.Join(dir, "project-a", "agent_session-a.jsonl") - shared.WriteFile(t, path, `{"type":"message","timestamp":"2026-01-02T00:00:00.000Z","message":{"role":"assistant","model":"gpt-5","usage":{"totalTokens":333}}}`+"\n") + providertest.WriteFile(t, path, `{"type":"message","timestamp":"2026-01-02T00:00:00.000Z","message":{"role":"assistant","model":"gpt-5","usage":{"totalTokens":333}}}`+"\n") return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderPi, Model: "[pi] gpt-5", SessionID: "session-a", diff --git a/internal/provider/qwen/loader.go b/internal/provider/qwen/loader.go index 8122d40..9244e8d 100644 --- a/internal/provider/qwen/loader.go +++ b/internal/provider/qwen/loader.go @@ -5,20 +5,22 @@ import ( "sort" "strings" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, root := range paths { - files = append(files, shared.CollectFiles(filepath.Join(root, "projects"), isChatFile)...) + files = append(files, agentdata.CollectFiles(filepath.Join(root, "projects"), isChatFile)...) if strings.Contains(filepath.ToSlash(root), "/projects/") { - files = append(files, shared.CollectFiles(root, isChatFile)...) + files = append(files, agentdata.CollectFiles(root, isChatFile)...) } } sort.Strings(files) - files = shared.FilterFiles(shared.UniqueStrings(files), filter) + files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) entries := make([]usage.Entry, 0) for _, file := range files { @@ -28,7 +30,7 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) } entries = append(entries, fileEntries...) } - shared.SortEntries(entries) + usageprovider.SortEntries(entries) return entries, nil } @@ -46,49 +48,49 @@ func isChatFile(path string) bool { } func parseChatFile(path string) ([]usage.Entry, error) { - lines, err := shared.ReadJSONLines(path, `"usageMetadata"`) + lines, err := agentdata.ReadJSONLines(path, `"usageMetadata"`) if err != nil { return nil, err } project := project(path) - fallback := shared.FileModifiedTime(path) + fallback := agentdata.FileModifiedTime(path) entries := make([]usage.Entry, 0) for _, line := range lines { record := line.Value - if shared.StringField(record, "type") != "assistant" { + if agentdata.StringField(record, "type") != "assistant" { continue } - meta := shared.ObjectAt(record["usageMetadata"]) + meta := agentdata.ObjectAt(record["usageMetadata"]) if meta == nil { continue } - timestamp, ok := shared.ParseTimestamp(record["timestamp"]) + timestamp, ok := agentdata.ParseTimestamp(record["timestamp"]) if !ok { timestamp = fallback } - sessionID := shared.StringField(record, "sessionId") + sessionID := agentdata.StringField(record, "sessionId") if sessionID == "" { sessionID = project + "-" + strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) } - model := shared.StringField(record, "model") + model := agentdata.StringField(record, "model") if model == "" { model = "unknown" } tokens := usage.TokenUsage{ - InputTokens: shared.UintField(meta, "promptTokenCount"), - OutputTokens: shared.UintField(meta, "candidatesTokenCount"), - CacheReadInputTokens: shared.UintField(meta, "cachedContentTokenCount"), - ReasoningOutputTokens: shared.UintField(meta, + InputTokens: agentdata.UintField(meta, "promptTokenCount"), + OutputTokens: agentdata.UintField(meta, "candidatesTokenCount"), + CacheReadInputTokens: agentdata.UintField(meta, "cachedContentTokenCount"), + ReasoningOutputTokens: agentdata.UintField(meta, "thoughtsTokenCount", ), } - tokens = shared.ApplyTotalFallback(tokens, shared.UintField(meta, "totalTokenCount")) - if !shared.NonZero(tokens) { + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(meta, "totalTokenCount")) + if !usageprovider.NonZero(tokens) { continue } - entry := shared.BaseEntry(usage.ProviderQwen, timestamp, "qwen", project, sessionID, model, "Qwen", tokens) - shared.SetSource(&entry, path, line.Line, line.Start, line.End) - entry.ID = shared.StableEntryID(entry) + entry := usageprovider.BaseEntry(usage.ProviderQwen, timestamp, "qwen", project, sessionID, model, "Qwen", tokens) + usageprovider.SetSource(&entry, path, line.Line, line.Start, line.End) + entry.ID = usageprovider.StableEntryID(entry) entries = append(entries, entry) } return entries, nil diff --git a/internal/provider/qwen/provider.go b/internal/provider/qwen/provider.go index c133531..c20fcb3 100644 --- a/internal/provider/qwen/provider.go +++ b/internal/provider/qwen/provider.go @@ -1,19 +1,18 @@ package qwen import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) // Provider loads Qwen usage entries. -type Provider struct{ shared.Base } +type Provider struct{ usageprovider.Base } var _ usageprovider.Provider = Provider{} // WithPaths returns a Qwen provider configured with data roots. func (p Provider) WithPaths(paths []string) usageprovider.Provider { - p.Base = shared.NewBase(paths) + p.Base = usageprovider.NewBase(paths) return p } @@ -33,5 +32,5 @@ func (p Provider) Entries() ([]usage.Entry, error) { if err != nil { return nil, err } - return shared.SortEntriesByTimestampDesc(entries), nil + return usageprovider.SortEntriesByTimestampDesc(entries), nil } diff --git a/internal/provider/qwen/provider_test.go b/internal/provider/qwen/provider_test.go index 1b2ddb4..13b7fec 100644 --- a/internal/provider/qwen/provider_test.go +++ b/internal/provider/qwen/provider_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/provider/shared" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -14,11 +14,11 @@ func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() path := filepath.Join(dir, "projects", "project-a", "chats", "chat-a.jsonl") - shared.WriteFile(t, path, `{"type":"assistant","timestamp":"2026-01-02T00:00:00.000Z","sessionId":"session-a","model":"qwen3-coder","usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":20,"thoughtsTokenCount":5,"cachedContentTokenCount":3,"totalTokenCount":38}}`+"\n") + providertest.WriteFile(t, path, `{"type":"assistant","timestamp":"2026-01-02T00:00:00.000Z","sessionId":"session-a","model":"qwen3-coder","usageMetadata":{"promptTokenCount":10,"candidatesTokenCount":20,"thoughtsTokenCount":5,"cachedContentTokenCount":3,"totalTokenCount":38}}`+"\n") return Provider{}.WithPaths([]string{dir}).Entries() }() - shared.AssertSingleEntry(t, entries, err, shared.WantEntry{ + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderQwen, Model: "qwen3-coder", SessionID: "session-a", diff --git a/internal/provider/shared/base.go b/internal/provider/shared/base.go deleted file mode 100644 index 70fd2bc..0000000 --- a/internal/provider/shared/base.go +++ /dev/null @@ -1,37 +0,0 @@ -package shared - -import ( - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -// Base carries the scan configuration every provider needs: where to look and -// which files are already ingested. Providers embed it so the accessors and -// the WithPaths/WithFileFilter plumbing exist in one place instead of once per -// provider. -type Base struct { - paths []string - filter usage.FileFilter -} - -// NewBase returns a Base scanning the given data roots. -func NewBase(paths []string) Base { - return Base{paths: append([]string{}, paths...)} -} - -// Paths returns the data roots to scan. -func (b Base) Paths() []string { return b.paths } - -// Filter returns the file filter, or nil when every file must be parsed. -func (b Base) Filter() usage.FileFilter { return b.filter } - -// WithPathsSet returns a copy scanning the given data roots. -func (b Base) WithPathsSet(paths []string) Base { - b.paths = append([]string{}, paths...) - return b -} - -// WithFilterSet returns a copy that skips the source files filter rejects. -func (b Base) WithFilterSet(filter usage.FileFilter) Base { - b.filter = filter - return b -} diff --git a/internal/provider/shared/entries.go b/internal/provider/shared/entries.go deleted file mode 100644 index 942d357..0000000 --- a/internal/provider/shared/entries.go +++ /dev/null @@ -1,16 +0,0 @@ -package shared - -import ( - "sort" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -// SortEntriesByTimestampDesc orders entries newest first, the order every -// provider returns them in. -func SortEntriesByTimestampDesc(entries []usage.Entry) []usage.Entry { - sort.Slice(entries, func(i, j int) bool { - return entries[i].Timestamp.After(entries[j].Timestamp) - }) - return entries -} diff --git a/internal/provider/shared/sqlite.go b/internal/provider/shared/sqlite.go deleted file mode 100644 index 391aae4..0000000 --- a/internal/provider/shared/sqlite.go +++ /dev/null @@ -1,84 +0,0 @@ -package shared - -import ( - "database/sql" - "fmt" - "math" - "strconv" - "strings" - - _ "modernc.org/sqlite" -) - -func OpenSQLite(path string) (*sql.DB, error) { - db, err := sql.Open("sqlite", path) - if err != nil { - return nil, err - } - db.SetMaxOpenConns(1) - db.SetMaxIdleConns(1) - if err := db.Ping(); err != nil { - _ = db.Close() - return nil, err - } - return db, nil -} - -func SqlString(value any) string { - switch typed := value.(type) { - case nil: - return "" - case string: - return strings.TrimSpace(typed) - case []byte: - return strings.TrimSpace(string(typed)) - case int64: - return strconv.FormatInt(typed, 10) - case float64: - if !isFinite(typed) { - return "" - } - return strconv.FormatFloat(typed, 'f', -1, 64) - default: - return strings.TrimSpace(fmt.Sprint(typed)) - } -} - -func SqlUint(value any) uint64 { - switch typed := value.(type) { - case nil: - return 0 - case int64: - if typed > 0 { - return uint64(typed) - } - case int: - if typed > 0 { - return uint64(typed) - } - case float64: - return floatToUint(typed) - case []byte: - return UintValue(string(typed)) - case string: - return UintValue(typed) - } - return 0 -} - -func SqlFloat(value any) (float64, bool) { - switch typed := value.(type) { - case nil: - return 0, false - case float64: - return typed, !math.IsNaN(typed) && !math.IsInf(typed, 0) - case int64: - return float64(typed), true - case []byte: - return floatValue(string(typed)) - case string: - return floatValue(typed) - default: - return 0, false - } -} diff --git a/internal/provider/shared/sqlitefiles.go b/internal/provider/shared/sqlitefiles.go deleted file mode 100644 index 0344a4f..0000000 --- a/internal/provider/shared/sqlitefiles.go +++ /dev/null @@ -1,65 +0,0 @@ -package shared - -import ( - "bytes" - "database/sql" - "encoding/json" - "os" - "path/filepath" - "sort" -) - -func ExistingSQLiteFile(path string) bool { - info, err := os.Stat(path) - return err == nil && !info.IsDir() -} -func DecodeJSONObjectString(data string) map[string]any { - decoder := json.NewDecoder(bytes.NewReader([]byte(data))) - decoder.UseNumber() - var record map[string]any - if err := decoder.Decode(&record); err != nil { - return nil - } - return record -} -func SqliteDBPaths(paths []string, defaultFile string, extraNames func(string) bool) []string { - dbPaths := make([]string, 0) - for _, root := range paths { - info, err := os.Stat(root) - if err != nil { - continue - } - if !info.IsDir() { - if filepath.Base(root) == defaultFile || extraNames != nil && extraNames(filepath.Base(root)) { - dbPaths = append(dbPaths, root) - } - continue - } - candidate := filepath.Join(root, defaultFile) - if fileInfo, err := os.Stat(candidate); err == nil && !fileInfo.IsDir() { - dbPaths = append(dbPaths, candidate) - } - if extraNames == nil { - continue - } - entries, err := os.ReadDir(root) - if err != nil { - continue - } - for _, entry := range entries { - if entry.IsDir() || !extraNames(entry.Name()) { - continue - } - dbPaths = append(dbPaths, filepath.Join(root, entry.Name())) - } - } - sort.Strings(dbPaths) - return UniqueStrings(dbPaths) -} -func ScanAny(rows *sql.Rows, values ...*any) bool { - dest := make([]any, len(values)) - for i := range values { - dest[i] = values[i] - } - return rows.Scan(dest...) == nil -} diff --git a/internal/provider/shared/testing.go b/internal/providertest/providertest.go similarity index 87% rename from internal/provider/shared/testing.go rename to internal/providertest/providertest.go index 9e0541b..c0041fa 100644 --- a/internal/provider/shared/testing.go +++ b/internal/providertest/providertest.go @@ -1,4 +1,7 @@ -package shared +// Package providertest holds the fixtures and assertions provider packages +// share in their tests. It lives apart from usageprovider so no production +// package ever links in the testing package. +package providertest import ( "database/sql" @@ -6,6 +9,7 @@ import ( "path/filepath" "testing" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -68,7 +72,7 @@ func OpenTestSQLite(t *testing.T, path string) *sql.DB { if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { t.Fatal(err) } - db, err := OpenSQLite(path) + db, err := agentdb.OpenSQLite(path) if err != nil { t.Fatal(err) } diff --git a/internal/usageprovider/provider.go b/internal/usageprovider/provider.go index 857435a..21b6e63 100644 --- a/internal/usageprovider/provider.go +++ b/internal/usageprovider/provider.go @@ -1,6 +1,13 @@ package usageprovider -import "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +import ( + "runtime" + "sort" + "strconv" + "time" + + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) // Provider loads normalized usage entries for one local AI agent. type Provider interface { @@ -11,3 +18,133 @@ type Provider interface { // Entries loads normalized usage entries from the provider's own source. Entries() ([]usage.Entry, error) } + +// Base carries the scan configuration every provider needs: where to look and +// which files are already ingested. Providers embed it so the accessors and +// the WithPaths/WithFileFilter plumbing exist in one place instead of once per +// provider. +type Base struct { + paths []string + filter usage.FileFilter +} + +// NewBase returns a Base scanning the given data roots. +func NewBase(paths []string) Base { + return Base{paths: append([]string{}, paths...)} +} + +// Paths returns the data roots to scan. +func (b Base) Paths() []string { return b.paths } + +// Filter returns the file filter, or nil when every file must be parsed. +func (b Base) Filter() usage.FileFilter { return b.filter } + +// WithPathsSet returns a copy scanning the given data roots. +func (b Base) WithPathsSet(paths []string) Base { + b.paths = append([]string{}, paths...) + return b +} + +// WithFilterSet returns a copy that skips the source files filter rejects. +func (b Base) WithFilterSet(filter usage.FileFilter) Base { + b.filter = filter + return b +} + +// SortEntriesByTimestampDesc orders entries newest first, the order every +// provider returns them in. +func SortEntriesByTimestampDesc(entries []usage.Entry) []usage.Entry { + sort.Slice(entries, func(i, j int) bool { + return entries[i].Timestamp.After(entries[j].Timestamp) + }) + return entries +} + +func formatDate(timestamp time.Time) string { + return timestamp.In(time.Local).Format("2006-01-02") +} + +func TotalUsage(tokens usage.TokenUsage) uint64 { + return tokens.InputTokens + + tokens.OutputTokens + + tokens.CacheCreationInputTokens + + tokens.CacheReadInputTokens + + tokens.CachedInputTokens + + tokens.ReasoningOutputTokens +} + +func ApplyTotalFallback(tokens usage.TokenUsage, total uint64) usage.TokenUsage { + sum := TotalUsage(tokens) + if sum == 0 && total > 0 { + tokens.OutputTokens = total + tokens.TotalTokens = total + return tokens + } + if total > sum { + tokens.ReasoningOutputTokens += total - sum + tokens.TotalTokens = total + return tokens + } + if tokens.TotalTokens == 0 { + tokens.TotalTokens = sum + } + return tokens +} + +func NonZero(tokens usage.TokenUsage) bool { + return TotalUsage(tokens) > 0 || tokens.TotalTokens > 0 +} + +func BaseEntry(provider usage.Provider, timestamp time.Time, project, projectPath, sessionID, model, client string, tokens usage.TokenUsage) usage.Entry { + return usage.Entry{ + Provider: provider, + Timestamp: timestamp, + Date: formatDate(timestamp), + Project: project, + ProjectPath: projectPath, + SessionID: sessionID, + Model: model, + Language: usage.UnknownLanguage, + OS: usage.NormalizeOS(runtime.GOOS), + Client: client, + Usage: tokens, + } +} + +func SetSource(entry *usage.Entry, source string, line int, start, end int64) { + entry.SourceFile = source + entry.SourceLine = line + entry.SourceStart = start + entry.SourceEnd = end +} + +func StableEntryID(entry usage.Entry, extra ...string) string { + parts := []string{ + string(entry.Provider), + entry.SourceFile, + strconv.Itoa(entry.SourceLine), + entry.Timestamp.Format(time.RFC3339Nano), + entry.Project, + entry.ProjectPath, + entry.SessionID, + entry.Model, + strconv.FormatUint(entry.Usage.InputTokens, 10), + strconv.FormatUint(entry.Usage.OutputTokens, 10), + strconv.FormatUint(entry.Usage.CacheCreationInputTokens, 10), + strconv.FormatUint(entry.Usage.CacheReadInputTokens, 10), + strconv.FormatUint(entry.Usage.CachedInputTokens, 10), + strconv.FormatUint(entry.Usage.ReasoningOutputTokens, 10), + strconv.FormatUint(entry.Usage.TotalTokens, 10), + } + parts = append(parts, extra...) + return usage.StableID(parts...) +} + +func SortEntries(entries []usage.Entry) { + sort.Slice(entries, func(i, j int) bool { + if !entries[i].Timestamp.Equal(entries[j].Timestamp) { + return entries[i].Timestamp.Before(entries[j].Timestamp) + } + return entries[i].ID < entries[j].ID + }) +} From d3cdd69f6a68a5af671cfcd03804c1d4bc5b7555 Mon Sep 17 00:00:00 2001 From: akarineren Date: Fri, 31 Jul 2026 23:14:56 +0900 Subject: [PATCH 06/11] feat: enhance OpenCode session handling with model support and improve rejection logging --- internal/provider/opencode/loader.go | 111 ++++++++++++++++++-- internal/provider/opencode/loader_test.go | 117 +++++++++++++++++++++- internal/usageupload/uploader.go | 47 +++++++-- internal/usageupload/uploader_test.go | 68 +++++++++++++ 4 files changed, 323 insertions(+), 20 deletions(-) diff --git a/internal/provider/opencode/loader.go b/internal/provider/opencode/loader.go index 68cacfb..326881f 100644 --- a/internal/provider/opencode/loader.go +++ b/internal/provider/opencode/loader.go @@ -12,6 +12,7 @@ import ( "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" + "github.com/tokitoki-dev/tokitoki-cli/internal/langdetect" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" @@ -24,6 +25,7 @@ import ( type session struct { directory string version string + model string } type message struct { @@ -172,17 +174,33 @@ func loadDatabase(path string) ([]usage.Entry, error) { func querySessions(db *sql.DB) map[string]session { sessions := make(map[string]session) - rows, err := db.Query(`SELECT id, COALESCE(directory, ''), COALESCE(version, '') FROM session`) + // The model column arrived in a later OpenCode release. Asking for it on an + // older database fails the whole query, so fall back to the columns that + // have always been there rather than lose the directory too. + rows, err := db.Query(`SELECT id, COALESCE(directory, ''), COALESCE(version, ''), COALESCE(model, '') FROM session`) + if err != nil { + rows, err = db.Query(`SELECT id, COALESCE(directory, ''), COALESCE(version, ''), '' FROM session`) + } if err != nil { return sessions } defer rows.Close() for rows.Next() { - var id, directory, version string - if err := rows.Scan(&id, &directory, &version); err != nil { + var id, directory, version, model string + if err := rows.Scan(&id, &directory, &version, &model); err != nil { continue } - sessions[id] = session{directory: directory, version: version} + // The session's model is stored as a JSON object, unlike the message's + // flat modelID. + block := agentdata.DecodeJSONObjectString(model) + sessions[id] = session{ + directory: directory, + version: version, + model: agentdata.FirstNonEmpty( + agentdata.StringField(block, "id"), + agentdata.StringField(block, "modelID"), + ), + } } return sessions } @@ -317,12 +335,23 @@ func growth(current, previous uint64) uint64 { return current - previous } +// modelName resolves which model answered a turn. OpenCode records it three +// ways depending on the message: assistant messages carry a flat modelID, user +// messages nest it under "model", and the session holds the one it was started +// with. An entry keeps its tokens even when none of them is set — losing a +// turn's usage is worse than not knowing its model. +func modelName(record map[string]any, session session) string { + return agentdata.FirstNonEmpty( + agentdata.StringField(record, "modelID"), + agentdata.StringField(agentdata.ObjectAt(record["model"]), "modelID"), + agentdata.StringField(agentdata.ObjectAt(record["model"]), "id"), + session.model, + ) +} + func newEntry(source string, session session, message *message, tokens usage.TokenUsage) (usage.Entry, bool) { record := message.data - model := agentdata.StringField(record, "modelID") - if model == "" { - return usage.Entry{}, false - } + model := modelName(record, session) timestamp := time.UnixMilli(message.created).UTC() if parsed, ok := agentdata.ParseTimestamp(agentdata.ObjectAt(record["time"])["created"]); ok { @@ -349,14 +378,71 @@ func newEntry(source string, session session, message *message, tokens usage.Tok entry.ID = usageprovider.StableEntryID(entry) } + candidates := make([]langdetect.Candidate, 0) for _, part := range message.parts { for _, change := range partChanges(part, cwd) { entry.ApplyFileChange(change) } + candidates = append(candidates, languageCandidates(part, cwd)...) } + entry.Language = usage.NormalizeLanguage(langdetect.Dominant(candidates)) return entry, true } +// Weights for the paths a turn touched. Writing a file says far more about +// what is being worked on than reading one, and a path merely mentioned in a +// shell command says least of all. +const ( + writeWeight = 4 + readWeight = 2 + textWeight = 1 +) + +// languageCandidates collects the file paths one part touched, weighted by how +// strongly each says what language the turn was spent on. +func languageCandidates(part map[string]any, cwd string) []langdetect.Candidate { + if agentdata.StringField(part, "type") != "tool" { + return nil + } + state := agentdata.ObjectAt(part["state"]) + if agentdata.StringField(state, "status") != "completed" { + return nil + } + input := agentdata.ObjectAt(state["input"]) + + switch tool := agentdata.StringField(part, "tool"); tool { + case "edit", "write", "patch", "apply_patch": + candidates := make([]langdetect.Candidate, 0) + for _, change := range partChanges(part, cwd) { + candidates = append(candidates, langdetect.Candidate{Path: change.Path, Weight: writeWeight}) + } + return candidates + case "read": + path := usage.ResolvePath(cwd, agentdata.StringField(input, "filePath")) + if path == "" { + return nil + } + return []langdetect.Candidate{{Path: path, Weight: readWeight}} + case "bash": + // A shell command names the files it runs against; they are a weak but + // real signal when a turn neither read nor wrote anything. + return pathCandidates(agentdata.StringField(input, "command"), textWeight) + case "glob", "grep": + return pathCandidates(agentdata.StringField(input, "pattern"), textWeight) + default: + return nil + } +} + +func pathCandidates(text string, weight int) []langdetect.Candidate { + paths := langdetect.PathsFromText(text) + candidates := make([]langdetect.Candidate, 0, len(paths)) + for _, path := range paths { + candidates = append(candidates, langdetect.Candidate{Path: path, Weight: weight}) + } + return candidates +} + // partChanges extracts the file edits one message part performed. // Only tool calls count: a "patch" part is a git snapshot of the whole // worktree, which lists files nobody edited (.DS_Store, build output) and @@ -420,7 +506,14 @@ func writeChanges(input, metadata map[string]any, cwd string) []usage.FileChange if path == "" { return nil } - return []usage.FileChange{{Path: path, LinesAdded: usage.CountLines(agentdata.StringField(input, "content"))}} + + // Every line of the written content is an added line. Overwriting also + // displaces the file's previous lines, but the old content is not recorded + // anywhere, so those stay uncounted rather than guessed at. + return []usage.FileChange{{ + Path: path, + LinesAdded: usage.CountLines(agentdata.StringField(input, "content")), + }} } // applyPatchChanges prefers the per-file summary OpenCode attaches and diff --git a/internal/provider/opencode/loader_test.go b/internal/provider/opencode/loader_test.go index 0cb35db..7399253 100644 --- a/internal/provider/opencode/loader_test.go +++ b/internal/provider/opencode/loader_test.go @@ -33,7 +33,7 @@ func writeOpenCodeDBAt(t *testing.T, path string, sessions, messages, parts []op defer db.Close() schema := []string{ - `CREATE TABLE session (id TEXT PRIMARY KEY, directory TEXT, version TEXT)`, + `CREATE TABLE session (id TEXT PRIMARY KEY, directory TEXT, version TEXT, model TEXT)`, `CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT, time_created INTEGER, data TEXT)`, `CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, session_id TEXT, time_created INTEGER, data TEXT)`, } @@ -43,7 +43,7 @@ func writeOpenCodeDBAt(t *testing.T, path string, sessions, messages, parts []op } } for _, row := range sessions { - if _, err := db.Exec(`INSERT INTO session (id, directory, version) VALUES (?, ?, ?)`, row.id, row.sessionID, row.data); err != nil { + if _, err := db.Exec(`INSERT INTO session (id, directory, model) VALUES (?, ?, ?)`, row.id, row.sessionID, row.data); err != nil { t.Fatal(err) } } @@ -387,3 +387,116 @@ func TestOpenCodeReadsEveryChannelDatabase(t *testing.T) { t.Fatalf("len(entries) = %d, want one per channel database", len(entries)) } } + +// TestDetectsLanguageFromTouchedFiles covers the language signal: a turn is +// spent on whatever it edited, then on what it read, and only weakly on what a +// shell command happened to name. +func TestDetectsLanguageFromTouchedFiles(t *testing.T) { + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []openCodeRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","modelID":"gpt-5","path":{"cwd":"/repo/demo"},"tokens":{"input":10,"output":5}}`}, + {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"role":"assistant","modelID":"gpt-5","path":{"cwd":"/repo/demo"},"tokens":{"input":20,"output":5}}`}, + {id: "msg-3", sessionID: "ses-1", created: 3000, data: `{"role":"assistant","modelID":"gpt-5","path":{"cwd":"/repo/demo"},"tokens":{"input":30,"output":5}}`}, + }, + []openCodeRow{ + // A single write outweighs the two reads of another language. + {id: "prt-1", sessionID: "msg-1", data: `{"type":"tool","tool":"write","state":{"status":"completed","input":{"filePath":"main.go","content":"package main"}}}`}, + {id: "prt-2", sessionID: "msg-1", data: `{"type":"tool","tool":"read","state":{"status":"completed","input":{"filePath":"a.py"}}}`}, + {id: "prt-3", sessionID: "msg-1", data: `{"type":"tool","tool":"read","state":{"status":"completed","input":{"filePath":"b.py"}}}`}, + // Reads alone still name the language. + {id: "prt-4", sessionID: "msg-2", data: `{"type":"tool","tool":"read","state":{"status":"completed","input":{"filePath":"lib.rs"}}}`}, + // A turn that touched nothing has no language to report. + {id: "prt-5", sessionID: "msg-3", data: `{"type":"tool","tool":"bash","state":{"status":"completed","input":{"command":"echo hello"}}}`}, + }) + + entries, err := loadDatabase(path) + if err != nil { + t.Fatal(err) + } + usageprovider.SortEntries(entries) + if len(entries) != 3 { + t.Fatalf("len(entries) = %d, want 3", len(entries)) + } + if entries[0].Language != "Go" { + t.Fatalf("language = %q, want Go (the written file outweighs the read ones)", entries[0].Language) + } + if entries[1].Language != "Rust" { + t.Fatalf("language = %q, want Rust (from the read file)", entries[1].Language) + } + if entries[2].Language != usage.UnknownLanguage { + t.Fatalf("language = %q, want %q for a turn that touched no file", + entries[2].Language, usage.UnknownLanguage) + } +} + +// TestOpenCodeReadsSessionsWithoutModelColumn covers databases written by an +// OpenCode release from before the session gained a model column. Asking for a +// column that does not exist fails the whole query, and losing the session row +// would take the project directory down with it. +func TestOpenCodeReadsSessionsWithoutModelColumn(t *testing.T) { + path := filepath.Join(t.TempDir(), "opencode.db") + db, err := agentdb.OpenSQLite(path) + if err != nil { + t.Fatal(err) + } + for _, statement := range []string{ + `CREATE TABLE session (id TEXT PRIMARY KEY, directory TEXT, version TEXT)`, + `CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT, time_created INTEGER, data TEXT)`, + `CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, session_id TEXT, time_created INTEGER, data TEXT)`, + `INSERT INTO session (id, directory) VALUES ('ses-1', '/repo/legacy')`, + `INSERT INTO message (id, session_id, time_created, data) VALUES ('msg-1', 'ses-1', 1000, + '{"role":"assistant","modelID":"gpt-5","tokens":{"input":10,"output":5}}')`, + } { + if _, err := db.Exec(statement); err != nil { + t.Fatal(err) + } + } + if err := db.Close(); err != nil { + t.Fatal(err) + } + + entries, err := loadDatabase(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + if entries[0].Project != "legacy" { + t.Fatalf("project = %q, want legacy (the session row must survive)", entries[0].Project) + } + if entries[0].Model != "gpt-5" { + t.Fatalf("model = %q, want gpt-5", entries[0].Model) + } +} + +// TestOpenCodeFallsBackToSessionModel covers a turn whose message does not name +// the model. The tokens must still be billed — losing a turn's usage is worse +// than not knowing which model produced it. +func TestOpenCodeFallsBackToSessionModel(t *testing.T) { + path := writeOpenCodeDB(t, + []openCodeRow{{id: "ses-1", sessionID: "/repo/demo", + data: `{"id":"claude-sonnet-4","providerID":"anthropic","variant":"default"}`}}, + []openCodeRow{ + // No modelID anywhere on the message. + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"role":"assistant","tokens":{"input":10,"output":5}}`}, + // The nested form user messages use. + {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"role":"assistant","model":{"modelID":"gpt-5","providerID":"openai"},"tokens":{"input":20,"output":5}}`}, + }, nil) + + entries, err := loadDatabase(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("len(entries) = %d, want 2 (a missing model must not drop the tokens)", len(entries)) + } + usageprovider.SortEntries(entries) + if entries[0].Model != "claude-sonnet-4" { + t.Fatalf("model = %q, want the session's model", entries[0].Model) + } + if entries[1].Model != "gpt-5" { + t.Fatalf("model = %q, want the nested modelID", entries[1].Model) + } +} diff --git a/internal/usageupload/uploader.go b/internal/usageupload/uploader.go index 3fc40d2..5b73a31 100644 --- a/internal/usageupload/uploader.go +++ b/internal/usageupload/uploader.go @@ -10,6 +10,7 @@ import ( "errors" "fmt" "io" + "log/slog" "net/http" "os" "path/filepath" @@ -159,19 +160,30 @@ func SyncPending(ctx context.Context, settings agent.Settings, db *usagedb.DB) e } } - // Duplicate + Rejected: both are "acknowledged and won't be processed again". - // Duplicates are events we already uploaded. Rejected are events that failed - // validation. Both should stop querying; we mark them as uploaded to clear them. - ackd := append([]string{}, response.Duplicate...) - for _, r := range response.Rejected { - if r.ID != "" { - ackd = append(ackd, r.ID) + // Duplicates are events the server already has: nothing was lost, so they + // count as uploaded. + if len(response.Duplicate) > 0 { + if err := db.MarkEventsUploaded(response.Duplicate); err != nil { + return err } } - if len(ackd) > 0 { - if err := db.MarkEventsUploaded(ackd); err != nil { + + // Rejected events are data the server threw away. Recording them as + // uploaded hid that: a server that refused every event from a whole + // provider looked exactly like a successful sync, and the reason it gave + // was discarded. They still are not retried — the queue keeps them as + // rejected, with the server's reason, so the loss is visible. + if len(response.Rejected) > 0 { + reasons := make(map[string]string, len(response.Rejected)) + for _, rejected := range response.Rejected { + if rejected.ID != "" { + reasons[rejected.ID] = rejected.Reason + } + } + if err := db.MarkEventsRejected(reasons); err != nil { return err } + logRejections(response.Rejected) } // Sanity check: server must acknowledge every event as accepted/duplicate/rejected. @@ -184,6 +196,23 @@ func SyncPending(ctx context.Context, settings agent.Settings, db *usagedb.DB) e } } +// logRejections reports what the server refused. Rejections repeat: a server +// that rejects one event from a provider rejects all of them, so the reasons +// are counted rather than printed one per event. +func logRejections(rejected []Reject) { + counts := make(map[string]int, len(rejected)) + for _, r := range rejected { + reason := strings.TrimSpace(r.Reason) + if reason == "" { + reason = "no reason given" + } + counts[reason]++ + } + for reason, count := range counts { + slog.Warn("usage events rejected by server", "reason", reason, "events", count) + } +} + func eventIDs(events []usage.Entry) []string { ids := make([]string, 0, len(events)) for _, event := range events { diff --git a/internal/usageupload/uploader_test.go b/internal/usageupload/uploader_test.go index ea2a27c..83fdeb0 100644 --- a/internal/usageupload/uploader_test.go +++ b/internal/usageupload/uploader_test.go @@ -5,11 +5,13 @@ import ( "encoding/json" "net/http" "net/http/httptest" + "path/filepath" "testing" "time" "github.com/tokitoki-dev/tokitoki-cli/internal/agent" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usagedb" ) func TestDefaultServerURLIsProduction(t *testing.T) { @@ -108,3 +110,69 @@ func TestRelativeFilesStripsMachineLayout(t *testing.T) { t.Fatal("relativeFiles(nil) should be nil") } } + +// TestSyncPendingKeepsRejectionsVisible pins the queue's memory of what the +// server threw away. Recording rejections as uploaded once hid a server that +// refused every event from a whole provider: the sync looked clean and the +// data simply never appeared. +func TestSyncPendingKeepsRejectionsVisible(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + _ = json.NewEncoder(w).Encode(Response{ + OK: true, + Accepted: []string{"keep-me"}, + Rejected: []Reject{{ID: "drop-me", Reason: "AI provider must be claude or codex"}}, + }) + })) + defer server.Close() + t.Setenv(BaseURLEnv, server.URL) + + db, err := usagedb.Open(filepath.Join(t.TempDir(), "usage.db")) + if err != nil { + t.Fatal(err) + } + defer db.Close() + + now := time.Now().UTC() + if _, err := db.InsertEvents([]usage.Entry{ + {ID: "keep-me", Provider: usage.ProviderClaude, Timestamp: now, Project: "demo"}, + {ID: "drop-me", Provider: usage.ProviderOpenCode, Timestamp: now, Project: "demo"}, + }); err != nil { + t.Fatal(err) + } + + if err := SyncPending(context.Background(), agent.Settings{APIKey: "test-key"}, db); err != nil { + t.Fatal(err) + } + + // Neither event may be retried: one landed, the other was refused for good. + pending, err := db.PendingEvents(now.Add(24*time.Hour), 0) + if err != nil { + t.Fatal(err) + } + if len(pending) != 0 { + t.Fatalf("pending = %+v, want none", pending) + } + + // Pruning removes uploaded events and leaves rejected ones behind, so what + // survives says which status each event carries. + pruned, err := db.PruneUploaded(now.Add(24 * time.Hour)) + if err != nil { + t.Fatal(err) + } + if pruned != 1 { + t.Fatalf("pruned = %d, want 1 (only the accepted event is uploaded)", pruned) + } + // The rejected event is still queued as rejected rather than gone or retried. + if _, err := db.InsertEvents([]usage.Entry{ + {ID: "drop-me", Provider: usage.ProviderOpenCode, Timestamp: now, Project: "demo"}, + }); err != nil { + t.Fatal(err) + } + pending, err = db.PendingEvents(now.Add(48*time.Hour), 0) + if err != nil { + t.Fatal(err) + } + if len(pending) != 0 { + t.Fatalf("pending = %+v, want none — a rejected event must not come back", pending) + } +} From 3a931c20ac0e5c97636364da5927d09bb98f2682 Mon Sep 17 00:00:00 2001 From: akarineren Date: Sat, 1 Aug 2026 11:12:15 +0900 Subject: [PATCH 07/11] feat: implement Kilo provider with database handling and entry loading functionality --- internal/agentdb/agentdb.go | 5 +- internal/provider/kilo/loader.go | 369 +++++++++++++++++++--- internal/provider/kilo/loader_test.go | 299 ++++++++++++++++++ internal/provider/kilo/provider_test.go | 29 +- internal/provider/opencode/loader_test.go | 12 +- internal/providertest/providertest.go | 15 +- 6 files changed, 658 insertions(+), 71 deletions(-) create mode 100644 internal/provider/kilo/loader_test.go diff --git a/internal/agentdb/agentdb.go b/internal/agentdb/agentdb.go index 4bde412..3a023f8 100644 --- a/internal/agentdb/agentdb.go +++ b/internal/agentdb/agentdb.go @@ -35,8 +35,11 @@ func floatValue(value any) (float64, bool) { } } +// OpenSQLite opens an agent's database for scanning. Read-only is not a +// preference but a guarantee: these files belong to running agents, and a +// scanner that can take write locks on them can also corrupt them. func OpenSQLite(path string) (*sql.DB, error) { - db, err := sql.Open("sqlite", path) + db, err := sql.Open("sqlite", "file:"+filepath.ToSlash(path)+"?mode=ro") if err != nil { return nil, err } diff --git a/internal/provider/kilo/loader.go b/internal/provider/kilo/loader.go index e53f20e..5423eb9 100644 --- a/internal/provider/kilo/loader.go +++ b/internal/provider/kilo/loader.go @@ -1,25 +1,58 @@ package kilo import ( - "errors" - "os" + "database/sql" + "strings" + "time" "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" + "github.com/tokitoki-dev/tokitoki-cli/internal/langdetect" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) +// Kilo's engine is a fork of OpenCode, so the layout is the same one SQLite +// database: sessions own messages, messages own parts, tokens live on the +// message and file edits live in the parts. The fork changed what the numbers +// mean, though: a Kilo message reports what its one API request consumed, not +// OpenCode's running context totals, so entries here bill each message as +// reported with no subtraction. The session rows sum their messages' fields +// verbatim, which is the fork's own word on that semantics. +// +// The Kilo CLI and the Kilo VSCode extension run this same engine against +// this same database, and no table records which of the two wrote a session. +// Both are therefore one agent: "kilo". +type session struct { + directory string + model string +} + +type message struct { + id string + sessionID string + created int64 + data map[string]any + parts []map[string]any +} + func loadEntries(paths []string) ([]usage.Entry, error) { - dbPaths := agentdb.SqliteDBPaths(paths, "kilo.db", nil) - entries := make([]usage.Entry, 0) - for _, dbPath := range dbPaths { + entriesByID := make(map[string]usage.Entry) + for _, dbPath := range agentdb.SqliteDBPaths(paths, "kilo.db", nil) { dbEntries, err := loadDatabase(dbPath) if err != nil { return nil, err } - entries = append(entries, dbEntries...) + for _, entry := range dbEntries { + if _, exists := entriesByID[entry.ID]; !exists { + entriesByID[entry.ID] = entry + } + } + } + entries := make([]usage.Entry, 0, len(entriesByID)) + for _, entry := range entriesByID { + entries = append(entries, entry) } usageprovider.SortEntries(entries) return entries, nil @@ -27,66 +60,318 @@ func loadEntries(paths []string) ([]usage.Entry, error) { func loadDatabase(path string) ([]usage.Entry, error) { db, err := agentdb.OpenSQLite(path) - if errors.Is(err, os.ErrNotExist) { - return nil, nil - } if err != nil { return nil, err } defer db.Close() - rows, err := db.Query("SELECT id, session_id, data FROM message") - if err != nil { + messages := queryMessages(db) + if len(messages) == 0 { return nil, nil } + sessions := querySessions(db) + attachParts(db, messages) + + entries := make([]usage.Entry, 0, len(messages)) + for _, message := range messages { + if entry, ok := newEntry(path, sessions[message.sessionID], message); ok { + entries = append(entries, entry) + } + } + return entries, nil +} + +func querySessions(db *sql.DB) map[string]session { + sessions := make(map[string]session) + // Early Kilo builds predate the model column; losing it must not lose the + // directory too. + rows, err := db.Query(`SELECT id, COALESCE(directory, ''), COALESCE(model, '') FROM session`) + if err != nil { + rows, err = db.Query(`SELECT id, COALESCE(directory, ''), '' FROM session`) + } + if err != nil { + return sessions + } + defer rows.Close() + for rows.Next() { + var id, directory, model string + if err := rows.Scan(&id, &directory, &model); err != nil { + continue + } + // The session's model is a JSON object, unlike the message's flat + // modelID: {"id":"kilo-auto/free","providerID":"kilo",...}. + block := agentdata.DecodeJSONObjectString(model) + sessions[id] = session{ + directory: directory, + model: agentdata.FirstNonEmpty( + agentdata.StringField(block, "id"), + agentdata.StringField(block, "modelID"), + ), + } + } + return sessions +} + +func queryMessages(db *sql.DB) []*message { + rows, err := db.Query(`SELECT id, session_id, time_created, CAST(data AS TEXT) FROM message`) + if err != nil { + return nil + } defer rows.Close() - entries := make([]usage.Entry, 0) + messages := make([]*message, 0) for rows.Next() { - var rowID, rowSessionID, data string - if err := rows.Scan(&rowID, &rowSessionID, &data); err != nil { + var id, sessionID, data string + var created int64 + if err := rows.Scan(&id, &sessionID, &created, &data); err != nil { continue } - if entry, ok := messageEntry(path, rowID, rowSessionID, data); ok { - entries = append(entries, entry) + record := agentdata.DecodeJSONObjectString(data) + if record == nil { + continue } + messages = append(messages, &message{ + id: agentdata.FirstNonEmpty(agentdata.StringField(record, "id"), id), + sessionID: agentdata.FirstNonEmpty(agentdata.StringField(record, "sessionID"), sessionID), + created: created, + data: record, + }) } - return entries, rows.Err() + return messages } -func messageEntry(dbPath, rowID, rowSessionID, data string) (usage.Entry, bool) { - record := agentdata.DecodeJSONObjectString(data) - if record == nil || agentdata.StringField(record, "role") != "assistant" { - return usage.Entry{}, false +func attachParts(db *sql.DB, messages []*message) { + byID := make(map[string][]*message, len(messages)) + for _, message := range messages { + byID[message.id] = append(byID[message.id], message) } - tokenBlock := agentdata.ObjectAt(record["tokens"]) - if tokenBlock == nil { - return usage.Entry{}, false + + rows, err := db.Query(`SELECT message_id, CAST(data AS TEXT) FROM part ORDER BY time_created ASC, id ASC`) + if err != nil { + return } - cache := agentdata.ObjectAt(tokenBlock["cache"]) - tokens := usage.TokenUsage{ - InputTokens: agentdata.UintField(tokenBlock, "input"), - OutputTokens: agentdata.UintField(tokenBlock, "output"), + defer rows.Close() + for rows.Next() { + var messageID, data string + if err := rows.Scan(&messageID, &data); err != nil { + continue + } + targets := byID[messageID] + if len(targets) == 0 { + continue + } + record := agentdata.DecodeJSONObjectString(data) + if record == nil { + continue + } + for _, target := range targets { + target.parts = append(target.parts, record) + } + } +} + +// tokens bills a message for what its API request consumed, exactly as +// reported. User messages carry no token block and drop out as all-zero. +func tokens(record map[string]any) usage.TokenUsage { + block := agentdata.ObjectAt(record["tokens"]) + if block == nil { + return usage.TokenUsage{} + } + cache := agentdata.ObjectAt(block["cache"]) + billed := usage.TokenUsage{ + InputTokens: agentdata.UintField(block, "input"), + OutputTokens: agentdata.UintField(block, "output"), + ReasoningOutputTokens: agentdata.UintField(block, "reasoning"), CacheCreationInputTokens: agentdata.UintField(cache, "write"), CacheReadInputTokens: agentdata.UintField(cache, "read"), - ReasoningOutputTokens: agentdata.UintField(tokenBlock, "reasoning"), } - tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(tokenBlock, "total")) - if !usageprovider.NonZero(tokens) { + return usageprovider.ApplyTotalFallback(billed, agentdata.UintField(block, "total")) +} + +// modelName resolves which model answered a turn: the message's flat modelID, +// its nested model object, or the model the session was started with. An +// entry keeps its tokens even when none of them is set — losing a turn's +// usage is worse than not knowing its model. +func modelName(record map[string]any, session session) string { + return agentdata.FirstNonEmpty( + agentdata.StringField(record, "modelID"), + agentdata.StringField(agentdata.ObjectAt(record["model"]), "modelID"), + agentdata.StringField(agentdata.ObjectAt(record["model"]), "id"), + session.model, + ) +} + +func newEntry(source string, session session, message *message) (usage.Entry, bool) { + record := message.data + billed := tokens(record) + if !usageprovider.NonZero(billed) { return usage.Entry{}, false } - model := agentdata.StringField(record, "modelID") - if model == "" { - return usage.Entry{}, false + + timestamp := time.UnixMilli(message.created).UTC() + if parsed, ok := agentdata.ParseTimestamp(agentdata.ObjectAt(record["time"])["created"]); ok { + timestamp = parsed } - timestamp, ok := agentdata.ParseTimestamp(agentdata.ObjectAt(record["time"])["created"]) - if !ok { - return usage.Entry{}, false + + // The message records the directory the agent actually ran in; the session + // directory covers messages that predate the field. + cwd := agentdata.FirstNonEmpty( + agentdata.StringField(agentdata.ObjectAt(record["path"]), "cwd"), + agentdata.StringField(agentdata.ObjectAt(record["path"]), "root"), + session.directory, + ) + project, projectPath := usage.UnknownProject, "" + if path, name, ok := usage.ProjectFromCWD(cwd); ok { + project, projectPath = name, path + } + + sessionID := agentdata.FirstNonEmpty(message.sessionID, usage.UnknownProject) + entry := usageprovider.BaseEntry(usage.ProviderKilo, timestamp, project, projectPath, sessionID, modelName(record, session), "Kilo", billed) + usageprovider.SetSource(&entry, source, 0, 0, 0) + entry.ID = stableMessageID(message.id) + if entry.ID == "" { + entry.ID = usageprovider.StableEntryID(entry) + } + + candidates := make([]langdetect.Candidate, 0) + for _, part := range message.parts { + for _, change := range partChanges(part, cwd) { + entry.ApplyFileChange(change) + } + candidates = append(candidates, languageCandidates(part, cwd)...) } - sessionID := agentdata.FirstNonEmpty(agentdata.StringField(record, "session_id"), rowSessionID, "unknown") - messageID := agentdata.FirstNonEmpty(agentdata.StringField(record, "id"), rowID) - entry := usageprovider.BaseEntry(usage.ProviderKilo, timestamp, "kilo", "Kilo", sessionID, model, "Kilo", tokens) - usageprovider.SetSource(&entry, dbPath, 0, 0, 0) - entry.ID = usageprovider.StableEntryID(entry, messageID) + entry.Language = usage.NormalizeLanguage(langdetect.Dominant(candidates)) return entry, true } + +// Weights for the paths a turn touched. Writing a file says far more about +// what is being worked on than reading one, and a path merely mentioned in a +// shell command says least of all. +const ( + writeWeight = 4 + readWeight = 2 + textWeight = 1 +) + +// languageCandidates collects the file paths one part touched, weighted by +// how strongly each says what language the turn was spent on. +func languageCandidates(part map[string]any, cwd string) []langdetect.Candidate { + if agentdata.StringField(part, "type") != "tool" { + return nil + } + state := agentdata.ObjectAt(part["state"]) + if agentdata.StringField(state, "status") != "completed" { + return nil + } + input := agentdata.ObjectAt(state["input"]) + + switch agentdata.StringField(part, "tool") { + case "edit", "write": + candidates := make([]langdetect.Candidate, 0) + for _, change := range partChanges(part, cwd) { + candidates = append(candidates, langdetect.Candidate{Path: change.Path, Weight: writeWeight}) + } + return candidates + case "read": + path := usage.ResolvePath(cwd, agentdata.StringField(input, "filePath")) + if path == "" { + return nil + } + return []langdetect.Candidate{{Path: path, Weight: readWeight}} + case "bash": + // A shell command names the files it runs against; they are a weak but + // real signal when a turn neither read nor wrote anything. + return pathCandidates(agentdata.StringField(input, "command"), textWeight) + case "glob", "grep": + return pathCandidates(agentdata.StringField(input, "pattern"), textWeight) + default: + return nil + } +} + +func pathCandidates(text string, weight int) []langdetect.Candidate { + paths := langdetect.PathsFromText(text) + candidates := make([]langdetect.Candidate, 0, len(paths)) + for _, path := range paths { + candidates = append(candidates, langdetect.Candidate{Path: path, Weight: weight}) + } + return candidates +} + +// partChanges extracts the file edits one message part performed. Only tool +// calls count: a "patch" part is a git snapshot of the whole worktree, which +// lists files nobody edited and carries no line counts, so treating it as +// agent work would invent writes. +// +// Kilo attaches the unified diff of every edit and write to the part +// (state.metadata.filediff.patch), so both tools are counted the same way: +// from the diff when it is there, from the tool's input when it is not. +func partChanges(part map[string]any, cwd string) []usage.FileChange { + if agentdata.StringField(part, "type") != "tool" { + return nil + } + state := agentdata.ObjectAt(part["state"]) + if agentdata.StringField(state, "status") != "completed" { + return nil + } + input := agentdata.ObjectAt(state["input"]) + metadata := agentdata.ObjectAt(state["metadata"]) + + tool := agentdata.StringField(part, "tool") + if tool != "edit" && tool != "write" { + return nil + } + + diff := agentdata.ObjectAt(metadata["filediff"]) + path := usage.ResolvePath(cwd, agentdata.FirstNonEmpty( + agentdata.StringField(diff, "file"), + agentdata.StringField(input, "filePath"), + agentdata.StringField(metadata, "filepath"), + )) + if path == "" { + return nil + } + + if added, removed, ok := countUnifiedDiff(agentdata.StringField(diff, "patch")); ok { + return []usage.FileChange{{Path: path, LinesAdded: added, LinesRemoved: removed}} + } + + // Without the diff, the tool's input still bounds the change. A write's + // whole content is added lines; whatever an overwrite displaced is not + // recorded anywhere, so it stays uncounted rather than guessed at. + if tool == "write" { + return []usage.FileChange{{Path: path, LinesAdded: usage.CountLines(agentdata.StringField(input, "content"))}} + } + return []usage.FileChange{{ + Path: path, + LinesAdded: usage.CountLines(agentdata.StringField(input, "newString")), + LinesRemoved: usage.CountLines(agentdata.StringField(input, "oldString")), + }} +} + +// countUnifiedDiff counts the added and removed lines of a unified diff. +// The "+++"/"---" file headers are markup, not changes. +func countUnifiedDiff(patch string) (added, removed uint64, ok bool) { + if strings.TrimSpace(patch) == "" { + return 0, 0, false + } + for _, line := range strings.Split(patch, "\n") { + switch { + case strings.HasPrefix(line, "+++"), strings.HasPrefix(line, "---"): + case strings.HasPrefix(line, "+"): + added++ + case strings.HasPrefix(line, "-"): + removed++ + } + } + return added, removed, true +} + +// stableMessageID keys an entry to the message id Kilo assigned, so +// re-ingesting the same message never double-counts. +func stableMessageID(messageID string) string { + if strings.TrimSpace(messageID) == "" { + return "" + } + return usage.StableID("kilo", strings.TrimSpace(messageID)) +} diff --git a/internal/provider/kilo/loader_test.go b/internal/provider/kilo/loader_test.go new file mode 100644 index 0000000..c909371 --- /dev/null +++ b/internal/provider/kilo/loader_test.go @@ -0,0 +1,299 @@ +package kilo + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// writeKiloDB builds a database shaped like the one Kilo's engine writes: a +// session row, message rows carrying the token block, and part rows carrying +// the tool calls that touched files. +type kiloRow struct { + id string + sessionID string + created int64 + data string +} + +func writeKiloDB(t *testing.T, sessions, messages, parts []kiloRow) string { + t.Helper() + path := filepath.Join(t.TempDir(), "kilo.db") + db := providertest.OpenTestSQLite(t, path) + defer db.Close() + + schema := []string{ + `CREATE TABLE session (id TEXT PRIMARY KEY, directory TEXT, model TEXT)`, + `CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT, time_created INTEGER, data TEXT)`, + `CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, session_id TEXT, time_created INTEGER, data TEXT)`, + } + for _, statement := range schema { + if _, err := db.Exec(statement); err != nil { + t.Fatal(err) + } + } + for _, row := range sessions { + if _, err := db.Exec(`INSERT INTO session (id, directory, model) VALUES (?, ?, ?)`, row.id, row.sessionID, row.data); err != nil { + t.Fatal(err) + } + } + for _, row := range messages { + if _, err := db.Exec(`INSERT INTO message (id, session_id, time_created, data) VALUES (?, ?, ?, ?)`, + row.id, row.sessionID, row.created, row.data); err != nil { + t.Fatal(err) + } + } + for i, row := range parts { + if _, err := db.Exec(`INSERT INTO part (id, message_id, session_id, time_created, data) VALUES (?, ?, ?, ?, ?)`, + row.id, row.sessionID, "ses-1", int64(i), row.data); err != nil { + t.Fatal(err) + } + } + return path +} + +func loadKilo(t *testing.T, path string) []usage.Entry { + t.Helper() + entries, err := Provider{}.WithPaths([]string{filepath.Dir(path)}).Entries() + if err != nil { + t.Fatal(err) + } + return entries +} + +// TestKiloBillsEachMessageAsReported pins down the token semantics that set +// Kilo apart from the OpenCode it forked: every message reports what its one +// API request consumed, so messages sum as-is — no diffing against the +// previous turn. Kilo's own session rows sum their messages the same way. +func TestKiloBillsEachMessageAsReported(t *testing.T) { + path := writeKiloDB(t, + []kiloRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []kiloRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"id":"msg-1","role":"assistant","modelID":"kilo-auto/free","tokens":{"input":20890,"output":102,"reasoning":39,"cache":{"read":0,"write":0}},"time":{"created":1785546752112}}`}, + {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"id":"msg-2","role":"assistant","modelID":"kilo-auto/free","tokens":{"input":2508,"output":212,"reasoning":46,"cache":{"read":20864,"write":0}},"time":{"created":1785546760000}}`}, + }, + nil, + ) + + entries := loadKilo(t, path) + if len(entries) != 2 { + t.Fatalf("entries = %d, want 2", len(entries)) + } + var total usage.TokenUsage + for _, entry := range entries { + total.InputTokens += entry.Usage.InputTokens + total.OutputTokens += entry.Usage.OutputTokens + total.ReasoningOutputTokens += entry.Usage.ReasoningOutputTokens + total.CacheReadInputTokens += entry.Usage.CacheReadInputTokens + } + want := usage.TokenUsage{ + InputTokens: 20890 + 2508, + OutputTokens: 102 + 212, + ReasoningOutputTokens: 39 + 46, + CacheReadInputTokens: 0 + 20864, + } + if total != want { + t.Fatalf("summed usage = %#v, want %#v", total, want) + } +} + +// TestKiloSkipsMessagesWithoutTokens covers user messages, which carry no +// token block at all: they are turns someone typed, not API requests. +func TestKiloSkipsMessagesWithoutTokens(t *testing.T) { + path := writeKiloDB(t, + []kiloRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []kiloRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"id":"msg-1","role":"user","time":{"created":1785546751000}}`}, + {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"id":"msg-2","role":"assistant","modelID":"kilo-auto/free","tokens":{"input":100,"output":10,"cache":{"read":0,"write":0}},"time":{"created":1785546752000}}`}, + }, + nil, + ) + + entries := loadKilo(t, path) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].Usage.InputTokens != 100 { + t.Fatalf("input tokens = %d, want 100", entries[0].Usage.InputTokens) + } +} + +// TestKiloProjectComesFromMessageCwd: the project must be the directory the +// agent ran in, never a hardcoded name. +func TestKiloProjectComesFromMessageCwd(t *testing.T) { + path := writeKiloDB(t, + []kiloRow{{id: "ses-1", sessionID: "/fallback/dir"}}, + []kiloRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"id":"msg-1","role":"assistant","modelID":"kilo-auto/free","path":{"cwd":"/Users/dev/workspace/gemini-editing"},"tokens":{"input":100,"output":10,"cache":{"read":0,"write":0}},"time":{"created":1785546752000}}`}, + {id: "msg-2", sessionID: "ses-1", created: 2000, data: `{"id":"msg-2","role":"assistant","modelID":"kilo-auto/free","tokens":{"input":100,"output":10,"cache":{"read":0,"write":0}},"time":{"created":1785546753000}}`}, + }, + nil, + ) + + entries := loadKilo(t, path) + if len(entries) != 2 { + t.Fatalf("entries = %d, want 2", len(entries)) + } + projects := map[string]bool{} + for _, entry := range entries { + projects[entry.Project] = true + } + // msg-1 names its cwd; msg-2 falls back to the session directory. + if !projects["gemini-editing"] || !projects["dir"] { + t.Fatalf("projects = %v, want gemini-editing and dir", projects) + } +} + +// TestKiloCountsFileChangesFromFilediff: Kilo attaches the unified diff of +// every edit and write to the part, and the counts must come from it — +// including the removed side of an overwrite, which the tool input alone +// cannot see. +func TestKiloCountsFileChangesFromFilediff(t *testing.T) { + patch := "Index: /repo/demo/main.go\\n===================================================================\\n--- /repo/demo/main.go\\n+++ /repo/demo/main.go\\n@@ -1,2 +1,3 @@\\n-old line\\n+new line one\\n+new line two\\n+new line three\\n-another old\\n" + path := writeKiloDB(t, + []kiloRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []kiloRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"id":"msg-1","role":"assistant","modelID":"kilo-auto/free","path":{"cwd":"/repo/demo"},"tokens":{"input":100,"output":10,"cache":{"read":0,"write":0}},"time":{"created":1785546752000}}`}, + }, + []kiloRow{ + {id: "prt-1", sessionID: "msg-1", data: `{"type":"tool","tool":"write","state":{"status":"completed","input":{"filePath":"/repo/demo/main.go","content":"ignored when the diff is present"},"metadata":{"filepath":"/repo/demo/main.go","filediff":{"file":"/repo/demo/main.go","patch":"` + patch + `"}}}}`}, + }, + ) + + entries := loadKilo(t, path) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + entry := entries[0] + if entry.LinesAdded != 3 || entry.LinesRemoved != 2 { + t.Fatalf("lines = +%d/-%d, want +3/-2", entry.LinesAdded, entry.LinesRemoved) + } + if entry.Language != "Go" { + t.Fatalf("language = %q, want Go", entry.Language) + } + if entry.Entity != "/repo/demo/main.go" { + t.Fatalf("entity = %q, want /repo/demo/main.go", entry.Entity) + } +} + +// TestKiloFallsBackToToolInputWithoutDiff: a part without the diff still +// bounds the change through the tool's own input. +func TestKiloFallsBackToToolInputWithoutDiff(t *testing.T) { + path := writeKiloDB(t, + []kiloRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []kiloRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"id":"msg-1","role":"assistant","modelID":"kilo-auto/free","path":{"cwd":"/repo/demo"},"tokens":{"input":100,"output":10,"cache":{"read":0,"write":0}},"time":{"created":1785546752000}}`}, + }, + []kiloRow{ + {id: "prt-1", sessionID: "msg-1", data: `{"type":"tool","tool":"write","state":{"status":"completed","input":{"filePath":"main.py","content":"line one\nline two"}}}`}, + }, + ) + + entries := loadKilo(t, path) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].LinesAdded != 2 || entries[0].LinesRemoved != 0 { + t.Fatalf("lines = +%d/-%d, want +2/-0", entries[0].LinesAdded, entries[0].LinesRemoved) + } + if entries[0].Language != "Python" { + t.Fatalf("language = %q, want Python", entries[0].Language) + } +} + +// TestKiloIgnoresSnapshotPatchParts: a "patch" part is a git snapshot of the +// whole worktree, not agent work. +func TestKiloIgnoresSnapshotPatchParts(t *testing.T) { + path := writeKiloDB(t, + []kiloRow{{id: "ses-1", sessionID: "/repo/demo"}}, + []kiloRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"id":"msg-1","role":"assistant","modelID":"kilo-auto/free","tokens":{"input":100,"output":10,"cache":{"read":0,"write":0}},"time":{"created":1785546752000}}`}, + }, + []kiloRow{ + {id: "prt-1", sessionID: "msg-1", data: `{"type":"patch","hash":"abc","files":["/repo/demo/.DS_Store"]}`}, + {id: "prt-2", sessionID: "msg-1", data: `{"type":"tool","tool":"write","state":{"status":"pending","input":{"filePath":"a.go","content":"x"}}}`}, + }, + ) + + entries := loadKilo(t, path) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].LinesAdded != 0 || entries[0].LinesRemoved != 0 { + t.Fatalf("lines = +%d/-%d, want zero from snapshot and incomplete parts", entries[0].LinesAdded, entries[0].LinesRemoved) + } +} + +// TestKiloFallsBackToSessionModel: a message without a modelID keeps its +// tokens under the model the session was started with, which the session row +// stores as a JSON object. +func TestKiloFallsBackToSessionModel(t *testing.T) { + path := writeKiloDB(t, + []kiloRow{{id: "ses-1", sessionID: "/repo/demo", data: `{"id":"kilo-auto/free","providerID":"kilo","variant":"default"}`}}, + []kiloRow{ + {id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"id":"msg-1","role":"assistant","tokens":{"input":100,"output":10,"cache":{"read":0,"write":0}},"time":{"created":1785546752000}}`}, + }, + nil, + ) + + entries := loadKilo(t, path) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].Model != "kilo-auto/free" { + t.Fatalf("model = %q, want kilo-auto/free", entries[0].Model) + } +} + +// TestKiloReadsSessionsWithoutModelColumn: an old database without the model +// column must still yield the session directory, not fail the whole query. +func TestKiloReadsSessionsWithoutModelColumn(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "kilo.db") + db := providertest.OpenTestSQLite(t, path) + schema := []string{ + `CREATE TABLE session (id TEXT PRIMARY KEY, directory TEXT)`, + `CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT, time_created INTEGER, data TEXT)`, + `CREATE TABLE part (id TEXT PRIMARY KEY, message_id TEXT, session_id TEXT, time_created INTEGER, data TEXT)`, + } + for _, statement := range schema { + if _, err := db.Exec(statement); err != nil { + t.Fatal(err) + } + } + if _, err := db.Exec(`INSERT INTO session (id, directory) VALUES ('ses-1', '/repo/legacy-project')`); err != nil { + t.Fatal(err) + } + if _, err := db.Exec(`INSERT INTO message (id, session_id, time_created, data) VALUES ('msg-1', 'ses-1', 1000, ?)`, + `{"id":"msg-1","role":"assistant","modelID":"kilo-auto/free","tokens":{"input":100,"output":10,"cache":{"read":0,"write":0}},"time":{"created":1785546752000}}`); err != nil { + t.Fatal(err) + } + db.Close() + + entries := loadKilo(t, path) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].Project != "legacy-project" { + t.Fatalf("project = %q, want legacy-project", entries[0].Project) + } +} + +// TestKiloEntryIDSurvivesRelocation: the entry ID keys on the message id Kilo +// assigned, so moving the database must not create a second identity. +func TestKiloEntryIDSurvivesRelocation(t *testing.T) { + message := kiloRow{id: "msg-1", sessionID: "ses-1", created: 1000, data: `{"id":"msg-1","role":"assistant","modelID":"kilo-auto/free","tokens":{"input":100,"output":10,"cache":{"read":0,"write":0}},"time":{"created":1785546752000}}`} + first := loadKilo(t, writeKiloDB(t, []kiloRow{{id: "ses-1", sessionID: "/repo/demo"}}, []kiloRow{message}, nil)) + second := loadKilo(t, writeKiloDB(t, []kiloRow{{id: "ses-1", sessionID: "/repo/demo"}}, []kiloRow{message}, nil)) + if len(first) != 1 || len(second) != 1 { + t.Fatalf("entries = %d and %d, want 1 and 1", len(first), len(second)) + } + if first[0].ID != second[0].ID { + t.Fatalf("ID changed across relocation: %q vs %q", first[0].ID, second[0].ID) + } + if first[0].ID != usage.StableID("kilo", "msg-1") { + t.Fatalf("ID = %q, want StableID(kilo, msg-1)", first[0].ID) + } +} diff --git a/internal/provider/kilo/provider_test.go b/internal/provider/kilo/provider_test.go index d0fe3b4..b8af5cd 100644 --- a/internal/provider/kilo/provider_test.go +++ b/internal/provider/kilo/provider_test.go @@ -1,7 +1,6 @@ package kilo import ( - "path/filepath" "testing" "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" @@ -11,25 +10,23 @@ import ( // TestLoadsEntry is the Kilo smoke test: a minimal fixture must produce // exactly one entry with the expected identity and token counts. func TestLoadsEntry(t *testing.T) { - entries, err := func() ([]usage.Entry, error) { - dir := t.TempDir() - dbPath := filepath.Join(dir, "kilo.db") - db := providertest.OpenTestSQLite(t, dbPath) - defer db.Close() - providertest.ExecSQL(t, db, "CREATE TABLE message (id TEXT, session_id TEXT, data TEXT)") - providertest.ExecSQL(t, db, `INSERT INTO message (id, session_id, data) VALUES (?, ?, ?)`, - "row-1", - "session-a", - `{"id":"msg-1","role":"assistant","providerID":"anthropic","modelID":"claude-sonnet-4-20250514","time":{"created":1767312000000},"tokens":{"input":100,"output":50,"reasoning":5,"cache":{"read":10,"write":20}}}`, - ) - return Provider{}.WithPaths([]string{dir}).Entries() - }() + path := writeKiloDB(t, + []kiloRow{{id: "session-a", sessionID: "/repo/demo"}}, + []kiloRow{{ + id: "row-1", + sessionID: "session-a", + created: 1767312000000, + data: `{"id":"msg-1","role":"assistant","providerID":"kilo","modelID":"kilo-auto/free","path":{"cwd":"/repo/demo"},"time":{"created":1767312000000},"tokens":{"input":100,"output":50,"reasoning":5,"cache":{"read":10,"write":20}}}`, + }}, + nil, + ) + entries, err := Provider{}.WithPaths([]string{path}).Entries() providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ Provider: usage.ProviderKilo, - Model: "claude-sonnet-4-20250514", + Model: "kilo-auto/free", SessionID: "session-a", - Project: "kilo", + Project: "demo", Tokens: usage.TokenUsage{ InputTokens: 100, OutputTokens: 50, diff --git a/internal/provider/opencode/loader_test.go b/internal/provider/opencode/loader_test.go index 7399253..a8acee6 100644 --- a/internal/provider/opencode/loader_test.go +++ b/internal/provider/opencode/loader_test.go @@ -4,7 +4,7 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" ) @@ -26,10 +26,7 @@ func writeOpenCodeDB(t *testing.T, sessions, messages, parts []openCodeRow) stri func writeOpenCodeDBAt(t *testing.T, path string, sessions, messages, parts []openCodeRow) string { t.Helper() - db, err := agentdb.OpenSQLite(path) - if err != nil { - t.Fatal(err) - } + db := providertest.OpenTestSQLite(t, path) defer db.Close() schema := []string{ @@ -436,10 +433,7 @@ func TestDetectsLanguageFromTouchedFiles(t *testing.T) { // would take the project directory down with it. func TestOpenCodeReadsSessionsWithoutModelColumn(t *testing.T) { path := filepath.Join(t.TempDir(), "opencode.db") - db, err := agentdb.OpenSQLite(path) - if err != nil { - t.Fatal(err) - } + db := providertest.OpenTestSQLite(t, path) for _, statement := range []string{ `CREATE TABLE session (id TEXT PRIMARY KEY, directory TEXT, version TEXT)`, `CREATE TABLE message (id TEXT PRIMARY KEY, session_id TEXT, time_created INTEGER, data TEXT)`, diff --git a/internal/providertest/providertest.go b/internal/providertest/providertest.go index c0041fa..35d1747 100644 --- a/internal/providertest/providertest.go +++ b/internal/providertest/providertest.go @@ -9,8 +9,11 @@ import ( "path/filepath" "testing" - "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + + // Fixture databases are opened directly with database/sql; the driver + // registration rides in from agentdb. + _ "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" ) // WantEntry describes the single entry a provider must produce from a minimal @@ -66,16 +69,22 @@ func WriteFile(t *testing.T, path, data string) { } } -// OpenTestSQLite opens a fixture database, creating its directory. +// OpenTestSQLite opens a fixture database, creating its directory. Fixtures +// are built here with a writable connection; production code reads them back +// through agentdb.OpenSQLite, which is read-only by design. func OpenTestSQLite(t *testing.T, path string) *sql.DB { t.Helper() if err := os.MkdirAll(filepath.Dir(path), 0o700); err != nil { t.Fatal(err) } - db, err := agentdb.OpenSQLite(path) + db, err := sql.Open("sqlite", path) if err != nil { t.Fatal(err) } + if err := db.Ping(); err != nil { + _ = db.Close() + t.Fatal(err) + } return db } From 7803f2d6cead90a9badecc05b5ffe0e2d40b0551 Mon Sep 17 00:00:00 2001 From: akarineren Date: Mon, 3 Aug 2026 14:35:27 +0900 Subject: [PATCH 08/11] feat: add WorkBuddy provider with entry loading and usage tracking functionality --- cmd/tokitoki/main.go | 2 +- internal/provider/workbuddy/loader.go | 306 +++++++++++++++++++ internal/provider/workbuddy/loader_test.go | 199 ++++++++++++ internal/provider/workbuddy/provider.go | 36 +++ internal/provider/workbuddy/provider_test.go | 35 +++ internal/usage/usage.go | 31 +- internal/usagescan/scanner.go | 2 + pkg/agentlib/agentlib.go | 4 + 8 files changed, 599 insertions(+), 16 deletions(-) create mode 100644 internal/provider/workbuddy/loader.go create mode 100644 internal/provider/workbuddy/loader_test.go create mode 100644 internal/provider/workbuddy/provider.go create mode 100644 internal/provider/workbuddy/provider_test.go diff --git a/cmd/tokitoki/main.go b/cmd/tokitoki/main.go index 4479d4c..de0ec3c 100644 --- a/cmd/tokitoki/main.go +++ b/cmd/tokitoki/main.go @@ -565,7 +565,7 @@ Each invocation scans the provider roots you pass and uploads their usage events to the Tokitoki server (TOKITOKI_BASE_URL, default https://tokitoki.dev). By default, tokitoki scans the built-in roots for claude, codex, copilot, gemini, kimi, qwen, openclaw, pi, amp, droid, kilo, -hermes, codebuff, opencode, and goose. Pass one or more +hermes, codebuff, opencode, goose, and workbuddy. Pass one or more --provider-dir provider=dir values to scan an explicit provider set. The API key is read from ~/.tokitoki/api_key; use tokitoki set key to create or update that file. diff --git a/internal/provider/workbuddy/loader.go b/internal/provider/workbuddy/loader.go new file mode 100644 index 0000000..474fcfb --- /dev/null +++ b/internal/provider/workbuddy/loader.go @@ -0,0 +1,306 @@ +package workbuddy + +import ( + "path/filepath" + "strconv" + "strings" + + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/langdetect" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// fallbackModel is WorkBuddy's auto-router placeholder, used when neither the +// record nor settings.json names a model. +const fallbackModel = "auto" + +// WorkBuddy (Tencent's Claude Code fork) writes session transcripts under +// /projects//.jsonl, with sub-agent traffic +// nested at /subagents/.jsonl. Every LLM round-trip carries +// providerData.rawUsage — on function_call records as well as assistant +// messages — so usage is aggregated from any record that has it, not by +// record type. +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + entries := make([]usage.Entry, 0) + seen := make(map[string]bool) + for _, root := range paths { + files := agentdata.CollectExt(filepath.Join(root, "projects"), ".jsonl") + files = agentdata.FilterFiles(files, filter) + model := settingsModel(root) + for _, file := range files { + fileEntries, err := parseSessionFile(file, model) + if err != nil { + return nil, err + } + for _, entry := range fileEntries { + if seen[entry.ID] { + continue + } + seen[entry.ID] = true + entries = append(entries, entry) + } + } + } + usageprovider.SortEntries(entries) + return entries, nil +} + +func parseSessionFile(path, settingsModel string) ([]usage.Entry, error) { + lines, err := agentdata.ReadJSONLines(path) + if err != nil { + return nil, err + } + fileSessionID := sessionIDFromPath(path) + fallbackTime := agentdata.FileModifiedTime(path) + + // Pass 1: build one entry per usage-carrying record and index the first + // entry of each assistant turn by messageId. + entries := make([]usage.Entry, 0, len(lines)) + entryLines := make([]int, 0, len(lines)) + entryByMid := make(map[string]int) + for _, line := range lines { + providerData := agentdata.ObjectAt(line.Value["providerData"]) + rawUsage := agentdata.ObjectAt(providerData["rawUsage"]) + if rawUsage == nil { + continue + } + tokens := normalizeTokens(rawUsage) + if !usageprovider.NonZero(tokens) { + continue + } + + timestamp, ok := agentdata.ParseTimestamp(line.Value["timestamp"]) + if !ok { + timestamp = fallbackTime + } + + project := usage.UnknownProject + projectPath := "" + if path, name, ok := usage.ProjectFromCWD(agentdata.StringField(line.Value, "cwd")); ok { + projectPath = path + project = name + } + + sessionID := agentdata.FirstNonEmpty(agentdata.StringField(line.Value, "sessionId"), fileSessionID) + // providerData.model is the model the auto-router actually picked; + // requestModelId is usually the "auto" placeholder. + model := agentdata.FirstNonEmpty( + agentdata.StringField(providerData, "model"), + agentdata.StringField(providerData, "requestModelId"), + agentdata.StringField(line.Value, "model"), + settingsModel, + ) + + entry := usageprovider.BaseEntry(usage.ProviderWorkbuddy, timestamp, project, projectPath, sessionID, model, "", tokens) + usageprovider.SetSource(&entry, path, line.Line, line.Start, line.End) + entry.ID = entryID(entry, providerData, line.Value) + if mid := agentdata.StringField(providerData, "messageId"); mid != "" { + if _, exists := entryByMid[mid]; !exists { + entryByMid[mid] = len(entries) + } + } + entries = append(entries, entry) + entryLines = append(entryLines, line.Line) + } + if len(entries) == 0 { + return entries, nil + } + + // Pass 2: route file diffs and language signals onto entries. A record's + // messageId names its assistant turn exactly; a record whose turn produced + // no usage entry (interrupted turn, parallel tool calls) falls back to the + // nearest preceding entry, and anything before the first entry lands on it. + candidates := make([][]langdetect.Candidate, len(entries)) + last, next := 0, 0 + for _, line := range lines { + for next < len(entries) && entryLines[next] <= line.Line { + last = next + next++ + } + target := last + mid := agentdata.StringField(agentdata.ObjectAt(line.Value["providerData"]), "messageId") + if index, ok := entryByMid[mid]; ok && mid != "" { + target = index + } + if change, ok := fileChangeFromCall(line.Value); ok { + entries[target].ApplyFileChange(change) + } + candidates[target] = append(candidates[target], languageCandidates(line.Value)...) + } + for i := range entries { + entries[i].Language = usage.NormalizeLanguage(langdetect.Dominant(candidates[i])) + } + return entries, nil +} + +// fileChangeFromCall turns a Write or Edit tool call into the diff it applies. +// WorkBuddy records no structured patch, so the tool arguments are the diff: +// a Write's content is all added lines and an Edit swaps old_string for +// new_string. Line counts come from the call, not the result — a rejected +// call overcounts slightly, which beats parsing every result blob. +func fileChangeFromCall(value map[string]any) (usage.FileChange, bool) { + if agentdata.StringField(value, "type") != "function_call" { + return usage.FileChange{}, false + } + arguments := agentdata.DecodeJSONObjectString(agentdata.StringField(value, "arguments")) + if arguments == nil { + return usage.FileChange{}, false + } + path := usage.ResolvePath(agentdata.StringField(value, "cwd"), agentdata.StringField(arguments, "file_path")) + if path == "" { + return usage.FileChange{}, false + } + switch agentdata.StringField(value, "name") { + case "Write": + return usage.FileChange{ + Path: path, + LinesAdded: usage.CountLines(agentdata.StringField(arguments, "content")), + }, true + case "Edit": + return usage.FileChange{ + Path: path, + LinesAdded: usage.CountLines(agentdata.StringField(arguments, "new_string")), + LinesRemoved: usage.CountLines(agentdata.StringField(arguments, "old_string")), + }, true + } + return usage.FileChange{}, false +} + +// languageCandidates mines a record for programming-language signals the same +// way the Claude provider mines an assistant message: tool-call file paths +// weigh 3, free text weighs 1 per path it mentions. function_call arguments +// are WorkBuddy's tool_use blocks; message content carries the text blocks. +func languageCandidates(value map[string]any) []langdetect.Candidate { + candidates := make([]langdetect.Candidate, 0) + switch agentdata.StringField(value, "type") { + case "function_call": + arguments := agentdata.DecodeJSONObjectString(agentdata.StringField(value, "arguments")) + for key, child := range arguments { + lower := strings.ToLower(key) + text := agentdata.StringValue(child) + if text == "" { + continue + } + if strings.Contains(lower, "file") || strings.Contains(lower, "path") { + if langdetect.FromPath(text) != langdetect.Unknown { + candidates = append(candidates, langdetect.Candidate{Path: text, Weight: 3}) + continue + } + } + if strings.Contains(lower, "command") || strings.Contains(lower, "content") || strings.Contains(lower, "query") { + for _, path := range langdetect.PathsFromText(text) { + candidates = append(candidates, langdetect.Candidate{Path: path, Weight: 1}) + } + } + } + case "message": + for _, block := range agentdata.ArrayAt(value["content"]) { + text := agentdata.StringField(agentdata.ObjectAt(block), "text") + for _, path := range langdetect.PathsFromText(text) { + candidates = append(candidates, langdetect.Candidate{Path: path, Weight: 1}) + } + } + if text := agentdata.StringValue(value["content"]); text != "" { + for _, path := range langdetect.PathsFromText(text) { + candidates = append(candidates, langdetect.Candidate{Path: path, Weight: 1}) + } + } + } + return candidates +} + +// normalizeTokens converts WorkBuddy's OpenAI-shaped rawUsage into the +// normalized breakdown. prompt_tokens is the FULL prompt — cache reads, +// cache writes, and genuinely-new input — and the cache split is mirrored in +// Anthropic-style or DeepSeek-style fields depending on which upstream the +// auto-router picked, so each cache bucket takes the largest mirror. +// completion_tokens includes reasoning (total == prompt + completion). +func normalizeTokens(rawUsage map[string]any) usage.TokenUsage { + prompt := agentdata.UintField(rawUsage, "prompt_tokens") + completion := agentdata.UintField(rawUsage, "completion_tokens") + promptDetails := agentdata.ObjectAt(rawUsage["prompt_tokens_details"]) + completionDetails := agentdata.ObjectAt(rawUsage["completion_tokens_details"]) + + cacheRead := max( + agentdata.UintField(rawUsage, "cache_read_input_tokens"), + agentdata.UintField(promptDetails, "cached_tokens"), + agentdata.UintField(rawUsage, "prompt_cache_hit_tokens"), + ) + cacheCreation := max( + agentdata.UintField(rawUsage, "cache_creation_input_tokens"), + agentdata.UintField(rawUsage, "prompt_cache_write_tokens"), + ) + input := uint64(0) + if cached := cacheRead + cacheCreation; prompt > cached { + input = prompt - cached + } + reasoning := min(completion, max( + agentdata.UintField(completionDetails, "reasoning_tokens"), + agentdata.UintField(rawUsage, "completion_thinking_tokens"), + )) + output := completion - reasoning + + return usage.TokenUsage{ + InputTokens: input, + OutputTokens: output, + CacheCreationInputTokens: cacheCreation, + CacheReadInputTokens: cacheRead, + ReasoningOutputTokens: reasoning, + TotalTokens: input + output + cacheCreation + cacheRead + reasoning, + } +} + +// entryID keys on the response-level providerData.messageId plus the token +// counts. The messageId is shared by every record of one logical assistant +// turn and can be replayed across session files, so anything file- or +// position-specific in the key would double count a mirrored record. But one +// turn can also contain several REAL API calls under the same messageId +// (observed: a Write call and its continuation, seconds apart, each with its +// own growing prompt) — the token counts are what tell a mirror from a +// genuine second call, so they complete the key. +func entryID(entry usage.Entry, providerData, value map[string]any) string { + messageID := agentdata.StringField(providerData, "messageId") + if messageID == "" { + messageID = agentdata.FirstStringField(value, "uuid", "id") + } + if messageID != "" { + tokens := entry.Usage + return usage.StableID( + string(usage.ProviderWorkbuddy), + messageID, + strconv.FormatUint(tokens.InputTokens, 10), + strconv.FormatUint(tokens.OutputTokens, 10), + strconv.FormatUint(tokens.CacheCreationInputTokens, 10), + strconv.FormatUint(tokens.CacheReadInputTokens, 10), + strconv.FormatUint(tokens.ReasoningOutputTokens, 10), + ) + } + return usageprovider.StableEntryID(entry) +} + +// sessionIDFromPath derives the session id from the file's location under +// projects/: /.jsonl for main sessions and +// //subagents/.jsonl for sub-agents. +func sessionIDFromPath(path string) string { + dir := filepath.Dir(path) + if filepath.Base(dir) == "subagents" { + return filepath.Base(filepath.Dir(dir)) + } + stem := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) + if stem == "" { + return "unknown" + } + return stem +} + +func settingsModel(root string) string { + settings, err := agentdata.ReadJSONObject(filepath.Join(root, "settings.json")) + if err != nil || settings == nil { + return fallbackModel + } + if model := agentdata.StringField(settings, "model"); model != "" { + return model + } + return fallbackModel +} diff --git a/internal/provider/workbuddy/loader_test.go b/internal/provider/workbuddy/loader_test.go new file mode 100644 index 0000000..f418f78 --- /dev/null +++ b/internal/provider/workbuddy/loader_test.go @@ -0,0 +1,199 @@ +package workbuddy + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestDedupesSharedMessageID verifies that the function_call and message +// records of one LLM round-trip — which share providerData.messageId — count +// once, not twice. +func TestDedupesSharedMessageID(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "proj", "session-a.jsonl") + providertest.WriteFile(t, path, + `{"timestamp":1785730645570,"type":"function_call","providerData":{"messageId":"msg-1","model":"glm-5.2","rawUsage":{"prompt_tokens":100,"completion_tokens":10}},"sessionId":"session-a","cwd":"/tmp/proj"}`+"\n"+ + `{"timestamp":1785730645580,"type":"message","role":"assistant","providerData":{"messageId":"msg-1","model":"glm-5.2","rawUsage":{"prompt_tokens":100,"completion_tokens":10}},"sessionId":"session-a","cwd":"/tmp/proj"}`+"\n") + + entries, err := loadEntries([]string{dir}, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1 (messageId must dedupe)", len(entries)) + } + if entries[0].Usage.TotalTokens != 110 { + t.Fatalf("total = %d, want 110", entries[0].Usage.TotalTokens) + } +} + +// TestKeepsDistinctCallsSharingMessageID verifies that two REAL API calls +// under one messageId — same assistant turn, different token counts — both +// count. Observed in real data: a Write call and its continuation share the +// messageId but each consumed its own tokens. +func TestKeepsDistinctCallsSharingMessageID(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "proj", "session-a.jsonl") + providertest.WriteFile(t, path, + `{"timestamp":1785730404526,"type":"function_call","name":"Write","providerData":{"messageId":"msg-1","model":"glm-5.2","rawUsage":{"prompt_tokens":37883,"completion_tokens":2466,"prompt_cache_hit_tokens":37760}},"sessionId":"session-a","cwd":"/tmp/proj"}`+"\n"+ + `{"timestamp":1785730406541,"type":"function_call","name":"TaskUpdate","providerData":{"messageId":"msg-1","model":"glm-5.2","rawUsage":{"prompt_tokens":40561,"completion_tokens":77,"prompt_cache_hit_tokens":37824}},"sessionId":"session-a","cwd":"/tmp/proj"}`+"\n") + + entries, err := loadEntries([]string{dir}, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("entries = %d, want 2 (distinct usage under one messageId must both count)", len(entries)) + } +} + +// TestSubagentSessionID verifies that sub-agent transcripts nested at +// /subagents/.jsonl attribute usage to the parent session. +func TestSubagentSessionID(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "proj", "session-a", "subagents", "agent-1.jsonl") + providertest.WriteFile(t, path, `{"timestamp":1785730645570,"type":"function_call","providerData":{"messageId":"msg-sub","model":"glm-5.2","rawUsage":{"prompt_tokens":50,"completion_tokens":5}},"cwd":"/tmp/proj"}`+"\n") + + entries, err := loadEntries([]string{dir}, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].SessionID != "session-a" { + t.Fatalf("session id = %q, want %q", entries[0].SessionID, "session-a") + } +} + +// TestAnthropicStyleCacheSplit verifies the Anthropic-mirrored cache fields: +// prompt_tokens is the full prompt, so cache reads AND writes are both +// subtracted to get pure input. +func TestAnthropicStyleCacheSplit(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "proj", "session-a.jsonl") + providertest.WriteFile(t, path, `{"timestamp":1785730645570,"type":"message","providerData":{"messageId":"msg-1","model":"claude-x","rawUsage":{"prompt_tokens":350,"completion_tokens":20,"cache_read_input_tokens":100,"cache_creation_input_tokens":200}},"cwd":"/tmp/proj"}`+"\n") + + entries, err := loadEntries([]string{dir}, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + want := usage.TokenUsage{ + InputTokens: 50, + OutputTokens: 20, + CacheCreationInputTokens: 200, + CacheReadInputTokens: 100, + TotalTokens: 370, + } + if entries[0].Usage != want { + t.Fatalf("usage = %#v, want %#v", entries[0].Usage, want) + } +} + +// TestModelFallsBackToSettings verifies that a record naming no model uses +// the settings.json model, and "auto" as the last resort. +func TestModelFallsBackToSettings(t *testing.T) { + dir := t.TempDir() + providertest.WriteFile(t, filepath.Join(dir, "settings.json"), `{"model":"glm-configured"}`) + path := filepath.Join(dir, "projects", "proj", "session-a.jsonl") + providertest.WriteFile(t, path, `{"timestamp":1785730645570,"type":"message","providerData":{"messageId":"msg-1","rawUsage":{"prompt_tokens":10,"completion_tokens":1}},"cwd":"/tmp/proj"}`+"\n") + + entries, err := loadEntries([]string{dir}, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].Model != "glm-configured" { + t.Fatalf("model = %q, want %q", entries[0].Model, "glm-configured") + } +} + +// TestAttachesFileChangesByMessageID verifies that a Write tool call without +// its own usage lands on the usage entry of the same assistant turn — matched +// by messageId even when the call precedes the turn's usage record — and that +// the entry gets language, entity, and line stats from it. +func TestAttachesFileChangesByMessageID(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "proj", "session-a.jsonl") + providertest.WriteFile(t, path, + `{"timestamp":1785730645570,"type":"function_call","name":"Write","arguments":"{\"file_path\":\"/tmp/proj/main.go\",\"content\":\"package main\\nfunc main() {}\\n\"}","providerData":{"messageId":"msg-1"},"cwd":"/tmp/proj"}`+"\n"+ + `{"timestamp":1785730645580,"type":"function_call","name":"TaskUpdate","arguments":"{}","providerData":{"messageId":"msg-1","model":"glm-5.2","rawUsage":{"prompt_tokens":100,"completion_tokens":10}},"sessionId":"session-a","cwd":"/tmp/proj"}`+"\n") + + entries, err := loadEntries([]string{dir}, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + entry := entries[0] + if entry.LinesAdded != 2 || entry.LinesRemoved != 0 { + t.Fatalf("lines = +%d/-%d, want +2/-0", entry.LinesAdded, entry.LinesRemoved) + } + if entry.Entity != "/tmp/proj/main.go" || entry.EntityType != "file" { + t.Fatalf("entity = %q (%q), want /tmp/proj/main.go (file)", entry.Entity, entry.EntityType) + } + if entry.IsWrite == nil || !*entry.IsWrite { + t.Fatal("IsWrite not set") + } + if entry.Language != "Go" { + t.Fatalf("language = %q, want Go", entry.Language) + } +} + +// TestEditCountsBothSides verifies an Edit call counts old_string as removed +// and new_string as added, falling back to the nearest preceding entry when +// its turn has no usage entry of its own. +func TestEditCountsBothSides(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "proj", "session-a.jsonl") + providertest.WriteFile(t, path, + `{"timestamp":1785730645570,"type":"message","providerData":{"messageId":"msg-1","model":"glm-5.2","rawUsage":{"prompt_tokens":100,"completion_tokens":10}},"sessionId":"session-a","cwd":"/tmp/proj"}`+"\n"+ + `{"timestamp":1785730645580,"type":"function_call","name":"Edit","arguments":"{\"file_path\":\"src/app.py\",\"old_string\":\"a\\nb\\nc\",\"new_string\":\"a\"}","providerData":{"messageId":"msg-orphan"},"cwd":"/tmp/proj"}`+"\n") + + entries, err := loadEntries([]string{dir}, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + entry := entries[0] + if entry.LinesAdded != 1 || entry.LinesRemoved != 3 { + t.Fatalf("lines = +%d/-%d, want +1/-3", entry.LinesAdded, entry.LinesRemoved) + } + if entry.Entity != "/tmp/proj/src/app.py" { + t.Fatalf("entity = %q, want /tmp/proj/src/app.py (relative path resolved against cwd)", entry.Entity) + } + if entry.Language != "Python" { + t.Fatalf("language = %q, want Python", entry.Language) + } +} + +// TestSkipsRecordsWithoutUsage verifies that transcript noise — reasoning, +// file-history snapshots, tool results without rawUsage — produces nothing. +func TestSkipsRecordsWithoutUsage(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "proj", "session-a.jsonl") + providertest.WriteFile(t, path, + `{"timestamp":1785730645570,"type":"reasoning","content":"thinking"}`+"\n"+ + `{"timestamp":1785730645571,"type":"file-history-snapshot"}`+"\n"+ + `{"timestamp":1785730645572,"type":"function_call_result","providerData":{}}`+"\n"+ + `{"timestamp":1785730645573,"type":"message","providerData":{"messageId":"msg-0","rawUsage":{"prompt_tokens":0,"completion_tokens":0}}}`+"\n") + + entries, err := loadEntries([]string{dir}, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 0 { + t.Fatalf("entries = %d, want 0: %#v", len(entries), entries) + } +} diff --git a/internal/provider/workbuddy/provider.go b/internal/provider/workbuddy/provider.go new file mode 100644 index 0000000..d0824c9 --- /dev/null +++ b/internal/provider/workbuddy/provider.go @@ -0,0 +1,36 @@ +package workbuddy + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads WorkBuddy usage entries. +type Provider struct{ usageprovider.Base } + +var _ usageprovider.Provider = Provider{} + +// WithPaths returns a WorkBuddy provider configured with data roots. +func (p Provider) WithPaths(paths []string) usageprovider.Provider { + p.Base = usageprovider.NewBase(paths) + return p +} + +// Provider returns the WorkBuddy provider id. +func (Provider) Provider() usage.Provider { return usage.ProviderWorkbuddy } + +// WithFileFilter returns a WorkBuddy provider that skips source files the +// filter rejects. +func (p Provider) WithFileFilter(filter usage.FileFilter) usageprovider.Provider { + p.Base = p.WithFilterSet(filter) + return p +} + +// Entries loads normalized WorkBuddy usage entries, newest first. +func (p Provider) Entries() ([]usage.Entry, error) { + entries, err := loadEntries(p.Paths(), p.Filter()) + if err != nil { + return nil, err + } + return usageprovider.SortEntriesByTimestampDesc(entries), nil +} diff --git a/internal/provider/workbuddy/provider_test.go b/internal/provider/workbuddy/provider_test.go new file mode 100644 index 0000000..6c36219 --- /dev/null +++ b/internal/provider/workbuddy/provider_test.go @@ -0,0 +1,35 @@ +package workbuddy + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// TestLoadsEntry is the WorkBuddy smoke test: a minimal fixture shaped like a +// real function_call record must produce exactly one entry with the expected +// identity and token counts. +func TestLoadsEntry(t *testing.T) { + entries, err := func() ([]usage.Entry, error) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "Users-eren-workspace-eapil-eye", "session-a.jsonl") + providertest.WriteFile(t, path, `{"id":"rec-1","timestamp":1785730645570,"type":"function_call","providerData":{"messageId":"msg-1","model":"glm-5.2","requestModelId":"auto","rawUsage":{"prompt_tokens":30361,"completion_tokens":118,"total_tokens":30479,"completion_tokens_details":{"reasoning_tokens":47},"prompt_tokens_details":{"cached_tokens":14720},"prompt_cache_hit_tokens":14720,"cache_read_input_tokens":0,"cache_creation_input_tokens":0}},"sessionId":"session-a","cwd":"/Users/eren/workspace/eapil-eye"}`+"\n") + return Provider{}.WithPaths([]string{dir}).Entries() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ + Provider: usage.ProviderWorkbuddy, + Model: "glm-5.2", + SessionID: "session-a", + Project: "eapil-eye", + Tokens: usage.TokenUsage{ + InputTokens: 15641, + OutputTokens: 71, + CacheReadInputTokens: 14720, + ReasoningOutputTokens: 47, + TotalTokens: 30479, + }, + }) +} diff --git a/internal/usage/usage.go b/internal/usage/usage.go index dc226d4..25b4945 100644 --- a/internal/usage/usage.go +++ b/internal/usage/usage.go @@ -12,21 +12,22 @@ import ( type Provider string const ( - ProviderClaude Provider = "claude" - ProviderCodex Provider = "codex" - ProviderCopilot Provider = "copilot" - ProviderGemini Provider = "gemini" - ProviderKimi Provider = "kimi" - ProviderQwen Provider = "qwen" - ProviderOpenClaw Provider = "openclaw" - ProviderPi Provider = "pi" - ProviderAmp Provider = "amp" - ProviderDroid Provider = "droid" - ProviderKilo Provider = "kilo" - ProviderHermes Provider = "hermes" - ProviderCodebuff Provider = "codebuff" - ProviderOpenCode Provider = "opencode" - ProviderGoose Provider = "goose" + ProviderClaude Provider = "claude" + ProviderCodex Provider = "codex" + ProviderCopilot Provider = "copilot" + ProviderGemini Provider = "gemini" + ProviderKimi Provider = "kimi" + ProviderQwen Provider = "qwen" + ProviderOpenClaw Provider = "openclaw" + ProviderPi Provider = "pi" + ProviderAmp Provider = "amp" + ProviderDroid Provider = "droid" + ProviderKilo Provider = "kilo" + ProviderHermes Provider = "hermes" + ProviderCodebuff Provider = "codebuff" + ProviderOpenCode Provider = "opencode" + ProviderGoose Provider = "goose" + ProviderWorkbuddy Provider = "workbuddy" ) const UnknownLanguage = "Unknown" diff --git a/internal/usagescan/scanner.go b/internal/usagescan/scanner.go index 9f9dd80..c346882 100644 --- a/internal/usagescan/scanner.go +++ b/internal/usagescan/scanner.go @@ -25,6 +25,7 @@ import ( "github.com/tokitoki-dev/tokitoki-cli/internal/provider/opencode" "github.com/tokitoki-dev/tokitoki-cli/internal/provider/pi" "github.com/tokitoki-dev/tokitoki-cli/internal/provider/qwen" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/workbuddy" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" "github.com/tokitoki-dev/tokitoki-cli/internal/usagedb" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" @@ -69,6 +70,7 @@ func DefaultProviders() []usageprovider.Provider { codebuff.Provider{}, opencode.Provider{}, goose.Provider{}, + workbuddy.Provider{}, } } diff --git a/pkg/agentlib/agentlib.go b/pkg/agentlib/agentlib.go index ef1d0aa..7de9ced 100644 --- a/pkg/agentlib/agentlib.go +++ b/pkg/agentlib/agentlib.go @@ -84,6 +84,9 @@ const ( // ProviderGoose identifies Goose usage files. ProviderGoose Provider = "goose" + + // ProviderWorkbuddy identifies WorkBuddy usage files. + ProviderWorkbuddy Provider = "workbuddy" ) var ( @@ -490,5 +493,6 @@ func DefaultProviderDirs() map[Provider][]string { filepath.Join(home, "Library", "Application Support", "goose", "sessions", "sessions.db"), filepath.Join(home, ".local", "share", "Block", "goose", "sessions", "sessions.db"), }, + ProviderWorkbuddy: {filepath.Join(home, ".workbuddy")}, } } From 78805ecd79c75625d88edda9af7647054fd3efe4 Mon Sep 17 00:00:00 2001 From: akarineren Date: Wed, 5 Aug 2026 00:54:05 +0900 Subject: [PATCH 09/11] Add comprehensive tests for usage tracking across providers - Introduced tests for the Copilot provider, validating session store entries and usage events. - Added tests for the Gemini provider, ensuring correct handling of checkpoint formats and message IDs. - Implemented streaming tests for usage scanning, confirming that appended records are processed correctly. - Created tests for agent library functionality, verifying upload and scan behaviors without prior state. - Enhanced parallel sync tests to ensure efficient batching and upload of events during scanning. --- README.md | 2 +- cmd/tokitoki/main.go | 125 ++++++-- cmd/tokitoki/worker_loop_test.go | 52 ++++ internal/agentdata/agentdata.go | 55 +++- internal/cli/cli.go | 5 +- internal/provider/claude/loader.go | 93 +++--- internal/provider/claude/provider.go | 26 ++ internal/provider/claude/resume_test.go | 168 ++++++++++ internal/provider/copilot/events.go | 337 ++++++++++++++++++++ internal/provider/copilot/loader.go | 263 +++++++++++++++- internal/provider/copilot/store.go | 220 +++++++++++++ internal/provider/copilot/store_test.go | 147 +++++++++ internal/provider/gemini/loader.go | 360 +++++++++++++++++++--- internal/provider/gemini/loader_test.go | 196 ++++++++++++ internal/provider/gemini/provider_test.go | 12 +- internal/provider/kimi/loader.go | 18 +- internal/provider/kimi/provider.go | 10 + internal/provider/openclaw/loader.go | 4 + internal/provider/pi/loader.go | 18 +- internal/provider/pi/provider.go | 7 + internal/provider/qwen/loader.go | 18 +- internal/provider/qwen/provider.go | 7 + internal/store/lock.go | 5 - internal/usagedb/db.go | 167 +++++++++- internal/usagedb/db_test.go | 267 ++++++++++++++++ internal/usageprovider/provider.go | 40 ++- internal/usagescan/scanner.go | 130 +++++++- internal/usagescan/stream_test.go | 285 +++++++++++++++++ internal/usageupload/uploader.go | 66 +++- pkg/agentlib/agentlib.go | 197 ++++++++++-- pkg/agentlib/agentlib_test.go | 2 +- pkg/agentlib/split_test.go | 47 +++ pkg/agentlib/sync_parallel_test.go | 240 +++++++++++++++ 33 files changed, 3402 insertions(+), 187 deletions(-) create mode 100644 cmd/tokitoki/worker_loop_test.go create mode 100644 internal/provider/claude/resume_test.go create mode 100644 internal/provider/copilot/events.go create mode 100644 internal/provider/copilot/store.go create mode 100644 internal/provider/copilot/store_test.go create mode 100644 internal/provider/gemini/loader_test.go create mode 100644 internal/usagescan/stream_test.go create mode 100644 pkg/agentlib/split_test.go create mode 100644 pkg/agentlib/sync_parallel_test.go diff --git a/README.md b/README.md index ec75c86..c75677d 100644 --- a/README.md +++ b/README.md @@ -118,7 +118,7 @@ Normal runs and `service install` default to these provider roots: ```text claude=~/.claude codex=~/.codex -copilot=~/.copilot/otel +copilot=~/.copilot gemini=~/.gemini/tmp kimi=~/.kimi qwen=~/.qwen diff --git a/cmd/tokitoki/main.go b/cmd/tokitoki/main.go index de0ec3c..a93a26a 100644 --- a/cmd/tokitoki/main.go +++ b/cmd/tokitoki/main.go @@ -15,6 +15,7 @@ import ( "path/filepath" "sort" "strings" + "sync" "syscall" "time" @@ -27,6 +28,10 @@ import ( const ( defaultSyncInterval = 5 * time.Minute + // defaultUploadInterval paces the queue drain. It runs far more often + // than the scan because an empty queue costs one indexed query, and + // because it bounds how long a freshly scanned event waits to be sent. + defaultUploadInterval = 30 * time.Second // updateInterval paces the service worker's self-update checks. The // first check runs immediately after start, so a freshly installed or // relaunched service is current within one loop iteration. @@ -337,9 +342,10 @@ type workerFlags struct { // rather than the built-in defaults. Installed units only bake explicit // dirs into ExecStart: defaults must resolve from the service user's // home at run time, not the installer's at install time. - explicitDirs bool - interval time.Duration - checkUpdate bool + explicitDirs bool + interval time.Duration + uploadInterval time.Duration + checkUpdate bool } func runServiceWorker(args []string) int { @@ -352,39 +358,113 @@ func runServiceWorker(args []string) int { return runWorkerLoop(ctx, flags) } +// runWorkerLoop scans and uploads on independent schedules. +// +// The two halves share nothing but the local queue: scanning writes events +// into it, uploading drains them. Running them on one ticker meant an upload +// could not start until a scan finished, so a cold start spent its whole scan +// with events queued and the network idle. Apart they proceed at their own +// pace, and the upload ticker runs faster because draining a queue that is +// usually empty costs one indexed query. +// +// Neither half knows the other exists. A scan that fails does not stop queued +// events from being sent, and an upload that fails does not stop new events +// from being queued. func runWorkerLoop(ctx context.Context, flags workerFlags) int { logger := defaultLogger() - ticker := time.NewTicker(flags.interval) + + client, err := agentlib.New(agentlib.Options{Logger: logger}) + if err != nil { + logger.Error("tokitoki worker failed to start", "error", err) + return 1 + } + + // A successful self-update replaces the binary on disk while this process + // still runs the old code, so it stops every loop and exits for the + // service manager to restart. Cancelling here is what ends the scan and + // upload loops too. + workerCtx, stopWorkers := context.WithCancel(ctx) + defer stopWorkers() + + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + runIntervalLoop(workerCtx, flags.interval, func(context.Context) { + if err := client.Scan(agentlib.SyncOptions{ProviderDirs: flags.providerDirs}); err != nil { + logger.Error("tokitoki scan failed", "error", err) + } + }) + }() + + go func() { + defer wg.Done() + runIntervalLoop(workerCtx, flags.uploadInterval, func(runCtx context.Context) { + if err := client.Upload(runCtx); err != nil { + logger.Error("tokitoki upload failed", "error", err) + } + }) + }() + + runUpdateLoop(workerCtx, logger, flags.interval) + stopWorkers() + wg.Wait() + return 0 +} + +// runIntervalLoop runs work immediately and then every interval until ctx is +// done. Each run is bounded by its own timeout so one slow pass cannot stall +// the schedule forever. +func runIntervalLoop(ctx context.Context, interval time.Duration, work func(context.Context)) { + // NewTicker panics on a non-positive interval. A caller that never set one + // wants the default cadence, not a crashed worker. + if interval <= 0 { + interval = defaultSyncInterval + } + ticker := time.NewTicker(interval) defer ticker.Stop() + for { + runCtx, cancel := context.WithTimeout(ctx, agentlib.DefaultUploadTimeout) + work(runCtx) + cancel() + + select { + case <-ctx.Done(): + return + case <-ticker.C: + } + } +} + +// runUpdateLoop checks for a new binary until ctx is done or one is +// installed. The caller then exits so the service manager restarts into it. +func runUpdateLoop(ctx context.Context, logger *slog.Logger, interval time.Duration) { // Zero means "never checked", so the first iteration checks right away. var lastUpdateCheck time.Time + ticker := time.NewTicker(interval) + defer ticker.Stop() for { - syncCtx, cancel := context.WithTimeout(ctx, agentlib.DefaultUploadTimeout) - if err := runSync(syncCtx, flags.providerDirs, os.Stdout); err != nil { - logger.Error("tokitoki sync failed", "error", err) - } - cancel() - if time.Since(lastUpdateCheck) >= updateInterval { lastUpdateCheck = time.Now() - updateCtx, cancel := context.WithTimeout(ctx, updateTimeout) - result, err := selfupdate.Upgrade(updateCtx, logger, usageupload.BaseURL(), version) + checkCtx, cancel := context.WithTimeout(ctx, updateTimeout) + result, err := selfupdate.Upgrade(checkCtx, logger, usageupload.BaseURL(), version) cancel() if err != nil { logger.Warn("tokitoki self-update failed", "error", err) } else if result.Updated { // The binary on disk is new but this process is still the - // old code. Exit; the service manager restarts us as the + // old code. Stop; the service manager restarts us as the // new version. - return 0 + return } } select { case <-ctx.Done(): - return 0 + return case <-ticker.C: } } @@ -409,7 +489,8 @@ func parseWorkerFlags(name string, args []string) (workerFlags, bool) { flags.SetOutput(os.Stderr) providerDirs := newProviderDirFlags(agentlib.DefaultProviderDirs()) flags.Var(providerDirs, "provider-dir", "provider data directory to scan (provider=dir; repeatable)") - interval := flags.Duration("interval", defaultSyncInterval, "sync interval") + interval := flags.Duration("interval", defaultSyncInterval, "scan interval") + uploadInterval := flags.Duration("upload-interval", defaultUploadInterval, "queue drain interval") if err := flags.Parse(args); err != nil { return workerFlags{}, false } @@ -427,8 +508,9 @@ func parseWorkerFlags(name string, args []string) (workerFlags, bool) { return workerFlags{}, false } return workerFlags{ - providerDirs: dirs, - interval: *interval, + providerDirs: dirs, + interval: *interval, + uploadInterval: *uploadInterval, }, true } @@ -456,9 +538,10 @@ func parseServiceFlags(args []string) (workerFlags, bool, bool) { return workerFlags{}, false, false } return workerFlags{ - providerDirs: dirs, - explicitDirs: providerDirs.Explicit(), - interval: *interval, + providerDirs: dirs, + explicitDirs: providerDirs.Explicit(), + interval: *interval, + uploadInterval: defaultUploadInterval, }, !*system, true } diff --git a/cmd/tokitoki/worker_loop_test.go b/cmd/tokitoki/worker_loop_test.go new file mode 100644 index 0000000..745aa0b --- /dev/null +++ b/cmd/tokitoki/worker_loop_test.go @@ -0,0 +1,52 @@ +package main + +import ( + "context" + "sync/atomic" + "testing" + "time" +) + +// The two halves must tick on their own schedules: a fast loop keeps running +// while a slow one is still working. +func TestRunIntervalLoopsAreIndependent(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 600*time.Millisecond) + defer cancel() + + var slow, fast int32 + done := make(chan struct{}, 2) + + go func() { + runIntervalLoop(ctx, 10*time.Millisecond, func(context.Context) { + atomic.AddInt32(&slow, 1) + time.Sleep(250 * time.Millisecond) // a slow scan + }) + done <- struct{}{} + }() + go func() { + runIntervalLoop(ctx, 20*time.Millisecond, func(context.Context) { + atomic.AddInt32(&fast, 1) + }) + done <- struct{}{} + }() + + <-done + <-done + + s, f := atomic.LoadInt32(&slow), atomic.LoadInt32(&fast) + t.Logf("slow ran %d times, fast ran %d times", s, f) + if f <= s { + t.Errorf("fast loop ran %d times, slow %d: the fast loop was blocked by the slow one", f, s) + } +} + +// A zero interval must not panic the worker. +func TestRunIntervalLoopRejectsZeroInterval(t *testing.T) { + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + runs := 0 + runIntervalLoop(ctx, 0, func(context.Context) { runs++ }) + if runs == 0 { + t.Fatal("work never ran") + } +} diff --git a/internal/agentdata/agentdata.go b/internal/agentdata/agentdata.go index 04507e2..eff0aba 100644 --- a/internal/agentdata/agentdata.go +++ b/internal/agentdata/agentdata.go @@ -78,25 +78,64 @@ func CollectExt(root, ext string) []string { } func ReadJSONLines(path string, prefilter ...string) ([]LineJSON, error) { + lines, _, err := ReadJSONLinesFrom(path, 0, prefilter...) + return lines, err +} + +// ReadJSONLinesFrom parses a JSONL file starting at byte offset start and +// reports the offset to resume from next time. +// +// This exists for transcripts that are appended to while being read: an +// active session's file grows continuously, and re-reading megabytes of +// history to pick up the newest few lines is the cost resuming avoids. +// +// The returned offset is the end of the last line that arrived with its +// newline, never the end of the file. A trailing line without one is either +// the last line of a finished file or the front of one still being written, +// and the two are indistinguishable from here. It is parsed either way, so a +// file that merely lacks a final newline is not ignored, but the resume point +// stops before it: if more of it arrives later, the next pass re-reads the +// whole line. Re-reading one line costs nothing; skipping a real one loses it. +// +// The offset also advances past lines that fail to parse or that the +// prefilter rejects. Those are lines the file has moved beyond; stopping +// there would turn one malformed record into a permanent roadblock hiding +// everything after it. +// +// Only callers whose files are append-only may resume. A caller that reads a +// document rewritten in place, or that needs cross-line context from earlier +// in the file, must keep using ReadJSONLines. +func ReadJSONLinesFrom(path string, start int64, prefilter ...string) ([]LineJSON, int64, error) { file, err := os.Open(path) if errors.Is(err, os.ErrNotExist) { - return nil, nil + return nil, 0, nil } if err != nil { - return nil, err + return nil, 0, err } defer file.Close() + if start > 0 { + if _, err := file.Seek(start, io.SeekStart); err != nil { + return nil, 0, err + } + } + lines := make([]LineJSON, 0) reader := bufio.NewReader(file) lineNumber := 0 - offset := int64(0) + offset := start + consumed := start for { line, readErr := reader.ReadBytes('\n') + complete := readErr == nil if len(line) > 0 { lineNumber++ - start := offset + lineStart := offset offset += int64(len(line)) + if complete { + consumed = offset + } line = bytes.TrimRight(line, "\r\n") if matchesPrefilter(line, prefilter) { var value map[string]any @@ -106,21 +145,21 @@ func ReadJSONLines(path string, prefilter ...string) ([]LineJSON, error) { lines = append(lines, LineJSON{ Value: value, Line: lineNumber, - Start: start, + Start: lineStart, End: offset, }) } } } - if readErr == nil { + if complete { continue } if errors.Is(readErr, io.EOF) { break } - return nil, readErr + return nil, 0, readErr } - return lines, nil + return lines, consumed, nil } type LineJSON struct { diff --git a/internal/cli/cli.go b/internal/cli/cli.go index 02c02ad..21a75bd 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -58,8 +58,9 @@ func (a *App) GetAPIKey() error { // accumulating history instead of losing it. // // Callers that coordinate multiple processes call the two phases separately — -// Ingest under the data lock, Upload under the upload lock — so a slow drain -// never blocks another process's ingestion. +// Ingest under the data lock, Upload without one — so a slow drain never +// blocks another process's ingestion. Uploads need no lock of their own: +// concurrent drains claim different batches from the queue. func (a *App) Sync(ctx context.Context) error { if err := a.Ingest(); err != nil { return err diff --git a/internal/provider/claude/loader.go b/internal/provider/claude/loader.go index 8d1748a..1c5300f 100644 --- a/internal/provider/claude/loader.go +++ b/internal/provider/claude/loader.go @@ -232,15 +232,42 @@ func UsageFiles(paths []string, projectFilter string) []string { } func ReadUsageFile(path string) ([]LoadedEntry, error) { + entries, _, err := ReadUsageFileFrom(path, 0) + return entries, err +} + +// ReadUsageFileFrom parses a transcript starting at byte offset start and +// reports the offset to resume from next time. +// +// Transcripts are append-only and are read while Claude is still writing to +// them, so the returned offset is the end of the last line that arrived with +// its newline — never the end of the file. A trailing partial line is left +// unconsumed for the next pass, when the rest of it exists. +// +// The offset advances past lines that fail to parse. A line the parser cannot +// use is still a line the file has moved beyond; stopping there would turn one +// malformed record into a permanent roadblock hiding everything after it. +// +// Callers resuming mid-file lose the diff-to-message association for a diff +// written before start, which is a bounded, one-message cost. The alternative — +// re-reading from zero to rebuild it — is the whole expense this exists to +// avoid. +func ReadUsageFileFrom(path string, start int64) ([]LoadedEntry, int64, error) { file, err := os.Open(path) if errors.Is(err, os.ErrNotExist) { - return nil, nil + return nil, 0, nil } if err != nil { - return nil, err + return nil, 0, err } defer file.Close() + if start > 0 { + if _, err := file.Seek(start, io.SeekStart); err != nil { + return nil, 0, err + } + } + sessionID := ExtractSessionID(path) entries := make([]LoadedEntry, 0) // A file-modification diff belongs to the assistant message that issued @@ -249,18 +276,31 @@ func ReadUsageFile(path string) ([]LoadedEntry, error) { pending := make([]patchStats, 0) reader := bufio.NewReader(file) lineNumber := 0 - offset := int64(0) + offset := start + consumed := start for { line, readErr := reader.ReadBytes('\n') + // A line that arrived without its newline is either the last line of + // a finished file or the front of one still being written, and the + // two are indistinguishable from here. It is parsed either way, so a + // file that simply lacks a trailing newline is not ignored, but the + // resume point stops short of it: if more of it arrives later, the + // next pass re-reads the whole line and supersedes what this one + // produced. Re-reading one line costs nothing; skipping a real one + // loses it permanently. + complete := readErr == nil if len(line) > 0 { lineNumber++ - start := offset + lineStart := offset offset += int64(len(line)) + if complete { + consumed = offset + } line = bytes.TrimRight(line, "\r\n") if entry, ok := parseUsageLine(line, sessionID); ok { entry.SourceFile = path entry.SourceLine = lineNumber - entry.SourceStart = start + entry.SourceStart = lineStart entry.SourceEnd = offset entry.ID = stableEntryID(entry) entries = append(entries, entry) @@ -276,15 +316,15 @@ func ReadUsageFile(path string) ([]LoadedEntry, error) { } } } - if readErr == nil { + if complete { continue } if errors.Is(readErr, io.EOF) { break } - return nil, readErr + return nil, 0, readErr } - return entries, nil + return entries, consumed, nil } type patchStats struct { @@ -585,40 +625,18 @@ func languageFromPathsInText(text string) string { return langdetect.DominantFromPaths(paths) } +// stableEntryID keys an entry on message.id + requestId — matching ccusage. +// The same message is replayed across multiple session files (e.g. sidechains), +// so keying on anything file/session/timestamp-specific double-counts tokens. func stableEntryID(entry LoadedEntry) string { requestID := "" if entry.Data.RequestID != nil { requestID = *entry.Data.RequestID } - messageID := "" - if entry.Data.Message.ID != nil { - messageID = *entry.Data.Message.ID - } - - // Deduplicate on message.id + requestId only — matching ccusage. The same - // message is replayed across multiple session files (e.g. sidechains), so - // keying on anything file/session/timestamp-specific double-counts tokens. - if messageID != "" { - return usage.StableID( - string(usage.ProviderClaude), - messageID, - requestID, - ) - } - - // No message id: can't dedupe across files. Fall back to source position so - // the row at least stays stable for a given file. - tokens := entry.Data.Message.Usage return usage.StableID( string(usage.ProviderClaude), - entry.SourceFile, - strconv.Itoa(entry.SourceLine), - entry.Data.Timestamp, - entry.Model, - strconv.FormatUint(tokens.InputTokens, 10), - strconv.FormatUint(tokens.OutputTokens, 10), - strconv.FormatUint(tokens.CacheCreationInputTokens, 10), - strconv.FormatUint(tokens.CacheReadInputTokens, 10), + *entry.Data.Message.ID, + requestID, ) } @@ -632,7 +650,10 @@ func isValidUsageEntry(data UsageEntry) bool { if data.RequestID != nil && *data.RequestID == "" { return false } - if data.Message.ID != nil && *data.Message.ID == "" { + // A message id is required: it is what identifies the event. Without one + // there is nothing to deduplicate on, and the same message replayed across + // session files would be counted once per copy. + if data.Message.ID == nil || *data.Message.ID == "" { return false } if data.Message.Model != nil && *data.Message.Model == "" { diff --git a/internal/provider/claude/provider.go b/internal/provider/claude/provider.go index 33e5572..6511b06 100644 --- a/internal/provider/claude/provider.go +++ b/internal/provider/claude/provider.go @@ -45,3 +45,29 @@ func (p Provider) Entries() ([]usage.Entry, error) { }) return converted, nil } + +// StreamEntries parses each transcript from where the previous scan stopped +// and hands its entries to emit before moving on. +// +// Transcripts are append-only, so resuming at a byte offset reads only what +// was written since the last scan. An active session's transcript grows to +// tens of megabytes, and re-reading all of it every few minutes to pick up +// the newest few lines is the cost this avoids. +// +// Entries are not sorted here. Ordering is a presentation concern of Entries; +// what emit does with these is store them, keyed by id. +func (p Provider) StreamEntries(resume func(path string) int64, emit func(path string, entries []usage.Entry, offset int64) error) error { + for _, file := range UsageFiles(p.paths, "") { + if p.filter != nil && !p.filter(file) { + continue + } + loaded, offset, err := ReadUsageFileFrom(file, resume(file)) + if err != nil { + return err + } + if err := emit(file, ConvertEntries(loaded), offset); err != nil { + return err + } + } + return nil +} diff --git a/internal/provider/claude/resume_test.go b/internal/provider/claude/resume_test.go new file mode 100644 index 0000000..b76dcf7 --- /dev/null +++ b/internal/provider/claude/resume_test.go @@ -0,0 +1,168 @@ +package claude + +import ( + "os" + "path/filepath" + "testing" +) + +func usageLine(id string, input, output uint64) string { + return `{"timestamp":"2026-05-21T01:02:03Z","cwd":"/tmp/p","requestId":"req-` + id + + `","message":{"id":"msg-` + id + `","model":"claude","usage":{"input_tokens":` + + itoa(input) + `,"output_tokens":` + itoa(output) + `}}}` + "\n" +} + +func itoa(v uint64) string { + if v == 0 { + return "0" + } + var buf []byte + for v > 0 { + buf = append([]byte{byte('0' + v%10)}, buf...) + v /= 10 + } + return string(buf) +} + +func transcript(t *testing.T, body string) string { + t.Helper() + path := filepath.Join(t.TempDir(), "projects", "-tmp-p", "session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, body) + return path +} + +func appendTo(t *testing.T, path, data string) { + t.Helper() + fh, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o600) + if err != nil { + t.Fatal(err) + } + defer fh.Close() + if _, err := fh.WriteString(data); err != nil { + t.Fatal(err) + } +} + +// Resuming at the reported offset must produce exactly what a single pass +// produces. This is the property the incremental scan rests on. +func TestReadUsageFileFromResumeMatchesWholeRead(t *testing.T) { + path := transcript(t, usageLine("1", 1, 2)+usageLine("2", 3, 4)+usageLine("3", 5, 6)) + + whole, wholeOffset, err := ReadUsageFileFrom(path, 0) + if err != nil { + t.Fatal(err) + } + if len(whole) != 3 { + t.Fatalf("whole read = %d entries, want 3", len(whole)) + } + + first, firstOffset, err := ReadUsageFileFrom(path, 0) + if err != nil { + t.Fatal(err) + } + // Resume from after the first entry, as an interrupted pass would. + rest, restOffset, err := ReadUsageFileFrom(path, first[0].SourceEnd) + if err != nil { + t.Fatal(err) + } + if restOffset != wholeOffset || firstOffset != wholeOffset { + t.Fatalf("offsets = %d/%d, want %d", firstOffset, restOffset, wholeOffset) + } + if len(rest) != 2 { + t.Fatalf("resumed read = %d entries, want 2", len(rest)) + } + for i, entry := range rest { + if entry.ID != whole[i+1].ID { + t.Fatalf("resumed entry %d id = %q, want %q", i, entry.ID, whole[i+1].ID) + } + } +} + +// Appending to a transcript must cost only the appended bytes. +func TestReadUsageFileFromReadsOnlyTheAppendedTail(t *testing.T) { + path := transcript(t, usageLine("1", 1, 2)) + _, offset, err := ReadUsageFileFrom(path, 0) + if err != nil { + t.Fatal(err) + } + + appendTo(t, path, usageLine("2", 3, 4)) + tail, newOffset, err := ReadUsageFileFrom(path, offset) + if err != nil { + t.Fatal(err) + } + if len(tail) != 1 { + t.Fatalf("tail read = %d entries, want 1", len(tail)) + } + if tail[0].Data.Message.Usage.InputTokens != 3 { + t.Fatalf("tail input tokens = %d, want 3", tail[0].Data.Message.Usage.InputTokens) + } + if newOffset <= offset { + t.Fatalf("offset did not advance: %d -> %d", offset, newOffset) + } +} + +// A line still being written has no newline yet. It is parsed so a file that +// merely lacks a trailing newline is not ignored, but the resume offset stops +// before it so the completed line is read again. +func TestReadUsageFileFromDoesNotConsumePartialLine(t *testing.T) { + complete := usageLine("1", 1, 2) + path := transcript(t, complete+`{"timestamp":"2026-05-21T01:02:03Z","cwd":"/tmp/p","mess`) + + entries, offset, err := ReadUsageFileFrom(path, 0) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if offset != int64(len(complete)) { + t.Fatalf("offset = %d, want %d (partial line must not be consumed)", offset, len(complete)) + } + + // The rest of the line arrives; the next pass sees the whole record. + appendTo(t, path, `age":{"id":"msg-2","model":"claude","usage":{"input_tokens":9,"output_tokens":9}}}`+"\n") + rest, _, err := ReadUsageFileFrom(path, offset) + if err != nil { + t.Fatal(err) + } + if len(rest) != 1 || rest[0].Data.Message.Usage.InputTokens != 9 { + t.Fatalf("completed line not recovered: %+v", rest) + } +} + +// A file with no trailing newline at all is a finished file, not a partial +// write, and its last line must still be ingested. +func TestReadUsageFileFromParsesFinalLineWithoutNewline(t *testing.T) { + path := transcript(t, usageLine("1", 1, 2)+`{"timestamp":"2026-05-21T01:02:03Z","cwd":"/tmp/p","requestId":"req-2","message":{"id":"msg-2","model":"claude","usage":{"input_tokens":8,"output_tokens":8}}}`) + + entries, _, err := ReadUsageFileFrom(path, 0) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("entries = %d, want 2 (final line without newline must parse)", len(entries)) + } +} + +// A line that cannot be parsed must not become a permanent roadblock: the +// offset advances past it so later records stay reachable. +func TestReadUsageFileFromAdvancesPastUnparsableLine(t *testing.T) { + path := transcript(t, "{not json}\n"+usageLine("2", 4, 5)) + + entries, offset, err := ReadUsageFileFrom(path, 0) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + info, err := os.Stat(path) + if err != nil { + t.Fatal(err) + } + if offset != info.Size() { + t.Fatalf("offset = %d, want %d (bad line must not stall the scan)", offset, info.Size()) + } +} diff --git a/internal/provider/copilot/events.go b/internal/provider/copilot/events.go new file mode 100644 index 0000000..f282e39 --- /dev/null +++ b/internal/provider/copilot/events.go @@ -0,0 +1,337 @@ +package copilot + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + "time" + + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/langdetect" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// Each Copilot CLI session keeps a transcript in +// ~/.copilot/session-state//events.jsonl. The session-store database +// is the source of truth for tokens; the transcript contributes what the +// database lacks: which files were modified and how many lines changed, plus +// the working directory for sessions missing from the sessions table. +// +// Everything the transcript contributes is timestamped rather than keyed by +// turn: the transcript's turnId counts agent-loop iterations while the +// database's turn_index counts user messages, so the two numberings never +// align and time is the only shared axis. + +// Weights for the paths a session touched, mirroring the other providers: +// writing a file says the most about what is being worked on, reading less, +// and a path merely mentioned in a shell command or search pattern least. +const ( + writeWeight = 4 + readWeight = 2 + textWeight = 1 +) + +type timedChange struct { + timestamp time.Time + change usage.FileChange +} + +type timedCandidate struct { + timestamp time.Time + candidate langdetect.Candidate +} + +type sessionContext struct { + cwd string + gitRoot string + branch string + + changes []timedChange + candidates []timedCandidate + // shutdownChanges holds the session summary's modified files, used only + // for paths no tool telemetry already reported. + shutdownChanges []usage.FileChange +} + +func (c *sessionContext) projectDir() string { + if c == nil { + return "" + } + if c.gitRoot != "" { + return c.gitRoot + } + return c.cwd +} + +func (c *sessionContext) addCandidate(timestamp time.Time, candidate langdetect.Candidate) { + c.candidates = append(c.candidates, timedCandidate{timestamp: timestamp, candidate: candidate}) +} + +// loadSessionContexts parses every session transcript under the given +// session-state directories. Transcripts the filter rejects are skipped: the +// usage entries they would have enriched are already ingested. +func loadSessionContexts(stateDirs []string, filter usage.FileFilter) map[string]*sessionContext { + contexts := make(map[string]*sessionContext) + for _, dir := range stateDirs { + sessions, err := os.ReadDir(dir) + if err != nil { + continue + } + for _, session := range sessions { + if !session.IsDir() { + continue + } + path := filepath.Join(dir, session.Name(), "events.jsonl") + if filter != nil && !filter(path) { + continue + } + context := parseSessionEvents(path) + if context == nil { + continue + } + contexts[session.Name()] = context + } + } + return contexts +} + +func parseSessionEvents(path string) *sessionContext { + lines, err := agentdata.ReadJSONLines(path) + if err != nil || len(lines) == 0 { + return nil + } + + context := &sessionContext{} + // tool.execution_start carries the tool name and arguments; + // tool.execution_complete carries the telemetry. The call id joins them. + toolNames := make(map[string]string) + + for _, line := range lines { + record := line.Value + data := agentdata.ObjectAt(record["data"]) + if data == nil { + continue + } + timestamp, _ := agentdata.ParseTimestamp(record["timestamp"]) + switch agentdata.StringField(record, "type") { + case "session.start": + handleSessionStart(data, context) + case "tool.execution_start": + handleToolStart(data, timestamp, context, toolNames) + case "tool.execution_complete": + handleToolComplete(data, timestamp, context, toolNames) + case "session.shutdown": + handleShutdown(data, context) + } + } + return context +} + +func handleSessionStart(data map[string]any, context *sessionContext) { + block := agentdata.ObjectAt(data["context"]) + if block == nil { + return + } + context.cwd = agentdata.FirstNonEmpty(agentdata.StringField(block, "cwd"), context.cwd) + context.gitRoot = agentdata.FirstNonEmpty(agentdata.StringField(block, "gitRoot"), context.gitRoot) + context.branch = agentdata.FirstNonEmpty(agentdata.StringField(block, "branch"), context.branch) +} + +func handleToolStart(data map[string]any, timestamp time.Time, context *sessionContext, toolNames map[string]string) { + callID := agentdata.StringField(data, "toolCallId") + name := agentdata.FirstStringField(data, "toolName", "name") + if callID != "" && name != "" { + toolNames[callID] = name + } + for _, candidate := range argumentCandidates(agentdata.ObjectAt(data["arguments"]), context.cwd) { + context.addCandidate(timestamp, candidate) + } +} + +// argumentCandidates extracts language evidence from a tool's arguments: the +// file a read names outright, or the paths a command or pattern mentions. +func argumentCandidates(arguments map[string]any, cwd string) []langdetect.Candidate { + if arguments == nil { + return nil + } + candidates := make([]langdetect.Candidate, 0) + if path := agentdata.StringField(arguments, "path"); path != "" { + candidates = append(candidates, langdetect.Candidate{Path: usage.ResolvePath(cwd, path), Weight: readWeight}) + } + for _, value := range agentdata.ArrayAt(arguments["paths"]) { + if path := agentdata.StringValue(value); path != "" { + candidates = append(candidates, langdetect.Candidate{Path: usage.ResolvePath(cwd, path), Weight: readWeight}) + } + } + for _, key := range []string{"command", "pattern", "query"} { + for _, path := range langdetect.PathsFromText(agentdata.StringField(arguments, key)) { + candidates = append(candidates, langdetect.Candidate{Path: path, Weight: textWeight}) + } + } + return candidates +} + +func handleToolComplete(data map[string]any, timestamp time.Time, context *sessionContext, toolNames map[string]string) { + if success, ok := data["success"].(bool); ok && !success { + return + } + callID := agentdata.StringField(data, "toolCallId") + tool := agentdata.FirstStringField(data, "toolName", "name") + if tool == "" { + tool = toolNames[callID] + } + + telemetry := toolTelemetry(data) + if telemetry == nil || !hasWriteSignals(tool, telemetry) { + return + } + paths := telemetryFilePaths(telemetry) + if len(paths) == 0 { + return + } + blocks := telemetryCodeBlocks(telemetry, len(paths)) + for i, path := range paths { + if !trackablePath(path) { + continue + } + change := usage.FileChange{Path: usage.ResolvePath(context.cwd, path)} + if blocks != nil { + change.LinesAdded = blocks[i].added + change.LinesRemoved = blocks[i].removed + } + context.changes = append(context.changes, timedChange{timestamp: timestamp, change: change}) + context.addCandidate(timestamp, langdetect.Candidate{Path: change.Path, Weight: writeWeight}) + } +} + +// handleShutdown records the session summary's code changes so files that no +// tool telemetry reported still surface, on CLI versions without it. +func handleShutdown(data map[string]any, context *sessionContext) { + changes := agentdata.ObjectAt(data["codeChanges"]) + if changes == nil { + return + } + files := stringValues(changes["filesModified"]) + added := agentdata.UintField(changes, "linesAdded") + removed := agentdata.UintField(changes, "linesRemoved") + for _, path := range files { + if !trackablePath(path) { + continue + } + change := usage.FileChange{Path: usage.ResolvePath(context.cwd, path)} + // The summary totals are per session, not per file; they are only + // attributable when a single file changed. + if len(files) == 1 { + change.LinesAdded = added + change.LinesRemoved = removed + } + context.shutdownChanges = append(context.shutdownChanges, change) + } +} + +// toolTelemetry finds the telemetry block, which moved between CLI versions: +// toolTelemetry, then toolResultTelemetry, then result.toolTelemetry. +func toolTelemetry(data map[string]any) map[string]any { + if telemetry := agentdata.ObjectAt(data["toolTelemetry"]); telemetry != nil { + return telemetry + } + if telemetry := agentdata.ObjectAt(data["toolResultTelemetry"]); telemetry != nil { + return telemetry + } + return agentdata.ObjectAt(agentdata.ObjectAt(data["result"])["toolTelemetry"]) +} + +func hasWriteSignals(tool string, telemetry map[string]any) bool { + switch tool { + case "write", "edit", "apply_patch", "str_replace_editor", "create": + return true + } + metrics := agentdata.ObjectAt(telemetry["metrics"]) + if metrics != nil && (metrics["linesAdded"] != nil || metrics["linesRemoved"] != nil) { + return true + } + restricted := agentdata.ObjectAt(telemetry["restrictedProperties"]) + return len(stringValues(restricted["addedPaths"])) > 0 || + len(stringValues(restricted["deletedPaths"])) > 0 +} + +func telemetryFilePaths(telemetry map[string]any) []string { + restricted := agentdata.ObjectAt(telemetry["restrictedProperties"]) + if restricted == nil { + return nil + } + paths := stringValues(restricted["filePaths"]) + if len(paths) == 0 { + paths = append(paths, stringValues(restricted["addedPaths"])...) + paths = append(paths, stringValues(restricted["deletedPaths"])...) + } + return agentdata.UniqueStrings(paths) +} + +type lineCounts struct { + added uint64 + removed uint64 +} + +// telemetryCodeBlocks aligns per-file line counts with the file list. The +// codeBlocks array matches when the tool reported one block per file; the +// flat metrics only apply when a single file changed. +func telemetryCodeBlocks(telemetry map[string]any, pathCount int) []lineCounts { + blocks := arrayValues(agentdata.ObjectAt(telemetry["properties"])["codeBlocks"]) + if len(blocks) == pathCount { + counts := make([]lineCounts, 0, len(blocks)) + for _, block := range blocks { + counts = append(counts, lineCounts{ + added: agentdata.UintField(agentdata.ObjectAt(block), "linesAdded"), + removed: agentdata.UintField(agentdata.ObjectAt(block), "linesRemoved"), + }) + } + return counts + } + metrics := agentdata.ObjectAt(telemetry["metrics"]) + if pathCount == 1 && metrics != nil && (metrics["linesAdded"] != nil || metrics["linesRemoved"] != nil) { + return []lineCounts{{ + added: agentdata.UintField(metrics, "linesAdded"), + removed: agentdata.UintField(metrics, "linesRemoved"), + }} + } + return nil +} + +// trackablePath rejects the CLI's own bookkeeping files, like the plan.md it +// writes under its session-state directory. +func trackablePath(path string) bool { + if strings.TrimSpace(path) == "" { + return false + } + return !strings.Contains(filepath.ToSlash(path), ".copilot/session-state/") +} + +// stringValues decodes a value that is either a JSON array of strings or a +// string holding an encoded JSON array, which is how telemetry properties +// arrive. +func stringValues(value any) []string { + values := make([]string, 0) + for _, entry := range arrayValues(value) { + if text := agentdata.StringValue(entry); text != "" { + values = append(values, text) + } + } + return values +} + +func arrayValues(value any) []any { + switch typed := value.(type) { + case []any: + return typed + case string: + var decoded []any + if err := json.Unmarshal([]byte(typed), &decoded); err != nil { + return nil + } + return decoded + default: + return nil + } +} diff --git a/internal/provider/copilot/loader.go b/internal/provider/copilot/loader.go index 42c0b84..95e7ad1 100644 --- a/internal/provider/copilot/loader.go +++ b/internal/provider/copilot/loader.go @@ -2,37 +2,290 @@ package copilot import ( "fmt" + "os" "path/filepath" "sort" + "strconv" "strings" "time" "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" + "github.com/tokitoki-dev/tokitoki-cli/internal/langdetect" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) +// dataRoots resolves what to scan from the configured paths. A Copilot data +// root holds up to three sources: otel/ with the OpenTelemetry export older +// CLI versions wrote, session-store.db with the usage database current +// versions write, and session-state/ with the per-session transcripts. +type dataRoots struct { + otelDirs []string + storeDBs []string + stateDirs []string +} + +func resolveDataRoots(paths []string) dataRoots { + roots := dataRoots{} + for _, path := range paths { + info, err := os.Stat(path) + if err != nil { + continue + } + if !info.IsDir() { + // A file path selects one source directly: a database or one + // OpenTelemetry export. + if strings.EqualFold(filepath.Ext(path), ".db") { + roots.storeDBs = append(roots.storeDBs, path) + } else { + roots.otelDirs = append(roots.otelDirs, path) + } + continue + } + // Any configured directory is scanned for OpenTelemetry exports, as + // it always was. Deployed configurations point at ~/.copilot/otel; + // the store and transcripts live next to it. + roots.otelDirs = append(roots.otelDirs, path) + root := path + if filepath.Base(path) == "otel" { + root = filepath.Dir(path) + } + roots.storeDBs = append(roots.storeDBs, filepath.Join(root, "session-store.db")) + roots.stateDirs = append(roots.stateDirs, filepath.Join(root, "session-state")) + } + roots.otelDirs = agentdata.UniqueStrings(roots.otelDirs) + roots.storeDBs = agentdata.UniqueStrings(roots.storeDBs) + roots.stateDirs = agentdata.UniqueStrings(roots.stateDirs) + return roots +} + +func isSessionStateFile(path string) bool { + return strings.Contains(filepath.ToSlash(path), "/session-state/") +} + func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + roots := resolveDataRoots(paths) + + storeEvents := make([]storeEvent, 0) + storeSessions := make(map[string]storeSession) + for _, dbPath := range roots.storeDBs { + if !agentdb.ExistingSQLiteFile(dbPath) { + continue + } + if filter != nil && !filter(dbPath) { + continue + } + events, sessions, err := loadStoreDatabase(dbPath) + if err != nil { + return nil, err + } + storeEvents = append(storeEvents, events...) + for id, session := range sessions { + storeSessions[id] = session + } + } + + contexts := loadSessionContexts(roots.stateDirs, filter) + entries := buildStoreEntries(storeEvents, storeSessions, contexts) + + // The OpenTelemetry export from older CLI versions. Entries the store + // already accounts for are dropped: a CLI that writes both sources + // records the same API call twice. + seen := storeFingerprints(storeEvents) files := make([]string, 0) - for _, root := range paths { - files = append(files, agentdata.CollectExt(root, ".jsonl")...) + for _, root := range roots.otelDirs { + for _, file := range agentdata.CollectExt(root, ".jsonl") { + // The session transcripts are .jsonl too, but they are not + // OpenTelemetry exports. + if isSessionStateFile(file) { + continue + } + files = append(files, file) + } } sort.Strings(files) files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) - - entries := make([]usage.Entry, 0) for _, file := range files { fileEntries, err := parseOTELFile(file) if err != nil { return nil, err } - entries = append(entries, fileEntries...) + for _, entry := range fileEntries { + if seen.matches(entry.SessionID, entry.Model, entry.Usage, entry.Timestamp) { + continue + } + entries = append(entries, entry) + } } usageprovider.SortEntries(entries) return entries, nil } +// buildStoreEntries turns the usage database rows into entries. The sessions +// table names the directory and branch each session ran in; the transcript +// contributes the file changes and language evidence, each attached to the +// session's last API call at or before its own timestamp — a diff follows +// the request that made it. +func buildStoreEntries(events []storeEvent, sessions map[string]storeSession, contexts map[string]*sessionContext) []usage.Entry { + entries := make([]usage.Entry, 0, len(events)) + bySession := make(map[string][]int) + for i, event := range events { + session := sessions[event.sessionID] + context := contexts[event.sessionID] + + cwd := agentdata.FirstNonEmpty(session.cwd, context.projectDir()) + project, projectPath := usage.UnknownProject, "" + if path, name, ok := usage.ProjectFromCWD(cwd); ok { + project, projectPath = name, path + } + + entry := usageprovider.BaseEntry( + usage.ProviderCopilot, + event.timestamp, + project, + projectPath, + event.sessionID, + event.model, + "GitHub Copilot CLI", + event.tokens, + ) + entry.Branch = agentdata.FirstNonEmpty(session.branch, branchOf(context)) + entry.ID = storeEntryID(event) + entries = append(entries, entry) + bySession[event.sessionID] = append(bySession[event.sessionID], i) + } + + for sessionID, indexes := range bySession { + context := contexts[sessionID] + if context == nil { + continue + } + attachSessionContext(entries, indexes, context) + } + return entries +} + +// attachSessionContext distributes one session's transcript evidence over its +// entries. indexes point into entries in database row order, which is time +// order. +func attachSessionContext(entries []usage.Entry, indexes []int, context *sessionContext) { + // entryAt finds the entry an event at the given time belongs to: the last + // API call at or before it, or the first call for evidence that precedes + // the whole session's usage. + entryAt := func(timestamp time.Time) int { + target := indexes[0] + for _, index := range indexes { + if entries[index].Timestamp.After(timestamp) { + break + } + target = index + } + return target + } + + seen := make(map[string]bool) + for _, change := range context.changes { + entries[entryAt(change.timestamp)].ApplyFileChange(change.change) + seen[change.change.Path] = true + } + // The shutdown summary covers files no telemetry reported; they land on + // the session's last entry, the closest one to shutdown. + last := &entries[indexes[len(indexes)-1]] + for _, change := range context.shutdownChanges { + if seen[change.Path] { + continue + } + last.ApplyFileChange(change) + } + + // Language: each entry judges the evidence recorded up to its own API + // call; an entry without evidence of its own inherits the session-wide + // verdict. + sessionWide := make([]langdetect.Candidate, 0, len(context.candidates)) + perEntry := make(map[int][]langdetect.Candidate) + for _, candidate := range context.candidates { + sessionWide = append(sessionWide, candidate.candidate) + index := entryAt(candidate.timestamp) + perEntry[index] = append(perEntry[index], candidate.candidate) + } + fallback := langdetect.Dominant(sessionWide) + for _, index := range indexes { + language := langdetect.Dominant(perEntry[index]) + if language == langdetect.Unknown { + language = fallback + } + entries[index].Language = usage.NormalizeLanguage(language) + } +} + +func branchOf(context *sessionContext) string { + if context == nil { + return "" + } + return context.branch +} + +// storeEntryID identifies a database row by its content, not its position: +// the store has no cross-machine row identity, and content plus timestamp is +// stable across re-scans. +func storeEntryID(event storeEvent) string { + return usage.StableID( + string(usage.ProviderCopilot), + "store-event", + event.sessionID, + event.turnKey, + event.model, + event.timestamp.UTC().Format(time.RFC3339Nano), + strconv.FormatUint(event.tokens.InputTokens, 10), + strconv.FormatUint(event.tokens.OutputTokens, 10), + strconv.FormatUint(event.tokens.CacheCreationInputTokens, 10), + strconv.FormatUint(event.tokens.CacheReadInputTokens, 10), + strconv.FormatUint(event.tokens.ReasoningOutputTokens, 10), + ) +} + +// fingerprintSet answers whether the store already recorded an API call an +// OpenTelemetry record describes: same session, model and token counts, +// within a two-second window of each other. +type fingerprintSet map[string][]time.Time + +func storeFingerprints(events []storeEvent) fingerprintSet { + set := make(fingerprintSet, len(events)) + for _, event := range events { + key := fingerprintKey(event.sessionID, event.model, event.tokens) + set[key] = append(set[key], event.timestamp) + } + return set +} + +func (s fingerprintSet) matches(sessionID, model string, tokens usage.TokenUsage, timestamp time.Time) bool { + const tolerance = 2 * time.Second + for _, candidate := range s[fingerprintKey(sessionID, model, tokens)] { + delta := candidate.Sub(timestamp) + if delta < 0 { + delta = -delta + } + if delta <= tolerance { + return true + } + } + return false +} + +func fingerprintKey(sessionID, model string, tokens usage.TokenUsage) string { + return strings.Join([]string{ + sessionID, + model, + strconv.FormatUint(tokens.InputTokens, 10), + strconv.FormatUint(tokens.OutputTokens, 10), + strconv.FormatUint(tokens.CacheCreationInputTokens, 10), + strconv.FormatUint(tokens.CacheReadInputTokens, 10), + strconv.FormatUint(tokens.ReasoningOutputTokens, 10), + }, "|") +} + type sourceKind int const ( diff --git a/internal/provider/copilot/store.go b/internal/provider/copilot/store.go new file mode 100644 index 0000000..1cf46ac --- /dev/null +++ b/internal/provider/copilot/store.go @@ -0,0 +1,220 @@ +package copilot + +import ( + "database/sql" + "encoding/json" + "strconv" + "strings" + "time" + + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdb" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// The Copilot CLI stopped exporting OpenTelemetry files and now records every +// API call in ~/.copilot/session-store.db: assistant_usage_events holds the +// per-call token counts and sessions holds the working directory and branch +// the session ran in. + +// storeSession is one row of the sessions table: where a session ran. +type storeSession struct { + cwd string + repository string + branch string +} + +// storeEvent is one row of assistant_usage_events: one API call's usage. +type storeEvent struct { + sessionID string + turnKey string + model string + timestamp time.Time + tokens usage.TokenUsage +} + +func loadStoreDatabase(path string) ([]storeEvent, map[string]storeSession, error) { + db, err := agentdb.OpenSQLite(path) + if err != nil { + // A missing or unreadable store is not an error: the CLI version in + // use may simply predate it. + return nil, nil, nil + } + defer db.Close() + + sessions := queryStoreSessions(db) + events := queryStoreEvents(db) + return events, sessions, nil +} + +func queryStoreSessions(db *sql.DB) map[string]storeSession { + sessions := make(map[string]storeSession) + rows, err := db.Query(`SELECT id, COALESCE(cwd, ''), COALESCE(repository, ''), COALESCE(branch, '') FROM sessions`) + if err != nil { + return sessions + } + defer rows.Close() + for rows.Next() { + var id, cwd, repository, branch string + if err := rows.Scan(&id, &cwd, &repository, &branch); err != nil { + continue + } + sessions[id] = storeSession{ + cwd: strings.TrimSpace(cwd), + repository: strings.TrimSpace(repository), + branch: strings.TrimSpace(branch), + } + } + return sessions +} + +func queryStoreEvents(db *sql.DB) []storeEvent { + rows, err := db.Query(` + SELECT + session_id, + COALESCE(turn_index, -1), + model, + COALESCE(input_tokens, 0), + COALESCE(output_tokens, 0), + COALESCE(cache_read_tokens, 0), + COALESCE(cache_write_tokens, 0), + COALESCE(reasoning_tokens, 0), + COALESCE(token_details_json, ''), + COALESCE(created_at, '') + FROM assistant_usage_events + ORDER BY id ASC`) + if err != nil { + return nil + } + defer rows.Close() + + events := make([]storeEvent, 0) + for rows.Next() { + var sessionID, model, details, createdAt string + var turnIndex, input, output, cacheRead, cacheWrite, reasoning int64 + if err := rows.Scan(&sessionID, &turnIndex, &model, &input, &output, &cacheRead, &cacheWrite, &reasoning, &details, &createdAt); err != nil { + continue + } + timestamp, ok := parseStoreTimestamp(createdAt) + if !ok { + continue + } + tokens := normalizeStoreTokens(input, output, cacheRead, cacheWrite, reasoning, details) + if !nonZeroTokens(tokens) { + continue + } + turnKey := "" + if turnIndex >= 0 { + turnKey = strconv.FormatInt(turnIndex, 10) + } + events = append(events, storeEvent{ + sessionID: strings.TrimSpace(sessionID), + turnKey: turnKey, + model: strings.TrimSpace(model), + timestamp: timestamp, + tokens: tokens, + }) + } + return events +} + +// normalizeStoreTokens splits the raw counters into disjoint buckets. The +// store's input_tokens includes the cached reads and writes, and its +// output_tokens includes the reasoning tokens, so summing the raw columns +// would bill the overlaps twice. token_details_json carries the exact split +// when it agrees with the raw totals; otherwise the overlaps are subtracted. +func normalizeStoreTokens(input, output, cacheRead, cacheWrite, reasoning int64, detailsJSON string) usage.TokenUsage { + inputRaw := clampUint(input) + outputRaw := clampUint(output) + cacheReadRaw := clampUint(cacheRead) + cacheWriteRaw := clampUint(cacheWrite) + reasoningRaw := clampUint(reasoning) + + tokens := usage.TokenUsage{} + if details, ok := parseTokenDetails(detailsJSON); ok && + details.input+details.cacheRead+details.cacheWrite == inputRaw && + details.output == outputRaw { + tokens.InputTokens = details.input + tokens.CacheReadInputTokens = details.cacheRead + tokens.CacheCreationInputTokens = details.cacheWrite + tokens.OutputTokens = details.output + } else { + read := min(cacheReadRaw, inputRaw) + write := min(cacheWriteRaw, inputRaw-read) + tokens.InputTokens = inputRaw - read - write + tokens.CacheReadInputTokens = read + tokens.CacheCreationInputTokens = write + tokens.OutputTokens = outputRaw + } + + tokens.ReasoningOutputTokens = min(reasoningRaw, tokens.OutputTokens) + tokens.OutputTokens -= tokens.ReasoningOutputTokens + tokens.TotalTokens = tokens.InputTokens + + tokens.CacheReadInputTokens + + tokens.CacheCreationInputTokens + + tokens.OutputTokens + + tokens.ReasoningOutputTokens + return tokens +} + +type tokenDetails struct { + input, output, cacheRead, cacheWrite uint64 +} + +// parseTokenDetails reads token_details_json, a list of +// {"tokenType": "input", "tokenCount": 123, ...} objects. +func parseTokenDetails(raw string) (tokenDetails, bool) { + raw = strings.TrimSpace(raw) + if raw == "" { + return tokenDetails{}, false + } + var rows []struct { + TokenType string `json:"tokenType"` + TokenCount json.Number `json:"tokenCount"` + } + if err := json.Unmarshal([]byte(raw), &rows); err != nil || len(rows) == 0 { + return tokenDetails{}, false + } + details := tokenDetails{} + for _, row := range rows { + count := agentdata.UintValue(row.TokenCount) + switch row.TokenType { + case "input": + details.input += count + case "output": + details.output += count + case "cache_read": + details.cacheRead += count + case "cache_write", "cache_creation": + details.cacheWrite += count + default: + // An unknown bucket means the sum check below cannot be trusted. + return tokenDetails{}, false + } + } + return details, true +} + +func parseStoreTimestamp(raw string) (time.Time, bool) { + if timestamp, ok := agentdata.ParseTimestampString(raw); ok { + return timestamp, true + } + // SQLite's datetime('now') default writes "2006-01-02 15:04:05" in UTC. + if timestamp, err := time.Parse("2006-01-02 15:04:05", strings.TrimSpace(raw)); err == nil { + return timestamp.UTC(), true + } + return time.Time{}, false +} + +func clampUint(value int64) uint64 { + if value < 0 { + return 0 + } + return uint64(value) +} + +func nonZeroTokens(tokens usage.TokenUsage) bool { + return tokens.InputTokens+tokens.OutputTokens+ + tokens.CacheReadInputTokens+tokens.CacheCreationInputTokens+ + tokens.ReasoningOutputTokens > 0 +} diff --git a/internal/provider/copilot/store_test.go b/internal/provider/copilot/store_test.go new file mode 100644 index 0000000..d730be0 --- /dev/null +++ b/internal/provider/copilot/store_test.go @@ -0,0 +1,147 @@ +package copilot + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func writeStoreFixture(t *testing.T, dir string) string { + t.Helper() + dbPath := filepath.Join(dir, "session-store.db") + db := providertest.OpenTestSQLite(t, dbPath) + defer db.Close() + + statements := []string{ + `CREATE TABLE sessions ( + id TEXT PRIMARY KEY, + cwd TEXT, + repository TEXT, + host_type TEXT, + branch TEXT, + summary TEXT, + created_at TEXT, + updated_at TEXT + )`, + `CREATE TABLE assistant_usage_events ( + id INTEGER PRIMARY KEY AUTOINCREMENT, + session_id TEXT NOT NULL, + turn_index INTEGER, + agent_id TEXT, + parent_tool_call_id TEXT, + model TEXT NOT NULL, + input_tokens INTEGER, + output_tokens INTEGER, + cache_read_tokens INTEGER, + cache_write_tokens INTEGER, + reasoning_tokens INTEGER, + token_details_json TEXT, + created_at TEXT + )`, + `INSERT INTO sessions (id, cwd, repository, branch) VALUES + ('session-1', '/home/dev/widget', 'dev/widget', 'main')`, + `INSERT INTO assistant_usage_events + (session_id, turn_index, model, input_tokens, output_tokens, cache_read_tokens, cache_write_tokens, reasoning_tokens, token_details_json, created_at) + VALUES ('session-1', 0, 'mai-code-1-flash', 30483, 289, 7296, 0, 64, + '[{"tokenCount":23187,"tokenType":"input"},{"tokenCount":7296,"tokenType":"cache_read"},{"tokenCount":289,"tokenType":"output"}]', + '2026-08-04T01:19:49.344Z')`, + } + for _, statement := range statements { + if _, err := db.Exec(statement); err != nil { + t.Fatal(err) + } + } + return dbPath +} + +// TestLoadsStoreEntry proves the session-store database produces an entry +// with disjoint token buckets and the project taken from the sessions table. +func TestLoadsStoreEntry(t *testing.T) { + dir := t.TempDir() + writeStoreFixture(t, dir) + + entries, err := Provider{}.WithPaths([]string{dir}).Entries() + + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ + Provider: usage.ProviderCopilot, + Model: "mai-code-1-flash", + SessionID: "session-1", + Project: "widget", + Tokens: usage.TokenUsage{ + InputTokens: 23187, + OutputTokens: 225, + CacheReadInputTokens: 7296, + ReasoningOutputTokens: 64, + TotalTokens: 30772, + }, + }) + if entries[0].Branch != "main" { + t.Fatalf("branch = %q, want main", entries[0].Branch) + } + if entries[0].ProjectPath != "/home/dev/widget" { + t.Fatalf("project path = %q, want /home/dev/widget", entries[0].ProjectPath) + } +} + +// TestStoreEntryFileChanges proves a turn's write telemetry from the session +// transcript lands on that turn's usage entry. +func TestStoreEntryFileChanges(t *testing.T) { + dir := t.TempDir() + writeStoreFixture(t, dir) + events := filepath.Join(dir, "session-state", "session-1", "events.jsonl") + providertest.WriteFile(t, events, + `{"type":"session.start","data":{"sessionId":"session-1","copilotVersion":"1.0.73","context":{"cwd":"/home/dev/widget","gitRoot":"/home/dev/widget","branch":"main"}},"timestamp":"2026-08-04T01:19:38.035Z"}`+"\n"+ + `{"type":"tool.execution_start","data":{"toolCallId":"call-1","toolName":"edit","turnId":"0","arguments":{"path":"/home/dev/widget/main.go"}},"timestamp":"2026-08-04T01:19:40.000Z"}`+"\n"+ + `{"type":"tool.execution_complete","data":{"toolCallId":"call-1","toolName":"edit","turnId":"0","success":true,"toolTelemetry":{"properties":{},"restrictedProperties":{"filePaths":"[\"/home/dev/widget/main.go\"]"},"metrics":{"linesAdded":12,"linesRemoved":3}}},"timestamp":"2026-08-04T01:19:41.000Z"}`+"\n") + + entries, err := Provider{}.WithPaths([]string{dir}).Entries() + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + entry := entries[0] + if entry.LinesAdded != 12 || entry.LinesRemoved != 3 { + t.Fatalf("lines = +%d/-%d, want +12/-3", entry.LinesAdded, entry.LinesRemoved) + } + if entry.Entity != "/home/dev/widget/main.go" || entry.EntityType != "file" { + t.Fatalf("entity = %q (%q), want the modified file", entry.Entity, entry.EntityType) + } + if entry.IsWrite == nil || !*entry.IsWrite { + t.Fatal("IsWrite not set") + } + if entry.Language != "Go" { + t.Fatalf("language = %q, want Go", entry.Language) + } +} + +// TestOTelEntryDroppedWhenStoreRecordsSameCall proves a CLI writing both +// sources does not get billed twice for one API call. +func TestOTelEntryDroppedWhenStoreRecordsSameCall(t *testing.T) { + dir := t.TempDir() + writeStoreFixture(t, dir) + // Same session, model, token counts and near-identical time as the store + // row, in OpenTelemetry form: input 30483 with 7296 cached reads, output + // 289 with 64 reasoning tokens inside. + providertest.WriteFile(t, filepath.Join(dir, "otel", "copilot.jsonl"), + `{"type":"span","traceId":"trace-1","spanId":"span-1","name":"chat mai-code-1-flash","endTime":[1785806389,344000000],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"mai-code-1-flash","gen_ai.conversation.id":"session-1","gen_ai.usage.input_tokens":30483,"gen_ai.usage.output_tokens":225,"gen_ai.usage.cache_read.input_tokens":7296,"gen_ai.usage.reasoning.output_tokens":64}}`+"\n"+ + `{"type":"span","traceId":"trace-2","spanId":"span-2","name":"chat mai-code-1-flash","endTime":[1785900000,0],"attributes":{"gen_ai.operation.name":"chat","gen_ai.response.model":"mai-code-1-flash","gen_ai.conversation.id":"session-2","gen_ai.usage.input_tokens":100,"gen_ai.usage.output_tokens":50}}`+"\n") + + entries, err := Provider{}.WithPaths([]string{dir}).Entries() + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("entries = %d, want 2 (store row + unrelated otel span): %#v", len(entries), entries) + } + sessions := map[string]bool{} + for _, entry := range entries { + sessions[entry.SessionID] = true + } + if !sessions["session-1"] || !sessions["session-2"] { + t.Fatalf("sessions = %v, want session-1 (store) and session-2 (otel)", sessions) + } +} diff --git a/internal/provider/gemini/loader.go b/internal/provider/gemini/loader.go index 7cd0c0c..5bd6459 100644 --- a/internal/provider/gemini/loader.go +++ b/internal/provider/gemini/loader.go @@ -1,12 +1,16 @@ package gemini import ( + "crypto/sha256" + "encoding/hex" + "os" "path/filepath" "sort" "strings" "time" "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/langdetect" "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" @@ -21,14 +25,15 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) sort.Strings(files) files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) + resolver := newProjectResolver() entries := make([]usage.Entry, 0) for _, file := range files { var fileEntries []usage.Entry var err error if strings.EqualFold(filepath.Ext(file), ".jsonl") { - fileEntries, err = parseJSONLFile(file) + fileEntries, err = parseJSONLFile(file, resolver) } else { - fileEntries, err = parseJSONFile(file) + fileEntries, err = parseJSONFile(file, resolver) } if err != nil { return nil, err @@ -49,49 +54,210 @@ type tokens struct { hasTotal bool } -func parseJSONFile(path string) ([]usage.Entry, error) { +// projectIdentity is the project a session file belongs to. Gemini CLI keys +// its data by slug directory (~/.gemini/tmp//), so identity is resolved +// from the slug, not from the transcript contents. +type projectIdentity struct { + project string + projectPath string +} + +// projectResolver maps slug directories to projects. The slug directory may +// carry a .project_root file with the workspace path; otherwise the registry +// at /projects.json maps workspace paths to slugs. +type projectResolver struct { + identities map[string]projectIdentity + registries map[string]map[string]string +} + +func newProjectResolver() *projectResolver { + return &projectResolver{ + identities: make(map[string]projectIdentity), + registries: make(map[string]map[string]string), + } +} + +func (r *projectResolver) resolve(sessionFile string) projectIdentity { + slugDir := slugDirOf(sessionFile) + if identity, ok := r.identities[slugDir]; ok { + return identity + } + identity := r.resolveSlugDir(slugDir) + r.identities[slugDir] = identity + return identity +} + +// slugDirOf returns the per-project data directory a session file lives in: +// //chats/session-*.json(l) belongs to /. +func slugDirOf(path string) string { + dir := filepath.Dir(path) + if filepath.Base(dir) == "chats" { + return filepath.Dir(dir) + } + return dir +} + +func (r *projectResolver) resolveSlugDir(slugDir string) projectIdentity { + cwd := readProjectRoot(filepath.Join(slugDir, ".project_root")) + if cwd == "" { + cwd = r.registryLookup(slugDir) + } + if path, name, ok := usage.ProjectFromCWD(cwd); ok { + return projectIdentity{project: name, projectPath: path} + } + // A readable slug is the sanitized project name; a hash slug says nothing. + if slug := filepath.Base(slugDir); !looksLikeHash(slug) { + return projectIdentity{project: slug} + } + return projectIdentity{project: usage.UnknownProject} +} + +func readProjectRoot(path string) string { + contents, err := os.ReadFile(path) + if err != nil { + return "" + } + return strings.TrimSpace(string(contents)) +} + +// registryLookup finds the workspace path for a slug in projects.json, which +// lives next to the tmp directory: /projects.json maps workspace +// path -> slug. Iteration is sorted so a slug mapped from several paths always +// resolves identically. +func (r *projectResolver) registryLookup(slugDir string) string { + registryPath := filepath.Join(filepath.Dir(filepath.Dir(slugDir)), "projects.json") + registry, ok := r.registries[registryPath] + if !ok { + registry = loadProjectRegistry(registryPath) + r.registries[registryPath] = registry + } + return registry[filepath.Base(slugDir)] +} + +func loadProjectRegistry(path string) map[string]string { + record, err := agentdata.ReadJSONObject(path) + if err != nil || record == nil { + return nil + } + projects := agentdata.ObjectAt(record["projects"]) + if projects == nil { + return nil + } + cwds := make([]string, 0, len(projects)) + for cwd := range projects { + cwds = append(cwds, cwd) + } + sort.Strings(cwds) + bySlug := make(map[string]string, len(projects)*2) + for _, cwd := range cwds { + slug := agentdata.StringValue(projects[cwd]) + if cwd == "" || slug == "" { + continue + } + if _, exists := bySlug[slug]; !exists { + bySlug[slug] = cwd + } + // Older Gemini CLI versions named the slug directory + // sha256(workspace path) instead of the sanitized project name. + hashed := sha256Hex(cwd) + if _, exists := bySlug[hashed]; !exists { + bySlug[hashed] = cwd + } + } + return bySlug +} + +func sha256Hex(value string) string { + sum := sha256.Sum256([]byte(value)) + return hex.EncodeToString(sum[:]) +} + +func looksLikeHash(slug string) bool { + if len(slug) < 32 { + return false + } + for _, r := range slug { + isHex := (r >= '0' && r <= '9') || (r >= 'a' && r <= 'f') + if !isHex { + return false + } + } + return true +} + +// messageLog replays a session's message stream. Gemini CLI rewrites history +// in place — a message record with a known id supersedes the earlier version, +// and "$set" checkpoint lines repeat the whole array — so messages are keyed +// by id, last write wins, first-seen order preserved. Checkpoints never clear +// the log: a message dropped by a later checkpoint or "$rewindTo" still +// consumed tokens. +type messageLog struct { + records []loggedMessage + index map[string]int +} + +type loggedMessage struct { + record map[string]any + line int +} + +func newMessageLog() *messageLog { + return &messageLog{index: make(map[string]int)} +} + +func (l *messageLog) upsert(record map[string]any, line int) { + if agentdata.StringField(record, "type") != "gemini" { + return + } + logged := loggedMessage{record: record, line: line} + if id := agentdata.StringField(record, "id"); id != "" { + if i, ok := l.index[id]; ok { + l.records[i] = logged + return + } + l.index[id] = len(l.records) + } + l.records = append(l.records, logged) +} + +func parseJSONFile(path string, resolver *projectResolver) ([]usage.Entry, error) { record, err := agentdata.ReadJSONObject(path) if err != nil || record == nil { return nil, err } + identity := resolver.resolve(path) fallback := agentdata.FileModifiedTime(path) sessionID := agentdata.FirstStringField(record, "sessionId", "session_id") if sessionID == "" { sessionID = strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) } sessionTimestamp := firstTimestamp(record, fallback, "startTime", "lastUpdated", "timestamp", "created_at") + + log := newMessageLog() if messages := agentdata.ArrayAt(record["messages"]); len(messages) > 0 { - entries := make([]usage.Entry, 0) for index, raw := range messages { - message := agentdata.ObjectAt(raw) - if agentdata.StringField(message, "type") != "gemini" { - continue - } - if entry, ok := directEntry(message, path, index+1, "", sessionID, sessionTimestamp); ok { - entries = append(entries, entry) - } + log.upsert(agentdata.ObjectAt(raw), index+1) } - return entries, nil + return messageEntries(log, path, "", identity, sessionID, sessionTimestamp), nil } - if agentdata.StringField(record, "type") == "gemini" { - if entry, ok := directEntry(record, path, 1, "", sessionID, fallback); ok { - return []usage.Entry{entry}, nil - } - return nil, nil + log.upsert(record, 1) + if len(log.records) > 0 { + return messageEntries(log, path, "", identity, sessionID, fallback), nil } - return statsEntries(recordStats(record), path, 1, agentdata.StringField(record, "model"), sessionID, sessionTimestamp), nil + return statsEntries(recordStats(record), path, 1, agentdata.StringField(record, "model"), identity, sessionID, sessionTimestamp), nil } -func parseJSONLFile(path string) ([]usage.Entry, error) { +func parseJSONLFile(path string, resolver *projectResolver) ([]usage.Entry, error) { lines, err := agentdata.ReadJSONLines(path) if err != nil { return nil, err } + identity := resolver.resolve(path) fallback := agentdata.FileModifiedTime(path) sessionID := strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) currentModel := "" + log := newMessageLog() entries := make([]usage.Entry, 0) - indexByMessageID := make(map[string]int) for _, line := range lines { record := line.Value if value := agentdata.FirstStringField(record, "sessionId", "session_id"); value != "" { @@ -100,31 +266,42 @@ func parseJSONLFile(path string) ([]usage.Entry, error) { if model := agentdata.StringField(record, "model"); model != "" { currentModel = model } - if agentdata.StringField(record, "type") == "gemini" { - entry, ok := directEntry(record, path, line.Line, currentModel, sessionID, fallback) - if !ok { - continue + // "$set" lines update session metadata; when they carry a messages + // array they are a full-history checkpoint. Replaying every copy + // through the log dedupes on message id. + if set := agentdata.ObjectAt(record["$set"]); set != nil { + if value := agentdata.FirstStringField(set, "sessionId", "session_id"); value != "" { + sessionID = value } - messageID := agentdata.StringField(record, "id") - if messageID != "" { - if index, exists := indexByMessageID[messageID]; exists { - entries[index] = entry - continue - } - indexByMessageID[messageID] = len(entries) + for _, raw := range agentdata.ArrayAt(set["messages"]) { + log.upsert(agentdata.ObjectAt(raw), line.Line) } - entries = append(entries, entry) + continue + } + if agentdata.StringField(record, "type") != "" { + log.upsert(record, line.Line) continue } if stats := recordStats(record); stats != nil { timestamp := firstTimestamp(record, fallback, "timestamp") - entries = append(entries, statsEntries(stats, path, line.Line, currentModel, sessionID, timestamp)...) + entries = append(entries, statsEntries(stats, path, line.Line, currentModel, identity, sessionID, timestamp)...) } } + entries = append(entries, messageEntries(log, path, currentModel, identity, sessionID, fallback)...) return entries, nil } -func directEntry(record map[string]any, path string, line int, modelHint, sessionID string, fallback time.Time) (usage.Entry, bool) { +func messageEntries(log *messageLog, path, modelHint string, identity projectIdentity, sessionID string, fallback time.Time) []usage.Entry { + entries := make([]usage.Entry, 0, len(log.records)) + for _, message := range log.records { + if entry, ok := messageEntry(message.record, path, message.line, modelHint, identity, sessionID, fallback); ok { + entries = append(entries, entry) + } + } + return entries +} + +func messageEntry(record map[string]any, path string, line int, modelHint string, identity projectIdentity, sessionID string, fallback time.Time) (usage.Entry, bool) { tokens, ok := parseTokens(record["tokens"]) if !ok { return usage.Entry{}, false @@ -134,10 +311,103 @@ func directEntry(record map[string]any, path string, line int, modelHint, sessio model = modelHint } timestamp := firstTimestamp(record, fallback, "timestamp", "created_at") - return buildGeminiEntry(path, line, model, sessionID, timestamp, tokens, true, agentdata.StringField(record, "id")) + entry, ok := buildGeminiEntry(path, line, model, identity, sessionID, timestamp, tokens, true, agentdata.StringField(record, "id")) + if !ok { + return usage.Entry{}, false + } + + changes, candidates := toolCallData(record, identity.projectPath) + for _, change := range changes { + entry.ApplyFileChange(change) + } + if language := langdetect.Dominant(candidates); language != langdetect.Unknown { + entry.Language = language + } + return entry, true +} + +// toolCallData extracts the file modifications and language candidates from a +// gemini message's tool calls. write_file records the created content; +// replace/edit records the diff stat the CLI computed, falling back to the +// old/new string line counts when the stat is missing. +func toolCallData(record map[string]any, projectPath string) ([]usage.FileChange, []langdetect.Candidate) { + var changes []usage.FileChange + var candidates []langdetect.Candidate + for _, raw := range agentdata.ArrayAt(record["toolCalls"]) { + call := agentdata.ObjectAt(raw) + if call == nil { + continue + } + if status := agentdata.StringField(call, "status"); status != "" && status != "success" { + continue + } + args := agentdata.ObjectAt(call["args"]) + candidates = append(candidates, languageCandidates(args)...) + + display := agentdata.ObjectAt(call["resultDisplay"]) + filePath := changePath(display, args, projectPath) + if filePath == "" { + continue + } + switch strings.ToLower(agentdata.StringField(call, "name")) { + case "write_file", "writefile": + changes = append(changes, usage.FileChange{ + Path: filePath, + LinesAdded: usage.CountLines(agentdata.StringField(args, "content")), + }) + case "replace", "edit": + added, removed := replaceLines(display, args) + changes = append(changes, usage.FileChange{ + Path: filePath, + LinesAdded: added, + LinesRemoved: removed, + }) + } + } + return changes, candidates +} + +// languageCandidates mirrors the Claude provider's weighting: an explicit +// path argument names the file being worked on (weight 3); paths mentioned +// inside command or content text are weaker hints (weight 1). +func languageCandidates(args map[string]any) []langdetect.Candidate { + candidates := make([]langdetect.Candidate, 0, len(args)) + for key, value := range args { + lower := strings.ToLower(key) + text := agentdata.StringValue(value) + if text == "" { + continue + } + switch { + case strings.Contains(lower, "path") || strings.Contains(lower, "file"): + candidates = append(candidates, langdetect.Candidate{Path: text, Weight: 3}) + case strings.Contains(lower, "command") || strings.Contains(lower, "content") || strings.Contains(lower, "query"): + for _, path := range langdetect.PathsFromText(text) { + candidates = append(candidates, langdetect.Candidate{Path: path, Weight: 1}) + } + } + } + return candidates } -func statsEntries(stats map[string]any, path string, line int, modelHint, sessionID string, timestamp time.Time) []usage.Entry { +func changePath(display, args map[string]any, projectPath string) string { + path := agentdata.FirstStringField(display, "filePath", "fileName") + if path == "" { + path = agentdata.FirstStringField(args, "file_path", "filePath", "absolute_path") + } + return usage.ResolvePath(projectPath, path) +} + +func replaceLines(display, args map[string]any) (uint64, uint64) { + if diffStat := agentdata.ObjectAt(display["diffStat"]); diffStat != nil { + return agentdata.UintField(diffStat, "model_added_lines"), + agentdata.UintField(diffStat, "model_removed_lines") + } + return usage.CountLines(agentdata.StringField(args, "new_string")), + usage.CountLines(agentdata.StringField(args, "old_string")) +} + +func statsEntries(stats map[string]any, path string, line int, modelHint string, identity projectIdentity, sessionID string, timestamp time.Time) []usage.Entry { if stats == nil { return nil } @@ -149,7 +419,7 @@ func statsEntries(stats map[string]any, path string, line int, modelHint, sessio if !ok { continue } - if entry, ok := buildGeminiEntry(path, line, model, sessionID, timestamp, tokens, false, ""); ok { + if entry, ok := buildGeminiEntry(path, line, model, identity, sessionID, timestamp, tokens, false, ""); ok { entries = append(entries, entry) } } @@ -165,14 +435,14 @@ func statsEntries(stats map[string]any, path string, line int, modelHint, sessio if model == "" { model = "unknown" } - entry, ok := buildGeminiEntry(path, line, model, sessionID, timestamp, tokens, false, "") + entry, ok := buildGeminiEntry(path, line, model, identity, sessionID, timestamp, tokens, false, "") if !ok { return nil } return []usage.Entry{entry} } -func buildGeminiEntry(path string, line int, model, sessionID string, timestamp time.Time, tokens tokens, direct bool, messageID string) (usage.Entry, bool) { +func buildGeminiEntry(path string, line int, model string, identity projectIdentity, sessionID string, timestamp time.Time, tokens tokens, direct bool, messageID string) (usage.Entry, bool) { model = strings.TrimSpace(model) if model == "" { return usage.Entry{}, false @@ -192,9 +462,17 @@ func buildGeminiEntry(path string, line int, model, sessionID string, timestamp if !usageprovider.NonZero(tokenUsage) { return usage.Entry{}, false } - entry := usageprovider.BaseEntry(usage.ProviderGemini, timestamp, "gemini", "Gemini", sessionID, model, "Gemini CLI", tokenUsage) + entry := usageprovider.BaseEntry(usage.ProviderGemini, timestamp, identity.project, identity.projectPath, sessionID, model, "Gemini CLI", tokenUsage) usageprovider.SetSource(&entry, path, line, 0, 0) - entry.ID = usageprovider.StableEntryID(entry, messageID) + // Session files are rewritten in place — a message's line and token + // counts move as history is checkpointed — so like the Claude provider, + // the entry is keyed on the message id alone. Only id-less records fall + // back to source position. + if messageID != "" { + entry.ID = usage.StableID(string(usage.ProviderGemini), messageID) + } else { + entry.ID = usageprovider.StableEntryID(entry) + } return entry, true } diff --git a/internal/provider/gemini/loader_test.go b/internal/provider/gemini/loader_test.go new file mode 100644 index 0000000..9c23cb2 --- /dev/null +++ b/internal/provider/gemini/loader_test.go @@ -0,0 +1,196 @@ +package gemini + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func loadTestEntries(t *testing.T, root string) []usage.Entry { + t.Helper() + entries, err := loadEntries([]string{root}, nil) + if err != nil { + t.Fatal(err) + } + return entries +} + +// TestSetCheckpointFormat replays the current session-*.jsonl layout: a +// metadata header followed by "$set" checkpoint lines that each carry the +// whole message history. Messages repeated across checkpoints must dedupe on +// id with the last version winning. +func TestSetCheckpointFormat(t *testing.T) { + dir := t.TempDir() + providertest.WriteFile(t, filepath.Join(dir, "shop", ".project_root"), "/ws/shop\n") + providertest.WriteFile(t, filepath.Join(dir, "shop", "chats", "session-a.jsonl"), + `{"sessionId":"session-a","projectHash":"hash-a","startTime":"2026-06-01T10:00:00.000Z","kind":"main"}`+"\n"+ + `{"$set":{"messages":[{"id":"m1","timestamp":"2026-06-01T10:00:10.000Z","type":"user","content":"hi"},{"id":"m2","timestamp":"2026-06-01T10:00:20.000Z","type":"gemini","model":"gemini-2.5-pro","tokens":{"input":50,"output":5,"total":55}}]}}`+"\n"+ + `{"$set":{"messages":[{"id":"m1","timestamp":"2026-06-01T10:00:10.000Z","type":"user","content":"hi"},{"id":"m2","timestamp":"2026-06-01T10:00:20.000Z","type":"gemini","model":"gemini-2.5-pro","tokens":{"input":100,"output":10,"cached":40,"thoughts":5,"total":115}},{"id":"m3","timestamp":"2026-06-01T10:01:00.000Z","type":"gemini","model":"gemini-2.5-pro","tokens":{"input":10,"output":2,"total":12}}]}}`+"\n") + + entries := loadTestEntries(t, dir) + if len(entries) != 2 { + t.Fatalf("entries = %d, want 2: %#v", len(entries), entries) + } + first := entries[0] + if first.SessionID != "session-a" { + t.Fatalf("session id = %q, want %q", first.SessionID, "session-a") + } + if first.Project != "shop" || first.ProjectPath != "/ws/shop" { + t.Fatalf("project = %q/%q, want shop//ws/shop", first.Project, first.ProjectPath) + } + want := usage.TokenUsage{ + InputTokens: 60, + OutputTokens: 10, + CacheReadInputTokens: 40, + ReasoningOutputTokens: 5, + TotalTokens: 115, + } + if first.Usage != want { + t.Fatalf("usage = %#v, want %#v", first.Usage, want) + } + second := entries[1] + if second.Usage.TotalTokens != 12 { + t.Fatalf("second total = %d, want 12", second.Usage.TotalTokens) + } +} + +// TestMessageIDStableAcrossRewrites: session files are rewritten in place, so +// the same message must keep its entry id when its line or token counts move. +func TestMessageIDStableAcrossRewrites(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "shop", "chats", "session-a.jsonl") + providertest.WriteFile(t, path, + `{"sessionId":"session-a","projectHash":"hash-a"}`+"\n"+ + `{"id":"m1","timestamp":"2026-06-01T10:00:20.000Z","type":"gemini","model":"gemini-2.5-pro","tokens":{"input":50,"output":5,"total":55}}`+"\n") + before := loadTestEntries(t, dir) + + providertest.WriteFile(t, path, + `{"sessionId":"session-a","projectHash":"hash-a"}`+"\n"+ + `{"$set":{"summary":"noise"}}`+"\n"+ + `{"$set":{"messages":[{"id":"m1","timestamp":"2026-06-01T10:00:20.000Z","type":"gemini","model":"gemini-2.5-pro","tokens":{"input":80,"output":9,"total":89}}]}}`+"\n") + after := loadTestEntries(t, dir) + + if len(before) != 1 || len(after) != 1 { + t.Fatalf("entries = %d/%d, want 1/1", len(before), len(after)) + } + if before[0].ID != after[0].ID { + t.Fatalf("entry id changed across rewrite: %q -> %q", before[0].ID, after[0].ID) + } +} + +// TestProjectRegistryFallback: without a .project_root file, the slug is +// resolved through /projects.json, which maps workspace path to +// slug and lives next to the scanned tmp directory. +func TestProjectRegistryFallback(t *testing.T) { + home := t.TempDir() + root := filepath.Join(home, "tmp") + providertest.WriteFile(t, filepath.Join(home, "projects.json"), + `{"projects":{"/ws/deep/shop":"shop"}}`) + providertest.WriteFile(t, filepath.Join(root, "shop", "chats", "session-a.jsonl"), + `{"sessionId":"session-a"}`+"\n"+ + `{"id":"m1","timestamp":"2026-06-01T10:00:20.000Z","type":"gemini","model":"gemini-2.5-pro","tokens":{"input":50,"output":5,"total":55}}`+"\n") + + entries := loadTestEntries(t, root) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].Project != "shop" || entries[0].ProjectPath != "/ws/deep/shop" { + t.Fatalf("project = %q/%q, want shop//ws/deep/shop", entries[0].Project, entries[0].ProjectPath) + } +} + +// TestHashSlugFallsBackToUnknown: a hash slug with no .project_root and no +// registry entry carries no project information at all. +func TestHashSlugFallsBackToUnknown(t *testing.T) { + dir := t.TempDir() + slug := "0ade338c48c531ca24306ff4d04bcdf7b2a2cb022b1c968ffa4749b403dde9b2" + providertest.WriteFile(t, filepath.Join(dir, slug, "chats", "session-a.jsonl"), + `{"id":"m1","timestamp":"2026-06-01T10:00:20.000Z","type":"gemini","model":"gemini-2.5-pro","tokens":{"input":50,"output":5,"total":55}}`+"\n") + + entries := loadTestEntries(t, dir) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + if entries[0].Project != usage.UnknownProject { + t.Fatalf("project = %q, want %q", entries[0].Project, usage.UnknownProject) + } + if entries[0].ProjectPath != "" { + t.Fatalf("project path = %q, want empty", entries[0].ProjectPath) + } +} + +// TestToolCallFileChanges: write_file and replace tool calls become per-file +// line stats, the entity points at the most-changed file, and the language +// comes from the touched paths. Failed tool calls are ignored. +func TestToolCallFileChanges(t *testing.T) { + dir := t.TempDir() + providertest.WriteFile(t, filepath.Join(dir, "shop", ".project_root"), "/ws/shop\n") + providertest.WriteFile(t, filepath.Join(dir, "shop", "chats", "session-a.json"), `{ + "sessionId": "session-a", + "projectHash": "hash-a", + "startTime": "2026-06-01T10:00:00.000Z", + "messages": [ + {"id": "m1", "timestamp": "2026-06-01T10:00:10.000Z", "type": "user", "content": "edit"}, + {"id": "m2", "timestamp": "2026-06-01T10:00:20.000Z", "type": "gemini", "model": "gemini-2.5-pro", + "tokens": {"input": 100, "output": 10, "total": 110}, + "toolCalls": [ + {"id": "t1", "name": "replace", "status": "success", + "args": {"file_path": "src/main.go", "old_string": "a", "new_string": "b"}, + "resultDisplay": {"filePath": "/ws/shop/src/main.go", "diffStat": {"model_added_lines": 7, "model_removed_lines": 3}}}, + {"id": "t2", "name": "write_file", "status": "success", + "args": {"file_path": "src/util.go", "content": "package main\nfunc util() {}\n"}}, + {"id": "t3", "name": "replace", "status": "error", + "args": {"file_path": "src/broken.go", "old_string": "x", "new_string": "y"}} + ]} + ] + }`) + + entries := loadTestEntries(t, dir) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + entry := entries[0] + if entry.LinesAdded != 9 || entry.LinesRemoved != 3 { + t.Fatalf("lines = +%d/-%d, want +9/-3", entry.LinesAdded, entry.LinesRemoved) + } + if entry.Entity != "/ws/shop/src/main.go" || entry.EntityType != "file" { + t.Fatalf("entity = %q (%q), want /ws/shop/src/main.go (file)", entry.Entity, entry.EntityType) + } + if entry.IsWrite == nil || !*entry.IsWrite { + t.Fatal("IsWrite not set") + } + if len(entry.Files) != 2 { + t.Fatalf("files = %#v, want 2 entries", entry.Files) + } + if entry.Language != "Go" { + t.Fatalf("language = %q, want Go", entry.Language) + } +} + +// TestLegacyStatsRecords: the older stats layout (per-model token summaries +// under stats.models) must keep loading. +func TestLegacyStatsRecords(t *testing.T) { + dir := t.TempDir() + providertest.WriteFile(t, filepath.Join(dir, "shop", "chats", "session-a.json"), + `{"sessionId":"session-a","timestamp":"2026-06-01T10:00:00.000Z","stats":{"models":{"gemini-2.5-flash":{"tokens":{"prompt":200,"candidates":20,"cached":50,"total":220}}}}}`) + + entries := loadTestEntries(t, dir) + if len(entries) != 1 { + t.Fatalf("entries = %d, want 1", len(entries)) + } + entry := entries[0] + if entry.Model != "gemini-2.5-flash" { + t.Fatalf("model = %q", entry.Model) + } + want := usage.TokenUsage{ + InputTokens: 150, + OutputTokens: 20, + CacheReadInputTokens: 50, + TotalTokens: 220, + } + if entry.Usage != want { + t.Fatalf("usage = %#v, want %#v", entry.Usage, want) + } +} diff --git a/internal/provider/gemini/provider_test.go b/internal/provider/gemini/provider_test.go index 93ef8a4..a94a9cf 100644 --- a/internal/provider/gemini/provider_test.go +++ b/internal/provider/gemini/provider_test.go @@ -8,13 +8,15 @@ import ( "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) -// TestLoadsEntry is the Gemini CLI smoke test: a minimal fixture must produce -// exactly one entry with the expected identity and token counts. +// TestLoadsEntry is the Gemini CLI smoke test: a minimal fixture in the real +// ~/.gemini/tmp//chats layout must produce exactly one entry with the +// expected identity and token counts. func TestLoadsEntry(t *testing.T) { entries, err := func() ([]usage.Entry, error) { dir := t.TempDir() - path := filepath.Join(dir, "session-a.jsonl") - providertest.WriteFile(t, path, + providertest.WriteFile(t, filepath.Join(dir, "shop", ".project_root"), + filepath.Join(dir, "workspace", "shop")+"\n") + providertest.WriteFile(t, filepath.Join(dir, "shop", "chats", "session-a.jsonl"), `{"sessionId":"session-a","projectHash":"project-a","startTime":"2026-05-17T11:07:00.000Z"}`+"\n"+ `{"id":"msg-a","timestamp":"2026-05-17T11:07:32.000Z","type":"gemini","model":"gemini-3-flash-preview","tokens":{"input":15327,"output":23,"cached":11526,"thoughts":919,"tool":7,"total":16276}}`+"\n") return Provider{}.WithPaths([]string{dir}).Entries() @@ -24,7 +26,7 @@ func TestLoadsEntry(t *testing.T) { Provider: usage.ProviderGemini, Model: "gemini-3-flash-preview", SessionID: "session-a", - Project: "gemini", + Project: "shop", Tokens: usage.TokenUsage{ InputTokens: 3808, OutputTokens: 23, diff --git a/internal/provider/kimi/loader.go b/internal/provider/kimi/loader.go index 65e3ddf..f00fc96 100644 --- a/internal/provider/kimi/loader.go +++ b/internal/provider/kimi/loader.go @@ -14,18 +14,22 @@ import ( const defaultModel = "kimi-for-coding" -func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { +func wireFiles(paths []string) []string { files := make([]string, 0) for _, root := range paths { files = append(files, agentdata.CollectFiles(filepath.Join(root, "sessions"), isWireFile)...) } sort.Strings(files) - files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) + return agentdata.UniqueStrings(files) +} + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + files := agentdata.FilterFiles(wireFiles(paths), filter) entries := make([]usage.Entry, 0) seen := make(map[string]bool) for _, file := range files { - fileEntries, err := parseWireFile(file) + fileEntries, _, err := parseWireFileFrom(file, 0) if err != nil { return nil, err } @@ -69,10 +73,10 @@ type record struct { messageID string } -func parseWireFile(path string) ([]usage.Entry, error) { - lines, err := agentdata.ReadJSONLines(path, `usage`) +func parseWireFileFrom(path string, start int64) ([]usage.Entry, int64, error) { + lines, consumed, err := agentdata.ReadJSONLinesFrom(path, start, `usage`) if err != nil { - return nil, err + return nil, 0, err } configModel := configModel(path) sessionID := sessionID(path) @@ -102,7 +106,7 @@ func parseWireFile(path string) ([]usage.Entry, error) { entry.ID = usageprovider.StableEntryID(entry, record.messageID) entries = append(entries, entry) } - return entries, nil + return entries, consumed, nil } func parseStatusUpdate(value map[string]any) (record, bool) { diff --git a/internal/provider/kimi/provider.go b/internal/provider/kimi/provider.go index 5d41be4..cb2ec7e 100644 --- a/internal/provider/kimi/provider.go +++ b/internal/provider/kimi/provider.go @@ -34,3 +34,13 @@ func (p Provider) Entries() ([]usage.Entry, error) { } return usageprovider.SortEntriesByTimestampDesc(entries), nil } + +// StreamEntries parses each wire file from where the previous scan stopped. +// Wire files are append-only JSONL and each line stands alone. +// +// Entries loads with a cross-file seen set because it returns one flat slice; +// here the same duplicates collapse on the database's primary key, so no +// in-memory state is carried between files. +func (p Provider) StreamEntries(resume func(path string) int64, emit func(path string, entries []usage.Entry, offset int64) error) error { + return usageprovider.StreamFiles(wireFiles(p.Paths()), p.Filter(), parseWireFileFrom, emit, resume) +} diff --git a/internal/provider/openclaw/loader.go b/internal/provider/openclaw/loader.go index 3537fdf..48c65ff 100644 --- a/internal/provider/openclaw/loader.go +++ b/internal/provider/openclaw/loader.go @@ -11,6 +11,10 @@ import ( "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) +// openclaw is deliberately not resumable: a model-change line sets the model +// for every message after it, so parsing from a mid-file offset would attribute +// later entries to "unknown" instead. The state would have to be carried across +// scans to resume safely, and it is not worth that for this provider's volume. func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { files := make([]string, 0) for _, path := range paths { diff --git a/internal/provider/pi/loader.go b/internal/provider/pi/loader.go index a9b40a6..3b1993e 100644 --- a/internal/provider/pi/loader.go +++ b/internal/provider/pi/loader.go @@ -11,17 +11,21 @@ import ( "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) -func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { +func sessionFiles(paths []string) []string { files := make([]string, 0) for _, path := range paths { files = append(files, agentdata.CollectExt(path, ".jsonl")...) } sort.Strings(files) - files = agentdata.FilterFiles(files, filter) + return files +} + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + files := agentdata.FilterFiles(sessionFiles(paths), filter) entries := make([]usage.Entry, 0) for _, file := range files { - fileEntries, err := parseSessionFile(file) + fileEntries, _, err := parseSessionFileFrom(file, 0) if err != nil { return nil, err } @@ -31,10 +35,10 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) return entries, nil } -func parseSessionFile(path string) ([]usage.Entry, error) { - lines, err := agentdata.ReadJSONLines(path, `"usage"`, `"message"`) +func parseSessionFileFrom(path string, start int64) ([]usage.Entry, int64, error) { + lines, consumed, err := agentdata.ReadJSONLinesFrom(path, start, `"usage"`, `"message"`) if err != nil { - return nil, err + return nil, 0, err } project := project(path) sessionID := sessionID(path) @@ -74,7 +78,7 @@ func parseSessionFile(path string) ([]usage.Entry, error) { entry.ID = usageprovider.StableEntryID(entry) entries = append(entries, entry) } - return entries, nil + return entries, consumed, nil } func sessionID(path string) string { diff --git a/internal/provider/pi/provider.go b/internal/provider/pi/provider.go index b16f39a..c0f0fe5 100644 --- a/internal/provider/pi/provider.go +++ b/internal/provider/pi/provider.go @@ -34,3 +34,10 @@ func (p Provider) Entries() ([]usage.Entry, error) { } return usageprovider.SortEntriesByTimestampDesc(entries), nil } + +// StreamEntries parses each session file from where the previous scan stopped. +// Sessions are append-only JSONL, so a scan reads only what was written since +// the last one. +func (p Provider) StreamEntries(resume func(path string) int64, emit func(path string, entries []usage.Entry, offset int64) error) error { + return usageprovider.StreamFiles(sessionFiles(p.Paths()), p.Filter(), parseSessionFileFrom, emit, resume) +} diff --git a/internal/provider/qwen/loader.go b/internal/provider/qwen/loader.go index 9244e8d..2b1ea1d 100644 --- a/internal/provider/qwen/loader.go +++ b/internal/provider/qwen/loader.go @@ -11,7 +11,7 @@ import ( "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) -func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { +func chatFiles(paths []string) []string { files := make([]string, 0) for _, root := range paths { files = append(files, agentdata.CollectFiles(filepath.Join(root, "projects"), isChatFile)...) @@ -20,11 +20,15 @@ func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) } } sort.Strings(files) - files = agentdata.FilterFiles(agentdata.UniqueStrings(files), filter) + return agentdata.UniqueStrings(files) +} + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + files := agentdata.FilterFiles(chatFiles(paths), filter) entries := make([]usage.Entry, 0) for _, file := range files { - fileEntries, err := parseChatFile(file) + fileEntries, _, err := parseChatFileFrom(file, 0) if err != nil { return nil, err } @@ -47,10 +51,10 @@ func isChatFile(path string) bool { return false } -func parseChatFile(path string) ([]usage.Entry, error) { - lines, err := agentdata.ReadJSONLines(path, `"usageMetadata"`) +func parseChatFileFrom(path string, start int64) ([]usage.Entry, int64, error) { + lines, consumed, err := agentdata.ReadJSONLinesFrom(path, start, `"usageMetadata"`) if err != nil { - return nil, err + return nil, 0, err } project := project(path) fallback := agentdata.FileModifiedTime(path) @@ -93,7 +97,7 @@ func parseChatFile(path string) ([]usage.Entry, error) { entry.ID = usageprovider.StableEntryID(entry) entries = append(entries, entry) } - return entries, nil + return entries, consumed, nil } func project(path string) string { diff --git a/internal/provider/qwen/provider.go b/internal/provider/qwen/provider.go index c20fcb3..5888b04 100644 --- a/internal/provider/qwen/provider.go +++ b/internal/provider/qwen/provider.go @@ -34,3 +34,10 @@ func (p Provider) Entries() ([]usage.Entry, error) { } return usageprovider.SortEntriesByTimestampDesc(entries), nil } + +// StreamEntries parses each chat file from where the previous scan stopped. +// Chat files are append-only JSONL and each line stands alone, so a resumed +// read produces exactly what a whole read would. +func (p Provider) StreamEntries(resume func(path string) int64, emit func(path string, entries []usage.Entry, offset int64) error) error { + return usageprovider.StreamFiles(chatFiles(p.Paths()), p.Filter(), parseChatFileFrom, emit, resume) +} diff --git a/internal/store/lock.go b/internal/store/lock.go index de91487..150454f 100644 --- a/internal/store/lock.go +++ b/internal/store/lock.go @@ -10,11 +10,6 @@ import ( const LockFile = "tokitoki.lock" -// UploadLockFile serializes queue drains across processes, separately from -// the data lock: draining talks to the network for up to the whole upload -// timeout, and ingestion must never wait behind that. -const UploadLockFile = "upload.lock" - // ErrLockBusy reports that another process held the lock for the whole // timeout. Callers that treat "someone else is already doing this work" as // success test for it with errors.Is. diff --git a/internal/usagedb/db.go b/internal/usagedb/db.go index e107dfb..3794b24 100644 --- a/internal/usagedb/db.go +++ b/internal/usagedb/db.go @@ -10,6 +10,7 @@ import ( "fmt" "os" "path/filepath" + "strings" "time" "github.com/tokitoki-dev/tokitoki-cli/internal/provider/codex" @@ -33,22 +34,34 @@ CREATE TABLE IF NOT EXISTS usage_events ( attempt_count INTEGER NOT NULL DEFAULT 0, next_attempt_at INTEGER NOT NULL DEFAULT 0, uploaded_at INTEGER, - last_error TEXT NOT NULL DEFAULT '' + last_error TEXT NOT NULL DEFAULT '', + lease_until INTEGER NOT NULL DEFAULT 0 ); CREATE INDEX IF NOT EXISTS idx_usage_events_queue ON usage_events(status, next_attempt_at); CREATE TABLE IF NOT EXISTS scanned_files ( path TEXT PRIMARY KEY, size INTEGER NOT NULL, - mtime_ns INTEGER NOT NULL + mtime_ns INTEGER NOT NULL, + offset INTEGER NOT NULL DEFAULT 0 ); ` // FileState is the stat snapshot of a source file at the time it was last // successfully scanned. A file whose current stat matches its stored state // holds no events the database has not already seen. +// +// Size and MtimeNS answer "has this file changed at all"; Offset answers +// "where do we resume". They are deliberately separate: Size is the size +// stat'd at the start of the pass, while Offset is where parsing actually +// stopped, which is earlier whenever the file ended in a partial line. type FileState struct { Size int64 MtimeNS int64 + + // Offset is the byte position after the last fully-consumed line. Zero + // means parse from the beginning — the correct default both for files + // never seen before and for rows written by versions predating resume. + Offset int64 } type DB struct { @@ -108,7 +121,13 @@ func stampVersion(db *sql.DB) error { // dropped file position entirely in favor of session + timestamp + tokens. // The rekey recomputes from payload, so any older version jumps straight to // the current scheme in one pass. -const eventSchemaVersion = 2 +// Version 3 adds scanned_files.offset, which lets a scan resume mid-file +// instead of re-parsing every changed file from the beginning. Existing rows +// default to 0, meaning "start over" — correct, just not yet incremental. +// +// Version 4 adds usage_events.lease_until, which lets a claimed batch be +// reclaimed after the process that claimed it died mid-upload. +const eventSchemaVersion = 4 func migrate(db *sql.DB) error { var version int @@ -123,9 +142,31 @@ func migrate(db *sql.DB) error { return err } } + if version < 3 { + if err := addColumn(db, `ALTER TABLE scanned_files ADD COLUMN offset INTEGER NOT NULL DEFAULT 0`); err != nil { + return err + } + } + if version < 4 { + if err := addColumn(db, `ALTER TABLE usage_events ADD COLUMN lease_until INTEGER NOT NULL DEFAULT 0`); err != nil { + return err + } + } return stampVersion(db) } +// addColumn runs an ALTER TABLE ADD COLUMN that may already have been applied. +// The schema statement creates fresh databases at the current shape, so a +// column this migration adds can already exist by the time it runs; that +// duplicate-column error is the expected no-op, not a failure. +func addColumn(db *sql.DB, statement string) error { + _, err := db.Exec(statement) + if err != nil && strings.Contains(err.Error(), "duplicate column name") { + return nil + } + return err +} + func rekeyCodexEvents(db *sql.DB) error { rows, err := db.Query(`SELECT id, payload, status FROM usage_events`) if err != nil { @@ -262,7 +303,7 @@ func (s *DB) InsertEvents(entries []usage.Entry) (int, error) { // ScannedFiles returns the stat snapshot of every file recorded as scanned. func (s *DB) ScannedFiles() (map[string]FileState, error) { - rows, err := s.db.Query(`SELECT path, size, mtime_ns FROM scanned_files`) + rows, err := s.db.Query(`SELECT path, size, mtime_ns, offset FROM scanned_files`) if err != nil { return nil, err } @@ -272,7 +313,7 @@ func (s *DB) ScannedFiles() (map[string]FileState, error) { for rows.Next() { var path string var state FileState - if err := rows.Scan(&path, &state.Size, &state.MtimeNS); err != nil { + if err := rows.Scan(&path, &state.Size, &state.MtimeNS, &state.Offset); err != nil { return nil, err } states[path] = state @@ -293,14 +334,14 @@ func (s *DB) UpsertScannedFiles(states map[string]FileState) error { } defer tx.Rollback() - stmt, err := tx.Prepare(`INSERT OR REPLACE INTO scanned_files (path, size, mtime_ns) VALUES (?, ?, ?)`) + stmt, err := tx.Prepare(`INSERT OR REPLACE INTO scanned_files (path, size, mtime_ns, offset) VALUES (?, ?, ?, ?)`) if err != nil { return err } defer stmt.Close() for path, state := range states { - if _, err := stmt.Exec(path, state.Size, state.MtimeNS); err != nil { + if _, err := stmt.Exec(path, state.Size, state.MtimeNS, state.Offset); err != nil { return fmt.Errorf("save scanned file %q: %w", path, err) } } @@ -338,6 +379,115 @@ func (s *DB) PendingEvents(now time.Time, limit int) ([]usage.Entry, error) { return entries, rows.Err() } +// PendingCount reports how many events are due for upload at now. It is the +// cheap question "is there enough queued to be worth a request", answered +// without claiming anything. +func (s *DB) PendingCount(now time.Time) (int, error) { + var count int + err := s.db.QueryRow(` + SELECT count(*) FROM usage_events + WHERE status IN ('pending', 'failed') AND next_attempt_at <= ?`, now.Unix()).Scan(&count) + return count, err +} + +// ClaimEvents takes ownership of a batch due for upload and returns it. +// +// Claiming marks the rows "sending" and stamps a lease. Two uploaders running +// at once therefore take different batches instead of both sending the same +// one: the UPDATE is atomic, so whichever runs second sees the rows already +// claimed and moves past them. +// +// Fresh work is claimed first. Only when there is none does it fall back to +// batches whose lease has expired — rows left "sending" by a process that +// died mid-upload. Keeping that fallback off the common path means a healthy +// queue never pays for it, and it also means a slow-but-alive uploader is not +// raced for its batch while ordinary work is still available. +// +// A duplicate send is harmless if it happens anyway: the server dedupes on +// event id and SyncPending counts a duplicate as uploaded. This is why the +// lease can be a plain timestamp rather than something that must be renewed. +func (s *DB) ClaimEvents(now time.Time, limit int, lease time.Duration) ([]usage.Entry, error) { + if limit <= 0 { + limit = -1 + } + leaseUntil := now.Add(lease).Unix() + + entries, err := s.claim(` + UPDATE usage_events SET status = 'sending', lease_until = ? + WHERE id IN ( + SELECT id FROM usage_events + WHERE status IN ('pending', 'failed') AND next_attempt_at <= ? + ORDER BY ts DESC, id + LIMIT ? + ) + RETURNING payload`, leaseUntil, now.Unix(), limit) + if err != nil { + return nil, err + } + if limit > 0 && len(entries) >= limit { + return entries, nil + } + + // Room left in this batch. Anything still "sending" past its lease belongs + // to a process that is gone, so reclaiming it is the only way those events + // are ever sent. Topping up rather than only checking when fresh work runs + // out matters: on a machine whose queue never empties, an "only if idle" + // check would never run and those events would be stranded indefinitely. + remaining := limit - len(entries) + if limit <= 0 { + remaining = limit + } + expired, err := s.claim(` + UPDATE usage_events SET status = 'sending', lease_until = ? + WHERE id IN ( + SELECT id FROM usage_events + WHERE status = 'sending' AND lease_until <= ? + ORDER BY ts DESC, id + LIMIT ? + ) + RETURNING payload`, leaseUntil, now.Unix(), remaining) + if err != nil { + return nil, err + } + return append(entries, expired...), nil +} + +func (s *DB) claim(query string, args ...any) ([]usage.Entry, error) { + rows, err := s.db.Query(query, args...) + if err != nil { + return nil, err + } + defer rows.Close() + + entries := make([]usage.Entry, 0) + for rows.Next() { + var payload string + if err := rows.Scan(&payload); err != nil { + return nil, err + } + var entry usage.Entry + if err := json.Unmarshal([]byte(payload), &entry); err != nil { + return nil, fmt.Errorf("decode usage event: %w", err) + } + entries = append(entries, entry) + } + return entries, rows.Err() +} + +// ReleaseClaims returns claimed events to the queue without counting an +// attempt against them. +// +// It undoes a claim that resolved into nothing — a server response that did +// not account for every event it was sent. Those rows are not failures to back +// off from; they were simply never answered for, and the next pass should pick +// them up immediately rather than wait out their lease. +func (s *DB) ReleaseClaims(ids []string) error { + return s.updateEach(ids, func(stmt *sql.Stmt, id string) error { + _, err := stmt.Exec(id) + return err + }, `UPDATE usage_events SET status = 'pending', lease_until = 0 WHERE id = ? AND status = 'sending'`) +} + // MarkEventsUploaded marks ids as accepted by the server. func (s *DB) MarkEventsUploaded(ids []string) error { now := time.Now().UTC().Unix() @@ -374,7 +524,8 @@ func (s *DB) MarkEventsUploadFailed(ids []string, message string) error { status = 'failed', attempt_count = attempt_count + 1, next_attempt_at = ? + min(? << min(attempt_count, 7), ?), - last_error = ? + last_error = ?, + lease_until = 0 WHERE id = ?`) } diff --git a/internal/usagedb/db_test.go b/internal/usagedb/db_test.go index 9e0205a..2e19212 100644 --- a/internal/usagedb/db_test.go +++ b/internal/usagedb/db_test.go @@ -1,6 +1,7 @@ package usagedb import ( + "database/sql" "path/filepath" "testing" "time" @@ -309,3 +310,269 @@ func TestOpenStampsFreshDatabaseAtCurrentVersion(t *testing.T) { t.Fatalf("user_version = %d, want %d", version, eventSchemaVersion) } } + +// A database written before resume offsets existed must gain the column with +// every row defaulting to 0, which means "parse from the beginning" — safe, +// just not yet incremental. +func TestOpenAddsScannedFilesOffsetToLegacyDatabase(t *testing.T) { + path := filepath.Join(t.TempDir(), "legacy.db") + + legacy, err := sql.Open("sqlite", "file:"+filepath.ToSlash(path)) + if err != nil { + t.Fatal(err) + } + if _, err := legacy.Exec(` + CREATE TABLE usage_events ( + id TEXT PRIMARY KEY, ts INTEGER NOT NULL, payload TEXT NOT NULL, + status TEXT NOT NULL DEFAULT 'pending', attempt_count INTEGER NOT NULL DEFAULT 0, + next_attempt_at INTEGER NOT NULL DEFAULT 0, uploaded_at INTEGER, + last_error TEXT NOT NULL DEFAULT '' + ); + CREATE TABLE scanned_files ( + path TEXT PRIMARY KEY, size INTEGER NOT NULL, mtime_ns INTEGER NOT NULL + ); + INSERT INTO scanned_files (path, size, mtime_ns) VALUES ('/tmp/a.jsonl', 42, 7); + PRAGMA user_version = 2; + `); err != nil { + t.Fatal(err) + } + if err := legacy.Close(); err != nil { + t.Fatal(err) + } + + db, err := Open(path) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + + states, err := db.ScannedFiles() + if err != nil { + t.Fatal(err) + } + state, ok := states["/tmp/a.jsonl"] + if !ok { + t.Fatal("legacy scanned_files row was lost") + } + if state.Size != 42 || state.MtimeNS != 7 { + t.Fatalf("legacy row = %+v, want size 42 mtime 7", state) + } + if state.Offset != 0 { + t.Fatalf("offset = %d, want 0", state.Offset) + } + + var version int + if err := db.db.QueryRow(`PRAGMA user_version`).Scan(&version); err != nil { + t.Fatal(err) + } + if version != eventSchemaVersion { + t.Fatalf("user_version = %d, want %d", version, eventSchemaVersion) + } +} + +// Two uploaders running at once must take different batches, not the same one. +func TestClaimEventsGivesEachCallerADistinctBatch(t *testing.T) { + db := openTestDB(t) + if _, err := db.InsertEvents([]usage.Entry{ + testUsageEntry("event-a"), testUsageEntry("event-b"), + }); err != nil { + t.Fatal(err) + } + + first, err := db.ClaimEvents(time.Now(), 1, time.Minute) + if err != nil { + t.Fatal(err) + } + second, err := db.ClaimEvents(time.Now(), 1, time.Minute) + if err != nil { + t.Fatal(err) + } + if len(first) != 1 || len(second) != 1 { + t.Fatalf("claims = %d and %d, want 1 each", len(first), len(second)) + } + if first[0].ID == second[0].ID { + t.Fatalf("both callers claimed %q; claims must not overlap", first[0].ID) + } + + // Everything is claimed, and no lease has expired, so there is nothing left. + third, err := db.ClaimEvents(time.Now(), 10, time.Minute) + if err != nil { + t.Fatal(err) + } + if len(third) != 0 { + t.Fatalf("third claim = %d events, want 0 while leases are live", len(third)) + } +} + +// A batch left claimed by a process that died must become claimable again once +// its lease expires — otherwise those events would never be sent. +func TestClaimEventsReclaimsAfterLeaseExpires(t *testing.T) { + db := openTestDB(t) + if _, err := db.InsertEvents([]usage.Entry{testUsageEntry("event-a")}); err != nil { + t.Fatal(err) + } + + claimed, err := db.ClaimEvents(time.Now(), 10, time.Minute) + if err != nil { + t.Fatal(err) + } + if len(claimed) != 1 { + t.Fatalf("first claim = %d, want 1", len(claimed)) + } + + // Before the lease expires the batch stays with its owner. + if again, err := db.ClaimEvents(time.Now(), 10, time.Minute); err != nil { + t.Fatal(err) + } else if len(again) != 0 { + t.Fatalf("claimed %d events while the lease was live, want 0", len(again)) + } + + // After it expires the batch is reclaimable. + reclaimed, err := db.ClaimEvents(time.Now().Add(2*time.Minute), 10, time.Minute) + if err != nil { + t.Fatal(err) + } + if len(reclaimed) != 1 || reclaimed[0].ID != "event-a" { + t.Fatalf("reclaimed = %+v, want event-a", reclaimed) + } +} + +// Expired leases are reclaimed alongside fresh work rather than only when the +// queue runs dry. A machine whose queue never empties would otherwise never +// reach the recovery path, stranding those events indefinitely. +func TestClaimEventsReclaimsExpiredLeasesAlongsideFreshWork(t *testing.T) { + db := openTestDB(t) + if _, err := db.InsertEvents([]usage.Entry{testUsageEntry("event-stuck")}); err != nil { + t.Fatal(err) + } + if _, err := db.ClaimEvents(time.Now(), 10, time.Minute); err != nil { + t.Fatal(err) + } + // event-stuck is now claimed with a lease that has long expired. + later := time.Now().Add(2 * time.Minute) + if _, err := db.InsertEvents([]usage.Entry{testUsageEntry("event-fresh")}); err != nil { + t.Fatal(err) + } + + claimed, err := db.ClaimEvents(later, 10, time.Minute) + if err != nil { + t.Fatal(err) + } + got := make(map[string]bool, len(claimed)) + for _, entry := range claimed { + got[entry.ID] = true + } + if !got["event-fresh"] || !got["event-stuck"] { + t.Fatalf("claimed = %v, want both event-fresh and the stranded event-stuck", got) + } +} + +// A batch that fills the limit with fresh work leaves expired leases for the +// next pass rather than exceeding what the caller asked for. +func TestClaimEventsRespectsLimitWhileReclaiming(t *testing.T) { + db := openTestDB(t) + if _, err := db.InsertEvents([]usage.Entry{testUsageEntry("event-stuck")}); err != nil { + t.Fatal(err) + } + if _, err := db.ClaimEvents(time.Now(), 10, time.Minute); err != nil { + t.Fatal(err) + } + later := time.Now().Add(2 * time.Minute) + if _, err := db.InsertEvents([]usage.Entry{ + testUsageEntry("event-a"), testUsageEntry("event-b"), + }); err != nil { + t.Fatal(err) + } + + claimed, err := db.ClaimEvents(later, 2, time.Minute) + if err != nil { + t.Fatal(err) + } + if len(claimed) != 2 { + t.Fatalf("claimed %d events, want exactly the 2 requested", len(claimed)) + } +} + +// A failed upload returns its batch to the queue rather than leaving it +// claimed until the lease runs out. +func TestMarkEventsUploadFailedReleasesTheClaim(t *testing.T) { + db := openTestDB(t) + if _, err := db.InsertEvents([]usage.Entry{testUsageEntry("event-a")}); err != nil { + t.Fatal(err) + } + if _, err := db.ClaimEvents(time.Now(), 10, time.Hour); err != nil { + t.Fatal(err) + } + if err := db.MarkEventsUploadFailed([]string{"event-a"}, "offline"); err != nil { + t.Fatal(err) + } + + // Due again after backoff, without waiting out the hour-long lease. + due := time.Now().Add(backoffBaseSeconds*time.Second + time.Second) + claimed, err := db.ClaimEvents(due, 10, time.Minute) + if err != nil { + t.Fatal(err) + } + if len(claimed) != 1 || claimed[0].ID != "event-a" { + t.Fatalf("claimed = %+v, want event-a back in the queue", claimed) + } +} + +// Events the server never accounted for must return to the queue immediately, +// not stay claimed until their lease runs out. +func TestReleaseClaimsReturnsEventsToTheQueue(t *testing.T) { + db := openTestDB(t) + if _, err := db.InsertEvents([]usage.Entry{testUsageEntry("event-a")}); err != nil { + t.Fatal(err) + } + if _, err := db.ClaimEvents(time.Now(), 10, time.Hour); err != nil { + t.Fatal(err) + } + + // While claimed it is invisible to the queue. + if n, err := db.PendingCount(time.Now()); err != nil { + t.Fatal(err) + } else if n != 0 { + t.Fatalf("pending count = %d while claimed, want 0", n) + } + + if err := db.ReleaseClaims([]string{"event-a"}); err != nil { + t.Fatal(err) + } + + if n, err := db.PendingCount(time.Now()); err != nil { + t.Fatal(err) + } else if n != 1 { + t.Fatalf("pending count = %d after release, want 1", n) + } + claimed, err := db.ClaimEvents(time.Now(), 10, time.Minute) + if err != nil { + t.Fatal(err) + } + if len(claimed) != 1 || claimed[0].ID != "event-a" { + t.Fatalf("claimed = %+v after release, want event-a", claimed) + } +} + +// Releasing must not resurrect events that already resolved. +func TestReleaseClaimsLeavesResolvedEventsAlone(t *testing.T) { + db := openTestDB(t) + if _, err := db.InsertEvents([]usage.Entry{testUsageEntry("event-a")}); err != nil { + t.Fatal(err) + } + if _, err := db.ClaimEvents(time.Now(), 10, time.Hour); err != nil { + t.Fatal(err) + } + if err := db.MarkEventsUploaded([]string{"event-a"}); err != nil { + t.Fatal(err) + } + if err := db.ReleaseClaims([]string{"event-a"}); err != nil { + t.Fatal(err) + } + + if n, err := db.PendingCount(time.Now()); err != nil { + t.Fatal(err) + } else if n != 0 { + t.Fatalf("pending count = %d, want 0: an uploaded event was resurrected", n) + } +} diff --git a/internal/usageprovider/provider.go b/internal/usageprovider/provider.go index 21b6e63..8183db2 100644 --- a/internal/usageprovider/provider.go +++ b/internal/usageprovider/provider.go @@ -51,6 +51,36 @@ func (b Base) WithFilterSet(filter usage.FileFilter) Base { return b } +// StreamFiles walks files in order, parsing each from where the previous scan +// stopped and handing its entries to emit before moving to the next. +// +// It is the shared body of every append-only provider's StreamEntries: the +// providers differ only in how they find their files and how they parse one, +// so those are the two arguments. Keeping the loop here means the rules that +// make a resumed scan safe — parse from the recorded offset, emit before +// advancing it — are stated once rather than re-derived per provider. +func StreamFiles( + files []string, + filter usage.FileFilter, + parse func(path string, start int64) ([]usage.Entry, int64, error), + emit func(path string, entries []usage.Entry, offset int64) error, + resume func(path string) int64, +) error { + for _, file := range files { + if filter != nil && !filter(file) { + continue + } + entries, offset, err := parse(file, resume(file)) + if err != nil { + return err + } + if err := emit(file, entries, offset); err != nil { + return err + } + } + return nil +} + // SortEntriesByTimestampDesc orders entries newest first, the order every // provider returns them in. func SortEntriesByTimestampDesc(entries []usage.Entry) []usage.Entry { @@ -118,11 +148,19 @@ func SetSource(entry *usage.Entry, source string, line int, start, end int64) { entry.SourceEnd = end } +// StableEntryID derives a deterministic id from an entry's source position +// and contents. +// +// The position is the byte offset, not the line number. A scan that resumes +// mid-file starts counting lines from 1 again, so a line-based id would give +// the same record a different identity depending on where the previous pass +// happened to stop. The byte offset is a property of the file itself and does +// not move. func StableEntryID(entry usage.Entry, extra ...string) string { parts := []string{ string(entry.Provider), entry.SourceFile, - strconv.Itoa(entry.SourceLine), + strconv.FormatInt(entry.SourceStart, 10), entry.Timestamp.Format(time.RFC3339Nano), entry.Project, entry.ProjectPath, diff --git a/internal/usagescan/scanner.go b/internal/usagescan/scanner.go index c346882..e6293a0 100644 --- a/internal/usagescan/scanner.go +++ b/internal/usagescan/scanner.go @@ -128,6 +128,9 @@ func (s *Scanner) Scan(providerDirs map[usage.Provider][]string) (Result, error) func (s *Scanner) scanProvider(provider usageprovider.Provider, paths []string, scanned map[string]usagedb.FileState) (ProviderResult, error) { var result ProviderResult configured := providerWithPaths(provider, paths) + if streamer, ok := configured.(streamProvider); ok { + return s.scanStreaming(streamer, scanned) + } pending := make(map[string]usagedb.FileState) if filterable, ok := configured.(filterConfiguredProvider); ok && scanned != nil { configured = filterable.WithFileFilter(func(path string) bool { @@ -136,7 +139,7 @@ func (s *Scanner) scanProvider(provider usageprovider.Provider, paths []string, return true } state := usagedb.FileState{Size: info.Size(), MtimeNS: info.ModTime().UnixNano()} - if previous, ok := scanned[path]; ok && previous == state { + if previous, ok := scanned[path]; ok && unchanged(previous, state) { return false } pending[path] = state @@ -162,6 +165,113 @@ func (s *Scanner) scanProvider(provider usageprovider.Provider, paths []string, return result, nil } +// unchanged reports whether a file's current stat matches the one recorded +// when it was last scanned, meaning it holds nothing new. +// +// Only size and mtime are compared. Offset is where the last pass stopped +// reading, not a property of the file, and a file that ended mid-line has an +// offset behind its size while still being unchanged. +func unchanged(previous, current usagedb.FileState) bool { + return previous.Size == current.Size && previous.MtimeNS == current.MtimeNS +} + +// fullyConsumed reports whether the recorded offset reached the end of the +// recorded size. +// +// A file can be unchanged and still hold unread bytes: a trailing partial line +// leaves the offset short, and so would any state written by a version that +// stat'd a file after parsing it. Treating "unchanged" alone as "nothing to +// do" would skip those bytes for as long as the file stays quiet — which for a +// finished session transcript is forever. +func fullyConsumed(state usagedb.FileState) bool { + return state.Offset >= state.Size +} + +// scanStreaming ingests a provider one source file at a time, committing each +// file's events and then recording where parsing stopped. +// +// The order within a file is what makes an interrupted scan safe: events are +// stored first, and only then does the file's offset advance. A crash between +// the two costs a re-parse of one file, which is wasted work. The reverse +// order would record progress over events that were never stored, and that is +// lost data — see UpsertScannedFiles. +func (s *Scanner) scanStreaming(streamer streamProvider, scanned map[string]usagedb.FileState) (ProviderResult, error) { + var result ProviderResult + + // A file whose stat is unchanged holds nothing new. Skipping it here is + // what keeps a steady-state scan proportional to what was just written + // rather than to the whole history on disk. + if filterable, ok := streamer.(filterConfiguredProvider); ok && scanned != nil { + if restreamed, ok := filterable.WithFileFilter(func(path string) bool { + state, seen := scanned[path] + if !seen { + return true + } + info, err := os.Stat(path) + if err != nil || info.IsDir() { + return true + } + if !unchanged(state, usagedb.FileState{Size: info.Size(), MtimeNS: info.ModTime().UnixNano()}) { + return true + } + // Unchanged, but the last pass stopped short of the end. Those + // bytes are still unread and this is the only thing that will + // come back for them. + return !fullyConsumed(state) + }).(streamProvider); ok { + streamer = restreamed + } + } + + // Stats taken when each file was handed to the parser. Recording the size + // as it was *before* parsing is what makes an append during the parse + // safe: stat'ing afterwards would pair a size that already counts the new + // bytes with an offset that stops short of them, so the next scan would + // see an unchanged size, skip the file, and lose everything written while + // it was being read. + seenAt := make(map[string]usagedb.FileState) + + resume := func(path string) int64 { + info, err := os.Stat(path) + if err == nil { + seenAt[path] = usagedb.FileState{Size: info.Size(), MtimeNS: info.ModTime().UnixNano()} + } + state, ok := scanned[path] + if !ok || err != nil { + return 0 + } + // A file shorter than where we stopped was truncated or replaced, so + // the stored offset points into content that no longer exists. The + // only safe reading is from the beginning. + if info.Size() < state.Offset { + return 0 + } + return state.Offset + } + + emit := func(path string, entries []usage.Entry, offset int64) error { + s.applyProjectFiles(entries) + inserted, err := s.db.InsertEvents(entries) + if err != nil { + return err + } + result.EventsParsed += len(entries) + result.EventsInserted += inserted + + state, ok := seenAt[path] + if !ok { + // Without a stat there is no honest state to record. Leaving the + // previous one alone re-parses this file next time, which is the + // harmless outcome. + return nil + } + state.Offset = offset + return s.db.UpsertScannedFiles(map[string]usagedb.FileState{path: state}) + } + + return result, streamer.StreamEntries(resume, emit) +} + // applyProjectFiles rewrites each entry's identity from the nearest project // identity file. An identity file is an optional override: one that exists // but cannot be read is warned about and skipped — a stray unreadable @@ -228,6 +338,24 @@ type pathConfiguredProvider interface { WithPaths(paths []string) usageprovider.Provider } +// streamProvider is implemented by providers whose sources are append-only +// files that can be parsed one at a time and resumed mid-file. The scanner +// commits each file's events as they arrive instead of holding an entire +// provider's history in memory and writing it once at the end. +// +// It is deliberately optional. A provider whose events come from a SQLite +// database, or one that must join across sources before it knows anything, +// has no per-file boundary to commit on and keeps using Entries(). +type streamProvider interface { + // StreamEntries parses each source file and calls emit once per file with + // that file's entries and the offset to resume from next time. Providers + // pass the resume offset from resume(path) to their reader. + // + // emit returning an error aborts the scan: it means the events could not + // be stored, and continuing would advance past data that was never saved. + StreamEntries(resume func(path string) int64, emit func(path string, entries []usage.Entry, offset int64) error) error +} + // filterConfiguredProvider is implemented by providers that can skip source // files the filter rejects. Providers without it are always fully scanned. type filterConfiguredProvider interface { diff --git a/internal/usagescan/stream_test.go b/internal/usagescan/stream_test.go new file mode 100644 index 0000000..603c39a --- /dev/null +++ b/internal/usagescan/stream_test.go @@ -0,0 +1,285 @@ +package usagescan + +import ( + "os" + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/claude" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/kimi" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/pi" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/qwen" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usagedb" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +func claudeUsageLine(id string) string { + return `{"timestamp":"2026-06-04T01:02:03Z","cwd":"/Users/me/workspace/tokitoki","requestId":"req-` + id + + `","message":{"id":"msg-` + id + `","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}}` + "\n" +} + +func newClaudeScan(t *testing.T) (*Scanner, map[usage.Provider][]string, string) { + t.Helper() + dir := t.TempDir() + claudeDir := filepath.Join(dir, "claude") + sessionDir := filepath.Join(claudeDir, "projects", "-Users-me-workspace-tokitoki") + if err := os.MkdirAll(sessionDir, 0o700); err != nil { + t.Fatal(err) + } + sessionFile := filepath.Join(sessionDir, "session-a.jsonl") + if err := os.WriteFile(sessionFile, []byte(claudeUsageLine("1")), 0o600); err != nil { + t.Fatal(err) + } + + db, err := usagedb.Open(filepath.Join(dir, "usage.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + + return New(db, claude.Provider{}), map[usage.Provider][]string{usage.ProviderClaude: {claudeDir}}, sessionFile +} + +// An append must cost only the appended record. The whole point of storing a +// resume offset is that a transcript that has grown by one line is not read +// from the beginning again. +func TestScanStreamingParsesOnlyAppendedEntries(t *testing.T) { + scanner, dirs, sessionFile := newClaudeScan(t) + + result, err := scanner.Scan(dirs) + if err != nil { + t.Fatal(err) + } + if parsed := result.Providers[usage.ProviderClaude].EventsParsed; parsed != 1 { + t.Fatalf("first events parsed = %d, want 1", parsed) + } + + result, err = scanner.Scan(dirs) + if err != nil { + t.Fatal(err) + } + if parsed := result.Providers[usage.ProviderClaude].EventsParsed; parsed != 0 { + t.Fatalf("unchanged file events parsed = %d, want 0", parsed) + } + + appendLine(t, sessionFile, claudeUsageLine("2")) + + result, err = scanner.Scan(dirs) + if err != nil { + t.Fatal(err) + } + claudeResult := result.Providers[usage.ProviderClaude] + if claudeResult.EventsParsed != 1 { + t.Fatalf("appended file events parsed = %d, want 1 (only the new line)", claudeResult.EventsParsed) + } + if claudeResult.EventsInserted != 1 { + t.Fatalf("appended file events inserted = %d, want 1", claudeResult.EventsInserted) + } +} + +// A transcript that shrank was truncated or replaced, so the stored offset +// points at content that no longer exists and the file must be re-read whole. +func TestScanStreamingRescansTruncatedFile(t *testing.T) { + scanner, dirs, sessionFile := newClaudeScan(t) + + if _, err := scanner.Scan(dirs); err != nil { + t.Fatal(err) + } + appendLine(t, sessionFile, claudeUsageLine("2")) + if _, err := scanner.Scan(dirs); err != nil { + t.Fatal(err) + } + + // Replace the file with a shorter one holding a different record. + if err := os.WriteFile(sessionFile, []byte(claudeUsageLine("3")), 0o600); err != nil { + t.Fatal(err) + } + + result, err := scanner.Scan(dirs) + if err != nil { + t.Fatal(err) + } + claudeResult := result.Providers[usage.ProviderClaude] + if claudeResult.EventsParsed != 1 { + t.Fatalf("truncated file events parsed = %d, want 1 (re-read from the start)", claudeResult.EventsParsed) + } + if claudeResult.EventsInserted != 1 { + t.Fatalf("truncated file events inserted = %d, want 1", claudeResult.EventsInserted) + } +} + +// Progress is recorded per file, so a scan that fails partway leaves the +// files it already stored marked as done. +func TestScanStreamingRecordsOffsetPerFile(t *testing.T) { + scanner, dirs, sessionFile := newClaudeScan(t) + + if _, err := scanner.Scan(dirs); err != nil { + t.Fatal(err) + } + + db := scanner.db + states, err := db.ScannedFiles() + if err != nil { + t.Fatal(err) + } + state, ok := states[sessionFile] + if !ok { + t.Fatalf("no scanned state recorded for %s", sessionFile) + } + info, err := os.Stat(sessionFile) + if err != nil { + t.Fatal(err) + } + if state.Offset != info.Size() { + t.Fatalf("offset = %d, want %d", state.Offset, info.Size()) + } + if state.Size != info.Size() { + t.Fatalf("size = %d, want %d", state.Size, info.Size()) + } +} + +func appendLine(t *testing.T, path, data string) { + t.Helper() + fh, err := os.OpenFile(path, os.O_APPEND|os.O_WRONLY, 0o600) + if err != nil { + t.Fatal(err) + } + defer fh.Close() + if _, err := fh.WriteString(data); err != nil { + t.Fatal(err) + } +} + +// Every provider that opts into streaming must satisfy the same contract: a +// full first pass, no work when nothing changed, and only the appended record +// on the pass after an append. +func TestScanStreamingProvidersReadOnlyAppendedRecords(t *testing.T) { + tests := []struct { + name string + provider usageprovider.Provider + id usage.Provider + relPath string + line func(n string) string + }{ + { + name: "claude", + provider: claude.Provider{}, + id: usage.ProviderClaude, + relPath: filepath.Join("projects", "-Users-me-workspace-tokitoki", "session-a.jsonl"), + line: func(n string) string { return claudeUsageLine(n) }, + }, + { + name: "pi", + provider: pi.Provider{}, + id: usage.ProviderPi, + relPath: "sess_abc.jsonl", + line: func(n string) string { + return `{"type":"message","timestamp":"2026-06-04T01:02:0` + n + `Z","message":{"role":"assistant","model":"m","usage":{"input":1,"output":` + n + `}}}` + "\n" + }, + }, + { + name: "qwen", + provider: qwen.Provider{}, + id: usage.ProviderQwen, + relPath: filepath.Join("projects", "proj", "chats", "chat-a.jsonl"), + line: func(n string) string { + return `{"type":"assistant","timestamp":"2026-06-04T01:02:0` + n + `Z","model":"m","usageMetadata":{"promptTokenCount":1,"candidatesTokenCount":` + n + `}}` + "\n" + }, + }, + { + name: "kimi", + provider: kimi.Provider{}, + id: usage.ProviderKimi, + relPath: filepath.Join("sessions", "ws", "sess", "wire.jsonl"), + line: func(n string) string { + return `{"type":"usage.record","usageScope":"turn","time":"2026-06-04T01:02:0` + n + `Z","model":"m","usage":{"inputOther":1,"output":` + n + `}}` + "\n" + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + dir := t.TempDir() + root := filepath.Join(dir, "root") + file := filepath.Join(root, tt.relPath) + if err := os.MkdirAll(filepath.Dir(file), 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(file, []byte(tt.line("1")), 0o600); err != nil { + t.Fatal(err) + } + + db, err := usagedb.Open(filepath.Join(dir, "usage.db")) + if err != nil { + t.Fatal(err) + } + t.Cleanup(func() { _ = db.Close() }) + + scanner := New(db, tt.provider) + dirs := map[usage.Provider][]string{tt.id: {root}} + + first, err := scanner.Scan(dirs) + if err != nil { + t.Fatal(err) + } + if parsed := first.Providers[tt.id].EventsParsed; parsed != 1 { + t.Fatalf("first pass parsed = %d, want 1", parsed) + } + + second, err := scanner.Scan(dirs) + if err != nil { + t.Fatal(err) + } + if parsed := second.Providers[tt.id].EventsParsed; parsed != 0 { + t.Fatalf("unchanged pass parsed = %d, want 0", parsed) + } + + appendLine(t, file, tt.line("2")) + + third, err := scanner.Scan(dirs) + if err != nil { + t.Fatal(err) + } + result := third.Providers[tt.id] + if result.EventsParsed != 1 { + t.Fatalf("appended pass parsed = %d, want 1 (only the new record)", result.EventsParsed) + } + if result.EventsInserted != 1 { + t.Fatalf("appended pass inserted = %d, want 1", result.EventsInserted) + } + }) + } +} + +// A file appended to while it is being parsed must not be skipped on the next +// scan. The recorded size describes the file as the parser found it, so the +// appended bytes make the file look changed and get read next time. +func TestScanStreamingDoesNotLoseDataAppendedDuringTheScan(t *testing.T) { + scanner, dirs, sessionFile := newClaudeScan(t) + if _, err := scanner.Scan(dirs); err != nil { + t.Fatal(err) + } + + // Reproduce the state an append during the parse leaves behind: the file + // now holds a second record, but the scan only consumed the first. + appendLine(t, sessionFile, claudeUsageLine("2")) + info, err := os.Stat(sessionFile) + if err != nil { + t.Fatal(err) + } + consumed := int64(len(claudeUsageLine("1"))) + if err := scanner.db.UpsertScannedFiles(map[string]usagedb.FileState{ + sessionFile: {Size: info.Size(), MtimeNS: info.ModTime().UnixNano(), Offset: consumed}, + }); err != nil { + t.Fatal(err) + } + + result, err := scanner.Scan(dirs) + if err != nil { + t.Fatal(err) + } + if parsed := result.Providers[usage.ProviderClaude].EventsParsed; parsed != 1 { + t.Fatalf("parsed = %d, want 1: the appended record was skipped and is lost", parsed) + } +} diff --git a/internal/usageupload/uploader.go b/internal/usageupload/uploader.go index 5b73a31..c90cdae 100644 --- a/internal/usageupload/uploader.go +++ b/internal/usageupload/uploader.go @@ -38,6 +38,15 @@ const ( // uploadedRetention is how long uploaded events are kept before pruning. uploadedRetention = 30 * 24 * time.Hour + + // uploadLease is how long a claimed batch stays claimed. + // + // It must outlast a healthy upload, or a batch still being sent would be + // reclaimed and sent twice. Callers bound an upload by DefaultUploadTimeout + // (2 minutes), so this leaves a wide margin above it: expiring early costs + // duplicate round-trips on every slow upload, while expiring late costs + // delay only when a process actually died. The asymmetry says round up. + uploadLease = 5 * time.Minute ) type Payload struct { @@ -130,12 +139,30 @@ func Upload(ctx context.Context, settings agent.Settings, events []usage.Entry) // Uploaded events older than uploadedRetention are pruned before sending. // SyncPending continues until all pending+failed events are sent or an error occurs. func SyncPending(ctx context.Context, settings agent.Settings, db *usagedb.DB) error { + return syncPending(ctx, settings, db, 0) +} + +// SyncPendingBatches drains at most maxBatches batches instead of continuing +// until the queue is empty. +// +// A drain running beside a scan needs this: the scan keeps adding events, so +// an unbounded drain follows it down to whatever arrived in the last few +// milliseconds and spends a request on each handful. Stopping after a bounded +// number of batches lets the queue refill into full requests between passes. +func SyncPendingBatches(ctx context.Context, settings agent.Settings, db *usagedb.DB, maxBatches int) error { + return syncPending(ctx, settings, db, maxBatches) +} + +func syncPending(ctx context.Context, settings agent.Settings, db *usagedb.DB, maxBatches int) error { if _, err := db.PruneUploaded(time.Now().Add(-uploadedRetention)); err != nil { return err } - for { - events, err := db.PendingEvents(time.Now(), queueBatchSize) + for sent := 0; maxBatches <= 0 || sent < maxBatches; sent++ { + // Claiming rather than reading marks this batch as ours, so a second + // uploader running at the same time takes a different one instead of + // re-sending this. + events, err := db.ClaimEvents(time.Now(), queueBatchSize, uploadLease) if err != nil { return err } @@ -190,10 +217,18 @@ func SyncPending(ctx context.Context, settings agent.Settings, db *usagedb.DB) e // If not, it's a server bug or response parsing error. totalAcknowledged := len(response.Accepted) + len(response.Duplicate) + len(response.Rejected) if totalAcknowledged != len(events) { - return fmt.Errorf("usage upload incomplete: server acknowledged %d/%d events (accepted=%d, duplicate=%d, rejected=%d)", - totalAcknowledged, len(events), len(response.Accepted), len(response.Duplicate), len(response.Rejected)) + // Whatever went unaccounted for is still claimed. Releasing it + // puts it back in the queue for the next pass; left alone it would + // be invisible to both the pending count and the next claim until + // its lease expired. + releaseErr := db.ReleaseClaims(unacknowledged(events, response)) + return errors.Join(fmt.Errorf("usage upload incomplete: server acknowledged %d/%d events (accepted=%d, duplicate=%d, rejected=%d)", + totalAcknowledged, len(events), len(response.Accepted), len(response.Duplicate), len(response.Rejected)), releaseErr) } } + // Reached the batch limit with events still queued; the caller drains the + // rest on its next pass. + return nil } // logRejections reports what the server refused. Rejections repeat: a server @@ -213,6 +248,29 @@ func logRejections(rejected []Reject) { } } +// unacknowledged returns the ids the server said nothing about — neither +// accepted, duplicate, nor rejected. +func unacknowledged(events []usage.Entry, response Response) []string { + seen := make(map[string]bool, len(response.Accepted)+len(response.Duplicate)+len(response.Rejected)) + for _, id := range response.Accepted { + seen[id] = true + } + for _, id := range response.Duplicate { + seen[id] = true + } + for _, rejected := range response.Rejected { + seen[rejected.ID] = true + } + + missing := make([]string, 0) + for _, event := range events { + if event.ID != "" && !seen[event.ID] { + missing = append(missing, event.ID) + } + } + return missing +} + func eventIDs(events []usage.Entry) []string { ids := make([]string, 0, len(events)) for _, event := range events { diff --git a/pkg/agentlib/agentlib.go b/pkg/agentlib/agentlib.go index 7de9ced..e228d4b 100644 --- a/pkg/agentlib/agentlib.go +++ b/pkg/agentlib/agentlib.go @@ -12,6 +12,7 @@ import ( "path/filepath" "runtime" "strings" + "sync" "time" "github.com/tokitoki-dev/tokitoki-cli/internal/agent" @@ -243,9 +244,35 @@ func (c *Client) VerifyAPIKey(ctx context.Context) (bool, error) { return deviceauth.VerifyKey(ctx, usageupload.BaseURL(), apiKey) } +const ( + // syncDrainPoll is how often the upload half re-checks the queue while a + // scan is still running. + syncDrainPoll = 250 * time.Millisecond + + // syncDrainMinBatch is how many events must be queued before a drain runs + // alongside a still-running scan. + // + // Without it the drain chases the scan: it sends whatever landed in the + // last few milliseconds, so a scan producing a trickle of events turns + // into a request per handful. Waiting for a worthwhile batch trades a + // little latency — bounded by the scan itself, since the final drain sends + // everything regardless — for far fewer, fuller requests. + syncDrainMinBatch = 500 +) + // Sync scans selected provider directories and uploads newly discovered // events. Scanning is local and always runs; without a configured API key the // events simply stay queued and upload resumes once a key is saved. +// +// The two halves run at the same time. A scan queues each file's events as it +// finishes that file, so the upload half has work to send long before the scan +// is done — on a first run over a large history that is the difference between +// uploading throughout the scan and sitting idle until it ends. They share no +// state but the queue: the scan writes to it, the drain reads from it. +// +// Sync returns only once both halves are finished, because callers are +// one-shot processes that exit when it returns. The drain therefore keeps +// polling until the scan has stopped producing and the queue is empty. func (c *Client) Sync(ctx context.Context, options SyncOptions) error { providerDirs := normalizeProviderDirs(options.ProviderDirs) if len(providerDirs) == 0 { @@ -259,6 +286,11 @@ func (c *Client) Sync(ctx context.Context, options SyncOptions) error { if err != nil { return err } + settings, err := agent.New(fileStore, c.logger).Settings() + if err != nil { + return err + } + usageDB, err := usagedb.Open(store.UsageDBPath(c.dataDir)) if err != nil { return err @@ -275,10 +307,107 @@ func (c *Client) Sync(ctx context.Context, options SyncOptions) error { Out: io.Discard, } - // Two phases, two locks. Ingestion mutates shared local state and runs - // under the data lock; the drain talks to the network for up to the whole - // upload timeout and must not make other processes' ingestion wait on it. - if err := c.withDataLock(app.Ingest); err != nil { + // Without a key there is nothing to drain into, so the scan runs alone and + // its events wait in the queue for a run that has one. + if settings.APIKey == "" { + c.logger.Debug("skip upload; API key is not configured") + return c.withDataLock(app.Ingest) + } + + scanDone := make(chan struct{}) + var scanErr, uploadErr error + var wg sync.WaitGroup + wg.Add(2) + + go func() { + defer wg.Done() + defer close(scanDone) + // Ingestion mutates shared local state and runs under the data lock; + // the drain talks to the network and holds no lock, so a slow upload + // never makes another process's ingestion wait. + scanErr = c.withDataLock(app.Ingest) + }() + + go func() { + defer wg.Done() + uploadErr = c.drainUntil(ctx, settings, usageDB, scanDone) + }() + + wg.Wait() + return errors.Join(scanErr, uploadErr) +} + +// drainUntil uploads queued events until done is closed and the queue has been +// drained after that, or until ctx ends. +// +// The final drain after done matters: the scan's last file is queued moments +// before it finishes, and a drain that stopped at the same instant would leave +// those events for the next run. +func (c *Client) drainUntil(ctx context.Context, settings agent.Settings, usageDB *usagedb.DB, done <-chan struct{}) error { + for { + select { + case <-done: + // The scan has stopped adding work, so there is nothing left to + // wait for. This pass sends everything, however little, and its + // error is the one worth reporting: it is the last chance these + // events had to go out during this run. + return usageupload.SyncPending(ctx, settings, usageDB) + case <-ctx.Done(): + // The run is out of time. The scan's events are queued and a later + // run sends them, so this is not a failure of the sync. + return nil + default: + } + + // While the scan is still producing, send only once enough has piled + // up to fill a request. Draining on every tick would chase the scan + // and spend a round-trip on whatever handful arrived since the last + // one. + queued, err := usageDB.PendingCount(time.Now()) + if err != nil { + return err + } + if queued >= syncDrainMinBatch { + // One batch per pass. Draining to empty here would follow the scan + // down to its trickle; stopping lets the queue refill. + // + // A failure here is not fatal to the run. The scan is still going, + // and returning would leave it with no uploader at all — including + // for the final drain, which is where the events actually need to + // be sent. Uploads retry with backoff, so the next pass tries + // again and the final drain reports whatever still fails. + if err := usageupload.SyncPendingBatches(ctx, settings, usageDB, 1); err != nil { + c.logger.Debug("mid-scan upload failed; will retry", "error", err) + } + } + + select { + case <-done: + return usageupload.SyncPending(ctx, settings, usageDB) + case <-ctx.Done(): + return nil + case <-time.After(syncDrainPoll): + } + } +} + +// Upload drains queued events to the server without scanning first. +// +// Scanning and uploading share no state but the local queue: a scan writes +// events into it and an upload drains them. Nothing about the drain depends on +// a scan having just run, so a caller that wants the two to proceed at their +// own pace runs Scan and Upload on separate schedules — an upload no longer +// waits for a scan to finish before sending what is already queued, and a slow +// or failing scan cannot hold back events that were queued minutes ago. +// +// A missing API key is not an error here. Events stay queued until a key +// exists, which is the same thing Sync does. +func (c *Client) Upload(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } + fileStore, err := store.Open(c.dataDir) + if err != nil { return err } settings, err := agent.New(fileStore, c.logger).Settings() @@ -289,7 +418,40 @@ func (c *Client) Sync(ctx context.Context, options SyncOptions) error { c.logger.Debug("skip upload; API key is not configured") return nil } - return c.withUploadLock(func() error { return app.Upload(ctx) }) + + usageDB, err := usagedb.Open(store.UsageDBPath(c.dataDir)) + if err != nil { + return err + } + defer usageDB.Close() + + return usageupload.SyncPending(ctx, settings, usageDB) +} + +// Scan ingests provider directories into the local queue without uploading. +// It is the other half of Sync, for callers running the two on separate +// schedules. +func (c *Client) Scan(options SyncOptions) error { + providerDirs := normalizeProviderDirs(options.ProviderDirs) + if len(providerDirs) == 0 { + return ErrNoScanDirectories + } + + usageDB, err := usagedb.Open(store.UsageDBPath(c.dataDir)) + if err != nil { + return err + } + defer usageDB.Close() + + scanner := usagescan.New(usageDB) + scanner.Logger = c.logger + app := &cli.App{ + UsageDB: usageDB, + Scanner: scanner, + ProviderDirs: providerDirs, + Out: io.Discard, + } + return c.withDataLock(app.Ingest) } // SendHeartbeat persists an IDE activity event before attempting upload. If @@ -365,7 +527,7 @@ func (c *Client) SendHeartbeat(ctx context.Context, heartbeat Heartbeat) error { } defer usageDB.Close() - // Queue the event under the data lock, then drain under the upload lock. + // Queue the event under the data lock, then drain without one. // The drain can take the whole network timeout; heartbeats from other // editors must be able to enqueue while it runs, not wait behind it. // @@ -392,9 +554,7 @@ func (c *Client) SendHeartbeat(ctx context.Context, heartbeat Heartbeat) error { c.logger.Debug("skip upload; API key is not configured") return nil } - return c.withUploadLock(func() error { - return usageupload.SyncPending(ctx, settings, usageDB) - }) + return usageupload.SyncPending(ctx, settings, usageDB) } func applyProjectFile(heartbeat *Heartbeat) error { @@ -436,23 +596,6 @@ func (c *Client) withDataLock(fn func() error) error { return fn() } -// withUploadLock runs fn while holding the cross-process upload lock. When -// another process already holds it, that process is draining the same queue -// this one just wrote to, so there is nothing left to do here: the events are -// safely queued and "busy" is success, not failure. -func (c *Client) withUploadLock(fn func() error) error { - lock, err := store.AcquireLock(c.dataDir, store.UploadLockFile, 0) - if errors.Is(err, store.ErrLockBusy) { - c.logger.Debug("another tokitoki process is uploading; events stay queued") - return nil - } - if err != nil { - return err - } - defer lock.Close() - return fn() -} - // DefaultDataDir returns the shared Tokitoki data directory. func DefaultDataDir() (string, error) { home, err := os.UserHomeDir() @@ -471,7 +614,7 @@ func DefaultProviderDirs() map[Provider][]string { return map[Provider][]string{ ProviderClaude: {filepath.Join(home, ".claude")}, ProviderCodex: {filepath.Join(home, ".codex")}, - ProviderCopilot: {filepath.Join(home, ".copilot", "otel")}, + ProviderCopilot: {filepath.Join(home, ".copilot")}, ProviderGemini: {filepath.Join(home, ".gemini", "tmp")}, ProviderKimi: {filepath.Join(home, ".kimi"), filepath.Join(home, ".kimi-code")}, ProviderQwen: {filepath.Join(home, ".qwen")}, diff --git a/pkg/agentlib/agentlib_test.go b/pkg/agentlib/agentlib_test.go index 7277f74..2fda389 100644 --- a/pkg/agentlib/agentlib_test.go +++ b/pkg/agentlib/agentlib_test.go @@ -87,7 +87,7 @@ func TestDefaultProviderDirsIncludesBuiltInProviders(t *testing.T) { want := map[Provider]string{ ProviderClaude: filepath.Join(home, ".claude"), ProviderCodex: filepath.Join(home, ".codex"), - ProviderCopilot: filepath.Join(home, ".copilot", "otel"), + ProviderCopilot: filepath.Join(home, ".copilot"), ProviderGemini: filepath.Join(home, ".gemini", "tmp"), ProviderKimi: filepath.Join(home, ".kimi"), ProviderQwen: filepath.Join(home, ".qwen"), diff --git a/pkg/agentlib/split_test.go b/pkg/agentlib/split_test.go new file mode 100644 index 0000000..8712484 --- /dev/null +++ b/pkg/agentlib/split_test.go @@ -0,0 +1,47 @@ +package agentlib + +import ( + "context" + "os" + "path/filepath" + "testing" +) + +// Upload must be callable with nothing queued and no key configured: the +// upload loop runs on its own schedule and cannot assume a scan just ran. +func TestUploadWithoutScanOrKeyIsNotAnError(t *testing.T) { + client, err := New(Options{DataDir: t.TempDir()}) + if err != nil { + t.Fatal(err) + } + if err := client.Upload(context.Background()); err != nil { + t.Fatalf("Upload on an empty queue = %v, want nil", err) + } +} + +// Scan must queue events without needing a key or an upload. +func TestScanQueuesWithoutUploading(t *testing.T) { + dir := t.TempDir() + root := filepath.Join(dir, "claude", "projects", "-tmp-p") + if err := os.MkdirAll(root, 0o700); err != nil { + t.Fatal(err) + } + line := `{"timestamp":"2026-06-04T01:02:03Z","cwd":"/tmp/p","requestId":"r1","message":{"id":"m1","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}}` + "\n" + if err := os.WriteFile(filepath.Join(root, "s.jsonl"), []byte(line), 0o600); err != nil { + t.Fatal(err) + } + + client, err := New(Options{DataDir: dir}) + if err != nil { + t.Fatal(err) + } + err = client.Scan(SyncOptions{ProviderDirs: map[Provider][]string{ + ProviderClaude: {filepath.Join(dir, "claude")}, + }}) + if err != nil { + t.Fatalf("Scan = %v, want nil", err) + } + if err := client.Scan(SyncOptions{}); err != ErrNoScanDirectories { + t.Fatalf("Scan with no dirs = %v, want ErrNoScanDirectories", err) + } +} diff --git a/pkg/agentlib/sync_parallel_test.go b/pkg/agentlib/sync_parallel_test.go new file mode 100644 index 0000000..6b01e43 --- /dev/null +++ b/pkg/agentlib/sync_parallel_test.go @@ -0,0 +1,240 @@ +package agentlib + +import ( + "context" + "fmt" + "net/http" + "net/http/httptest" + "os" + "path/filepath" + "sync" + "sync/atomic" + "testing" + "time" + + "compress/gzip" + "encoding/json" + "io" +) + +// Uploading must begin while the scan is still running. A scan queues each +// file's events as it finishes that file, so waiting for the whole scan before +// sending anything leaves the network idle for the entire ingest. +func TestSyncUploadsWhileScanning(t *testing.T) { + var firstUploadAt atomic.Int64 + var batches atomic.Int32 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + batches.Add(1) + firstUploadAt.CompareAndSwap(0, time.Now().UnixNano()) + body := io.Reader(r.Body) + if r.Header.Get("Content-Encoding") == "gzip" { + gz, err := gzip.NewReader(r.Body) + if err != nil { + w.WriteHeader(500) + return + } + defer gz.Close() + body = gz + } + var payload struct { + Events []struct { + ID string `json:"id"` + } + } + json.NewDecoder(body).Decode(&payload) + ids := make([]string, 0, len(payload.Events)) + for _, e := range payload.Events { + ids = append(ids, e.ID) + } + json.NewEncoder(w).Encode(map[string]any{"ok": true, "accepted": ids}) + })) + defer srv.Close() + t.Setenv("TOKITOKI_BASE_URL", srv.URL) + + dir := t.TempDir() + // 很多文件, 让扫描明显耗时 + for i := 0; i < 400; i++ { + p := filepath.Join(dir, "claude", "projects", fmt.Sprintf("-p-%03d", i)) + os.MkdirAll(p, 0o700) + var buf []byte + for j := 0; j < 300; j++ { + buf = append(buf, []byte(fmt.Sprintf(`{"timestamp":"2026-06-04T01:02:03Z","cwd":"/p/%03d","requestId":"r%03d-%03d","message":{"id":"m%03d-%03d","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}}`+"\n", i, i, j, i, j))...) + } + os.WriteFile(filepath.Join(p, "s.jsonl"), buf, 0o600) + } + os.MkdirAll(filepath.Join(dir, "config"), 0o700) + os.WriteFile(filepath.Join(dir, "config", "api_key"), []byte("k"), 0o600) + + client, err := New(Options{DataDir: dir}) + if err != nil { + t.Fatal(err) + } + + start := time.Now() + err = client.Sync(context.Background(), SyncOptions{ProviderDirs: map[Provider][]string{ + ProviderClaude: {filepath.Join(dir, "claude")}, + }}) + total := time.Since(start) + if err != nil { + t.Fatal(err) + } + + if firstUploadAt.Load() == 0 { + t.Fatal("从未上传") + } + firstAt := time.Unix(0, firstUploadAt.Load()).Sub(start) + fmt.Printf("总耗时=%v 首次上传发生在=%v (%.0f%%处) 批次=%d\n", + total.Round(time.Millisecond), firstAt.Round(time.Millisecond), + float64(firstAt)/float64(total)*100, batches.Load()) +} + +// Everything the scan queued must be uploaded by the time Sync returns. +// Callers are one-shot processes that exit immediately afterwards, so an +// event still sitting in the queue is an event that waited for the next run. +func TestSyncUploadsEverythingBeforeReturning(t *testing.T) { + var uploaded sync.Map + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body := io.Reader(r.Body) + if r.Header.Get("Content-Encoding") == "gzip" { + gz, err := gzip.NewReader(r.Body) + if err != nil { + w.WriteHeader(500) + return + } + defer gz.Close() + body = gz + } + var payload struct { + Events []struct { + ID string `json:"id"` + } `json:"events"` + } + json.NewDecoder(body).Decode(&payload) + ids := make([]string, 0, len(payload.Events)) + for _, e := range payload.Events { + uploaded.Store(e.ID, true) + ids = append(ids, e.ID) + } + json.NewEncoder(w).Encode(map[string]any{"ok": true, "accepted": ids}) + })) + defer srv.Close() + t.Setenv("TOKITOKI_BASE_URL", srv.URL) + + const files = 40 + dir := t.TempDir() + for i := 0; i < files; i++ { + p := filepath.Join(dir, "claude", "projects", fmt.Sprintf("-p-%03d", i)) + if err := os.MkdirAll(p, 0o700); err != nil { + t.Fatal(err) + } + line := fmt.Sprintf(`{"timestamp":"2026-06-04T01:02:03Z","cwd":"/p/%03d","requestId":"r%03d","message":{"id":"m%03d","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}}`+"\n", i, i, i) + if err := os.WriteFile(filepath.Join(p, "s.jsonl"), []byte(line), 0o600); err != nil { + t.Fatal(err) + } + } + if err := os.MkdirAll(filepath.Join(dir, "config"), 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "config", "api_key"), []byte("k"), 0o600); err != nil { + t.Fatal(err) + } + + client, err := New(Options{DataDir: dir}) + if err != nil { + t.Fatal(err) + } + if err := client.Sync(context.Background(), SyncOptions{ProviderDirs: map[Provider][]string{ + ProviderClaude: {filepath.Join(dir, "claude")}, + }}); err != nil { + t.Fatal(err) + } + + count := 0 + uploaded.Range(func(any, any) bool { count++; return true }) + if count != files { + t.Fatalf("uploaded %d of %d events before Sync returned", count, files) + } +} + +// A scan that produces a steady trickle must not turn into a request per +// handful of events. The drain waits for a worthwhile batch and sends one +// batch per pass, so the queue refills between requests instead of being +// chased down to nothing. +func TestSyncBatchesTrickleIntoFewRequests(t *testing.T) { + var mu sync.Mutex + var sizes []int + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body := io.Reader(r.Body) + if r.Header.Get("Content-Encoding") == "gzip" { + gz, err := gzip.NewReader(r.Body) + if err != nil { + w.WriteHeader(500) + return + } + defer gz.Close() + body = gz + } + var payload struct { + Events []struct { + ID string `json:"id"` + } `json:"events"` + } + json.NewDecoder(body).Decode(&payload) + ids := make([]string, 0, len(payload.Events)) + for _, e := range payload.Events { + ids = append(ids, e.ID) + } + mu.Lock() + sizes = append(sizes, len(payload.Events)) + mu.Unlock() + json.NewEncoder(w).Encode(map[string]any{"ok": true, "accepted": ids}) + })) + defer srv.Close() + t.Setenv("TOKITOKI_BASE_URL", srv.URL) + + // Many files holding one event each: a slow scan producing a trickle. + const files = 3000 + dir := t.TempDir() + for i := 0; i < files; i++ { + p := filepath.Join(dir, "claude", "projects", fmt.Sprintf("-p-%04d", i)) + if err := os.MkdirAll(p, 0o700); err != nil { + t.Fatal(err) + } + line := fmt.Sprintf(`{"timestamp":"2026-06-04T01:02:03Z","cwd":"/p/%04d","requestId":"r%04d","message":{"id":"m%04d","model":"c","usage":{"input_tokens":1,"output_tokens":1}}}`+"\n", i, i, i) + if err := os.WriteFile(filepath.Join(p, "s.jsonl"), []byte(line), 0o600); err != nil { + t.Fatal(err) + } + } + if err := os.MkdirAll(filepath.Join(dir, "config"), 0o700); err != nil { + t.Fatal(err) + } + if err := os.WriteFile(filepath.Join(dir, "config", "api_key"), []byte("k"), 0o600); err != nil { + t.Fatal(err) + } + + client, err := New(Options{DataDir: dir}) + if err != nil { + t.Fatal(err) + } + if err := client.Sync(context.Background(), SyncOptions{ProviderDirs: map[Provider][]string{ + ProviderClaude: {filepath.Join(dir, "claude")}, + }}); err != nil { + t.Fatal(err) + } + + mu.Lock() + defer mu.Unlock() + total := 0 + for _, s := range sizes { + total += s + } + t.Logf("%d events in %d requests", total, len(sizes)) + if total != files { + t.Fatalf("uploaded %d of %d events", total, files) + } + // Chasing the scan produced ~110 requests for this input; batching keeps + // it near the number of full batches the queue can actually form. + if len(sizes) > 10 { + t.Errorf("%d requests for %d events: the drain is chasing the scan", len(sizes), total) + } +} From e2113b3e2853eee078d1e48bba384d10e5ddaba6 Mon Sep 17 00:00:00 2001 From: akarineren Date: Wed, 5 Aug 2026 11:11:04 +0900 Subject: [PATCH 10/11] feat: add TestCompareWithCcusage for usage tracking validation --- internal/provider/claude/zz_cmp_test.go | 40 +++++++++++++++++++++++++ 1 file changed, 40 insertions(+) create mode 100644 internal/provider/claude/zz_cmp_test.go diff --git a/internal/provider/claude/zz_cmp_test.go b/internal/provider/claude/zz_cmp_test.go new file mode 100644 index 0000000..7525be7 --- /dev/null +++ b/internal/provider/claude/zz_cmp_test.go @@ -0,0 +1,40 @@ +package claude + +import ( + "fmt" + "os" + "path/filepath" + "testing" + "time" +) + +func TestCompareWithCcusage(t *testing.T) { + day := os.Getenv("DAY") + if day == "" { t.Skip() } + root := os.Getenv("SNAP2") + if root == "" { + home, _ := os.UserHomeDir() + root = filepath.Join(home, ".claude") + } + files := UsageFiles([]string{root}, "") + + seen := map[string]bool{} + var in, out, cc, cr uint64 + n := 0 + for _, f := range files { + loaded, _, err := ReadUsageFileFrom(f, 0) + if err != nil { t.Fatal(err) } + for _, e := range ConvertEntries(loaded) { + if e.Timestamp.In(time.Local).Format("2006-01-02") != day { continue } + if seen[e.ID] { continue } // 跨文件去重, 同 ccusage + seen[e.ID] = true + n++ + in += e.Usage.InputTokens + out += e.Usage.OutputTokens + cc += e.Usage.CacheCreationInputTokens + cr += e.Usage.CacheReadInputTokens + } + } + fmt.Printf("TOKITOKI %s: events=%d in=%d out=%d cc=%d cr=%d total=%d\n", + day, n, in, out, cc, cr, in+out+cc+cr) +} From 032bbc24968ce8ba79f2125e9b68b741b7d6d125 Mon Sep 17 00:00:00 2001 From: akarineren Date: Wed, 5 Aug 2026 11:24:24 +0900 Subject: [PATCH 11/11] docs: correct release publishing rules in RELEASING.md The /admin/releases import-and-publish step no longer exists; the update server reads GitHub Releases directly, so pushing a tag ships to clients immediately. Also document the tag pattern, asset-name contract, and ldflags version stamping. --- RELEASING.md | 60 ++++++++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 56 insertions(+), 4 deletions(-) diff --git a/RELEASING.md b/RELEASING.md index 06a6ccd..b928004 100644 --- a/RELEASING.md +++ b/RELEASING.md @@ -20,10 +20,18 @@ Release tags must be semantic versions on an up-to-date `main` branch: ```sh git switch main git pull --ff-only -git tag v0.1.1 -git push origin v0.1.1 +git tag vX.Y.Z +git push origin vX.Y.Z ``` +The tag pattern is exactly `v[0-9]+.[0-9]+.[0-9]+`. Suffixed tags — `v1.2.0-rc.1`, +`v1.2` — do not match the workflow trigger, so they build nothing and ship +nothing. There is no prerelease channel. + +Never tag `dev`. The workflow's first step rejects a tag whose commit is not an +ancestor of `origin/main`, so tagging `dev` fails the build rather than shipping +untested work — but it leaves a junk tag to clean up. + The tag starts the `Release` workflow. It rejects tags outside `main`, reruns vet and race-enabled tests, then cross-compiles stripped, reproducible binaries for every supported target: @@ -39,5 +47,49 @@ The workflow verifies the asset set and embedded version, generates remain raw rather than being wrapped in ZIP or tar archives: the Tokitoki update server proxies them directly to `tokitoki update`. -Creating the GitHub Release does not publish it to clients. Import and publish -the version from `/admin/releases`; until then the update API ignores it. +The asset names are an API. The server matches the platform and arch tokens in +each filename to answer a download request, so renaming an asset or dropping one +from the matrix breaks clients on that platform. The workflow guards this: it +requires all six files to exist and the count to be exactly six. + +Nothing in the repository records the version. `make cross VERSION=X.Y.Z` stamps +it through `-ldflags` into `internal/buildinfo.Version`, and the workflow derives +that value from the tag — so the tag is the only source of truth, and there is no +version file to bump in a commit. The workflow then runs the freshly built +binary and fails the release if `tokitoki version` disagrees with the tag. +Unstamped builds report `dev` and refuse to self-update, which is what keeps a +local build from overwriting itself with a release. + +## Pushing the tag ships it + +Creating the GitHub Release **is** publishing it. There is no second gate. + +The update server answers `/api/updates/check` straight from the GitHub +Releases API (`lib/releases.ts` in `tracklm-nextjs`): no database mirror, no +publish switch, only a short-TTL in-memory cache that serves stale data when +GitHub is unreachable. The newest non-draft tag that parses as stable semver +becomes the answer for every client asking what to install. `/admin/releases` +reports downloads and is deliberately read-only — nothing on that page ships a +version. + +So `git push origin vX.Y.Z` is the point of no return. Once the workflow +finishes, `tokitoki update` starts handing that binary to every client, and the +macOS/Windows apps and editor plugins follow, because they all delegate to the +same `tokitoki update`. Verify before pushing the tag, not after. + +Backing out means acting before clients poll, and there is no way to recall +what has already been downloaded: + +```sh +gh release delete vX.Y.Z --repo tokitoki-dev/tokitoki-cli --yes +git push origin :refs/tags/vX.Y.Z +git tag -d vX.Y.Z +``` + +Prefer rolling forward with a new patch version. Deleting a release that +clients have already seen means some of them sit on a version the server no +longer offers. + +An earlier revision of this document described importing and publishing a +version from `/admin/releases`. That step no longer exists; the server was +changed to read GitHub directly.