refactor(sandbox): use structured commands

Author: ramon-langchain

internal/cmd/sandbox_box.go

@@ -417,60 +417,59 @@ var sandboxStopCommand = structured.Command[struct{}]{
 `),
 }
 
-func newSandboxExecCmd() *cobra.Command {
-	cmd := &cobra.Command{
-		Use:   "exec <name> -- <command>",
-		Short: "Execute a command inside a sandbox",
-		Long: `Execute a one-off command inside a running sandbox and print its output.
+type sandboxExecInput struct{}
+
+var sandboxExecCommand = structured.Command[sandboxExecInput]{
+	Use:   "exec <name> -- <command>",
+	Short: "Execute a command inside a sandbox",
+	Long: `Execute a one-off command inside a running sandbox and print its output.
 
 Examples:
   langsmith sandbox exec my-vm -- uname -a
   langsmith sandbox exec my-vm -- ls -la /
   langsmith sandbox exec my-vm -- cat /etc/os-release`,
-		Args:               cobra.MinimumNArgs(1),
-		DisableFlagParsing: false,
-		RunE: func(cmd *cobra.Command, args []string) error {
-			name := args[0]
-
-			// Everything after "--" is the command.
-			cmdArgs := cmd.ArgsLenAtDash()
-			if cmdArgs < 0 || cmdArgs >= len(args) {
-				return fmt.Errorf("usage: langsmith sandbox exec <name> -- <command>")
-			}
-			command := args[cmdArgs:]
-			if len(command) == 0 {
-				return fmt.Errorf("no command specified")
-			}
+	Args:         cobra.MinimumNArgs(1),
+	Input:        func(cmd *cobra.Command) sandboxExecInput { return sandboxExecInput{} },
+	CustomOutput: true,
+	Action: func(ctx context.Context, cmd *cobra.Command, in sandboxExecInput, args []string) (any, error) {
+		name := args[0]
 
-			ctx := cmd.Context()
-			if ctx == nil {
-				ctx = context.Background()
-			}
-			c, err := cmdutil.GetClient(cmd)
-			if err != nil {
-				return err
-			}
+		cmdArgs := cmd.ArgsLenAtDash()
+		if cmdArgs < 0 || cmdArgs >= len(args) {
+			return nil, fmt.Errorf("usage: langsmith sandbox exec <name> -- <command>")
+		}
+		command := args[cmdArgs:]
+		if len(command) == 0 {
+			return nil, fmt.Errorf("no command specified")
+		}
 
-			result, err := c.SDK.Sandboxes.Boxes.Run(ctx, name, langsmith.SandboxBoxRunParams{
-				Command: langsmith.F(sandboxShellCommand(command)),
-			})
-			if err != nil {
-				return fmt.Errorf("execute: %w", err)
-			}
+		c, err := cmdutil.GetClient(cmd)
+		if err != nil {
+			return nil, err
+		}
 
-			if result.Stdout != "" {
-				fmt.Print(result.Stdout)
-			}
-			if result.Stderr != "" {
-				fmt.Fprint(os.Stderr, result.Stderr)
-			}
-			if result.ExitCode != 0 {
-				os.Exit(int(result.ExitCode))
-			}
-			return nil
-		},
-	}
-	return cmd
+		result, err := c.SDK.Sandboxes.Boxes.Run(ctx, name, langsmith.SandboxBoxRunParams{
+			Command: langsmith.F(sandboxShellCommand(command)),
+		})
+		if err != nil {
+			return nil, fmt.Errorf("execute: %w", err)
+		}
+
+		if result.Stdout != "" {
+			fmt.Print(result.Stdout)
+		}
+		if result.Stderr != "" {
+			fmt.Fprint(os.Stderr, result.Stderr)
+		}
+		if result.ExitCode != 0 {
+			os.Exit(int(result.ExitCode))
+		}
+		return nil, nil
+	},
+}
+
+func newSandboxExecCmd() *cobra.Command {
+	return sandboxExecCommand.Cobra()
 }
 
 func sandboxShellCommand(args []string) string {

internal/cmd/sandbox_console.go

@@ -11,16 +11,21 @@ import (
 	"sync"
 	"syscall"
 
+	"github.com/langchain-ai/langsmith-cli/internal/structured"
 	langsmith "github.com/langchain-ai/langsmith-go"
 	"github.com/spf13/cobra"
 	"golang.org/x/term"
 )
 
-func newSandboxConsoleCmd() *cobra.Command {
-	cmd := &cobra.Command{
-		Use:   "console <name>",
-		Short: "Open an interactive shell inside a sandbox",
-		Long: `Open an interactive terminal session inside a running sandbox.
+type sandboxConsoleInput struct {
+	Shell           string
+	ForwardSSHAgent bool
+}
+
+var sandboxConsoleCommand = structured.Command[*sandboxConsoleInput]{
+	Use:   "console <name>",
+	Short: "Open an interactive shell inside a sandbox",
+	Long: `Open an interactive terminal session inside a running sandbox.
 
 Connects via WebSocket to the sandbox daemon and allocates a PTY,
 giving you a full interactive shell (bash by default).
@@ -29,18 +34,21 @@ Examples:
   langsmith sandbox console my-vm
   langsmith sandbox console my-vm --shell /bin/sh
   langsmith sandbox console my-vm --forward-ssh-agent`,
-		Args: cobra.ExactArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
-			shell, _ := cmd.Flags().GetString("shell")
-			forwardSSHAgent, _ := cmd.Flags().GetBool("forward-ssh-agent")
-			return runConsole(args[0], shell, forwardSSHAgent)
-		},
-	}
-
-	cmd.Flags().String("shell", "", "Shell to use (default: sandbox default, usually /bin/bash)")
-	cmd.Flags().Bool("forward-ssh-agent", false, "Forward the local SSH agent (SSH_AUTH_SOCK) into the sandbox")
+	Args: cobra.ExactArgs(1),
+	Input: func(cmd *cobra.Command) *sandboxConsoleInput {
+		in := &sandboxConsoleInput{}
+		cmd.Flags().StringVar(&in.Shell, "shell", in.Shell, "Shell to use (default: sandbox default, usually /bin/bash)")
+		cmd.Flags().BoolVar(&in.ForwardSSHAgent, "forward-ssh-agent", in.ForwardSSHAgent, "Forward the local SSH agent (SSH_AUTH_SOCK) into the sandbox")
+		return in
+	},
+	CustomOutput: true,
+	Action: func(ctx context.Context, cmd *cobra.Command, in *sandboxConsoleInput, args []string) (any, error) {
+		return nil, runConsole(args[0], in.Shell, in.ForwardSSHAgent)
+	},
+}
 
-	return cmd
+func newSandboxConsoleCmd() *cobra.Command {
+	return sandboxConsoleCommand.Cobra()
 }
 
 func runConsole(name, shell string, forwardSSHAgent bool) error {

internal/cmd/sandbox_console_windows.go

@@ -1,18 +1,23 @@
 package cmd
 
 import (
+	"context"
 	"fmt"
 
+	"github.com/langchain-ai/langsmith-cli/internal/structured"
 	"github.com/spf13/cobra"
 )
 
+var sandboxConsoleCommand = structured.Command[struct{}]{
+	Use:          "console <name>",
+	Short:        "Open an interactive shell inside a sandbox (not supported on Windows)",
+	Args:         cobra.ExactArgs(1),
+	CustomOutput: true,
+	Action: func(ctx context.Context, cmd *cobra.Command, in struct{}, args []string) (any, error) {
+		return nil, fmt.Errorf("sandbox console is not supported on Windows; use SSH instead: langsmith sandbox ssh-setup %s", args[0])
+	},
+}
+
 func newSandboxConsoleCmd() *cobra.Command {
-	return &cobra.Command{
-		Use:   "console <name>",
-		Short: "Open an interactive shell inside a sandbox (not supported on Windows)",
-		Args:  cobra.ExactArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
-			return fmt.Errorf("sandbox console is not supported on Windows; use SSH instead: langsmith sandbox ssh-setup %s", args[0])
-		},
-	}
+	return sandboxConsoleCommand.Cobra()
 }

internal/cmd/sandbox_ssh.go

@@ -10,17 +10,19 @@ import (
 
 	"github.com/langchain-ai/langsmith-cli/internal/client"
 	"github.com/langchain-ai/langsmith-cli/internal/cmdutil"
+	"github.com/langchain-ai/langsmith-cli/internal/structured"
 	"github.com/langchain-ai/langsmith-go"
 	"github.com/spf13/cobra"
 )
 
-func newSandboxSSHSetupCmd() *cobra.Command {
-	var identity string
+type sandboxSSHSetupInput struct {
+	Identity string
+}
 
-	cmd := &cobra.Command{
-		Use:   "ssh-setup <name>",
-		Short: "Upload your SSH public key and configure ~/.ssh/config for a sandbox",
-		Long: `Upload your SSH public key to a running sandbox so you can connect
+var sandboxSSHSetupCommand = structured.Command[*sandboxSSHSetupInput]{
+	Use:   "ssh-setup <name>",
+	Short: "Upload your SSH public key and configure ~/.ssh/config for a sandbox",
+	Long: `Upload your SSH public key to a running sandbox so you can connect
 with standard SSH tools (ssh, scp, rsync, sftp).
 
 This command uploads your key, fetches the host key, writes a
@@ -30,15 +32,20 @@ immediately connect with: ssh sandbox-<name>
 Examples:
   langsmith sandbox ssh-setup my-sandbox
   langsmith sandbox ssh-setup my-sandbox --identity ~/.ssh/id_ed25519.pub`,
-		Args: cobra.ExactArgs(1),
-		RunE: func(cmd *cobra.Command, args []string) error {
-			return runSSHSetup(cmd, args[0], identity)
-		},
-	}
-
-	cmd.Flags().StringVar(&identity, "identity", "", "Path to SSH public key (default: auto-detect)")
+	Args: cobra.ExactArgs(1),
+	Input: func(cmd *cobra.Command) *sandboxSSHSetupInput {
+		in := &sandboxSSHSetupInput{}
+		cmd.Flags().StringVar(&in.Identity, "identity", in.Identity, "Path to SSH public key (default: auto-detect)")
+		return in
+	},
+	CustomOutput: true,
+	Action: func(ctx context.Context, cmd *cobra.Command, in *sandboxSSHSetupInput, args []string) (any, error) {
+		return nil, runSSHSetup(cmd, args[0], in.Identity)
+	},
+}
 
-	return cmd
+func newSandboxSSHSetupCmd() *cobra.Command {
+	return sandboxSSHSetupCommand.Cobra()
 }
 
 func runSSHSetup(cmd *cobra.Command, name, identity string) error {

internal/cmd/sandbox_test.go

@@ -9,6 +9,7 @@ import (
 	"testing"
 
 	langsmith "github.com/langchain-ai/langsmith-go"
+	"github.com/spf13/cobra"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -157,6 +158,65 @@ func TestSandboxCreateCmd_AllowsNoArgs(t *testing.T) {
 	require.NotContains(t, body, "snapshot_id")
 }
 
+func TestSandboxExecCmd_PositionalNameAndCommandSeparator(t *testing.T) {
+	cmd := newSandboxExecCmd()
+	require.NoError(t, cmd.Args(cmd, []string{"my-vm", "echo", "hi"}))
+	require.Error(t, cmd.Args(cmd, []string{}))
+}
+
+func TestSandboxExecCmd_WritesCommandOutputDirectly(t *testing.T) {
+	var body map[string]any
+	ts := newTestServer(t, func(w http.ResponseWriter, r *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		switch {
+		case r.Method == http.MethodGet && r.URL.Path == "/v2/sandboxes/boxes/my-vm":
+			_, err := w.Write([]byte(`{"id":"box-id","name":"my-vm","status":"ready","dataplane_url":"` + tsURL(t, r) + `"}`))
+			require.NoError(t, err)
+		case r.Method == http.MethodPost && r.URL.Path == "/execute":
+			require.NoError(t, json.NewDecoder(r.Body).Decode(&body))
+			_, err := w.Write([]byte(`{"stdout":"hello\n","stderr":"","exit_code":0}`))
+			require.NoError(t, err)
+		default:
+			t.Fatalf("unexpected request %s %s", r.Method, r.URL.Path)
+		}
+	})
+
+	var out string
+	var err error
+	stdout := captureStdout(t, func() {
+		out, err = executeCommand(t, "--api-key", "test-key", "--api-url", ts.URL, "sandbox", "exec", "my-vm", "--", "echo", "hello")
+	})
+
+	require.NoError(t, err)
+	assert.Empty(t, out)
+	assert.Equal(t, "hello\n", stdout)
+	assert.Equal(t, "'echo' 'hello'", body["command"])
+}
+
+func tsURL(t *testing.T, r *http.Request) string {
+	t.Helper()
+	return "http://" + r.Host
+}
+
+func TestSandboxCustomOutputCommands_Flags(t *testing.T) {
+	tests := []struct {
+		name string
+		cmd  *cobra.Command
+		want []string
+	}{
+		{"console", newSandboxConsoleCmd(), []string{"shell", "forward-ssh-agent"}},
+		{"tunnel", newSandboxTunnelCmd(), []string{"url", "name", "remote-port", "local-port", "stdio", "log-level"}},
+		{"ssh-setup", newSandboxSSHSetupCmd(), []string{"identity"}},
+	}
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			for _, name := range tc.want {
+				require.NotNil(t, tc.cmd.Flags().Lookup(name), "flag --%s not found", name)
+			}
+		})
+	}
+}
+
 func TestSandboxTunnelCmd_PositionalNameOrURL(t *testing.T) {
 	cmd := newSandboxTunnelCmd()
 	// Should accept 0 or 1 args