Skip to content

Commit 326c024

Browse files
authored
feat: add workspace SSH execution tool for AI SDK (#18924)
# Add SSH Command Execution Tool for Coder Workspaces This PR adds a new AI tool `coder_workspace_ssh_exec` that allows executing commands in Coder workspaces via SSH. The tool provides functionality similar to the `coder ssh <workspace> <command>` CLI command. Key features: - Executes commands in workspaces via SSH and returns the output and exit code - Automatically starts workspaces if they're stopped - Waits for the agent to be ready before executing commands - Trims leading and trailing whitespace from command output - Supports various workspace identifier formats: - `workspace` (uses current user) - `owner/workspace` - `owner--workspace` - `workspace.agent` (specific agent) - `owner/workspace.agent` The implementation includes: - A new tool definition with schema and handler - Helper functions for workspace and agent discovery - Workspace name normalization to handle different input formats - Comprehensive test coverage including integration tests This tool enables AI assistants to execute commands in user workspaces, making it possible to automate tasks and provide more interactive assistance. <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Introduced the ability to execute bash commands inside a Coder workspace via SSH, supporting multiple workspace identification formats. * **Tests** * Added comprehensive unit and integration tests for executing bash commands in workspaces, including input validation, output handling, and error scenarios. * **Chores** * Registered the new bash execution tool in the global tools list. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
1 parent 75c1240 commit 326c024

File tree

4 files changed

+533
-2
lines changed

4 files changed

+533
-2
lines changed

codersdk/toolsdk/bash.go

Lines changed: 295 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,295 @@
1+
package toolsdk
2+
3+
import (
4+
"context"
5+
"errors"
6+
"fmt"
7+
"io"
8+
"strings"
9+
10+
gossh "golang.org/x/crypto/ssh"
11+
"golang.org/x/xerrors"
12+
13+
"github.com/coder/aisdk-go"
14+
15+
"github.com/coder/coder/v2/cli/cliui"
16+
"github.com/coder/coder/v2/codersdk"
17+
"github.com/coder/coder/v2/codersdk/workspacesdk"
18+
)
19+
20+
type WorkspaceBashArgs struct {
21+
Workspace string `json:"workspace"`
22+
Command string `json:"command"`
23+
}
24+
25+
type WorkspaceBashResult struct {
26+
Output string `json:"output"`
27+
ExitCode int `json:"exit_code"`
28+
}
29+
30+
var WorkspaceBash = Tool[WorkspaceBashArgs, WorkspaceBashResult]{
31+
Tool: aisdk.Tool{
32+
Name: ToolNameWorkspaceBash,
33+
Description: `Execute a bash command in a Coder workspace.
34+
35+
This tool provides the same functionality as the 'coder ssh <workspace> <command>' CLI command.
36+
It automatically starts the workspace if it's stopped and waits for the agent to be ready.
37+
The output is trimmed of leading and trailing whitespace.
38+
39+
The workspace parameter supports various formats:
40+
- workspace (uses current user)
41+
- owner/workspace
42+
- owner--workspace
43+
- workspace.agent (specific agent)
44+
- owner/workspace.agent
45+
46+
Examples:
47+
- workspace: "my-workspace", command: "ls -la"
48+
- workspace: "john/dev-env", command: "git status"
49+
- workspace: "my-workspace.main", command: "docker ps"`,
50+
Schema: aisdk.Schema{
51+
Properties: map[string]any{
52+
"workspace": map[string]any{
53+
"type": "string",
54+
"description": "The workspace name in format [owner/]workspace[.agent]. If owner is not specified, the authenticated user is used.",
55+
},
56+
"command": map[string]any{
57+
"type": "string",
58+
"description": "The bash command to execute in the workspace.",
59+
},
60+
},
61+
Required: []string{"workspace", "command"},
62+
},
63+
},
64+
Handler: func(ctx context.Context, deps Deps, args WorkspaceBashArgs) (WorkspaceBashResult, error) {
65+
if args.Workspace == "" {
66+
return WorkspaceBashResult{}, xerrors.New("workspace name cannot be empty")
67+
}
68+
if args.Command == "" {
69+
return WorkspaceBashResult{}, xerrors.New("command cannot be empty")
70+
}
71+
72+
// Normalize workspace input to handle various formats
73+
workspaceName := NormalizeWorkspaceInput(args.Workspace)
74+
75+
// Find workspace and agent
76+
_, workspaceAgent, err := findWorkspaceAndAgent(ctx, deps.coderClient, workspaceName)
77+
if err != nil {
78+
return WorkspaceBashResult{}, xerrors.Errorf("failed to find workspace: %w", err)
79+
}
80+
81+
// Wait for agent to be ready
82+
err = cliui.Agent(ctx, nil, workspaceAgent.ID, cliui.AgentOptions{
83+
FetchInterval: 0,
84+
Fetch: deps.coderClient.WorkspaceAgent,
85+
FetchLogs: deps.coderClient.WorkspaceAgentLogsAfter,
86+
Wait: true, // Always wait for startup scripts
87+
})
88+
if err != nil {
89+
return WorkspaceBashResult{}, xerrors.Errorf("agent not ready: %w", err)
90+
}
91+
92+
// Create workspace SDK client for agent connection
93+
wsClient := workspacesdk.New(deps.coderClient)
94+
95+
// Dial agent
96+
conn, err := wsClient.DialAgent(ctx, workspaceAgent.ID, &workspacesdk.DialAgentOptions{
97+
BlockEndpoints: false,
98+
})
99+
if err != nil {
100+
return WorkspaceBashResult{}, xerrors.Errorf("failed to dial agent: %w", err)
101+
}
102+
defer conn.Close()
103+
104+
// Wait for connection to be reachable
105+
if !conn.AwaitReachable(ctx) {
106+
return WorkspaceBashResult{}, xerrors.New("agent connection not reachable")
107+
}
108+
109+
// Create SSH client
110+
sshClient, err := conn.SSHClient(ctx)
111+
if err != nil {
112+
return WorkspaceBashResult{}, xerrors.Errorf("failed to create SSH client: %w", err)
113+
}
114+
defer sshClient.Close()
115+
116+
// Create SSH session
117+
session, err := sshClient.NewSession()
118+
if err != nil {
119+
return WorkspaceBashResult{}, xerrors.Errorf("failed to create SSH session: %w", err)
120+
}
121+
defer session.Close()
122+
123+
// Execute command and capture output
124+
output, err := session.CombinedOutput(args.Command)
125+
outputStr := strings.TrimSpace(string(output))
126+
127+
if err != nil {
128+
// Check if it's an SSH exit error to get the exit code
129+
var exitErr *gossh.ExitError
130+
if errors.As(err, &exitErr) {
131+
return WorkspaceBashResult{
132+
Output: outputStr,
133+
ExitCode: exitErr.ExitStatus(),
134+
}, nil
135+
}
136+
// For other errors, return exit code 1
137+
return WorkspaceBashResult{
138+
Output: outputStr,
139+
ExitCode: 1,
140+
}, nil
141+
}
142+
143+
return WorkspaceBashResult{
144+
Output: outputStr,
145+
ExitCode: 0,
146+
}, nil
147+
},
148+
}
149+
150+
// findWorkspaceAndAgent finds workspace and agent by name with auto-start support
151+
func findWorkspaceAndAgent(ctx context.Context, client *codersdk.Client, workspaceName string) (codersdk.Workspace, codersdk.WorkspaceAgent, error) {
152+
// Parse workspace name to extract workspace and agent parts
153+
parts := strings.Split(workspaceName, ".")
154+
var agentName string
155+
if len(parts) >= 2 {
156+
agentName = parts[1]
157+
workspaceName = parts[0]
158+
}
159+
160+
// Get workspace
161+
workspace, err := namedWorkspace(ctx, client, workspaceName)
162+
if err != nil {
163+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
164+
}
165+
166+
// Auto-start workspace if needed
167+
if workspace.LatestBuild.Transition != codersdk.WorkspaceTransitionStart {
168+
if workspace.LatestBuild.Transition == codersdk.WorkspaceTransitionDelete {
169+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q is deleted", workspace.Name)
170+
}
171+
if workspace.LatestBuild.Job.Status == codersdk.ProvisionerJobFailed {
172+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q is in failed state", workspace.Name)
173+
}
174+
if workspace.LatestBuild.Status != codersdk.WorkspaceStatusStopped {
175+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace must be started; was unable to autostart as the last build job is %q, expected %q",
176+
workspace.LatestBuild.Status, codersdk.WorkspaceStatusStopped)
177+
}
178+
179+
// Start workspace
180+
build, err := client.CreateWorkspaceBuild(ctx, workspace.ID, codersdk.CreateWorkspaceBuildRequest{
181+
Transition: codersdk.WorkspaceTransitionStart,
182+
})
183+
if err != nil {
184+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("failed to start workspace: %w", err)
185+
}
186+
187+
// Wait for build to complete
188+
if build.Job.CompletedAt == nil {
189+
err := cliui.WorkspaceBuild(ctx, io.Discard, client, build.ID)
190+
if err != nil {
191+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("failed to wait for build completion: %w", err)
192+
}
193+
}
194+
195+
// Refresh workspace after build completes
196+
workspace, err = client.Workspace(ctx, workspace.ID)
197+
if err != nil {
198+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
199+
}
200+
}
201+
202+
// Find agent
203+
workspaceAgent, err := getWorkspaceAgent(workspace, agentName)
204+
if err != nil {
205+
return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, err
206+
}
207+
208+
return workspace, workspaceAgent, nil
209+
}
210+
211+
// getWorkspaceAgent finds the specified agent in the workspace
212+
func getWorkspaceAgent(workspace codersdk.Workspace, agentName string) (codersdk.WorkspaceAgent, error) {
213+
resources := workspace.LatestBuild.Resources
214+
215+
var agents []codersdk.WorkspaceAgent
216+
var availableNames []string
217+
218+
for _, resource := range resources {
219+
for _, agent := range resource.Agents {
220+
availableNames = append(availableNames, agent.Name)
221+
agents = append(agents, agent)
222+
}
223+
}
224+
225+
if len(agents) == 0 {
226+
return codersdk.WorkspaceAgent{}, xerrors.Errorf("workspace %q has no agents", workspace.Name)
227+
}
228+
229+
if agentName != "" {
230+
for _, agent := range agents {
231+
if agent.Name == agentName || agent.ID.String() == agentName {
232+
return agent, nil
233+
}
234+
}
235+
return codersdk.WorkspaceAgent{}, xerrors.Errorf("agent not found by name %q, available agents: %v", agentName, availableNames)
236+
}
237+
238+
if len(agents) == 1 {
239+
return agents[0], nil
240+
}
241+
242+
return codersdk.WorkspaceAgent{}, xerrors.Errorf("multiple agents found, please specify the agent name, available agents: %v", availableNames)
243+
}
244+
245+
// namedWorkspace gets a workspace by owner/name or just name
246+
func namedWorkspace(ctx context.Context, client *codersdk.Client, identifier string) (codersdk.Workspace, error) {
247+
// Parse owner and workspace name
248+
parts := strings.SplitN(identifier, "/", 2)
249+
var owner, workspaceName string
250+
251+
if len(parts) == 2 {
252+
owner = parts[0]
253+
workspaceName = parts[1]
254+
} else {
255+
owner = "me"
256+
workspaceName = identifier
257+
}
258+
259+
// Handle -- separator format (convert to / format)
260+
if strings.Contains(identifier, "--") && !strings.Contains(identifier, "/") {
261+
dashParts := strings.SplitN(identifier, "--", 2)
262+
if len(dashParts) == 2 {
263+
owner = dashParts[0]
264+
workspaceName = dashParts[1]
265+
}
266+
}
267+
268+
return client.WorkspaceByOwnerAndName(ctx, owner, workspaceName, codersdk.WorkspaceOptions{})
269+
}
270+
271+
// NormalizeWorkspaceInput converts workspace name input to standard format.
272+
// Handles the following input formats:
273+
// - workspace → workspace
274+
// - workspace.agent → workspace.agent
275+
// - owner/workspace → owner/workspace
276+
// - owner--workspace → owner/workspace
277+
// - owner/workspace.agent → owner/workspace.agent
278+
// - owner--workspace.agent → owner/workspace.agent
279+
// - agent.workspace.owner → owner/workspace.agent (Coder Connect format)
280+
func NormalizeWorkspaceInput(input string) string {
281+
// Handle the special Coder Connect format: agent.workspace.owner
282+
// This format uses only dots and has exactly 3 parts
283+
if strings.Count(input, ".") == 2 && !strings.Contains(input, "/") && !strings.Contains(input, "--") {
284+
parts := strings.Split(input, ".")
285+
if len(parts) == 3 {
286+
// Convert agent.workspace.owner → owner/workspace.agent
287+
return fmt.Sprintf("%s/%s.%s", parts[2], parts[1], parts[0])
288+
}
289+
}
290+
291+
// Convert -- separator to / separator for consistency
292+
normalized := strings.ReplaceAll(input, "--", "/")
293+
294+
return normalized
295+
}

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy