Skip to content

Commit e976eea

Browse files
ThomasK33ThomasK33
committed
feat(toolsdk): add SSH exec tool
Change-Id: I61f694a89e33c60ab6e5a68b6773755bff1840a4 Signed-off-by: Thomas Kosiewski <tk@coder.com>
1 parent 7b06fc7 commit e976eea

File tree

4 files changed

+488
-0
lines changed

4 files changed

+488
-0
lines changed

‎codersdk/toolsdk/ssh.go

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

0 commit comments

Comments
 (0)