Skip to content

fix: don't close shared session store in runtime.Close #2879

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
dgageot wants to merge 2 commits into
base: main
Choose a base branch
from
+131 −18
Open

fix: don't close shared session store in runtime.Close #2879

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
fdf18eb
fix: don't close shared session store in runtime.Close (#2872)
dgageot May 22, 2026
5e988e0
fix: ensure visibility of b.store in localBackend.Close
dgageot May 24, 2026
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
60 changes: 59 additions & 1 deletion cmd/root/backend.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"sync"

"github.com/docker/docker-agent/pkg/config"
pathx "github.com/docker/docker-agent/pkg/path"
"github.com/docker/docker-agent/pkg/runtime"
"github.com/docker/docker-agent/pkg/session"
"github.com/docker/docker-agent/pkg/teamloader"
Expand All @@ -28,6 +29,11 @@ type backend interface {
CreateSession(ctx context.Context, loaded *teamloader.LoadResult, req runtime.CreateSessionRequest) (runtime.Runtime, *session.Session, func(), error)

Spawner(rt runtime.Runtime) tui.SessionSpawner

// Close releases backend-owned resources (e.g., the shared session
// store). It is called once when the embedder is shutting down,
// after every per-session cleanup has run.
Close() error
}

// selectBackend picks the backend implied by the current flags.
Expand All @@ -43,9 +49,20 @@ func (f *runExecFlags) selectBackend(agentFileName string) (backend, error) {
}

// localBackend builds the in-process runtime and session.
//
// The session store is owned by the backend (not by individual
// runtimes) because the TUI's session spawner reuses the same store
// across spawned sessions. Closing it inside a per-session cleanup
// would break later session lookups (issue #2872). The store is
// lazily opened on the first CreateSession call and closed once when
// Close is invoked.
type localBackend struct {
flags *runExecFlags
agentSource config.Source

storeOnce sync.Once
storeErr error
store session.Store
}

func (b *localBackend) LoadTeamRequest() runtime.LoadTeamRequest {
Expand All @@ -60,8 +77,34 @@ func (b *localBackend) CreateSessionRequest(workingDir string) runtime.CreateSes
return b.flags.createSessionRequest(workingDir)
}

// sessionStore returns the backend-owned session store, opening it on
// first use. The store is shared by the initial runtime and by every
// runtime spawned by [localBackend.Spawner].
func (b *localBackend) sessionStore(req runtime.CreateSessionRequest) (session.Store, error) {
b.storeOnce.Do(func() {
Copy link
Copy Markdown
Contributor

@aheritier aheritier May 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Sticky sync.Once failure (low severity): if ExpandHomeDir or NewSQLiteSessionStore errors here, storeOnce is permanently consumed and every subsequent call to sessionStore() will return (nil, storeErr) — even if the underlying problem has been corrected (e.g., a retry with a different req). In the current call graph CreateSession is called exactly once, so this is not a practical issue. Flagging in case future callers add retry logic.

sessionDB, err := pathx.ExpandHomeDir(req.SessionDB)
if err != nil {
b.storeErr = err
return
}
store, err := session.NewSQLiteSessionStore(sessionDB)
if err != nil {
b.storeErr = fmt.Errorf("creating session store: %w", err)
return
}
b.store = store
})
return b.store, b.storeErr
}

