Merge pull request #67 from bnema/fix/context-usage-cache-only

fix(status): refresh context usage for cached turns

4 files · +60 -4

M packages/agent/modes/bot/runner.go +2 -2


@@ -161,8 +161,8 @@ func (r *Runner) runTurn(ctx context.Context, t queuedTurn) {
161	161		replyBuilder.WriteString(e.Delta)
162	162		case core.EvUsage:
163	163		r.mu.Lock()
164		-	if e.Usage.InputTokens > 0 {
165		-	r.lastCtxInput = e.Usage.InputTokens + e.Usage.CacheReadTokens + e.Usage.CacheWriteTokens
	164	+	if contextUsed := e.Usage.InputTokens + e.Usage.CacheReadTokens + e.Usage.CacheWriteTokens; contextUsed > 0 {
	165	+	r.lastCtxInput = contextUsed
166	166		}
167	167		r.mu.Unlock()
168	168		case core.EvAssistantMessage:


@@ -161,8 +161,8 @@ func (r *Runner) runTurn(ctx context.Context, t queuedTurn) {
161		replyBuilder.WriteString(e.Delta)	161		replyBuilder.WriteString(e.Delta)
162		case core.EvUsage:	162		case core.EvUsage:
163		r.mu.Lock()	163		r.mu.Lock()
164	-	if e.Usage.InputTokens > 0 {	164	+	if contextUsed := e.Usage.InputTokens + e.Usage.CacheReadTokens + e.Usage.CacheWriteTokens; contextUsed > 0 {
165	-	r.lastCtxInput = e.Usage.InputTokens + e.Usage.CacheReadTokens + e.Usage.CacheWriteTokens	165	+	r.lastCtxInput = contextUsed
166		}	166		}
167		r.mu.Unlock()	167		r.mu.Unlock()
168		case core.EvAssistantMessage:	168		case core.EvAssistantMessage:

M packages/agent/modes/bot/runner_test.go +29 -0


@@ -103,3 +103,32 @@ func TestHandleMessageClaimsDrainSlotBeforeSpawningDrainer(t *testing.T) {
103	103
104	104		close(client.release)
105	105		}
	106	+
	107	+	type cacheOnlyUsageClient struct{}
	108	+
	109	+	func (cacheOnlyUsageClient) Name() string { return "cache-only-usage" }
	110	+
	111	+	func (cacheOnlyUsageClient) Stream(context.Context, provider.Request) (<-chan provider.Event, error) {
	112	+	events := make(chan provider.Event, 2)
	113	+	events <- provider.EventUsage{Usage: provider.Usage{
	114	+	CacheReadTokens: 400,
	115	+	CacheWriteTokens: 50,
	116	+	}}
	117	+	events <- provider.EventDone{Stop: provider.StopEnd}
	118	+	close(events)
	119	+	return events, nil
	120	+	}
	121	+
	122	+	func TestRunTurnUpdatesContextUsageFromCacheOnlyTurn(t *testing.T) {
	123	+	runner := NewRunner(testAdapter{}, core.NewAgent(cacheOnlyUsageClient{}, "test-model", "", nil), Config{})
	124	+	runner.lastCtxInput = 99
	125	+
	126	+	runner.runTurn(context.Background(), queuedTurn{channelID: "channel", prompt: "prompt"})
	127	+
	128	+	runner.mu.Lock()
	129	+	got := runner.lastCtxInput
	130	+	runner.mu.Unlock()
	131	+	if got != 450 {
	132	+	t.Fatalf("context usage = %d, want 450", got)
	133	+	}
	134	+	}


@@ -103,3 +103,32 @@ func TestHandleMessageClaimsDrainSlotBeforeSpawningDrainer(t *testing.T) {
103		103
104	close(client.release)	104		close(client.release)
105	}	105		}
		106	+
		107	+	type cacheOnlyUsageClient struct{}
		108	+
		109	+	func (cacheOnlyUsageClient) Name() string { return "cache-only-usage" }
		110	+
		111	+	func (cacheOnlyUsageClient) Stream(context.Context, provider.Request) (<-chan provider.Event, error) {
		112	+	events := make(chan provider.Event, 2)
		113	+	events <- provider.EventUsage{Usage: provider.Usage{
		114	+	CacheReadTokens: 400,
		115	+	CacheWriteTokens: 50,
		116	+	}}
		117	+	events <- provider.EventDone{Stop: provider.StopEnd}
		118	+	close(events)
		119	+	return events, nil
		120	+	}
		121	+
		122	+	func TestRunTurnUpdatesContextUsageFromCacheOnlyTurn(t *testing.T) {
		123	+	runner := NewRunner(testAdapter{}, core.NewAgent(cacheOnlyUsageClient{}, "test-model", "", nil), Config{})
		124	+	runner.lastCtxInput = 99
		125	+
		126	+	runner.runTurn(context.Background(), queuedTurn{channelID: "channel", prompt: "prompt"})
		127	+
		128	+	runner.mu.Lock()
		129	+	got := runner.lastCtxInput
		130	+	runner.mu.Unlock()
		131	+	if got != 450 {
		132	+	t.Fatalf("context usage = %d, want 450", got)
		133	+	}
		134	+	}

M packages/agent/modes/interactive.go +2 -2


@@ -8109,8 +8109,8 @@ func (i *Interactive) handleEvent(ev core.AgentEvent) {
8109	8109		}
8110	8110		case core.EvUsage:
8111	8111		i.cumUsage = e.Cumulative
8112		-	if e.Usage.InputTokens > 0 {
8113		-	i.lastCtxInput = e.Usage.InputTokens + e.Usage.CacheReadTokens + e.Usage.CacheWriteTokens
	8112	+	if contextUsed := e.Usage.InputTokens + e.Usage.CacheReadTokens + e.Usage.CacheWriteTokens; contextUsed > 0 {
	8113	+	i.lastCtxInput = contextUsed
8114	8114		}
8115	8115		case core.EvTurnEnd:
8116	8116		if e.Stop == provider.StopAborted {


@@ -8109,8 +8109,8 @@ func (i *Interactive) handleEvent(ev core.AgentEvent) {
8109		}	8109		}
8110		case core.EvUsage:	8110		case core.EvUsage:
8111		i.cumUsage = e.Cumulative	8111		i.cumUsage = e.Cumulative
8112	-	if e.Usage.InputTokens > 0 {	8112	+	if contextUsed := e.Usage.InputTokens + e.Usage.CacheReadTokens + e.Usage.CacheWriteTokens; contextUsed > 0 {
8113	-	i.lastCtxInput = e.Usage.InputTokens + e.Usage.CacheReadTokens + e.Usage.CacheWriteTokens	8113	+	i.lastCtxInput = contextUsed
8114		}	8114		}
8115		case core.EvTurnEnd:	8115		case core.EvTurnEnd:
8116		if e.Stop == provider.StopAborted {	8116		if e.Stop == provider.StopAborted {

A packages/agent/modes/interactive_usage_test.go +27 -0

@@ -0,0 +1,27 @@
	1	+	package modes
	2	+
	3	+	import (
	4	+	"testing"
	5	+
	6	+	"github.com/bnema/zut/packages/core"
	7	+	"github.com/bnema/zut/packages/provider"
	8	+	)
	9	+
	10	+	func TestHandleEventUpdatesContextUsageFromCacheOnlyTurn(t *testing.T) {
	11	+	interactive := NewInteractive(InteractiveConfig{})
	12	+	interactive.lastCtxInput = 99
	13	+
	14	+	interactive.handleEvent(core.EvUsage{Usage: provider.Usage{
	15	+	CacheReadTokens: 400,
	16	+	CacheWriteTokens: 50,
	17	+	}})
	18	+
	19	+	if got := interactive.lastCtxInput; got != 450 {
	20	+	t.Fatalf("context usage = %d, want 450", got)
	21	+	}
	22	+
	23	+	interactive.handleEvent(core.EvUsage{})
	24	+	if got := interactive.lastCtxInput; got != 450 {
	25	+	t.Fatalf("empty usage cleared context usage to %d, want 450", got)
	26	+	}
	27	+	}

@@ -0,0 +1,27 @@
	1	+	package modes
	2	+
	3	+	import (
	4	+	"testing"
	5	+
	6	+	"github.com/bnema/zut/packages/core"
	7	+	"github.com/bnema/zut/packages/provider"
	8	+	)
	9	+
	10	+	func TestHandleEventUpdatesContextUsageFromCacheOnlyTurn(t *testing.T) {
	11	+	interactive := NewInteractive(InteractiveConfig{})
	12	+	interactive.lastCtxInput = 99
	13	+
	14	+	interactive.handleEvent(core.EvUsage{Usage: provider.Usage{
	15	+	CacheReadTokens: 400,
	16	+	CacheWriteTokens: 50,
	17	+	}})
	18	+
	19	+	if got := interactive.lastCtxInput; got != 450 {
	20	+	t.Fatalf("context usage = %d, want 450", got)
	21	+	}
	22	+
	23	+	interactive.handleEvent(core.EvUsage{})
	24	+	if got := interactive.lastCtxInput; got != 450 {
	25	+	t.Fatalf("empty usage cleared context usage to %d, want 450", got)
	26	+	}
	27	+	}