From 004d013792130b798a022d8399bddfed2b2f431c Mon Sep 17 00:00:00 2001 From: David Gageot Date: Fri, 22 May 2026 17:34:08 +0200 Subject: [PATCH] fix: don't close shared session store in runtime.Close (#2872) --- cmd/root/backend.go | 56 ++++++++++++++++++++++++++++++++++++++- cmd/root/run.go | 18 +++++-------- pkg/runtime/close_test.go | 53 ++++++++++++++++++++++++++++++++++++ pkg/runtime/runtime.go | 18 +++++++++---- 4 files changed, 127 insertions(+), 18 deletions(-) create mode 100644 pkg/runtime/close_test.go diff --git a/cmd/root/backend.go b/cmd/root/backend.go index 181cde40d..8785e63a6 100644 --- a/cmd/root/backend.go +++ b/cmd/root/backend.go @@ -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" @@ -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. @@ -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 { @@ -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() { + 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 @@ -83,6 +126,13 @@ func (b *localBackend) Spawner(rt runtime.Runtime) tui.SessionSpawner { return b.flags.createSessionSpawner(b.agentSource, rt.SessionStore()) } +func (b *localBackend) Close() error { + if b.store == nil { + return nil + } + return b.store.Close() +} + // remoteBackend talks to a docker-agent server. type remoteBackend struct { flags *runExecFlags @@ -142,3 +192,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 +} diff --git a/cmd/root/run.go b/cmd/root/run.go index 74f5ff784..87f2ca0ab 100644 --- a/cmd/root/run.go +++ b/cmd/root/run.go @@ -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" @@ -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 { @@ -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 @@ -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 diff --git a/pkg/runtime/close_test.go b/pkg/runtime/close_test.go new file mode 100644 index 000000000..f6bd631e9 --- /dev/null +++ b/pkg/runtime/close_test.go @@ -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) +} diff --git a/pkg/runtime/runtime.go b/pkg/runtime/runtime.go index 3ea8441fa..1339d7728 100644 --- a/pkg/runtime/runtime.go +++ b/pkg/runtime/runtime.go @@ -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 } @@ -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 }