monitor: Move commands to a separated package

This commit moves monitor commands to `monior/commands` package.
Commands still need access to the `monitor` object and buildx controller so this
commit enables this via `Monitor` interface stored in `monitor/types`.

Signed-off-by: Kohei Tokunaga <ktokunaga.mail@gmail.com>

monitor/commands/attach.go

@@ -0,0 +1,71 @@
+package commands
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	"github.com/docker/buildx/monitor/types"
+	"github.com/pkg/errors"
+)
+
+type AttachCmd struct {
+	m types.Monitor
+
+	stdout io.WriteCloser
+}
+
+func NewAttachCmd(m types.Monitor, stdout io.WriteCloser) types.Command {
+	return &AttachCmd{m, stdout}
+}
+
+func (cm *AttachCmd) Info() types.CommandInfo {
+	return types.CommandInfo{HelpMessage: "attach to a buildx server or a process in the container"}
+}
+
+func (cm *AttachCmd) Exec(ctx context.Context, args []string) error {
+	if len(args) < 2 {
+		return errors.Errorf("server name must be passed")
+	}
+	ref := args[1]
+	var id string
+
+	isProcess, err := isProcessID(ctx, cm.m, ref)
+	if err == nil && isProcess {
+		cm.m.Attach(ctx, ref)
+		id = ref
+	}
+	if id == "" {
+		refs, err := cm.m.List(ctx)
+		if err != nil {
+			return errors.Errorf("failed to get the list of sessions: %v", err)
+		}
+		found := false
+		for _, s := range refs {
+			if s == ref {
+				found = true
+				break
+			}
+		}
+		if !found {
+			return errors.Errorf("unknown ID: %q", ref)
+		}
+		cm.m.Detach() // Finish existing attach
+		cm.m.AttachSession(ref)
+	}
+	fmt.Fprintf(cm.stdout, "Attached to process %q. Press Ctrl-a-c to switch to the new container\n", id)
+	return nil
+}
+
+func isProcessID(ctx context.Context, c types.Monitor, ref string) (bool, error) {
+	infos, err := c.ListProcesses(ctx)
+	if err != nil {
+		return false, err
+	}
+	for _, p := range infos {
+		if p.ProcessID == ref {
+			return true, nil
+		}
+	}
+	return false, nil
+}

monitor/commands/disconnect.go

@@ -0,0 +1,38 @@
+package commands
+
+import (
+	"context"
+
+	"github.com/docker/buildx/monitor/types"
+	"github.com/pkg/errors"
+)
+
+type DisconnectCmd struct {
+	m types.Monitor
+}
+
+func NewDisconnectCmd(m types.Monitor) types.Command {
+	return &DisconnectCmd{m}
+}
+
+func (cm *DisconnectCmd) Info() types.CommandInfo {
+	return types.CommandInfo{HelpMessage: "disconnect a client from a buildx server. Specific session ID can be specified an arg"}
+}
+
+func (cm *DisconnectCmd) Exec(ctx context.Context, args []string) error {
+	target := cm.m.AttachedSessionID()
+	if len(args) >= 2 {
+		target = args[1]
+	}
+	isProcess, err := isProcessID(ctx, cm.m, target)
+	if err == nil && isProcess {
+		if err := cm.m.DisconnectProcess(ctx, target); err != nil {
+			return errors.Errorf("disconnecting from process failed %v", target)
+		}
+		return nil
+	}
+	if err := cm.m.DisconnectSession(ctx, target); err != nil {
+		return errors.Errorf("disconnecting from session failed: %v", err)
+	}
+	return nil
+}

monitor/commands/exec.go

