Skip to content

Commit 8593fda

Browse files
authored
feat: show token/elapsed summary after one-shot text runs (#288)
Text-mode --print previously discarded the stream usage event. Capture the final usage and render a muted token/elapsed footer to stderr on completion, so interactive one-shot runs report cost while stdout stays pure for scripts. JSON/stream-json output is unchanged.
1 parent 7637bc8 commit 8593fda

2 files changed

Lines changed: 110 additions & 0 deletions

File tree

cmd/chat_print.go

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,6 +83,8 @@ func runPrint(text string) error {
8383

8484
var printed strings.Builder
8585
var countdownShown bool
86+
var lastUsage *engine.StreamUsage
87+
started := time.Now()
8688
for ev := range ch {
8789
switch ev.Type {
8890
case "content":
@@ -118,6 +120,9 @@ func runPrint(text string) error {
118120
_, _ = fmt.Fprintf(os.Stderr, "%s %s\n", auditTint("["+ev.ToolName+"]", infoSky), content)
119121
}
120122
case "usage":
123+
if ev.Usage != nil {
124+
lastUsage = ev.Usage
125+
}
121126
if outputFormat == "stream-json" && ev.Usage != nil {
122127
writePrintUsageEvent(sessionID, ev.Usage)
123128
}
@@ -132,6 +137,7 @@ func runPrint(text string) error {
132137
if !strings.HasSuffix(printed.String(), "\n") {
133138
fmt.Println()
134139
}
140+
printTextUsageFooter(lastUsage, started)
135141
case "json":
136142
writePrintResult(printed.String(), sessionID, sess, false, nil)
137143
case "stream-json":
@@ -148,6 +154,7 @@ func runPrint(text string) error {
148154
if !strings.HasSuffix(printed.String(), "\n") {
149155
fmt.Println()
150156
}
157+
printTextUsageFooter(lastUsage, started)
151158
case "json":
152159
writePrintResult(printed.String(), sessionID, sess, false, nil)
153160
case "stream-json":
@@ -180,6 +187,21 @@ func writePrintUsageEvent(sessionID string, usage *engine.StreamUsage) {
180187
fmt.Println(string(data))
181188
}
182189

190+
// printTextUsageFooter renders a muted token/elapsed summary to stderr after a
191+
// one-shot text-mode run. It writes to stderr so stdout stays pure for scripts,
192+
// and is skipped entirely when no usage event was received.
193+
func printTextUsageFooter(usage *engine.StreamUsage, started time.Time) {
194+
if usage == nil {
195+
return
196+
}
197+
parts := []string{fmt.Sprintf("%d in · %d out", usage.PromptTokens, usage.CompletionTokens)}
198+
if usage.CacheReadTokens > 0 || usage.CacheWriteTokens > 0 {
199+
parts = append(parts, fmt.Sprintf("cache %d read · %d write", usage.CacheReadTokens, usage.CacheWriteTokens))
200+
}
201+
parts = append(parts, time.Since(started).Round(time.Second).String())
202+
_, _ = fmt.Fprintf(os.Stderr, "%s\n", auditTint("tokens: "+strings.Join(parts, " · "), textMuted))
203+
}
204+
183205
func writePrintResult(result, sessionID string, sess *engine.Session, isError bool, errors []string) {
184206
event := map[string]interface{}{
185207
"type": "result",

cmd/chat_print_usage_test.go

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,88 @@
1+
package cmd
2+
3+
import (
4+
"os"
5+
"strings"
6+
"testing"
7+
"time"
8+
9+
"github.com/GrayCodeAI/graycode-cli/internal/engine"
10+
)
11+
12+
// captureStderr runs fn with os.Stderr redirected to a pipe and returns the
13+
// captured bytes.
14+
func captureStderr(t *testing.T, fn func()) string {
15+
t.Helper()
16+
old := os.Stderr
17+
r, w, err := os.Pipe()
18+
if err != nil {
19+
t.Fatalf("pipe: %v", err)
20+
}
21+
os.Stderr = w
22+
defer func() { os.Stderr = old }()
23+
24+
fn()
25+
_ = w.Close()
26+
27+
buf := make([]byte, 4096)
28+
n, _ := r.Read(buf)
29+
_ = r.Close()
30+
return string(buf[:n])
31+
}
32+
33+
func TestPrintTextUsageFooter(t *testing.T) {
34+
started := time.Now().Add(-2 * time.Second)
35+
36+
t.Run("renders token and elapsed summary", func(t *testing.T) {
37+
got := captureStderr(t, func() {
38+
printTextUsageFooter(&engine.StreamUsage{
39+
PromptTokens: 100,
40+
CompletionTokens: 50,
41+
}, started)
42+
})
43+
if !strings.Contains(got, "100 in · 50 out") {
44+
t.Errorf("footer missing token counts: %q", got)
45+
}
46+
if !strings.Contains(got, "tokens:") {
47+
t.Errorf("footer missing prefix: %q", got)
48+
}
49+
if !strings.Contains(got, "2s") {
50+
t.Errorf("footer missing elapsed: %q", got)
51+
}
52+
})
53+
54+
t.Run("includes cache when nonzero", func(t *testing.T) {
55+
got := captureStderr(t, func() {
56+
printTextUsageFooter(&engine.StreamUsage{
57+
PromptTokens: 10,
58+
CompletionTokens: 5,
59+
CacheReadTokens: 90,
60+
CacheWriteTokens: 7,
61+
}, started)
62+
})
63+
if !strings.Contains(got, "cache 90 read · 7 write") {
64+
t.Errorf("footer missing cache summary: %q", got)
65+
}
66+
})
67+
68+
t.Run("omits cache when zero", func(t *testing.T) {
69+
got := captureStderr(t, func() {
70+
printTextUsageFooter(&engine.StreamUsage{
71+
PromptTokens: 10,
72+
CompletionTokens: 5,
73+
}, started)
74+
})
75+
if strings.Contains(got, "cache") {
76+
t.Errorf("footer should omit zero cache: %q", got)
77+
}
78+
})
79+
80+
t.Run("skips output when usage is nil", func(t *testing.T) {
81+
got := captureStderr(t, func() {
82+
printTextUsageFooter(nil, started)
83+
})
84+
if got != "" {
85+
t.Errorf("expected no output for nil usage, got: %q", got)
86+
}
87+
})
88+
}

0 commit comments

Comments
 (0)