refactor: extract SnapshotController so the runtime no longer brokers /undo

cmd/root/run.go

@@ -19,6 +19,8 @@ import (
 	"github.com/docker/docker-agent/pkg/app"
 	"github.com/docker/docker-agent/pkg/cli"
 	"github.com/docker/docker-agent/pkg/config"
+	"github.com/docker/docker-agent/pkg/hooks"
+	"github.com/docker/docker-agent/pkg/hooks/builtins"
 	"github.com/docker/docker-agent/pkg/paths"
 	"github.com/docker/docker-agent/pkg/permissions"
 	"github.com/docker/docker-agent/pkg/profiling"
@@ -66,6 +68,16 @@ type runExecFlags struct {
 	// from user config settings. Nil when no global permissions are configured.
 	globalPermissions *permissions.Checker
 	snapshotsEnabled  bool
+
+	// snapshotController is the [builtins.SnapshotController] for the
+	// initial App: it is wired into the initial runtime as an
+	// auto-injector and into the App via app.WithSnapshotController so
+	// /undo, /snapshots, /reset drive the same instance that captures
+	// the checkpoints. Sub-runtimes created by [createSessionSpawner]
+	// build their own controller (and registry) so each spawned
+	// session has independent snapshot state; that controller is local
+	// to the spawner closure and never reaches this field.
+	snapshotController builtins.SnapshotController
 }
 
 func newRunCmd() *cobra.Command {
@@ -329,6 +341,32 @@ func (f *runExecFlags) createRemoteRuntimeAndSession(ctx context.Context, origin
 	return remoteRt, sess, nil
 }
 
+// snapshotRuntimeOpts wires the snapshot builtin into a runtime.
+// Returns the [runtime.Opt]s that hand the registry and the
+// [builtins.SnapshotController] auto-injector to the runtime, plus
+// the controller itself for the embedder to pass to the App via
+// [app.WithSnapshotController]. When snapshots aren't enabled,
+// returns no opts and a nil controller so callers don't have to
+// branch on f.snapshotsEnabled themselves.
+//
+// A fresh registry is created here rather than reused across runtimes
+// so the spawner-created sub-runtimes get their own snapshot state
+// (each spawned session has independent /undo history).
+func (f *runExecFlags) snapshotRuntimeOpts() ([]runtime.Opt, builtins.SnapshotController, error) {
+	if !f.snapshotsEnabled {
+		return nil, nil, nil
+	}
+	reg := hooks.NewRegistry()
+	ctrl, err := builtins.RegisterSnapshot(reg, true)
+	if err != nil {
+		return nil, nil, fmt.Errorf("register snapshot builtin: %w", err)
+	}
+	return []runtime.Opt{
+		runtime.WithHooksRegistry(reg),
+		runtime.WithAutoInjector(ctrl),
+	}, ctrl, nil
+}
+
 func (f *runExecFlags) createLocalRuntimeAndSession(ctx context.Context, loadResult *teamloader.LoadResult) (runtime.Runtime, *session.Session, error) {
 	t := loadResult.Team
 
@@ -364,19 +402,22 @@ func (f *runExecFlags) createLocalRuntimeAndSession(ctx context.Context, loadRes
 		AgentDefaultModels: loadResult.AgentDefaultModels,
 	}
 
+	rtOpts, ctrl, err := f.snapshotRuntimeOpts()
+	if err != nil {
+		return nil, nil, err
+	}
 	runtimeOpts := []runtime.Opt{
 		runtime.WithSessionStore(sessStore),
 		runtime.WithCurrentAgent(agentName),
 		runtime.WithTracer(otel.Tracer(AppName)),
 		runtime.WithModelSwitcherConfig(modelSwitcherCfg),
 	}
-	if f.snapshotsEnabled {
-		runtimeOpts = append(runtimeOpts, runtime.WithSnapshots(true))
-	}
+	runtimeOpts = append(runtimeOpts, rtOpts...)
 	localRt, err := runtime.New(t, runtimeOpts...)
 	if err != nil {
 		return nil, nil, fmt.Errorf("creating runtime: %w", err)
 	}
+	f.snapshotController = ctrl
 
 	var sess *session.Session
 	if f.sessionID != "" {
@@ -509,6 +550,9 @@ func (f *runExecFlags) buildAppOpts(args []string) ([]app.Opt, error) {
 	if f.exitAfterResponse {
 		opts = append(opts, app.WithExitAfterFirstResponse())
 	}
+	if f.snapshotController != nil {
+		opts = append(opts, app.WithSnapshotController(f.snapshotController))
+	}
 	return opts, nil
 }
 
@@ -560,15 +604,17 @@ func (f *runExecFlags) createSessionSpawner(agentSource config.Source, sessStore
 		}
 
 		// Create the local runtime
+		rtOpts, ctrl, err := f.snapshotRuntimeOpts()
+		if err != nil {
+			return nil, nil, nil, err
+		}
 		runtimeOpts := []runtime.Opt{
 			runtime.WithSessionStore(sessStore),
 			runtime.WithCurrentAgent(agt.Name()),
 			runtime.WithTracer(otel.Tracer(AppName)),
 			runtime.WithModelSwitcherConfig(modelSwitcherCfg),
 		}
-		if f.snapshotsEnabled {
-			runtimeOpts = append(runtimeOpts, runtime.WithSnapshots(true))
-		}
+		runtimeOpts = append(runtimeOpts, rtOpts...)
 		localRt, err := runtime.New(t, runtimeOpts...)
 		if err != nil {
 			return nil, nil, nil, err
@@ -587,6 +633,9 @@ func (f *runExecFlags) createSessionSpawner(agentSource config.Source, sessStore
 		if gen := localRt.TitleGenerator(); gen != nil {
 			appOpts = append(appOpts, app.WithTitleGenerator(gen))
 		}
+		if ctrl != nil {
+			appOpts = append(appOpts, app.WithSnapshotController(ctrl))
+		}
 
 		a := app.New(spawnCtx, localRt, newSess, appOpts...)
 

lint/hook_builtins_registered.go

@@ -8,7 +8,10 @@ import (
 
 // HookBuiltinsRegistered enforces that every builtin-name constant declared
 // under pkg/hooks/builtins/ is wired into the package's Register() function
-// in pkg/hooks/builtins/builtins.go.
+// in pkg/hooks/builtins/builtins.go — with one exception: the snapshot
+// builtin ships its own entry point ([builtins.RegisterSnapshot]) because
+// it returns a [SnapshotController] for embedders, so its declaration in
+// snapshot.go is intentionally not registered through Register().
 //
 // Each in-process builtin lives in its own file with a name constant and an
 // implementation:
@@ -71,12 +74,12 @@ var HookBuiltinsRegistered = &cop.Func{
 }
 
 // exportedBuiltinNames returns the identifiers of every exported `const Name = "..."`
-// declaration in pkg/hooks/builtins/, excluding builtins.go itself and any
-// test files (which is not where new builtins land).
+// declaration in pkg/hooks/builtins/, excluding builtins.go itself, snapshot.go
+// (which has its own RegisterSnapshot entry point), and any test files.
 func exportedBuiltinNames(p *cop.Pass) ([]string, error) {
 	files, err := p.ParseDir(".", cop.ParseDirOptions{
 		SkipTests: true,
-		SkipFiles: []string{"builtins.go"},
+		SkipFiles: []string{"builtins.go", "snapshot.go"},
 	})
 	if err != nil {
 		return nil, err

pkg/app/app.go

@@ -21,6 +21,7 @@ import (
 	"github.com/docker/docker-agent/pkg/chat"
 	"github.com/docker/docker-agent/pkg/cli"
 	"github.com/docker/docker-agent/pkg/config/types"
+	"github.com/docker/docker-agent/pkg/hooks/builtins"
 	"github.com/docker/docker-agent/pkg/runtime"
 	"github.com/docker/docker-agent/pkg/session"
 	"github.com/docker/docker-agent/pkg/sessiontitle"
@@ -40,10 +41,11 @@ type App struct {
 	events                 chan tea.Msg
 	throttleDuration       time.Duration
 	cancel                 context.CancelFunc
-	currentAgentModel      string                  // Tracks the current agent's model ID from AgentInfoEvent
-	exitAfterFirstResponse bool                    // Exit TUI after first assistant response completes
-	titleGenerating        atomic.Bool             // True when title generation is in progress
-	titleGen               *sessiontitle.Generator // Title generator for local runtime (nil for remote)
+	currentAgentModel      string                      // Tracks the current agent's model ID from AgentInfoEvent
+	exitAfterFirstResponse bool                        // Exit TUI after first assistant response completes
+	titleGenerating        atomic.Bool                 // True when title generation is in progress
+	titleGen               *sessiontitle.Generator     // Title generator for local runtime (nil for remote)
+	snapshotController     builtins.SnapshotController // Drives /undo, /snapshots, /reset; nil for runtimes that don't capture snapshots
 }
 
 // Opt is an option for creating a new App.
@@ -87,6 +89,19 @@ func WithTitleGenerator(gen *sessiontitle.Generator) Opt {
 	}
 }
 
+// WithSnapshotController plumbs in the [builtins.SnapshotController]
+// the App uses to drive /undo, /snapshots, /reset. Pass the same
+// controller to the runtime via runtime.WithAutoInjector so the
+// instance that captures the checkpoints is the one the TUI commands
+// drive. Pass nil (or omit the option) for runtimes that don't capture
+// snapshots; the App then reports SnapshotsEnabled()==false and the
+// related commands silently no-op.
+func WithSnapshotController(c builtins.SnapshotController) Opt {
+	return func(a *App) {
+		a.snapshotController = c
+	}
+}
+
 func New(ctx context.Context, rt runtime.Runtime, sess *session.Session, opts ...Opt) *App {
 	app := &App{
 		runtime:          rt,

pkg/app/app_test.go

@@ -8,6 +8,7 @@ import (
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"github.com/docker/docker-agent/pkg/hooks"
 	"github.com/docker/docker-agent/pkg/hooks/builtins"
 	"github.com/docker/docker-agent/pkg/runtime"
 	"github.com/docker/docker-agent/pkg/session"
@@ -17,12 +18,13 @@ import (
 	mcptools "github.com/docker/docker-agent/pkg/tools/mcp"
 )
 
-// mockRuntime is a minimal mock for testing App without a real runtime
+// mockRuntime is a minimal mock for testing App without a real runtime.
+// Snapshot operations are NOT modeled here: they are driven through a
+// [builtins.SnapshotController] passed to the App via WithSnapshotController,
+// so the mock runtime stays small and focused on the runtime.Runtime
+// surface.
 type mockRuntime struct {
-	store     session.Store
-	undoFiles int
-	undoOK    bool
-	undoErr   error
+	store session.Store
 }
 
 func (m *mockRuntime) CurrentAgentInfo(ctx context.Context) runtime.CurrentAgentInfo {
@@ -78,18 +80,35 @@ func (m *mockRuntime) Close() error                            { return nil }
 func (m *mockRuntime) Stop()                                   {}
 func (m *mockRuntime) Steer(_ runtime.QueuedMessage) error     { return nil }
 func (m *mockRuntime) FollowUp(_ runtime.QueuedMessage) error  { return nil }
-func (m *mockRuntime) UndoLastSnapshot(context.Context, *session.Session) (int, bool, error) {
-	return m.undoFiles, m.undoOK, m.undoErr
-}
-func (m *mockRuntime) SnapshotsEnabled() bool                                 { return true }
-func (m *mockRuntime) ListSnapshots(*session.Session) []builtins.SnapshotInfo { return nil }
-func (m *mockRuntime) ResetSnapshot(context.Context, *session.Session, int) (int, bool, error) {
-	return m.undoFiles, m.undoOK, m.undoErr
-}
 
 // Verify mockRuntime implements runtime.Runtime
 var _ runtime.Runtime = (*mockRuntime)(nil)
 
+// stubSnapshotController is a tiny SnapshotController used by the app
+// tests to drive /undo without spinning up a real shadow-git
+// repository. enabled gates SnapshotsEnabled(), and the (files, ok,
+// err) tuple is returned verbatim from UndoLast / Reset so each test
+// can assert the result-shaping logic in [snapshotResult].
+type stubSnapshotController struct {
+	enabled bool
+	files   int
+	ok      bool
+	err     error
+}
+
+func (s *stubSnapshotController) Enabled() bool { return s.enabled }
+func (s *stubSnapshotController) UndoLast(context.Context, string, string) (int, bool, error) {
+	return s.files, s.ok, s.err
+}
+
+func (s *stubSnapshotController) List(string) []builtins.SnapshotInfo { return nil }
+func (s *stubSnapshotController) Reset(context.Context, string, string, int) (int, bool, error) {
+	return s.files, s.ok, s.err
+}
+func (s *stubSnapshotController) AutoInject(*hooks.Config) {}
+
+var _ builtins.SnapshotController = (*stubSnapshotController)(nil)
+
 func TestApp_NewSession_PreservesToolsApproved(t *testing.T) {
 	t.Parallel()
 
@@ -245,7 +264,9 @@ func TestApp_UndoLastSnapshot(t *testing.T) {
 	t.Parallel()
 
 	ctx := t.Context()
-	app := New(ctx, &mockRuntime{undoFiles: 2, undoOK: true}, session.New())
+	app := New(ctx, &mockRuntime{}, session.New(),
+		WithSnapshotController(&stubSnapshotController{enabled: true, files: 2, ok: true}),
+	)
 	result, err := app.UndoLastSnapshot(ctx)
 	require.NoError(t, err)
 	assert.Equal(t, 2, result.RestoredFiles)
@@ -255,17 +276,36 @@ func TestApp_UndoLastSnapshot_NoSnapshot(t *testing.T) {
 	t.Parallel()
 
 	ctx := t.Context()
-	app := New(ctx, &mockRuntime{}, session.New())
+	app := New(ctx, &mockRuntime{}, session.New(),
+		WithSnapshotController(&stubSnapshotController{enabled: true}),
+	)
 	_, err := app.UndoLastSnapshot(ctx)
 	assert.ErrorIs(t, err, ErrNothingToUndo)
 }
 
+func TestApp_UndoLastSnapshot_NoController(t *testing.T) {
+	t.Parallel()
+
+	// Without a SnapshotController the App reports nothing to undo,
+	// so the same UI affordance can light up regardless of which
+	// runtime the embedder paired the App with.
+	ctx := t.Context()
+	app := New(ctx, &mockRuntime{}, session.New())
+	_, err := app.UndoLastSnapshot(ctx)
+	require.ErrorIs(t, err, ErrNothingToUndo)
+	assert.False(t, app.SnapshotsEnabled())
+}
+
 func TestApp_SnapshotsEnabled_DoesNotRequireSession(t *testing.T) {
 	t.Parallel()
 
-	// SnapshotsEnabled answers a runtime-capability question; it must not
-	// silently return false just because no session is attached.
-	app := &App{runtime: &mockRuntime{}, session: nil}
+	// SnapshotsEnabled answers a controller-capability question; it
+	// must not silently return false just because no session is attached.
+	app := &App{
+		runtime:            &mockRuntime{},
+		session:            nil,
+		snapshotController: &stubSnapshotController{enabled: true},
+	}
 	assert.True(t, app.SnapshotsEnabled())
 }