@@ -0,0 +1,44 @@
+package commands
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	controllerapi "github.com/docker/buildx/controller/pb"
+	"github.com/docker/buildx/monitor/types"
+	"github.com/pkg/errors"
+)
+
+type ExecCmd struct {
+	m types.Monitor
+
+	invokeConfig controllerapi.InvokeConfig
+	stdout       io.WriteCloser
+}
+
+func NewExecCmd(m types.Monitor, invokeConfig controllerapi.InvokeConfig, stdout io.WriteCloser) types.Command {
+	return &ExecCmd{m, invokeConfig, stdout}
+}
+
+func (cm *ExecCmd) Info() types.CommandInfo {
+	return types.CommandInfo{HelpMessage: "execute a process in the interactive container"}
+}
+
+func (cm *ExecCmd) Exec(ctx context.Context, args []string) error {
+	if len(args) < 2 {
+		return errors.Errorf("command must be passed")
+	}
+	cfg := controllerapi.InvokeConfig{
+		Entrypoint: []string{args[1]},
+		Cmd:        args[2:],
+		// TODO: support other options as well via flags
+		Env:  cm.invokeConfig.Env,
+		User: cm.invokeConfig.User,
+		Cwd:  cm.invokeConfig.Cwd,
+		Tty:  true,
+	}
+	pid := cm.m.Exec(ctx, cfg)
+	fmt.Fprintf(cm.stdout, "Process %q started. Press Ctrl-a-c to switch to that process.\n", pid)
+	return nil
+}

monitor/commands/kill.go

@@ -0,0 +1,27 @@
+package commands
+
+import (
+	"context"
+
+	"github.com/docker/buildx/monitor/types"
+	"github.com/pkg/errors"
+)
+
+type KillCmd struct {
+	m types.Monitor
+}
+
+func NewKillCmd(m types.Monitor) types.Command {
+	return &KillCmd{m}
+}
+
+func (cm *KillCmd) Info() types.CommandInfo {
+	return types.CommandInfo{HelpMessage: "kill buildx server"}
+}
+
+func (cm *KillCmd) Exec(ctx context.Context, args []string) error {
+	if err := cm.m.Kill(ctx); err != nil {
+		return errors.Errorf("failed to kill: %v", err)
+	}
+	return nil
+}

monitor/commands/list.go

@@ -0,0 +1,40 @@
+package commands
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"sort"
+	"text/tabwriter"
+
+	"github.com/docker/buildx/monitor/types"
+)
+
+type ListCmd struct {
+	m types.Monitor
+
+	stdout io.WriteCloser
+}
+
+func NewListCmd(m types.Monitor, stdout io.WriteCloser) types.Command {
+	return &ListCmd{m, stdout}
+}
+
+func (cm *ListCmd) Info() types.CommandInfo {
+	return types.CommandInfo{HelpMessage: "list buildx sessions"}
+}
+
+func (cm *ListCmd) Exec(ctx context.Context, args []string) error {
+	refs, err := cm.m.List(ctx)
+	if err != nil {
+		return err
+	}
+	sort.Strings(refs)
+	tw := tabwriter.NewWriter(cm.stdout, 1, 8, 1, '\t', 0)
+	fmt.Fprintln(tw, "ID\tCURRENT_SESSION")
+	for _, k := range refs {
+		fmt.Fprintf(tw, "%-20s\t%v\n", k, k == cm.m.AttachedSessionID())
+	}
+	tw.Flush()
+	return nil
+}

monitor/commands/ps.go

@@ -0,0 +1,37 @@
+package commands
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"text/tabwriter"
+
+	"github.com/docker/buildx/monitor/types"
+)
+
+type PsCmd struct {
+	m      types.Monitor
+	stdout io.WriteCloser
+}
+
+func NewPsCmd(m types.Monitor, stdout io.WriteCloser) types.Command {
+	return &PsCmd{m, stdout}
+}
+
+func (cm *PsCmd) Info() types.CommandInfo {
+	return types.CommandInfo{HelpMessage: `list processes invoked by "exec". Use "attach" to attach IO to that process`}
+}
+
+func (cm *PsCmd) Exec(ctx context.Context, args []string) error {
+	plist, err := cm.m.ListProcesses(ctx)
+	if err != nil {
+		return err
+	}
+	tw := tabwriter.NewWriter(cm.stdout, 1, 8, 1, '\t', 0)
+	fmt.Fprintln(tw, "PID\tCURRENT_SESSION\tCOMMAND")
+	for _, p := range plist {
+		fmt.Fprintf(tw, "%-20s\t%v\t%v\n", p.ProcessID, p.ProcessID == cm.m.AttachedPID(), append(p.InvokeConfig.Entrypoint, p.InvokeConfig.Cmd...))
+	}
+	tw.Flush()
+	return nil
+}

monitor/commands/reload.go

