Merge pull request #1396 from dgageot/costs

Pass message level costs in --remote mode and `cagent exec --json`

Author: dgageot

gen/cagent/v1/cagent.pb.go

@@ -377,6 +377,20 @@ func toolCallResultToProto(r *tools.ToolCallResult) *cagentv1.ToolCallResult {
 	}
 }
 
+func messageUsageToProto(m *runtime.MessageUsage) *cagentv1.LastMessageUsage {
+	if m == nil {
+		return nil
+	}
+	return &cagentv1.LastMessageUsage{
+		InputTokens:       m.InputTokens,
+		OutputTokens:      m.OutputTokens,
+		CachedInputTokens: m.CachedInputTokens,
+		CacheWriteTokens:  m.CacheWriteTokens,
+		Cost:              m.Cost,
+		Model:             m.Model,
+	}
+}
+
 func runtimeEventToProto(event runtime.Event) *cagentv1.Event {
 	switch e := event.(type) {
 	case *runtime.UserMessageEvent:
@@ -505,6 +519,7 @@ func runtimeEventToProto(event runtime.Event) *cagentv1.Event {
 						ContextLength: e.Usage.ContextLength,
 						ContextLimit:  e.Usage.ContextLimit,
 						Cost:          e.Usage.Cost,
+						LastMessage:   messageUsageToProto(e.Usage.LastMessage),
 					},
 					AgentName: e.AgentName,
 				},

gen/cagent/v1/cagentv1connect/cagent.connect.go

@@ -13,6 +13,7 @@ import (
 	cagentv1 "github.com/docker/cagent/gen/proto/cagent/v1"
 	"github.com/docker/cagent/gen/proto/cagent/v1/cagentv1connect"
 	"github.com/docker/cagent/pkg/api"
+	"github.com/docker/cagent/pkg/chat"
 	"github.com/docker/cagent/pkg/config/latest"
 	"github.com/docker/cagent/pkg/session"
 	"github.com/docker/cagent/pkg/tools"
@@ -345,6 +346,7 @@ func (c *ConnectRPCClient) convertProtoEventToRuntimeEvent(e *cagentv1.Event) Ev
 				ContextLength: ev.TokenUsage.Usage.ContextLength,
 				ContextLimit:  ev.TokenUsage.Usage.ContextLimit,
 				Cost:          ev.TokenUsage.Usage.Cost,
+				LastMessage:   convertProtoMessageUsage(ev.TokenUsage.Usage.LastMessage),
 			}
 		}
 		return &TokenUsageEvent{
@@ -526,6 +528,22 @@ func convertProtoToolCall(tc *cagentv1.ToolCall) tools.ToolCall {
 	}
 }
 
