diff --git a/cmd/entire/cli/agent/agent.go b/cmd/entire/cli/agent/agent.go index eb85bf7e8..71815d01b 100644 --- a/cmd/entire/cli/agent/agent.go +++ b/cmd/entire/cli/agent/agent.go @@ -5,6 +5,7 @@ package agent import ( "context" + "encoding/json" "io" "os/exec" @@ -188,6 +189,31 @@ type TokenCalculator interface { CalculateTokenUsage(transcriptData []byte, fromOffset int) (*TokenUsage, error) } +// OutOfBandTokenSource provides token usage from a source other than the +// transcript. Antigravity is the only agent that needs this: agy never writes +// token data into its transcript or hook payloads — its title/statusline pipe +// is the only surface, captured to disk by `entire hooks antigravity +// title-tee` (see agent/antigravity/statusline.go). +// +// Flow: the lifecycle calls SnapshotTokenBaseline at TurnStart and stores the +// opaque baseline in PrePromptState; at TurnEnd (when transcript-based +// calculation yields nothing) it calls CalculateTokenUsageSince to get the +// checkpoint-scoped delta — the same cumulative-totals-minus-baseline pattern +// Codex uses, sourced out-of-band. +type OutOfBandTokenSource interface { + Agent + + // SnapshotTokenBaseline returns an opaque, agent-defined marker of the + // current cumulative token position for the session. A nil baseline with + // nil error means "no usage observed yet" (delta will count from zero). + SnapshotTokenBaseline(ctx context.Context, sessionID string) (json.RawMessage, error) + + // CalculateTokenUsageSince computes usage between the baseline and now. + // A nil result with nil error means no data is available (degrade to no + // token counts, never to an error). + CalculateTokenUsageSince(ctx context.Context, sessionID string, baseline json.RawMessage) (*TokenUsage, error) +} + // TextGenerator is an optional interface for agents whose CLI supports // non-interactive text generation (e.g., claude --print). // Used for AI-powered metadata generation (trail titles, summaries). diff --git a/cmd/entire/cli/agent/antigravity/antigravity.go b/cmd/entire/cli/agent/antigravity/antigravity.go index e557d2adb..853106416 100644 --- a/cmd/entire/cli/agent/antigravity/antigravity.go +++ b/cmd/entire/cli/agent/antigravity/antigravity.go @@ -21,6 +21,8 @@ type AntigravityAgent struct { CommandRunner agent.TextCommandRunner } +var _ agent.OutOfBandTokenSource = (*AntigravityAgent)(nil) + // NewAntigravityAgent creates a new AntigravityAgent instance. func NewAntigravityAgent() agent.Agent { return &AntigravityAgent{} diff --git a/cmd/entire/cli/agent/antigravity/antigravity_test.go b/cmd/entire/cli/agent/antigravity/antigravity_test.go index 4dfa34459..3147707fd 100644 --- a/cmd/entire/cli/agent/antigravity/antigravity_test.go +++ b/cmd/entire/cli/agent/antigravity/antigravity_test.go @@ -50,6 +50,7 @@ func TestDetectPresence(t *testing.T) { t.Run("hooks installed", func(t *testing.T) { tempDir := t.TempDir() t.Chdir(tempDir) + t.Setenv(configDirEnv, t.TempDir()) ag := &AntigravityAgent{} if _, err := ag.InstallHooks(context.Background(), false, false); err != nil { diff --git a/cmd/entire/cli/agent/antigravity/hooks.go b/cmd/entire/cli/agent/antigravity/hooks.go index 967feebad..eee694417 100644 --- a/cmd/entire/cli/agent/antigravity/hooks.go +++ b/cmd/entire/cli/agent/antigravity/hooks.go @@ -10,6 +10,7 @@ import ( "github.com/entireio/cli/cmd/entire/cli/agent" "github.com/entireio/cli/cmd/entire/cli/jsonutil" + "github.com/entireio/cli/cmd/entire/cli/logging" "github.com/entireio/cli/cmd/entire/cli/paths" ) @@ -59,6 +60,20 @@ func (a *AntigravityAgent) InstallHooks(ctx context.Context, localDev bool, forc candidate := buildEntireHookConfig(cmdPrefix, localDev) + // Title tee: agy's only token-usage surface (same payload as the + // statusline script). Run this BEFORE the idempotency early-return: the + // title slot lives in agy's GLOBAL settings.json, independent of this + // repo's .agents/hooks.json. If repo hooks already match but the global + // slot is missing or stale (upgrade from a pre-title-tee version, a failed + // first install, or `entire agent add` without --force), re-running setup + // must still repair it — otherwise the doctor's "re-run setup" hint is a + // no-op. InstallTitleTee is itself idempotent. Best-effort: a failure to + // claim the global slot must not fail repo-level hook setup. + if err := InstallTitleTee(localDev); err != nil { + logging.Warn(ctx, "failed to install antigravity title tee", + "error", err.Error()) + } + // Idempotency check: compare candidate against existing "entire" entry by // re-marshaling both to compact JSON for a stable comparison. if !force { diff --git a/cmd/entire/cli/agent/antigravity/hooks_test.go b/cmd/entire/cli/agent/antigravity/hooks_test.go index f423d1caa..33c6352c2 100644 --- a/cmd/entire/cli/agent/antigravity/hooks_test.go +++ b/cmd/entire/cli/agent/antigravity/hooks_test.go @@ -12,6 +12,7 @@ import ( func TestInstallHooks_FreshRepo(t *testing.T) { tmpDir := t.TempDir() t.Chdir(tmpDir) + t.Setenv(configDirEnv, t.TempDir()) a := &AntigravityAgent{} n, err := a.InstallHooks(context.Background(), false, false) @@ -49,6 +50,7 @@ func TestInstallHooks_FreshRepo(t *testing.T) { func TestInstallHooks_Idempotent(t *testing.T) { tmpDir := t.TempDir() t.Chdir(tmpDir) + t.Setenv(configDirEnv, t.TempDir()) a := &AntigravityAgent{} @@ -74,6 +76,7 @@ func TestInstallHooks_Idempotent(t *testing.T) { func TestInstallHooks_PreservesForeignHooks(t *testing.T) { tmpDir := t.TempDir() t.Chdir(tmpDir) + t.Setenv(configDirEnv, t.TempDir()) // Pre-seed .agents/hooks.json with a foreign entry agentsDir := filepath.Join(tmpDir, ".agents") @@ -127,6 +130,7 @@ func TestInstallHooks_PreservesForeignHooks(t *testing.T) { func TestUninstallHooks_LeavesForeignHooks(t *testing.T) { tmpDir := t.TempDir() t.Chdir(tmpDir) + t.Setenv(configDirEnv, t.TempDir()) a := &AntigravityAgent{} @@ -187,6 +191,7 @@ func TestUninstallHooks_LeavesForeignHooks(t *testing.T) { func TestInstallHooks_LocalDevWritesQuotedSubshell(t *testing.T) { tmpDir := t.TempDir() t.Chdir(tmpDir) + t.Setenv(configDirEnv, t.TempDir()) a := &AntigravityAgent{} if _, err := a.InstallHooks(context.Background(), true, false); err != nil { @@ -211,6 +216,7 @@ func TestInstallHooks_LocalDevWritesQuotedSubshell(t *testing.T) { func TestAreHooksInstalled(t *testing.T) { tmpDir := t.TempDir() t.Chdir(tmpDir) + t.Setenv(configDirEnv, t.TempDir()) a := &AntigravityAgent{} @@ -226,3 +232,40 @@ func TestAreHooksInstalled(t *testing.T) { t.Error("AreHooksInstalled() = false after install, want true") } } + +// TestInstallHooks_IdempotentStillRepairsTitleTee guards the regression where +// the repo-hooks idempotency early-return skipped the global title-tee install, +// leaving Antigravity checkpoints without token counts after an upgrade or a +// failed first title install (and making the doctor's "re-run setup" hint a +// no-op). Re-running InstallHooks must repair the missing global slot even when +// the repo's .agents/hooks.json already matches. +func TestInstallHooks_IdempotentStillRepairsTitleTee(t *testing.T) { + tmpDir := t.TempDir() + t.Chdir(tmpDir) + t.Setenv(configDirEnv, t.TempDir()) + + a := &AntigravityAgent{} + if _, err := a.InstallHooks(context.Background(), false, false); err != nil { + t.Fatalf("first InstallHooks: %v", err) + } + if !TitleTeeInstalled() { + t.Fatal("title tee should be installed after the first InstallHooks") + } + + // Simulate a missing/stale global slot while repo hooks remain correct. + if err := UninstallTitleTee(); err != nil { + t.Fatalf("UninstallTitleTee: %v", err) + } + if TitleTeeInstalled() { + t.Fatal("precondition: title tee should be gone before the idempotent re-install") + } + + // Second install hits the repo-hooks idempotency early-return, but must + // still re-install the missing global title tee. + if _, err := a.InstallHooks(context.Background(), false, false); err != nil { + t.Fatalf("second InstallHooks: %v", err) + } + if !TitleTeeInstalled() { + t.Error("idempotent InstallHooks must repair the missing title tee") + } +} diff --git a/cmd/entire/cli/agent/antigravity/statusline.go b/cmd/entire/cli/agent/antigravity/statusline.go new file mode 100644 index 000000000..deb50d26e --- /dev/null +++ b/cmd/entire/cli/agent/antigravity/statusline.go @@ -0,0 +1,322 @@ +package antigravity + +import ( + "bufio" + "bytes" + "context" + "encoding/json" + "fmt" + "os" + "path/filepath" + "time" + + "github.com/entireio/cli/cmd/entire/cli/agent" +) + +// agy's title/statusline hook pipes a state JSON to the configured command on +// every agent state change. The context_window object is the ONLY surface +// where agy exposes token usage — it never appears in transcripts or +// lifecycle hook payloads. AppendStatusSnapshot persists those snapshots so +// the lifecycle can compute per-checkpoint deltas later. +// +// Totals are cumulative per conversation; current_usage is the latest API call. + +// statusDirEnv overrides the snapshot cache directory (tests, ops). +const statusDirEnv = "ENTIRE_ANTIGRAVITY_STATUS_DIR" + +// statusRetention is how long snapshot files for other conversations are kept. +const statusRetention = 14 * 24 * time.Hour + +// statusCurrentUsage mirrors context_window.current_usage in the payload. +type statusCurrentUsage struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + CacheCreationInputTokens int `json:"cache_creation_input_tokens"` + CacheReadInputTokens int `json:"cache_read_input_tokens"` +} + +// statusContextWindow mirrors context_window in the payload. +type statusContextWindow struct { + TotalInputTokens int `json:"total_input_tokens"` + TotalOutputTokens int `json:"total_output_tokens"` + ContextWindowSize int `json:"context_window_size,omitempty"` + CurrentUsage *statusCurrentUsage `json:"current_usage,omitempty"` +} + +// statusSnapshot is one persisted line in .jsonl. +type statusSnapshot struct { + Timestamp string `json:"ts"` + ConversationID string `json:"conversation_id"` + ContextWindow statusContextWindow `json:"context_window"` +} + +// statuslinePayload is the subset of agy's state JSON we consume. +type statuslinePayload struct { + ConversationID string `json:"conversation_id"` + ContextWindow *statusContextWindow `json:"context_window"` +} + +// statusDir returns the directory used to store snapshot files. +// It honours the ENTIRE_ANTIGRAVITY_STATUS_DIR env override (tests, ops), +// otherwise uses /entire/antigravity/status. +func statusDir() (string, error) { + if override := os.Getenv(statusDirEnv); override != "" { + return override, nil + } + cacheDir, err := os.UserCacheDir() + if err != nil { + return "", fmt.Errorf("antigravity status: resolve user cache dir: %w", err) + } + return filepath.Join(cacheDir, "entire", "antigravity", "status"), nil +} + +// statusFilePath returns the path for the JSONL snapshot file of a conversation. +// filepath.Base guards against path traversal in the conversation ID. +func statusFilePath(conversationID string) (string, error) { + dir, err := statusDir() + if err != nil { + return "", err + } + return filepath.Join(dir, filepath.Base(conversationID)+".jsonl"), nil +} + +// AppendStatusSnapshot parses an agy state-JSON payload and appends a snapshot +// to the per-conversation JSONL file. The hot path never returns an error for +// malformed input — only for genuine I/O failures after the file has been opened. +func AppendStatusSnapshot(payload []byte) error { + var p statuslinePayload + if err := json.Unmarshal(payload, &p); err != nil { + return nil + } + if p.ConversationID == "" || p.ContextWindow == nil { + return nil // missing required fields — silently skip + } + + dir, err := statusDir() + if err != nil { + return err + } + isNew := false + filePath := filepath.Join(dir, filepath.Base(p.ConversationID)+".jsonl") + + if _, statErr := os.Stat(filePath); os.IsNotExist(statErr) { + isNew = true + } + + if err := os.MkdirAll(dir, 0o750); err != nil { + return fmt.Errorf("antigravity status: mkdir: %w", err) + } + + // Dedup: compare compact JSON of the new context_window against the last + // persisted line's context_window. readLastContextWindow streams the file + // keeping only the final line (O(1) memory); the file stays small because + // this very dedup suppresses unchanged snapshots. + newCWBytes, err := json.Marshal(p.ContextWindow) + if err != nil { + return nil + } + + if !isNew { + lastCW, readErr := readLastContextWindow(filePath) + if readErr == nil && lastCW != nil { + lastCWBytes, marshalErr := json.Marshal(lastCW) + if marshalErr == nil && bytes.Equal(newCWBytes, lastCWBytes) { + return nil // duplicate — skip + } + } + } + + snap := statusSnapshot{ + Timestamp: time.Now().UTC().Format(time.RFC3339Nano), + ConversationID: p.ConversationID, + ContextWindow: *p.ContextWindow, + } + line, err := json.Marshal(snap) + if err != nil { + return nil + } + + //nolint:gosec // filePath is derived from filepath.Base(conversationID) + f, err := os.OpenFile(filePath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600) + if err != nil { + return fmt.Errorf("antigravity status: open: %w", err) + } + defer func() { _ = f.Close() }() + + line = append(line, '\n') + if _, err := f.Write(line); err != nil { + return fmt.Errorf("antigravity status: write: %w", err) + } + + // Best-effort prune of stale files for other conversations when we first + // create the active file (avoids per-append overhead). + if isNew { + pruneStaleStatusFiles(dir, p.ConversationID) + } + + return nil +} + +// SnapshotTokenBaseline returns the latest persisted snapshot for the +// conversation, or nil if none exists yet. A nil baseline is exact only for a +// genuinely fresh conversation; a resumed conversation whose title-tee shim +// hasn't written a snapshot before the first TurnStart will over-count the +// prior cumulative total on that first tracked turn. +func (a *AntigravityAgent) SnapshotTokenBaseline(_ context.Context, sessionID string) (json.RawMessage, error) { + snaps, err := readStatusSnapshots(sessionID) + if err != nil || len(snaps) == 0 { + return nil, nil //nolint:nilerr // degrade to "no baseline"; never block the turn + } + raw, err := json.Marshal(snaps[len(snaps)-1]) + if err != nil { + return nil, nil //nolint:nilerr // ditto + } + return raw, nil +} + +// CalculateTokenUsageSince computes the delta between the baseline snapshot +// and the latest persisted snapshot. +// +// Exact: InputTokens/OutputTokens (cumulative totals minus baseline totals). +// Best-effort: cache fields and APICallCount, derived from the snapshot lines +// appended after the baseline timestamp (the dedup writer appends ~one line +// per API response, but lines can be missed between agent state changes). +func (a *AntigravityAgent) CalculateTokenUsageSince(_ context.Context, sessionID string, baseline json.RawMessage) (*agent.TokenUsage, error) { + snaps, err := readStatusSnapshots(sessionID) + if err != nil || len(snaps) == 0 { + return nil, nil //nolint:nilerr,nilnil // no data -> no token counts, never an error + } + + var base statusSnapshot + if len(baseline) > 0 { + _ = json.Unmarshal(baseline, &base) //nolint:errcheck // unparseable baseline -> zero baseline + } + + latest := snaps[len(snaps)-1] + usage := &agent.TokenUsage{ + InputTokens: max(0, latest.ContextWindow.TotalInputTokens-base.ContextWindow.TotalInputTokens), + OutputTokens: max(0, latest.ContextWindow.TotalOutputTokens-base.ContextWindow.TotalOutputTokens), + } + + // The strictly-after (.After, not >=) filter is load-bearing for + // multi-turn correctness: turn N+1's baseline IS turn N's latest snapshot, + // so excluding the equal-timestamp boundary line prevents re-counting it. + // Changing this to >= would double-count the boundary line every turn. + baseTS, baseTSErr := time.Parse(time.RFC3339Nano, base.Timestamp) + for _, s := range snaps { + // If baseTS is unparseable we count cache/apicalls over all lines; accepted because input/output remain exact via the totals delta. + if base.Timestamp != "" && baseTSErr == nil { + ts, parseErr := time.Parse(time.RFC3339Nano, s.Timestamp) + if parseErr != nil || !ts.After(baseTS) { + continue + } + } + usage.APICallCount++ + if cu := s.ContextWindow.CurrentUsage; cu != nil { + usage.CacheCreationTokens += cu.CacheCreationInputTokens + usage.CacheReadTokens += cu.CacheReadInputTokens + } + } + + if usage.InputTokens == 0 && usage.OutputTokens == 0 && usage.CacheCreationTokens == 0 && usage.CacheReadTokens == 0 { + return nil, nil //nolint:nilnil // nothing observed this turn + } + return usage, nil +} + +// readLastContextWindow streams the JSONL file line-by-line and returns the +// context_window of the final non-empty line (O(1) memory, O(n) I/O), or nil +// if the file has no usable line. Used for dedup comparison. +func readLastContextWindow(filePath string) (*statusContextWindow, error) { + //nolint:gosec // filePath is derived from filepath.Base(conversationID) + f, err := os.Open(filePath) + if err != nil { + return nil, fmt.Errorf("antigravity status: open for dedup: %w", err) + } + defer func() { _ = f.Close() }() + + scanner := bufio.NewScanner(f) + scanner.Buffer(make([]byte, 1024*1024), 1024*1024) + + var lastLine string + for scanner.Scan() { + if line := scanner.Text(); line != "" { + lastLine = line + } + } + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("antigravity status: scan for dedup: %w", err) + } + if lastLine == "" { + return nil, nil //nolint:nilnil // no lines yet — caller handles nil gracefully + } + + var snap statusSnapshot + if err := json.Unmarshal([]byte(lastLine), &snap); err != nil { + return nil, nil //nolint:nilnil // malformed last line — treat as no prior snapshot + } + return &snap.ContextWindow, nil +} + +// readStatusSnapshots reads all valid snapshot lines from the JSONL file for +// the given conversationID. A missing file returns nil, nil (not an error). +func readStatusSnapshots(conversationID string) ([]statusSnapshot, error) { + filePath, err := statusFilePath(conversationID) + if err != nil { + return nil, err + } + + //nolint:gosec // filePath is derived from filepath.Base(conversationID) + f, err := os.Open(filePath) + if os.IsNotExist(err) { + return nil, nil + } + if err != nil { + return nil, fmt.Errorf("antigravity status: open for read: %w", err) + } + defer func() { _ = f.Close() }() + + scanner := bufio.NewScanner(f) + scanner.Buffer(make([]byte, 1024*1024), 1024*1024) + + var snaps []statusSnapshot + for scanner.Scan() { + line := scanner.Text() + if line == "" { + continue + } + var snap statusSnapshot + if err := json.Unmarshal([]byte(line), &snap); err != nil { + continue // skip malformed lines + } + snaps = append(snaps, snap) + } + if err := scanner.Err(); err != nil { + return nil, fmt.Errorf("antigravity status: scan: %w", err) + } + return snaps, nil +} + +// pruneStaleStatusFiles removes JSONL files in dir that are not the active +// conversation and whose mtime is older than statusRetention. Best-effort: +// errors are silently ignored. +func pruneStaleStatusFiles(dir, activeConversationID string) { + activeFile := filepath.Base(activeConversationID) + ".jsonl" + entries, err := os.ReadDir(dir) + if err != nil { + return + } + cutoff := time.Now().Add(-statusRetention) + for _, entry := range entries { + if entry.IsDir() || entry.Name() == activeFile { + continue + } + info, err := entry.Info() + if err != nil { + continue + } + if info.ModTime().Before(cutoff) { + _ = os.Remove(filepath.Join(dir, entry.Name())) + } + } +} diff --git a/cmd/entire/cli/agent/antigravity/statusline_test.go b/cmd/entire/cli/agent/antigravity/statusline_test.go new file mode 100644 index 000000000..7ee5d682a --- /dev/null +++ b/cmd/entire/cli/agent/antigravity/statusline_test.go @@ -0,0 +1,592 @@ +package antigravity + +import ( + "context" + "encoding/json" + "os" + "path/filepath" + "testing" + "time" +) + +// Note: these tests use t.Setenv and/or t.Chdir, so t.Parallel() is not called. + +func TestAppendStatusSnapshot_WritesLineKeyedByConversation(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + payload := []byte(`{"conversation_id":"conv-1","agent_state":"working","context_window":{"total_input_tokens":1000,"total_output_tokens":50,"context_window_size":200000,"current_usage":{"input_tokens":900,"output_tokens":50,"cache_creation_input_tokens":100,"cache_read_input_tokens":800}}}`) + + if err := AppendStatusSnapshot(payload); err != nil { + t.Fatalf("AppendStatusSnapshot: %v", err) + } + + snaps, err := readStatusSnapshots("conv-1") + if err != nil { + t.Fatalf("readStatusSnapshots: %v", err) + } + if len(snaps) != 1 { + t.Fatalf("got %d snapshots, want 1", len(snaps)) + } + + s := snaps[0] + if s.ContextWindow.TotalInputTokens != 1000 { + t.Errorf("TotalInputTokens = %d, want 1000", s.ContextWindow.TotalInputTokens) + } + if s.ContextWindow.TotalOutputTokens != 50 { + t.Errorf("TotalOutputTokens = %d, want 50", s.ContextWindow.TotalOutputTokens) + } + if s.ContextWindow.CurrentUsage == nil { + t.Fatal("CurrentUsage is nil") + } + if s.ContextWindow.CurrentUsage.CacheReadInputTokens != 800 { + t.Errorf("CacheReadInputTokens = %d, want 800", s.ContextWindow.CurrentUsage.CacheReadInputTokens) + } + if s.Timestamp == "" { + t.Error("Timestamp is empty") + } +} + +func TestAppendStatusSnapshot_DedupsUnchangedContextWindow(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + // First payload + p1 := []byte(`{"conversation_id":"conv-2","agent_state":"working","context_window":{"total_input_tokens":1000,"total_output_tokens":50}}`) + // Second payload: same context_window, different agent_state + p2 := []byte(`{"conversation_id":"conv-2","agent_state":"idle","context_window":{"total_input_tokens":1000,"total_output_tokens":50}}`) + // Third payload: different context_window + p3 := []byte(`{"conversation_id":"conv-2","agent_state":"working","context_window":{"total_input_tokens":2000,"total_output_tokens":100}}`) + + for _, p := range [][]byte{p1, p2, p3} { + if err := AppendStatusSnapshot(p); err != nil { + t.Fatalf("AppendStatusSnapshot: %v", err) + } + } + + snaps, err := readStatusSnapshots("conv-2") + if err != nil { + t.Fatalf("readStatusSnapshots: %v", err) + } + if len(snaps) != 2 { + t.Errorf("got %d snapshots, want 2 (dedup should have skipped p2)", len(snaps)) + } +} + +func TestAppendStatusSnapshot_IgnoresGarbageAndMissingFields(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + for _, payload := range []string{"not json", "{}", `{"conversation_id":"conv-3"}`} { + if err := AppendStatusSnapshot([]byte(payload)); err != nil { + t.Errorf("AppendStatusSnapshot(%q) returned error: %v", payload, err) + } + } + + // dir must be empty (no snapshot files written) + entries, err := os.ReadDir(dir) + if err != nil { + t.Fatalf("ReadDir: %v", err) + } + if len(entries) != 0 { + t.Errorf("expected empty dir, got %d entries", len(entries)) + } +} + +func TestReadStatusSnapshots_SkipsMalformedLines(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + // Write a file manually with valid / garbage / valid lines + validLine1, err := json.Marshal(statusSnapshot{ + Timestamp: "2026-01-01T00:00:00Z", + ConversationID: "conv-4", + ContextWindow: statusContextWindow{TotalInputTokens: 10, TotalOutputTokens: 1}, + }) + if err != nil { + t.Fatal(err) + } + validLine2, err := json.Marshal(statusSnapshot{ + Timestamp: "2026-01-01T00:01:00Z", + ConversationID: "conv-4", + ContextWindow: statusContextWindow{TotalInputTokens: 20, TotalOutputTokens: 2}, + }) + if err != nil { + t.Fatal(err) + } + + filePath := filepath.Join(dir, "conv-4.jsonl") + content := string(validLine1) + "\nGARBAGE LINE\n" + string(validLine2) + "\n" + if err := os.WriteFile(filePath, []byte(content), 0o600); err != nil { + t.Fatal(err) + } + + snaps, err := readStatusSnapshots("conv-4") + if err != nil { + t.Fatalf("readStatusSnapshots: %v", err) + } + if len(snaps) != 2 { + t.Errorf("got %d snapshots, want 2 (malformed line skipped)", len(snaps)) + } +} + +func TestReadStatusSnapshots_MissingFileReturnsEmpty(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + snaps, err := readStatusSnapshots("no-such-conv") + if err != nil { + t.Fatalf("readStatusSnapshots: %v", err) + } + if len(snaps) != 0 { + t.Errorf("got %d snapshots, want 0", len(snaps)) + } +} + +func BenchmarkAppendStatusSnapshot_GrownFile(b *testing.B) { + dir := b.TempDir() + b.Setenv(statusDirEnv, dir) + + // Seed 500 distinct snapshots + for i := range 500 { + payload, err := json.Marshal(map[string]any{ + "conversation_id": "bench-conv", + "agent_state": "working", + "context_window": map[string]any{ + "total_input_tokens": 1000 + i, + "total_output_tokens": 50 + i, + }, + }) + if err != nil { + b.Fatal(err) + } + if err := AppendStatusSnapshot(payload); err != nil { + b.Fatalf("seed %d: %v", i, err) + } + } + + // The duplicate payload matches the last seeded snapshot (dedup path) + dupPayload, err := json.Marshal(map[string]any{ + "conversation_id": "bench-conv", + "agent_state": "working", + "context_window": map[string]any{ + "total_input_tokens": 1000 + 499, + "total_output_tokens": 50 + 499, + }, + }) + if err != nil { + b.Fatal(err) + } + + b.ResetTimer() + for range b.N { + if err := AppendStatusSnapshot(dupPayload); err != nil { + b.Fatal(err) + } + } +} + +func TestAppendStatusSnapshot_PrunesStaleFilesOnCreate(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + // Pre-seed a stale file for another conversation (mtime older than retention) + stale := filepath.Join(dir, "old-conv.jsonl") + if err := os.WriteFile(stale, []byte("{}\n"), 0o600); err != nil { + t.Fatal(err) + } + old := time.Now().Add(-statusRetention - time.Hour) + if err := os.Chtimes(stale, old, old); err != nil { + t.Fatal(err) + } + // And a fresh file for a third conversation that must survive + fresh := filepath.Join(dir, "fresh-conv.jsonl") + if err := os.WriteFile(fresh, []byte("{}\n"), 0o600); err != nil { + t.Fatal(err) + } + + // First write for a new conversation triggers the prune + payload := []byte(`{"conversation_id":"conv-new","context_window":{"total_input_tokens":1}}`) + if err := AppendStatusSnapshot(payload); err != nil { + t.Fatal(err) + } + + if _, err := os.Stat(stale); !os.IsNotExist(err) { + t.Errorf("stale file should be pruned, stat err = %v", err) + } + if _, err := os.Stat(fresh); err != nil { + t.Errorf("fresh file must survive prune: %v", err) + } + if _, err := os.Stat(filepath.Join(dir, "conv-new.jsonl")); err != nil { + t.Errorf("active file must exist: %v", err) + } +} + +// writeSnapshotFixture writes the given snapshots as JSONL to the snapshot file +// for conversationID, using the statusDirEnv override already set by the test. +func writeSnapshotFixture(t *testing.T, conversationID string, snaps []statusSnapshot) { + t.Helper() + path, err := statusFilePath(conversationID) + if err != nil { + t.Fatalf("statusFilePath: %v", err) + } + if err := os.MkdirAll(filepath.Dir(path), 0o750); err != nil { + t.Fatalf("mkdir: %v", err) + } + var buf []byte + for _, s := range snaps { + line, marshalErr := json.Marshal(s) + if marshalErr != nil { + t.Fatalf("marshal snapshot: %v", marshalErr) + } + buf = append(buf, line...) + buf = append(buf, '\n') + } + if err := os.WriteFile(path, buf, 0o600); err != nil { + t.Fatalf("write fixture: %v", err) + } +} + +func TestCalculateTokenUsageSince_DeltaFromBaseline(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + snaps := []statusSnapshot{ + { + Timestamp: "2026-06-03T10:00:00.000000000Z", + ConversationID: "c1", + ContextWindow: statusContextWindow{ + TotalInputTokens: 1000, + TotalOutputTokens: 100, + CurrentUsage: &statusCurrentUsage{CacheCreationInputTokens: 200, CacheReadInputTokens: 700}, + }, + }, + { + Timestamp: "2026-06-03T10:05:00.000000000Z", + ConversationID: "c1", + ContextWindow: statusContextWindow{ + TotalInputTokens: 3000, + TotalOutputTokens: 250, + CurrentUsage: &statusCurrentUsage{CacheCreationInputTokens: 50, CacheReadInputTokens: 1900}, + }, + }, + { + Timestamp: "2026-06-03T10:06:00.000000000Z", + ConversationID: "c1", + ContextWindow: statusContextWindow{ + TotalInputTokens: 4500, + TotalOutputTokens: 400, + CurrentUsage: &statusCurrentUsage{CacheCreationInputTokens: 0, CacheReadInputTokens: 1500}, + }, + }, + } + writeSnapshotFixture(t, "c1", snaps) + + baseline, err := json.Marshal(snaps[0]) + if err != nil { + t.Fatalf("marshal baseline: %v", err) + } + + a := &AntigravityAgent{} + usage, err := a.CalculateTokenUsageSince(context.Background(), "c1", baseline) + if err != nil { + t.Fatalf("CalculateTokenUsageSince: %v", err) + } + if usage == nil { + t.Fatal("usage is nil") + } + if usage.InputTokens != 3500 { + t.Errorf("InputTokens = %d, want 3500", usage.InputTokens) + } + if usage.OutputTokens != 300 { + t.Errorf("OutputTokens = %d, want 300", usage.OutputTokens) + } + if usage.CacheCreationTokens != 50 { + t.Errorf("CacheCreationTokens = %d, want 50", usage.CacheCreationTokens) + } + if usage.CacheReadTokens != 3400 { + t.Errorf("CacheReadTokens = %d, want 3400", usage.CacheReadTokens) + } + if usage.APICallCount != 2 { + t.Errorf("APICallCount = %d, want 2", usage.APICallCount) + } +} + +func TestCalculateTokenUsageSince_NilBaselineCountsFromZero(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + snaps := []statusSnapshot{ + { + Timestamp: "2026-06-03T10:00:00.000000000Z", + ConversationID: "c2", + ContextWindow: statusContextWindow{ + TotalInputTokens: 1000, + TotalOutputTokens: 100, + CurrentUsage: &statusCurrentUsage{CacheReadInputTokens: 700}, + }, + }, + } + writeSnapshotFixture(t, "c2", snaps) + + a := &AntigravityAgent{} + usage, err := a.CalculateTokenUsageSince(context.Background(), "c2", nil) + if err != nil { + t.Fatalf("CalculateTokenUsageSince: %v", err) + } + if usage == nil { + t.Fatal("usage is nil") + } + if usage.InputTokens != 1000 { + t.Errorf("InputTokens = %d, want 1000", usage.InputTokens) + } + if usage.OutputTokens != 100 { + t.Errorf("OutputTokens = %d, want 100", usage.OutputTokens) + } +} + +func TestCalculateTokenUsageSince_NoDataReturnsNilNil(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + a := &AntigravityAgent{} + usage, err := a.CalculateTokenUsageSince(context.Background(), "missing-conv", nil) + if err != nil { + t.Fatalf("CalculateTokenUsageSince: %v", err) + } + if usage != nil { + t.Errorf("usage = %+v, want nil", usage) + } +} + +func TestSnapshotTokenBaseline_ReturnsLatestSnapshot(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + snaps := []statusSnapshot{ + { + Timestamp: "2026-06-03T10:00:00.000000000Z", + ConversationID: "c3", + ContextWindow: statusContextWindow{TotalInputTokens: 1000}, + }, + { + Timestamp: "2026-06-03T10:05:00.000000000Z", + ConversationID: "c3", + ContextWindow: statusContextWindow{TotalInputTokens: 2000}, + }, + } + writeSnapshotFixture(t, "c3", snaps) + + a := &AntigravityAgent{} + baseline, err := a.SnapshotTokenBaseline(context.Background(), "c3") + if err != nil { + t.Fatalf("SnapshotTokenBaseline: %v", err) + } + if len(baseline) == 0 { + t.Fatal("baseline is empty") + } + var snap statusSnapshot + if err := json.Unmarshal(baseline, &snap); err != nil { + t.Fatalf("unmarshal baseline: %v", err) + } + if snap.ContextWindow.TotalInputTokens != 2000 { + t.Errorf("baseline TotalInputTokens = %d, want 2000", snap.ContextWindow.TotalInputTokens) + } +} + +func TestSnapshotTokenBaseline_EmptyStoreReturnsNil(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + a := &AntigravityAgent{} + baseline, err := a.SnapshotTokenBaseline(context.Background(), "no-such-conv") + if err != nil { + t.Fatalf("SnapshotTokenBaseline: %v", err) + } + if baseline != nil { + t.Errorf("baseline = %v, want nil", baseline) + } +} + +func TestCalculateTokenUsageSince_ClampsWhenTotalsGoBackwards(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + // Simulate conversation-id reuse where cumulative totals reset lower than the baseline. + snaps := []statusSnapshot{ + { + Timestamp: "2026-06-03T10:05:00.000000000Z", + ConversationID: "c1", + ContextWindow: statusContextWindow{TotalInputTokens: 500, TotalOutputTokens: 30}, + }, + } + writeSnapshotFixture(t, "c1", snaps) + + // Baseline has higher totals than the latest snapshot — totals went backwards. + baseSnap := statusSnapshot{ + Timestamp: "2026-06-03T10:00:00.000000000Z", + ContextWindow: statusContextWindow{TotalInputTokens: 5000, TotalOutputTokens: 400}, + } + baseline, err := json.Marshal(baseSnap) + if err != nil { + t.Fatalf("marshal baseline: %v", err) + } + + a := &AntigravityAgent{} + usage, err := a.CalculateTokenUsageSince(context.Background(), "c1", baseline) + if err != nil { + t.Fatalf("CalculateTokenUsageSince: %v", err) + } + // The single line has no current_usage, so once max(0, ...) clamps the + // negative input/output deltas to zero, every field is zero and the method + // returns (nil, nil) — the all-zero "nothing observed" path. If the clamp + // were removed, the negative deltas would make usage non-nil, so asserting + // nil here pins the clamp behavior directly. + if usage != nil { + t.Errorf("usage = %+v, want nil (negative deltas clamped to zero -> all-zero -> nil)", usage) + } +} + +func TestCalculateTokenUsageSince_UnparseableBaselineCountsAllLines(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + snaps := []statusSnapshot{ + { + Timestamp: "2026-06-03T10:00:00.000000000Z", + ConversationID: "c1", + ContextWindow: statusContextWindow{ + TotalInputTokens: 1000, + TotalOutputTokens: 100, + CurrentUsage: &statusCurrentUsage{CacheReadInputTokens: 700}, + }, + }, + { + Timestamp: "2026-06-03T10:05:00.000000000Z", + ConversationID: "c1", + ContextWindow: statusContextWindow{ + TotalInputTokens: 3000, + TotalOutputTokens: 250, + CurrentUsage: &statusCurrentUsage{CacheReadInputTokens: 900}, + }, + }, + } + writeSnapshotFixture(t, "c1", snaps) + + // Baseline with an unparseable timestamp — documents accepted degradation: + // input/output stay exact via the totals subtraction, but cache/apicall + // counts cover all lines rather than only post-baseline lines. + baseline := json.RawMessage(`{"ts":"not-a-timestamp","context_window":{"total_input_tokens":1000,"total_output_tokens":100}}`) + + a := &AntigravityAgent{} + usage, err := a.CalculateTokenUsageSince(context.Background(), "c1", baseline) + if err != nil { + t.Fatalf("CalculateTokenUsageSince: %v", err) + } + if usage == nil { + t.Fatal("usage is nil") + } + if usage.InputTokens != 2000 { + t.Errorf("InputTokens = %d, want 2000 (3000-1000)", usage.InputTokens) + } + if usage.OutputTokens != 150 { + t.Errorf("OutputTokens = %d, want 150 (250-100)", usage.OutputTokens) + } + // Both lines are counted because the baseline timestamp didn't parse. + if usage.APICallCount != 2 { + t.Errorf("APICallCount = %d, want 2 (all lines counted)", usage.APICallCount) + } + if usage.CacheReadTokens != 1600 { + t.Errorf("CacheReadTokens = %d, want 1600 (700+900)", usage.CacheReadTokens) + } +} + +// TestAppendStatusSnapshot_CurrentUsageChangeIsNotDeduped proves that two +// payloads with IDENTICAL total_input_tokens/total_output_tokens but DIFFERENT +// current_usage are NOT deduped: both must be persisted. The delta calculation +// sums per-line cache fields from current_usage, so collapsing two lines that +// differ only in current_usage would silently drop cache accounting. +func TestAppendStatusSnapshot_CurrentUsageChangeIsNotDeduped(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + // Same totals {1000,100}, different current_usage cache_read. + a := []byte(`{"conversation_id":"conv-cu","agent_state":"working","context_window":{"total_input_tokens":1000,"total_output_tokens":100,"current_usage":{"cache_read_input_tokens":700}}}`) + b := []byte(`{"conversation_id":"conv-cu","agent_state":"working","context_window":{"total_input_tokens":1000,"total_output_tokens":100,"current_usage":{"cache_read_input_tokens":900}}}`) + + for _, p := range [][]byte{a, b} { + if err := AppendStatusSnapshot(p); err != nil { + t.Fatalf("AppendStatusSnapshot: %v", err) + } + } + + snaps, err := readStatusSnapshots("conv-cu") + if err != nil { + t.Fatalf("readStatusSnapshots: %v", err) + } + if len(snaps) != 2 { + t.Fatalf("got %d snapshots, want 2 (current_usage change must not be deduped)", len(snaps)) + } +} + +// TestCalculateTokenUsageSince_NoDoubleCountWhenBaselineIsLatest proves the +// load-bearing strictly-.After filter prevents double-counting across turns. +// Simulating turn N+1: use turn N's LATEST snapshot as the baseline and append +// nothing new. CalculateTokenUsageSince must return (nil, nil) — zero delta, +// APICallCount 0 — because no snapshot is strictly after the baseline timestamp +// and totals − baseline = 0. +func TestCalculateTokenUsageSince_NoDoubleCountWhenBaselineIsLatest(t *testing.T) { + dir := t.TempDir() + t.Setenv(statusDirEnv, dir) + + snaps := []statusSnapshot{ + { + Timestamp: "2026-06-03T10:00:00.000000000Z", + ConversationID: "c1", + ContextWindow: statusContextWindow{ + TotalInputTokens: 1000, + TotalOutputTokens: 100, + CurrentUsage: &statusCurrentUsage{CacheCreationInputTokens: 200, CacheReadInputTokens: 700}, + }, + }, + { + Timestamp: "2026-06-03T10:05:00.000000000Z", + ConversationID: "c1", + ContextWindow: statusContextWindow{ + TotalInputTokens: 3000, + TotalOutputTokens: 250, + CurrentUsage: &statusCurrentUsage{CacheCreationInputTokens: 50, CacheReadInputTokens: 1900}, + }, + }, + } + writeSnapshotFixture(t, "c1", snaps) + + a := &AntigravityAgent{} + + // Turn N: full usage from nil baseline. + full, err := a.CalculateTokenUsageSince(context.Background(), "c1", nil) + if err != nil { + t.Fatalf("CalculateTokenUsageSince (full): %v", err) + } + if full == nil || full.InputTokens != 3000 { + t.Fatalf("full usage = %+v, want InputTokens 3000", full) + } + + // Capture the latest snapshot (T2 line) exactly as the lifecycle does. + baseline, err := a.SnapshotTokenBaseline(context.Background(), "c1") + if err != nil { + t.Fatalf("SnapshotTokenBaseline: %v", err) + } + if len(baseline) == 0 { + t.Fatal("baseline is empty") + } + + // Turn N+1: nothing new appended. Delta from the latest baseline is zero. + usage, err := a.CalculateTokenUsageSince(context.Background(), "c1", baseline) + if err != nil { + t.Fatalf("CalculateTokenUsageSince (delta): %v", err) + } + if usage != nil { + t.Errorf("usage = %+v, want nil (no snapshot strictly after baseline; zero delta)", usage) + } +} diff --git a/cmd/entire/cli/agent/antigravity/title_install.go b/cmd/entire/cli/agent/antigravity/title_install.go new file mode 100644 index 000000000..8009ab560 --- /dev/null +++ b/cmd/entire/cli/agent/antigravity/title_install.go @@ -0,0 +1,283 @@ +package antigravity + +import ( + "context" + "encoding/json" + "fmt" + "os" + "os/exec" + "path/filepath" + "strings" + + "github.com/entireio/cli/cmd/entire/cli/jsonutil" +) + +// agy reads its window-title command from the GLOBAL config +// ~/.gemini/antigravity-cli/settings.json — a single slot: +// +// {"title": {"type": "command", "command": ""}} +// +// We occupy that slot with the title-tee shim (the title script receives the +// same state JSON as the statusline script — agy's only token-usage surface). +// A pre-existing user command is preserved INSIDE the shim invocation via +// --wrap '', making the config self-describing: uninstall restores +// the original without any backup file. Because the slot is global, per-repo +// `entire disable` does NOT uninstall it (other repos may rely on it); only +// agent removal does. + +// configDirEnv overrides the agy config directory (tests). +const configDirEnv = "ENTIRE_ANTIGRAVITY_CONFIG_DIR" + +const titleTeeMarker = "hooks antigravity title-tee" + +type titleConfig struct { + Type string `json:"type"` + Command string `json:"command"` +} + +// agyConfigDir returns the agy config directory, honouring the override env var. +func agyConfigDir() (string, error) { + if dir := os.Getenv(configDirEnv); dir != "" { + return dir, nil + } + home, err := os.UserHomeDir() + if err != nil { + return "", fmt.Errorf("failed to resolve home dir: %w", err) + } + return filepath.Join(home, ".gemini", "antigravity-cli"), nil +} + +// titleTeeCommand returns the full shell command string for the title-tee shim. +// If original is non-empty, the original command is wrapped via --wrap. +// +// localDev note: unlike the repo-scoped lifecycle hooks (.agents/hooks.json), +// the title command lives in agy's GLOBAL settings.json and is invoked from +// whatever directory agy runs in. A runtime "$(git rev-parse --show-toplevel)" +// would resolve against the wrong repo — or fail entirely outside one — so we +// resolve the repo root at install time and bake in the absolute main.go path. +// If resolution fails, we fall back to the production "entire ..." form, which +// resolves the dev binary via $PATH. +func titleTeeCommand(localDev bool, original string) string { + base := "entire hooks antigravity title-tee" + if localDev { + if mainPath := localDevMainPath(); mainPath != "" { + base = "go run " + shellSingleQuote(mainPath) + " hooks antigravity title-tee" + } + } + if original == "" { + return base + } + return base + " --wrap " + shellSingleQuote(original) +} + +// localDevMainPath resolves the absolute path to cmd/entire/main.go in the +// current repo at call time, or "" if not inside a git repo. +func localDevMainPath() string { + out, err := exec.CommandContext(context.Background(), "git", "rev-parse", "--show-toplevel").Output() + if err != nil { + return "" + } + root := strings.TrimSpace(string(out)) + if root == "" { + return "" + } + return filepath.Join(root, "cmd", "entire", "main.go") +} + +// shellSingleQuote wraps s in POSIX single quotes. Embedded single quotes are +// rewritten with the standard close-escape-reopen technique (see the +// strings.ReplaceAll below) so the result is safe inside a single-quoted shell +// argument. +func shellSingleQuote(s string) string { + return "'" + strings.ReplaceAll(s, "'", `'\''`) + "'" +} + +// InstallTitleTee installs the title-tee shim into agy's global settings.json. +// If a user's own title command is already present, it is preserved via --wrap. +// The call is idempotent: if our marker is already in the command, it returns nil. +func InstallTitleTee(localDev bool) error { + cfgDir, err := agyConfigDir() + if err != nil { + return err + } + settingsPath := filepath.Join(cfgDir, "settings.json") + + rawFile, err := readAgySettings(settingsPath) + if err != nil { + return err + } + + // Parse existing title entry (if any). Unparseable → treat as absent. + var existing titleConfig + if raw, ok := rawFile["title"]; ok { + _ = json.Unmarshal(raw, &existing) //nolint:errcheck // treat unparseable as absent + } + + // Idempotency: already contains our marker. + if strings.Contains(existing.Command, titleTeeMarker) { + return nil + } + + // Build new title config, wrapping any pre-existing command. + cfg := titleConfig{ + Type: "command", + Command: titleTeeCommand(localDev, existing.Command), + } + + cfgBytes, err := jsonutil.MarshalWithNoHTMLEscape(cfg) + if err != nil { + return fmt.Errorf("failed to marshal title config: %w", err) + } + rawFile["title"] = cfgBytes + + return writeAgySettings(rawFile, settingsPath) +} + +// TitleTeeInstalled reports whether agy's global settings.json declares a +// title command containing the title-tee marker. It is used by `entire doctor` +// to warn when Antigravity hooks are installed in a repo but the global title +// slot — agy's only token-usage surface — has not been claimed, which would +// leave token counts missing from checkpoints. A missing or unparseable +// settings file reports false. +func TitleTeeInstalled() bool { + cfgDir, err := agyConfigDir() + if err != nil { + return false + } + settingsPath := filepath.Join(cfgDir, "settings.json") + + rawFile, err := readAgySettings(settingsPath) + if err != nil { + return false + } + + raw, ok := rawFile["title"] + if !ok { + return false + } + + var existing titleConfig + if err := json.Unmarshal(raw, &existing); err != nil { + return false + } + return strings.Contains(existing.Command, titleTeeMarker) +} + +// UninstallTitleTee removes or restores the title entry in agy's global settings.json: +// - bare tee (no --wrap) → delete "title" key +// - tee with --wrap 'X' → restore X +// - any other (foreign) cmd → leave untouched +// - missing settings file → no-op +func UninstallTitleTee() error { + cfgDir, err := agyConfigDir() + if err != nil { + return err + } + settingsPath := filepath.Join(cfgDir, "settings.json") + + // Missing file → nothing to uninstall. + if _, err := os.Stat(settingsPath); os.IsNotExist(err) { + return nil + } + + rawFile, err := readAgySettings(settingsPath) + if err != nil { + return err + } + + raw, ok := rawFile["title"] + if !ok { + return nil // no title key — nothing to do + } + + var existing titleConfig + if err := json.Unmarshal(raw, &existing); err != nil { + return nil //nolint:nilerr // unparseable title entry — leave it alone rather than destroying user data + } + + // Not our command → leave untouched. + if !strings.Contains(existing.Command, titleTeeMarker) { + return nil + } + + wrapped, hasWrap := extractWrappedCommand(existing.Command) + if hasWrap { + // Restore the original command. + restored := titleConfig{ + Type: "command", + Command: wrapped, + } + restoredBytes, err := jsonutil.MarshalWithNoHTMLEscape(restored) + if err != nil { + return fmt.Errorf("failed to marshal restored title config: %w", err) + } + rawFile["title"] = restoredBytes + } else { + // Only delete the key when it is exactly one of the canonical bare-tee + // strings we would have written ourselves. Anything else containing the + // marker is either a user-authored wrapper (e.g. "my-wrapper.sh 'entire + // hooks antigravity title-tee'") or a corrupted command — leaving it + // alone is always safer than deleting the user's config. + bare := existing.Command == titleTeeCommand(false, "") || + existing.Command == titleTeeCommand(true, "") + if !bare { + return nil + } + delete(rawFile, "title") + } + + return writeAgySettings(rawFile, settingsPath) +} + +// extractWrappedCommand parses the --wrap '' portion of a title-tee +// command string. It returns the original command and true if found and valid, +// or ("", false) otherwise. +func extractWrappedCommand(command string) (string, bool) { + const wrapFlag = " --wrap " + idx := strings.Index(command, wrapFlag) + if idx < 0 { + return "", false + } + rest := strings.TrimSpace(command[idx+len(wrapFlag):]) + if len(rest) < 2 || rest[0] != '\'' || rest[len(rest)-1] != '\'' { + return "", false + } + // Strip outer single quotes and reverse the '\'' escaping. + inner := rest[1 : len(rest)-1] + return strings.ReplaceAll(inner, `'\''`, "'"), true +} + +// readAgySettings reads and parses settings.json into a raw map. +// A missing file returns an empty map (not an error). +func readAgySettings(settingsPath string) (map[string]json.RawMessage, error) { + rawFile := make(map[string]json.RawMessage) + data, err := os.ReadFile(settingsPath) //nolint:gosec // path is constructed from config dir + fixed filename + if os.IsNotExist(err) { + return rawFile, nil + } + if err != nil { + return nil, fmt.Errorf("failed to read agy settings: %w", err) + } + if err := json.Unmarshal(data, &rawFile); err != nil { + return nil, fmt.Errorf("failed to parse agy settings: %w", err) + } + return rawFile, nil +} + +// writeAgySettings marshals rawFile and writes it to settingsPath, creating +// parent directories as needed. +func writeAgySettings(rawFile map[string]json.RawMessage, settingsPath string) error { + if err := os.MkdirAll(filepath.Dir(settingsPath), 0o750); err != nil { + return fmt.Errorf("failed to create agy config directory: %w", err) + } + + output, err := jsonutil.MarshalIndentWithNewline(rawFile, "", " ") + if err != nil { + return fmt.Errorf("failed to marshal agy settings: %w", err) + } + + if err := os.WriteFile(settingsPath, output, 0o600); err != nil { + return fmt.Errorf("failed to write agy settings: %w", err) + } + return nil +} diff --git a/cmd/entire/cli/agent/antigravity/title_install_test.go b/cmd/entire/cli/agent/antigravity/title_install_test.go new file mode 100644 index 000000000..533c70ab2 --- /dev/null +++ b/cmd/entire/cli/agent/antigravity/title_install_test.go @@ -0,0 +1,323 @@ +package antigravity + +import ( + "encoding/json" + "os" + "path/filepath" + "strings" + "testing" +) + +// writeAgySettingsFile writes content to /settings.json. +func writeAgySettingsFile(t *testing.T, dir, content string) { + t.Helper() + if err := os.MkdirAll(dir, 0o750); err != nil { + t.Fatalf("writeAgySettingsFile: mkdir: %v", err) + } + if err := os.WriteFile(filepath.Join(dir, "settings.json"), []byte(content), 0o600); err != nil { + t.Fatalf("writeAgySettingsFile: write: %v", err) + } +} + +// readTitleCommand parses /settings.json and returns the title.command value, +// or "" if the file is absent or the key is not present. +func readTitleCommand(t *testing.T, dir string) string { + t.Helper() + data, err := os.ReadFile(filepath.Join(dir, "settings.json")) + if err != nil { + return "" + } + var s struct { + Title *struct { + Type string `json:"type"` + Command string `json:"command"` + } `json:"title"` + Theme string `json:"theme"` + } + if err := json.Unmarshal(data, &s); err != nil { + t.Fatalf("readTitleCommand: unmarshal: %v", err) + } + if s.Title == nil { + return "" + } + return s.Title.Command +} + +// mustJSON marshals s to a JSON string value (quoted). +func mustJSON(t *testing.T, s string) []byte { + t.Helper() + b, err := json.Marshal(s) + if err != nil { + t.Fatalf("mustJSON: %v", err) + } + return b +} + +func TestInstallTitle_FreshConfig(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + if err := InstallTitleTee(false); err != nil { + t.Fatalf("InstallTitleTee: %v", err) + } + + got := readTitleCommand(t, cfgDir) + want := "entire hooks antigravity title-tee" + if got != want { + t.Errorf("title.command = %q, want %q", got, want) + } +} + +func TestInstallTitle_WrapsExistingCommand(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + writeAgySettingsFile(t, cfgDir, `{"theme":"dark","title":{"type":"command","command":"~/bin/my-status.sh"}}`) + + if err := InstallTitleTee(false); err != nil { + t.Fatalf("InstallTitleTee: %v", err) + } + + got := readTitleCommand(t, cfgDir) + want := "entire hooks antigravity title-tee --wrap '~/bin/my-status.sh'" + if got != want { + t.Errorf("title.command = %q, want %q", got, want) + } + + // Unknown-key preservation: "theme" must still be present in raw file. + raw, err := os.ReadFile(filepath.Join(cfgDir, "settings.json")) + if err != nil { + t.Fatalf("read settings.json: %v", err) + } + if !strings.Contains(string(raw), `"theme"`) { + t.Error(`settings.json lost "theme" key after install`) + } +} + +func TestInstallTitle_Idempotent(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + if err := InstallTitleTee(false); err != nil { + t.Fatalf("first InstallTitleTee: %v", err) + } + first := readTitleCommand(t, cfgDir) + + if err := InstallTitleTee(false); err != nil { + t.Fatalf("second InstallTitleTee: %v", err) + } + second := readTitleCommand(t, cfgDir) + + if first != second { + t.Errorf("idempotency: first=%q second=%q", first, second) + } +} + +func TestUninstallTitle_RestoresOriginal(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + writeAgySettingsFile(t, cfgDir, `{"title":{"type":"command","command":"entire hooks antigravity title-tee --wrap '~/bin/my-status.sh'"}}`) + + if err := UninstallTitleTee(); err != nil { + t.Fatalf("UninstallTitleTee: %v", err) + } + + got := readTitleCommand(t, cfgDir) + want := "~/bin/my-status.sh" + if got != want { + t.Errorf("title.command after uninstall = %q, want %q", got, want) + } +} + +func TestUninstallTitle_RemovesBareTee(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + writeAgySettingsFile(t, cfgDir, `{"title":{"type":"command","command":"entire hooks antigravity title-tee"}}`) + + if err := UninstallTitleTee(); err != nil { + t.Fatalf("UninstallTitleTee: %v", err) + } + + got := readTitleCommand(t, cfgDir) + if got != "" { + t.Errorf("title.command after bare-tee uninstall = %q, want empty", got) + } +} + +func TestUninstallTitle_LeavesForeignCommandAlone(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + writeAgySettingsFile(t, cfgDir, `{"title":{"type":"command","command":"~/bin/my-status.sh"}}`) + + if err := UninstallTitleTee(); err != nil { + t.Fatalf("UninstallTitleTee: %v", err) + } + + got := readTitleCommand(t, cfgDir) + want := "~/bin/my-status.sh" + if got != want { + t.Errorf("title.command after foreign-command uninstall = %q, want %q", got, want) + } +} + +func TestUninstallTitle_LeavesUserWrappedTeeAlone(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + // User-authored wrapper that happens to contain the marker string. + const cmd = "my-wrapper.sh 'entire hooks antigravity title-tee'" + writeAgySettingsFile(t, cfgDir, `{"title":{"type":"command","command":"`+cmd+`"}}`) + + if err := UninstallTitleTee(); err != nil { + t.Fatalf("UninstallTitleTee: %v", err) + } + + got := readTitleCommand(t, cfgDir) + if got != cmd { + t.Errorf("title.command = %q, want %q (user wrapper should be left alone)", got, cmd) + } +} + +func TestUninstallTitle_LeavesMalformedWrapAlone(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + // Contains the marker + a --wrap flag but unquoted (malformed) — safer to leave than delete. + const cmd = "entire hooks antigravity title-tee --wrap unquoted" + writeAgySettingsFile(t, cfgDir, `{"title":{"type":"command","command":"`+cmd+`"}}`) + + if err := UninstallTitleTee(); err != nil { + t.Fatalf("UninstallTitleTee: %v", err) + } + + got := readTitleCommand(t, cfgDir) + if got != cmd { + t.Errorf("title.command = %q, want %q (malformed wrap should be left alone)", got, cmd) + } +} + +func TestInstallTitle_WrapsCommandContainingWrapSubstring(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + // Original command itself contains "--wrap fancy" — must survive round-trip. + const original = "~/bin/title.sh --wrap fancy" + origJSON := mustJSON(t, original) + content := `{"title":{"type":"command","command":` + string(origJSON) + `}}` + writeAgySettingsFile(t, cfgDir, content) + + if err := InstallTitleTee(false); err != nil { + t.Fatalf("InstallTitleTee: %v", err) + } + + // Installed command should wrap the original. + installed := readTitleCommand(t, cfgDir) + want := "entire hooks antigravity title-tee --wrap '~/bin/title.sh --wrap fancy'" + if installed != want { + t.Errorf("installed title.command = %q, want %q", installed, want) + } + + // Uninstall should restore the original exactly. + if err := UninstallTitleTee(); err != nil { + t.Fatalf("UninstallTitleTee: %v", err) + } + got := readTitleCommand(t, cfgDir) + if got != original { + t.Errorf("round-trip: got %q, want %q", got, original) + } +} + +func TestUninstallTitle_RemovesBareLocalDevTee(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + // Exactly the localDev tee command — must be removed. + localDevCmd := titleTeeCommand(true, "") + cmdJSON := mustJSON(t, localDevCmd) + content := `{"title":{"type":"command","command":` + string(cmdJSON) + `}}` + writeAgySettingsFile(t, cfgDir, content) + + if err := UninstallTitleTee(); err != nil { + t.Fatalf("UninstallTitleTee: %v", err) + } + + got := readTitleCommand(t, cfgDir) + if got != "" { + t.Errorf("title.command after localDev bare-tee uninstall = %q, want empty", got) + } +} + +func TestInstallUninstall_RoundTripsEmbeddedSingleQuotes(t *testing.T) { + // No t.Parallel — uses t.Setenv + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + + original := `echo 'hi there' | awk '{print $1}'` + origJSON := mustJSON(t, original) + content := `{"title":{"type":"command","command":` + string(origJSON) + `}}` + writeAgySettingsFile(t, cfgDir, content) + + if err := InstallTitleTee(false); err != nil { + t.Fatalf("InstallTitleTee: %v", err) + } + if err := UninstallTitleTee(); err != nil { + t.Fatalf("UninstallTitleTee: %v", err) + } + + got := readTitleCommand(t, cfgDir) + if got != original { + t.Errorf("round-trip: got %q, want %q", got, original) + } +} + +// TestTitleTeeInstalled covers the three states the doctor check cares about: +// our marker present (true), no title key (false), and a foreign command (false). +func TestTitleTeeInstalled(t *testing.T) { + t.Run("configured", func(t *testing.T) { + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + if err := InstallTitleTee(false); err != nil { + t.Fatalf("InstallTitleTee: %v", err) + } + if !TitleTeeInstalled() { + t.Error("TitleTeeInstalled() = false, want true after InstallTitleTee") + } + }) + + t.Run("absent", func(t *testing.T) { + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + // No settings.json at all. + if TitleTeeInstalled() { + t.Error("TitleTeeInstalled() = true, want false with no settings file") + } + // settings.json with no title key. + writeAgySettingsFile(t, cfgDir, `{"theme":"dark"}`) + if TitleTeeInstalled() { + t.Error("TitleTeeInstalled() = true, want false with no title key") + } + }) + + t.Run("foreign command", func(t *testing.T) { + cfgDir := t.TempDir() + t.Setenv(configDirEnv, cfgDir) + writeAgySettingsFile(t, cfgDir, + `{"title":{"type":"command","command":"my-own-title-script.sh"}}`) + if TitleTeeInstalled() { + t.Error("TitleTeeInstalled() = true, want false for a foreign command") + } + }) +} diff --git a/cmd/entire/cli/agent/capabilities.go b/cmd/entire/cli/agent/capabilities.go index b4677f7bf..f74890dc8 100644 --- a/cmd/entire/cli/agent/capabilities.go +++ b/cmd/entire/cli/agent/capabilities.go @@ -90,6 +90,24 @@ func AsTokenCalculator(ag Agent) (TokenCalculator, bool) { return tc, true } +// AsOutOfBandTokenSource returns the agent as OutOfBandTokenSource if supported. +// External (CapabilityDeclarer) agents are excluded because the out-of-band +// store is fed by a built-in shim subcommand they cannot provide, and +// DeclaredCaps has no field for this capability to opt into. +func AsOutOfBandTokenSource(ag Agent) (OutOfBandTokenSource, bool) { + if ag == nil { + return nil, false + } + src, ok := ag.(OutOfBandTokenSource) + if !ok { + return nil, false + } + if _, isDeclarer := ag.(CapabilityDeclarer); isDeclarer { + return nil, false + } + return src, true +} + // AsTextGenerator returns the agent as TextGenerator if it both // implements the interface and (for CapabilityDeclarer agents) has declared the capability. func AsTextGenerator(ag Agent) (TextGenerator, bool) { diff --git a/cmd/entire/cli/agent/capabilities_test.go b/cmd/entire/cli/agent/capabilities_test.go index 5e884e08a..e99fa2a99 100644 --- a/cmd/entire/cli/agent/capabilities_test.go +++ b/cmd/entire/cli/agent/capabilities_test.go @@ -2,6 +2,7 @@ package agent import ( "context" + "encoding/json" "io" "testing" @@ -83,6 +84,15 @@ func (m *mockFullAgent) GenerateText(context.Context, string, string) (string, e return "", nil } +// OutOfBandTokenSource (mockFullAgent is a CapabilityDeclarer, so AsOutOfBandTokenSource +// must still exclude it — verifies the built-in-only gate). +func (m *mockFullAgent) SnapshotTokenBaseline(context.Context, string) (json.RawMessage, error) { + return nil, nil +} +func (m *mockFullAgent) CalculateTokenUsageSince(context.Context, string, json.RawMessage) (*TokenUsage, error) { + return nil, nil //nolint:nilnil // test mock +} + // TranscriptCompactor func (m *mockFullAgent) CompactTranscript(context.Context, string) (*CompactedTranscript, error) { return &CompactedTranscript{}, nil @@ -108,6 +118,19 @@ func (m *mockBuiltinPromptAgent) ExtractPrompts(string, int) ([]string, error) { return []string{"test prompt"}, nil } +// mockBuiltinOOBAgent is a built-in agent that implements OutOfBandTokenSource +// but NOT CapabilityDeclarer. +type mockBuiltinOOBAgent struct { + mockBaseAgent +} + +func (m *mockBuiltinOOBAgent) SnapshotTokenBaseline(context.Context, string) (json.RawMessage, error) { + return nil, nil +} +func (m *mockBuiltinOOBAgent) CalculateTokenUsageSince(context.Context, string, json.RawMessage) (*TokenUsage, error) { + return nil, nil //nolint:nilnil // test mock +} + // --- Tests --- func TestAsHookSupport(t *testing.T) { @@ -360,6 +383,44 @@ func TestAsSubagentAwareExtractor(t *testing.T) { }) } +func TestAsOutOfBandTokenSource(t *testing.T) { + t.Parallel() + + t.Run("nil agent", func(t *testing.T) { + t.Parallel() + _, ok := AsOutOfBandTokenSource(nil) + if ok { + t.Error("expected false for nil agent") + } + }) + + t.Run("not implemented", func(t *testing.T) { + t.Parallel() + _, ok := AsOutOfBandTokenSource(&mockBaseAgent{}) + if ok { + t.Error("expected false for agent not implementing OutOfBandTokenSource") + } + }) + + t.Run("builtin agent", func(t *testing.T) { + t.Parallel() + src, ok := AsOutOfBandTokenSource(&mockBuiltinOOBAgent{}) + if !ok || src == nil { + t.Error("expected true for built-in agent implementing OutOfBandTokenSource") + } + }) + + t.Run("capability declarer excluded", func(t *testing.T) { + t.Parallel() + // mockFullAgent implements the interface but is a CapabilityDeclarer + // (external agent), so it must be excluded. + _, ok := AsOutOfBandTokenSource(&mockFullAgent{}) + if ok { + t.Error("expected false for CapabilityDeclarer agent") + } + }) +} + func TestAsPromptExtractor(t *testing.T) { t.Parallel() diff --git a/cmd/entire/cli/doctor.go b/cmd/entire/cli/doctor.go index fa85fa9b4..aff5bda59 100644 --- a/cmd/entire/cli/doctor.go +++ b/cmd/entire/cli/doctor.go @@ -10,6 +10,7 @@ import ( "time" "charm.land/huh/v2" + "github.com/entireio/cli/cmd/entire/cli/agent/antigravity" "github.com/entireio/cli/cmd/entire/cli/agent/codex" "github.com/entireio/cli/cmd/entire/cli/checkpoint" "github.com/entireio/cli/cmd/entire/cli/paths" @@ -97,6 +98,9 @@ func runSessionsFix(cmd *cobra.Command, force bool) error { // Agent-specific: Codex hook trust state. checkCodexHookTrust(cmd) + // Agent-specific: Antigravity title-tee (token-usage surface). + checkAntigravityTitleTee(cmd) + // Stuck sessions // Load all session states states, err := strategy.ListSessionStates(ctx) @@ -434,6 +438,30 @@ func checkCodexHookTrust(cmd *cobra.Command) { } } +// checkAntigravityTitleTee warns when Antigravity hooks are installed in this +// repo but agy's global title slot has NOT been claimed by the title-tee shim. +// agy exposes token usage only through the title/statusline state JSON, so a +// missing title-tee means token counts will be absent from every Antigravity +// checkpoint. Stays silent when Antigravity hooks aren't installed here. +// Warn-only. +func checkAntigravityTitleTee(cmd *cobra.Command) { + ag := &antigravity.AntigravityAgent{} + if !ag.AreHooksInstalled(cmd.Context()) { + return + } + + w := cmd.OutOrStdout() + if antigravity.TitleTeeInstalled() { + fmt.Fprintln(w, "✓ Antigravity title-tee: OK") + return + } + + fmt.Fprintln(w, "Antigravity title-tee: NOT CONFIGURED") + fmt.Fprintln(w, " agy's title command isn't routed through Entire, so token counts") + fmt.Fprintln(w, " will be missing for Antigravity checkpoints.") + fmt.Fprintln(w, " Re-run agent setup (`entire agent add`) to configure it.") +} + // canDeleteShadowBranch checks if a shadow branch can be safely deleted. // Returns true if no other sessions (besides excludeSessionID) need this branch. func canDeleteShadowBranch(ctx context.Context, shadowBranch, excludeSessionID string) (bool, error) { diff --git a/cmd/entire/cli/doctor_test.go b/cmd/entire/cli/doctor_test.go index f57f27d45..65bd48c8b 100644 --- a/cmd/entire/cli/doctor_test.go +++ b/cmd/entire/cli/doctor_test.go @@ -549,3 +549,66 @@ trusted_hash = "sha256:ccc" require.Contains(t, out, "entire enable") require.NotContains(t, out, "Codex hook trust: REVIEW NEEDED") } + +// antigravityHooksJSON returns a minimal .agents/hooks.json declaring the +// Entire PreInvocation hook, enough for AreHooksInstalled to report true. +func antigravityHooksJSON() string { + return `{"entire":{"PreInvocation":[{"type":"command","command":"entire hooks antigravity pre-invocation"}]}}` +} + +// TestCheckAntigravityTitleTee_SilentWhenHooksNotInstalled stays quiet when +// the repo has no Antigravity hooks — nothing to check. +func TestCheckAntigravityTitleTee_SilentWhenHooksNotInstalled(t *testing.T) { + dir := setupGitRepoForPhaseTest(t) + t.Chdir(dir) + t.Setenv("ENTIRE_ANTIGRAVITY_CONFIG_DIR", filepath.Join(t.TempDir(), "agy")) + + cmd, stdout := newTestCmd(t) + checkAntigravityTitleTee(cmd) + require.NotContains(t, stdout.String(), "Antigravity title-tee") +} + +// TestCheckAntigravityTitleTee_OKWhenConfigured reports OK when hooks are +// installed and agy's title slot routes through the title-tee shim. +func TestCheckAntigravityTitleTee_OKWhenConfigured(t *testing.T) { + dir := setupGitRepoForPhaseTest(t) + t.Chdir(dir) + + agentsDir := filepath.Join(dir, ".agents") + require.NoError(t, os.MkdirAll(agentsDir, 0o750)) + require.NoError(t, os.WriteFile(filepath.Join(agentsDir, "hooks.json"), + []byte(antigravityHooksJSON()), 0o600)) + + cfgDir := filepath.Join(t.TempDir(), "agy") + require.NoError(t, os.MkdirAll(cfgDir, 0o750)) + require.NoError(t, os.WriteFile(filepath.Join(cfgDir, "settings.json"), + []byte(`{"title":{"type":"command","command":"entire hooks antigravity title-tee"}}`), 0o600)) + t.Setenv("ENTIRE_ANTIGRAVITY_CONFIG_DIR", cfgDir) + + cmd, stdout := newTestCmd(t) + checkAntigravityTitleTee(cmd) + require.Contains(t, stdout.String(), "✓ Antigravity title-tee: OK") +} + +// TestCheckAntigravityTitleTee_WarnsWhenNotConfigured surfaces the missing +// token-usage surface when hooks are installed but the title slot is unclaimed. +func TestCheckAntigravityTitleTee_WarnsWhenNotConfigured(t *testing.T) { + dir := setupGitRepoForPhaseTest(t) + t.Chdir(dir) + + agentsDir := filepath.Join(dir, ".agents") + require.NoError(t, os.MkdirAll(agentsDir, 0o750)) + require.NoError(t, os.WriteFile(filepath.Join(agentsDir, "hooks.json"), + []byte(antigravityHooksJSON()), 0o600)) + + // Empty agy config dir — no title slot claimed. + t.Setenv("ENTIRE_ANTIGRAVITY_CONFIG_DIR", filepath.Join(t.TempDir(), "agy")) + + cmd, stdout := newTestCmd(t) + checkAntigravityTitleTee(cmd) + + out := stdout.String() + require.Contains(t, out, "Antigravity title-tee: NOT CONFIGURED") + require.Contains(t, out, "token counts") + require.Contains(t, out, "entire agent add") +} diff --git a/cmd/entire/cli/hooks_antigravity_title.go b/cmd/entire/cli/hooks_antigravity_title.go new file mode 100644 index 000000000..69ce23e5b --- /dev/null +++ b/cmd/entire/cli/hooks_antigravity_title.go @@ -0,0 +1,60 @@ +package cli + +import ( + "bytes" + "io" + "os/exec" + + "github.com/entireio/cli/cmd/entire/cli/agent/antigravity" + "github.com/entireio/cli/cmd/entire/cli/logging" + "github.com/spf13/cobra" +) + +// newAntigravityTitleTeeCmd implements `entire hooks antigravity title-tee`. +// +// Antigravity invokes the configured title command on every agent state +// change, piping a state JSON (the only agy surface exposing token usage — +// same payload as the statusline script) to stdin. This command tees that +// JSON into the snapshot store and, with --wrap, pipes it through to the +// user's original title command so their window title is preserved. +// +// Contract: NEVER exit non-zero and NEVER write noise to stdout — stdout is +// rendered verbatim as the terminal window title. It also must work outside +// git repos and without entire being enabled (the title config is global). +func newAntigravityTitleTeeCmd() *cobra.Command { + var wrap string + cmd := &cobra.Command{ + Use: "title-tee", + Short: "Tee agy state JSON (title/statusline payload) into the token snapshot store", + Hidden: true, + // NoArgs is safe: agy invokes the title command with stdin only, never positional args. + Args: cobra.NoArgs, + RunE: func(cmd *cobra.Command, _ []string) error { + payload, err := io.ReadAll(cmd.InOrStdin()) + if err != nil { + return nil //nolint:nilerr // never break agy's title rendering + } + if err := antigravity.AppendStatusSnapshot(payload); err != nil { + // Best-effort capture: a persistent I/O failure (disk full, + // unwritable cache dir) leaves a breadcrumb without ever + // affecting stdout or the exit code. + logging.Debug(cmd.Context(), "antigravity title-tee: snapshot append failed", "error", err.Error()) + } + + if wrap == "" { + return nil + } + // wrap is the user's own original title command, preserved from + // settings.json and round-tripped via shellSingleQuote, so running + // it under `sh -c` is intentional — not an external-input injection surface. + wrapped := exec.CommandContext(cmd.Context(), "sh", "-c", wrap) + wrapped.Stdin = bytes.NewReader(payload) + wrapped.Stdout = cmd.OutOrStdout() + wrapped.Stderr = cmd.ErrOrStderr() + _ = wrapped.Run() //nolint:errcheck // a failing user title script must not fail the tee + return nil + }, + } + cmd.Flags().StringVar(&wrap, "wrap", "", "original title command to chain after capturing") + return cmd +} diff --git a/cmd/entire/cli/hooks_antigravity_title_test.go b/cmd/entire/cli/hooks_antigravity_title_test.go new file mode 100644 index 000000000..25ebbedfb --- /dev/null +++ b/cmd/entire/cli/hooks_antigravity_title_test.go @@ -0,0 +1,84 @@ +package cli + +import ( + "bytes" + "os" + "path/filepath" + "strings" + "testing" +) + +// Note: these tests use t.Setenv, so t.Parallel() is not called. + +func TestTitleTee_WritesSnapshotAndStaysSilent(t *testing.T) { + dir := t.TempDir() + t.Setenv("ENTIRE_ANTIGRAVITY_STATUS_DIR", dir) + + payload := `{"conversation_id":"conv-9","agent_state":"working","context_window":{"total_input_tokens":500,"total_output_tokens":25,"context_window_size":100000,"current_usage":{"input_tokens":400,"output_tokens":25,"cache_creation_input_tokens":50,"cache_read_input_tokens":300}}}` + + cmd := newAntigravityTitleTeeCmd() + + var out bytes.Buffer + cmd.SetOut(&out) + cmd.SetIn(strings.NewReader(payload)) + + if err := cmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + // stdout must be empty — agy renders it verbatim as the window title + if out.Len() != 0 { + t.Errorf("stdout not empty: %q", out.String()) + } + + // snapshot file must exist + snapFile := filepath.Join(dir, "conv-9.jsonl") + if _, err := os.Stat(snapFile); err != nil { + t.Errorf("snapshot file not created: %v", err) + } +} + +func TestTitleTee_WrapStillCapturesSnapshot(t *testing.T) { + dir := t.TempDir() + t.Setenv("ENTIRE_ANTIGRAVITY_STATUS_DIR", dir) + + payload := `{"conversation_id":"conv-11","agent_state":"working","context_window":{"total_input_tokens":700,"total_output_tokens":40}}` + + cmd := newAntigravityTitleTeeCmd() + cmd.SetArgs([]string{"--wrap", "cat"}) + + var out bytes.Buffer + cmd.SetOut(&out) + cmd.SetIn(strings.NewReader(payload)) + + if err := cmd.Execute(); err != nil { + t.Fatalf("Execute: %v", err) + } + + // --wrap cat must pipe the payload through verbatim... + if got, want := strings.TrimSpace(out.String()), strings.TrimSpace(payload); got != want { + t.Errorf("stdout = %q, want %q", got, want) + } + + // ...AND the snapshot must still be captured. + snapFile := filepath.Join(dir, "conv-11.jsonl") + if _, err := os.Stat(snapFile); err != nil { + t.Errorf("snapshot file not created under --wrap: %v", err) + } +} + +func TestTitleTee_GarbageInputAndFailingWrapNeverError(t *testing.T) { + dir := t.TempDir() + t.Setenv("ENTIRE_ANTIGRAVITY_STATUS_DIR", dir) + + cmd := newAntigravityTitleTeeCmd() + cmd.SetArgs([]string{"--wrap", "false"}) + + var out bytes.Buffer + cmd.SetOut(&out) + cmd.SetIn(strings.NewReader("not json")) + + if err := cmd.Execute(); err != nil { + t.Fatalf("Execute returned error: %v", err) + } +} diff --git a/cmd/entire/cli/hooks_cmd.go b/cmd/entire/cli/hooks_cmd.go index d819cd6c6..c6cfc1473 100644 --- a/cmd/entire/cli/hooks_cmd.go +++ b/cmd/entire/cli/hooks_cmd.go @@ -68,7 +68,13 @@ func newHooksCmd() *cobra.Command { continue } if handler, ok := agent.AsHookSupport(ag); ok { - cmd.AddCommand(newAgentHooksCmd(agentName, handler)) + sub := newAgentHooksCmd(agentName, handler) + // title-tee is not a lifecycle verb: it runs globally (outside + // git repos, without the enabled check) and owns its stdout. + if agentName == agent.AgentNameAntigravity { + sub.AddCommand(newAntigravityTitleTeeCmd()) + } + cmd.AddCommand(sub) } } diff --git a/cmd/entire/cli/integration_test/antigravity_test.go b/cmd/entire/cli/integration_test/antigravity_test.go index 09e914f59..9ee222ac3 100644 --- a/cmd/entire/cli/integration_test/antigravity_test.go +++ b/cmd/entire/cli/integration_test/antigravity_test.go @@ -11,7 +11,9 @@ import ( "slices" "testing" + "github.com/entireio/cli/cmd/entire/cli/checkpoint" "github.com/entireio/cli/cmd/entire/cli/execx" + "github.com/entireio/cli/cmd/entire/cli/paths" "github.com/entireio/cli/cmd/entire/cli/testutil" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -119,7 +121,144 @@ func TestAntigravity_FullEventFlow(t *testing.T) { "PreToolUse(write_to_file foo.txt) should have populated files_touched; got %v", got) } +// TestAntigravity_TokenUsageInCheckpointMetadata proves the entire.io UI +// contract end-to-end: agy's out-of-band token counts (the title-script +// context_window, captured into the snapshot store) flow through TurnStart +// baseline → TurnEnd delta → SaveStep → SessionState → condensation → the +// committed checkpoint metadata on entire/checkpoints/v1. +// +// The snapshot totals are cumulative per conversation, so the per-turn delta is +// (latest − baseline). Baseline here is 1000/100 (input/output) captured at +// TurnStart; the latest before Stop is 4500/400 → expected delta 3500/300. +func TestAntigravity_TokenUsageInCheckpointMetadata(t *testing.T) { + t.Parallel() + env := NewFeatureBranchEnv(t) + + statusDir := t.TempDir() + configDir := t.TempDir() + conversationID := "antigravity-tokens-it" + extraEnv := []string{ + "ENTIRE_ANTIGRAVITY_STATUS_DIR=" + statusDir, + "ENTIRE_ANTIGRAVITY_CONFIG_DIR=" + configDir, + } + + snapshotPath := filepath.Join(statusDir, conversationID+".jsonl") + appendSnapshotLine := func(line string) { + f, err := os.OpenFile(snapshotPath, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600) + require.NoError(t, err) + _, werr := f.WriteString(line + "\n") + require.NoError(t, f.Close()) + require.NoError(t, werr) + } + + transcriptPath := filepath.Join(env.RepoDir, ".gemini", "antigravity-cli", + "brain", conversationID, ".system_generated", "logs", "transcript.jsonl") + require.NoError(t, os.MkdirAll(filepath.Dir(transcriptPath), 0o750)) + require.NoError(t, os.WriteFile(transcriptPath, + []byte(`{"step_index":0,"source":"USER_EXPLICIT","type":"USER_INPUT","status":"DONE","content":"create tok.txt"}`+"\n"), + 0o600)) + + common := map[string]any{ + "conversationId": conversationID, + "workspacePaths": []string{env.RepoDir}, + "transcriptPath": transcriptPath, + "artifactDirectoryPath": filepath.Join(env.RepoDir, ".gemini", "antigravity-cli", "artifacts"), + } + + // 1. Seed the BASELINE snapshot before TurnStart captures it. + appendSnapshotLine(`{"ts":"2026-06-03T10:00:00Z","conversation_id":"` + conversationID + + `","context_window":{"total_input_tokens":1000,"total_output_tokens":100,` + + `"current_usage":{"cache_read_input_tokens":700}}}`) + + // 2. TurnStart (invocationNum 0) → baseline captured = 1000/100. + preInv := mergeMaps(common, map[string]any{ + "invocationNum": 0, + "initialNumSteps": 1, + }) + require.NoError(t, runAntigravityHookWithEnv(t, env.RepoDir, "pre-invocation", preInv, extraEnv), + "pre-invocation should capture the token baseline") + + // 3. PreToolUse writing a real file, so the turn has a working-tree change + // to checkpoint and condense. + env.WriteFile("tok.txt", "token test contents\n") + preTU := mergeMaps(common, map[string]any{ + "toolCall": map[string]any{ + "name": "write_to_file", + "args": map[string]any{"TargetFile": "tok.txt", "Overwrite": false}, + }, + "stepIdx": 1, + }) + require.NoError(t, runAntigravityHookWithEnv(t, env.RepoDir, "pre-tool-use", preTU, extraEnv), + "pre-tool-use should record tok.txt in files_touched") + + // 4. Append the SECOND (latest) snapshot: cumulative totals grew. + appendSnapshotLine(`{"ts":"2026-06-03T10:05:00Z","conversation_id":"` + conversationID + + `","context_window":{"total_input_tokens":4500,"total_output_tokens":400,` + + `"current_usage":{"cache_read_input_tokens":1500,"cache_creation_input_tokens":50}}}`) + + // 5. Stop (fullyIdle true) → TurnEnd → OOB delta 3500/300 → SaveStep. + stopIdle := mergeMaps(common, map[string]any{ + "executionNum": 1, + "terminationReason": "model_stop", + "error": "", + "fullyIdle": true, + }) + require.NoError(t, runAntigravityHookWithEnv(t, env.RepoDir, "stop", stopIdle, extraEnv), + "stop hook should finalize the turn and accumulate token usage") + + // Sanity: SessionState shows the accumulated delta before commit. This + // isolates the OOB capture+delta+SaveStep path from condensation. + statePath := filepath.Join(env.RepoDir, ".git", "entire-sessions", conversationID+".json") + stateBytes, err := os.ReadFile(statePath) + require.NoError(t, err) + var state struct { + TokenUsage *struct { + InputTokens int `json:"input_tokens"` + OutputTokens int `json:"output_tokens"` + } `json:"token_usage"` + } + require.NoError(t, json.Unmarshal(stateBytes, &state)) + require.NotNil(t, state.TokenUsage, "session state should record token usage after stop") + assert.Equal(t, 3500, state.TokenUsage.InputTokens, "session state input tokens") + assert.Equal(t, 300, state.TokenUsage.OutputTokens, "session state output tokens") + + // 6. Commit: prepare-commit-msg adds the Entire-Checkpoint trailer, then + // post-commit condenses the session (with its token usage) onto + // entire/checkpoints/v1. + env.GitCommitWithShadowHooks("Add tok.txt", "tok.txt") + + // 7. Read the committed checkpoint metadata and assert the token usage + // reached both the per-session metadata.json and the summary aggregate. + checkpointID := env.GetLatestCheckpointID() + + sessionMeta, found := env.ReadFileFromBranch(paths.MetadataBranchName, SessionMetadataPath(checkpointID)) + require.True(t, found, "per-session metadata.json should exist on the metadata branch") + var sessionMetadata checkpoint.CommittedMetadata + require.NoError(t, json.Unmarshal([]byte(sessionMeta), &sessionMetadata)) + require.NotNil(t, sessionMetadata.TokenUsage, "committed per-session metadata should carry token_usage") + assert.Equal(t, 3500, sessionMetadata.TokenUsage.InputTokens, "committed per-session input tokens") + assert.Equal(t, 300, sessionMetadata.TokenUsage.OutputTokens, "committed per-session output tokens") + + summaryRaw, found := env.ReadFileFromBranch(paths.MetadataBranchName, CheckpointSummaryPath(checkpointID)) + require.True(t, found, "CheckpointSummary metadata.json should exist on the metadata branch") + var summary checkpoint.CheckpointSummary + require.NoError(t, json.Unmarshal([]byte(summaryRaw), &summary)) + require.NotNil(t, summary.TokenUsage, "CheckpointSummary should carry aggregated token_usage") + assert.Equal(t, 3500, summary.TokenUsage.InputTokens, "summary aggregate input tokens") + assert.Equal(t, 300, summary.TokenUsage.OutputTokens, "summary aggregate output tokens") +} + func runAntigravityHook(t *testing.T, repoDir, hookName string, input map[string]any) error { + t.Helper() + return runAntigravityHookWithEnv(t, repoDir, hookName, input, nil) +} + +// runAntigravityHookWithEnv runs an Antigravity hook subprocess with extraEnv +// appended to the isolated git env. The override (e.g. +// ENTIRE_ANTIGRAVITY_STATUS_DIR) must be on the child's env to reach the +// lifecycle code that reads token snapshots, since hooks run as a subprocess +// of the real entire binary. +func runAntigravityHookWithEnv(t *testing.T, repoDir, hookName string, input map[string]any, extraEnv []string) error { t.Helper() inputJSON, err := json.Marshal(input) require.NoError(t, err) @@ -127,7 +266,7 @@ func runAntigravityHook(t *testing.T, repoDir, hookName string, input map[string cmd := execx.NonInteractive(context.Background(), getTestBinary(), "hooks", "antigravity", hookName) cmd.Dir = repoDir cmd.Stdin = bytes.NewReader(inputJSON) - cmd.Env = testutil.GitIsolatedEnv() + cmd.Env = append(testutil.GitIsolatedEnv(), extraEnv...) output, runErr := cmd.CombinedOutput() t.Logf("antigravity hook %s output: %s", hookName, output) return runErr diff --git a/cmd/entire/cli/lifecycle.go b/cmd/entire/cli/lifecycle.go index beb2824cf..787fef80b 100644 --- a/cmd/entire/cli/lifecycle.go +++ b/cmd/entire/cli/lifecycle.go @@ -9,6 +9,7 @@ package cli import ( "context" + "encoding/json" "errors" "fmt" "log/slog" @@ -723,6 +724,25 @@ func handleLifecycleTurnEnd(ctx context.Context, ag agent.Agent, event *agent.Ev // Calculate token usage - prefer SubagentAwareExtractor to include subagent tokens tokenUsage := agent.CalculateTokenUsage(ctx, ag, transcriptData, transcriptLinesAtStart, subagentsDir) + // Out-of-band fallback: Antigravity exposes token usage only via its + // title/statusline pipe (captured by the title-tee shim), never in the + // transcript. Delta = current cumulative totals minus the TurnStart + // baseline stored in PrePromptState. + if tokenUsage == nil { + if src, ok := agent.AsOutOfBandTokenSource(ag); ok { + var baseline json.RawMessage + if preState != nil { + baseline = preState.TokenBaseline + } + oobUsage, oobErr := src.CalculateTokenUsageSince(ctx, sessionID, baseline) + if oobErr != nil { + logging.Warn(logCtx, "failed to compute out-of-band token usage", slog.String("error", oobErr.Error())) + } else { + tokenUsage = oobUsage + } + } + } + // Build fully-populated step context and delegate to strategy stepCtx := strategy.StepContext{ SessionID: sessionID, diff --git a/cmd/entire/cli/setup.go b/cmd/entire/cli/setup.go index 15532dd7a..0c16ffa85 100644 --- a/cmd/entire/cli/setup.go +++ b/cmd/entire/cli/setup.go @@ -10,6 +10,7 @@ import ( "strings" "github.com/entireio/cli/cmd/entire/cli/agent" + "github.com/entireio/cli/cmd/entire/cli/agent/antigravity" "github.com/entireio/cli/cmd/entire/cli/agent/external" "github.com/entireio/cli/cmd/entire/cli/agent/types" "github.com/entireio/cli/cmd/entire/cli/interactive" @@ -1189,6 +1190,16 @@ func runRemoveAgent(ctx context.Context, w io.Writer, name string) error { return fmt.Errorf("failed to remove %s hooks: %w", ag.Type(), err) } + // Antigravity's title tee lives in agy's GLOBAL settings.json, not in + // this repo — only remove it when the user removes the agent itself, + // never on per-repo disable. + if ag.Name() == agent.AgentNameAntigravity { + if err := antigravity.UninstallTitleTee(); err != nil { + logging.Warn(ctx, "failed to uninstall antigravity title tee", + "error", err.Error()) + } + } + fmt.Fprintf(w, "Removed %s hooks.\n", ag.Type()) return nil } diff --git a/cmd/entire/cli/state.go b/cmd/entire/cli/state.go index 09d5e4994..99a0d98ef 100644 --- a/cmd/entire/cli/state.go +++ b/cmd/entire/cli/state.go @@ -49,6 +49,13 @@ type PrePromptState struct { // Deprecated: LastTranscriptLineCount is the oldest name for transcript position. // Migrated to TranscriptOffset on load. LastTranscriptLineCount int `json:"last_transcript_line_count,omitempty"` + + // TokenBaseline is an opaque, agent-defined token position captured at turn + // start for OutOfBandTokenSource agents (currently Antigravity). At TurnEnd + // the lifecycle passes it back to CalculateTokenUsageSince to compute the + // checkpoint-scoped token delta. Empty for agents with transcript-embedded + // token data. + TokenBaseline json.RawMessage `json:"token_baseline,omitempty"` } // PreUntrackedFiles returns the untracked files list, or nil if the receiver is nil. @@ -137,6 +144,18 @@ func CapturePrePromptState(ctx context.Context, ag agent.Agent, sessionID, sessi TranscriptOffset: transcriptOffset, } + // Out-of-band token baseline: agents whose token usage lives outside the + // transcript (Antigravity) snapshot their cumulative token position now so + // TurnEnd can compute the per-checkpoint delta. + if src, ok := agent.AsOutOfBandTokenSource(ag); ok { + baseline, blErr := src.SnapshotTokenBaseline(ctx, sessionID) + if blErr != nil { + logging.Warn(logging.WithComponent(ctx, "state"), "failed to snapshot out-of-band token baseline", "error", blErr.Error()) + } else { + state.TokenBaseline = baseline + } + } + data, err := jsonutil.MarshalIndentWithNewline(state, "", " ") if err != nil { return fmt.Errorf("failed to marshal state: %w", err) diff --git a/cmd/entire/cli/strategy/manual_commit_condensation.go b/cmd/entire/cli/strategy/manual_commit_condensation.go index 3e58960cd..84755d459 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation.go @@ -144,6 +144,18 @@ func (s *ManualCommitStrategy) CondenseSession(ctx context.Context, repo *git.Re extractSessionDataSpan.End() extractDuration := time.Since(extractStart) + // Out-of-band token fallback: OutOfBandTokenSource agents (e.g. Antigravity) + // have their token usage captured out-of-band and accumulated into + // SessionState.TokenUsage at SaveStep time; the transcript recompute yields + // nil for them. Gate on the purpose-built capability so only those agents + // inherit the accumulated state value — transcript-based agents keep their + // recomputed, checkpoint-scoped values. + if !hasTokenUsageData(sessionData.TokenUsage) && hasTokenUsageData(state.TokenUsage) { + if _, ok := agent.AsOutOfBandTokenSource(ag); ok { + sessionData.TokenUsage = state.TokenUsage + } + } + // Backfill session state token usage from the freshly-extracted transcript. // Copilot CLI writes session.shutdown after the hooks return, so by condensation // time we can recover the authoritative full-session total from the transcript diff --git a/cmd/entire/cli/strategy/manual_commit_condensation_test.go b/cmd/entire/cli/strategy/manual_commit_condensation_test.go index b91b0dbb6..877b8cee5 100644 --- a/cmd/entire/cli/strategy/manual_commit_condensation_test.go +++ b/cmd/entire/cli/strategy/manual_commit_condensation_test.go @@ -22,6 +22,7 @@ import ( "github.com/go-git/go-git/v6/plumbing" // Register agents so GetByAgentType works in tests. + _ "github.com/entireio/cli/cmd/entire/cli/agent/antigravity" _ "github.com/entireio/cli/cmd/entire/cli/agent/claudecode" _ "github.com/entireio/cli/cmd/entire/cli/agent/copilotcli" _ "github.com/entireio/cli/cmd/entire/cli/agent/cursor" @@ -528,3 +529,192 @@ func TestCondenseSession_TagsCheckpointSummaryWithHasInvestigation(t *testing.T) require.Equal(t, "0123456789ab", meta.InvestigateRunID, "per-session InvestigateRunID") require.Equal(t, "Why is checkout flaky?", meta.InvestigateTopic, "per-session InvestigateTopic") } + +// TestCondenseSession_OutOfBandTokenFallback verifies that for an agent without +// transcript-embedded token data (Antigravity is not a TokenCalculator), a +// populated SessionState.TokenUsage flows through to the per-session +// CommittedMetadata.token_usage on the metadata branch. This is the out-of-band +// fallback: agy accumulates per-turn deltas into SessionState.TokenUsage at +// SaveStep time, and the transcript recompute yields nil, so without the +// fallback the UI would show no token counts. +// +// Tests in this file use t.Chdir for CWD-based git resolution, so this +// cannot be a parallel test. +func TestCondenseSession_OutOfBandTokenFallback(t *testing.T) { + dir := setupGitRepo(t) + t.Chdir(dir) + + repo, err := git.PlainOpen(dir) + require.NoError(t, err) + + s := &ManualCommitStrategy{} + sessionID := "2026-06-03-antigravity-tokens" + + metadataDir := ".entire/metadata/" + sessionID + metadataDirAbs := filepath.Join(dir, metadataDir) + require.NoError(t, os.MkdirAll(metadataDirAbs, 0o755)) + + // Antigravity transcripts carry no token usage, so the condensation + // recompute yields nil — exactly the case the fallback handles. + transcript := `{"type":"human","message":{"content":"add tokens"}} +{"type":"assistant","message":{"content":"Done."}} +` + require.NoError(t, os.WriteFile(filepath.Join(metadataDirAbs, paths.TranscriptFileName), []byte(transcript), 0o644)) + + trackedFile := filepath.Join(dir, "test.txt") + require.NoError(t, os.WriteFile(trackedFile, []byte("agent-modified content"), 0o644)) + + require.NoError(t, s.SaveStep(context.Background(), StepContext{ + SessionID: sessionID, + AgentType: agent.AgentTypeAntigravity, + ModifiedFiles: []string{"test.txt"}, + MetadataDir: metadataDir, + MetadataDirAbs: metadataDirAbs, + CommitMessage: "Antigravity checkpoint 1", + AuthorName: "Test", + AuthorEmail: "test@test.com", + })) + + state, err := s.loadSessionState(context.Background(), sessionID) + require.NoError(t, err) + require.Equal(t, agent.AgentTypeAntigravity, state.AgentType, "session must be tagged as Antigravity") + + // Simulate the lifecycle accumulating an out-of-band token delta into + // SessionState.TokenUsage (what SaveStep does with StepContext.TokenUsage + // for OutOfBandTokenSource agents). + state.TokenUsage = &agent.TokenUsage{ + InputTokens: 3500, + OutputTokens: 300, + CacheCreationTokens: 50, + CacheReadTokens: 3400, + APICallCount: 2, + } + require.NoError(t, SaveSessionState(context.Background(), state)) + + checkpointID := id.MustCheckpointID("ddee00112233") + result, err := s.CondenseSession(context.Background(), repo, checkpointID, state, nil) + require.NoError(t, err) + require.False(t, result.Skipped, "condensation must not skip when files are touched") + + ref, err := repo.Reference(plumbing.NewBranchReferenceName(paths.MetadataBranchName), true) + require.NoError(t, err) + commit, err := repo.CommitObject(ref.Hash()) + require.NoError(t, err) + tree, err := commit.Tree() + require.NoError(t, err) + + checkpointTree, err := tree.Tree(checkpointID.Path()) + require.NoError(t, err) + + // Per-session metadata must round-trip the out-of-band token usage. + sessionMeta, err := checkpointTree.File(checkpointID.Path() + "/0/" + paths.MetadataFileName) + if err != nil { + subtree, subErr := checkpointTree.Tree("0") + require.NoError(t, subErr) + sessionMeta, err = subtree.File(paths.MetadataFileName) + require.NoError(t, err) + } + sessionBytes, err := sessionMeta.Contents() + require.NoError(t, err) + var meta checkpoint.CommittedMetadata + require.NoError(t, json.Unmarshal([]byte(sessionBytes), &meta)) + + require.NotNil(t, meta.TokenUsage, "per-session token_usage must be populated from the out-of-band fallback") + require.Equal(t, 3500, meta.TokenUsage.InputTokens, "InputTokens") + require.Equal(t, 300, meta.TokenUsage.OutputTokens, "OutputTokens") + require.Equal(t, 50, meta.TokenUsage.CacheCreationTokens, "CacheCreationTokens") + require.Equal(t, 3400, meta.TokenUsage.CacheReadTokens, "CacheReadTokens") + require.Equal(t, 2, meta.TokenUsage.APICallCount, "APICallCount") +} + +// TestCondenseSession_NonOOBAgentDoesNotInheritStateTokens is the negative +// branch of the out-of-band fallback: a non-OutOfBandTokenSource agent (Cursor) +// must NOT inherit SessionState.TokenUsage when the transcript recompute yields +// nil. The fallback gate (agent.AsOutOfBandTokenSource) is what excludes such +// agents — without it, the populated state value would leak into per-session +// CommittedMetadata.token_usage. This test must FAIL if the gate is reverted to +// the old "not a TokenCalculator" form (Cursor is not a TokenCalculator, so the +// old gate would copy the value). +// +// Tests in this file use t.Chdir for CWD-based git resolution, so this +// cannot be a parallel test. +func TestCondenseSession_NonOOBAgentDoesNotInheritStateTokens(t *testing.T) { + dir := setupGitRepo(t) + t.Chdir(dir) + + repo, err := git.PlainOpen(dir) + require.NoError(t, err) + + s := &ManualCommitStrategy{} + sessionID := "2026-06-03-cursor-tokens" + + metadataDir := ".entire/metadata/" + sessionID + metadataDirAbs := filepath.Join(dir, metadataDir) + require.NoError(t, os.MkdirAll(metadataDirAbs, 0o755)) + + // Token-less transcript: the condensation recompute yields nil TokenUsage, + // mirroring the positive test so the only behavioral difference is the + // agent's OutOfBandTokenSource capability. + transcript := `{"type":"human","message":{"content":"add tokens"}} +{"type":"assistant","message":{"content":"Done."}} +` + require.NoError(t, os.WriteFile(filepath.Join(metadataDirAbs, paths.TranscriptFileName), []byte(transcript), 0o644)) + + trackedFile := filepath.Join(dir, "test.txt") + require.NoError(t, os.WriteFile(trackedFile, []byte("agent-modified content"), 0o644)) + + require.NoError(t, s.SaveStep(context.Background(), StepContext{ + SessionID: sessionID, + AgentType: agent.AgentTypeCursor, + ModifiedFiles: []string{"test.txt"}, + MetadataDir: metadataDir, + MetadataDirAbs: metadataDirAbs, + CommitMessage: "Cursor checkpoint 1", + AuthorName: "Test", + AuthorEmail: "test@test.com", + })) + + state, err := s.loadSessionState(context.Background(), sessionID) + require.NoError(t, err) + require.Equal(t, agent.AgentTypeCursor, state.AgentType, "session must be tagged as Cursor") + + // Populate SessionState.TokenUsage anyway. A non-OOB agent must NOT inherit + // this — the fallback gate excludes it. + state.TokenUsage = &agent.TokenUsage{ + InputTokens: 3500, + OutputTokens: 300, + CacheCreationTokens: 50, + CacheReadTokens: 3400, + APICallCount: 2, + } + require.NoError(t, SaveSessionState(context.Background(), state)) + + checkpointID := id.MustCheckpointID("ddee44556677") + result, err := s.CondenseSession(context.Background(), repo, checkpointID, state, nil) + require.NoError(t, err) + require.False(t, result.Skipped, "condensation must not skip when files are touched") + + ref, err := repo.Reference(plumbing.NewBranchReferenceName(paths.MetadataBranchName), true) + require.NoError(t, err) + commit, err := repo.CommitObject(ref.Hash()) + require.NoError(t, err) + tree, err := commit.Tree() + require.NoError(t, err) + + checkpointTree, err := tree.Tree(checkpointID.Path()) + require.NoError(t, err) + + sessionMeta, err := checkpointTree.File(checkpointID.Path() + "/0/" + paths.MetadataFileName) + if err != nil { + subtree, subErr := checkpointTree.Tree("0") + require.NoError(t, subErr) + sessionMeta, err = subtree.File(paths.MetadataFileName) + require.NoError(t, err) + } + sessionBytes, err := sessionMeta.Contents() + require.NoError(t, err) + var meta checkpoint.CommittedMetadata + require.NoError(t, json.Unmarshal([]byte(sessionBytes), &meta)) + + require.Nil(t, meta.TokenUsage, "non-OOB agent must NOT inherit SessionState.TokenUsage; per-session token_usage must be nil") +}