banger/internal/cli/commands_vm_session.go

package cli

import (
	"context"
	"errors"
	"fmt"
	"io"
	"net"
	"strings"

	"banger/internal/api"
	"banger/internal/model"
	"banger/internal/sessionstream"

	"github.com/spf13/cobra"
)

func (d *deps) newVMSessionCommand() *cobra.Command {
	cmd := &cobra.Command{
		Use:   "session",
		Short: "Manage long-lived guest commands inside a VM",
		Long:  "Start, inspect, stop, and attach to daemon-managed guest commands. Pipe-mode sessions expose live stdio for interactive protocols. Attach is exclusive and currently uses a same-host local bridge.",
		RunE:  helpNoArgs,
	}
	cmd.AddCommand(
		d.newVMSessionStartCommand(),
		d.newVMSessionListCommand(),
		d.newVMSessionShowCommand(),
		d.newVMSessionLogsCommand(),
		d.newVMSessionStopCommand(),
		d.newVMSessionKillCommand(),
		d.newVMSessionAttachCommand(),
		d.newVMSessionSendCommand(),
	)
	return cmd
}

func (d *deps) newVMSessionStartCommand() *cobra.Command {
	var name string
	var cwd string
	var stdinMode string
	var envPairs []string
	var tagPairs []string
	var requiredCommands []string
	cmd := &cobra.Command{
		Use:               "start <id-or-name> <command> [args...]",
		Short:             "Start a managed guest command",
		Long:              "Start a daemon-managed guest command. The daemon verifies that the guest working directory exists and that the requested command is present in guest PATH before launch. Use --stdin-mode pipe when you need live attach.",
		Args:              minArgsUsage(2, "usage: banger vm session start <id-or-name> [flags] -- <command> [args...]"),
		ValidArgsFunction: d.completeVMNameOnlyAtPos0,
		Example: strings.TrimSpace(`
  banger vm session start devbox --name planner --cwd /root/repo --stdin-mode pipe --require-command git -- pi --mode rpc --no-session
  banger vm session start devbox --name shell --stdin-mode pipe -- bash -lc 'exec bash'
`),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := d.ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			env, err := parseKeyValuePairs(envPairs)
			if err != nil {
				return err
			}
			tags, err := parseKeyValuePairs(tagPairs)
			if err != nil {
				return err
			}
			result, err := d.guestSessionStart(cmd.Context(), layout.SocketPath, api.GuestSessionStartParams{
				VMIDOrName:       args[0],
				Name:             name,
				Command:          args[1],
				Args:             append([]string(nil), args[2:]...),
				CWD:              cwd,
				Env:              env,
				StdinMode:        stdinMode,
				Tags:             tags,
				RequiredCommands: append([]string(nil), requiredCommands...),
			})
			if err != nil {
				return err
			}
			if err := printGuestSessionSummary(cmd.OutOrStdout(), result.Session); err != nil {
				return err
			}
			if result.Session.Status == model.GuestSessionStatusFailed && strings.TrimSpace(result.Session.LaunchMessage) != "" {
				_, _ = fmt.Fprintf(cmd.ErrOrStderr(), "warning: session failed at %s: %s\n", result.Session.LaunchStage, result.Session.LaunchMessage)
			}
			return nil
		},
	}
	cmd.Flags().StringVar(&name, "name", "", "session name")
	cmd.Flags().StringVar(&cwd, "cwd", "", "guest working directory; must already exist")
	cmd.Flags().StringVar(&stdinMode, "stdin-mode", string(model.GuestSessionStdinClosed), "stdin mode: closed or pipe (pipe enables attach)")
	cmd.Flags().StringArrayVar(&envPairs, "env", nil, "environment entry in KEY=VALUE form")
	cmd.Flags().StringArrayVar(&tagPairs, "tag", nil, "session tag in KEY=VALUE form")
	cmd.Flags().StringArrayVar(&requiredCommands, "require-command", nil, "extra guest command that must exist in PATH before launch; repeatable")
	return cmd
}

func (d *deps) newVMSessionListCommand() *cobra.Command {
	return &cobra.Command{
		Use:               "list <id-or-name>",
		Aliases:           []string{"ls"},
		Short:             "List managed guest commands for a VM",
		Args:              exactArgsUsage(1, "usage: banger vm session list <id-or-name>"),
		ValidArgsFunction: d.completeVMNameOnlyAtPos0,
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := d.ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := d.guestSessionList(cmd.Context(), layout.SocketPath, args[0])
			if err != nil {
				return err
			}
			return printGuestSessionTable(cmd.OutOrStdout(), result.Sessions)
		},
	}
}

