From d9faf2bd366dea93230cefac1dabe11739c62157 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Arnaud=20He=CC=81ritier?= <arnaud.heritier@docker.com>
Date: Sat, 11 Jul 2026 18:27:56 +0200
Subject: [PATCH] fix(runtime): make startupInfoEmitted an atomic.Bool guarded
 by CAS (#3593)

---
 pkg/runtime/runtime.go               | 12 ++--
 pkg/runtime/startupinfo_race_test.go | 91 ++++++++++++++++++++++++++++
 2 files changed, 98 insertions(+), 5 deletions(-)
 create mode 100644 pkg/runtime/startupinfo_race_test.go

diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go
index fa6a5365f..1e0fa996a 100644
--- a/pkg/runtime/runtime.go
+++ b/pkg/runtime/runtime.go
@@ -224,7 +224,7 @@ type LocalRuntime struct {
 	managedOAuth              bool
 	unmanagedOAuthRedirectURI string
 	nonInteractive            bool
-	startupInfoEmitted        bool                   // Track if startup info has been emitted to avoid unnecessary duplication
+	startupInfoEmitted        atomic.Bool            // Track if startup info has been emitted to avoid unnecessary duplication
 	elicitationRequestCh      chan ElicitationResult // Channel for receiving elicitation responses
 	elicitation               elicitationBridge      // Owns the per-stream events channel for outbound elicitation requests
 	sessionStore              session.Store
@@ -1421,7 +1421,7 @@ func (r *LocalRuntime) PermissionsInfo() *PermissionsInfo {
 // This should be called when replacing a session to allow re-emission of
 // agent, team, and toolset info to the UI.
 func (r *LocalRuntime) ResetStartupInfo() {
-	r.startupInfoEmitted = false
+	r.startupInfoEmitted.Store(false)
 }
 
 // OnToolsChanged registers a handler that is called when an MCP toolset
@@ -1526,11 +1526,13 @@ func (r *LocalRuntime) EmitAgentInfo(ctx context.Context, events EventSink) {
 // When sess is non-nil and contains token data, a TokenUsageEvent is also emitted so that the
 // sidebar can display context usage percentage on session restore.
 func (r *LocalRuntime) EmitStartupInfo(ctx context.Context, sess *session.Session, events EventSink) {
-	// Prevent duplicate emissions
-	if r.startupInfoEmitted {
+	// CompareAndSwap makes the check-and-set atomic: App.Start emits from a
+	// spawned goroutine while reEmitStartupInfo (e.g. on /new session) resets
+	// and re-emits from another, and a plain bool check-then-set race would
+	// let both goroutines through or miss an emission.
+	if !r.startupInfoEmitted.CompareAndSwap(false, true) {
 		return
 	}
-	r.startupInfoEmitted = true
 
 	a := r.CurrentAgent()
 
diff --git a/pkg/runtime/startupinfo_race_test.go b/pkg/runtime/startupinfo_race_test.go
new file mode 100644
index 000000000..5bd8e44f0
--- /dev/null
+++ b/pkg/runtime/startupinfo_race_test.go
@@ -0,0 +1,91 @@
+package runtime
+
+import (
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/docker/docker-agent/pkg/agent"
+	"github.com/docker/docker-agent/pkg/session"
+	"github.com/docker/docker-agent/pkg/team"
+)
+
+// TestStartupInfoEmittedConcurrentStartAndReEmit pins the fix for
+// startupInfoEmitted: it used to be a plain bool checked and set with two
+// separate statements ("if r.startupInfoEmitted { return }; r.startupInfoEmitted
+// = true"). App.Start emits startup info from a spawned goroutine while
+// ResetStartupInfo + EmitStartupInfo (the reEmitStartupInfo path, e.g. on
+// /new session) reset and re-emit from another goroutine — the
+// check-then-set was racy and could let both goroutines through, or drop
+// both. The flag must be an atomic.Bool driven by CompareAndSwap so exactly
+// one EmitStartupInfo call wins each round.
+func TestStartupInfoEmittedConcurrentStartAndReEmit(t *testing.T) {
+	t.Parallel()
+
+	prov := &mockProvider{id: "test/mock-model"}
+	root := agent.New("root", "test", agent.WithModel(prov))
+	tm := team.New(team.WithAgents(root))
+
+	rt, err := NewLocalRuntime(t.Context(), tm, WithModelStore(mockModelStore{}))
+	require.NoError(t, err)
+
+	sess := session.New()
+
+	var wg sync.WaitGroup
+	for range 50 {
+		wg.Go(func() {
+			events := make(chan Event, 16)
+			rt.EmitStartupInfo(t.Context(), sess, NewChannelSink(events))
+			close(events)
+			for range events {
+			}
+		})
+		wg.Go(func() {
+			rt.ResetStartupInfo()
+		})
+	}
+	wg.Wait()
+}
+
+// TestStartupInfoEmittedSingleEmissionWithoutReset pins the "avoid
+// unnecessary duplication" contract: concurrent EmitStartupInfo calls with
+// no ResetStartupInfo interleaved must result in exactly one goroutine
+// actually emitting events.
+func TestStartupInfoEmittedSingleEmissionWithoutReset(t *testing.T) {
+	t.Parallel()
+
+	prov := &mockProvider{id: "test/mock-model"}
+	root := agent.New("root", "test", agent.WithModel(prov))
+	tm := team.New(team.WithAgents(root))
+
+	rt, err := NewLocalRuntime(t.Context(), tm, WithModelStore(mockModelStore{}))
+	require.NoError(t, err)
+
+	sess := session.New()
+
+	var mu sync.Mutex
+	var emissions int
+	var wg sync.WaitGroup
+	for range 50 {
+		wg.Go(func() {
+			events := make(chan Event, 16)
+			rt.EmitStartupInfo(t.Context(), sess, NewChannelSink(events))
+			close(events)
+			n := 0
+			for range events {
+				n++
+			}
+			if n > 0 {
+				mu.Lock()
+				emissions++
+				mu.Unlock()
+			}
+		})
+	}
+	wg.Wait()
+
+	if emissions != 1 {
+		t.Errorf("expected exactly 1 emission across concurrent calls, got %d", emissions)
+	}
+}