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/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. diff --git a/cmd/tokitoki/main.go b/cmd/tokitoki/main.go index 4479d4c..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 } @@ -565,7 +648,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/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/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/agentusage/helpers.go b/internal/agentdata/agentdata.go similarity index 52% rename from internal/agentusage/helpers.go rename to internal/agentdata/agentdata.go index 9a26e53..eff0aba 100644 --- a/internal/agentusage/helpers.go +++ b/internal/agentdata/agentdata.go @@ -1,4 +1,6 @@ -package agentusage +// 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" @@ -20,7 +21,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 +56,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,64 +70,103 @@ 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) { + 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() - lines := make([]lineJSON, 0) + 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 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: 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 { - 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 +178,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 +195,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 +216,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,17 +241,17 @@ 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 { 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 +267,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)) { @@ -243,48 +283,33 @@ 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 } } return 0 } -func floatValue(value any) (float64, bool) { +func ParseTimestamp(value any) (time.Time, 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: - 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,13 +321,13 @@ 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) { - if !isFinite(value) || value <= 0 { +func TimestampFromFloat(value float64) (time.Time, bool) { + if !IsFinite(value) || value <= 0 { return time.Time{}, false } if value < 100_000_000_000 { @@ -311,7 +336,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) } @@ -336,13 +361,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 +378,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() @@ -365,96 +390,7 @@ 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 { +func UniqueStrings(values []string) []string { if len(values) == 0 { return nil } @@ -470,3 +406,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..3a023f8 --- /dev/null +++ b/internal/agentdb/agentdb.go @@ -0,0 +1,160 @@ +// 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 + } +} + +// 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", "file:"+filepath.ToSlash(path)+"?mode=ro") + 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/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/copilot.go b/internal/agentusage/copilot.go deleted file mode 100644 index e2aaefe..0000000 --- a/internal/agentusage/copilot.go +++ /dev/null @@ -1,419 +0,0 @@ -package agentusage - -import ( - "fmt" - "path/filepath" - "sort" - "strings" - "time" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func loadCopilotEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { - files := make([]string, 0) - for _, root := range paths { - files = append(files, collectExt(root, ".jsonl")...) - } - sort.Strings(files) - files = filterFiles(uniqueStrings(files), filter) - - entries := make([]usage.Entry, 0) - for _, file := range files { - fileEntries, err := parseCopilotOTELFile(file) - if err != nil { - return nil, err - } - entries = append(entries, fileEntries...) - } - sortEntries(entries) - return entries, nil -} - -type copilotSource int - -const ( - copilotChatSpan copilotSource = iota - copilotInferenceLog - copilotAgentTurnLog - copilotAgentSummarySpan -) - -type copilotCandidate struct { - source copilotSource - traceID string - responseID string - sessionID string - model string - timestamp time.Time - tokens usage.TokenUsage - dedupKey string - sourceFile string - sourceLine int - sourceStart, sourceEnd int64 -} - -type copilotTraceContext struct { - model string - sessionID string - sessionIDPriority int -} - -func parseCopilotOTELFile(path string) ([]usage.Entry, error) { - lines, err := readJSONLines(path, `"attributes"`) - if err != nil { - return nil, err - } - contexts := copilotTraceContexts(lines) - fallback := fileModifiedTime(path) - candidates := make([]copilotCandidate, 0) - for index, line := range lines { - if candidate, ok := copilotRecordCandidate(path, line, index, fallback, contexts); ok { - candidates = append(candidates, candidate) - } - } - sets := copilotCandidateSets(candidates) - entries := make([]usage.Entry, 0) - for _, candidate := range candidates { - if !shouldEmitCopilot(candidate, sets) { - continue - } - entry := baseEntry( - usage.ProviderCopilot, - candidate.timestamp, - "copilot", - "GitHub Copilot CLI", - candidate.sessionID, - candidate.model, - "GitHub Copilot CLI", - candidate.tokens, - ) - setSource(&entry, candidate.sourceFile, candidate.sourceLine, candidate.sourceStart, candidate.sourceEnd) - entry.ID = stableEntryID(entry, candidate.dedupKey) - entries = append(entries, entry) - } - return entries, nil -} - -func copilotTraceContexts(lines []lineJSON) map[string]copilotTraceContext { - contexts := make(map[string]copilotTraceContext) - for _, line := range lines { - record := line.value - traceID := copilotTraceID(record) - if traceID == "" { - continue - } - attrs := 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") - } - if sessionID, priority := copilotBestSession(attrs); sessionID != "" && priority > context.sessionIDPriority { - context.sessionID = sessionID - context.sessionIDPriority = priority - } - contexts[traceID] = context - } - 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"]) - if attrs == nil { - return copilotCandidate{}, false - } - source, ok := copilotRecordSource(record, attrs) - if !ok { - return copilotCandidate{}, false - } - input := uintField(attrs, "gen_ai.usage.input_tokens") - cacheRead := uintField(attrs, "gen_ai.usage.cache_read.input_tokens") - if cacheRead <= input { - input -= cacheRead - } else { - input = 0 - } - 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"), - CacheReadInputTokens: cacheRead, - ReasoningOutputTokens: 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 - } - traceID := copilotTraceID(record) - context := contexts[traceID] - model := firstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") - if model == "" { - model = context.model - } - if model == "" { - model = "unknown" - } - sessionID, _ := copilotBestSession(attrs) - if sessionID == "" { - sessionID = context.sessionID - } - if sessionID == "" { - sessionID = traceID - } - if sessionID == "" { - sessionID = "unknown-session" - } - timestamp, ok := copilotTimestamp(record) - if !ok { - timestamp = fallback - } - responseID := stringField(attrs, "gen_ai.response.id") - return copilotCandidate{ - source: source, - traceID: traceID, - responseID: responseID, - sessionID: sessionID, - model: model, - timestamp: timestamp, - tokens: tokens, - dedupKey: copilotDedupKey(source, record, attrs, traceID, sessionID, timestamp, index), - sourceFile: path, - sourceLine: line.line, - sourceStart: line.start, - sourceEnd: line.end, - }, true -} - -func copilotRecordSource(record, attrs map[string]any) (copilotSource, bool) { - switch { - case copilotIsChatSpan(record, attrs): - return copilotChatSpan, true - case copilotIsInferenceLog(record, attrs): - return copilotInferenceLog, true - case copilotIsAgentTurnLog(record, attrs): - return copilotAgentTurnLog, true - case copilotIsAgentSummarySpan(record, attrs): - return copilotAgentSummarySpan, true - default: - return 0, false - } -} - -func copilotIsSpan(record map[string]any) bool { - if stringField(record, "type") == "span" { - return true - } - if stringField(record, "name") == "" { - return false - } - return stringField(record, "spanId") != "" || - 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 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 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 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 copilotBody(record map[string]any) string { - if body := stringField(record, "body"); body != "" { - return body - } - return stringField(record, "_body") -} - -func copilotTraceID(record map[string]any) string { - if traceID := stringField(record, "traceId"); traceID != "" { - return traceID - } - return stringField(objectAt(record["spanContext"]), "traceId") -} - -func copilotSpanID(record map[string]any) string { - if spanID := stringField(record, "spanId"); spanID != "" { - return spanID - } - return stringField(objectAt(record["spanContext"]), "spanId") -} - -func copilotBestSession(attrs map[string]any) (string, int) { - candidates := []struct { - key string - priority int - }{ - {"gen_ai.conversation.id", 3}, - {"copilot_chat.session_id", 3}, - {"copilot_chat.chat_session_id", 3}, - {"session.id", 3}, - {"github.copilot.interaction_id", 2}, - {"gen_ai.response.id", 1}, - } - bestValue := "" - bestPriority := 0 - for _, candidate := range candidates { - if value := stringField(attrs, candidate.key); value != "" && candidate.priority > bestPriority { - bestValue = value - bestPriority = candidate.priority - } - } - return bestValue, bestPriority -} - -func copilotTimestamp(record map[string]any) (time.Time, bool) { - for _, key := range []string{"endTime", "startTime", "hrTime", "_hrTime", "time"} { - if timestamp, ok := timestampFromParts(record[key]); ok { - return timestamp, true - } - } - for _, key := range []string{"timestamp", "observedTimestamp"} { - if timestamp, ok := parseTimestamp(record[key]); ok { - return timestamp, true - } - } - if raw := 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) - switch source { - case copilotChatSpan, copilotAgentSummarySpan: - if traceID != "" && spanID != "" { - return traceID + ":" + spanID - } - return fmt.Sprintf("span:%s:%d:%d", sessionID, timestamp.UnixMilli(), index) - case copilotInferenceLog: - if traceID != "" && spanID != "" { - return "log:" + traceID + ":" + spanID - } - return fmt.Sprintf("log:%s:%d:%d", sessionID, timestamp.UnixMilli(), index) - case copilotAgentTurnLog: - turnIndex := uintField(attrs, "turn.index", "copilot_chat.turn.index") - turn := fmt.Sprintf("idx-%d", index) - if turnIndex > 0 { - turn = fmt.Sprintf("%d", turnIndex) - } - if traceID != "" { - return "agent-turn:" + traceID + ":" + turn - } - return "agent-turn:" + sessionID + ":" + turn + fmt.Sprintf(":%d", index) - default: - return fmt.Sprintf("%s:%d", filepath.Base(copilotSourceName(source)), index) - } -} - -func copilotSourceName(source copilotSource) string { - switch source { - case copilotChatSpan: - return "chat" - case copilotInferenceLog: - return "inference" - case copilotAgentTurnLog: - return "agent-turn" - case copilotAgentSummarySpan: - return "agent-summary" - default: - return "unknown" - } -} - -type copilotSets struct { - chatTraces map[string]bool - inferenceTraces map[string]bool - agentTurnTraces map[string]bool - chatResponses map[string]bool - inferenceResponses map[string]bool - agentTurnResponses map[string]bool -} - -func copilotCandidateSets(candidates []copilotCandidate) copilotSets { - sets := copilotSets{ - chatTraces: make(map[string]bool), - inferenceTraces: make(map[string]bool), - agentTurnTraces: make(map[string]bool), - chatResponses: make(map[string]bool), - inferenceResponses: make(map[string]bool), - agentTurnResponses: make(map[string]bool), - } - for _, candidate := range candidates { - if candidate.traceID != "" { - switch candidate.source { - case copilotChatSpan: - sets.chatTraces[candidate.traceID] = true - case copilotInferenceLog: - sets.inferenceTraces[candidate.traceID] = true - case copilotAgentTurnLog: - sets.agentTurnTraces[candidate.traceID] = true - } - } - if candidate.responseID != "" { - switch candidate.source { - case copilotChatSpan: - sets.chatResponses[candidate.responseID] = true - case copilotInferenceLog: - sets.inferenceResponses[candidate.responseID] = true - case copilotAgentTurnLog: - sets.agentTurnResponses[candidate.responseID] = true - } - } - } - return sets -} - -func shouldEmitCopilot(candidate copilotCandidate, sets copilotSets) bool { - traceMatch := func(values map[string]bool) bool { - return candidate.traceID != "" && values[candidate.traceID] - } - responseMatch := func(values map[string]bool) bool { - return candidate.responseID != "" && values[candidate.responseID] - } - switch candidate.source { - case copilotChatSpan: - return true - case copilotInferenceLog: - return !traceMatch(sets.chatTraces) && !responseMatch(sets.chatResponses) - case copilotAgentTurnLog: - return !traceMatch(sets.chatTraces) && - !traceMatch(sets.inferenceTraces) && - !responseMatch(sets.chatResponses) && - !responseMatch(sets.inferenceResponses) - case copilotAgentSummarySpan: - return !traceMatch(sets.chatTraces) && - !traceMatch(sets.inferenceTraces) && - !traceMatch(sets.agentTurnTraces) && - !responseMatch(sets.chatResponses) && - !responseMatch(sets.inferenceResponses) && - !responseMatch(sets.agentTurnResponses) - default: - return false - } -} diff --git a/internal/agentusage/gemini.go b/internal/agentusage/gemini.go deleted file mode 100644 index 8fa999a..0000000 --- a/internal/agentusage/gemini.go +++ /dev/null @@ -1,250 +0,0 @@ -package agentusage - -import ( - "path/filepath" - "sort" - "strings" - "time" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func loadGeminiEntries(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")...) - } - sort.Strings(files) - files = filterFiles(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) - } else { - fileEntries, err = parseGeminiJSONFile(file) - } - if err != nil { - return nil, err - } - entries = append(entries, fileEntries...) - } - sortEntries(entries) - return entries, nil -} - -type geminiTokens struct { - input uint64 - output uint64 - cached uint64 - thoughts uint64 - tool uint64 - total uint64 - hasTotal bool -} - -func parseGeminiJSONFile(path string) ([]usage.Entry, error) { - record, err := readJSONObject(path) - if err != nil || record == nil { - return nil, err - } - fallback := fileModifiedTime(path) - sessionID := 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 { - entries := make([]usage.Entry, 0) - for index, raw := range messages { - message := objectAt(raw) - if stringField(message, "type") != "gemini" { - continue - } - if entry, ok := geminiDirectEntry(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 { - return []usage.Entry{entry}, nil - } - return nil, nil - } - return geminiStatsEntries(recordStats(record), path, 1, stringField(record, "model"), sessionID, sessionTimestamp), nil -} - -func parseGeminiJSONLFile(path string) ([]usage.Entry, error) { - lines, err := readJSONLines(path) - if err != nil { - return nil, err - } - fallback := 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 != "" { - sessionID = value - } - if model := stringField(record, "model"); model != "" { - currentModel = model - } - if stringField(record, "type") == "gemini" { - entry, ok := geminiDirectEntry(record, path, line.line, currentModel, sessionID, fallback) - if !ok { - continue - } - messageID := stringField(record, "id") - if messageID != "" { - if index, exists := indexByMessageID[messageID]; exists { - entries[index] = entry - continue - } - indexByMessageID[messageID] = len(entries) - } - entries = append(entries, entry) - continue - } - if stats := recordStats(record); stats != nil { - timestamp := firstTimestamp(record, fallback, "timestamp") - entries = append(entries, geminiStatsEntries(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"]) - if !ok { - return usage.Entry{}, false - } - model := 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")) -} - -func geminiStatsEntries(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 { - entries := make([]usage.Entry, 0) - for model, raw := range models { - data := objectAt(raw) - tokens, ok := parseGeminiTokens(data["tokens"]) - if !ok { - continue - } - if entry, ok := buildGeminiEntry(path, line, model, sessionID, timestamp, tokens, false, ""); ok { - entries = append(entries, entry) - } - } - if len(entries) > 0 { - return entries - } - } - tokens, ok := parseGeminiTokens(stats) - if !ok { - return nil - } - model := modelHint - if model == "" { - model = "unknown" - } - entry, ok := buildGeminiEntry(path, line, model, 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 geminiTokens, direct bool, messageID string) (usage.Entry, bool) { - model = strings.TrimSpace(model) - if model == "" { - return usage.Entry{}, false - } - input, cacheRead := normalizeGeminiInput(tokens, direct) - tokenUsage := usage.TokenUsage{ - InputTokens: input + tokens.tool, - OutputTokens: tokens.output, - CacheReadInputTokens: cacheRead, - ReasoningOutputTokens: tokens.thoughts, - } - if tokens.hasTotal { - tokenUsage = applyTotalFallback(tokenUsage, tokens.total) - } else if tokenUsage.TotalTokens == 0 { - tokenUsage.TotalTokens = totalUsage(tokenUsage) - } - if !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) - return entry, true -} - -func parseGeminiTokens(raw any) (geminiTokens, bool) { - record := objectAt(raw) - if record == nil { - return geminiTokens{}, 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.hasTotal = tokens.total > 0 - return tokens, true -} - -func normalizeGeminiInput(tokens geminiTokens, direct bool) (uint64, uint64) { - if !direct { - cachedPortion := tokens.input - if tokens.cached < cachedPortion { - cachedPortion = tokens.cached - } - return tokens.input - cachedPortion, tokens.cached - } - inclusiveTotal := tokens.input + tokens.output + tokens.thoughts + tokens.tool - exclusiveTotal := inclusiveTotal + tokens.cached - if tokens.cached > 0 && tokens.hasTotal && tokens.total == inclusiveTotal && tokens.total != exclusiveTotal { - cachedPortion := tokens.input - if tokens.cached < cachedPortion { - cachedPortion = tokens.cached - } - return tokens.input - cachedPortion, tokens.cached - } - return tokens.input, tokens.cached -} - -func recordStats(record map[string]any) map[string]any { - if stats := objectAt(record["stats"]); stats != nil { - return stats - } - result := objectAt(record["result"]) - return 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 { - return timestamp - } - } - return fallback -} 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/opencode.go b/internal/agentusage/opencode.go deleted file mode 100644 index a0c2b4b..0000000 --- a/internal/agentusage/opencode.go +++ /dev/null @@ -1,225 +0,0 @@ -package agentusage - -import ( - "bytes" - "encoding/json" - "os" - "path/filepath" - "sort" - "strings" - "time" - - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -func loadOpenCodeEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { - entriesByID := make(map[string]usage.Entry) - for _, root := range paths { - rootEntries, err := loadOpenCodeRoot(root, filter) - if err != nil { - return nil, err - } - for _, entry := range rootEntries { - if entry.ID == "" { - entry.ID = stableEntryID(entry) - } - 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) - } - sortEntries(entries) - return entries, nil -} - -func loadOpenCodeRoot(root string, filter usage.FileFilter) ([]usage.Entry, error) { - info, err := os.Stat(root) - if err != nil { - return nil, nil - } - if !info.IsDir() { - switch strings.ToLower(filepath.Ext(root)) { - 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 nil, nil - } - - entries := make([]usage.Entry, 0) - seenIDs := make(map[string]bool) - if dbPath := openCodeDBPath(root); dbPath != "" { - dbEntries, err := loadOpenCodeDatabase(dbPath) - if err != nil { - return nil, err - } - for _, entry := range dbEntries { - if entry.ID != "" { - seenIDs[entry.ID] = true - } - entries = append(entries, entry) - } - } - - files := collectExt(filepath.Join(root, "storage", "message"), ".json") - sort.Strings(files) - files = filterFiles(files, filter) - for _, file := range files { - stem := strings.TrimSuffix(filepath.Base(file), filepath.Ext(file)) - if seenIDs[stableOpenCodeMessageID(stem)] { - continue - } - entry, ok, err := parseOpenCodeMessageFile(file, "", "") - if err != nil { - return nil, err - } - if !ok { - continue - } - if entry.ID != "" && seenIDs[entry.ID] { - continue - } - if entry.ID != "" { - seenIDs[entry.ID] = true - } - 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) - if err != nil { - return "" - } - candidates := make([]string, 0) - for _, entry := range entries { - name := entry.Name() - if entry.IsDir() || !isOpenCodeChannelDB(name) { - continue - } - candidates = append(candidates, filepath.Join(root, name)) - } - sort.Strings(candidates) - if len(candidates) == 0 { - return "" - } - return candidates[0] -} - -func isOpenCodeChannelDB(name string) bool { - if !strings.HasPrefix(name, "opencode-") || !strings.HasSuffix(name, ".db") { - return false - } - channel := strings.TrimSuffix(strings.TrimPrefix(name, "opencode-"), ".db") - if channel == "" { - return false - } - for _, ch := range channel { - if !(ch >= 'a' && ch <= 'z' || ch >= 'A' && ch <= 'Z' || ch >= '0' && ch <= '9' || ch == '_' || ch == '-') { - return false - } - } - return true -} - -func loadOpenCodeDatabase(path string) ([]usage.Entry, error) { - db, err := openSQLite(path) - 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 := openCodeMessageEntry(path, decodeJSONObjectString(data), rowID, rowSessionID, 0); ok { - entries = append(entries, entry) - } - } - return entries, rows.Err() -} - -func parseOpenCodeMessageFile(path, rowID, rowSessionID string) (usage.Entry, bool, error) { - data, err := os.ReadFile(path) - if err != nil { - return usage.Entry{}, false, err - } - 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 - } - 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 - } - 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"), - } - tokens = applyTotalFallback(tokens, uintField(tokenBlock, "total")) - if !nonZero(tokens) { - return usage.Entry{}, false - } - model := stringField(record, "modelID") - if model == "" { - return usage.Entry{}, false - } - timestamp := time.Unix(0, 0).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) - if entry.ID == "" { - entry.ID = stableEntryID(entry) - } - return entry, true -} - -func stableOpenCodeMessageID(messageID string) string { - if strings.TrimSpace(messageID) == "" { - return "" - } - return usage.StableID("opencode", strings.TrimSpace(messageID)) -} diff --git a/internal/agentusage/pi.go b/internal/agentusage/pi.go deleted file mode 100644 index 304eeac..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 "unknown" -} 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 d906e79..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: "unknown", - 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","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", - 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 36c2a40..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 "unknown" -} diff --git a/internal/agentusage/sqlite.go b/internal/agentusage/sqlite.go deleted file mode 100644 index 96d1871..0000000 --- a/internal/agentusage/sqlite.go +++ /dev/null @@ -1,84 +0,0 @@ -package agentusage - -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/claudeusage/loader_test.go b/internal/claudeusage/loader_test.go deleted file mode 100644 index fa1387b..0000000 --- a/internal/claudeusage/loader_test.go +++ /dev/null @@ -1,274 +0,0 @@ -package claudeusage - -import ( - "os" - "path/filepath" - "testing" - "time" -) - -func TestUsageFilesLimitsDiscoveryToProjectFilter(t *testing.T) { - dir := t.TempDir() - projectA := filepath.Join(dir, "projects", "project-a", "session-a") - projectB := filepath.Join(dir, "projects", "project-b", "session-b") - mkdirAll(t, projectA) - mkdirAll(t, projectB) - writeFile(t, filepath.Join(projectA, "a.jsonl"), "{}") - writeFile(t, filepath.Join(projectB, "b.jsonl"), "{}") - - files := UsageFiles([]string{dir}, "project-a") - - if len(files) != 1 { - t.Fatalf("len(files) = %d, want 1", len(files)) - } - if got := files[0]; !containsPathSegment(got, "project-a") { - t.Fatalf("file = %q, want project-a path", got) - } -} - -func TestUsageFilesFallsBackForNonSegmentProjectFilter(t *testing.T) { - dir := t.TempDir() - projectA := filepath.Join(dir, "projects", "project-a", "session-a") - projectB := filepath.Join(dir, "projects", "project-b", "session-b") - mkdirAll(t, projectA) - mkdirAll(t, projectB) - writeFile(t, filepath.Join(projectA, "a.jsonl"), "{}") - writeFile(t, filepath.Join(projectB, "b.jsonl"), "{}") - - files := UsageFiles([]string{dir}, "project-a/session-a") - - if len(files) != 2 { - t.Fatalf("len(files) = %d, want 2", len(files)) - } -} - -func TestProjectPathSegmentRejectsUnsafeValues(t *testing.T) { - cases := map[string]bool{ - "": false, - ".": false, - "..": false, - "project-a/session-a": false, - `project-a\session-a`: false, - "project-a": true, - } - - for value, want := range cases { - if got := isProjectPathSegment(value); got != want { - t.Fatalf("isProjectPathSegment(%q) = %v, want %v", value, got, want) - } - } -} - -func TestExtractSessionParts(t *testing.T) { - tests := []struct { - name string - path string - wantSessionID string - wantProjectPath string - }{ - { - name: "modern", - path: "/home/me/.claude/projects/project-a/session-a.jsonl", - wantSessionID: "session-a", - wantProjectPath: "project-a", - }, - { - name: "nested", - path: "/home/me/.claude/projects/project-a/session-a/chat.jsonl", - wantSessionID: "session-a", - wantProjectPath: "project-a", - }, - { - name: "subagent", - path: "/home/me/.claude/projects/project-a/session-a/subagents/worker.jsonl", - wantSessionID: "session-a", - wantProjectPath: "project-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"), - }, - } - - for _, tt := range tests { - t.Run(tt.name, func(t *testing.T) { - sessionID, projectPath := ExtractSessionParts(tt.path) - if 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" - - if got := ExtractProject(path); got != "relink" { - t.Fatalf("ExtractProject() = %q, want relink", got) - } -} - -func TestHasUnsupportedNullField(t *testing.T) { - rejected := [][]byte{ - []byte(`{"message":{"usage":{"speed":null}}}`), - []byte(`{"message":{"model":null,"usage":{"input_tokens":0}}}`), - []byte(`{"sessionId":null,"message":{"usage":{"input_tokens":0}}}`), - } - for _, line := range rejected { - if !hasUnsupportedNullField(line) { - t.Fatalf("hasUnsupportedNullField(%s) = false, want true", line) - } - } - - allowed := []byte(`{"message":{"content":null,"usage":{"input_tokens":0}}}`) - if hasUnsupportedNullField(allowed) { - t.Fatalf("hasUnsupportedNullField(%s) = true, want false", allowed) - } -} - -func TestReadUsageFileParsesUsageLines(t *testing.T) { - path := filepath.Join(t.TempDir(), "projects", "project-a", "session-a", "chat.jsonl") - 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"}}} -`) - - entries, err := ReadUsageFile(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.Project != "project-a" { - t.Fatalf("project = %q, want project-a", entry.Project) - } - if entry.SessionID != "session-a" { - t.Fatalf("sessionID = %q, want session-a", entry.SessionID) - } - if entry.Model != "claude-sonnet-4-20250514-fast" { - t.Fatalf("model = %q, want fast suffix", entry.Model) - } - if entry.Language != "Unknown" { - t.Fatalf("language = %q, want Unknown", entry.Language) - } - if entry.Date != "2026-05-21" { - t.Fatalf("date = %q, want 2026-05-21", entry.Date) - } - if got := tokenTotal(entry.Data.Message.Usage); got != 20 { - t.Fatalf("tokenTotal = %d, want 20", got) - } -} - -func TestReadUsageFileInfersLanguageFromToolUseFilePath(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","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1},"content":[{"type":"tool_use","name":"Read","input":{"file_path":"/repo/internal/server/server.go"}}]}}`) - - 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].Language != "Go" { - t.Fatalf("language = %q, want Go", entries[0].Language) - } -} - -func TestReadUsageFileDoesNotInferLanguageFromCodeFenceWithoutFilePath(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\",\"message\":{\"id\":\"msg-1\",\"model\":\"claude\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1},\"content\":[{\"type\":\"text\",\"text\":\"```tsx\\nexport default function Page() {}\\n```\"}]}}\n") - - 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].Language != "Unknown" { - t.Fatalf("language = %q, want Unknown", entries[0].Language) - } -} - -func TestReadUsageFileSkipsUnsupportedSpeed(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","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1,"speed":"turbo"}}}`) - - entries, err := ReadUsageFile(path) - if err != nil { - t.Fatal(err) - } - if len(entries) != 0 { - t.Fatalf("len(entries) = %d, want 0", len(entries)) - } -} - -func TestLoadEntriesDeduplicatesByMessageAndRequest(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "projects", "project-a", "session-a.jsonl") - mkdirAll(t, filepath.Dir(path)) - writeFile(t, path, ` -{"timestamp":"2026-05-21T01:02:03Z","requestId":"req-1","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}} -{"timestamp":"2026-05-21T01:02:04Z","requestId":"req-1","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":10,"output_tokens":1}}} -`) - - entries, err := LoadEntriesFromPaths([]string{dir}, "", nil) - if err != nil { - t.Fatal(err) - } - if len(entries) != 1 { - t.Fatalf("len(entries) = %d, want 1", len(entries)) - } - if got := entries[0].Data.Message.Usage.InputTokens; got != 10 { - t.Fatalf("input tokens = %d, want replacement with larger usage", got) - } -} - -func TestUsageLimitResetTimeFromLine(t *testing.T) { - isAPIError := true - line := []byte(`{"timestamp":"2026-05-21T01:02:03Z","isApiErrorMessage":true,"message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1},"content":"Claude AI usage limit reached|1779325200"}}`) - - reset := usageLimitResetTimeFromLine(line, &isAPIError) - if reset == nil { - t.Fatal("reset = nil, want timestamp") - } - if want := time.Unix(1779325200, 0).UTC(); !reset.Equal(want) { - t.Fatalf("reset = %s, want %s", reset, want) - } -} - -func mkdirAll(t *testing.T, path string) { - t.Helper() - if err := os.MkdirAll(path, 0o700); err != nil { - t.Fatal(err) - } -} - -func writeFile(t *testing.T, path, data string) { - t.Helper() - if err := os.WriteFile(path, []byte(data), 0o600); err != nil { - t.Fatal(err) - } -} - -func containsPathSegment(path, segment string) bool { - for _, part := range pathParts(path) { - if part == segment { - return true - } - } - return false -} 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/codexusage/loader.go b/internal/codexusage/loader.go deleted file mode 100644 index 9b90bd7..0000000 --- a/internal/codexusage/loader.go +++ /dev/null @@ -1,402 +0,0 @@ -package codexusage - -import ( - "bufio" - "bytes" - "encoding/json" - "errors" - "io" - "os" - "path/filepath" - "runtime" - "sort" - "strconv" - "strings" - "time" - - "github.com/tokitoki-dev/tokitoki-cli/internal/langdetect" - "github.com/tokitoki-dev/tokitoki-cli/internal/usage" -) - -var ErrNoDataDirs = errors.New("no valid Codex data directories found") - -type codexLine struct { - Timestamp string `json:"timestamp"` - Type string `json:"type"` - Payload json.RawMessage `json:"payload"` -} - -type sessionMetaPayload struct { - ID string `json:"id"` - CWD string `json:"cwd"` - Originator string `json:"originator"` -} - -type turnContextPayload struct { - CWD string `json:"cwd"` - Model string `json:"model"` -} - -type eventPayload struct { - Type string `json:"type"` - Info struct { - LastTokenUsage *tokenUsagePayload `json:"last_token_usage"` - } `json:"info"` -} - -type tokenUsagePayload struct { - InputTokens uint64 `json:"input_tokens"` - CachedInputTokens uint64 `json:"cached_input_tokens"` - OutputTokens uint64 `json:"output_tokens"` - ReasoningOutputTokens uint64 `json:"reasoning_output_tokens"` - TotalTokens uint64 `json:"total_tokens"` -} - -func LoadEntriesFromPaths(paths []string, projectFilter string, fileFilter usage.FileFilter) ([]usage.Entry, error) { - files := UsageFiles(paths) - entries := make([]usage.Entry, 0) - for _, file := range files { - if fileFilter != nil && !fileFilter(file) { - continue - } - fileEntries, err := ReadUsageFile(file) - if err != nil { - return nil, err - } - for _, entry := range fileEntries { - if projectFilter != "" && entry.Project != projectFilter && entry.ProjectPath != projectFilter { - continue - } - entries = append(entries, entry) - } - } - return entries, nil -} - -func UsageFiles(paths []string) []string { - files := make([]string, 0) - for _, path := range paths { - collectJSONLFiles(filepath.Join(path, "sessions"), &files) - collectJSONLFiles(filepath.Join(path, "archived_sessions"), &files) - } - sort.Strings(files) - return files -} - -func ReadUsageFile(path string) ([]usage.Entry, error) { - file, err := os.Open(path) - if errors.Is(err, os.ErrNotExist) { - return nil, nil - } - if err != nil { - return nil, err - } - defer file.Close() - - state := fileState{ - sessionID: sessionIDFromFilename(path), - } - entries := make([]usage.Entry, 0) - reader := bufio.NewReader(file) - lineNumber := 0 - offset := int64(0) - for { - line, readErr := reader.ReadBytes('\n') - if len(line) > 0 { - lineNumber++ - start := offset - offset += int64(len(line)) - line = bytes.TrimRight(line, "\r\n") - if entry, ok := parseLine(line, &state); ok { - entry.SourceFile = path - entry.SourceLine = lineNumber - entry.SourceStart = start - entry.SourceEnd = offset - entry.ID = stableEntryID(entry) - entries = append(entries, entry) - } - } - if readErr == nil { - continue - } - if errors.Is(readErr, io.EOF) { - break - } - return nil, readErr - } - return entries, nil -} - -type fileState struct { - sessionID string - projectPath string - model string - language string - client string -} - -func parseLine(line []byte, state *fileState) (usage.Entry, bool) { - if !bytes.Contains(line, []byte(`"type"`)) { - return usage.Entry{}, false - } - - var envelope codexLine - if err := json.Unmarshal(line, &envelope); err != nil { - return usage.Entry{}, false - } - - if language := languageFromPayload(envelope.Payload); language != langdetect.Unknown { - state.language = language - } - - switch envelope.Type { - case "session_meta": - var payload sessionMetaPayload - if err := json.Unmarshal(envelope.Payload, &payload); err != nil { - return usage.Entry{}, false - } - if strings.TrimSpace(payload.ID) != "" { - state.sessionID = payload.ID - } - if strings.TrimSpace(payload.CWD) != "" { - state.projectPath = payload.CWD - } - if client := usage.NormalizeClient(usage.ProviderCodex, payload.Originator); client != "" { - state.client = client - } - return usage.Entry{}, false - case "turn_context": - var payload turnContextPayload - if err := json.Unmarshal(envelope.Payload, &payload); err != nil { - return usage.Entry{}, false - } - if strings.TrimSpace(payload.CWD) != "" { - state.projectPath = payload.CWD - } - if strings.TrimSpace(payload.Model) != "" { - state.model = payload.Model - } - return usage.Entry{}, false - case "event_msg": - var payload eventPayload - if err := json.Unmarshal(envelope.Payload, &payload); err != nil { - return usage.Entry{}, false - } - if payload.Type != "token_count" || payload.Info.LastTokenUsage == nil { - return usage.Entry{}, false - } - timestamp, err := time.Parse(time.RFC3339Nano, envelope.Timestamp) - 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 - } - - return usage.Entry{ - Provider: usage.ProviderCodex, - Timestamp: timestamp, - Date: timestamp.In(time.Local).Format("2006-01-02"), - Project: projectName(state.projectPath), - ProjectPath: state.projectPath, - SessionID: state.sessionID, - Model: state.model, - 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 - default: - return usage.Entry{}, false - } -} - -func stateLanguage(state *fileState) string { - return usage.NormalizeLanguage(state.language) -} - -func languageFromPayload(raw json.RawMessage) string { - if len(raw) == 0 { - return langdetect.Unknown - } - - var payload map[string]any - if err := json.Unmarshal(raw, &payload); err != nil { - return langdetect.Unknown - } - candidates := candidatesFromGenericValue(payload, 1) - - if arguments, ok := payload["arguments"].(string); ok { - var parsed any - if err := json.Unmarshal([]byte(arguments), &parsed); err == nil { - candidates = append(candidates, candidatesFromGenericValue(parsed, 3)...) - } else { - candidates = append(candidates, candidatesFromText(arguments, 1)...) - } - } - - if input, ok := payload["input"].(string); ok { - var parsed any - if err := json.Unmarshal([]byte(input), &parsed); err == nil { - candidates = append(candidates, candidatesFromGenericValue(parsed, 3)...) - } else { - candidates = append(candidates, candidatesFromText(input, 1)...) - } - } - - return langdetect.Dominant(candidates) -} - -func candidatesFromGenericValue(value any, weight int) []langdetect.Candidate { - candidates := make([]langdetect.Candidate, 0) - switch typed := value.(type) { - case map[string]any: - for key, child := range typed { - if langdetect.FromPath(key) != langdetect.Unknown { - candidates = append(candidates, langdetect.Candidate{Path: key, Weight: weight}) - } - - lowerKey := strings.ToLower(key) - switch { - case strings.Contains(lowerKey, "file") || strings.Contains(lowerKey, "path"): - candidates = append(candidates, candidatesFromPathValue(child, weight+2)...) - case strings.Contains(lowerKey, "command") || lowerKey == "cmd" || strings.Contains(lowerKey, "query") || strings.Contains(lowerKey, "content"): - candidates = append(candidates, candidatesFromTextValue(child, weight)...) - default: - candidates = append(candidates, candidatesFromGenericValue(child, weight)...) - } - } - case []any: - for _, child := range typed { - candidates = append(candidates, candidatesFromGenericValue(child, weight)...) - } - case string: - if langdetect.FromPath(typed) != langdetect.Unknown { - candidates = append(candidates, langdetect.Candidate{Path: typed, Weight: weight}) - } - } - return candidates -} - -func candidatesFromPathValue(value any, weight int) []langdetect.Candidate { - switch typed := value.(type) { - case string: - if langdetect.FromPath(typed) != langdetect.Unknown { - return []langdetect.Candidate{{Path: typed, Weight: weight}} - } - return candidatesFromText(typed, 1) - case []any: - candidates := make([]langdetect.Candidate, 0, len(typed)) - for _, child := range typed { - candidates = append(candidates, candidatesFromPathValue(child, weight)...) - } - return candidates - default: - return candidatesFromGenericValue(value, weight) - } -} - -func candidatesFromTextValue(value any, weight int) []langdetect.Candidate { - text, ok := value.(string) - if !ok { - return nil - } - return candidatesFromText(text, weight) -} - -func candidatesFromText(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 -} - -func stableEntryID(entry usage.Entry) string { - return usage.StableID( - string(usage.ProviderCodex), - entry.SourceFile, - strconv.Itoa(entry.SourceLine), - entry.Timestamp.Format(time.RFC3339Nano), - entry.Model, - strconv.FormatUint(entry.Usage.InputTokens, 10), - strconv.FormatUint(entry.Usage.CachedInputTokens, 10), - strconv.FormatUint(entry.Usage.OutputTokens, 10), - strconv.FormatUint(entry.Usage.ReasoningOutputTokens, 10), - strconv.FormatUint(entry.Usage.TotalTokens, 10), - ) -} - -func projectName(path string) string { - clean := strings.TrimSpace(path) - if clean == "" { - return "unknown" - } - base := filepath.Base(filepath.Clean(clean)) - if base == "." || base == string(filepath.Separator) || base == "" { - return "unknown" - } - return base -} - -func sessionIDFromFilename(path string) string { - base := filepath.Base(path) - sessionID := strings.TrimSuffix(base, filepath.Ext(base)) - sessionID = strings.TrimPrefix(sessionID, "rollout-") - if strings.TrimSpace(sessionID) == "" { - return "unknown" - } - return sessionID -} - -func collectJSONLFiles(dir string, files *[]string) { - entries, err := os.ReadDir(dir) - if err != nil { - return - } - for _, entry := range entries { - path := filepath.Join(dir, entry.Name()) - if entry.IsDir() { - collectJSONLFiles(path, files) - continue - } - if strings.EqualFold(filepath.Ext(path), ".jsonl") { - *files = append(*files, path) - } - } -} - -func dirExists(path string) bool { - info, err := os.Stat(path) - return err == nil && info.IsDir() -} - -func expandHomePath(raw string) string { - home, err := os.UserHomeDir() - if err != nil || home == "" { - return raw - } - if raw == "~" { - return home - } - if strings.HasPrefix(raw, "~/") { - return filepath.Join(home, strings.TrimPrefix(raw, "~/")) - } - return raw -} diff --git a/internal/codexusage/loader_test.go b/internal/codexusage/loader_test.go deleted file mode 100644 index f51e3c6..0000000 --- a/internal/codexusage/loader_test.go +++ /dev/null @@ -1,148 +0,0 @@ -package codexusage - -import ( - "os" - "path/filepath" - "testing" -) - -func TestReadUsageFileParsesTokenCountEvents(t *testing.T) { - path := filepath.Join(t.TempDir(), "sessions", "2026", "06", "03", "rollout-session-a.jsonl") - mkdirAll(t, filepath.Dir(path)) - writeFile(t, path, ` -{"timestamp":"2026-06-03T01:02:03Z","type":"session_meta","payload":{"id":"session-a","cwd":"/Users/me/workspace/tokitoki"}} -{"timestamp":"2026-06-03T01:02:04Z","type":"turn_context","payload":{"cwd":"/Users/me/workspace/tokitoki","model":"gpt-5.2-codex"}} -{"timestamp":"2026-06-03T01:02:05Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":10,"reasoning_output_tokens":3,"total_tokens":110},"last_token_usage":{"input_tokens":40,"cached_input_tokens":8,"output_tokens":5,"reasoning_output_tokens":2,"total_tokens":45}}}} -{"timestamp":"2026-06-03T01:02:06Z","type":"event_msg","payload":{"type":"agent_message","message":"ignored"}} -`) - - entries, err := ReadUsageFile(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.Project != "tokitoki" { - t.Fatalf("project = %q, want tokitoki", entry.Project) - } - if entry.ProjectPath != "/Users/me/workspace/tokitoki" { - t.Fatalf("project path = %q, want cwd", entry.ProjectPath) - } - if entry.SessionID != "session-a" { - t.Fatalf("session id = %q, want session-a", entry.SessionID) - } - if entry.Model != "gpt-5.2-codex" { - t.Fatalf("model = %q, want gpt-5.2-codex", entry.Model) - } - 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) - } - if entry.Usage.CacheReadInputTokens != 8 { - t.Fatalf("cache read tokens = %d, want 8 (cached portion)", entry.Usage.CacheReadInputTokens) - } - if entry.Usage.ReasoningOutputTokens != 2 { - t.Fatalf("reasoning output tokens = %d, want 2", entry.Usage.ReasoningOutputTokens) - } - if entry.Usage.TotalTokens != 45 { - t.Fatalf("total tokens = %d, want 45", entry.Usage.TotalTokens) - } -} - -func TestReadUsageFileInfersLanguageFromPriorToolPayload(t *testing.T) { - path := filepath.Join(t.TempDir(), "sessions", "2026", "06", "03", "rollout-session-a.jsonl") - mkdirAll(t, filepath.Dir(path)) - writeFile(t, path, ` -{"timestamp":"2026-06-03T01:02:03Z","type":"session_meta","payload":{"id":"session-a","cwd":"/Users/me/workspace/tokitoki"}} -{"timestamp":"2026-06-03T01:02:04Z","type":"response_item","payload":{"type":"function_call","name":"exec_command","arguments":"{\"cmd\":\"sed -n '1,20p' internal/httpapi/server.go\",\"workdir\":\"/Users/me/workspace/tokitoki\"}"}} -{"timestamp":"2026-06-03T01:02:05Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"output_tokens":2,"total_tokens":12}}}} -{"timestamp":"2026-06-03T01:02:06Z","type":"event_msg","payload":{"type":"patch_apply_end","changes":{"/Users/me/workspace/app/page.tsx":{"status":"modified"}}}} -{"timestamp":"2026-06-03T01:02:07Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":20,"output_tokens":3,"total_tokens":23}}}} -`) - - entries, err := ReadUsageFile(path) - if err != nil { - t.Fatal(err) - } - if len(entries) != 2 { - t.Fatalf("len(entries) = %d, want 2", len(entries)) - } - if entries[0].Language != "Go" { - t.Fatalf("first language = %q, want Go", entries[0].Language) - } - if entries[1].Language != "TypeScript" { - t.Fatalf("second language = %q, want TypeScript", entries[1].Language) - } -} - -func TestUsageFilesIncludesSessionsAndArchivedSessions(t *testing.T) { - dir := t.TempDir() - active := filepath.Join(dir, "sessions", "2026", "06", "03", "active.jsonl") - archived := filepath.Join(dir, "archived_sessions", "archived.jsonl") - mkdirAll(t, filepath.Dir(active)) - mkdirAll(t, filepath.Dir(archived)) - writeFile(t, active, "{}") - writeFile(t, archived, "{}") - - files := UsageFiles([]string{dir}) - - if len(files) != 2 { - t.Fatalf("len(files) = %d, want 2", len(files)) - } -} - -func TestLoadEntriesFiltersByProjectOrProjectPath(t *testing.T) { - dir := t.TempDir() - path := filepath.Join(dir, "sessions", "2026", "06", "03", "rollout-session-a.jsonl") - mkdirAll(t, filepath.Dir(path)) - writeFile(t, path, ` -{"timestamp":"2026-06-03T01:02:03Z","type":"session_meta","payload":{"id":"session-a","cwd":"/Users/me/workspace/tokitoki"}} -{"timestamp":"2026-06-03T01:02:05Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}}} -`) - - entries, err := LoadEntriesFromPaths([]string{dir}, "tokitoki", nil) - if err != nil { - t.Fatal(err) - } - if len(entries) != 1 { - t.Fatalf("len(entries) = %d, want 1", len(entries)) - } - - entries, err = LoadEntriesFromPaths([]string{dir}, "/Users/me/workspace/tokitoki", nil) - if err != nil { - t.Fatal(err) - } - if len(entries) != 1 { - t.Fatalf("len(entries by path) = %d, want 1", len(entries)) - } - - entries, err = LoadEntriesFromPaths([]string{dir}, "other", nil) - if err != nil { - t.Fatal(err) - } - if len(entries) != 0 { - t.Fatalf("len(entries for other) = %d, want 0", len(entries)) - } -} - -func mkdirAll(t *testing.T, path string) { - t.Helper() - if err := os.MkdirAll(path, 0o700); err != nil { - t.Fatal(err) - } -} - -func writeFile(t *testing.T, path, data string) { - t.Helper() - if err := os.WriteFile(path, []byte(data), 0o600); err != nil { - t.Fatal(err) - } -} 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/provider/amp/loader.go b/internal/provider/amp/loader.go new file mode 100644 index 0000000..f36057b --- /dev/null +++ b/internal/provider/amp/loader.go @@ -0,0 +1,185 @@ +package amp + +import ( + "encoding/json" + "os" + "path/filepath" + "sort" + "strconv" + + "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 { + if info, err := os.Stat(path); err == nil && !info.IsDir() { + if filepath.Ext(path) == ".json" { + files = append(files, path) + } + continue + } + files = append(files, agentdata.CollectExt(filepath.Join(path, "threads"), ".json")...) + if filepath.Base(path) == "threads" { + files = append(files, agentdata.CollectExt(path, ".json")...) + } + } + sort.Strings(files) + files = agentdata.FilterFiles(agentdata.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...) + } + usageprovider.SortEntries(entries) + return entries, nil +} + +func parseThreadFile(path string) ([]usage.Entry, error) { + thread, err := agentdata.ReadJSONObject(path) + if err != nil || thread == nil { + return nil, err + } + threadID := agentdata.StringField(thread, "id") + if threadID == "" { + return nil, nil + } + 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 + } + } + 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 := agentdata.ObjectAt(raw) + if event == nil { + continue + } + timestamp, ok := agentdata.ParseTimestamp(event["timestamp"]) + if !ok { + continue + } + model := agentdata.StringField(event, "model") + if model == "" { + continue + } + tokenBlock := agentdata.ObjectAt(event["tokens"]) + if tokenBlock == nil { + continue + } + cache := cacheTokens[int64Value(event["toMessageId"])] + tokens := usage.TokenUsage{ + InputTokens: agentdata.UintField(tokenBlock, "input"), + OutputTokens: agentdata.UintField(tokenBlock, "output"), + CacheCreationInputTokens: cache.cacheCreation, + CacheReadInputTokens: cache.cacheRead, + } + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(tokenBlock, "total")) + if !usageprovider.NonZero(tokens) { + continue + } + 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 +} + +func messageEntries(path, threadID string, messages []any) []usage.Entry { + entries := make([]usage.Entry, 0) + for index, raw := range messages { + message := agentdata.ObjectAt(raw) + if message == nil || agentdata.StringValue(message["role"]) != "assistant" { + continue + } + usageBlock := agentdata.ObjectAt(message["usage"]) + if usageBlock == nil { + continue + } + timestamp, ok := agentdata.ParseTimestamp(usageBlock["timestamp"]) + if !ok { + timestamp, ok = agentdata.ParseTimestamp(message["timestamp"]) + } + if !ok { + continue + } + model := agentdata.StringField(usageBlock, "model") + if model == "" { + model = agentdata.StringValue(message["model"]) + } + if model == "" { + continue + } + tokens := usage.TokenUsage{ + InputTokens: agentdata.UintField(usageBlock, "inputTokens"), + OutputTokens: agentdata.UintField(usageBlock, "outputTokens"), + CacheCreationInputTokens: agentdata.UintField(usageBlock, "cacheCreationInputTokens"), + CacheReadInputTokens: agentdata.UintField(usageBlock, "cacheReadInputTokens"), + } + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(usageBlock, "totalTokens")) + if !usageprovider.NonZero(tokens) { + continue + } + 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 +} + +type cache struct { + cacheCreation uint64 + cacheRead uint64 +} + +func cacheTokens(messages []any) map[int64]cache { + tokens := make(map[int64]cache) + for _, raw := range messages { + message := agentdata.ObjectAt(raw) + if message == nil || agentdata.StringValue(message["role"]) != "assistant" { + continue + } + id := int64Value(message["messageId"]) + if id == 0 { + continue + } + usageBlock := agentdata.ObjectAt(message["usage"]) + tokens[id] = cache{ + cacheCreation: agentdata.UintField(usageBlock, "cacheCreationInputTokens"), + cacheRead: agentdata.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..79772e5 --- /dev/null +++ b/internal/provider/amp/provider.go @@ -0,0 +1,36 @@ +package amp + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Amp usage entries. +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 = usageprovider.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 usageprovider.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..e16318e --- /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/providertest" + "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") + 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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 64% rename from internal/claudeusage/loader.go rename to internal/provider/claude/loader.go index 85520bb..1c5300f 100644 --- a/internal/claudeusage/loader.go +++ b/internal/provider/claude/loader.go @@ -1,4 +1,4 @@ -package claudeusage +package claude import ( "bufio" @@ -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"` @@ -70,21 +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"` - 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 { @@ -109,22 +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, + 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, @@ -208,69 +232,193 @@ 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() - project := ExtractProject(path) - sessionID, projectPath := ExtractSessionParts(path) + 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 + // 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) + 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") - entry, ok := parseUsageLine(line, project, sessionID, projectPath) - if ok { + 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) + 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 { + if complete { continue } if errors.Is(readErr, io.EOF) { break } - return nil, readErr + return nil, 0, readErr } - return entries, nil + return entries, consumed, nil } -func ExtractProject(path string) string { - parts := pathParts(path) - for i, part := range parts { - if part != "projects" { - continue +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 } - if i+1 >= len(parts) || strings.TrimSpace(parts[i+1]) == "" { - return "unknown" + } + 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 } - projectPath := normalizeClaudeProjectPathParts([]string{parts[i+1]}) - project := filepath.Base(filepath.Clean(projectPath)) - if strings.TrimSpace(project) == "" || project == "." || project == string(filepath.Separator) { - return "unknown" + return patchStats{file: result.FilePath, added: usage.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 project } - return "unknown" + return stats, 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 +436,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" - } - 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 relative[len(relative)-2] } - 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 +467,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 := usage.ProjectFromCWD(*data.CWD); ok { + projectPath = path + project = name + } + } + model := "" if data.Message.Model != nil && *data.Message.Model != "" { model = *data.Message.Model @@ -369,6 +492,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 +507,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 } @@ -496,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, ) } @@ -543,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 == "" { @@ -756,4 +866,3 @@ func pathParts(path string) []string { }) return parts } - diff --git a/internal/provider/claude/loader_test.go b/internal/provider/claude/loader_test.go new file mode 100644 index 0000000..03410bc --- /dev/null +++ b/internal/provider/claude/loader_test.go @@ -0,0 +1,476 @@ +package claude + +import ( + "os" + "path/filepath" + "testing" + "time" + + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func TestUsageFilesLimitsDiscoveryToProjectFilter(t *testing.T) { + dir := t.TempDir() + projectA := filepath.Join(dir, "projects", "project-a", "session-a") + projectB := filepath.Join(dir, "projects", "project-b", "session-b") + mkdirAll(t, projectA) + mkdirAll(t, projectB) + writeFile(t, filepath.Join(projectA, "a.jsonl"), "{}") + writeFile(t, filepath.Join(projectB, "b.jsonl"), "{}") + + files := UsageFiles([]string{dir}, "project-a") + + if len(files) != 1 { + t.Fatalf("len(files) = %d, want 1", len(files)) + } + if got := files[0]; !containsPathSegment(got, "project-a") { + t.Fatalf("file = %q, want project-a path", got) + } +} + +func TestUsageFilesFallsBackForNonSegmentProjectFilter(t *testing.T) { + dir := t.TempDir() + projectA := filepath.Join(dir, "projects", "project-a", "session-a") + projectB := filepath.Join(dir, "projects", "project-b", "session-b") + mkdirAll(t, projectA) + mkdirAll(t, projectB) + writeFile(t, filepath.Join(projectA, "a.jsonl"), "{}") + writeFile(t, filepath.Join(projectB, "b.jsonl"), "{}") + + files := UsageFiles([]string{dir}, "project-a/session-a") + + if len(files) != 2 { + t.Fatalf("len(files) = %d, want 2", len(files)) + } +} + +func TestProjectPathSegmentRejectsUnsafeValues(t *testing.T) { + cases := map[string]bool{ + "": false, + ".": false, + "..": false, + "project-a/session-a": false, + `project-a\session-a`: false, + "project-a": true, + } + + for value, want := range cases { + if got := isProjectPathSegment(value); got != want { + t.Fatalf("isProjectPathSegment(%q) = %v, want %v", value, got, want) + } + } +} + +func TestExtractSessionID(t *testing.T) { + tests := []struct { + name string + path string + wantSessionID string + }{ + { + 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", + }, + { + 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", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + if sessionID := ExtractSessionID(tt.path); sessionID != tt.wantSessionID { + t.Fatalf("sessionID = %q, want %q", sessionID, tt.wantSessionID) + } + }) + } +} + +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 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)) + 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}}}`) + + 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 := usage.ProjectFromCWD(cwd); ok { + t.Fatalf("usage.ProjectFromCWD(%q) ok = true, want false", cwd) + } + } +} + +func TestHasUnsupportedNullField(t *testing.T) { + rejected := [][]byte{ + []byte(`{"message":{"usage":{"speed":null}}}`), + []byte(`{"message":{"model":null,"usage":{"input_tokens":0}}}`), + []byte(`{"sessionId":null,"message":{"usage":{"input_tokens":0}}}`), + } + for _, line := range rejected { + if !hasUnsupportedNullField(line) { + t.Fatalf("hasUnsupportedNullField(%s) = false, want true", line) + } + } + + allowed := []byte(`{"message":{"content":null,"usage":{"input_tokens":0}}}`) + if hasUnsupportedNullField(allowed) { + t.Fatalf("hasUnsupportedNullField(%s) = true, want false", allowed) + } +} + +func TestReadUsageFileParsesUsageLines(t *testing.T) { + path := filepath.Join(t.TempDir(), "projects", "project-a", "session-a", "chat.jsonl") + 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","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) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + + entry := entries[0] + if entry.Project != "project-a" { + t.Fatalf("project = %q, want project-a", entry.Project) + } + if entry.SessionID != "session-a" { + t.Fatalf("sessionID = %q, want session-a", entry.SessionID) + } + if entry.Model != "claude-sonnet-4-20250514-fast" { + t.Fatalf("model = %q, want fast suffix", entry.Model) + } + if entry.Language != "Unknown" { + t.Fatalf("language = %q, want Unknown", entry.Language) + } + if entry.Date != "2026-05-21" { + t.Fatalf("date = %q, want 2026-05-21", entry.Date) + } + if got := tokenTotal(entry.Data.Message.Usage); got != 20 { + t.Fatalf("tokenTotal = %d, want 20", got) + } +} + +func TestReadUsageFileInfersLanguageFromToolUseFilePath(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","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1},"content":[{"type":"tool_use","name":"Read","input":{"file_path":"/repo/internal/server/server.go"}}]}}`) + + 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].Language != "Go" { + t.Fatalf("language = %q, want Go", entries[0].Language) + } +} + +func TestReadUsageFileDoesNotInferLanguageFromCodeFenceWithoutFilePath(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\",\"message\":{\"id\":\"msg-1\",\"model\":\"claude\",\"usage\":{\"input_tokens\":1,\"output_tokens\":1},\"content\":[{\"type\":\"text\",\"text\":\"```tsx\\nexport default function Page() {}\\n```\"}]}}\n") + + 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].Language != "Unknown" { + t.Fatalf("language = %q, want Unknown", entries[0].Language) + } +} + +func TestReadUsageFileSkipsUnsupportedSpeed(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","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1,"speed":"turbo"}}}`) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 0 { + t.Fatalf("len(entries) = %d, want 0", len(entries)) + } +} + +func TestLoadEntriesDeduplicatesByMessageAndRequest(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "projects", "project-a", "session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, ` +{"timestamp":"2026-05-21T01:02:03Z","requestId":"req-1","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1}}} +{"timestamp":"2026-05-21T01:02:04Z","requestId":"req-1","message":{"id":"msg-1","model":"claude","usage":{"input_tokens":10,"output_tokens":1}}} +`) + + entries, err := LoadEntriesFromPaths([]string{dir}, "", nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + if got := entries[0].Data.Message.Usage.InputTokens; got != 10 { + t.Fatalf("input tokens = %d, want replacement with larger usage", got) + } +} + +func TestUsageLimitResetTimeFromLine(t *testing.T) { + isAPIError := true + line := []byte(`{"timestamp":"2026-05-21T01:02:03Z","isApiErrorMessage":true,"message":{"id":"msg-1","model":"claude","usage":{"input_tokens":1,"output_tokens":1},"content":"Claude AI usage limit reached|1779325200"}}`) + + reset := usageLimitResetTimeFromLine(line, &isAPIError) + if reset == nil { + t.Fatal("reset = nil, want timestamp") + } + if want := time.Unix(1779325200, 0).UTC(); !reset.Equal(want) { + t.Fatalf("reset = %s, want %s", reset, want) + } +} + +func mkdirAll(t *testing.T, path string) { + t.Helper() + if err := os.MkdirAll(path, 0o700); err != nil { + t.Fatal(err) + } +} + +func writeFile(t *testing.T, path, data string) { + t.Helper() + if err := os.WriteFile(path, []byte(data), 0o600); err != nil { + t.Fatal(err) + } +} + +func containsPathSegment(path, segment string) bool { + for _, part := range pathParts(path) { + if part == segment { + return true + } + } + 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/claudeusage/provider.go b/internal/provider/claude/provider.go similarity index 54% rename from internal/claudeusage/provider.go rename to internal/provider/claude/provider.go index 6a96afc..6511b06 100644 --- a/internal/claudeusage/provider.go +++ b/internal/provider/claude/provider.go @@ -1,4 +1,4 @@ -package claudeusage +package claude import ( "sort" @@ -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/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) +} diff --git a/internal/agentusage/codebuff.go b/internal/provider/codebuff/loader.go similarity index 58% rename from internal/agentusage/codebuff.go rename to internal/provider/codebuff/loader.go index 6bac979..a88c36e 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,13 @@ import ( "strings" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) -type codebuffUsage struct { +type tokenUsage struct { model string inputTokens uint64 outputTokens uint64 @@ -22,17 +25,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 = agentdata.FilterFiles(agentdata.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 +47,11 @@ func loadCodebuffEntries(paths []string, filter usage.FileFilter) ([]usage.Entry for _, entry := range entriesByID { entries = append(entries, entry) } - sortEntries(entries) + usageprovider.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 +66,12 @@ func collectCodebuffChatFiles(root string) []string { if filepath.Base(root) != "projects" { projectRoot = filepath.Join(root, "projects") } - return collectFiles(projectRoot, func(path string) bool { + return agentdata.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 +82,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 := agentdata.FileModifiedTime(path) entries := make([]usage.Entry, 0) for index, raw := range messages { - message := objectAt(raw) - if !isCodebuffAssistant(message) { + message := agentdata.ObjectAt(raw) + if !isAssistant(message) { continue } - parsedUsage := extractCodebuffUsage(message) + parsedUsage := extractUsage(message) tokens := usage.TokenUsage{ InputTokens: parsedUsage.inputTokens, OutputTokens: parsedUsage.outputTokens, @@ -97,16 +100,16 @@ func parseCodebuffChatFile(path string) ([]usage.Entry, error) { ReasoningOutputTokens: parsedUsage.extraTotalTokens, } if tokens.TotalTokens == 0 { - tokens.TotalTokens = totalUsage(tokens) + tokens.TotalTokens = usageprovider.TotalUsage(tokens) } - if !nonZero(tokens) { + if !usageprovider.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 +117,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 := 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 } -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 +143,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 := agentdata.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 := agentdata.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 = agentdata.StringField(metadata, "model") + mergeCodebuffUsage(&usage, parseUsageObject(metadata["usage"])) + mergeCodebuffUsage(&usage, parseUsageObject(agentdata.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 := agentdata.ArrayAt(agentdata.ObjectAt(agentdata.ObjectAt(agentdata.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 := agentdata.ObjectAt(history[i]) + if agentdata.StringField(entry, "role") != "assistant" { continue } - providerOptions := objectAt(entry["providerOptions"]) + providerOptions := agentdata.ObjectAt(entry["providerOptions"]) if providerOptions == nil { continue } - entryUsage := parseCodebuffUsageObject(providerOptions["usage"]) - codebuff := objectAt(providerOptions["codebuff"]) + entryUsage := parseUsageObject(providerOptions["usage"]) + codebuff := agentdata.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 = agentdata.StringField(codebuff, "model") } } - if codebuffUsageHasTokens(entryUsage) || entryUsage.model != "" { + if usageHasTokens(entryUsage) || entryUsage.model != "" { found = true } mergeCodebuffUsage(&usage, entryUsage) @@ -194,27 +197,27 @@ func codebuffRunStateUsage(metadata map[string]any) *codebuffUsage { return &usage } -func parseCodebuffUsageObject(value any) codebuffUsage { - record := objectAt(value) +func parseUsageObject(value any) tokenUsage { + record := agentdata.ObjectAt(value) if record == nil { - return codebuffUsage{} + return tokenUsage{} } - parsed := codebuffUsage{ - model: stringField(record, "model"), + parsed := tokenUsage{ + 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(objectAt(record["promptTokensDetails"]), "cachedTokens")) - parsed.cacheReadInputTokens = maxUint64(parsed.cacheReadInputTokens, firstUint(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 = 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 @@ -223,7 +226,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 +247,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 := agentdata.ParseTimestamp(message["timestamp"]); ok { return timestamp, true } - if timestamp, ok := parseTimestamp(message["createdAt"]); ok { + if timestamp, ok := agentdata.ParseTimestamp(message["createdAt"]); ok { return timestamp, true } - return parseTimestamp(objectAt(message["metadata"])["timestamp"]) + return agentdata.ParseTimestamp(agentdata.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 +271,21 @@ func parseCodebuffChatTimestamp(chatID string) (time.Time, bool) { clock = clock[:index] + ":" + clock[index+1:] } } - return parseTimestampString(date + "T" + clock) + return agentdata.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 := agentdata.StringField(message, "id"); id != "" { return "codebuff:" + sessionID + ":" + id } - return stableEntryID(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 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 new file mode 100644 index 0000000..433054f --- /dev/null +++ b/internal/provider/codebuff/provider.go @@ -0,0 +1,36 @@ +package codebuff + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Codebuff usage entries. +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 = usageprovider.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 usageprovider.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..0d5a516 --- /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/providertest" + "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") + 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() + }() + + 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", + Project: "codebuff", + Tokens: usage.TokenUsage{ + InputTokens: 100, + OutputTokens: 50, + CacheCreationInputTokens: 20, + CacheReadInputTokens: 10, + TotalTokens: 180, + }, + }) +} diff --git a/internal/provider/codex/loader.go b/internal/provider/codex/loader.go new file mode 100644 index 0000000..c9d2c7f --- /dev/null +++ b/internal/provider/codex/loader.go @@ -0,0 +1,702 @@ +package codex + +import ( + "bufio" + "bytes" + "encoding/json" + "errors" + "io" + "os" + "path/filepath" + "runtime" + "sort" + "strconv" + "strings" + "time" + + "github.com/tokitoki-dev/tokitoki-cli/internal/langdetect" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +var ErrNoDataDirs = errors.New("no valid Codex data directories found") + +type codexLine struct { + Timestamp string `json:"timestamp"` + Type string `json:"type"` + Payload json.RawMessage `json:"payload"` +} + +type sessionMetaPayload struct { + ID string `json:"id"` + CWD string `json:"cwd"` + Originator string `json:"originator"` +} + +type turnContextPayload struct { + CWD string `json:"cwd"` + Model string `json:"model"` +} + +type eventPayload struct { + 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"` +} + +type tokenUsagePayload struct { + InputTokens uint64 `json:"input_tokens"` + CachedInputTokens uint64 `json:"cached_input_tokens"` + OutputTokens uint64 `json:"output_tokens"` + ReasoningOutputTokens uint64 `json:"reasoning_output_tokens"` + TotalTokens uint64 `json:"total_tokens"` +} + +func LoadEntriesFromPaths(paths []string, projectFilter string, fileFilter usage.FileFilter) ([]usage.Entry, error) { + files := UsageFiles(paths) + entries := make([]usage.Entry, 0) + for _, file := range files { + if fileFilter != nil && !fileFilter(file) { + continue + } + fileEntries, err := ReadUsageFile(file) + if err != nil { + return nil, err + } + for _, entry := range fileEntries { + if projectFilter != "" && entry.Project != projectFilter && entry.ProjectPath != projectFilter { + continue + } + entries = append(entries, entry) + } + } + return entries, nil +} + +func UsageFiles(paths []string) []string { + files := make([]string, 0) + for _, path := range paths { + collectJSONLFiles(filepath.Join(path, "sessions"), &files) + collectJSONLFiles(filepath.Join(path, "archived_sessions"), &files) + } + sort.Strings(files) + return files +} + +func ReadUsageFile(path string) ([]usage.Entry, error) { + file, err := os.Open(path) + if errors.Is(err, os.ErrNotExist) { + return nil, nil + } + if err != nil { + return nil, err + } + defer file.Close() + + state := fileState{ + sessionID: sessionIDFromFilename(path), + } + entries := make([]usage.Entry, 0) + reader := bufio.NewReader(file) + lineNumber := 0 + offset := int64(0) + for { + line, readErr := reader.ReadBytes('\n') + if len(line) > 0 { + lineNumber++ + start := offset + offset += int64(len(line)) + line = bytes.TrimRight(line, "\r\n") + if entry, ok := parseLine(line, &state); ok { + entry.SourceFile = path + entry.SourceLine = lineNumber + entry.SourceStart = start + entry.SourceEnd = offset + entries = append(entries, entry) + } + } + if readErr == nil { + continue + } + if errors.Is(readErr, io.EOF) { + break + } + return nil, readErr + } + return entries, nil +} + +type fileState struct { + sessionID string + projectPath string + 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) { + if !bytes.Contains(line, []byte(`"type"`)) { + return usage.Entry{}, false + } + + var envelope codexLine + if err := json.Unmarshal(line, &envelope); err != nil { + return usage.Entry{}, false + } + + if language := languageFromPayload(envelope.Payload); language != langdetect.Unknown { + state.language = language + } + + 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 { + return usage.Entry{}, false + } + if strings.TrimSpace(payload.ID) != "" { + state.sessionID = payload.ID + } + if strings.TrimSpace(payload.CWD) != "" { + state.projectPath = payload.CWD + } + if client := usage.NormalizeClient(usage.ProviderCodex, payload.Originator); client != "" { + state.client = client + } + return usage.Entry{}, false + case "turn_context": + var payload turnContextPayload + if err := json.Unmarshal(envelope.Payload, &payload); err != nil { + return usage.Entry{}, false + } + if strings.TrimSpace(payload.CWD) != "" { + state.projectPath = payload.CWD + } + if strings.TrimSpace(payload.Model) != "" { + state.model = payload.Model + } + return usage.Entry{}, false + case "event_msg": + var payload eventPayload + 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 + } + timestamp, err := time.Parse(time.RFC3339Nano, envelope.Timestamp) + if err != nil { + return usage.Entry{}, false + } + 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 + } + + entry := usage.Entry{ + Provider: usage.ProviderCodex, + Timestamp: timestamp, + Date: timestamp.In(time.Local).Format("2006-01-02"), + Project: projectName(state.projectPath), + ProjectPath: state.projectPath, + SessionID: state.sessionID, + Model: state.model, + Language: stateLanguage(state), + OS: usage.NormalizeOS(runtime.GOOS), + Client: state.client, + 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 + } +} + +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 + } + + var payload map[string]any + if err := json.Unmarshal(raw, &payload); err != nil { + return langdetect.Unknown + } + candidates := candidatesFromGenericValue(payload, 1) + + if arguments, ok := payload["arguments"].(string); ok { + var parsed any + if err := json.Unmarshal([]byte(arguments), &parsed); err == nil { + candidates = append(candidates, candidatesFromGenericValue(parsed, 3)...) + } else { + candidates = append(candidates, candidatesFromText(arguments, 1)...) + } + } + + if input, ok := payload["input"].(string); ok { + var parsed any + if err := json.Unmarshal([]byte(input), &parsed); err == nil { + candidates = append(candidates, candidatesFromGenericValue(parsed, 3)...) + } else { + candidates = append(candidates, candidatesFromText(input, 1)...) + } + } + + return langdetect.Dominant(candidates) +} + +func candidatesFromGenericValue(value any, weight int) []langdetect.Candidate { + candidates := make([]langdetect.Candidate, 0) + switch typed := value.(type) { + case map[string]any: + for key, child := range typed { + if langdetect.FromPath(key) != langdetect.Unknown { + candidates = append(candidates, langdetect.Candidate{Path: key, Weight: weight}) + } + + lowerKey := strings.ToLower(key) + switch { + case strings.Contains(lowerKey, "file") || strings.Contains(lowerKey, "path"): + candidates = append(candidates, candidatesFromPathValue(child, weight+2)...) + case strings.Contains(lowerKey, "command") || lowerKey == "cmd" || strings.Contains(lowerKey, "query") || strings.Contains(lowerKey, "content"): + candidates = append(candidates, candidatesFromTextValue(child, weight)...) + default: + candidates = append(candidates, candidatesFromGenericValue(child, weight)...) + } + } + case []any: + for _, child := range typed { + candidates = append(candidates, candidatesFromGenericValue(child, weight)...) + } + case string: + if langdetect.FromPath(typed) != langdetect.Unknown { + candidates = append(candidates, langdetect.Candidate{Path: typed, Weight: weight}) + } + } + return candidates +} + +func candidatesFromPathValue(value any, weight int) []langdetect.Candidate { + switch typed := value.(type) { + case string: + if langdetect.FromPath(typed) != langdetect.Unknown { + return []langdetect.Candidate{{Path: typed, Weight: weight}} + } + return candidatesFromText(typed, 1) + case []any: + candidates := make([]langdetect.Candidate, 0, len(typed)) + for _, child := range typed { + candidates = append(candidates, candidatesFromPathValue(child, weight)...) + } + return candidates + default: + return candidatesFromGenericValue(value, weight) + } +} + +func candidatesFromTextValue(value any, weight int) []langdetect.Candidate { + text, ok := value.(string) + if !ok { + return nil + } + return candidatesFromText(text, weight) +} + +func candidatesFromText(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 +} + +// 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), + entry.SessionID, + entry.Timestamp.Format(time.RFC3339Nano), + entry.Model, + strconv.FormatUint(entry.Usage.InputTokens, 10), + strconv.FormatUint(entry.Usage.CachedInputTokens, 10), + strconv.FormatUint(entry.Usage.OutputTokens, 10), + strconv.FormatUint(entry.Usage.ReasoningOutputTokens, 10), + strconv.FormatUint(entry.Usage.TotalTokens, 10), + ) +} + +func projectName(path string) string { + clean := strings.TrimSpace(path) + if clean == "" { + return usage.UnknownProject + } + base := filepath.Base(filepath.Clean(clean)) + if base == "." || base == string(filepath.Separator) || base == "" { + return usage.UnknownProject + } + return base +} + +func sessionIDFromFilename(path string) string { + base := filepath.Base(path) + sessionID := strings.TrimSuffix(base, filepath.Ext(base)) + sessionID = strings.TrimPrefix(sessionID, "rollout-") + if strings.TrimSpace(sessionID) == "" { + return "unknown" + } + return sessionID +} + +func collectJSONLFiles(dir string, files *[]string) { + entries, err := os.ReadDir(dir) + if err != nil { + return + } + for _, entry := range entries { + path := filepath.Join(dir, entry.Name()) + if entry.IsDir() { + collectJSONLFiles(path, files) + continue + } + if strings.EqualFold(filepath.Ext(path), ".jsonl") { + *files = append(*files, path) + } + } +} + +func dirExists(path string) bool { + info, err := os.Stat(path) + return err == nil && info.IsDir() +} + +func expandHomePath(raw string) string { + home, err := os.UserHomeDir() + if err != nil || home == "" { + return raw + } + if raw == "~" { + return home + } + if strings.HasPrefix(raw, "~/") { + return filepath.Join(home, strings.TrimPrefix(raw, "~/")) + } + return raw +} diff --git a/internal/provider/codex/loader_test.go b/internal/provider/codex/loader_test.go new file mode 100644 index 0000000..9ace368 --- /dev/null +++ b/internal/provider/codex/loader_test.go @@ -0,0 +1,361 @@ +package codex + +import ( + "os" + "path/filepath" + "testing" +) + +func TestReadUsageFileParsesTokenCountEvents(t *testing.T) { + path := filepath.Join(t.TempDir(), "sessions", "2026", "06", "03", "rollout-session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, ` +{"timestamp":"2026-06-03T01:02:03Z","type":"session_meta","payload":{"id":"session-a","cwd":"/Users/me/workspace/tokitoki"}} +{"timestamp":"2026-06-03T01:02:04Z","type":"turn_context","payload":{"cwd":"/Users/me/workspace/tokitoki","model":"gpt-5.2-codex"}} +{"timestamp":"2026-06-03T01:02:05Z","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"cached_input_tokens":20,"output_tokens":10,"reasoning_output_tokens":3,"total_tokens":110},"last_token_usage":{"input_tokens":40,"cached_input_tokens":8,"output_tokens":5,"reasoning_output_tokens":2,"total_tokens":45}}}} +{"timestamp":"2026-06-03T01:02:06Z","type":"event_msg","payload":{"type":"agent_message","message":"ignored"}} +`) + + entries, err := ReadUsageFile(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.Project != "tokitoki" { + t.Fatalf("project = %q, want tokitoki", entry.Project) + } + if entry.ProjectPath != "/Users/me/workspace/tokitoki" { + t.Fatalf("project path = %q, want cwd", entry.ProjectPath) + } + if entry.SessionID != "session-a" { + t.Fatalf("session id = %q, want session-a", entry.SessionID) + } + if entry.Model != "gpt-5.2-codex" { + t.Fatalf("model = %q, want gpt-5.2-codex", entry.Model) + } + if entry.Language != "Unknown" { + t.Fatalf("language = %q, want Unknown", entry.Language) + } + // 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 != 20 { + t.Fatalf("cache read tokens = %d, want 20 (cached portion)", entry.Usage.CacheReadInputTokens) + } + if entry.Usage.ReasoningOutputTokens != 3 { + t.Fatalf("reasoning output tokens = %d, want 3", entry.Usage.ReasoningOutputTokens) + } + 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) + } +} + +func TestReadUsageFileInfersLanguageFromPriorToolPayload(t *testing.T) { + path := filepath.Join(t.TempDir(), "sessions", "2026", "06", "03", "rollout-session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, ` +{"timestamp":"2026-06-03T01:02:03Z","type":"session_meta","payload":{"id":"session-a","cwd":"/Users/me/workspace/tokitoki"}} +{"timestamp":"2026-06-03T01:02:04Z","type":"response_item","payload":{"type":"function_call","name":"exec_command","arguments":"{\"cmd\":\"sed -n '1,20p' internal/httpapi/server.go\",\"workdir\":\"/Users/me/workspace/tokitoki\"}"}} +{"timestamp":"2026-06-03T01:02:05Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":10,"output_tokens":2,"total_tokens":12}}}} +{"timestamp":"2026-06-03T01:02:06Z","type":"event_msg","payload":{"type":"patch_apply_end","changes":{"/Users/me/workspace/app/page.tsx":{"status":"modified"}}}} +{"timestamp":"2026-06-03T01:02:07Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":20,"output_tokens":3,"total_tokens":23}}}} +`) + + entries, err := ReadUsageFile(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("len(entries) = %d, want 2", len(entries)) + } + if entries[0].Language != "Go" { + t.Fatalf("first language = %q, want Go", entries[0].Language) + } + if entries[1].Language != "TypeScript" { + t.Fatalf("second language = %q, want TypeScript", entries[1].Language) + } +} + +func TestUsageFilesIncludesSessionsAndArchivedSessions(t *testing.T) { + dir := t.TempDir() + active := filepath.Join(dir, "sessions", "2026", "06", "03", "active.jsonl") + archived := filepath.Join(dir, "archived_sessions", "archived.jsonl") + mkdirAll(t, filepath.Dir(active)) + mkdirAll(t, filepath.Dir(archived)) + writeFile(t, active, "{}") + writeFile(t, archived, "{}") + + files := UsageFiles([]string{dir}) + + if len(files) != 2 { + t.Fatalf("len(files) = %d, want 2", len(files)) + } +} + +func TestLoadEntriesFiltersByProjectOrProjectPath(t *testing.T) { + dir := t.TempDir() + path := filepath.Join(dir, "sessions", "2026", "06", "03", "rollout-session-a.jsonl") + mkdirAll(t, filepath.Dir(path)) + writeFile(t, path, ` +{"timestamp":"2026-06-03T01:02:03Z","type":"session_meta","payload":{"id":"session-a","cwd":"/Users/me/workspace/tokitoki"}} +{"timestamp":"2026-06-03T01:02:05Z","type":"event_msg","payload":{"type":"token_count","info":{"last_token_usage":{"input_tokens":1,"output_tokens":2,"total_tokens":3}}}} +`) + + entries, err := LoadEntriesFromPaths([]string{dir}, "tokitoki", nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries) = %d, want 1", len(entries)) + } + + entries, err = LoadEntriesFromPaths([]string{dir}, "/Users/me/workspace/tokitoki", nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 1 { + t.Fatalf("len(entries by path) = %d, want 1", len(entries)) + } + + entries, err = LoadEntriesFromPaths([]string{dir}, "other", nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 0 { + t.Fatalf("len(entries for other) = %d, want 0", len(entries)) + } +} + +func mkdirAll(t *testing.T, path string) { + t.Helper() + if err := os.MkdirAll(path, 0o700); err != nil { + t.Fatal(err) + } +} + +func writeFile(t *testing.T, path, data string) { + t.Helper() + if err := os.WriteFile(path, []byte(data), 0o600); err != nil { + t.Fatal(err) + } +} + +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() + 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) + } + 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(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 len(entries) != 1 || entries[0].IsWrite != nil { + t.Fatalf("failed patch counted: %+v", entries[0]) + } +} 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/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/filter_test.go b/internal/provider/copilot/filter_test.go new file mode 100644 index 0000000..ce0b240 --- /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/providertest" +) + +// 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") + 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). + 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/provider/copilot/loader.go b/internal/provider/copilot/loader.go new file mode 100644 index 0000000..95e7ad1 --- /dev/null +++ b/internal/provider/copilot/loader.go @@ -0,0 +1,675 @@ +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 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) + for _, file := range files { + fileEntries, err := parseOTELFile(file) + if err != nil { + return nil, err + } + 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 ( + copilotChatSpan sourceKind = iota + copilotInferenceLog + copilotAgentTurnLog + copilotAgentSummarySpan +) + +type candidate struct { + source sourceKind + traceID string + responseID string + sessionID string + model string + timestamp time.Time + tokens usage.TokenUsage + dedupKey string + sourceFile string + sourceLine int + sourceStart, sourceEnd int64 +} + +type traceContext struct { + model string + sessionID string + sessionIDPriority int +} + +func parseOTELFile(path string) ([]usage.Entry, error) { + lines, err := agentdata.ReadJSONLines(path, `"attributes"`) + if err != nil { + return nil, err + } + contexts := traceContexts(lines) + fallback := agentdata.FileModifiedTime(path) + candidates := make([]candidate, 0) + for index, line := range lines { + if candidate, ok := recordCandidate(path, line, index, fallback, contexts); ok { + candidates = append(candidates, candidate) + } + } + sets := candidateSets(candidates) + entries := make([]usage.Entry, 0) + for _, candidate := range candidates { + if !shouldEmitCopilot(candidate, sets) { + continue + } + entry := usageprovider.BaseEntry( + usage.ProviderCopilot, + candidate.timestamp, + "copilot", + "GitHub Copilot CLI", + candidate.sessionID, + candidate.model, + "GitHub Copilot CLI", + candidate.tokens, + ) + 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 []agentdata.LineJSON) map[string]traceContext { + contexts := make(map[string]traceContext) + for _, line := range lines { + record := line.Value + traceID := traceID(record) + if traceID == "" { + continue + } + attrs := agentdata.ObjectAt(record["attributes"]) + if attrs == nil { + continue + } + context := contexts[traceID] + if context.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 + context.sessionIDPriority = priority + } + contexts[traceID] = context + } + return contexts +} + +func recordCandidate(path string, line agentdata.LineJSON, index int, fallback time.Time, contexts map[string]traceContext) (candidate, bool) { + record := line.Value + attrs := agentdata.ObjectAt(record["attributes"]) + if attrs == nil { + return candidate{}, false + } + source, ok := recordSource(record, attrs) + if !ok { + return candidate{}, false + } + 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 { + input = 0 + } + tokens := usage.TokenUsage{ + InputTokens: input, + 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: agentdata.UintField(attrs, "gen_ai.usage.reasoning.output_tokens", "gen_ai.usage.reasoning_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 := agentdata.FirstStringField(attrs, "gen_ai.response.model", "gen_ai.request.model") + if model == "" { + model = context.model + } + if model == "" { + model = "unknown" + } + sessionID, _ := bestSession(attrs) + if sessionID == "" { + sessionID = context.sessionID + } + if sessionID == "" { + sessionID = traceID + } + if sessionID == "" { + sessionID = "unknown-session" + } + timestamp, ok := timestamp(record) + if !ok { + timestamp = fallback + } + responseID := agentdata.StringField(attrs, "gen_ai.response.id") + return candidate{ + source: source, + traceID: traceID, + responseID: responseID, + sessionID: sessionID, + model: model, + timestamp: timestamp, + tokens: tokens, + dedupKey: dedupKey(source, record, attrs, traceID, sessionID, timestamp, index), + sourceFile: path, + sourceLine: line.Line, + sourceStart: line.Start, + sourceEnd: line.End, + }, true +} + +func recordSource(record, attrs map[string]any) (sourceKind, bool) { + switch { + case isChatSpan(record, attrs): + return copilotChatSpan, true + case isInferenceLog(record, attrs): + return copilotInferenceLog, true + case isAgentTurnLog(record, attrs): + return copilotAgentTurnLog, true + case isAgentSummarySpan(record, attrs): + return copilotAgentSummarySpan, true + default: + return 0, false + } +} + +func isSpan(record map[string]any) bool { + if agentdata.StringField(record, "type") == "span" { + return true + } + if agentdata.StringField(record, "name") == "" { + return false + } + return agentdata.StringField(record, "spanId") != "" || + agentdata.StringField(record, "traceId") != "" || + record["startTime"] != nil || + record["endTime"] != nil || + record["duration"] != nil || + record["kind"] != nil +} + +func isChatSpan(record, attrs map[string]any) bool { + return isSpan(record) && + (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) && + (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) && + (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) && + (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 := agentdata.StringField(record, "body"); body != "" { + return body + } + return agentdata.StringField(record, "_body") +} + +func traceID(record map[string]any) string { + if traceID := agentdata.StringField(record, "traceId"); traceID != "" { + return traceID + } + return agentdata.StringField(agentdata.ObjectAt(record["spanContext"]), "traceId") +} + +func spanID(record map[string]any) string { + if spanID := agentdata.StringField(record, "spanId"); spanID != "" { + return spanID + } + return agentdata.StringField(agentdata.ObjectAt(record["spanContext"]), "spanId") +} + +func bestSession(attrs map[string]any) (string, int) { + candidates := []struct { + key string + priority int + }{ + {"gen_ai.conversation.id", 3}, + {"copilot_chat.session_id", 3}, + {"copilot_chat.chat_session_id", 3}, + {"session.id", 3}, + {"github.copilot.interaction_id", 2}, + {"gen_ai.response.id", 1}, + } + bestValue := "" + bestPriority := 0 + for _, candidate := range candidates { + if value := agentdata.StringField(attrs, candidate.key); value != "" && candidate.priority > bestPriority { + bestValue = value + bestPriority = candidate.priority + } + } + return bestValue, bestPriority +} + +func timestamp(record map[string]any) (time.Time, bool) { + for _, key := range []string{"endTime", "startTime", "hrTime", "_hrTime", "time"} { + if timestamp, ok := agentdata.TimestampFromParts(record[key]); ok { + return timestamp, true + } + } + for _, key := range []string{"timestamp", "observedTimestamp"} { + if timestamp, ok := agentdata.ParseTimestamp(record[key]); ok { + return timestamp, true + } + } + if raw := agentdata.UintValue(record["timeUnixNano"]); raw > 0 { + return time.UnixMilli(int64(raw / 1_000_000)), true + } + return time.Time{}, false +} + +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 != "" { + return traceID + ":" + spanID + } + return fmt.Sprintf("span:%s:%d:%d", sessionID, timestamp.UnixMilli(), index) + case copilotInferenceLog: + if traceID != "" && spanID != "" { + return "log:" + traceID + ":" + spanID + } + return fmt.Sprintf("log:%s:%d:%d", sessionID, timestamp.UnixMilli(), index) + case copilotAgentTurnLog: + turnIndex := agentdata.UintField(attrs, "turn.index", "copilot_chat.turn.index") + turn := fmt.Sprintf("idx-%d", index) + if turnIndex > 0 { + turn = fmt.Sprintf("%d", turnIndex) + } + if traceID != "" { + return "agent-turn:" + traceID + ":" + turn + } + return "agent-turn:" + sessionID + ":" + turn + fmt.Sprintf(":%d", index) + default: + return fmt.Sprintf("%s:%d", filepath.Base(sourceName(source)), index) + } +} + +func sourceName(source sourceKind) string { + switch source { + case copilotChatSpan: + return "chat" + case copilotInferenceLog: + return "inference" + case copilotAgentTurnLog: + return "agent-turn" + case copilotAgentSummarySpan: + return "agent-summary" + default: + return "unknown" + } +} + +type candidateSetMap struct { + chatTraces map[string]bool + inferenceTraces map[string]bool + agentTurnTraces map[string]bool + chatResponses map[string]bool + inferenceResponses map[string]bool + agentTurnResponses map[string]bool +} + +func candidateSets(candidates []candidate) candidateSetMap { + sets := candidateSetMap{ + chatTraces: make(map[string]bool), + inferenceTraces: make(map[string]bool), + agentTurnTraces: make(map[string]bool), + chatResponses: make(map[string]bool), + inferenceResponses: make(map[string]bool), + agentTurnResponses: make(map[string]bool), + } + for _, candidate := range candidates { + if candidate.traceID != "" { + switch candidate.source { + case copilotChatSpan: + sets.chatTraces[candidate.traceID] = true + case copilotInferenceLog: + sets.inferenceTraces[candidate.traceID] = true + case copilotAgentTurnLog: + sets.agentTurnTraces[candidate.traceID] = true + } + } + if candidate.responseID != "" { + switch candidate.source { + case copilotChatSpan: + sets.chatResponses[candidate.responseID] = true + case copilotInferenceLog: + sets.inferenceResponses[candidate.responseID] = true + case copilotAgentTurnLog: + sets.agentTurnResponses[candidate.responseID] = true + } + } + } + return sets +} + +func shouldEmitCopilot(candidate candidate, sets candidateSetMap) bool { + traceMatch := func(values map[string]bool) bool { + return candidate.traceID != "" && values[candidate.traceID] + } + responseMatch := func(values map[string]bool) bool { + return candidate.responseID != "" && values[candidate.responseID] + } + switch candidate.source { + case copilotChatSpan: + return true + case copilotInferenceLog: + return !traceMatch(sets.chatTraces) && !responseMatch(sets.chatResponses) + case copilotAgentTurnLog: + return !traceMatch(sets.chatTraces) && + !traceMatch(sets.inferenceTraces) && + !responseMatch(sets.chatResponses) && + !responseMatch(sets.inferenceResponses) + case copilotAgentSummarySpan: + return !traceMatch(sets.chatTraces) && + !traceMatch(sets.inferenceTraces) && + !traceMatch(sets.agentTurnTraces) && + !responseMatch(sets.chatResponses) && + !responseMatch(sets.inferenceResponses) && + !responseMatch(sets.agentTurnResponses) + default: + return false + } +} diff --git a/internal/provider/copilot/provider.go b/internal/provider/copilot/provider.go new file mode 100644 index 0000000..8b1d6bf --- /dev/null +++ b/internal/provider/copilot/provider.go @@ -0,0 +1,36 @@ +package copilot + +import ( + "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{ 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 = usageprovider.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 usageprovider.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..71eacc3 --- /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/providertest" + "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") + 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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/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/agentusage/droid.go b/internal/provider/droid/loader.go similarity index 62% rename from internal/agentusage/droid.go rename to internal/provider/droid/loader.go index fe6a4a7..99d8a57 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,23 @@ import ( "sort" "strings" + "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 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, agentdata.CollectFiles(root, isSettingsFile)...) } sort.Strings(files) - files = filterFiles(uniqueStrings(files), filter) + files = agentdata.FilterFiles(agentdata.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 +31,58 @@ func loadDroidEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, e entries = append(entries, entry) } } - sortEntries(entries) + usageprovider.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 := agentdata.ReadJSONObject(path) if err != nil || settings == nil { return usage.Entry{}, false, err } - usageBlock := objectAt(settings["tokenUsage"]) + usageBlock := agentdata.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: 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(stringField(settings, "providerLock")) - model := normalizeDroidModel(stringField(settings, "model")) + provider := normalizeDroidProvider(agentdata.StringField(settings, "providerLock")) + model := normalizeDroidModel(agentdata.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 := agentdata.ParseTimestamp(settings["providerLockTimestamp"]) if !ok { - timestamp = fileModifiedTime(path) + timestamp = agentdata.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 := 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 } @@ -141,7 +144,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 +159,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 +176,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..19bd22a --- /dev/null +++ b/internal/provider/droid/provider.go @@ -0,0 +1,36 @@ +package droid + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Droid usage entries. +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 = usageprovider.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 usageprovider.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..1a1615d --- /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/providertest" + "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") + 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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/provider/gemini/loader.go b/internal/provider/gemini/loader.go new file mode 100644 index 0000000..5bd6459 --- /dev/null +++ b/internal/provider/gemini/loader.go @@ -0,0 +1,531 @@ +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" +) + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + files := make([]string, 0) + for _, root := range paths { + files = append(files, agentdata.CollectExt(root, ".json")...) + files = append(files, agentdata.CollectExt(root, ".jsonl")...) + } + 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, resolver) + } else { + fileEntries, err = parseJSONFile(file, resolver) + } + if err != nil { + return nil, err + } + entries = append(entries, fileEntries...) + } + usageprovider.SortEntries(entries) + return entries, nil +} + +type tokens struct { + input uint64 + output uint64 + cached uint64 + thoughts uint64 + tool uint64 + total uint64 + hasTotal bool +} + +// 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 { + for index, raw := range messages { + log.upsert(agentdata.ObjectAt(raw), index+1) + } + return messageEntries(log, path, "", identity, sessionID, sessionTimestamp), 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"), identity, sessionID, sessionTimestamp), nil +} + +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) + for _, line := range lines { + record := line.Value + if value := agentdata.FirstStringField(record, "sessionId", "session_id"); value != "" { + sessionID = value + } + if model := agentdata.StringField(record, "model"); model != "" { + currentModel = model + } + // "$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 + } + for _, raw := range agentdata.ArrayAt(set["messages"]) { + log.upsert(agentdata.ObjectAt(raw), line.Line) + } + 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, identity, sessionID, timestamp)...) + } + } + entries = append(entries, messageEntries(log, path, currentModel, identity, sessionID, fallback)...) + return entries, nil +} + +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 + } + model := agentdata.StringField(record, "model") + if model == "" { + model = modelHint + } + timestamp := firstTimestamp(record, fallback, "timestamp", "created_at") + 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 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 + } + if models := agentdata.ObjectAt(stats["models"]); models != nil { + entries := make([]usage.Entry, 0) + for model, raw := range models { + data := agentdata.ObjectAt(raw) + tokens, ok := parseTokens(data["tokens"]) + if !ok { + continue + } + if entry, ok := buildGeminiEntry(path, line, model, identity, sessionID, timestamp, tokens, false, ""); ok { + entries = append(entries, entry) + } + } + if len(entries) > 0 { + return entries + } + } + tokens, ok := parseTokens(stats) + if !ok { + return nil + } + model := modelHint + if model == "" { + model = "unknown" + } + 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 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 + } + input, cacheRead := normalizeGeminiInput(tokens, direct) + tokenUsage := usage.TokenUsage{ + InputTokens: input + tokens.tool, + OutputTokens: tokens.output, + CacheReadInputTokens: cacheRead, + ReasoningOutputTokens: tokens.thoughts, + } + if tokens.hasTotal { + tokenUsage = usageprovider.ApplyTotalFallback(tokenUsage, tokens.total) + } else if tokenUsage.TotalTokens == 0 { + tokenUsage.TotalTokens = usageprovider.TotalUsage(tokenUsage) + } + if !usageprovider.NonZero(tokenUsage) { + return usage.Entry{}, false + } + entry := usageprovider.BaseEntry(usage.ProviderGemini, timestamp, identity.project, identity.projectPath, sessionID, model, "Gemini CLI", tokenUsage) + usageprovider.SetSource(&entry, path, line, 0, 0) + // 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 +} + +func parseTokens(raw any) (tokens, bool) { + record := agentdata.ObjectAt(raw) + if record == nil { + return tokens{}, false + } + tokens := 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 +} + +func normalizeGeminiInput(tokens tokens, direct bool) (uint64, uint64) { + if !direct { + cachedPortion := tokens.input + if tokens.cached < cachedPortion { + cachedPortion = tokens.cached + } + return tokens.input - cachedPortion, tokens.cached + } + inclusiveTotal := tokens.input + tokens.output + tokens.thoughts + tokens.tool + exclusiveTotal := inclusiveTotal + tokens.cached + if tokens.cached > 0 && tokens.hasTotal && tokens.total == inclusiveTotal && tokens.total != exclusiveTotal { + cachedPortion := tokens.input + if tokens.cached < cachedPortion { + cachedPortion = tokens.cached + } + return tokens.input - cachedPortion, tokens.cached + } + return tokens.input, tokens.cached +} + +func recordStats(record map[string]any) map[string]any { + if stats := agentdata.ObjectAt(record["stats"]); stats != nil { + return 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 := agentdata.ParseTimestamp(record[key]); ok { + return timestamp + } + } + return fallback +} 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.go b/internal/provider/gemini/provider.go new file mode 100644 index 0000000..c5ecb3c --- /dev/null +++ b/internal/provider/gemini/provider.go @@ -0,0 +1,36 @@ +package gemini + +import ( + "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{ 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 = usageprovider.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 usageprovider.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..a94a9cf --- /dev/null +++ b/internal/provider/gemini/provider_test.go @@ -0,0 +1,38 @@ +package gemini + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +// 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() + 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.WantEntry{ + Provider: usage.ProviderGemini, + Model: "gemini-3-flash-preview", + SessionID: "session-a", + Project: "shop", + 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 50% rename from internal/agentusage/goose.go rename to internal/provider/goose/loader.go index b36e054..10fef6a 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,19 @@ import ( "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/usageprovider" + "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 +31,14 @@ func loadGooseEntries(paths []string) ([]usage.Entry, error) { entries = append(entries, entry) } } - sortEntries(entries) + usageprovider.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 agentdb.ExistingSQLiteFile(root) { dbPaths = append(dbPaths, root) continue } @@ -43,17 +47,17 @@ func gooseDBPaths(paths []string) []string { filepath.Join(root, "sessions", "sessions.db"), filepath.Join(root, "data", "sessions", "sessions.db"), } { - if existingSQLiteFile(candidate) { + if agentdb.ExistingSQLiteFile(candidate) { dbPaths = append(dbPaths, candidate) } } } sort.Strings(dbPaths) - return uniqueStrings(dbPaths) + return agentdata.UniqueStrings(dbPaths) } -func loadGooseDatabase(path string) ([]usage.Entry, error) { - db, err := openSQLite(path) +func loadDatabase(path string) ([]usage.Entry, error) { + db, err := agentdb.OpenSQLite(path) if err != nil { return nil, err } @@ -74,29 +78,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 !agentdb.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 := agentdb.SqlString(id) + model := modelName(agentdb.SqlString(modelConfig)) if sessionID == "" || model == "" { return usage.Entry{}, false } - timestamp, ok := gooseTimestamp(sqlString(createdAt)) + timestamp, ok := timestamp(agentdb.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(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, @@ -104,34 +108,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 = usageprovider.TotalUsage(tokens) + if !usageprovider.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 := 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 gooseModelName(config string) string { - record := decodeJSONObjectString(config) - return stringField(record, "model_name") +func modelName(config string) string { + record := agentdata.DecodeJSONObjectString(config) + return agentdata.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 := agentdata.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 agentdata.ParseTimestampString(value[:10] + "T" + value[11:] + "Z") } if len(value) == 10 && value[4] == '-' && value[7] == '-' { - return 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 new file mode 100644 index 0000000..4faab0f --- /dev/null +++ b/internal/provider/goose/provider.go @@ -0,0 +1,29 @@ +package goose + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Goose usage entries. +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 = usageprovider.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 usageprovider.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..7be40f0 --- /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/providertest" + "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 := providertest.OpenTestSQLite(t, dbPath) + defer db.Close() + providertest.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 + )`) + 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 (?, ?, ?, ?, ?, ?, ?)`, + "session-a", `{"model_name":"claude-sonnet-4-20250514"}`, "anthropic", "2026-05-01 01:02:03", 180, 100, 50, + ) + return Provider{}.WithPaths([]string{dbPath}).Entries() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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..a80725e --- /dev/null +++ b/internal/provider/hermes/loader.go @@ -0,0 +1,101 @@ +package hermes + +import ( + "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/usageprovider" + + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +func loadEntries(paths []string) ([]usage.Entry, error) { + dbPaths := agentdb.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...) + } + usageprovider.SortEntries(entries) + return entries, nil +} + +func loadDatabase(path string) ([]usage.Entry, error) { + db, err := agentdb.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 !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) + 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 := agentdb.SqlString(sessionRaw) + model := strings.TrimSpace(agentdb.SqlString(modelRaw)) + if sessionID == "" || model == "" { + return usage.Entry{}, false + } + timestamp, ok := timestamp(startedAt) + if !ok { + return usage.Entry{}, false + } + tokens := usage.TokenUsage{ + 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 = usageprovider.TotalUsage(tokens) + } + if !usageprovider.NonZero(tokens) { + return usage.Entry{}, false + } + 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 := agentdata.ParseTimestamp(value); ok { + return parsed, true + } + if number, ok := agentdb.SqlFloat(value); ok { + return agentdata.TimestampFromFloat(number) + } + 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 new file mode 100644 index 0000000..4686392 --- /dev/null +++ b/internal/provider/hermes/provider.go @@ -0,0 +1,29 @@ +package hermes + +import ( + "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{ 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 = usageprovider.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 usageprovider.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..9e33013 --- /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/providertest" + "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 := providertest.OpenTestSQLite(t, dbPath) + defer db.Close() + providertest.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 + )`) + 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 + ) 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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..5423eb9 --- /dev/null +++ b/internal/provider/kilo/loader.go @@ -0,0 +1,377 @@ +package kilo + +import ( + "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) { + 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 + } + 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 +} + +func loadDatabase(path string) ([]usage.Entry, error) { + db, err := agentdb.OpenSQLite(path) + if err != nil { + return nil, err + } + defer db.Close() + + 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() + + 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 := 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 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) + } + + rows, err := db.Query(`SELECT message_id, CAST(data AS TEXT) FROM part ORDER BY time_created ASC, id ASC`) + if err != nil { + return + } + 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"), + } + 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 + } + + timestamp := time.UnixMilli(message.created).UTC() + 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 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)...) + } + 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.go b/internal/provider/kilo/provider.go new file mode 100644 index 0000000..c7c3c4b --- /dev/null +++ b/internal/provider/kilo/provider.go @@ -0,0 +1,29 @@ +package kilo + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Kilo usage entries. +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 = usageprovider.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 usageprovider.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..b8af5cd --- /dev/null +++ b/internal/provider/kilo/provider_test.go @@ -0,0 +1,39 @@ +package kilo + +import ( + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" + "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) { + 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: "kilo-auto/free", + SessionID: "session-a", + Project: "demo", + 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..f00fc96 --- /dev/null +++ b/internal/provider/kimi/loader.go @@ -0,0 +1,205 @@ +package kimi + +import ( + "path/filepath" + "sort" + "strings" + "time" + + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" +) + +const defaultModel = "kimi-for-coding" + +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) + 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 := parseWireFileFrom(file, 0) + 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 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 parseWireFileFrom(path string, start int64) ([]usage.Entry, int64, error) { + lines, consumed, err := agentdata.ReadJSONLinesFrom(path, start, `usage`) + if err != nil { + return nil, 0, err + } + configModel := configModel(path) + sessionID := sessionID(path) + fallback := agentdata.FileModifiedTime(path) + entries := make([]usage.Entry, 0) + for _, line := range lines { + var record record + var ok bool + if agentdata.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 := 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, consumed, nil +} + +func parseStatusUpdate(value map[string]any) (record, bool) { + message := agentdata.ObjectAt(value["message"]) + if agentdata.StringField(message, "type") != "StatusUpdate" { + return record{}, false + } + payload := agentdata.ObjectAt(message["payload"]) + tokenUsage := agentdata.ObjectAt(payload["token_usage"]) + if tokenUsage == nil { + return record{}, false + } + tokens := usage.TokenUsage{ + 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 = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(tokenUsage, "total")) + if !usageprovider.NonZero(tokens) { + return record{}, false + } + 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 agentdata.StringField(value, "usageScope") != "turn" { + return record{}, false + } + tokenUsage := agentdata.ObjectAt(value["usage"]) + if tokenUsage == nil { + return record{}, false + } + tokens := usage.TokenUsage{ + InputTokens: agentdata.UintField(tokenUsage, "inputOther"), + OutputTokens: agentdata.UintField(tokenUsage, "output"), + CacheCreationInputTokens: agentdata.UintField(tokenUsage, "inputCacheCreation"), + CacheReadInputTokens: agentdata.UintField(tokenUsage, "inputCacheRead"), + } + tokens = usageprovider.ApplyTotalFallback(tokens, 0) + if !usageprovider.NonZero(tokens) { + return record{}, false + } + record := record{ + tokens: tokens, + model: strings.TrimPrefix(agentdata.StringField(value, "model"), "kimi-code/"), + } + record.timestamp, record.hasTime = agentdata.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 := agentdata.ReadJSONObject(filepath.Join(root, "config.json")) + if err != nil || config == nil { + return defaultModel + } + 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 new file mode 100644 index 0000000..cb2ec7e --- /dev/null +++ b/internal/provider/kimi/provider.go @@ -0,0 +1,46 @@ +package kimi + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Kimi usage entries. +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 = usageprovider.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 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/kimi/provider_test.go b/internal/provider/kimi/provider_test.go new file mode 100644 index 0000000..017c923 --- /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/providertest" + "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() + providertest.WriteFile(t, filepath.Join(dir, "config.json"), `{"model":"kimi-k2"}`) + path := filepath.Join(dir, "sessions", "group", "session-a", "wire.jsonl") + 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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..48c65ff --- /dev/null +++ b/internal/provider/openclaw/loader.go @@ -0,0 +1,139 @@ +package openclaw + +import ( + "path/filepath" + "sort" + "strings" + + "github.com/tokitoki-dev/tokitoki-cli/internal/agentdata" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" + + "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 { + files = append(files, agentdata.CollectFiles(path, isSessionFile)...) + } + sort.Strings(files) + files = agentdata.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...) + } + usageprovider.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 := agentdata.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 := agentdata.ObjectAt(record["data"]) + if source == nil { + source = record + } + if model := agentdata.FirstStringField(source, "modelId", "model"); model != "" { + currentModel = model + } + if provider := agentdata.StringField(source, "provider"); provider != "" { + currentProvider = provider + } + continue + } + if agentdata.StringField(record, "type") != "message" { + continue + } + message := agentdata.ObjectAt(record["message"]) + if agentdata.StringField(message, "role") != "assistant" { + continue + } + usageBlock := agentdata.ObjectAt(message["usage"]) + if usageBlock == nil { + continue + } + timestamp, ok := agentdata.ParseTimestamp(message["timestamp"]) + if !ok { + timestamp, ok = agentdata.ParseTimestamp(record["timestamp"]) + } + if !ok { + timestamp = agentdata.FileModifiedTime(path) + } + model := agentdata.FirstStringField(message, "modelId", "model") + if model == "" { + model = currentModel + } + if model == "" { + model = "unknown" + } + provider := agentdata.StringField(message, "provider") + if provider == "" { + provider = currentProvider + } + tokens := usage.TokenUsage{ + InputTokens: agentdata.UintField(usageBlock, "input"), + OutputTokens: agentdata.UintField(usageBlock, "output"), + CacheCreationInputTokens: agentdata.UintField(usageBlock, "cacheWrite"), + CacheReadInputTokens: agentdata.UintField(usageBlock, "cacheRead"), + } + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(usageBlock, "totalTokens")) + if !usageprovider.NonZero(tokens) { + continue + } + 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 agentdata.StringField(record, "type") == "model_change" { + return true + } + return agentdata.StringField(record, "type") == "custom" && agentdata.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..f64fe06 --- /dev/null +++ b/internal/provider/openclaw/provider.go @@ -0,0 +1,36 @@ +package openclaw + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads OpenClaw usage entries. +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 = usageprovider.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 usageprovider.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..4890459 --- /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/providertest" + "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") + 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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/provider/opencode/loader.go b/internal/provider/opencode/loader.go new file mode 100644 index 0000000..326881f --- /dev/null +++ b/internal/provider/opencode/loader.go @@ -0,0 +1,675 @@ +package opencode + +import ( + "bytes" + "database/sql" + "encoding/json" + "os" + "path/filepath" + "sort" + "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" +) + +// 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 session struct { + directory string + version string + model string +} + +type message struct { + id string + sessionID string + created int64 + data map[string]any + parts []map[string]any +} + +func loadEntries(paths []string, filter usage.FileFilter) ([]usage.Entry, error) { + entriesByID := make(map[string]usage.Entry) + for _, root := range paths { + rootEntries, err := loadRoot(root, filter) + if err != nil { + return nil, err + } + for _, entry := range rootEntries { + if entry.ID == "" { + entry.ID = usageprovider.StableEntryID(entry) + } + 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 +} + +func loadRoot(root string, filter usage.FileFilter) ([]usage.Entry, error) { + info, err := os.Stat(root) + if err != nil { + return nil, nil + } + if !info.IsDir() { + switch strings.ToLower(filepath.Ext(root)) { + case ".db": + return loadDatabase(root) + case ".json": + return loadLegacyFile(root) + } + return nil, nil + } + + entries := make([]usage.Entry, 0) + seenIDs := make(map[string]bool) + for _, dbPath := range dbPaths(root) { + dbEntries, err := loadDatabase(dbPath) + if err != nil { + return nil, err + } + for _, entry := range dbEntries { + if entry.ID != "" { + seenIDs[entry.ID] = true + } + entries = append(entries, entry) + } + } + + // Older OpenCode releases wrote one JSON file per message instead of the + // database. Their users keep those files, so keep reading them. + files := agentdata.CollectExt(filepath.Join(root, "storage", "message"), ".json") + sort.Strings(files) + files = agentdata.FilterFiles(files, filter) + for _, file := range files { + stem := strings.TrimSuffix(filepath.Base(file), filepath.Ext(file)) + if seenIDs[stableMessageID(stem)] { + continue + } + legacy, err := loadLegacyFile(file) + if err != nil { + return nil, err + } + for _, entry := range legacy { + if entry.ID != "" && seenIDs[entry.ID] { + continue + } + if entry.ID != "" { + seenIDs[entry.ID] = true + } + entries = append(entries, entry) + } + } + usageprovider.SortEntries(entries) + return entries, nil +} + +// 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 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 agentdb.ExistingSQLiteFile(candidate) { + paths = append(paths, candidate) + } + } + sort.Strings(paths) + return paths +} + +func loadDatabase(path string) ([]usage.Entry, error) { + db, err := agentdb.OpenSQLite(path) + if err != nil { + return nil, err + } + defer db.Close() + + sessions := querySessions(db) + messages := queryMessages(db) + if len(messages) == 0 { + return nil, nil + } + 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][]*message) + for _, message := range messages { + bySession[message.sessionID] = append(bySession[message.sessionID], message) + } + + 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, sessionEntries(path, sessions[sessionID], sessionMessages)...) + } + return entries, nil +} + +func querySessions(db *sql.DB) map[string]session { + sessions := make(map[string]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, model string + if err := rows.Scan(&id, &directory, &version, &model); err != nil { + continue + } + // 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 +} + +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([]*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 := 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 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) + } + + rows, err := db.Query(`SELECT message_id, CAST(data AS TEXT) FROM part ORDER BY time_created ASC, id ASC`) + if err != nil { + return + } + 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) + } + } +} + +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 := tokens(message.data) + if !usageprovider.NonZero(reported) { + continue + } + billed, ok := billTokens(reported, previous) + previous = reported + if !ok { + continue + } + entry, ok := newEntry(source, session, message, billed) + if !ok { + continue + } + entries = append(entries, entry) + } + return entries +} + +// 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. billTokens is what +// turns this mixture into one turn's cost. +func tokens(record map[string]any) usage.TokenUsage { + block := agentdata.ObjectAt(record["tokens"]) + if block == nil { + return usage.TokenUsage{} + } + cache := agentdata.ObjectAt(block["cache"]) + return 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"), + } +} + +// 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 billTokens(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 !usageprovider.NonZero(billed) { + return usage.TokenUsage{}, false + } + return usageprovider.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 +} + +// 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 := modelName(record, session) + + timestamp := time.UnixMilli(message.created).UTC() + 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 := 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.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 = 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 +// carries no line counts, so treating it as agent work invents writes. +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"]) + + switch agentdata.StringField(part, "tool") { + case "edit": + return editChanges(input, metadata, cwd) + case "write": + return writeChanges(input, metadata, cwd) + case "patch", "apply_patch": + return applyPatchChanges(input, metadata, cwd) + default: + return nil + } +} + +// 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, agentdata.StringField(input, "filePath")) + + 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 { + return []usage.FileChange{{Path: path, LinesAdded: added, LinesRemoved: removed}} + } + } + + added := usage.CountLines(agentdata.StringField(input, "newString")) + removed := usage.CountLines(agentdata.StringField(input, "oldString")) + if path == "" && added == 0 && removed == 0 { + return nil + } + return []usage.FileChange{{Path: path, LinesAdded: added, LinesRemoved: removed}} +} + +// 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 writeChanges(input, metadata map[string]any, cwd string) []usage.FileChange { + path := usage.ResolvePath(cwd, agentdata.FirstNonEmpty( + agentdata.StringField(input, "filePath"), + agentdata.StringField(metadata, "filepath"), + )) + if path == "" { + return nil + } + + // 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 +// falls back to counting the raw patch envelope itself. +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 := agentdata.ObjectAt(raw) + if file == nil { + continue + } + 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: agentdata.UintField(file, "additions"), + LinesRemoved: agentdata.UintField(file, "deletions"), + }) + } + if len(changes) > 0 { + return changes + } + } + + patchText := agentdata.FirstNonEmpty( + agentdata.StringField(input, "patchText"), + agentdata.StringField(input, "patch"), + agentdata.StringField(metadata, "patch"), + ) + return parsePatchEnvelope(patchText, cwd) +} + +// parsePatchEnvelope counts the added and removed lines of each file in +// an apply_patch envelope. One envelope can carry any number of files. +func parsePatchEnvelope(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, patchHeaderPath(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 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) + } + } + return "" +} + +// loadLegacyFile reads one message from the pre-database storage layout, +// where the message and its parts each lived in their own JSON file. +func loadLegacyFile(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 := &message{ + id: agentdata.StringField(record, "id"), + sessionID: agentdata.StringField(record, "sessionID"), + data: record, + 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. + tokens, ok := billTokens(tokens(record), usage.TokenUsage{}) + if !ok { + return nil, nil + } + entry, ok := newEntry(path, session{}, message, tokens) + if !ok { + return nil, nil + } + entry.SourceLine = 1 + return []usage.Entry{entry}, nil +} + +// readLegacyParts loads the part files stored alongside a legacy +// message: storage/message//.json pairs with +// storage/part//*.json. +func readLegacyParts(messagePath, messageID string) []map[string]any { + if messageID == "" { + return nil + } + storage := filepath.Dir(filepath.Dir(filepath.Dir(messagePath))) + files := agentdata.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 +} + +// 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 stableMessageID(messageID string) string { + if strings.TrimSpace(messageID) == "" { + return "" + } + return usage.StableID("opencode", strings.TrimSpace(messageID)) +} diff --git a/internal/provider/opencode/loader_test.go b/internal/provider/opencode/loader_test.go new file mode 100644 index 0000000..a8acee6 --- /dev/null +++ b/internal/provider/opencode/loader_test.go @@ -0,0 +1,496 @@ +package opencode + +import ( + "path/filepath" + "testing" + + "github.com/tokitoki-dev/tokitoki-cli/internal/providertest" + "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 +// 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 := providertest.OpenTestSQLite(t, path) + defer db.Close() + + schema := []string{ + `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)`, + } + 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 +} + +// 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 := loadDatabase(path) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + t.Fatalf("len(entries) = %d, want 2", len(entries)) + } + usageprovider.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 := 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)) + } + 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) + } +} + +// 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 := loadDatabase(path) + if err != nil { + t.Fatal(err) + } + usageprovider.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 := 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].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 := loadDatabase(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 := loadDatabase(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 := loadDatabase(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 := loadDatabase(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 := loadDatabase(first) + if err != nil { + t.Fatal(err) + } + 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 != stableMessageID("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 := loadRoot(root, nil) + if err != nil { + t.Fatal(err) + } + if len(entries) != 2 { + 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 := 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)`, + `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/provider/opencode/provider.go b/internal/provider/opencode/provider.go new file mode 100644 index 0000000..31dd4fc --- /dev/null +++ b/internal/provider/opencode/provider.go @@ -0,0 +1,36 @@ +package opencode + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads OpenCode usage entries. +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 = usageprovider.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 usageprovider.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..7716a5c --- /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/providertest" + "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") + 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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..3b1993e --- /dev/null +++ b/internal/provider/pi/loader.go @@ -0,0 +1,103 @@ +package pi + +import ( + "path/filepath" + "sort" + "strings" + + "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 sessionFiles(paths []string) []string { + files := make([]string, 0) + for _, path := range paths { + files = append(files, agentdata.CollectExt(path, ".jsonl")...) + } + sort.Strings(files) + 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 := parseSessionFileFrom(file, 0) + if err != nil { + return nil, err + } + entries = append(entries, fileEntries...) + } + usageprovider.SortEntries(entries) + return entries, nil +} + +func parseSessionFileFrom(path string, start int64) ([]usage.Entry, int64, error) { + lines, consumed, err := agentdata.ReadJSONLinesFrom(path, start, `"usage"`, `"message"`) + if err != nil { + return nil, 0, err + } + project := project(path) + sessionID := sessionID(path) + entries := make([]usage.Entry, 0) + for _, line := range lines { + if typ := agentdata.StringField(line.Value, "type"); typ != "" && typ != "message" { + continue + } + message := agentdata.ObjectAt(line.Value["message"]) + if agentdata.StringField(message, "role") != "assistant" { + continue + } + usageBlock := agentdata.ObjectAt(message["usage"]) + if usageBlock == nil { + continue + } + timestamp, ok := agentdata.ParseTimestamp(line.Value["timestamp"]) + if !ok { + continue + } + tokens := usage.TokenUsage{ + InputTokens: agentdata.UintField(usageBlock, "input"), + OutputTokens: agentdata.UintField(usageBlock, "output"), + CacheCreationInputTokens: agentdata.UintField(usageBlock, "cacheWrite"), + CacheReadInputTokens: agentdata.UintField(usageBlock, "cacheRead"), + } + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(usageBlock, "totalTokens")) + if !usageprovider.NonZero(tokens) { + continue + } + model := agentdata.StringField(message, "model") + if model != "" { + model = "[pi] " + model + } + 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, consumed, 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..c0f0fe5 --- /dev/null +++ b/internal/provider/pi/provider.go @@ -0,0 +1,43 @@ +package pi + +import ( + "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{ 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 = usageprovider.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 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/pi/provider_test.go b/internal/provider/pi/provider_test.go new file mode 100644 index 0000000..72993e1 --- /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/providertest" + "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") + 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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..2b1ea1d --- /dev/null +++ b/internal/provider/qwen/loader.go @@ -0,0 +1,111 @@ +package qwen + +import ( + "path/filepath" + "sort" + "strings" + + "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 chatFiles(paths []string) []string { + files := make([]string, 0) + for _, root := range paths { + files = append(files, agentdata.CollectFiles(filepath.Join(root, "projects"), isChatFile)...) + if strings.Contains(filepath.ToSlash(root), "/projects/") { + files = append(files, agentdata.CollectFiles(root, isChatFile)...) + } + } + sort.Strings(files) + 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 := parseChatFileFrom(file, 0) + if err != nil { + return nil, err + } + entries = append(entries, fileEntries...) + } + usageprovider.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 parseChatFileFrom(path string, start int64) ([]usage.Entry, int64, error) { + lines, consumed, err := agentdata.ReadJSONLinesFrom(path, start, `"usageMetadata"`) + if err != nil { + return nil, 0, err + } + project := project(path) + fallback := agentdata.FileModifiedTime(path) + entries := make([]usage.Entry, 0) + for _, line := range lines { + record := line.Value + if agentdata.StringField(record, "type") != "assistant" { + continue + } + meta := agentdata.ObjectAt(record["usageMetadata"]) + if meta == nil { + continue + } + timestamp, ok := agentdata.ParseTimestamp(record["timestamp"]) + if !ok { + timestamp = fallback + } + sessionID := agentdata.StringField(record, "sessionId") + if sessionID == "" { + sessionID = project + "-" + strings.TrimSuffix(filepath.Base(path), filepath.Ext(path)) + } + model := agentdata.StringField(record, "model") + if model == "" { + model = "unknown" + } + tokens := usage.TokenUsage{ + InputTokens: agentdata.UintField(meta, "promptTokenCount"), + OutputTokens: agentdata.UintField(meta, "candidatesTokenCount"), + CacheReadInputTokens: agentdata.UintField(meta, "cachedContentTokenCount"), + ReasoningOutputTokens: agentdata.UintField(meta, + "thoughtsTokenCount", + ), + } + tokens = usageprovider.ApplyTotalFallback(tokens, agentdata.UintField(meta, "totalTokenCount")) + if !usageprovider.NonZero(tokens) { + continue + } + 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, consumed, 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..5888b04 --- /dev/null +++ b/internal/provider/qwen/provider.go @@ -0,0 +1,43 @@ +package qwen + +import ( + "github.com/tokitoki-dev/tokitoki-cli/internal/usage" + "github.com/tokitoki-dev/tokitoki-cli/internal/usageprovider" +) + +// Provider loads Qwen usage entries. +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 = usageprovider.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 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/provider/qwen/provider_test.go b/internal/provider/qwen/provider_test.go new file mode 100644 index 0000000..13b7fec --- /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/providertest" + "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") + 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() + }() + + providertest.AssertSingleEntry(t, entries, err, providertest.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/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/providertest/providertest.go b/internal/providertest/providertest.go new file mode 100644 index 0000000..35d1747 --- /dev/null +++ b/internal/providertest/providertest.go @@ -0,0 +1,97 @@ +// 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" + "os" + "path/filepath" + "testing" + + "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 +// 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. 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 := sql.Open("sqlite", path) + if err != nil { + t.Fatal(err) + } + if err := db.Ping(); err != nil { + _ = db.Close() + 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/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/usage/usage.go b/internal/usage/usage.go index c7c0fcf..25b4945 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" @@ -11,25 +12,40 @@ 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" +// 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. @@ -45,6 +61,94 @@ 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"` +} + +// 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"` @@ -67,15 +171,23 @@ 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"` - Raw map[string]any `json:"raw,omitempty"` - Usage TokenUsage `json:"usage"` + 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". + 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"` } // 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 4e18f99..3794b24 100644 --- a/internal/usagedb/db.go +++ b/internal/usagedb/db.go @@ -10,8 +10,10 @@ import ( "fmt" "os" "path/filepath" + "strings" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/codex" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" _ "modernc.org/sqlite" ) @@ -32,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 { @@ -66,13 +80,181 @@ 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) } + // 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. 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. +// 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 + if err := db.QueryRow(`PRAGMA user_version`).Scan(&version); err != nil { + return err + } + if version >= eventSchemaVersion { + return nil + } + if version < 2 { + if err := rekeyCodexEvents(db); err != nil { + 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 { + 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 := codex.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 +283,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) @@ -120,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 } @@ -130,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 @@ -151,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) } } @@ -196,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() @@ -232,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 79639e4..2e19212 100644 --- a/internal/usagedb/db_test.go +++ b/internal/usagedb/db_test.go @@ -1,10 +1,12 @@ package usagedb import ( + "database/sql" "path/filepath" "testing" "time" + "github.com/tokitoki-dev/tokitoki-cli/internal/provider/codex" "github.com/tokitoki-dev/tokitoki-cli/internal/usage" ) @@ -209,3 +211,368 @@ func testUsageEntry(id string) usage.Entry { }, } } + +func TestOpenRekeysCodexEventsToSemanticIDs(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 != 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) + } + + // 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) + } +} + +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) + } +} + +// 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 857435a..8183db2 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,171 @@ 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 +} + +// 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 { + 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 +} + +// 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.FormatInt(entry.SourceStart, 10), + 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 + }) +} diff --git a/internal/usagescan/scanner.go b/internal/usagescan/scanner.go index 8f70bc5..e6293a0 100644 --- a/internal/usagescan/scanner.go +++ b/internal/usagescan/scanner.go @@ -9,10 +9,23 @@ 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/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" @@ -42,21 +55,22 @@ 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{}, + workbuddy.Provider{}, } } @@ -114,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 { @@ -122,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 @@ -148,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 @@ -214,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 f93b913..c90cdae 100644 --- a/internal/usageupload/uploader.go +++ b/internal/usageupload/uploader.go @@ -10,8 +10,10 @@ import ( "errors" "fmt" "io" + "log/slog" "net/http" "os" + "path/filepath" "runtime" "strings" "time" @@ -36,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 { @@ -54,12 +65,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 +81,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 { @@ -125,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 } @@ -155,29 +187,88 @@ 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. // 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 +// 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) + } +} + +// 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 { @@ -295,12 +386,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 +424,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..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) { @@ -73,3 +75,104 @@ 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") + } +} + +// 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) + } +} diff --git a/pkg/agentlib/agentlib.go b/pkg/agentlib/agentlib.go index 99f741f..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" @@ -84,6 +85,9 @@ const ( // ProviderGoose identifies Goose usage files. ProviderGoose Provider = "goose" + + // ProviderWorkbuddy identifies WorkBuddy usage files. + ProviderWorkbuddy Provider = "workbuddy" ) var ( @@ -240,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 { @@ -256,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 @@ -272,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() @@ -286,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 @@ -362,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. // @@ -389,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 { @@ -433,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() @@ -468,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")}, @@ -479,11 +625,17 @@ 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"), filepath.Join(home, ".local", "share", "Block", "goose", "sessions", "sessions.db"), }, + ProviderWorkbuddy: {filepath.Join(home, ".workbuddy")}, } } 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) + } +}