func (d *deps) newVMSessionShowCommand() *cobra.Command {
	return &cobra.Command{
		Use:               "show <id-or-name> <session>",
		Short:             "Show managed guest command details",
		Args:              exactArgsUsage(2, "usage: banger vm session show <id-or-name> <session>"),
		ValidArgsFunction: d.completeSessionNames,
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := d.ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := d.guestSessionGet(cmd.Context(), layout.SocketPath, api.GuestSessionRefParams{VMIDOrName: args[0], SessionIDOrName: args[1]})
			if err != nil {
				return err
			}
			return printJSON(cmd.OutOrStdout(), result.Session)
		},
	}
}

func (d *deps) newVMSessionLogsCommand() *cobra.Command {
	var stream string
	var tailLines int
	cmd := &cobra.Command{
		Use:               "logs <id-or-name> <session>",
		Short:             "Show stdout or stderr for a guest session",
		Args:              exactArgsUsage(2, "usage: banger vm session logs [--stream stdout|stderr] [-n LINES] <id-or-name> <session>"),
		ValidArgsFunction: d.completeSessionNames,
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := d.ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := d.guestSessionLogs(cmd.Context(), layout.SocketPath, api.GuestSessionLogsParams{VMIDOrName: args[0], SessionIDOrName: args[1], Stream: stream, TailLines: tailLines})
			if err != nil {
				return err
			}
			_, err = fmt.Fprint(cmd.OutOrStdout(), result.Content)
			return err
		},
	}
	cmd.Flags().StringVar(&stream, "stream", "stdout", "log stream to read")
	cmd.Flags().IntVarP(&tailLines, "lines", "n", 200, "number of lines to tail")
	return cmd
}

func (d *deps) newVMSessionStopCommand() *cobra.Command {
	return &cobra.Command{
		Use:               "stop <id-or-name> <session>",
		Short:             "Send SIGTERM to a guest session",
		Args:              exactArgsUsage(2, "usage: banger vm session stop <id-or-name> <session>"),
		ValidArgsFunction: d.completeSessionNames,
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := d.ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := d.guestSessionStop(cmd.Context(), layout.SocketPath, api.GuestSessionRefParams{VMIDOrName: args[0], SessionIDOrName: args[1]})
			if err != nil {
				return err
			}
			return printGuestSessionSummary(cmd.OutOrStdout(), result.Session)
		},
	}
}

func (d *deps) newVMSessionKillCommand() *cobra.Command {
	return &cobra.Command{
		Use:               "kill <id-or-name> <session>",
		Short:             "Send SIGKILL to a guest session",
		Args:              exactArgsUsage(2, "usage: banger vm session kill <id-or-name> <session>"),
		ValidArgsFunction: d.completeSessionNames,
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := d.ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := d.guestSessionKill(cmd.Context(), layout.SocketPath, api.GuestSessionRefParams{VMIDOrName: args[0], SessionIDOrName: args[1]})
			if err != nil {
				return err
			}
			return printGuestSessionSummary(cmd.OutOrStdout(), result.Session)
		},
	}
}

func (d *deps) newVMSessionAttachCommand() *cobra.Command {
	return &cobra.Command{
		Use:               "attach <id-or-name> <session>",
		Short:             "Attach local stdio to an attachable guest session",
		Long:              "Attach local stdio to a pipe-mode session through a daemon-created local Unix socket bridge. Only one active attach is allowed at a time, and the client must run on the same host as the daemon.",
		Args:              exactArgsUsage(2, "usage: banger vm session attach <id-or-name> <session>"),
		ValidArgsFunction: d.completeSessionNames,
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := d.ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			result, err := d.guestSessionAttachBegin(cmd.Context(), layout.SocketPath, api.GuestSessionAttachBeginParams{VMIDOrName: args[0], SessionIDOrName: args[1]})
			if err != nil {
				return err
			}
			socketPath := strings.TrimSpace(result.SocketPath)
			if socketPath == "" && result.TransportKind == "unix_socket" {
				socketPath = strings.TrimSpace(result.TransportTarget)
			}
			return runGuestSessionAttach(cmd.Context(), cmd.InOrStdin(), cmd.OutOrStdout(), cmd.ErrOrStderr(), socketPath)
		},
	}
}