+func convertProtoMessageUsage(m *cagentv1.LastMessageUsage) *MessageUsage {
+	if m == nil {
+		return nil
+	}
+	return &MessageUsage{
+		Usage: chat.Usage{
+			InputTokens:       m.InputTokens,
+			OutputTokens:      m.OutputTokens,
+			CachedInputTokens: m.CachedInputTokens,
+			CacheWriteTokens:  m.CacheWriteTokens,
+		},
+		Cost:  m.Cost,
+		Model: m.Model,
+	}
+}
+
 func convertProtoTool(t *cagentv1.Tool) tools.Tool {
 	if t == nil {
 		return tools.Tool{}

gen/proto/cagent/v1/cagent.pb.go

@@ -3,6 +3,7 @@ package runtime
 import (
 	"cmp"
 
+	"github.com/docker/cagent/pkg/chat"
 	"github.com/docker/cagent/pkg/tools"
 )
 
@@ -194,14 +195,27 @@ type TokenUsageEvent struct {
 }
 
 type Usage struct {
-	InputTokens   int64   `json:"input_tokens"`
-	OutputTokens  int64   `json:"output_tokens"`
-	ContextLength int64   `json:"context_length"`
-	ContextLimit  int64   `json:"context_limit"`
-	Cost          float64 `json:"cost"`
+	InputTokens   int64         `json:"input_tokens"`
+	OutputTokens  int64         `json:"output_tokens"`
+	ContextLength int64         `json:"context_length"`
+	ContextLimit  int64         `json:"context_limit"`
+	Cost          float64       `json:"cost"`
+	LastMessage   *MessageUsage `json:"last_message,omitempty"`
+}
+
+// MessageUsage contains per-message usage data to include in TokenUsageEvent.
+// It embeds chat.Usage and adds Cost and Model fields.
+type MessageUsage struct {
+	chat.Usage
+	Cost  float64
+	Model string
 }
 
 func TokenUsage(sessionID, agentName string, inputTokens, outputTokens, contextLength, contextLimit int64, cost float64) Event {
+	return TokenUsageWithMessage(sessionID, agentName, inputTokens, outputTokens, contextLength, contextLimit, cost, nil)
+}
+
+func TokenUsageWithMessage(sessionID, agentName string, inputTokens, outputTokens, contextLength, contextLimit int64, cost float64, msgUsage *MessageUsage) Event {
 	return &TokenUsageEvent{
 		Type:      "token_usage",
 		SessionID: sessionID,
@@ -211,6 +225,7 @@ func TokenUsage(sessionID, agentName string, inputTokens, outputTokens, contextL
 			InputTokens:   inputTokens,
 			OutputTokens:  outputTokens,
 			Cost:          cost,
+			LastMessage:   msgUsage,
 		},
 		AgentContext: AgentContext{AgentName: agentName},
 	}

pkg/connectrpc/server.go

@@ -811,6 +811,7 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 
 			// Add assistant message to conversation history, but skip empty assistant messages
 			// Providers reject assistant messages that have neither content nor tool calls.
+			var msgUsage *MessageUsage
 			if strings.TrimSpace(res.Content) != "" || len(res.Calls) > 0 {
 				// Build tool definitions for the tool calls
 				var toolDefs []tools.Tool
@@ -855,14 +856,23 @@ func (r *LocalRuntime) RunStream(ctx context.Context, sess *session.Session) <-c
 					Cost:              messageCost,
 				}
 
+				// Build per-message usage for the event
+				if res.Usage != nil {
+					msgUsage = &MessageUsage{
+						Usage: *res.Usage,
+						Cost:  messageCost,
+						Model: messageModel,
+					}
+				}
+
 				sess.AddMessage(session.NewAgentMessage(a, &assistantMessage))
 				r.saveSession(ctx, sess)
 				slog.Debug("Added assistant message to session", "agent", a.Name(), "total_messages", len(sess.GetAllMessages()))
 			} else {
 				slog.Debug("Skipping empty assistant message (no content and no tool calls)", "agent", a.Name())
 			}
 
-			events <- TokenUsage(sess.ID, r.currentAgent, sess.InputTokens, sess.OutputTokens, sess.InputTokens+sess.OutputTokens, contextLimit, sess.Cost)
+			events <- TokenUsageWithMessage(sess.ID, r.currentAgent, sess.InputTokens, sess.OutputTokens, sess.InputTokens+sess.OutputTokens, contextLimit, sess.Cost, msgUsage)
 
 			r.processToolCalls(ctx, sess, res.Calls, agentTools, events)
 

pkg/runtime/connectrpc_client.go

@@ -219,7 +219,10 @@ func TestSimple(t *testing.T) {
 		UserMessage("Hi"),
 		StreamStarted(sess.ID, "root"),
 		AgentChoice("root", "Hello"),
-		TokenUsage(sess.ID, "root", 3, 2, 5, 0, 0),
+		TokenUsageWithMessage(sess.ID, "root", 3, 2, 5, 0, 0, &MessageUsage{
+			Usage: chat.Usage{InputTokens: 3, OutputTokens: 2},
+			Model: "test/mock-model",
+		}),
 		StreamStopped(sess.ID, "root"),
 	}
 
@@ -251,7 +254,10 @@ func TestMultipleContentChunks(t *testing.T) {
 		AgentChoice("root", "how "),
 		AgentChoice("root", "are "),
 		AgentChoice("root", "you?"),
-		TokenUsage(sess.ID, "root", 8, 12, 20, 0, 0),
+		TokenUsageWithMessage(sess.ID, "root", 8, 12, 20, 0, 0, &MessageUsage{
+			Usage: chat.Usage{InputTokens: 8, OutputTokens: 12},
+			Model: "test/mock-model",
+		}),
 		StreamStopped(sess.ID, "root"),
 	}
 
@@ -279,7 +285,10 @@ func TestWithReasoning(t *testing.T) {
 		AgentChoiceReasoning("root", "Let me think about this..."),
 		AgentChoiceReasoning("root", " I should respond politely."),
 		AgentChoice("root", "Hello, how can I help you?"),
-		TokenUsage(sess.ID, "root", 10, 15, 25, 0, 0),
+		TokenUsageWithMessage(sess.ID, "root", 10, 15, 25, 0, 0, &MessageUsage{
+			Usage: chat.Usage{InputTokens: 10, OutputTokens: 15},
+			Model: "test/mock-model",
+		}),
 		StreamStopped(sess.ID, "root"),
 	}
 
@@ -309,7 +318,10 @@ func TestMixedContentAndReasoning(t *testing.T) {
 		AgentChoice("root", "Hello!"),
 		AgentChoiceReasoning("root", " I should be friendly"),
 		AgentChoice("root", " How can I help you today?"),
-		TokenUsage(sess.ID, "root", 15, 20, 35, 0, 0),
+		TokenUsageWithMessage(sess.ID, "root", 15, 20, 35, 0, 0, &MessageUsage{
+			Usage: chat.Usage{InputTokens: 15, OutputTokens: 20},
+			Model: "test/mock-model",
+		}),
 		StreamStopped(sess.ID, "root"),
 	}
 

pkg/runtime/event.go

@@ -99,6 +99,20 @@ type Session struct {
 	// Sub-sessions are not persisted as standalone entries; they are embedded
 	// within the parent session's Messages array.
 	ParentID string `json:"-"`
+
+	// MessageUsageHistory stores per-message usage data for remote mode.
+	// In remote mode, messages are managed server-side, so we track usage separately.
+	// This is not persisted (json:"-") as it's only needed for the current session display.
+	MessageUsageHistory []MessageUsageRecord `json:"-"`
+}
+
+// MessageUsageRecord stores usage data for a single assistant message.
+// Used in remote mode where messages aren't stored in the client-side session.
+type MessageUsageRecord struct {
+	AgentName string     `json:"agent_name"`
+	Model     string     `json:"model"`
+	Cost      float64    `json:"cost"`
+	Usage     chat.Usage `json:"usage"`
 }
 
 // Permission mode constants
@@ -300,6 +314,35 @@ func (s *Session) getLastMessageContentByRole(role chat.MessageRole) string {
 	return ""
 }
 
+// UpdateLastAssistantMessageUsage updates the usage and cost fields of the last assistant message.
+// This is used in remote mode to populate per-message cost data from TokenUsageEvent.
+func (s *Session) UpdateLastAssistantMessageUsage(usage *chat.Usage, cost float64, model string) {
+	for i := len(s.Messages) - 1; i >= 0; i-- {
+		if s.Messages[i].IsMessage() && s.Messages[i].Message.Message.Role == chat.MessageRoleAssistant {
+			s.Messages[i].Message.Message.Usage = usage
+			s.Messages[i].Message.Message.Cost = cost
+			if model != "" {
+				s.Messages[i].Message.Message.Model = model
+			}
+			return
+		}
+	}
+}
+
+// AddMessageUsageRecord appends a usage record for remote mode where messages aren't stored locally.
+// This enables the /cost dialog to show per-message breakdown even when using a remote runtime.
+func (s *Session) AddMessageUsageRecord(agentName, model string, cost float64, usage *chat.Usage) {
+	if usage == nil {
+		return
+	}
+	s.MessageUsageHistory = append(s.MessageUsageHistory, MessageUsageRecord{
+		AgentName: agentName,
+		Model:     model,
+		Cost:      cost,
+		Usage:     *usage,
+	})
+}
+
 type Opt func(s *Session)
 
 func WithUserMessage(content string) Opt {

pkg/runtime/runtime.go

@@ -217,3 +217,104 @@ func TestGetMessages_CacheControlWithSummary(t *testing.T) {
 	// Verify checkpoint #2 is on date
 	assert.Contains(t, messages[checkpointIndices[1]].Content, "Today's date", "checkpoint #2 should be on date message")
 }
+
+func TestUpdateLastAssistantMessageUsage(t *testing.T) {
+	testAgent := &agent.Agent{}
+
+	s := New()
+
+	// Add user message
+	s.AddMessage(NewAgentMessage(testAgent, &chat.Message{
+		Role:    chat.MessageRoleUser,
+		Content: "hello",
+	}))
+
+	// Add assistant message without usage
+	s.AddMessage(NewAgentMessage(testAgent, &chat.Message{
+		Role:    chat.MessageRoleAssistant,
+		Content: "response",
+	}))
+
+	// Update the last assistant message with usage data
+	usage := &chat.Usage{
+		InputTokens:       100,
+		OutputTokens:      50,
+		CachedInputTokens: 10,
+	}
+	s.UpdateLastAssistantMessageUsage(usage, 0.005, "gpt-4")
+
+	// Verify the update
+	messages := s.GetAllMessages()
+	assert.Len(t, messages, 2)
+
+	lastMsg := messages[1]
+	assert.Equal(t, chat.MessageRoleAssistant, lastMsg.Message.Role)
+	assert.NotNil(t, lastMsg.Message.Usage)
+	assert.Equal(t, int64(100), lastMsg.Message.Usage.InputTokens)
+	assert.Equal(t, int64(50), lastMsg.Message.Usage.OutputTokens)
+	assert.Equal(t, int64(10), lastMsg.Message.Usage.CachedInputTokens)
+	assert.InEpsilon(t, 0.005, lastMsg.Message.Cost, 0.0001)
+	assert.Equal(t, "gpt-4", lastMsg.Message.Model)
+}
+
+func TestUpdateLastAssistantMessageUsage_NoAssistantMessage(t *testing.T) {
+	testAgent := &agent.Agent{}
+
+	s := New()
+
+	// Add only user message
+	s.AddMessage(NewAgentMessage(testAgent, &chat.Message{
+		Role:    chat.MessageRoleUser,
+		Content: "hello",
+	}))
+
+	// Should not panic when no assistant message exists
+	usage := &chat.Usage{InputTokens: 100}
+	s.UpdateLastAssistantMessageUsage(usage, 0.01, "model")
+
+	// Verify nothing changed
+	messages := s.GetAllMessages()
+	assert.Len(t, messages, 1)
+	assert.Equal(t, chat.MessageRoleUser, messages[0].Message.Role)
+}
+
+func TestUpdateLastAssistantMessageUsage_UpdatesOnlyLast(t *testing.T) {
+	testAgent := &agent.Agent{}
+
+	s := New()
+
+	// Add multiple assistant messages
+	s.AddMessage(NewAgentMessage(testAgent, &chat.Message{
+		Role:    chat.MessageRoleAssistant,
+		Content: "first response",
+		Usage:   &chat.Usage{InputTokens: 10},
+	}))
+
+	s.AddMessage(NewAgentMessage(testAgent, &chat.Message{
+		Role:    chat.MessageRoleUser,
+		Content: "follow up",
+	}))
+
+	s.AddMessage(NewAgentMessage(testAgent, &chat.Message{
+		Role:    chat.MessageRoleAssistant,
+		Content: "second response",
+	}))
+
+	// Update usage - should only affect the last assistant message
+	usage := &chat.Usage{InputTokens: 200}
+	s.UpdateLastAssistantMessageUsage(usage, 0.02, "new-model")
+
+	// Verify only the last assistant message was updated
+	messages := s.GetAllMessages()
+	assert.Len(t, messages, 3)
+
+	// First assistant message should keep original usage
+	assert.NotNil(t, messages[0].Message.Usage)
+	assert.Equal(t, int64(10), messages[0].Message.Usage.InputTokens)
+
+	// Last assistant message should have new usage
+	assert.NotNil(t, messages[2].Message.Usage)
+	assert.Equal(t, int64(200), messages[2].Message.Usage.InputTokens)
+	assert.InEpsilon(t, 0.02, messages[2].Message.Cost, 0.0001)
+	assert.Equal(t, "new-model", messages[2].Message.Model)
+}