diff --git a/cmd/chat_print.go b/cmd/chat_print.go index 6fefebee..aef907a6 100644 --- a/cmd/chat_print.go +++ b/cmd/chat_print.go @@ -83,6 +83,8 @@ func runPrint(text string) error { var printed strings.Builder var countdownShown bool + var lastUsage *engine.StreamUsage + started := time.Now() for ev := range ch { switch ev.Type { case "content": @@ -118,6 +120,9 @@ func runPrint(text string) error { _, _ = fmt.Fprintf(os.Stderr, "%s %s\n", auditTint("["+ev.ToolName+"]", infoSky), content) } case "usage": + if ev.Usage != nil { + lastUsage = ev.Usage + } if outputFormat == "stream-json" && ev.Usage != nil { writePrintUsageEvent(sessionID, ev.Usage) } @@ -132,6 +137,7 @@ func runPrint(text string) error { if !strings.HasSuffix(printed.String(), "\n") { fmt.Println() } + printTextUsageFooter(lastUsage, started) case "json": writePrintResult(printed.String(), sessionID, sess, false, nil) case "stream-json": @@ -148,6 +154,7 @@ func runPrint(text string) error { if !strings.HasSuffix(printed.String(), "\n") { fmt.Println() } + printTextUsageFooter(lastUsage, started) case "json": writePrintResult(printed.String(), sessionID, sess, false, nil) case "stream-json": @@ -180,6 +187,21 @@ func writePrintUsageEvent(sessionID string, usage *engine.StreamUsage) { fmt.Println(string(data)) } +// printTextUsageFooter renders a muted token/elapsed summary to stderr after a +// one-shot text-mode run. It writes to stderr so stdout stays pure for scripts, +// and is skipped entirely when no usage event was received. +func printTextUsageFooter(usage *engine.StreamUsage, started time.Time) { + if usage == nil { + return + } + parts := []string{fmt.Sprintf("%d in · %d out", usage.PromptTokens, usage.CompletionTokens)} + if usage.CacheReadTokens > 0 || usage.CacheWriteTokens > 0 { + parts = append(parts, fmt.Sprintf("cache %d read · %d write", usage.CacheReadTokens, usage.CacheWriteTokens)) + } + parts = append(parts, time.Since(started).Round(time.Second).String()) + _, _ = fmt.Fprintf(os.Stderr, "%s\n", auditTint("tokens: "+strings.Join(parts, " · "), textMuted)) +} + func writePrintResult(result, sessionID string, sess *engine.Session, isError bool, errors []string) { event := map[string]interface{}{ "type": "result", diff --git a/cmd/chat_print_usage_test.go b/cmd/chat_print_usage_test.go new file mode 100644 index 00000000..28220278 --- /dev/null +++ b/cmd/chat_print_usage_test.go @@ -0,0 +1,88 @@ +package cmd + +import ( + "os" + "strings" + "testing" + "time" + + "github.com/GrayCodeAI/graycode-cli/internal/engine" +) + +// captureStderr runs fn with os.Stderr redirected to a pipe and returns the +// captured bytes. +func captureStderr(t *testing.T, fn func()) string { + t.Helper() + old := os.Stderr + r, w, err := os.Pipe() + if err != nil { + t.Fatalf("pipe: %v", err) + } + os.Stderr = w + defer func() { os.Stderr = old }() + + fn() + _ = w.Close() + + buf := make([]byte, 4096) + n, _ := r.Read(buf) + _ = r.Close() + return string(buf[:n]) +} + +func TestPrintTextUsageFooter(t *testing.T) { + started := time.Now().Add(-2 * time.Second) + + t.Run("renders token and elapsed summary", func(t *testing.T) { + got := captureStderr(t, func() { + printTextUsageFooter(&engine.StreamUsage{ + PromptTokens: 100, + CompletionTokens: 50, + }, started) + }) + if !strings.Contains(got, "100 in · 50 out") { + t.Errorf("footer missing token counts: %q", got) + } + if !strings.Contains(got, "tokens:") { + t.Errorf("footer missing prefix: %q", got) + } + if !strings.Contains(got, "2s") { + t.Errorf("footer missing elapsed: %q", got) + } + }) + + t.Run("includes cache when nonzero", func(t *testing.T) { + got := captureStderr(t, func() { + printTextUsageFooter(&engine.StreamUsage{ + PromptTokens: 10, + CompletionTokens: 5, + CacheReadTokens: 90, + CacheWriteTokens: 7, + }, started) + }) + if !strings.Contains(got, "cache 90 read · 7 write") { + t.Errorf("footer missing cache summary: %q", got) + } + }) + + t.Run("omits cache when zero", func(t *testing.T) { + got := captureStderr(t, func() { + printTextUsageFooter(&engine.StreamUsage{ + PromptTokens: 10, + CompletionTokens: 5, + }, started) + }) + if strings.Contains(got, "cache") { + t.Errorf("footer should omit zero cache: %q", got) + } + }) + + t.Run("skips output when usage is nil", func(t *testing.T) { + got := captureStderr(t, func() { + printTextUsageFooter(nil, started) + }) + if got != "" { + t.Errorf("expected no output for nil usage, got: %q", got) + } + }) +}