From d903980256c3ecb8259b9b34f90111e30a09f1a1 Mon Sep 17 00:00:00 2001
From: ollie <lkn@ollie>
Date: Sun, 26 Jul 2026 19:01:19 +0200
Subject: [PATCH] persist usage/cost data across session restarts

---
 pkg/agent/commands.go  |  5 +----
 pkg/agent/core_test.go |  4 ++--
 pkg/agent/session.go   | 50 ++++++++++++++++++++++++++++++++----------
 3 files changed, 42 insertions(+), 17 deletions(-)

diff --git a/pkg/agent/commands.go b/pkg/agent/commands.go
index f4a660b..1310ceb 100644
--- a/pkg/agent/commands.go
+++ b/pkg/agent/commands.go
@@ -411,10 +411,7 @@ func (s *agent) handleCommand(ctx context.Context, input string) bool {
 				s.emit(infoEvent(fmt.Sprintf("error: %v", err)))
 				return
 			}
-			s.session = RestoreSession(ps.Messages)
-			if ps.TaskState != nil {
-				s.session.TaskState = ps.TaskState
-			}
+			s.session = RestoreSession(&ps)
 			s.emit(infoEvent(fmt.Sprintf("resumed session %s (%d messages)", name, len(ps.Messages))))
 		},
 
diff --git a/pkg/agent/core_test.go b/pkg/agent/core_test.go
index 991a29a..e31d3e2 100644
--- a/pkg/agent/core_test.go
+++ b/pkg/agent/core_test.go
@@ -2177,7 +2177,7 @@ func TestRestoreSession_RoundTrip(t *testing.T) {
 		t.Errorf("ps.ID=%q ps.Agent=%q", ps.ID, ps.Agent)
 	}
 
-	restored := RestoreSession(ps.Messages)
+	restored := RestoreSession(&ps)
 	if restored.goal != "first user message" {
 		t.Errorf("goal = %q; want %q", restored.goal, "first user message")
 	}
@@ -2192,7 +2192,7 @@ func TestRestoreSession_GoalFromFirstUserMessage(t *testing.T) {
 		{Role: "user", Content: "the real goal"},
 		{Role: "user", Content: "second user msg"},
 	}
-	s := RestoreSession(msgs)
+	s := RestoreSession(&PersistedSession{Messages: msgs})
 	if s.goal != "the real goal" {
 		t.Errorf("goal = %q; want %q", s.goal, "the real goal")
 	}
diff --git a/pkg/agent/session.go b/pkg/agent/session.go
index 7c4adbf..68ab243 100644
--- a/pkg/agent/session.go
+++ b/pkg/agent/session.go
@@ -43,6 +43,15 @@ type PersistedSession struct {
 	Remote    string            `json:"remote,omitempty"`
 	Messages  []backend.Message `json:"messages"`
 	TaskState *TaskState        `json:"taskState,omitempty"`
+	// Usage and cost tracking, persisted across restarts.
+	TotalInputTokens         int     `json:"totalInputTokens,omitempty"`
+	TotalCachedInputTokens   int     `json:"totalCachedInputTokens,omitempty"`
+	TotalCacheCreationTokens int     `json:"totalCacheCreationTokens,omitempty"`
+	TotalOutputTokens        int     `json:"totalOutputTokens,omitempty"`
+	TotalRequests            int     `json:"totalRequests,omitempty"`
+	Estimated                bool    `json:"estimated,omitempty"`
+	LastTurnCostUSD          float64 `json:"lastTurnCostUSD,omitempty"`
+	SessionCostUSD           float64 `json:"sessionCostUSD,omitempty"`
 }
 
 // TaskState is a compact structured overlay that summarizes the agent's
@@ -92,14 +101,22 @@ func (s *Session) saveTo(path, id, agentName, cwd string) error {
 // saveToFull writes session state including backend/model info.
 func (s *Session) saveToFull(path, id, agentName, backendName, modelName, cwd, remote string) error {
 	ps := PersistedSession{
-		ID:        id,
-		Agent:     agentName,
-		Backend:   backendName,
-		Model:     modelName,
-		CWD:       cwd,
-		Remote:    remote,
-		Messages:  s.messages,
-		TaskState: s.TaskState,
+		ID:                       id,
+		Agent:                    agentName,
+		Backend:                  backendName,
+		Model:                    modelName,
+		CWD:                      cwd,
+		Remote:                   remote,
+		Messages:                 s.messages,
+		TaskState:                s.TaskState,
+		TotalInputTokens:         s.TotalInputTokens,
+		TotalCachedInputTokens:   s.TotalCachedInputTokens,
+		TotalCacheCreationTokens: s.TotalCacheCreationTokens,
+		TotalOutputTokens:        s.TotalOutputTokens,
+		TotalRequests:            s.TotalRequests,
+		Estimated:                s.Estimated,
+		LastTurnCostUSD:          s.LastTurnCostUSD,
+		SessionCostUSD:           s.SessionCostUSD,
 	}
 	data, err := json.Marshal(ps)
 	if err != nil {
@@ -122,9 +139,20 @@ func LoadPersistedSession(path string) (*PersistedSession, error) {
 }
 
 // RestoreSession reconstructs a Session from a persisted message list.
-func RestoreSession(messages []backend.Message) *Session {
-	s := &Session{messages: messages}
-	for _, m := range messages {
+func RestoreSession(ps *PersistedSession) *Session {
+	s := &Session{
+		messages:                 ps.Messages,
+		TaskState:                ps.TaskState,
+		TotalInputTokens:         ps.TotalInputTokens,
+		TotalCachedInputTokens:   ps.TotalCachedInputTokens,
+		TotalCacheCreationTokens: ps.TotalCacheCreationTokens,
+		TotalOutputTokens:        ps.TotalOutputTokens,
+		TotalRequests:            ps.TotalRequests,
+		Estimated:                ps.Estimated,
+		LastTurnCostUSD:          ps.LastTurnCostUSD,
+		SessionCostUSD:           ps.SessionCostUSD,
+	}
+	for _, m := range ps.Messages {
 		if m.Role == "user" {
 			s.goal = m.Content
 			break