diff --git a/pkg/app/app.go b/pkg/app/app.go index 3f5b37197..80ba295cf 100644 --- a/pkg/app/app.go +++ b/pkg/app/app.go @@ -161,11 +161,15 @@ func (a *App) Start(ctx context.Context) { // Emit startup info (agent, team, tools) through the events channel. // This runs in the background so the TUI can start immediately while // slow operations (like MCP tool loading) complete asynchronously. + // Snapshot the session on this goroutine: ReplaceSession may swap + // a.session concurrently, and it re-emits startup info for the new + // session itself. + sess := a.session go func() { startupEvents := make(chan runtime.Event, 10) go func() { defer close(startupEvents) - a.runtime.EmitStartupInfo(ctx, a.session, runtime.NewChannelSink(startupEvents)) + a.runtime.EmitStartupInfo(ctx, sess, runtime.NewChannelSink(startupEvents)) }() for event := range startupEvents { select { @@ -415,6 +419,9 @@ func (a *App) SkillCommandFork(_ context.Context, input string) (skillName, task // SKILL.md body. Companion of SkillCommandFork. func (a *App) RunSkillFork(ctx context.Context, cancel context.CancelFunc, skillName, task string, _ []messages.Attachment) { a.cancel = cancel + // Snapshot the session like Run does: the goroutines below outlive any + // concurrent ReplaceSession and must keep working against this session. + sess := a.session // Mirrors App.Run's drain loop: forward events to the App bus and // always let StreamStoppedEvent through, even after ctx cancellation, @@ -424,7 +431,7 @@ func (a *App) RunSkillFork(ctx context.Context, cancel context.CancelFunc, skill var failed atomic.Bool go func() { defer close(events) - result, err := a.runtime.RunSkillFork(ctx, a.session, skillstool.RunSkillArgs{ + result, err := a.runtime.RunSkillFork(ctx, sess, skillstool.RunSkillArgs{ Name: skillName, Task: task, }, runtime.NewChannelSink(events)) @@ -478,7 +485,7 @@ func (a *App) RunSkillFork(ctx context.Context, cancel context.CancelFunc, skill } if !sawStop { - a.synthesizeStreamStopped(ctx, cmp.Or(lastSessionID, a.session.ID), agentName, failed.Load()) + a.synthesizeStreamStopped(ctx, cmp.Or(lastSessionID, sess.ID), agentName, failed.Load()) } }() } @@ -1229,8 +1236,11 @@ func (a *App) NewSession() { // through the events channel so the sidebar updates. func (a *App) reEmitStartupInfo(ctx context.Context) { a.runtime.ResetStartupInfo() + // Snapshot before handing off to the background goroutine so a later + // ReplaceSession cannot race with this read. + sess := a.session a.pumpToEvents(ctx, func(sink runtime.EventSink) { - a.runtime.EmitStartupInfo(ctx, a.session, sink) + a.runtime.EmitStartupInfo(ctx, sess, sink) }) } diff --git a/pkg/app/session_lifecycle_test.go b/pkg/app/session_lifecycle_test.go index 2b6070223..83f34268b 100644 --- a/pkg/app/session_lifecycle_test.go +++ b/pkg/app/session_lifecycle_test.go @@ -16,6 +16,7 @@ import ( "github.com/docker/docker-agent/pkg/session" "github.com/docker/docker-agent/pkg/sessiontitle" "github.com/docker/docker-agent/pkg/tools" + skillstool "github.com/docker/docker-agent/pkg/tools/builtin/skills" ) type sessionCaptureRuntime struct { @@ -35,6 +36,75 @@ func (r *sessionCaptureRuntime) RunStream(_ context.Context, sess *session.Sessi return ch } +// backgroundSessionCaptureRuntime records the session handed to the runtime +// entry points App drives from background goroutines, and holds each call +// until release is closed so a concurrent ReplaceSession can be interleaved. +type backgroundSessionCaptureRuntime struct { + mockRuntime + + started chan *session.Session + release chan struct{} +} + +func (r *backgroundSessionCaptureRuntime) EmitStartupInfo(_ context.Context, sess *session.Session, _ runtime.EventSink) { + r.started <- sess + <-r.release +} + +func (r *backgroundSessionCaptureRuntime) RunSkillFork(_ context.Context, sess *session.Session, _ skillstool.RunSkillArgs, _ runtime.EventSink) (*tools.ToolCallResult, error) { + r.started <- sess + <-r.release + return nil, nil +} + +// TestAppBackgroundWorkSnapshotsSessionBeforeReplace covers every App entry +// point that hands a.session to a background goroutine: the goroutine must +// receive the session that was current when it was spawned, and must not +// read the field itself, which races with ReplaceSession (#4229). Under +// -race a field read from the goroutine fails this test deterministically. +func TestAppBackgroundWorkSnapshotsSessionBeforeReplace(t *testing.T) { + for _, entryPoint := range []struct { + name string + run func(*App, context.Context, context.CancelFunc) + }{ + {name: "Start", run: func(app *App, ctx context.Context, _ context.CancelFunc) { + app.Start(ctx) + }}, + {name: "reEmitStartupInfo", run: func(app *App, ctx context.Context, _ context.CancelFunc) { + app.reEmitStartupInfo(ctx) + }}, + {name: "RunSkillFork", run: func(app *App, ctx context.Context, cancel context.CancelFunc) { + app.RunSkillFork(ctx, cancel, "skill", "task", nil) + }}, + } { + t.Run(entryPoint.name, func(t *testing.T) { + oldSession := session.New() + rt := &backgroundSessionCaptureRuntime{ + started: make(chan *session.Session, 4), + release: make(chan struct{}), + } + app := &App{ + runtime: rt, + session: oldSession, + events: make(chan tea.Msg, 16), + } + ctx, cancel := context.WithCancel(t.Context()) + entryPoint.run(app, ctx, cancel) + + // Replace the session right away, before the background goroutine + // has necessarily started running; ReplaceSession re-emits startup + // info for the new session, so two sessions reach the runtime. + newSession := session.New() + app.ReplaceSession(t.Context(), newSession) + first, second := <-rt.started, <-rt.started + close(rt.release) + + assert.ElementsMatch(t, []*session.Session{oldSession, newSession}, []*session.Session{first, second}, + "background work must run against the session current when it was spawned") + }) + } +} + func TestAppRunKeepsWorkScopedToOriginalSession(t *testing.T) { for _, entryPoint := range []struct { name string