func (d *deps) newVMSessionSendCommand() *cobra.Command {
	var message string
	cmd := &cobra.Command{
		Use:               "send <id-or-name> <session>",
		Short:             "Write bytes to a running guest session's stdin pipe",
		Long:              "Write a payload to the stdin pipe of a running pipe-mode guest session without holding the exclusive attach. Use --message for an inline JSONL string, or pipe bytes via stdin when --message is omitted. A trailing newline is appended to --message values that lack one.",
		Args:              exactArgsUsage(2, "usage: banger vm session send <id-or-name> <session> [--message '<json>']"),
		ValidArgsFunction: d.completeSessionNames,
		Example: strings.TrimSpace(`
  banger vm session send devbox planner --message '{"type":"abort"}'
  banger vm session send devbox planner --message '{"type":"steer","message":"Focus on src/"}'
  echo '{"type":"prompt","prompt":"Summarize."}' | banger vm session send devbox planner
`),
		RunE: func(cmd *cobra.Command, args []string) error {
			layout, _, err := d.ensureDaemon(cmd.Context())
			if err != nil {
				return err
			}
			var payload []byte
			if message != "" {
				payload = []byte(message)
				if len(payload) > 0 && payload[len(payload)-1] != '\n' {
					payload = append(payload, '\n')
				}
			} else {
				payload, err = io.ReadAll(cmd.InOrStdin())
				if err != nil {
					return fmt.Errorf("read stdin: %w", err)
				}
			}
			result, err := d.guestSessionSend(cmd.Context(), layout.SocketPath, api.GuestSessionSendParams{
				VMIDOrName:      args[0],
				SessionIDOrName: args[1],
				Payload:         payload,
			})
			if err != nil {
				return err
			}
			_, err = fmt.Fprintf(cmd.OutOrStdout(), "sent %d bytes to session %s\n", result.BytesWritten, result.Session.Name)
			return err
		},
	}
	cmd.Flags().StringVar(&message, "message", "", "JSONL message to send; a trailing newline is appended if absent")
	return cmd
}

func parseKeyValuePairs(values []string) (map[string]string, error) {
	if len(values) == 0 {
		return nil, nil
	}
	result := make(map[string]string, len(values))
	for _, value := range values {
		key, raw, ok := strings.Cut(value, "=")
		if !ok || strings.TrimSpace(key) == "" {
			return nil, fmt.Errorf("invalid key=value entry %q", value)
		}
		result[strings.TrimSpace(key)] = raw
	}
	return result, nil
}

func runGuestSessionAttach(ctx context.Context, stdin io.Reader, stdout, stderr io.Writer, socketPath string) error {
	conn, err := (&net.Dialer{}).DialContext(ctx, "unix", socketPath)
	if err != nil {
		return err
	}
	defer conn.Close()
	writeErrCh := make(chan error, 1)
	go func() {
		writeErrCh <- streamGuestSessionAttachInput(conn, stdin)
	}()
	for {
		channel, payload, err := sessionstream.ReadFrame(conn)
		if err != nil {
			if ctx.Err() != nil {
				return ctx.Err()
			}
			if errors.Is(err, io.EOF) {
				return nil
			}
			return err
		}
		switch channel {
		case sessionstream.ChannelStdout:
			if _, err := stdout.Write(payload); err != nil {
				return err
			}
		case sessionstream.ChannelStderr:
			if _, err := stderr.Write(payload); err != nil {
				return err
			}
		case sessionstream.ChannelControl:
			message, err := sessionstream.ReadControl(payload)
			if err != nil {
				return err
			}
			switch message.Type {
			case "exit":
				if message.ExitCode != nil && *message.ExitCode != 0 {
					return fmt.Errorf("guest session exited with code %d", *message.ExitCode)
				}
				return nil
			case "error":
				if strings.TrimSpace(message.Error) == "" {
					return errors.New("guest session attach failed")
				}
				return errors.New(message.Error)
			}
		}
		select {
		case err := <-writeErrCh:
			if err != nil {
				return err
			}
		default:
		}
	}
}

func streamGuestSessionAttachInput(conn net.Conn, stdin io.Reader) error {
	if stdin == nil {
		return sessionstream.WriteControl(conn, sessionstream.ControlMessage{Type: "eof"})
	}
	buffer := make([]byte, 32*1024)
	for {
		n, err := stdin.Read(buffer)
		if n > 0 {
			if writeErr := sessionstream.WriteFrame(conn, sessionstream.ChannelStdin, buffer[:n]); writeErr != nil {
				return writeErr
			}
		}
		if err != nil {
			if errors.Is(err, io.EOF) {
				return sessionstream.WriteControl(conn, sessionstream.ControlMessage{Type: "eof"})
			}
			return err
		}
	}
}