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

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
+}

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

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)
+}

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
 }