Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions cmd/entire/cli/agent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -68,7 +68,7 @@ type Agent interface {
ResolveSessionFile(sessionDir, agentSessionID string) string

// ReadSession reads session data from agent's storage.
ReadSession(input *HookInput) (*AgentSession, error)
ReadSession(ctx context.Context, input *HookInput) (*AgentSession, error)

// WriteSession writes session data for resumption.
WriteSession(ctx context.Context, session *AgentSession) error
Expand Down Expand Up @@ -184,7 +184,7 @@ type TokenCalculator interface {
Agent

// CalculateTokenUsage computes token usage from the transcript starting at the given offset.
CalculateTokenUsage(transcriptData []byte, fromOffset int) (*TokenUsage, error)
CalculateTokenUsage(ctx context.Context, transcriptData []byte, fromOffset int) (*TokenUsage, error)
}

// TextGenerator is an optional interface for agents whose CLI supports
Expand Down
4 changes: 3 additions & 1 deletion cmd/entire/cli/agent/agent_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,7 +42,9 @@ func (m *mockAgent) ResolveSessionFile(sessionDir, agentSessionID string) string
}

//nolint:nilnil // Mock implementation
func (m *mockAgent) ReadSession(_ *HookInput) (*AgentSession, error) { return nil, nil }
func (m *mockAgent) ReadSession(_ context.Context, _ *HookInput) (*AgentSession, error) {
return nil, nil
}
func (m *mockAgent) WriteSession(_ context.Context, _ *AgentSession) error { return nil }
func (m *mockAgent) FormatResumeCommand(_ string) string { return "" }

