coder
diff --git a/‎codersdk/toolsdk/ssh.go
Lines changed: 307 additions & 0 deletions b/‎codersdk/toolsdk/ssh.go
Lines changed: 307 additions & 0 deletions
@@ -0,0 +1,307 @@
+package toolsdk
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+
+	gossh "golang.org/x/crypto/ssh"
+	"golang.org/x/xerrors"
+
+	"github.com/coder/aisdk-go"
+
+	"github.com/coder/coder/v2/cli/cliui"
+	"github.com/coder/coder/v2/codersdk"
+	"github.com/coder/coder/v2/codersdk/workspacesdk"
+)
+
+type WorkspaceSSHExecArgs struct {
+	Workspace string `json:"workspace"`
+	Command   string `json:"command"`
+}
+
+type WorkspaceSSHExecResult struct {
+	Output   string `json:"output"`
+	ExitCode int    `json:"exit_code"`
+}
+
+var WorkspaceSSHExec = Tool[WorkspaceSSHExecArgs, WorkspaceSSHExecResult]{
+	Tool: aisdk.Tool{
+		Name: ToolNameWorkspaceSSHExec,
+		Description: `Execute a command in a Coder workspace via SSH.
+
+This tool provides the same functionality as the 'coder ssh <workspace> <command>' CLI command.
+It automatically starts the workspace if it's stopped and waits for the agent to be ready.
+The output is trimmed of leading and trailing whitespace.
+
+The workspace parameter supports various formats:
+- workspace (uses current user)
+- owner/workspace
+- owner--workspace
+- workspace.agent (specific agent)
+- owner/workspace.agent
+
+Examples:
+- workspace: "my-workspace", command: "ls -la"
+- workspace: "john/dev-env", command: "git status"
+- workspace: "my-workspace.main", command: "docker ps"`,
+		Schema: aisdk.Schema{
+			Properties: map[string]any{
+				"workspace": map[string]any{
+					"type":        "string",
+					"description": "The workspace name in format [owner/]workspace[.agent]. If owner is not specified, the authenticated user is used.",
+				},
+				"command": map[string]any{
+					"type":        "string",
+					"description": "The command to execute in the workspace.",
+				},
+			},
+			Required: []string{"workspace", "command"},
+		},
+	},
+	Handler: func(ctx context.Context, deps Deps, args WorkspaceSSHExecArgs) (WorkspaceSSHExecResult, error) {
+		if args.Workspace == "" {
+			return WorkspaceSSHExecResult{}, xerrors.New("workspace name cannot be empty")
+		}
+		if args.Command == "" {
+			return WorkspaceSSHExecResult{}, xerrors.New("command cannot be empty")
+		}
+
+		// Normalize workspace input to handle various formats
+		workspaceName := NormalizeWorkspaceInput(args.Workspace)
+
+		// Find workspace and agent
+		_, workspaceAgent, err := findWorkspaceAndAgentWithAutostart(ctx, deps.coderClient, workspaceName)
+		if err != nil {
+			return WorkspaceSSHExecResult{}, xerrors.Errorf("failed to find workspace: %w", err)
+		}
+
+		// Wait for agent to be ready
+		err = cliui.Agent(ctx, nil, workspaceAgent.ID, cliui.AgentOptions{
+			FetchInterval: 0,
+			Fetch:         deps.coderClient.WorkspaceAgent,
+			FetchLogs:     deps.coderClient.WorkspaceAgentLogsAfter,
+			Wait:          true, // Always wait for startup scripts
+		})
+		if err != nil {
+			return WorkspaceSSHExecResult{}, xerrors.Errorf("agent not ready: %w", err)
+		}
+
+		// Create workspace SDK client for agent connection
+		wsClient := workspacesdk.New(deps.coderClient)
+
+		// Dial agent
+		conn, err := wsClient.DialAgent(ctx, workspaceAgent.ID, &workspacesdk.DialAgentOptions{
+			BlockEndpoints: false,
+		})
+		if err != nil {
+			return WorkspaceSSHExecResult{}, xerrors.Errorf("failed to dial agent: %w", err)
+		}
+		defer conn.Close()
+
+		// Wait for connection to be reachable
+		conn.AwaitReachable(ctx)
+
+		// Create SSH client
+		sshClient, err := conn.SSHClient(ctx)
+		if err != nil {
+			return WorkspaceSSHExecResult{}, xerrors.Errorf("failed to create SSH client: %w", err)
+		}
+		defer sshClient.Close()
+
+		// Create SSH session
+		session, err := sshClient.NewSession()
+		if err != nil {
+			return WorkspaceSSHExecResult{}, xerrors.Errorf("failed to create SSH session: %w", err)
+		}
+		defer session.Close()
+
+		// Execute command and capture output
+		output, err := session.CombinedOutput(args.Command)
+		outputStr := strings.TrimSpace(string(output))
+
+		if err != nil {
+			// Check if it's an SSH exit error to get the exit code
+			var exitErr *gossh.ExitError
+			if errors.As(err, &exitErr) {
+				return WorkspaceSSHExecResult{
+					Output:   outputStr,
+					ExitCode: exitErr.ExitStatus(),
+				}, nil
+			}
+			// For other errors, return exit code 1
+			return WorkspaceSSHExecResult{
+				Output:   outputStr,
+				ExitCode: 1,
+			}, nil
+		}
+
+		return WorkspaceSSHExecResult{
+			Output:   outputStr,
+			ExitCode: 0,
+		}, nil
+	},
+}
+
+// findWorkspaceAndAgentWithAutostart finds workspace and agent by name and auto-starts if needed
+func findWorkspaceAndAgentWithAutostart(ctx context.Context, client *codersdk.Client, workspaceName string) (codersdk.Workspace, codersdk.WorkspaceAgent, error) {
+	return findWorkspaceAndAgent(ctx, client, workspaceName)
+}
+
+// findWorkspaceAndAgent finds workspace and agent by name with auto-start support
+func findWorkspaceAndAgent(ctx context.Context, client *codersdk.Client, workspaceName string) (codersdk.Workspace, codersdk.WorkspaceAgent, error) {
+	// Parse workspace name to extract workspace and agent parts
+	parts := strings.Split(workspaceName, ".")
+	var agentName string
+	if len(parts) >= 2 {
+		agentName = parts[1]
+		workspaceName = parts[0]
+	}
+
+	// Get workspace
+	workspace, err := namedWorkspace(ctx, client, workspaceName)
+	if err != nil {
+		return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
+	}
+
+	// Auto-start workspace if needed
+	if workspace.LatestBuild.Transition != codersdk.WorkspaceTransitionStart {
+		if workspace.LatestBuild.Transition == codersdk.WorkspaceTransitionDelete {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q is deleted", workspace.Name)
+		}
+		if workspace.LatestBuild.Job.Status == codersdk.ProvisionerJobFailed {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q is in failed state", workspace.Name)
+		}
+		if workspace.LatestBuild.Status != codersdk.WorkspaceStatusStopped {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace must be started; was unable to autostart as the last build job is %q, expected %q",
+				workspace.LatestBuild.Status, codersdk.WorkspaceStatusStopped)
+		}
+
+		// Start workspace
+		build, err := client.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{
+			Transition: codersdk.WorkspaceTransitionStart,
+		})
+		if err != nil {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("failed to start workspace: %w", err)
+		}
+
+		// Wait for build to complete
+		for {
+			build, err = client.WorkspaceBuild(ctx, build.ID)
+			if err != nil {
+				return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("failed to get build status: %w", err)
+			}
+			if build.Job.CompletedAt != nil {
+				break
+			}
+			// Small delay before checking again
+			select {
+			case <-ctx.Done():
+				return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, ctx.Err()
+			default:
+			}
+		}
+
+		// Refresh workspace after build completes
+		workspace, err = client.Workspace(ctx, workspace.ID)
+		if err != nil {
+			return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
+		}
+	}
+
+	// Find agent
+	workspaceAgent, err := getWorkspaceAgent(workspace, agentName)
+	if err != nil {
+		return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
+	}
+
+	return workspace, workspaceAgent, nil
+}
+
+// getWorkspaceAgent finds the specified agent in the workspace
+func getWorkspaceAgent(workspace codersdk.Workspace, agentName string) (codersdk.WorkspaceAgent, error) {
+	resources := workspace.LatestBuild.Resources
+
+	var agents []codersdk.WorkspaceAgent
+	var availableNames []string
+
+	for _, resource := range resources {
+		for _, agent := range resource.Agents {
+			availableNames = append(availableNames, agent.Name)
+			agents = append(agents, agent)
+		}
+	}
+
+	if len(agents) == 0 {
+		return codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q has no agents", workspace.Name)
+	}
+
+	if agentName != "" {
+		for _, agent := range agents {
+			if agent.Name == agentName || agent.ID.String() == agentName {
+				return agent, nil
+			}
+		}
+		return codersdk.WorkspaceAgent{}, xerrors.Errorf("agent not found by name %q, available agents: %v", agentName, availableNames)
+	}
+
+	if len(agents) == 1 {
+		return agents[0], nil
+	}
+
+	return codersdk.WorkspaceAgent{}, xerrors.Errorf("multiple agents found, please specify the agent name, available agents: %v", availableNames)
+}
+
+// namedWorkspace gets a workspace by owner/name or just name
+func namedWorkspace(ctx context.Context, client *codersdk.Client, identifier string) (codersdk.Workspace, error) {
+	// Parse owner and workspace name
+	parts := strings.SplitN(identifier, "/", 2)
+	var owner, workspaceName string
+
+	if len(parts) == 2 {
+		owner = parts[0]
+		workspaceName = parts[1]
+	} else {
+		owner = "me"
+		workspaceName = identifier
+	}
+
+	// Handle -- separator format (convert to / format)
+	if strings.Contains(identifier, "--") && !strings.Contains(identifier, "/") {
+		dashParts := strings.SplitN(identifier, "--", 2)
+		if len(dashParts) == 2 {
+			owner = dashParts[0]
+			workspaceName = dashParts[1]
+		}
+	}
+
+	return client.WorkspaceByOwnerAndName(ctx, owner, workspaceName, codersdk.WorkspaceOptions{})
+}
+
+// NormalizeWorkspaceInput converts workspace name input to standard format
+// Handles formats like: workspace, workspace.agent, owner/workspace, owner--workspace, etc.
+func NormalizeWorkspaceInput(input string) string {
+	// This matches the logic from cli/ssh.go
+	// Split on "/", "--", and "."
+	workspaceNameRe := strings.NewReplacer("/", ":", "--", ":", ".", ":")
+	parts := strings.Split(workspaceNameRe.Replace(input), ":")
+
+	switch len(parts) {
+	case 1:
+		return input // "workspace"
+	case 2:
+		if strings.Contains(input, ".") {
+			return fmt.Sprintf("%s.%s", parts[0], parts[1]) // "workspace.agent"
+		}
+		return fmt.Sprintf("%s/%s", parts[0], parts[1]) // "owner/workspace"
+	case 3:
+		// If the only separator is a dot, it's the Coder Connect format
+		if !strings.Contains(input, "/") && !strings.Contains(input, "--") {
+			return fmt.Sprintf("%s/%s.%s", parts[2], parts[1], parts[0]) // "owner/workspace.agent"
+		}
+		return fmt.Sprintf("%s/%s.%s", parts[0], parts[1], parts[2]) // "owner/workspace.agent"
+	default:
+		return input // Fallback
+	}
+}