From b8187fa9d5588a02f63381abf4bb029d5741c65d Mon Sep 17 00:00:00 2001 From: akarineren Date: Wed, 29 Jul 2026 23:57:14 +0900 Subject: [PATCH 1/2] Refactor API key error handling: introduce sentinel error for missing API key and update related tests; enhance Kimi provider to support new usage record format --- cmd/tokitoki/main.go | 10 ++ cmd/tokitoki/main_test.go | 8 +- internal/agentusage/kimi.go | 156 ++++++++++++++++++++------ internal/agentusage/providers_test.go | 22 ++++ internal/cli/cli.go | 10 +- internal/cli/cli_test.go | 10 +- pkg/agentlib/agentlib.go | 7 +- pkg/agentlib/agentlib_test.go | 3 + 8 files changed, 181 insertions(+), 45 deletions(-) diff --git a/cmd/tokitoki/main.go b/cmd/tokitoki/main.go index 1a90436..4479d4c 100644 --- a/cmd/tokitoki/main.go +++ b/cmd/tokitoki/main.go @@ -5,6 +5,7 @@ package main import ( "context" "encoding/json" + "errors" "flag" "fmt" "io" @@ -596,8 +597,17 @@ func defaultLogger() *slog.Logger { return slog.New(slog.NewTextHandler(os.Stderr, &slog.HandlerOptions{Level: slog.LevelWarn})) } +// exitNoAPIKey marks the one failure a caller can act on: no key is +// configured. Front-ends prompt for a key on this code and treat every other +// non-zero exit as a transient problem to log and retry, instead of guessing +// from the error text. +const exitNoAPIKey = 3 + func fail(logger *slog.Logger, err error) int { logger.Error("tokitoki failed", "error", err) + if errors.Is(err, agentlib.ErrMissingAPIKey) { + return exitNoAPIKey + } return 1 } diff --git a/cmd/tokitoki/main_test.go b/cmd/tokitoki/main_test.go index b70b75e..db86be1 100644 --- a/cmd/tokitoki/main_test.go +++ b/cmd/tokitoki/main_test.go @@ -80,12 +80,14 @@ func TestRunSetKeyRejectsMissingKey(t *testing.T) { } } -func TestRunGetKeyReturnsErrorWhenMissing(t *testing.T) { +// A missing key exits with its own code so front-ends can prompt for one +// without pattern-matching the error text. +func TestRunGetKeyReturnsNoAPIKeyCodeWhenMissing(t *testing.T) { home := t.TempDir() t.Setenv("HOME", home) - if code := run([]string{"get", "key"}); code != 1 { - t.Fatalf("run(get key) = %d, want 1", code) + if code := run([]string{"get", "key"}); code != exitNoAPIKey { + t.Fatalf("run(get key) = %d, want %d", code, exitNoAPIKey) } } diff --git a/internal/agentusage/kimi.go b/internal/agentusage/kimi.go index b24ca7f..4dc6a6e 100644 --- a/internal/agentusage/kimi.go +++ b/internal/agentusage/kimi.go @@ -4,10 +4,13 @@ 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 { @@ -40,67 +43,156 @@ func isKimiWireFile(path string) bool { return false } parts := strings.Split(filepath.ToSlash(filepath.Clean(path)), "/") - for i := 0; i+3 < len(parts); i++ { - if parts[i] == "sessions" && i+3 == len(parts)-1 { + 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, `"StatusUpdate"`, `"token_usage"`) + lines, err := readJSONLines(path, `usage`) if err != nil { return nil, err } - model := kimiModel(path) - sessionID := filepath.Base(filepath.Dir(path)) - if sessionID == "" || sessionID == "." { - sessionID = "unknown" - } + configModel := kimiConfigModel(path) + sessionID := kimiSessionID(path) fallback := fileModifiedTime(path) entries := make([]usage.Entry, 0) for _, line := range lines { - message := objectAt(line.value["message"]) - if stringField(message, "type") != "StatusUpdate" { - continue + var record kimiRecord + var ok bool + if stringField(line.value, "type") == "usage.record" { + record, ok = parseKimiUsageRecord(line.value) + } else { + record, ok = parseKimiStatusUpdate(line.value) } - payload := objectAt(message["payload"]) - tokenUsage := objectAt(payload["token_usage"]) - if tokenUsage == nil { + if !ok { continue } - timestamp, ok := parseTimestamp(line.value["timestamp"]) - if !ok { + timestamp := record.timestamp + if !record.hasTime { timestamp = fallback } - 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) { - continue + model := record.model + if model == "" { + model = configModel } - messageID := stringField(payload, "message_id") - entry := baseEntry(usage.ProviderKimi, timestamp, "kimi", "Kimi", sessionID, model, "Kimi", tokens) + 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, messageID) + entry.ID = stableEntryID(entry, record.messageID) entries = append(entries, entry) } return entries, nil } -func kimiModel(path string) string { - root := filepath.Dir(filepath.Dir(filepath.Dir(filepath.Dir(path)))) +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 "kimi-for-coding" + return kimiDefaultModel } if model := stringField(config, "model"); model != "" { return model } - return "kimi-for-coding" + return kimiDefaultModel } diff --git a/internal/agentusage/providers_test.go b/internal/agentusage/providers_test.go index c3f0a04..d906e79 100644 --- a/internal/agentusage/providers_test.go +++ b/internal/agentusage/providers_test.go @@ -108,6 +108,28 @@ func TestProvidersLoadEntries(t *testing.T) { 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) { diff --git a/internal/cli/cli.go b/internal/cli/cli.go index a1f61b1..02c02ad 100644 --- a/internal/cli/cli.go +++ b/internal/cli/cli.go @@ -16,6 +16,12 @@ import ( "github.com/tokitoki-dev/tokitoki-cli/internal/usageupload" ) +// ErrNoAPIKey reports that no API key is configured. Callers distinguish it +// from every other failure — a missing key is a thing the user fixes, a +// network or server error is not — so it travels as a sentinel rather than as +// text for someone to pattern-match on. +var ErrNoAPIKey = errors.New("API key is not configured in ~/.tokitoki/api_key") + type App struct { Agent *agent.Agent UsageDB *usagedb.DB @@ -37,7 +43,7 @@ func (a *App) GetAPIKey() error { return err } if settings.APIKey == "" { - return errors.New("API key is not configured in ~/.tokitoki/api_key") + return ErrNoAPIKey } _, err = fmt.Fprintf(a.Out, "%s\n", settings.APIKey) return err @@ -85,7 +91,7 @@ func (a *App) Upload(ctx context.Context) error { return err } if settings.APIKey == "" { - return errors.New("API key is required in ~/.tokitoki/api_key") + return ErrNoAPIKey } if err := usageupload.SyncPending(ctx, settings, a.UsageDB); err != nil { return err diff --git a/internal/cli/cli_test.go b/internal/cli/cli_test.go index dd8481b..6ca6a01 100644 --- a/internal/cli/cli_test.go +++ b/internal/cli/cli_test.go @@ -3,9 +3,9 @@ package cli import ( "bytes" "context" + "errors" "io" "log/slog" - "strings" "testing" "github.com/tokitoki-dev/tokitoki-cli/internal/agent" @@ -34,8 +34,8 @@ func TestIngestWorksWithoutAPIKey(t *testing.T) { func TestUploadRequiresAPIKey(t *testing.T) { app := newApp(t) err := app.Upload(context.Background()) - if err == nil || !strings.Contains(err.Error(), "API key is required") { - t.Fatalf("Upload() error = %v, want API key requirement", err) + if !errors.Is(err, ErrNoAPIKey) { + t.Fatalf("Upload() error = %v, want ErrNoAPIKey", err) } } @@ -73,8 +73,8 @@ func TestGetAPIKeyWritesSavedKey(t *testing.T) { func TestGetAPIKeyRequiresConfiguredKey(t *testing.T) { app := newApp(t) err := app.GetAPIKey() - if err == nil || !strings.Contains(err.Error(), "API key is not configured") { - t.Fatalf("GetAPIKey() error = %v, want missing key error", err) + if !errors.Is(err, ErrNoAPIKey) { + t.Fatalf("GetAPIKey() error = %v, want ErrNoAPIKey", err) } } diff --git a/pkg/agentlib/agentlib.go b/pkg/agentlib/agentlib.go index 58958ff..99f741f 100644 --- a/pkg/agentlib/agentlib.go +++ b/pkg/agentlib/agentlib.go @@ -88,8 +88,9 @@ const ( var ( // ErrMissingAPIKey is returned when the shared data directory does not have - // a configured API key. - ErrMissingAPIKey = errors.New("API key is not configured in ~/.tokitoki/api_key") + // a configured API key. It aliases the inner package's sentinel so a key + // error raised anywhere in the stack compares equal here. + ErrMissingAPIKey = cli.ErrNoAPIKey // ErrNoScanDirectories is returned when a sync call has no provider // directory to scan. @@ -469,7 +470,7 @@ func DefaultProviderDirs() map[Provider][]string { ProviderCodex: {filepath.Join(home, ".codex")}, ProviderCopilot: {filepath.Join(home, ".copilot", "otel")}, ProviderGemini: {filepath.Join(home, ".gemini", "tmp")}, - ProviderKimi: {filepath.Join(home, ".kimi")}, + ProviderKimi: {filepath.Join(home, ".kimi"), filepath.Join(home, ".kimi-code")}, ProviderQwen: {filepath.Join(home, ".qwen")}, ProviderOpenClaw: {filepath.Join(home, ".openclaw"), filepath.Join(home, ".clawdbot"), filepath.Join(home, ".moltbot"), filepath.Join(home, ".moldbot")}, ProviderPi: {filepath.Join(home, ".pi", "agent", "sessions")}, diff --git a/pkg/agentlib/agentlib_test.go b/pkg/agentlib/agentlib_test.go index cc3dfd3..7277f74 100644 --- a/pkg/agentlib/agentlib_test.go +++ b/pkg/agentlib/agentlib_test.go @@ -106,6 +106,9 @@ func TestDefaultProviderDirsIncludesBuiltInProviders(t *testing.T) { t.Fatalf("%s dirs = %#v, want first dir %q", provider, got, dir) } } + if got := dirs[ProviderKimi]; len(got) != 2 || got[1] != filepath.Join(home, ".kimi-code") { + t.Fatalf("kimi dirs = %#v, want .kimi and .kimi-code", got) + } if got := dirs[ProviderOpenClaw]; len(got) != 4 { t.Fatalf("openclaw dirs = %#v, want four defaults", got) } From 8d65f2c78de848bb39129be7b8399a506f902414 Mon Sep 17 00:00:00 2001 From: akarineren Date: Thu, 30 Jul 2026 21:35:07 +0900 Subject: [PATCH 2/2] Enhance language detection for C-family headers: add support for ".h" files and implement tests for header language resolution --- internal/langdetect/langdetect.go | 56 +++++++++++++++++++++----- internal/langdetect/langdetect_test.go | 48 ++++++++++++++++++++++ 2 files changed, 95 insertions(+), 9 deletions(-) diff --git a/internal/langdetect/langdetect.go b/internal/langdetect/langdetect.go index f2edff5..a7dabf7 100644 --- a/internal/langdetect/langdetect.go +++ b/internal/langdetect/langdetect.go @@ -99,7 +99,6 @@ var extensionLanguages = map[string]string{ ".groovy": "Groovy", ".gsp": "Gosu", ".gs": "Gosu", - ".h": "C/C++ Header", ".h++": "C++", ".haml": "Haml", ".hh": "C++", @@ -183,18 +182,33 @@ var extensionLanguages = map[string]string{ ".zsh": "Bash", } +// cFamilyLanguages lists the languages a ".h" header can belong to, in the +// order preferred on a tie. A header inherits the dominant one from the +// surrounding candidates; with no context it defaults to C. +var cFamilyLanguages = []string{"C++", "C", "Objective-C", "Objective-C++"} + +func normalizedBase(path string) string { + return strings.ToLower(filepath.Base(strings.Trim(strings.TrimSpace(path), `"'`))) +} + +func isCHeader(path string) bool { + return filepath.Ext(normalizedBase(path)) == ".h" +} + func FromPath(path string) string { - path = strings.TrimSpace(path) - if path == "" { + base := normalizedBase(path) + if base == "" || base == "." { return Unknown } - base := strings.ToLower(filepath.Base(strings.Trim(path, `"'`))) if language, ok := filenameLanguages[base]; ok { return language } - ext := strings.ToLower(filepath.Ext(base)) + ext := filepath.Ext(base) + if ext == ".h" { + return "C" + } if language, ok := extensionLanguages[ext]; ok { return language } @@ -210,6 +224,21 @@ func DominantFromPaths(paths []string) string { return Dominant(candidates) } +// headerLanguage resolves ".h" files to the dominant C-family language among +// the other candidates, the same way wakatime-cli resolves headers by their +// sibling files. Defaults to C when no context exists. +func headerLanguage(scores map[string]int) string { + best := "C" + bestWeight := 0 + for _, language := range cFamilyLanguages { + if scores[language] > bestWeight { + best = language + bestWeight = scores[language] + } + } + return best +} + func Dominant(candidates []Candidate) string { type score struct { language string @@ -217,18 +246,27 @@ func Dominant(candidates []Candidate) string { } scores := map[string]int{} + headerWeight := 0 for _, candidate := range candidates { - language := FromPath(candidate.Path) - if language == Unknown { - continue - } weight := candidate.Weight if weight <= 0 { weight = 1 } + if isCHeader(candidate.Path) { + headerWeight += weight + continue + } + language := FromPath(candidate.Path) + if language == Unknown { + continue + } scores[language] += weight } + if headerWeight > 0 { + scores[headerLanguage(scores)] += headerWeight + } + if len(scores) == 0 { return Unknown } diff --git a/internal/langdetect/langdetect_test.go b/internal/langdetect/langdetect_test.go index 4202bcd..6ce8128 100644 --- a/internal/langdetect/langdetect_test.go +++ b/internal/langdetect/langdetect_test.go @@ -10,6 +10,7 @@ func TestFromPathUsesFilenameAndExtensionRules(t *testing.T) { "/repo/CMakeLists.txt": "CMake", "/repo/Dockerfile": "Docker", "/repo/README.md": "Markdown", + "/repo/include/util.h": "C", "/repo/unknown.nopeext": Unknown, } @@ -31,6 +32,53 @@ func TestDominantWeightsCandidates(t *testing.T) { } } +func TestDominantResolvesHeadersToContextLanguage(t *testing.T) { + tests := []struct { + name string + candidates []Candidate + want string + }{ + { + name: "headers inherit C++ from session", + candidates: []Candidate{ + {Path: "/repo/src/engine.cpp", Weight: 1}, + {Path: "/repo/src/engine.h", Weight: 3}, + {Path: "/repo/src/render.h", Weight: 3}, + }, + want: "C++", + }, + { + name: "headers inherit C from session", + candidates: []Candidate{ + {Path: "/repo/src/main.c", Weight: 1}, + {Path: "/repo/src/main.h", Weight: 3}, + }, + want: "C", + }, + { + name: "headers inherit Objective-C from session", + candidates: []Candidate{ + {Path: "/repo/App/AppDelegate.m", Weight: 1}, + {Path: "/repo/App/AppDelegate.h", Weight: 3}, + }, + want: "Objective-C", + }, + { + name: "headers alone default to C", + candidates: []Candidate{ + {Path: "/repo/include/util.h", Weight: 1}, + }, + want: "C", + }, + } + + for _, tt := range tests { + if got := Dominant(tt.candidates); got != tt.want { + t.Fatalf("%s: Dominant = %q, want %q", tt.name, got, tt.want) + } + } +} + func TestPathsFromTextExtractsKnownFilePaths(t *testing.T) { paths := PathsFromText(`sed -n '1,20p' internal/httpapi/server.go && cat app/page.tsx`) if len(paths) != 2 {