Expand Down
16 changes: 10 additions & 6 deletions cmd/entire/cli/agent/capabilities_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -23,11 +23,13 @@ func (m *mockBaseAgent) ReadTranscript(string) ([]byte, error) { return n
func (m *mockBaseAgent) ChunkTranscript(context.Context, []byte, int) ([][]byte, error) {
return nil, nil
}
func (m *mockBaseAgent) ReassembleTranscript([][]byte) ([]byte, error) { return nil, nil }
func (m *mockBaseAgent) GetSessionID(*HookInput) string { return "" }
func (m *mockBaseAgent) GetSessionDir(string) (string, error) { return "", nil }
func (m *mockBaseAgent) ResolveSessionFile(string, string) string { return "" }
func (m *mockBaseAgent) ReadSession(*HookInput) (*AgentSession, error) { return nil, nil } //nolint:nilnil // test mock
func (m *mockBaseAgent) ReassembleTranscript([][]byte) ([]byte, error) { return nil, nil }
func (m *mockBaseAgent) GetSessionID(*HookInput) string { return "" }
func (m *mockBaseAgent) GetSessionDir(string) (string, error) { return "", nil }
func (m *mockBaseAgent) ResolveSessionFile(string, string) string { return "" }
func (m *mockBaseAgent) ReadSession(context.Context, *HookInput) (*AgentSession, error) {
return nil, nil //nolint:nilnil // test mock
}
func (m *mockBaseAgent) WriteSession(context.Context, *AgentSession) error { return nil }
func (m *mockBaseAgent) FormatResumeCommand(string) string { return "" }

Expand Down Expand Up @@ -76,7 +78,9 @@ func (m *mockFullAgent) ExtractSummary(string) (string, error) { return "
func (m *mockFullAgent) PrepareTranscript(context.Context, string) error { return nil }

// TokenCalculator
func (m *mockFullAgent) CalculateTokenUsage([]byte, int) (*TokenUsage, error) { return nil, nil } //nolint:nilnil // test mock
func (m *mockFullAgent) CalculateTokenUsage(context.Context, []byte, int) (*TokenUsage, error) {
return nil, nil //nolint:nilnil // test mock
}

// TextGenerator
func (m *mockFullAgent) GenerateText(context.Context, string, string) (string, error) {
Expand Down
6 changes: 3 additions & 3 deletions cmd/entire/cli/agent/claudecode/claude.go
Original file line number Diff line number Diff line change
Expand Up @@ -121,7 +121,7 @@ func (c *ClaudeCodeAgent) GetSessionBaseDir() (string, error) {
// ReadSession reads a session from Claude's storage (JSONL transcript file).
// The session data is stored in NativeData as raw JSONL bytes.
// ModifiedFiles is computed by parsing the transcript.
func (c *ClaudeCodeAgent) ReadSession(input *agent.HookInput) (*agent.AgentSession, error) {
func (c *ClaudeCodeAgent) ReadSession(_ context.Context, input *agent.HookInput) (*agent.AgentSession, error) {
if input.SessionRef == "" {
return nil, errors.New("session reference (transcript path) is required")
}
Expand Down Expand Up @@ -251,8 +251,8 @@ func (c *ClaudeCodeAgent) FindCheckpointUUID(session *agent.AgentSession, toolUs

// ReadSessionFromPath is a convenience method that reads a session directly from a file path.
// This is useful when you have the path but not a HookInput.
func (c *ClaudeCodeAgent) ReadSessionFromPath(transcriptPath, sessionID string) (*agent.AgentSession, error) {
return c.ReadSession(&agent.HookInput{
func (c *ClaudeCodeAgent) ReadSessionFromPath(ctx context.Context, transcriptPath, sessionID string) (*agent.AgentSession, error) {
return c.ReadSession(ctx, &agent.HookInput{
SessionID: sessionID,
SessionRef: transcriptPath,
})
Expand Down
2 changes: 1 addition & 1 deletion cmd/entire/cli/agent/claudecode/lifecycle.go
Original file line number Diff line number Diff line change
Expand Up @@ -90,7 +90,7 @@ func (c *ClaudeCodeAgent) PrepareTranscript(ctx context.Context, sessionRef stri
}

// CalculateTokenUsage computes token usage from the transcript starting at the given line offset.
func (c *ClaudeCodeAgent) CalculateTokenUsage(transcriptData []byte, fromOffset int) (*agent.TokenUsage, error) {
func (c *ClaudeCodeAgent) CalculateTokenUsage(_ context.Context, transcriptData []byte, fromOffset int) (*agent.TokenUsage, error) {
return c.CalculateTotalTokenUsage(transcriptData, fromOffset, "")
}

Expand Down
2 changes: 1 addition & 1 deletion cmd/entire/cli/agent/codex/codex.go
Original file line number Diff line number Diff line change
Expand Up @@ -119,7 +119,7 @@ func (c *CodexAgent) ResolveRestoredSessionFile(sessionDir, agentSessionID strin
func (c *CodexAgent) ProtectedDirs() []string { return []string{".codex"} }

// ReadSession reads a session from Codex's storage (JSONL rollout file).
func (c *CodexAgent) ReadSession(input *agent.HookInput) (*agent.AgentSession, error) {
func (c *CodexAgent) ReadSession(_ context.Context, input *agent.HookInput) (*agent.AgentSession, error) {
if input.SessionRef == "" {
return nil, errors.New("session reference (transcript path) is required")
}
Expand Down
4 changes: 2 additions & 2 deletions cmd/entire/cli/agent/codex/codex_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ func TestCodexAgent_ReadSession(t *testing.T) {
ag := &CodexAgent{}
path := writeSampleRollout(t)

session, err := ag.ReadSession(&agent.HookInput{
session, err := ag.ReadSession(t.Context(), &agent.HookInput{
SessionID: "codex-session-1",
SessionRef: path,
})
Expand All @@ -145,7 +145,7 @@ func TestCodexAgent_ReadSession_InvalidSessionMeta(t *testing.T) {
path := filepath.Join(dir, "rollout.jsonl")
require.NoError(t, os.WriteFile(path, []byte(`{"timestamp":"2026-03-25T11:31:11.754Z","type":"response_item","payload":{"type":"message"}}`), 0o600))

_, err := ag.ReadSession(&agent.HookInput{
_, err := ag.ReadSession(t.Context(), &agent.HookInput{
SessionID: "codex-session-1",
SessionRef: path,
})
Expand Down
3 changes: 2 additions & 1 deletion cmd/entire/cli/agent/codex/transcript.go
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ package codex
import (
"bufio"
"bytes"
"context"
"encoding/json"
"errors"
"fmt"
Expand Down Expand Up @@ -197,7 +198,7 @@ func extractFilesFromApplyPatch(input string) []string {
// CalculateTokenUsage computes token usage from the transcript starting at the given line offset.
// Codex reports cumulative total_token_usage, so we compute the delta between the last
// token_count at/before the offset and the last token_count after the offset.
func (c *CodexAgent) CalculateTokenUsage(transcriptData []byte, fromOffset int) (*agent.TokenUsage, error) {
func (c *CodexAgent) CalculateTokenUsage(_ context.Context, transcriptData []byte, fromOffset int) (*agent.TokenUsage, error) {
var baselineUsage *tokenUsageData // last token_count at or before offset
var lastUsage *tokenUsageData // last token_count after offset
apiCalls := 0
Expand Down
6 changes: 3 additions & 3 deletions cmd/entire/cli/agent/codex/transcript_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -98,7 +98,7 @@ func TestCalculateTokenUsage(t *testing.T) {
ag := &CodexAgent{}

// From offset 0 (no baseline), should return full cumulative total
usage, err := ag.CalculateTokenUsage([]byte(sampleRollout), 0)
usage, err := ag.CalculateTokenUsage(t.Context(), []byte(sampleRollout), 0)
require.NoError(t, err)
require.NotNil(t, usage)

Expand All @@ -115,7 +115,7 @@ func TestCalculateTokenUsage_WithOffset(t *testing.T) {

// Skip past first token_count (line 4) — baseline is {5000, 4000, 100}
// Last after offset is {15000, 12000, 300} → delta = {10000, 8000, 200}
usage, err := ag.CalculateTokenUsage([]byte(sampleRollout), 4)
usage, err := ag.CalculateTokenUsage(t.Context(), []byte(sampleRollout), 4)
require.NoError(t, err)
require.NotNil(t, usage)

Expand All @@ -130,7 +130,7 @@ func TestCalculateTokenUsage_NoData(t *testing.T) {
t.Parallel()
ag := &CodexAgent{}

usage, err := ag.CalculateTokenUsage([]byte(`{"timestamp":"t","type":"session_meta","payload":{}}`), 0)
usage, err := ag.CalculateTokenUsage(t.Context(), []byte(`{"timestamp":"t","type":"session_meta","payload":{}}`), 0)
require.NoError(t, err)
require.Nil(t, usage)
}
Expand Down
2 changes: 1 addition & 1 deletion cmd/entire/cli/agent/copilotcli/copilotcli.go
Original file line number Diff line number Diff line change
Expand Up @@ -85,7 +85,7 @@ func (c *CopilotCLIAgent) ProtectedDirs() []string {
}

// ReadSession reads a session from Copilot CLI's storage (JSONL transcript file).
func (c *CopilotCLIAgent) ReadSession(input *agent.HookInput) (*agent.AgentSession, error) {
func (c *CopilotCLIAgent) ReadSession(_ context.Context, input *agent.HookInput) (*agent.AgentSession, error) {
if input.SessionRef == "" {
return nil, errors.New("session reference (transcript path) is required")
}
Expand Down
12 changes: 6 additions & 6 deletions cmd/entire/cli/agent/copilotcli/copilotcli_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -156,7 +156,7 @@ func TestReadSession_Success(t *testing.T) {
SessionRef: transcriptPath,
}

session, err := ag.ReadSession(input)
session, err := ag.ReadSession(t.Context(), input)
if err != nil {
t.Fatalf("ReadSession() error = %v", err)
}
Expand Down Expand Up @@ -189,7 +189,7 @@ func TestReadSession_NativeDataMatchesFile(t *testing.T) {
SessionRef: transcriptPath,
}

session, err := ag.ReadSession(input)
session, err := ag.ReadSession(t.Context(), input)
if err != nil {
t.Fatalf("ReadSession() error = %v", err)
}
Expand All @@ -209,7 +209,7 @@ func TestReadSession_EmptySessionRef(t *testing.T) {
ag := &CopilotCLIAgent{}
input := &agent.HookInput{SessionID: "sess-no-ref"}

_, err := ag.ReadSession(input)
_, err := ag.ReadSession(t.Context(), input)
if err == nil {
t.Fatal("ReadSession() should error when SessionRef is empty")
}
Expand All @@ -223,7 +223,7 @@ func TestReadSession_MissingFile(t *testing.T) {
SessionRef: "/nonexistent/path/events.jsonl",
}

_, err := ag.ReadSession(input)
_, err := ag.ReadSession(t.Context(), input)
if err == nil {
t.Fatal("ReadSession() should error when transcript file doesn't exist")
}
Expand Down Expand Up @@ -271,7 +271,7 @@ func TestWriteSession_RoundTrip(t *testing.T) {
SessionID: "roundtrip-session",
SessionRef: transcriptPath,
}
session, err := ag.ReadSession(input)
session, err := ag.ReadSession(t.Context(), input)
if err != nil {
t.Fatalf("ReadSession() error = %v", err)
}
Expand Down Expand Up @@ -562,7 +562,7 @@ func TestReadTranscript_MatchesReadSession(t *testing.T) {
t.Fatalf("ReadTranscript() error = %v", err)
}

session, err := ag.ReadSession(&agent.HookInput{
session, err := ag.ReadSession(t.Context(), &agent.HookInput{
SessionID: "compare-session",
SessionRef: transcriptPath,
})
Expand Down
11 changes: 6 additions & 5 deletions cmd/entire/cli/agent/copilotcli/transcript.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"encoding/json"
"fmt"
"io"
"log/slog"
"os"

"github.com/entireio/cli/cmd/entire/cli/agent"
Expand Down Expand Up @@ -267,19 +268,19 @@ const eventTypeSessionShutdown = "session.shutdown"
// contains session-wide aggregates. For sliced transcripts, session.shutdown
// would overcount because it is not checkpoint-scoped, so we fall back to
// summing assistant.message outputTokens within the slice.
func (c *CopilotCLIAgent) CalculateTokenUsage(transcriptData []byte, fromOffset int) (*agent.TokenUsage, error) {
func (c *CopilotCLIAgent) CalculateTokenUsage(ctx context.Context, transcriptData []byte, fromOffset int) (*agent.TokenUsage, error) {
events, err := parseEventsFromOffset(transcriptData, fromOffset)
if err != nil {
return nil, fmt.Errorf("failed to parse transcript for token usage: %w", err)
}

return extractTokenUsageFromEvents(events, fromOffset == 0), nil
return extractTokenUsageFromEvents(ctx, events, fromOffset == 0), nil
}

// extractTokenUsageFromEvents extracts token usage from parsed events.
// Prefers session.shutdown aggregate only when the caller is looking at the
// full transcript; otherwise falls back to per-message outputTokens.
func extractTokenUsageFromEvents(events []copilotEvent, preferSessionShutdown bool) *agent.TokenUsage {
func extractTokenUsageFromEvents(ctx context.Context, events []copilotEvent, preferSessionShutdown bool) *agent.TokenUsage {
if preferSessionShutdown {
// session.shutdown is authoritative, but only for full-session totals.
for i := len(events) - 1; i >= 0; i-- {
Expand All @@ -289,8 +290,8 @@ func extractTokenUsageFromEvents(events []copilotEvent, preferSessionShutdown bo

var data sessionShutdownData
if err := json.Unmarshal(events[i].Data, &data); err != nil {
logging.Debug(context.Background(), "copilot-cli: session.shutdown data unmarshal failed",
"err", err)
logging.Debug(ctx, "copilot-cli: session.shutdown data unmarshal failed",
slog.String("err", err.Error()))
continue
}

Expand Down
12 changes: 6 additions & 6 deletions cmd/entire/cli/agent/copilotcli/transcript_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -382,7 +382,7 @@ func TestCalculateTokenUsage_SessionShutdown(t *testing.T) {
)
content := []byte(strings.Join(lines, "\n") + "\n")

usage, err := ag.CalculateTokenUsage(content, 0)
usage, err := ag.CalculateTokenUsage(t.Context(), content, 0)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
Expand Down Expand Up @@ -412,7 +412,7 @@ func TestCalculateTokenUsage_MultiModel(t *testing.T) {
}
content := []byte(strings.Join(lines, "\n") + "\n")

usage, err := ag.CalculateTokenUsage(content, 0)
usage, err := ag.CalculateTokenUsage(t.Context(), content, 0)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
Expand Down Expand Up @@ -445,7 +445,7 @@ func TestCalculateTokenUsage_EmptyModelMetrics(t *testing.T) {
}
content := []byte(strings.Join(lines, "\n") + "\n")

usage, err := ag.CalculateTokenUsage(content, 0)
usage, err := ag.CalculateTokenUsage(t.Context(), content, 0)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
Expand All @@ -468,7 +468,7 @@ func TestCalculateTokenUsage_FallbackToAssistantMessages(t *testing.T) {
}
content := []byte(strings.Join(lines, "\n") + "\n")

usage, err := ag.CalculateTokenUsage(content, 0)
usage, err := ag.CalculateTokenUsage(t.Context(), content, 0)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
Expand All @@ -487,7 +487,7 @@ func TestCalculateTokenUsage_EmptyTranscript(t *testing.T) {
t.Parallel()
ag := &CopilotCLIAgent{}

usage, err := ag.CalculateTokenUsage([]byte{}, 0)
usage, err := ag.CalculateTokenUsage(t.Context(), []byte{}, 0)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
Expand All @@ -509,7 +509,7 @@ func TestCalculateTokenUsage_WithOffset(t *testing.T) {
}
content := []byte(strings.Join(lines, "\n") + "\n")

usage, err := ag.CalculateTokenUsage(content, 1)
usage, err := ag.CalculateTokenUsage(t.Context(), content, 1)
if err != nil {
t.Fatalf("unexpected error: %v", err)
}
Expand Down
2 changes: 1 addition & 1 deletion cmd/entire/cli/agent/cursor/cursor.go
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,7 @@ func (c *CursorAgent) GetSessionBaseDir() (string, error) {
// Note: ModifiedFiles is left empty because Cursor's transcript does not contain
// tool_use blocks for file detection. TranscriptAnalyzer extracts prompts and
// summaries; file detection relies on git status.
func (c *CursorAgent) ReadSession(input *agent.HookInput) (*agent.AgentSession, error) {
func (c *CursorAgent) ReadSession(_ context.Context, input *agent.HookInput) (*agent.AgentSession, error) {
if input.SessionRef == "" {
return nil, errors.New("session reference (transcript path) is required")
}
Expand Down
12 changes: 6 additions & 6 deletions cmd/entire/cli/agent/cursor/cursor_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -239,7 +239,7 @@ func TestReadSession_Success(t *testing.T) {
SessionRef: transcriptPath,
}

session, err := ag.ReadSession(input)
session, err := ag.ReadSession(t.Context(), input)
if err != nil {
t.Fatalf("ReadSession() error = %v", err)
}
Expand Down Expand Up @@ -272,7 +272,7 @@ func TestReadSession_NativeDataMatchesFile(t *testing.T) {
SessionRef: transcriptPath,
}

session, err := ag.ReadSession(input)
session, err := ag.ReadSession(t.Context(), input)
if err != nil {
t.Fatalf("ReadSession() error = %v", err)
}
Expand All @@ -298,7 +298,7 @@ func TestReadSession_ModifiedFilesEmpty(t *testing.T) {
SessionRef: transcriptPath,
}

session, err := ag.ReadSession(input)
session, err := ag.ReadSession(t.Context(), input)
if err != nil {
t.Fatalf("ReadSession() error = %v", err)
}
Expand All @@ -315,7 +315,7 @@ func TestReadSession_EmptySessionRef(t *testing.T) {
ag := &CursorAgent{}
input := &agent.HookInput{SessionID: "sess-no-ref"}

_, err := ag.ReadSession(input)
_, err := ag.ReadSession(t.Context(), input)
if err == nil {
t.Fatal("ReadSession() should error when SessionRef is empty")
}
Expand All @@ -329,7 +329,7 @@ func TestReadSession_MissingFile(t *testing.T) {
SessionRef: "/nonexistent/path/transcript.jsonl",
}

_, err := ag.ReadSession(input)
_, err := ag.ReadSession(t.Context(), input)
if err == nil {
t.Fatal("ReadSession() should error when transcript file doesn't exist")
}
Expand Down Expand Up @@ -377,7 +377,7 @@ func TestWriteSession_RoundTrip(t *testing.T) {
SessionID: "roundtrip-session",
SessionRef: transcriptPath,
}
session, err := ag.ReadSession(input)
session, err := ag.ReadSession(t.Context(), input)
if err != nil {
t.Fatalf("ReadSession() error = %v", err)
}
Expand Down
2 changes: 1 addition & 1 deletion cmd/entire/cli/agent/cursor/lifecycle_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -627,7 +627,7 @@ func TestReadTranscript_MatchesReadSession(t *testing.T) {
}

// ReadSession
session, err := ag.ReadSession(&agent.HookInput{
session, err := ag.ReadSession(t.Context(), &agent.HookInput{
SessionID: "compare-session",
SessionRef: transcriptPath,
})
Expand Down
Loading
Loading