@@ -0,0 +1,76 @@
+package commands
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	controllererrors "github.com/docker/buildx/controller/errdefs"
+	controllerapi "github.com/docker/buildx/controller/pb"
+	"github.com/docker/buildx/monitor/types"
+	"github.com/docker/buildx/util/progress"
+	"github.com/pkg/errors"
+)
+
+type ReloadCmd struct {
+	m types.Monitor
+
+	stdout   io.WriteCloser
+	progress *progress.Printer
+
+	options      *controllerapi.BuildOptions
+	invokeConfig controllerapi.InvokeConfig
+}
+
+func NewReloadCmd(m types.Monitor, stdout io.WriteCloser, progress *progress.Printer, options *controllerapi.BuildOptions, invokeConfig controllerapi.InvokeConfig) types.Command {
+	return &ReloadCmd{m, stdout, progress, options, invokeConfig}
+}
+
+func (cm *ReloadCmd) Info() types.CommandInfo {
+	return types.CommandInfo{HelpMessage: "reloads the context and build it"}
+}
+
+func (cm *ReloadCmd) Exec(ctx context.Context, args []string) error {
+	var bo *controllerapi.BuildOptions
+	if cm.m.AttachedSessionID() != "" {
+		// Rebuilding an existing session; Restore the build option used for building this session.
+		res, err := cm.m.Inspect(ctx)
+		if err != nil {
+			fmt.Printf("failed to inspect the current build session: %v\n", err)
+		} else {
+			bo = res.Options
+		}
+	} else {
+		bo = cm.options
+	}
+	if bo == nil {
+		return errors.Errorf("no build option is provided")
+	}
+	if cm.m.AttachedSessionID() != "" {
+		if err := cm.m.Disconnect(ctx); err != nil {
+			fmt.Println("disconnect error", err)
+		}
+	}
+	var resultUpdated bool
+	cm.progress.Unpause()
+	ref, _, err := cm.m.Build(ctx, *bo, nil, cm.progress) // TODO: support stdin, hold build ref
+	cm.progress.Pause()
+	if err != nil {
+		var be *controllererrors.BuildError
+		if errors.As(err, &be) {
+			ref = be.Ref
+			resultUpdated = true
+		} else {
+			fmt.Printf("failed to reload: %v\n", err)
+		}
+	} else {
+		resultUpdated = true
+	}
+	cm.m.AttachSession(ref)
+	if resultUpdated {
+		// rollback the running container with the new result
+		id := cm.m.Rollback(ctx, cm.invokeConfig)
+		fmt.Fprintf(cm.stdout, "Interactive container was restarted with process %q. Press Ctrl-a-c to switch to the new container\n", id)
+	}
+	return nil
+}

monitor/commands/rollback.go

@@ -0,0 +1,43 @@
+package commands
+
+import (
+	"context"
+	"fmt"
+	"io"
+
+	controllerapi "github.com/docker/buildx/controller/pb"
+	"github.com/docker/buildx/monitor/types"
+)
+
+type RollbackCmd struct {
+	m types.Monitor
+
+	invokeConfig controllerapi.InvokeConfig
+	stdout       io.WriteCloser
+}
+
+func NewRollbackCmd(m types.Monitor, invokeConfig controllerapi.InvokeConfig, stdout io.WriteCloser) types.Command {
+	return &RollbackCmd{m, invokeConfig, stdout}
+}
+
+func (cm *RollbackCmd) Info() types.CommandInfo {
+	return types.CommandInfo{HelpMessage: "re-runs the interactive container with initial rootfs contents"}
+}
+
+func (cm *RollbackCmd) Exec(ctx context.Context, args []string) error {
+	cfg := cm.invokeConfig
+	if len(args) >= 2 {
+		cmds := args[1:]
+		if cmds[0] == "--init" {
+			cfg.Initial = true
+			cmds = cmds[1:]
+		}
+		if len(cmds) > 0 {
+			cfg.Entrypoint = []string{cmds[0]}
+			cfg.Cmd = cmds[1:]
+		}
+	}
+	id := cm.m.Rollback(ctx, cfg)
+	fmt.Fprintf(cm.stdout, "Interactive container was restarted with process %q. Press Ctrl-a-c to switch to the new container\n", id)
+	return nil
+}

monitor/monitor.go

