Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions pkg/runtime/runtime.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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()

Expand Down
91 changes: 91 additions & 0 deletions pkg/runtime/startupinfo_race_test.go
Original file line number Diff line number Diff line change
@@ -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)
}
}
Loading