|
| 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 | +typeWorkspaceSSHExecArgsstruct { |
| 20 | +Workspacestring`json:"workspace"` |
| 21 | +Commandstring`json:"command"` |
| 22 | +} |
| 23 | + |
| 24 | +typeWorkspaceSSHExecResultstruct { |
| 25 | +Outputstring`json:"output"` |
| 26 | +ExitCodeint`json:"exit_code"` |
| 27 | +} |
| 28 | + |
| 29 | +varWorkspaceSSHExec=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,depsDeps,argsWorkspaceSSHExecArgs) (WorkspaceSSHExecResult,error) { |
| 64 | +ifargs.Workspace=="" { |
| 65 | +returnWorkspaceSSHExecResult{},xerrors.New("workspace name cannot be empty") |
| 66 | +} |
| 67 | +ifargs.Command=="" { |
| 68 | +returnWorkspaceSSHExecResult{},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 | +iferr!=nil { |
| 77 | +returnWorkspaceSSHExecResult{},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 | +iferr!=nil { |
| 88 | +returnWorkspaceSSHExecResult{},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 | +iferr!=nil { |
| 99 | +returnWorkspaceSSHExecResult{},xerrors.Errorf("failed to dial agent: %w",err) |
| 100 | +} |
| 101 | +deferconn.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 | +iferr!=nil { |
| 109 | +returnWorkspaceSSHExecResult{},xerrors.Errorf("failed to create SSH client: %w",err) |
| 110 | +} |
| 111 | +defersshClient.Close() |
| 112 | + |
| 113 | +// Create SSH session |
| 114 | +session,err:=sshClient.NewSession() |
| 115 | +iferr!=nil { |
| 116 | +returnWorkspaceSSHExecResult{},xerrors.Errorf("failed to create SSH session: %w",err) |
| 117 | +} |
| 118 | +defersession.Close() |
| 119 | + |
| 120 | +// Execute command and capture output |
| 121 | +output,err:=session.CombinedOutput(args.Command) |
| 122 | +outputStr:=strings.TrimSpace(string(output)) |
| 123 | + |
| 124 | +iferr!=nil { |
| 125 | +// Check if it's an SSH exit error to get the exit code |
| 126 | +varexitErr*gossh.ExitError |
| 127 | +iferrors.As(err,&exitErr) { |
| 128 | +returnWorkspaceSSHExecResult{ |
| 129 | +Output:outputStr, |
| 130 | +ExitCode:exitErr.ExitStatus(), |
| 131 | +},nil |
| 132 | +} |
| 133 | +// For other errors, return exit code 1 |
| 134 | +returnWorkspaceSSHExecResult{ |
| 135 | +Output:outputStr, |
| 136 | +ExitCode:1, |
| 137 | +},nil |
| 138 | +} |
| 139 | + |
| 140 | +returnWorkspaceSSHExecResult{ |
| 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 | +funcfindWorkspaceAndAgentWithAutostart(ctx context.Context,client*codersdk.Client,workspaceNamestring) (codersdk.Workspace, codersdk.WorkspaceAgent,error) { |
| 149 | +returnfindWorkspaceAndAgent(ctx,client,workspaceName) |
| 150 | +} |
| 151 | + |
| 152 | +// findWorkspaceAndAgent finds workspace and agent by name with auto-start support |
| 153 | +funcfindWorkspaceAndAgent(ctx context.Context,client*codersdk.Client,workspaceNamestring) (codersdk.Workspace, codersdk.WorkspaceAgent,error) { |
| 154 | +// Parse workspace name to extract workspace and agent parts |
| 155 | +parts:=strings.Split(workspaceName,".") |
| 156 | +varagentNamestring |
| 157 | +iflen(parts)>=2 { |
| 158 | +agentName=parts[1] |
| 159 | +workspaceName=parts[0] |
| 160 | +} |
| 161 | + |
| 162 | +// Get workspace |
| 163 | +workspace,err:=namedWorkspace(ctx,client,workspaceName) |
| 164 | +iferr!=nil { |
| 165 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},err |
| 166 | +} |
| 167 | + |
| 168 | +// Auto-start workspace if needed |
| 169 | +ifworkspace.LatestBuild.Transition!=codersdk.WorkspaceTransitionStart { |
| 170 | +ifworkspace.LatestBuild.Transition==codersdk.WorkspaceTransitionDelete { |
| 171 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},xerrors.Errorf("workspace %q is deleted",workspace.Name) |
| 172 | +} |
| 173 | +ifworkspace.LatestBuild.Job.Status==codersdk.ProvisionerJobFailed { |
| 174 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},xerrors.Errorf("workspace %q is in failed state",workspace.Name) |
| 175 | +} |
| 176 | +ifworkspace.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 | +iferr!=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 | +iferr!=nil { |
| 193 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},xerrors.Errorf("failed to get build status: %w",err) |
| 194 | +} |
| 195 | +ifbuild.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 | +iferr!=nil { |
| 209 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},err |
| 210 | +} |
| 211 | +} |
| 212 | + |
| 213 | +// Find agent |
| 214 | +workspaceAgent,err:=getWorkspaceAgent(workspace,agentName) |
| 215 | +iferr!=nil { |
| 216 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},err |
| 217 | +} |
| 218 | + |
| 219 | +returnworkspace,workspaceAgent,nil |
| 220 | +} |
| 221 | + |
| 222 | +// getWorkspaceAgent finds the specified agent in the workspace |
| 223 | +funcgetWorkspaceAgent(workspace codersdk.Workspace,agentNamestring) (codersdk.WorkspaceAgent,error) { |
| 224 | +resources:=workspace.LatestBuild.Resources |
| 225 | + |
| 226 | +varagents []codersdk.WorkspaceAgent |
| 227 | +varavailableNames []string |
| 228 | + |
| 229 | +for_,resource:=rangeresources { |
| 230 | +for_,agent:=rangeresource.Agents { |
| 231 | +availableNames=append(availableNames,agent.Name) |
| 232 | +agents=append(agents,agent) |
| 233 | +} |
| 234 | +} |
| 235 | + |
| 236 | +iflen(agents)==0 { |
| 237 | +return codersdk.WorkspaceAgent{},xerrors.Errorf("workspace %q has no agents",workspace.Name) |
| 238 | +} |
| 239 | + |
| 240 | +ifagentName!="" { |
| 241 | +for_,agent:=rangeagents { |
| 242 | +ifagent.Name==agentName||agent.ID.String()==agentName { |
| 243 | +returnagent,nil |
| 244 | +} |
| 245 | +} |
| 246 | +return codersdk.WorkspaceAgent{},xerrors.Errorf("agent not found by name %q, available agents: %v",agentName,availableNames) |
| 247 | +} |
| 248 | + |
| 249 | +iflen(agents)==1 { |
| 250 | +returnagents[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 | +funcnamedWorkspace(ctx context.Context,client*codersdk.Client,identifierstring) (codersdk.Workspace,error) { |
| 258 | +// Parse owner and workspace name |
| 259 | +parts:=strings.SplitN(identifier,"/",2) |
| 260 | +varowner,workspaceNamestring |
| 261 | + |
| 262 | +iflen(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 | +ifstrings.Contains(identifier,"--")&&!strings.Contains(identifier,"/") { |
| 272 | +dashParts:=strings.SplitN(identifier,"--",2) |
| 273 | +iflen(dashParts)==2 { |
| 274 | +owner=dashParts[0] |
| 275 | +workspaceName=dashParts[1] |
| 276 | +} |
| 277 | +} |
| 278 | + |
| 279 | +returnclient.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 | +funcNormalizeWorkspaceInput(inputstring)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 | +switchlen(parts) { |
| 291 | +case1: |
| 292 | +returninput// "workspace" |
| 293 | +case2: |
| 294 | +ifstrings.Contains(input,".") { |
| 295 | +returnfmt.Sprintf("%s.%s",parts[0],parts[1])// "workspace.agent" |
| 296 | +} |
| 297 | +returnfmt.Sprintf("%s/%s",parts[0],parts[1])// "owner/workspace" |
| 298 | +case3: |
| 299 | +// If the only separator is a dot, it's the Coder Connect format |
| 300 | +if!strings.Contains(input,"/")&&!strings.Contains(input,"--") { |
| 301 | +returnfmt.Sprintf("%s/%s.%s",parts[2],parts[1],parts[0])// "owner/workspace.agent" |
| 302 | +} |
| 303 | +returnfmt.Sprintf("%s/%s.%s",parts[0],parts[1],parts[2])// "owner/workspace.agent" |
| 304 | +default: |
| 305 | +returninput// Fallback |
| 306 | +} |
| 307 | +} |