Merge pull request #3659 from docker/fix/3591-session-scalar-locking

fix(session): lock title and usage scalar access

Author: Sayt-0

pkg/a2a/adapter.go

@@ -91,7 +91,7 @@ func runDockerAgent(ctx agent.InvocationContext, t *team.Team, agentName string,
 				session.WithNonInteractive(true),
 				session.WithWorkingDir(workingDir),
 			)
-			sess.Title = "A2A Session " + sessionID
+			sess.SetTitle("A2A Session " + sessionID)
 		}
 
 		// Create runtime

pkg/acp/agent.go

@@ -214,7 +214,7 @@ func (a *Agent) NewSession(ctx context.Context, params acp.NewSessionRequest) (a
 		session.WithMaxToolResultTokens(defaultAgent.MaxToolResultTokens()),
 		session.WithWorkingDir(workingDir),
 	)
-	sess.Title = "ACP Session " + sess.ID
+	sess.SetTitle("ACP Session " + sess.ID)
 
 	if err := a.sessionStore.AddSession(ctx, sess); err != nil {
 		return acp.NewSessionResponse{}, fmt.Errorf("failed to persist session: %w", err)

pkg/app/app.go

@@ -484,7 +484,7 @@ func (a *App) Run(ctx context.Context, cancel context.CancelFunc, message string
 	a.cancel = cancel
 
 	// If this is the first message and no title exists, start local title generation
-	if a.session.Title == "" && a.titleGen != nil {
+	if a.session.TitleSnapshot() == "" && a.titleGen != nil {
 		a.titleGenerating.Store(true)
 		go a.generateTitle(ctx, []string{message})
 	}
@@ -764,7 +764,7 @@ func (a *App) RunWithMessage(ctx context.Context, cancel context.CancelFunc, msg
 	a.cancel = cancel
 
 	// If this is the first message and no title exists, start local title generation
-	if a.session.Title == "" && a.titleGen != nil {
+	if a.session.TitleSnapshot() == "" && a.titleGen != nil {
 		a.titleGenerating.Store(true)
 		// Extract text content from the message for title generation
 		userMessage := msg.Message.Content

pkg/cli/runner.go

@@ -90,7 +90,7 @@ func Run(ctx context.Context, out *Printer, cfg Config, rt runtime.Runtime, sess
 		ctx = telemetry.WithClient(ctx, telemetryClient)
 	}
 
-	sess.Title = "Running agent"
+	sess.SetTitle("Running agent")
 	// If the last received event was an error, return it. That way the exit code
 	// will be non-zero if the agent failed.
 	var lastErr error

pkg/evaluation/eval.go

@@ -360,7 +360,7 @@ func (r *Runner) runSingleEval(ctx context.Context, evalSess *InputSession) (Res
 
 	// Re-apply the display title in case a session_title event overrode it.
 	// This ensures repeated evals retain their '#N' suffix in stored sessions.
-	result.Session.Title = title
+	result.Session.SetTitle(title)
 
 	if len(expectedToolCalls) > 0 || len(actualToolCalls) > 0 {
 		result.ToolCallsScore = toolCallF1Score(expectedToolCalls, actualToolCalls)

pkg/evaluation/save.go

@@ -211,15 +211,20 @@ func SessionFromEvents(events []map[string]any, title string, questions []string
 		case "token_usage":
 			// Update session token usage
 			if usage, ok := event["usage"].(map[string]any); ok {
+				// Route the write through the locked setter, keeping the
+				// previous value for any field absent from the event. sess is
+				// local to this import, so the pre-read cannot race.
+				sessInput, sessOutput, sessCost := sess.TokensAndCost()
 				if inputTokens, ok := usage["input_tokens"].(float64); ok {
-					sess.InputTokens = int64(inputTokens)
+					sessInput = int64(inputTokens)
 				}
 				if outputTokens, ok := usage["output_tokens"].(float64); ok {
-					sess.OutputTokens = int64(outputTokens)
+					sessOutput = int64(outputTokens)
 				}
 				if cost, ok := usage["cost"].(float64); ok {
-					sess.Cost = cost
+					sessCost = cost
 				}
+				sess.SetTokensAndCost(sessInput, sessOutput, sessCost)
 				// Extract per-message usage if available
 				if lastMsg, ok := usage["last_message"].(map[string]any); ok {
 					currentUsage = parseMessageUsage(lastMsg)
@@ -251,7 +256,7 @@ func SessionFromEvents(events []map[string]any, title string, questions []string
 		case "session_title":
 			// Update session title if provided (may override the one from eval config)
 			if eventTitle, ok := event["title"].(string); ok && eventTitle != "" {
-				sess.Title = eventTitle
+				sess.SetTitle(eventTitle)
 			}
 
 		case "stream_stopped":

pkg/evaluation/save_test.go

@@ -468,6 +468,71 @@ func TestSessionFromEventsTokenUsage(t *testing.T) {
 	assert.InDelta(t, 0.005, sess.Cost, 0.0001)
 }
 
+// TestSessionFromEventsPartialTokenUsage pins the merge semantics of
+// consecutive token_usage events: a later event that omits some usage
+// fields must preserve the previously recorded values for the omitted
+// fields instead of resetting them to zero.
+func TestSessionFromEventsPartialTokenUsage(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name       string
+		second     map[string]any
+		wantInput  int64
+		wantOutput int64
+		wantCost   float64
+	}{
+		{
+			name:       "second event omits input_tokens and cost",
+			second:     map[string]any{"output_tokens": float64(80)},
+			wantInput:  100,
+			wantOutput: 80,
+			wantCost:   0.005,
+		},
+		{
+			name:       "second event omits output_tokens",
+			second:     map[string]any{"input_tokens": float64(200), "cost": 0.01},
+			wantInput:  200,
+			wantOutput: 50,
+			wantCost:   0.01,
+		},
+		{
+			name:       "second event omits every field",
+			second:     map[string]any{},
+			wantInput:  100,
+			wantOutput: 50,
+			wantCost:   0.005,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			t.Parallel()
+
+			events := []map[string]any{
+				{"type": "agent_choice", "content": "Answer"},
+				{
+					"type": "token_usage",
+					"usage": map[string]any{
+						"input_tokens":  float64(100),
+						"output_tokens": float64(50),
+						"cost":          0.005,
+					},
+				},
+				{"type": "token_usage", "usage": tt.second},
+				{"type": "stream_stopped"},
+			}
+
+			sess := SessionFromEvents(events, "test", []string{"question"})
+
+			gotInput, gotOutput, gotCost := sess.TokensAndCost()
+			assert.Equal(t, tt.wantInput, gotInput)
+			assert.Equal(t, tt.wantOutput, gotOutput)
+			assert.InDelta(t, tt.wantCost, gotCost, 0.0001)
+		})
+	}
+}
+
 func TestParseToolCall(t *testing.T) {
 	t.Parallel()
 

pkg/runtime/event.go

@@ -404,10 +404,13 @@ func (e *TokenUsageEvent) GetSessionID() string { return e.SessionID }
 // compactionThreshold is the agent's configured auto-compaction trigger
 // fraction (0 or omitted when unknown), passed along verbatim for UI gauges.
 func SessionUsage(sess *session.Session, contextLimit int64, compactionThreshold ...float64) *Usage {
+	// Usage() snapshots both counters under sess.mu so a concurrent
+	// SetUsage/ApplyCompaction cannot tear the pair.
+	input, output := sess.Usage()
 	u := &Usage{
-		InputTokens:   sess.InputTokens,
-		OutputTokens:  sess.OutputTokens,
-		ContextLength: sess.InputTokens + sess.OutputTokens,
+		InputTokens:   input,
+		OutputTokens:  output,
+		ContextLength: input + output,
 		ContextLimit:  contextLimit,
 		Cost:          sess.OwnCost(),
 	}

pkg/runtime/hooks.go

@@ -502,10 +502,11 @@ func (r *LocalRuntime) executeBeforeCompactionHooks(
 	contextLimit int64,
 	events EventSink,
 ) *hooks.Result {
+	inputTokens, outputTokens := sess.Usage()
 	return r.dispatchHook(ctx, a, hooks.EventBeforeCompaction, &hooks.Input{
 		SessionID:        sess.ID,
-		InputTokens:      sess.InputTokens,
-		OutputTokens:     sess.OutputTokens,
+		InputTokens:      inputTokens,
+		OutputTokens:     outputTokens,
 		ContextLimit:     contextLimit,
 		CompactionReason: reason,
 	}, events)

pkg/runtime/loop.go

@@ -525,7 +525,8 @@ func (r *LocalRuntime) runStreamLoop(ctx context.Context, sess *session.Session,
 		// trigger and the UI context gauge (issue #3241); it equals the primary
 		// window unless a smaller compaction model is configured.
 		contextLimit := r.effectiveContextLimit(ctx, a, r.resolveContextLimit(ctx, model, modelID))
-		if contextLimit > 0 && r.sessionCompactionEnabled(a) && compaction.ShouldCompact(sess.InputTokens, sess.OutputTokens, 0, contextLimit, a.CompactionThreshold()) {
+		inputTokens, outputTokens := sess.Usage()
+		if contextLimit > 0 && r.sessionCompactionEnabled(a) && compaction.ShouldCompact(inputTokens, outputTokens, 0, contextLimit, a.CompactionThreshold()) {
 			r.compactWithReason(ctx, sess, "", compactionReasonThreshold, sink)
 		}
 
@@ -1161,17 +1162,18 @@ func (r *LocalRuntime) compactIfNeeded(
 		addedTokens += estimator.EstimateMessageTokens(&newMessages[i].Message)
 	}
 
-	if !compaction.ShouldCompact(sess.InputTokens, sess.OutputTokens, addedTokens, contextLimit, a.CompactionThreshold()) {
+	inputTokens, outputTokens := sess.Usage()
+	if !compaction.ShouldCompact(inputTokens, outputTokens, addedTokens, contextLimit, a.CompactionThreshold()) {
 		return
 	}
 
 	slog.InfoContext(ctx, "Proactive compaction: tool results pushed estimated context past the compaction threshold",
 		"agent", a.Name(),
-		"input_tokens", sess.InputTokens,
-		"output_tokens", sess.OutputTokens,
+		"input_tokens", inputTokens,
+		"output_tokens", outputTokens,
 		"added_estimated_tokens", addedTokens,
 		"estimator_scale", estimator.Scale(),
-		"estimated_total", sess.InputTokens+sess.OutputTokens+addedTokens,
+		"estimated_total", inputTokens+outputTokens+addedTokens,
 		"context_limit", contextLimit,
 	)
 	r.compactWithReason(ctx, sess, "", compactionReasonThreshold, events)