From 76f5baae4ab33713649e46432135883a4f3229aa Mon Sep 17 00:00:00 2001 From: Omar Alani Date: Sat, 23 May 2026 17:43:48 -0500 Subject: [PATCH 1/3] feat(assistant): add context seam breakdown --- internal/assistant/context_build.go | 20 ++-- internal/assistant/context_build_test.go | 2 + internal/assistant/context_usage.go | 13 +++ internal/assistant/lifecycle.go | 3 +- internal/assistant/runtime.go | 3 + internal/assistant/runtime_events_test.go | 8 +- internal/assistant/runtime_test.go | 8 +- internal/assistant/usage.go | 12 ++- internal/assistant/usage_events.go | 1 + internal/assistant/usage_test.go | 50 ++++++++-- internal/model/usage.go | 17 +++- internal/terminal/app.go | 2 +- internal/terminal/autocomplete.go | 1 + internal/terminal/commands.go | 5 +- internal/terminal/context_commands.go | 51 ++++++++++ internal/terminal/render_parity_test.go | 1 + internal/terminal/render_test.go | 3 + internal/terminal/token_usage.go | 12 +++ internal/terminal/token_usage_export_test.go | 4 + internal/terminal/token_usage_test.go | 99 ++++++++++++++++++-- 20 files changed, 281 insertions(+), 34 deletions(-) create mode 100644 internal/assistant/context_usage.go create mode 100644 internal/terminal/context_commands.go diff --git a/internal/assistant/context_build.go b/internal/assistant/context_build.go index 02307fe..a65d8ca 100644 --- a/internal/assistant/context_build.go +++ b/internal/assistant/context_build.go @@ -151,16 +151,19 @@ func (runtime *Runtime) modelContextBase( } func initialContextBuildResult(base *modelContextBase, selectedModel *model.Model) *contextBuildResult { + breakdown := contextBreakdown(base.SystemTokens, base.SkillTokens, base.HistoryTokens, nil) + return &contextBuildResult{ Contributions: []contextContribution{}, Messages: base.Messages, - Breakdown: contextBreakdown(base.SystemTokens, base.SkillTokens, base.HistoryTokens, nil), + Breakdown: breakdown, SystemPrompt: base.SystemPrompt, Usage: estimateContextBuildUsage( base.SystemPrompt, base.Messages, nil, selectedModel, + breakdown, ), } } @@ -170,18 +173,19 @@ func recalculateContextBuildResult( base *modelContextBase, selectedModel *model.Model, ) { - result.Usage = estimateContextBuildUsage( - base.SystemPrompt, - base.Messages, - result.Contributions, - selectedModel, - ) result.Breakdown = contextBreakdown( base.SystemTokens, base.SkillTokens, base.HistoryTokens, result.Contributions, ) + result.Usage = estimateContextBuildUsage( + result.SystemPrompt, + base.Messages, + result.Contributions, + selectedModel, + result.Breakdown, + ) } func estimateContextBuildUsage( @@ -189,6 +193,7 @@ func estimateContextBuildUsage( messages []database.MessageEntity, contributions []contextContribution, selectedModel *model.Model, + breakdown map[string]int, ) model.TokenUsage { inputTokens := estimateInputTokens(systemPrompt, messages) for index := range contributions { @@ -196,6 +201,7 @@ func estimateContextBuildUsage( } return model.TokenUsage{ + Breakdown: cloneIntMapForUsage(breakdown), ContextWindow: selectedModel.ContextWindow, ContextTokens: inputTokens, InputTokens: inputTokens, diff --git a/internal/assistant/context_build_test.go b/internal/assistant/context_build_test.go index 8efe405..44af426 100644 --- a/internal/assistant/context_build_test.go +++ b/internal/assistant/context_build_test.go @@ -42,6 +42,8 @@ end) assert.Contains(t, client.request.SystemPrompt, "") assert.Contains(t, client.request.SystemPrompt, "project-note") assert.Contains(t, client.request.SystemPrompt, "Always mention extension context") + require.NotNil(t, client.request.Usage.Breakdown) + assert.Greater(t, client.request.Usage.Breakdown["extensions"], 0) } func TestRuntime_ContextBuildRejectsOversizedExtensionContributions(t *testing.T) { diff --git a/internal/assistant/context_usage.go b/internal/assistant/context_usage.go new file mode 100644 index 0000000..8950be7 --- /dev/null +++ b/internal/assistant/context_usage.go @@ -0,0 +1,13 @@ +package assistant + +func cloneIntMapForUsage(values map[string]int) map[string]int { + if len(values) == 0 { + return nil + } + cloned := make(map[string]int, len(values)) + for key, value := range values { + cloned[key] = value + } + + return cloned +} diff --git a/internal/assistant/lifecycle.go b/internal/assistant/lifecycle.go index b2358b4..1858e6f 100644 --- a/internal/assistant/lifecycle.go +++ b/internal/assistant/lifecycle.go @@ -236,7 +236,7 @@ func contextBuildLifecyclePayload( lifecycleCWDKey: cwd, jsonSessionIDKey: sessionID, "message_count": len(base.Messages), - "breakdown": cloneIntMap(result.Breakdown), + jsonBreakdownKey: cloneIntMap(result.Breakdown), "contributions": []any{}, "max_contribution_tokens": contextContributionMaxTokens, "system_tokens": base.SystemTokens, @@ -311,6 +311,7 @@ func entryLifecyclePayload(entry *database.EntryEntity) map[string]any { func tokenUsageLifecyclePayload(usage model.TokenUsage) map[string]any { return map[string]any{ + jsonBreakdownKey: cloneIntMap(usage.Breakdown), jsonContextTokensKey: usage.ContextTokens, jsonContextWindowKey: usage.ContextWindow, jsonInputTokensKey: usage.InputTokens, diff --git a/internal/assistant/runtime.go b/internal/assistant/runtime.go index 43ee5c8..94ea24d 100644 --- a/internal/assistant/runtime.go +++ b/internal/assistant/runtime.go @@ -787,6 +787,7 @@ func (runtime *Runtime) modelResponse( sessionID, contextResult.SystemPrompt, cwd, + contextResult.Usage, registry, onEvent, ) @@ -812,6 +813,7 @@ func (runtime *Runtime) modelCompletionRequest( sessionID string, systemPrompt string, cwd string, + usage model.TokenUsage, registry *tool.Registry, onEvent func(StreamEvent), ) *CompletionRequest { @@ -826,6 +828,7 @@ func (runtime *Runtime) modelCompletionRequest( CWD: cwd, Auth: auth, Messages: messages, + Usage: usage, Model: *selectedModel, } } diff --git a/internal/assistant/runtime_events_test.go b/internal/assistant/runtime_events_test.go index 173f696..050e54e 100644 --- a/internal/assistant/runtime_events_test.go +++ b/internal/assistant/runtime_events_test.go @@ -29,7 +29,13 @@ func TestRuntime_ProviderLifecyclePublishesReactiveEvents(t *testing.T) { Text: testCompletionText, Thinking: nil, ToolEvents: nil, - Usage: model.TokenUsage{InputTokens: 9, OutputTokens: 3, ContextTokens: 9, ContextWindow: 100}, + Usage: model.TokenUsage{ + Breakdown: nil, + ContextWindow: 100, + ContextTokens: 9, + InputTokens: 9, + OutputTokens: 3, + }, }, err: nil, }) diff --git a/internal/assistant/runtime_test.go b/internal/assistant/runtime_test.go index 0190d4c..2ba381d 100644 --- a/internal/assistant/runtime_test.go +++ b/internal/assistant/runtime_test.go @@ -573,7 +573,13 @@ func (testCompletionClient) Complete( Text: "test assistant response for " + request.Messages[len(request.Messages)-1].Content, Thinking: nil, ToolEvents: nil, - Usage: model.TokenUsage{InputTokens: 12, OutputTokens: 4, ContextTokens: 12, ContextWindow: 1000}, + Usage: model.TokenUsage{ + Breakdown: nil, + ContextWindow: 1000, + ContextTokens: 12, + InputTokens: 12, + OutputTokens: 4, + }, }, nil } diff --git a/internal/assistant/usage.go b/internal/assistant/usage.go index 28e6a16..5979f69 100644 --- a/internal/assistant/usage.go +++ b/internal/assistant/usage.go @@ -46,6 +46,10 @@ func mergeUsage(estimated, reported model.TokenUsage) model.TokenUsage { if reported.OutputTokens > 0 { usage.OutputTokens = reported.OutputTokens } + if len(usage.Breakdown) == 0 && len(reported.Breakdown) > 0 { + usage.Breakdown = cloneIntMapForUsage(reported.Breakdown) + } + return usage } @@ -57,7 +61,13 @@ func usageFromObject(value any) model.TokenUsage { input := usageInputTokens(object) output := intFromAny(firstPresent(object, jsonOutputTokensKey, "completion_tokens")) - return model.TokenUsage{ContextWindow: 0, ContextTokens: 0, InputTokens: input, OutputTokens: output} + return model.TokenUsage{ + Breakdown: nil, + ContextWindow: 0, + ContextTokens: 0, + InputTokens: input, + OutputTokens: output, + } } func usageInputTokens(object map[string]any) int { diff --git a/internal/assistant/usage_events.go b/internal/assistant/usage_events.go index c2947cc..a8707e7 100644 --- a/internal/assistant/usage_events.go +++ b/internal/assistant/usage_events.go @@ -17,6 +17,7 @@ func (runtime *Runtime) emitUsage(ctx context.Context, onEvent func(StreamEvent) Text: "", }) payload := map[string]any{ + jsonBreakdownKey: cloneIntMap(usage.Breakdown), jsonContextWindowKey: usage.ContextWindow, jsonContextTokensKey: usage.ContextTokens, jsonInputTokensKey: usage.InputTokens, diff --git a/internal/assistant/usage_test.go b/internal/assistant/usage_test.go index ccfa3d1..ba4bf2a 100644 --- a/internal/assistant/usage_test.go +++ b/internal/assistant/usage_test.go @@ -22,7 +22,10 @@ func TestUsageFromObjectParsesProviderShapes(t *testing.T) { "input_tokens": float64(123), jsonOutputTokensKey: float64(45), }, - expected: model.TokenUsage{ContextWindow: 0, ContextTokens: 0, InputTokens: 123, OutputTokens: 45}, + expected: model.TokenUsage{ + Breakdown: nil, ContextWindow: 0, ContextTokens: 0, + InputTokens: 123, OutputTokens: 45, + }, }, { name: "chat completions", @@ -30,7 +33,10 @@ func TestUsageFromObjectParsesProviderShapes(t *testing.T) { "prompt_tokens": json.Number("77"), "completion_tokens": json.Number("9"), }, - expected: model.TokenUsage{ContextWindow: 0, ContextTokens: 0, InputTokens: 77, OutputTokens: 9}, + expected: model.TokenUsage{ + Breakdown: nil, ContextWindow: 0, ContextTokens: 0, + InputTokens: 77, OutputTokens: 9, + }, }, { name: "total tokens does not become input tokens", @@ -38,7 +44,10 @@ func TestUsageFromObjectParsesProviderShapes(t *testing.T) { "total_tokens": json.Number("120"), jsonOutputTokensKey: json.Number("20"), }, - expected: model.TokenUsage{ContextWindow: 0, ContextTokens: 0, InputTokens: 100, OutputTokens: 20}, + expected: model.TokenUsage{ + Breakdown: nil, ContextWindow: 0, ContextTokens: 0, + InputTokens: 100, OutputTokens: 20, + }, }, } @@ -54,10 +63,17 @@ func TestUsageFromObjectParsesProviderShapes(t *testing.T) { func TestMergeUsagePreservesEstimatedContextWindow(t *testing.T) { t.Parallel() - estimated := model.TokenUsage{ContextWindow: 1000, ContextTokens: 200, InputTokens: 200, OutputTokens: 0} - reported := model.TokenUsage{ContextWindow: 0, ContextTokens: 0, InputTokens: 150, OutputTokens: 25} + estimated := model.TokenUsage{ + Breakdown: nil, ContextWindow: 1000, ContextTokens: 200, + InputTokens: 200, OutputTokens: 0, + } + reported := model.TokenUsage{ + Breakdown: nil, ContextWindow: 0, ContextTokens: 0, + InputTokens: 150, OutputTokens: 25, + } assert.Equal(t, model.TokenUsage{ + Breakdown: nil, ContextWindow: 1000, ContextTokens: 200, InputTokens: 150, @@ -68,10 +84,17 @@ func TestMergeUsagePreservesEstimatedContextWindow(t *testing.T) { func TestMergeUsageNeverShrinksEstimatedContext(t *testing.T) { t.Parallel() - estimated := model.TokenUsage{ContextWindow: 100_000, ContextTokens: 14_000, InputTokens: 14_000, OutputTokens: 0} - reported := model.TokenUsage{ContextWindow: 0, ContextTokens: 12_000, InputTokens: 12_000, OutputTokens: 700} + estimated := model.TokenUsage{ + Breakdown: nil, ContextWindow: 100_000, ContextTokens: 14_000, + InputTokens: 14_000, OutputTokens: 0, + } + reported := model.TokenUsage{ + Breakdown: nil, ContextWindow: 0, ContextTokens: 12_000, + InputTokens: 12_000, OutputTokens: 700, + } assert.Equal(t, model.TokenUsage{ + Breakdown: nil, ContextWindow: 100_000, ContextTokens: 14_000, InputTokens: 12_000, @@ -82,10 +105,17 @@ func TestMergeUsageNeverShrinksEstimatedContext(t *testing.T) { func TestMergeUsageDoesNotPromoteProviderTotalToContext(t *testing.T) { t.Parallel() - estimated := model.TokenUsage{ContextWindow: 272_000, ContextTokens: 0, InputTokens: 0, OutputTokens: 0} - reported := model.TokenUsage{ContextWindow: 0, ContextTokens: 0, InputTokens: 13_000_000, OutputTokens: 100} + estimated := model.TokenUsage{ + Breakdown: nil, ContextWindow: 272_000, ContextTokens: 0, + InputTokens: 0, OutputTokens: 0, + } + reported := model.TokenUsage{ + Breakdown: nil, ContextWindow: 0, ContextTokens: 0, + InputTokens: 13_000_000, OutputTokens: 100, + } assert.Equal(t, model.TokenUsage{ + Breakdown: nil, ContextWindow: 272_000, ContextTokens: 0, InputTokens: 13_000_000, @@ -106,6 +136,7 @@ func TestParseSSEResultPreservesUsageWhenItemsProvideText(t *testing.T) { result, err := parseSSEResult(strings.NewReader(stream), nil) require.NoError(t, err) assert.Equal(t, model.TokenUsage{ + Breakdown: nil, ContextWindow: 0, ContextTokens: 0, InputTokens: 12, @@ -128,6 +159,7 @@ func TestParseSSEResultPreservesUsageAcrossLaterResponseEvents(t *testing.T) { result, err := parseSSEResult(strings.NewReader(stream), nil) require.NoError(t, err) assert.Equal(t, model.TokenUsage{ + Breakdown: nil, ContextWindow: 0, ContextTokens: 0, InputTokens: 12, diff --git a/internal/model/usage.go b/internal/model/usage.go index 81f3414..8e0a038 100644 --- a/internal/model/usage.go +++ b/internal/model/usage.go @@ -2,15 +2,22 @@ package model // TokenUsage tracks model context and request/response token counts. type TokenUsage struct { - ContextWindow int `json:"context_window,omitempty"` - ContextTokens int `json:"context_tokens,omitempty"` - InputTokens int `json:"input_tokens,omitempty"` - OutputTokens int `json:"output_tokens,omitempty"` + Breakdown map[string]int `json:"breakdown,omitempty"` + ContextWindow int `json:"context_window,omitempty"` + ContextTokens int `json:"context_tokens,omitempty"` + InputTokens int `json:"input_tokens,omitempty"` + OutputTokens int `json:"output_tokens,omitempty"` } // EmptyTokenUsage returns a zero-value token usage with explicit fields. func EmptyTokenUsage() TokenUsage { - return TokenUsage{ContextWindow: 0, ContextTokens: 0, InputTokens: 0, OutputTokens: 0} + return TokenUsage{ + Breakdown: nil, + ContextWindow: 0, + ContextTokens: 0, + InputTokens: 0, + OutputTokens: 0, + } } // TotalTokens returns input plus output tokens reported for the turn. diff --git a/internal/terminal/app.go b/internal/terminal/app.go index 27059da..2137613 100644 --- a/internal/terminal/app.go +++ b/internal/terminal/app.go @@ -139,8 +139,8 @@ type App struct { composerBuffer extension.BufferState messageLineCacheState messageLineCacheState streamingBlockLineCacheState messageLineCacheState - selection mouseSelection tokenUsage model.TokenUsage + selection mouseSelection promptSequence uint64 workFrame int lastMessageMaxRows int diff --git a/internal/terminal/autocomplete.go b/internal/terminal/autocomplete.go index 170acde..b0d3396 100644 --- a/internal/terminal/autocomplete.go +++ b/internal/terminal/autocomplete.go @@ -21,6 +21,7 @@ func slashSuggestions() []slashSuggestion { {Name: "changelog", Description: "open changelog"}, {Name: "clone", Description: "clone current session"}, {Name: "compact", Description: "compact conversation context"}, + {Name: "context", Description: "show context token breakdown"}, {Name: "copy", Description: "copy the last assistant message"}, {Name: "export", Description: "export current session"}, {Name: "fork", Description: "fork current session"}, diff --git a/internal/terminal/commands.go b/internal/terminal/commands.go index cfa717f..9a90842 100644 --- a/internal/terminal/commands.go +++ b/internal/terminal/commands.go @@ -48,7 +48,7 @@ func (app *App) openCommandPanel(ctx context.Context, command string) bool { } func (app *App) runSessionCommand(ctx context.Context, command, args, original string) (bool, error) { - if handler, ok := app.sessionCommandHandlers(ctx, args)[command]; ok { + if handler, ok := app.sessionCommandHandlers(ctx, args, original)[command]; ok { return false, handler() } if handler, ok := app.sessionCommandNotifications(ctx, command); ok { @@ -61,9 +61,10 @@ func (app *App) runSessionCommand(ctx context.Context, command, args, original s return false, nil } -func (app *App) sessionCommandHandlers(ctx context.Context, args string) map[string]func() error { +func (app *App) sessionCommandHandlers(ctx context.Context, args, original string) map[string]func() error { return map[string]func() error{ "clone": func() error { return app.cloneSession(ctx, args) }, + "context": func() error { return app.showContextInfo(ctx, original) }, "copy": func() error { return app.copyLastAssistantMessage(ctx) }, "fork": func() error { return app.newSession(ctx, args) }, commandLogin: func() error { return app.loginCommand(ctx, args) }, diff --git a/internal/terminal/context_commands.go b/internal/terminal/context_commands.go new file mode 100644 index 0000000..72e7811 --- /dev/null +++ b/internal/terminal/context_commands.go @@ -0,0 +1,51 @@ +package terminal + +import ( + "context" + "fmt" + "sort" + "strings" + + "github.com/omarluq/librecode/internal/database" +) + +func (app *App) showContextInfo(ctx context.Context, original string) error { + if !app.tokenUsage.HasAny() { + app.sendPrompt(ctx, original) + return nil + } + + lines := []string{"context:"} + if summary := formatContextUsage(app.tokenUsage); summary != "" { + lines = append(lines, "- "+summary) + } + breakdownLines := contextBreakdownLines(app.tokenUsage.Breakdown) + if len(breakdownLines) > 0 { + lines = append(lines, "- breakdown:") + lines = append(lines, breakdownLines...) + } + app.addMessage(database.RoleCustom, strings.Join(lines, "\n")) + + return nil +} + +func contextBreakdownLines(breakdown map[string]int) []string { + if len(breakdown) == 0 { + return nil + } + keys := make([]string, 0, len(breakdown)) + for key := range breakdown { + keys = append(keys, key) + } + sort.Strings(keys) + + lines := make([]string, 0, len(keys)) + for _, key := range keys { + if breakdown[key] <= 0 { + continue + } + lines = append(lines, fmt.Sprintf(" - %s: %s", key, compactCount(breakdown[key]))) + } + + return lines +} diff --git a/internal/terminal/render_parity_test.go b/internal/terminal/render_parity_test.go index 201a42e..c053d64 100644 --- a/internal/terminal/render_parity_test.go +++ b/internal/terminal/render_parity_test.go @@ -66,6 +66,7 @@ func TestRenderParityStatuslineTokenUsage(t *testing.T) { app := newRenderTestApp(t) app.tokenUsage = model.TokenUsage{ + Breakdown: nil, ContextWindow: 1000, ContextTokens: 250, InputTokens: 0, diff --git a/internal/terminal/render_test.go b/internal/terminal/render_test.go index 3679e73..bf34b00 100644 --- a/internal/terminal/render_test.go +++ b/internal/terminal/render_test.go @@ -364,6 +364,7 @@ func TestApplyPromptResponsePreservesLargerStreamedContextUsage(t *testing.T) { app := newRenderTestApp(t) app.applyTokenUsage(&model.TokenUsage{ + Breakdown: nil, ContextWindow: 100_000, ContextTokens: 14_000, InputTokens: 14_000, @@ -378,6 +379,7 @@ func TestApplyPromptResponsePreservesLargerStreamedContextUsage(t *testing.T) { Thinking: nil, ToolEvents: nil, Usage: model.TokenUsage{ + Breakdown: nil, ContextWindow: 100_000, ContextTokens: 12_000, InputTokens: 12_000, @@ -387,6 +389,7 @@ func TestApplyPromptResponsePreservesLargerStreamedContextUsage(t *testing.T) { }, 0) assert.Equal(t, model.TokenUsage{ + Breakdown: nil, ContextWindow: 100_000, ContextTokens: 14_000, InputTokens: 0, diff --git a/internal/terminal/token_usage.go b/internal/terminal/token_usage.go index 4ac3baf..560e5b2 100644 --- a/internal/terminal/token_usage.go +++ b/internal/terminal/token_usage.go @@ -20,6 +20,9 @@ func mergeTerminalUsage(current, next model.TokenUsage) model.TokenUsage { if next.ContextTokens > current.ContextTokens { current.ContextTokens = next.ContextTokens } + if len(next.Breakdown) > 0 { + current.Breakdown = cloneTokenBreakdown(next.Breakdown) + } return current } @@ -58,6 +61,15 @@ func formatContextUsage(usage model.TokenUsage) string { } } +func cloneTokenBreakdown(values map[string]int) map[string]int { + cloned := make(map[string]int, len(values)) + for key, value := range values { + cloned[key] = value + } + + return cloned +} + func compactCount(value int) string { if value >= 1_000_000 { return fmt.Sprintf("%.1fm", float64(value)/1_000_000) diff --git a/internal/terminal/token_usage_export_test.go b/internal/terminal/token_usage_export_test.go index 40e82f7..58e41f8 100644 --- a/internal/terminal/token_usage_export_test.go +++ b/internal/terminal/token_usage_export_test.go @@ -9,6 +9,10 @@ func MergeTerminalUsageForTest(current, next model.TokenUsage) model.TokenUsage return mergeTerminalUsage(current, next) } +func ContextBreakdownLinesForTest(breakdown map[string]int) []string { + return contextBreakdownLines(breakdown) +} + func NewAppForTest() *App { return newApp(nil, &RunOptions{ Extensions: nil, diff --git a/internal/terminal/token_usage_test.go b/internal/terminal/token_usage_test.go index eb0a25d..330760d 100644 --- a/internal/terminal/token_usage_test.go +++ b/internal/terminal/token_usage_test.go @@ -9,13 +9,32 @@ import ( "github.com/omarluq/librecode/internal/terminal" ) +const ( + testBreakdownExtensions = "extensions" + testBreakdownHistory = "history" + testBreakdownSystem = "system" +) + func TestMergeTerminalUsageIgnoresInputOutputTokens(t *testing.T) { t.Parallel() - current := model.TokenUsage{ContextWindow: 1_000_000, ContextTokens: 0, InputTokens: 0, OutputTokens: 0} - next := model.TokenUsage{ContextWindow: 0, ContextTokens: 0, InputTokens: 12_000, OutputTokens: 700} + current := model.TokenUsage{ + Breakdown: nil, + ContextWindow: 1_000_000, + ContextTokens: 0, + InputTokens: 0, + OutputTokens: 0, + } + next := model.TokenUsage{ + Breakdown: nil, + ContextWindow: 0, + ContextTokens: 0, + InputTokens: 12_000, + OutputTokens: 700, + } assert.Equal(t, model.TokenUsage{ + Breakdown: nil, ContextWindow: 1_000_000, ContextTokens: 0, InputTokens: 0, @@ -26,10 +45,23 @@ func TestMergeTerminalUsageIgnoresInputOutputTokens(t *testing.T) { func TestMergeTerminalUsagePreservesEstimatedContext(t *testing.T) { t.Parallel() - current := model.TokenUsage{ContextWindow: 1_000_000, ContextTokens: 17_000, InputTokens: 17_000, OutputTokens: 0} - next := model.TokenUsage{ContextWindow: 0, ContextTokens: 0, InputTokens: 12_000, OutputTokens: 700} + current := model.TokenUsage{ + Breakdown: nil, + ContextWindow: 1_000_000, + ContextTokens: 17_000, + InputTokens: 17_000, + OutputTokens: 0, + } + next := model.TokenUsage{ + Breakdown: nil, + ContextWindow: 0, + ContextTokens: 0, + InputTokens: 12_000, + OutputTokens: 700, + } assert.Equal(t, model.TokenUsage{ + Breakdown: nil, ContextWindow: 1_000_000, ContextTokens: 17_000, InputTokens: 17_000, @@ -41,7 +73,13 @@ func TestResetMessagesClearsTokenUsage(t *testing.T) { t.Parallel() app := terminal.NewAppForTest() - app.SetTokenUsageForTest(model.TokenUsage{ContextWindow: 1000, ContextTokens: 250, InputTokens: 0, OutputTokens: 0}) + app.SetTokenUsageForTest(model.TokenUsage{ + Breakdown: nil, + ContextWindow: 1000, + ContextTokens: 250, + InputTokens: 0, + OutputTokens: 0, + }) app.ResetMessagesForTest() @@ -52,10 +90,59 @@ func TestTruncateMessagesClearsTokenUsage(t *testing.T) { t.Parallel() app := terminal.NewAppForTest() - app.SetTokenUsageForTest(model.TokenUsage{ContextWindow: 1000, ContextTokens: 250, InputTokens: 0, OutputTokens: 0}) + app.SetTokenUsageForTest(model.TokenUsage{ + Breakdown: nil, + ContextWindow: 1000, + ContextTokens: 250, + InputTokens: 0, + OutputTokens: 0, + }) app.AddMessageForTest("user", "hello") app.TruncateMessagesForTest(0) assert.Equal(t, model.EmptyTokenUsage(), app.TokenUsageForTest()) } + +func TestMergeTerminalUsageKeepsBreakdown(t *testing.T) { + t.Parallel() + + current := model.TokenUsage{ + Breakdown: map[string]int{testBreakdownSystem: 10}, + ContextWindow: 1000, + ContextTokens: 20, + InputTokens: 20, + OutputTokens: 0, + } + next := model.TokenUsage{ + Breakdown: map[string]int{ + testBreakdownExtensions: 5, + testBreakdownHistory: 15, + testBreakdownSystem: 10, + }, + ContextWindow: 0, + ContextTokens: 30, + InputTokens: 30, + OutputTokens: 0, + } + + merged := terminal.MergeTerminalUsageForTest(current, next) + + assert.Equal(t, map[string]int{ + testBreakdownExtensions: 5, + testBreakdownHistory: 15, + testBreakdownSystem: 10, + }, merged.Breakdown) +} + +func TestContextBreakdownLinesSortsAndSkipsEmptyValues(t *testing.T) { + t.Parallel() + + lines := terminal.ContextBreakdownLinesForTest(map[string]int{ + testBreakdownExtensions: 0, + testBreakdownHistory: 1200, + testBreakdownSystem: 50, + }) + + assert.Equal(t, []string{" - history: 1.2k", " - system: 50"}, lines) +} From f82aefcd817dd284e55b7f7e33d7fe67bbd37dd8 Mon Sep 17 00:00:00 2001 From: Omar Alani Date: Sat, 23 May 2026 17:52:45 -0500 Subject: [PATCH 2/3] fix(assistant): include context usage fields --- go.sum | 4 ---- internal/assistant/anthropic_internal_test.go | 1 + internal/assistant/client.go | 2 ++ 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/go.sum b/go.sum index ae04f53..df7ac25 100644 --- a/go.sum +++ b/go.sum @@ -105,12 +105,8 @@ github.com/samber/oops v1.21.0 h1:18atcO4oEigNFuGXqr3NZWZ6P0XOSEXyBSAMXdQRxTc= github.com/samber/oops v1.21.0/go.mod h1:Hsm/sKPxtCfPh0w/cE3xVoRfSiE1joDRiStPAsmG9bo= github.com/samber/ro v0.3.0 h1:fxrdIL9yuA6JiGQuPE/xSXovowdG6+WcMBCCEL00wVk= github.com/samber/ro v0.3.0/go.mod h1:eInj5R1BbXfGoT1ef0HIO5Qie0wlPkkyL0koOaEmfNM= -github.com/samber/ro/plugins/fsnotify v0.0.0-20260516194255-a8c153943435 h1:xy3P6prHbHBCBe/pn5Qjd61k0dBnZmCN64ob9g5W27c= -github.com/samber/ro/plugins/fsnotify v0.0.0-20260516194255-a8c153943435/go.mod h1:39S5OD7epUKdNVSeg4tpy916q1SjxgIlMUN0twGI6sM= github.com/samber/ro/plugins/fsnotify v0.0.0-20260522200928-fb86b5fc464e h1:eOxaWpgXm/ygKqAz2Avvei898lBD5YOe9twefAFYEZE= github.com/samber/ro/plugins/fsnotify v0.0.0-20260522200928-fb86b5fc464e/go.mod h1:QtQE2XaB1KYIWnUotePRq36qZi7goPjtR/pm7gy2XWU= -github.com/samber/ro/plugins/signal v0.0.0-20260516194255-a8c153943435 h1:tySZEg7gguzOtqdOxOoe7EU4G34aPOb77d3MU/wJsG8= -github.com/samber/ro/plugins/signal v0.0.0-20260516194255-a8c153943435/go.mod h1:/zJel4/bMNX08spdRGRZvUWz1BE+iqXSZij3hDps3/A= github.com/samber/ro/plugins/signal v0.0.0-20260522200928-fb86b5fc464e h1:ehkGCX5UBJoGHHN3KMbAjt/bJLBe5asb0gY0KibWSmU= github.com/samber/ro/plugins/signal v0.0.0-20260522200928-fb86b5fc464e/go.mod h1:qZvrMezFzFNkX7V/Qh01B79tVW/Nw8TIGLjsrSdccy0= github.com/samber/slog-common v0.21.0 h1:Wo2hTly1Br5RjYqX/BTWJJeDnTE85oWk/7vqlpZuAUc= diff --git a/internal/assistant/anthropic_internal_test.go b/internal/assistant/anthropic_internal_test.go index 5bfb76a..cb4607b 100644 --- a/internal/assistant/anthropic_internal_test.go +++ b/internal/assistant/anthropic_internal_test.go @@ -201,6 +201,7 @@ func testCompletionRequestAuth(args ...string) *CompletionRequest { CWD: "", Auth: testRequestAuth(apiKey), Messages: nil, + Usage: model.EmptyTokenUsage(), Model: model.Model{ ThinkingLevelMap: nil, Headers: nil, diff --git a/internal/assistant/client.go b/internal/assistant/client.go index 04b0050..f094d91 100644 --- a/internal/assistant/client.go +++ b/internal/assistant/client.go @@ -50,6 +50,7 @@ const ( jsonAssistantRole = "assistant" jsonToolRole = "tool" jsonCommandKey = "command" + jsonBreakdownKey = "breakdown" jsonReadToolName = "read" jsonBashToolName = "bash" jsonEditToolName = "edit" @@ -92,6 +93,7 @@ type CompletionRequest struct { CWD string `json:"cwd"` Auth model.RequestAuth `json:"auth"` Messages []database.MessageEntity `json:"messages"` + Usage model.TokenUsage `json:"usage"` Model model.Model `json:"model"` } From 11a1208250bd3929cb857dd8bc508c0bf7b00d77 Mon Sep 17 00:00:00 2001 From: Omar Alani Date: Sat, 23 May 2026 18:05:29 -0500 Subject: [PATCH 3/3] fix(terminal): keep context command local --- internal/model/usage_test.go | 49 ++++++++++++++++++++ internal/terminal/context_commands.go | 7 ++- internal/terminal/token_usage_export_test.go | 15 ++++++ internal/terminal/token_usage_test.go | 40 ++++++++++++++++ 4 files changed, 110 insertions(+), 1 deletion(-) create mode 100644 internal/model/usage_test.go diff --git a/internal/model/usage_test.go b/internal/model/usage_test.go new file mode 100644 index 0000000..66e9370 --- /dev/null +++ b/internal/model/usage_test.go @@ -0,0 +1,49 @@ +package model_test + +import ( + "testing" + + "github.com/stretchr/testify/assert" + + "github.com/omarluq/librecode/internal/model" +) + +func TestTokenUsageHelpers(t *testing.T) { + t.Parallel() + + usage := model.TokenUsage{ + Breakdown: nil, + ContextWindow: 100, + ContextTokens: 25, + InputTokens: 10, + OutputTokens: 15, + } + + assert.Equal(t, 25, usage.TotalTokens()) + assert.True(t, usage.HasAny()) + assert.Equal(t, 25, usage.ContextPercent()) +} + +func TestTokenUsageContextPercentCapsAtHundred(t *testing.T) { + t.Parallel() + + usage := model.TokenUsage{ + Breakdown: nil, + ContextWindow: 100, + ContextTokens: 250, + InputTokens: 0, + OutputTokens: 0, + } + + assert.Equal(t, 100, usage.ContextPercent()) +} + +func TestEmptyTokenUsageHasNoUsage(t *testing.T) { + t.Parallel() + + usage := model.EmptyTokenUsage() + + assert.False(t, usage.HasAny()) + assert.Equal(t, 0, usage.TotalTokens()) + assert.Equal(t, 0, usage.ContextPercent()) +} diff --git a/internal/terminal/context_commands.go b/internal/terminal/context_commands.go index 72e7811..5ee7d0f 100644 --- a/internal/terminal/context_commands.go +++ b/internal/terminal/context_commands.go @@ -10,7 +10,7 @@ import ( ) func (app *App) showContextInfo(ctx context.Context, original string) error { - if !app.tokenUsage.HasAny() { + if !app.tokenUsage.HasAny() && !isContextCommand(original) { app.sendPrompt(ctx, original) return nil } @@ -29,6 +29,11 @@ func (app *App) showContextInfo(ctx context.Context, original string) error { return nil } +func isContextCommand(original string) bool { + trimmed := strings.TrimSpace(original) + return trimmed == "/context" || strings.HasPrefix(trimmed, "/context ") +} + func contextBreakdownLines(breakdown map[string]int) []string { if len(breakdown) == 0 { return nil diff --git a/internal/terminal/token_usage_export_test.go b/internal/terminal/token_usage_export_test.go index 58e41f8..81f4965 100644 --- a/internal/terminal/token_usage_export_test.go +++ b/internal/terminal/token_usage_export_test.go @@ -1,6 +1,8 @@ package terminal import ( + "context" + "github.com/omarluq/librecode/internal/database" "github.com/omarluq/librecode/internal/model" ) @@ -46,3 +48,16 @@ func (app *App) TruncateMessagesForTest(length int) { func (app *App) AddMessageForTest(role, content string) { app.addMessage(database.Role(role), content) } + +func (app *App) ShowContextInfoForTest(original string) error { + return app.showContextInfo(context.Background(), original) +} + +func (app *App) MessageContentsForTest() []string { + contents := make([]string, 0, len(app.messages)) + for _, message := range app.messages { + contents = append(contents, message.Content) + } + + return contents +} diff --git a/internal/terminal/token_usage_test.go b/internal/terminal/token_usage_test.go index 330760d..30ea1b1 100644 --- a/internal/terminal/token_usage_test.go +++ b/internal/terminal/token_usage_test.go @@ -1,9 +1,11 @@ package terminal_test import ( + "strings" "testing" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/omarluq/librecode/internal/model" "github.com/omarluq/librecode/internal/terminal" @@ -146,3 +148,41 @@ func TestContextBreakdownLinesSortsAndSkipsEmptyValues(t *testing.T) { assert.Equal(t, []string{" - history: 1.2k", " - system: 50"}, lines) } + +func TestShowContextInfoHandlesContextCommandWithoutUsage(t *testing.T) { + t.Parallel() + + app := terminal.NewAppForTest() + + require.NoError(t, app.ShowContextInfoForTest("/context")) + + messages := app.MessageContentsForTest() + require.NotEmpty(t, messages) + assert.Equal(t, "context:", messages[len(messages)-1]) +} + +func TestShowContextInfoDisplaysSummaryAndBreakdown(t *testing.T) { + t.Parallel() + + app := terminal.NewAppForTest() + app.SetTokenUsageForTest(model.TokenUsage{ + Breakdown: map[string]int{ + testBreakdownExtensions: 0, + testBreakdownHistory: 1200, + testBreakdownSystem: 50, + }, + ContextWindow: 1000, + ContextTokens: 250, + InputTokens: 0, + OutputTokens: 0, + }) + + require.NoError(t, app.ShowContextInfoForTest("/context now")) + + messages := app.MessageContentsForTest() + require.NotEmpty(t, messages) + message := messages[len(messages)-1] + assert.Contains(t, message, "context:") + assert.Contains(t, message, "- ctx 250/1.0k 25%") + assert.True(t, strings.Contains(message, "- breakdown:\n - history: 1.2k\n - system: 50")) +}