func (b *localBackend) CreateSession(ctx context.Context, loaded *teamloader.LoadResult, req runtime.CreateSessionRequest) (runtime.Runtime, *session.Session, func(), error) {
rt, sess, err := b.flags.createLocalRuntimeAndSession(ctx, loaded, req)
store, err := b.sessionStore(req)
if err != nil {
stopToolSets(loaded.Team)
return nil, nil, nil, err
}

rt, sess, err := b.flags.createLocalRuntimeAndSession(ctx, loaded, req, store)
if err != nil {
stopToolSets(loaded.Team)
return nil, nil, nil, err
Expand All @@ -83,6 +126,17 @@ func (b *localBackend) Spawner(rt runtime.Runtime) tui.SessionSpawner {
return b.flags.createSessionSpawner(b.agentSource, rt.SessionStore())
}

func (b *localBackend) Close() error {
// Ensure any in-progress sessionStore initialization is observed
// before reading b.store. If sessionStore was never called, the Do
// runs an empty closure and b.store stays nil.
b.storeOnce.Do(func() {})
if b.store == nil {
Copy link
Copy Markdown

@docker-agent docker-agent May 22, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[MEDIUM] Potential data race: Close() reads b.store outside sync.Once

localBackend.Close() reads b.store directly (line 130) without any synchronization, while sessionStore() writes b.store inside storeOnce.Do(...). The Go memory model only guarantees that the write to b.store inside the Do closure is visible after Do returns — it does not provide a happens-before guarantee for concurrent reads of b.store that happen outside of Do.

If Close() races with an in-progress storeOnce.Do() call (e.g., a session is still being initialized when shutdown begins), the Go race detector would flag this.

Suggested fix — either use storeOnce.Do to query the store in Close() as well, or protect b.store with a mutex:

func (b *localBackend) Close() error {
    b.storeOnce.Do(func() {}) // ensure the write, if any, has completed
    if b.store == nil {
        return nil
    }
    return b.store.Close()
}

Or, more robustly, add a sync.Mutex to guard access to b.store in both sessionStore and Close.

Copy link
Copy Markdown
Contributor

@aheritier aheritier May 23, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nil guard is correct and defensive. b.store can only be nil here if CreateSession was never called (e.g., an early-exit code path) or if storeOnce recorded an error without setting b.store. The check is not racy: b.Close() is invoked from the defer registered after selectBackend(), by which point any concurrent Do() invocations have completed. Technically redundant on the happy path, but harmless and good defensive practice.

return nil
}
return b.store.Close()
}

// remoteBackend talks to a docker-agent server.
type remoteBackend struct {
flags *runExecFlags
Expand Down Expand Up @@ -142,3 +196,7 @@ func (b *remoteBackend) CreateSession(ctx context.Context, _ *teamloader.LoadRes
func (b *remoteBackend) Spawner(runtime.Runtime) tui.SessionSpawner {
return nil
}

func (b *remoteBackend) Close() error {
return nil
}
18 changes: 6 additions & 12 deletions cmd/root/run.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,6 @@ import (
"github.com/docker/docker-agent/pkg/config"
"github.com/docker/docker-agent/pkg/hooks"
"github.com/docker/docker-agent/pkg/hooks/builtins"
pathx "github.com/docker/docker-agent/pkg/path"
"github.com/docker/docker-agent/pkg/paths"
"github.com/docker/docker-agent/pkg/permissions"
"github.com/docker/docker-agent/pkg/profiling"
Expand Down Expand Up @@ -264,6 +263,11 @@ func (f *runExecFlags) runOrExec(ctx context.Context, out *cli.Printer, args []s
if err != nil {
return err
}
defer func() {
if err := b.Close(); err != nil {
slog.ErrorContext(ctx, "Failed to close backend", "error", err)
}
}()

loadResult, err := b.LoadTeam(ctx, b.LoadTeamRequest())
if err != nil {
Expand Down Expand Up @@ -372,7 +376,7 @@ func (f *runExecFlags) snapshotRuntimeOpts() ([]runtime.Opt, builtins.SnapshotCo
}, ctrl, nil
}

func (f *runExecFlags) createLocalRuntimeAndSession(ctx context.Context, loadResult *teamloader.LoadResult, req runtime.CreateSessionRequest) (runtime.Runtime, *session.Session, error) {
func (f *runExecFlags) createLocalRuntimeAndSession(ctx context.Context, loadResult *teamloader.LoadResult, req runtime.CreateSessionRequest, sessStore session.Store) (runtime.Runtime, *session.Session, error) {
t := loadResult.Team

// Merge user-level global permissions into the team's checker so the
Expand All @@ -387,16 +391,6 @@ func (f *runExecFlags) createLocalRuntimeAndSession(ctx context.Context, loadRes
}
agentName := agt.Name()

sessionDB, err := pathx.ExpandHomeDir(req.SessionDB)
if err != nil {
return nil, nil, err
}

sessStore, err := session.NewSQLiteSessionStore(sessionDB)
if err != nil {
return nil, nil, fmt.Errorf("creating session store: %w", err)
}

rtOpts, ctrl, err := f.snapshotRuntimeOpts()
if err != nil {
return nil, nil, err
Expand Down
53 changes: 53 additions & 0 deletions pkg/runtime/close_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
package runtime

import (
"sync/atomic"
"testing"

"github.com/stretchr/testify/assert"
"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"
)

// trackingStore wraps an in-memory store so a test can detect Close() calls.
type trackingStore struct {
session.Store

closes atomic.Int32
}

func (s *trackingStore) Close() error {
s.closes.Add(1)
return s.Store.Close()
}

// TestRuntimeClose_DoesNotCloseExternalSessionStore is a regression test for
// the bug fixed in #2872: closing one runtime would close the SQLite session
// store shared with other runtimes (e.g. spawned by the TUI's session
// browser), making subsequent reads fail with "sql: database is closed".
func TestRuntimeClose_DoesNotCloseExternalSessionStore(t *testing.T) {
store := &trackingStore{Store: session.NewInMemorySessionStore()}

prov := &mockProvider{id: "test/mock-model", stream: &mockStream{}}
root := agent.New("root", "You are a test agent", agent.WithModel(prov))
tm := team.New(team.WithAgents(root))

rt, err := New(tm,
WithSessionCompaction(false),
WithModelStore(mockModelStore{}),
WithSessionStore(store),
)
require.NoError(t, err)

require.NoError(t, rt.Close())
assert.Zero(t, store.closes.Load(),
"runtime.Close() must not close an externally supplied session store; "+
"the store may be shared with other runtimes")

// The store is still usable after the runtime is closed.
_, err = store.GetSessionSummaries(t.Context())
require.NoError(t, err)
}
18 changes: 13 additions & 5 deletions pkg/runtime/runtime.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,12 @@ type Runtime interface {
// [ErrUnsupported] for runtimes that don't expose pause control.
TogglePause(ctx context.Context) (paused bool, err error)

// Close releases resources held by the runtime (e.g., session store connections).
// Close releases resources held by the runtime (e.g., background
// agents and pending lifecycles). The session store is *not* closed
// here: it is supplied by the embedder via WithSessionStore and may
// be shared with other runtimes (e.g. the TUI's spawner reuses the
// same store across spawned sessions). Embedders that own the store
// must close it themselves.
Close() error
}

Expand Down Expand Up @@ -925,12 +930,15 @@ func (r *LocalRuntime) SessionStore() session.Store {
return r.sessionStore
}

// Close releases resources held by the runtime, including the session store.
// Close releases resources held by the runtime. The session store is
// *not* closed here: it is provided by the embedder via
// WithSessionStore and may be shared with other runtimes (e.g. the
// TUI's session spawner reuses the same store across spawned
// sessions, so closing it here would break later session lookups).
// Embedders that own the store are responsible for closing it once
// when their process is shutting down.
func (r *LocalRuntime) Close() error {
r.bgAgents.StopAll()
if r.sessionStore != nil {
return r.sessionStore.Close()
}
return nil
}

Expand Down
Loading