Skip to content

feat: add workspace SSH execution tool for AI SDK #18924

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jul 21, 2025
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
295 changes: 295 additions & 0 deletions codersdk/toolsdk/bash.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,295 @@
package toolsdk

import (
"context"
"errors"
"fmt"
"io"
"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 WorkspaceBashArgs struct {
Workspace string `json:"workspace"`
Command string `json:"command"`
}

type WorkspaceBashResult struct {
Output string `json:"output"`
ExitCode int `json:"exit_code"`
}

var WorkspaceBash = Tool[WorkspaceBashArgs, WorkspaceBashResult]{
Tool: aisdk.Tool{
Name: ToolNameWorkspaceBash,
Description: `Execute a bash command in a Coder workspace.

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 bash command to execute in the workspace.",
},
},
Required: []string{"workspace", "command"},
},
},
Handler: func(ctx context.Context, deps Deps, args WorkspaceBashArgs) (WorkspaceBashResult, error) {
if args.Workspace == "" {
return WorkspaceBashResult{}, xerrors.New("workspace name cannot be empty")
}
if args.Command == "" {
return WorkspaceBashResult{}, xerrors.New("command cannot be empty")
}

// Normalize workspace input to handle various formats
workspaceName := NormalizeWorkspaceInput(args.Workspace)

// Find workspace and agent
_, workspaceAgent, err := findWorkspaceAndAgent(ctx, deps.coderClient, workspaceName)
if err != nil {
return WorkspaceBashResult{}, 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 WorkspaceBashResult{}, 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 WorkspaceBashResult{}, xerrors.Errorf("failed to dial agent: %w", err)
}
defer conn.Close()

// Wait for connection to be reachable
if !conn.AwaitReachable(ctx) {
return WorkspaceBashResult{}, xerrors.New("agent connection not reachable")
}

// Create SSH client
sshClient, err := conn.SSHClient(ctx)
if err != nil {
return WorkspaceBashResult{}, xerrors.Errorf("failed to create SSH client: %w", err)
}
defer sshClient.Close()

// Create SSH session
session, err := sshClient.NewSession()
if err != nil {
return WorkspaceBashResult{}, 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 WorkspaceBashResult{
Output: outputStr,
ExitCode: exitErr.ExitStatus(),
}, nil
}
// For other errors, return exit code 1
return WorkspaceBashResult{
Output: outputStr,
ExitCode: 1,
}, nil
}

return WorkspaceBashResult{
Output: outputStr,
ExitCode: 0,
}, nil
},
}

// 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
if build.Job.CompletedAt == nil {
err := cliui.WorkspaceBuild(ctx, io.Discard, client, build.ID)
if err != nil {
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("failed to wait for build completion: %w", err)
}
}

// 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 the following input formats:
// - workspace → workspace
// - workspace.agent → workspace.agent
// - owner/workspace → owner/workspace
// - owner--workspace → owner/workspace
// - owner/workspace.agent → owner/workspace.agent
// - owner--workspace.agent → owner/workspace.agent
// - agent.workspace.owner → owner/workspace.agent (Coder Connect format)
func NormalizeWorkspaceInput(input string) string {
// Handle the special Coder Connect format: agent.workspace.owner
// This format uses only dots and has exactly 3 parts
if strings.Count(input, ".") == 2 && !strings.Contains(input, "/") && !strings.Contains(input, "--") {
parts := strings.Split(input, ".")
if len(parts) == 3 {
// Convert agent.workspace.owner → owner/workspace.agent
return fmt.Sprintf("%s/%s.%s", parts[2], parts[1], parts[0])
}
}

// Convert -- separator to / separator for consistency
normalized := strings.ReplaceAll(input, "--", "/")

return normalized
}
Loading
Loading