@@ -12,30 +12,18 @@ import (
 
 	"github.com/containerd/console"
 	"github.com/docker/buildx/controller/control"
-	controllererrors "github.com/docker/buildx/controller/errdefs"
 	controllerapi "github.com/docker/buildx/controller/pb"
+	"github.com/docker/buildx/monitor/commands"
+	"github.com/docker/buildx/monitor/types"
 	"github.com/docker/buildx/util/ioset"
 	"github.com/docker/buildx/util/progress"
+	"github.com/moby/buildkit/client"
 	"github.com/moby/buildkit/identity"
 	"github.com/pkg/errors"
 	"github.com/sirupsen/logrus"
 	"golang.org/x/term"
 )
 
-const helpMessage = `
-Available commands are:
-  reload     reloads the context and build it.
-  rollback   re-runs the interactive container with initial rootfs contents.
-  list       list buildx sessions.
-  attach     attach to a buildx server or a process in the container.
-  exec       execute a process in the interactive container.
-  ps         list processes invoked by "exec". Use "attach" to attach IO to that process.
-  disconnect disconnect a client from a buildx server. Specific session ID can be specified an arg.
-  kill       kill buildx server.
-  exit       exits monitor.
-  help       shows this message.
-`
-
 // RunMonitor provides an interactive session for running and managing containers via specified IO.
 func RunMonitor(ctx context.Context, curRef string, options *controllerapi.BuildOptions, invokeConfig controllerapi.InvokeConfig, c control.BuildxController, stdin io.ReadCloser, stdout io.WriteCloser, stderr console.File, progress *progress.Printer) error {
 	defer func() {
@@ -75,7 +63,8 @@ func RunMonitor(ctx context.Context, curRef string, options *controllerapi.Build
 	invokeForwarder := ioset.NewForwarder()
 	invokeForwarder.SetIn(&containerIn)
 	m := &monitor{
-		invokeIO: invokeForwarder,
+		controllerRef: newControllerRef(c, curRef),
+		invokeIO:      invokeForwarder,
 		muxIO: ioset.NewMuxIO(ioset.In{
 			Stdin:  io.NopCloser(stdin),
 			Stdout: nopCloser{stdout},
@@ -87,16 +76,30 @@ func RunMonitor(ctx context.Context, curRef string, options *controllerapi.Build
 			}
 			return "Switched IO\n"
 		}),
-		invokeFunc: c.Invoke,
 	}
 
 	// Start container automatically
 	fmt.Fprintf(stdout, "Launching interactive container. Press Ctrl-a-c to switch to monitor console\n")
 	invokeConfig.Rollback = false
 	invokeConfig.Initial = false
-	id := m.rollback(ctx, curRef, invokeConfig)
+	id := m.Rollback(ctx, invokeConfig)
 	fmt.Fprintf(stdout, "Interactive container was restarted with process %q. Press Ctrl-a-c to switch to the new container\n", id)
 
+	registeredCommands := map[string]types.Command{
+		"reload":     commands.NewReloadCmd(m, stdout, progress, options, invokeConfig),
+		"rollback":   commands.NewRollbackCmd(m, invokeConfig, stdout),
+		"list":       commands.NewListCmd(m, stdout),
+		"disconnect": commands.NewDisconnectCmd(m),
+		"kill":       commands.NewKillCmd(m),
+		"attach":     commands.NewAttachCmd(m, stdout),
+		"exec":       commands.NewExecCmd(m, invokeConfig, stdout),
+		"ps":         commands.NewPsCmd(m, stdout),
+	}
+	additionalHelpMessages := map[string]string{
+		"help": "shows this message",
+		"exit": "exits monitor",
+	}
+
 	// Serve monitor commands
 	monitorForwarder := ioset.NewForwarder()
 	monitorForwarder.SetIn(&monitorIn)
@@ -126,174 +129,29 @@ func RunMonitor(ctx context.Context, curRef string, options *controllerapi.Build
 				if len(args) == 0 {
 					continue
 				}
+
+				// Builtin commands
 				switch args[0] {
 				case "":
 					// nop
-				case "reload":
-					var bo *controllerapi.BuildOptions
-					if curRef != "" {
-						// Rebuilding an existing session; Restore the build option used for building this session.
-						res, err := c.Inspect(ctx, curRef)
-						if err != nil {
-							fmt.Printf("failed to inspect the current build session: %v\n", err)
-						} else {
-							bo = res.Options
-						}
-					} else {
-						bo = options
-					}
-					if bo == nil {
-						fmt.Println("reload: no build option is provided")
-						continue
-					}
-					if curRef != "" {
-						if err := c.Disconnect(ctx, curRef); err != nil {
-							fmt.Println("disconnect error", err)
-						}
-					}
-					var resultUpdated bool
-					progress.Unpause()
-					ref, _, err := c.Build(ctx, *bo, nil, progress) // TODO: support stdin, hold build ref
-					progress.Pause()
-					if err != nil {
-						var be *controllererrors.BuildError
-						if errors.As(err, &be) {
-							curRef = be.Ref
-							resultUpdated = true
-						} else {
-							fmt.Printf("failed to reload: %v\n", err)
-						}
-					} else {
-						curRef = ref
-						resultUpdated = true
-					}
-					if resultUpdated {
-						// rollback the running container with the new result
-						id := m.rollback(ctx, curRef, invokeConfig)
-						fmt.Fprintf(stdout, "Interactive container was restarted with process %q. Press Ctrl-a-c to switch to the new container\n", id)
-					}
-				case "rollback":
-					cfg := invokeConfig
-					if len(args) >= 2 {
-						cmds := args[1:]
-						if cmds[0] == "--init" {
-							cfg.Initial = true
-							cmds = cmds[1:]
-						}
-						if len(cmds) > 0 {
-							cfg.Entrypoint = []string{cmds[0]}
-							cfg.Cmd = cmds[1:]
-						}
-					}
-					id := m.rollback(ctx, curRef, cfg)
-					fmt.Fprintf(stdout, "Interactive container was restarted with process %q. Press Ctrl-a-c to switch to the new container\n", id)
-				case "list":
-					refs, err := c.List(ctx)
-					if err != nil {
-						fmt.Printf("failed to list: %v\n", err)
-					}
-					sort.Strings(refs)
-					tw := tabwriter.NewWriter(stdout, 1, 8, 1, '\t', 0)
-					fmt.Fprintln(tw, "ID\tCURRENT_SESSION")
-					for _, k := range refs {
-						fmt.Fprintf(tw, "%-20s\t%v\n", k, k == curRef)
-					}
-					tw.Flush()
-				case "disconnect":
-					target := curRef
-					if len(args) >= 2 {
-						target = args[1]
-					}
-					isProcess, err := isProcessID(ctx, c, curRef, target)
-					if err == nil && isProcess {
-						if err := c.DisconnectProcess(ctx, curRef, target); err != nil {
-							fmt.Printf("disconnect process failed %v\n", target)
-							continue
-						}
-						continue
-					}
-					if err := c.Disconnect(ctx, target); err != nil {
-						fmt.Println("disconnect error", err)
-					}
-				case "kill":
-					if err := c.Kill(ctx); err != nil {
-						fmt.Printf("failed to kill: %v\n", err)
-					}
-				case "attach":
-					if len(args) < 2 {
-						fmt.Println("attach: server name must be passed")
-						continue
-					}
-					ref := args[1]
-					var id string
-
-					isProcess, err := isProcessID(ctx, c, curRef, ref)
-					if err == nil && isProcess {
-						m.attach(ctx, curRef, ref)
-						id = ref
-					}
-					if id == "" {
-						refs, err := c.List(ctx)
-						if err != nil {
-							fmt.Printf("failed to get the list of sessions: %v\n", err)
-							continue
-						}
-						found := false
-						for _, s := range refs {
-							if s == ref {
-								found = true
-								break
-							}
-						}
-						if !found {
-							fmt.Printf("unknown ID: %q\n", ref)
-							continue
-						}
-						if m.invokeCancel != nil {
-							m.invokeCancel() // Finish existing attach
-						}
-						curRef = ref
-					}
-					fmt.Fprintf(stdout, "Attached to process %q. Press Ctrl-a-c to switch to the new container\n", id)
-				case "exec":
-					if len(args) < 2 {
-						fmt.Println("attach: command must be passed")
-						continue
-					}
-					if curRef == "" {
-						fmt.Println("attach to a session first")
-						continue
-					}
-					cfg := controllerapi.InvokeConfig{
-						Entrypoint: []string{args[1]},
-						Cmd:        args[2:],
-						// TODO: support other options as well via flags
-						Env:  invokeConfig.Env,
-						User: invokeConfig.User,
-						Cwd:  invokeConfig.Cwd,
-						Tty:  true,
-					}
-					pid := m.exec(ctx, curRef, cfg)
-					fmt.Fprintf(stdout, "Process %q started. Press Ctrl-a-c to switch to that process.\n", pid)
-				case "ps":
-					plist, err := c.ListProcesses(ctx, curRef)
-					if err != nil {
-						fmt.Println("cannot list process:", err)
-						continue
-					}
-					tw := tabwriter.NewWriter(stdout, 1, 8, 1, '\t', 0)
-					fmt.Fprintln(tw, "PID\tCURRENT_SESSION\tCOMMAND")
-					for _, p := range plist {
-						fmt.Fprintf(tw, "%-20s\t%v\t%v\n", p.ProcessID, p.ProcessID == m.attachedPid.Load(), append(p.InvokeConfig.Entrypoint, p.InvokeConfig.Cmd...))
-					}
-					tw.Flush()
+					continue
 				case "exit":
 					return
 				case "help":
-					fmt.Fprint(stdout, helpMessage)
+					printHelpMessage(stdout, registeredCommands, additionalHelpMessages)
+					continue
 				default:
-					fmt.Printf("unknown command: %q\n", l)
-					fmt.Fprint(stdout, helpMessage)
+				}
+
+				// Registered commands
+				cmdname := args[0]
+				if cm, ok := registeredCommands[cmdname]; ok {
+					if err := cm.Exec(ctx, args); err != nil {
+						fmt.Fprintf(stdout, "%s: %v\n", cmdname, err)
+					}
+				} else {
+					fmt.Fprintf(stdout, "monitor: unknown command: %q\n", l)
+					printHelpMessage(stdout, registeredCommands, additionalHelpMessages)
 				}
 			}
 		}()
@@ -310,41 +168,88 @@ func RunMonitor(ctx context.Context, curRef string, options *controllerapi.Build
 	}
 }
 
+func printHelpMessage(out io.Writer, registeredCommands map[string]types.Command, additional map[string]string) {
+	var names []string
+	for name := range registeredCommands {
+		names = append(names, name)
+	}
+	for name := range additional {
+		names = append(names, name)
+	}
+	sort.Strings(names)
+	fmt.Fprint(out, "Available commands are:\n")
+	w := new(tabwriter.Writer)
+	w.Init(out, 0, 8, 0, '\t', 0)
+	for _, name := range names {
+		var mes string
+		if c, ok := registeredCommands[name]; ok {
+			mes = c.Info().HelpMessage
+		} else if m, ok := additional[name]; ok {
+			mes = m
+		} else {
+			continue
+		}
+		fmt.Fprintln(w, "  "+name+"\t"+mes)
+	}
+	w.Flush()
+}
+
 type readWriter struct {
 	io.Reader
 	io.Writer
 }
 
 type monitor struct {
+	*controllerRef
+
 	muxIO        *ioset.MuxIO
 	invokeIO     *ioset.Forwarder
-	invokeFunc   func(ctx context.Context, ref, pid string, cfg controllerapi.InvokeConfig, in io.ReadCloser, out io.WriteCloser, err io.WriteCloser) error
 	invokeCancel func()
 	attachedPid  atomic.Value
 }
 
-func (m *monitor) rollback(ctx context.Context, ref string, cfg controllerapi.InvokeConfig) string {
+func (m *monitor) DisconnectSession(ctx context.Context, targetID string) error {
+	return m.controllerRef.raw().Disconnect(ctx, targetID)
+}
+
+func (m *monitor) AttachSession(ref string) {
+	m.controllerRef.ref.Store(ref)
+}
+
+func (m *monitor) AttachedSessionID() string {
+	return m.controllerRef.ref.Load().(string)
+}
+
+func (m *monitor) Rollback(ctx context.Context, cfg controllerapi.InvokeConfig) string {
 	pid := identity.NewID()
 	cfg1 := cfg
 	cfg1.Rollback = true
-	return m.startInvoke(ctx, ref, pid, cfg1)
+	return m.startInvoke(ctx, pid, cfg1)
 }
 
-func (m *monitor) exec(ctx context.Context, ref string, cfg controllerapi.InvokeConfig) string {
-	return m.startInvoke(ctx, ref, identity.NewID(), cfg)
+func (m *monitor) Exec(ctx context.Context, cfg controllerapi.InvokeConfig) string {
+	return m.startInvoke(ctx, identity.NewID(), cfg)
 }
 
-func (m *monitor) attach(ctx context.Context, ref, pid string) {
-	m.startInvoke(ctx, ref, pid, controllerapi.InvokeConfig{})
+func (m *monitor) Attach(ctx context.Context, pid string) {
+	m.startInvoke(ctx, pid, controllerapi.InvokeConfig{})
 }
 
-func (m *monitor) close() {
+func (m *monitor) Detach() {
 	if m.invokeCancel != nil {
-		m.invokeCancel()
+		m.invokeCancel() // Finish existing attach
 	}
 }
 
-func (m *monitor) startInvoke(ctx context.Context, ref, pid string, cfg controllerapi.InvokeConfig) string {
+func (m *monitor) AttachedPID() string {
+	return m.attachedPid.Load().(string)
+}
+
+func (m *monitor) close() {
+	m.Detach()
+}
+
+func (m *monitor) startInvoke(ctx context.Context, pid string, cfg controllerapi.InvokeConfig) string {
 	if m.invokeCancel != nil {
 		m.invokeCancel() // Finish existing attach
 	}
@@ -353,7 +258,7 @@ func (m *monitor) startInvoke(ctx context.Context, ref, pid string, cfg controll
 	}
 	go func() {
 		// Start a new invoke
-		if err := m.invoke(ctx, ref, pid, cfg); err != nil {
+		if err := m.invoke(ctx, pid, cfg); err != nil {
 			logrus.Debugf("invoke error: %v", err)
 		}
 		if pid == m.attachedPid.Load() {
@@ -364,13 +269,13 @@ func (m *monitor) startInvoke(ctx context.Context, ref, pid string, cfg controll
 	return pid
 }
 
-func (m *monitor) invoke(ctx context.Context, ref, pid string, cfg controllerapi.InvokeConfig) error {
+func (m *monitor) invoke(ctx context.Context, pid string, cfg controllerapi.InvokeConfig) error {
 	m.muxIO.Enable(1)
 	defer m.muxIO.Disable(1)
 	if err := m.muxIO.SwitchTo(1); err != nil {
 		return errors.Errorf("failed to switch to process IO: %v", err)
 	}
-	if ref == "" {
+	if m.AttachedSessionID() == "" {
 		return nil
 	}
 	invokeCtx, invokeCancel := context.WithCancel(ctx)
@@ -390,27 +295,89 @@ func (m *monitor) invoke(ctx context.Context, ref, pid string, cfg controllerapi
 	defer invokeCancelAndDetachFn()
 	m.invokeCancel = invokeCancelAndDetachFn
 
-	err := m.invokeFunc(invokeCtx, ref, pid, cfg, containerIn.Stdin, containerIn.Stdout, containerIn.Stderr)
+	err := m.controllerRef.invoke(invokeCtx, pid, cfg, containerIn.Stdin, containerIn.Stdout, containerIn.Stderr)
 	close(waitInvokeDoneCh)
 
 	return err
 }
 
-type nopCloser struct {
-	io.Writer
+func newControllerRef(c control.BuildxController, ref string) *controllerRef {
+	cr := &controllerRef{c: c}
+	cr.ref.Store(ref)
+	return cr
 }
 
-func (c nopCloser) Close() error { return nil }
+type controllerRef struct {
+	c   control.BuildxController
+	ref atomic.Value
+}
 
-func isProcessID(ctx context.Context, c control.BuildxController, curRef, ref string) (bool, error) {
-	infos, err := c.ListProcesses(ctx, curRef)
+func (c *controllerRef) raw() control.BuildxController {
+	return c.c
+}
+
+func (c *controllerRef) getRef() (string, error) {
+	ref := c.ref.Load()
+	if ref == "" {
+		return "", errors.Errorf("client is not attached to a session")
+	}
+	return ref.(string), nil
+}
+
+func (c *controllerRef) Build(ctx context.Context, options controllerapi.BuildOptions, in io.ReadCloser, progress progress.Writer) (ref string, resp *client.SolveResponse, err error) {
+	return c.c.Build(ctx, options, in, progress)
+}
+
+func (c *controllerRef) invoke(ctx context.Context, pid string, options controllerapi.InvokeConfig, ioIn io.ReadCloser, ioOut io.WriteCloser, ioErr io.WriteCloser) error {
+	ref, err := c.getRef()
 	if err != nil {
-		return false, err
+		return err
 	}
-	for _, p := range infos {
-		if p.ProcessID == ref {
-			return true, nil
-		}
+	return c.c.Invoke(ctx, ref, pid, options, ioIn, ioOut, ioErr)
+}
+
+func (c *controllerRef) Kill(ctx context.Context) error {
+	return c.c.Kill(ctx)
+}
+
+func (c *controllerRef) List(ctx context.Context) (refs []string, _ error) {
+	return c.c.List(ctx)
+}
+
+func (c *controllerRef) ListProcesses(ctx context.Context) (infos []*controllerapi.ProcessInfo, retErr error) {
+	ref, err := c.getRef()
+	if err != nil {
+		return nil, err
+	}
+	return c.c.ListProcesses(ctx, ref)
+}
+
+func (c *controllerRef) DisconnectProcess(ctx context.Context, pid string) error {
+	ref, err := c.getRef()
+	if err != nil {
+		return err
+	}
+	return c.c.DisconnectProcess(ctx, ref, pid)
+}
+
+func (c *controllerRef) Inspect(ctx context.Context) (*controllerapi.InspectResponse, error) {
+	ref, err := c.getRef()
+	if err != nil {
+		return nil, err
 	}
-	return false, nil
+	return c.c.Inspect(ctx, ref)
 }
+
+func (c *controllerRef) Disconnect(ctx context.Context) error {
+	ref, err := c.getRef()
+	if err != nil {
+		return err
+	}
+	return c.c.Disconnect(ctx, ref)
+}
+
+type nopCloser struct {
+	io.Writer
+}
+
+func (c nopCloser) Close() error { return nil }

monitor/types/types.go

@@ -0,0 +1,75 @@
+package types
+
+import (
+	"context"
+	"io"
+
+	controllerapi "github.com/docker/buildx/controller/pb"
+	"github.com/docker/buildx/util/progress"
+	"github.com/moby/buildkit/client"
+)
+
+// Monitor provides APIs for attaching and controlling the buildx server.
+type Monitor interface {
+	// Rollback re-runs the interactive container with initial rootfs contents.
+	Rollback(ctx context.Context, cfg controllerapi.InvokeConfig) string
+
+	// Rollback executes a process in the interactive container.
+	Exec(ctx context.Context, cfg controllerapi.InvokeConfig) string
+
+	// Attach attaches IO to a process in the container.
+	Attach(ctx context.Context, pid string)
+
+	// AttachedPID returns ID of the attached process.
+	AttachedPID() string
+
+	// Detach detaches IO from the container.
+	Detach()
+
+	// DisconnectSession finishes the specified session.
+	DisconnectSession(ctx context.Context, targetID string) error
+
+	// AttachSession attaches the monitor to the specified session.
+	AttachSession(ref string)
+
+	// AttachedSessionID returns the ID of the attached session.
+	AttachedSessionID() string
+
+	// Build executes the specified build and returns an ID of the session for debugging that build.
+	Build(ctx context.Context, options controllerapi.BuildOptions, in io.ReadCloser, progress progress.Writer) (ref string, resp *client.SolveResponse, err error)
+
+	// Kill kills the buildx server.
+	Kill(ctx context.Context) error
+
+	// List lists sessions.
+	List(ctx context.Context) (refs []string, _ error)
+
+	// ListPrcesses lists processes in the attached session.
+	ListProcesses(ctx context.Context) (infos []*controllerapi.ProcessInfo, retErr error)
+
+	// DisconnectProcess finishes the specified process.
+	DisconnectProcess(ctx context.Context, pid string) error
+
+	// Inspect returns information about the attached build.
+	Inspect(ctx context.Context) (*controllerapi.InspectResponse, error)
+
+	// Disconnect finishes the attached session.
+	Disconnect(ctx context.Context) error
+}
+
+// CommandInfo is information about a command.
+type CommandInfo struct {
+
+	// HelpMessage is the message printed to the console when "help" command is invoked.
+	HelpMessage string
+}
+
+// Command represents a command for debugging.
+type Command interface {
+
+	// Exec executes the command.
+	Exec(ctx context.Context, args []string) error
+
+	// Info returns information of the command.
+	Info() CommandInfo
+}