|
| 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 | +typeWorkspaceBashArgsstruct { |
| 21 | +Workspacestring`json:"workspace"` |
| 22 | +Commandstring`json:"command"` |
| 23 | +} |
| 24 | + |
| 25 | +typeWorkspaceBashResultstruct { |
| 26 | +Outputstring`json:"output"` |
| 27 | +ExitCodeint`json:"exit_code"` |
| 28 | +} |
| 29 | + |
| 30 | +varWorkspaceBash=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,depsDeps,argsWorkspaceBashArgs) (WorkspaceBashResult,error) { |
| 65 | +ifargs.Workspace=="" { |
| 66 | +returnWorkspaceBashResult{},xerrors.New("workspace name cannot be empty") |
| 67 | +} |
| 68 | +ifargs.Command=="" { |
| 69 | +returnWorkspaceBashResult{},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 | +iferr!=nil { |
| 78 | +returnWorkspaceBashResult{},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 | +iferr!=nil { |
| 89 | +returnWorkspaceBashResult{},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 | +iferr!=nil { |
| 100 | +returnWorkspaceBashResult{},xerrors.Errorf("failed to dial agent: %w",err) |
| 101 | +} |
| 102 | +deferconn.Close() |
| 103 | + |
| 104 | +// Wait for connection to be reachable |
| 105 | +if!conn.AwaitReachable(ctx) { |
| 106 | +returnWorkspaceBashResult{},xerrors.New("agent connection not reachable") |
| 107 | +} |
| 108 | + |
| 109 | +// Create SSH client |
| 110 | +sshClient,err:=conn.SSHClient(ctx) |
| 111 | +iferr!=nil { |
| 112 | +returnWorkspaceBashResult{},xerrors.Errorf("failed to create SSH client: %w",err) |
| 113 | +} |
| 114 | +defersshClient.Close() |
| 115 | + |
| 116 | +// Create SSH session |
| 117 | +session,err:=sshClient.NewSession() |
| 118 | +iferr!=nil { |
| 119 | +returnWorkspaceBashResult{},xerrors.Errorf("failed to create SSH session: %w",err) |
| 120 | +} |
| 121 | +defersession.Close() |
| 122 | + |
| 123 | +// Execute command and capture output |
| 124 | +output,err:=session.CombinedOutput(args.Command) |
| 125 | +outputStr:=strings.TrimSpace(string(output)) |
| 126 | + |
| 127 | +iferr!=nil { |
| 128 | +// Check if it's an SSH exit error to get the exit code |
| 129 | +varexitErr*gossh.ExitError |
| 130 | +iferrors.As(err,&exitErr) { |
| 131 | +returnWorkspaceBashResult{ |
| 132 | +Output:outputStr, |
| 133 | +ExitCode:exitErr.ExitStatus(), |
| 134 | +},nil |
| 135 | +} |
| 136 | +// For other errors, return exit code 1 |
| 137 | +returnWorkspaceBashResult{ |
| 138 | +Output:outputStr, |
| 139 | +ExitCode:1, |
| 140 | +},nil |
| 141 | +} |
| 142 | + |
| 143 | +returnWorkspaceBashResult{ |
| 144 | +Output:outputStr, |
| 145 | +ExitCode:0, |
| 146 | +},nil |
| 147 | +}, |
| 148 | +} |
| 149 | + |
| 150 | +// findWorkspaceAndAgent finds workspace and agent by name with auto-start support |
| 151 | +funcfindWorkspaceAndAgent(ctx context.Context,client*codersdk.Client,workspaceNamestring) (codersdk.Workspace, codersdk.WorkspaceAgent,error) { |
| 152 | +// Parse workspace name to extract workspace and agent parts |
| 153 | +parts:=strings.Split(workspaceName,".") |
| 154 | +varagentNamestring |
| 155 | +iflen(parts)>=2 { |
| 156 | +agentName=parts[1] |
| 157 | +workspaceName=parts[0] |
| 158 | +} |
| 159 | + |
| 160 | +// Get workspace |
| 161 | +workspace,err:=namedWorkspace(ctx,client,workspaceName) |
| 162 | +iferr!=nil { |
| 163 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},err |
| 164 | +} |
| 165 | + |
| 166 | +// Auto-start workspace if needed |
| 167 | +ifworkspace.LatestBuild.Transition!=codersdk.WorkspaceTransitionStart { |
| 168 | +ifworkspace.LatestBuild.Transition==codersdk.WorkspaceTransitionDelete { |
| 169 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},xerrors.Errorf("workspace %q is deleted",workspace.Name) |
| 170 | +} |
| 171 | +ifworkspace.LatestBuild.Job.Status==codersdk.ProvisionerJobFailed { |
| 172 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},xerrors.Errorf("workspace %q is in failed state",workspace.Name) |
| 173 | +} |
| 174 | +ifworkspace.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 | +iferr!=nil { |
| 184 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},xerrors.Errorf("failed to start workspace: %w",err) |
| 185 | +} |
| 186 | + |
| 187 | +// Wait for build to complete |
| 188 | +ifbuild.Job.CompletedAt==nil { |
| 189 | +err:=cliui.WorkspaceBuild(ctx,io.Discard,client,build.ID) |
| 190 | +iferr!=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 | +iferr!=nil { |
| 198 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},err |
| 199 | +} |
| 200 | +} |
| 201 | + |
| 202 | +// Find agent |
| 203 | +workspaceAgent,err:=getWorkspaceAgent(workspace,agentName) |
| 204 | +iferr!=nil { |
| 205 | +return codersdk.Workspace{}, codersdk.WorkspaceAgent{},err |
| 206 | +} |
| 207 | + |
| 208 | +returnworkspace,workspaceAgent,nil |
| 209 | +} |
| 210 | + |
| 211 | +// getWorkspaceAgent finds the specified agent in the workspace |
| 212 | +funcgetWorkspaceAgent(workspace codersdk.Workspace,agentNamestring) (codersdk.WorkspaceAgent,error) { |
| 213 | +resources:=workspace.LatestBuild.Resources |
| 214 | + |
| 215 | +varagents []codersdk.WorkspaceAgent |
| 216 | +varavailableNames []string |
| 217 | + |
| 218 | +for_,resource:=rangeresources { |
| 219 | +for_,agent:=rangeresource.Agents { |
| 220 | +availableNames=append(availableNames,agent.Name) |
| 221 | +agents=append(agents,agent) |
| 222 | +} |
| 223 | +} |
| 224 | + |
| 225 | +iflen(agents)==0 { |
| 226 | +return codersdk.WorkspaceAgent{},xerrors.Errorf("workspace %q has no agents",workspace.Name) |
| 227 | +} |
| 228 | + |
| 229 | +ifagentName!="" { |
| 230 | +for_,agent:=rangeagents { |
| 231 | +ifagent.Name==agentName||agent.ID.String()==agentName { |
| 232 | +returnagent,nil |
| 233 | +} |
| 234 | +} |
| 235 | +return codersdk.WorkspaceAgent{},xerrors.Errorf("agent not found by name %q, available agents: %v",agentName,availableNames) |
| 236 | +} |
| 237 | + |
| 238 | +iflen(agents)==1 { |
| 239 | +returnagents[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 | +funcnamedWorkspace(ctx context.Context,client*codersdk.Client,identifierstring) (codersdk.Workspace,error) { |
| 247 | +// Parse owner and workspace name |
| 248 | +parts:=strings.SplitN(identifier,"/",2) |
| 249 | +varowner,workspaceNamestring |
| 250 | + |
| 251 | +iflen(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 | +ifstrings.Contains(identifier,"--")&&!strings.Contains(identifier,"/") { |
| 261 | +dashParts:=strings.SplitN(identifier,"--",2) |
| 262 | +iflen(dashParts)==2 { |
| 263 | +owner=dashParts[0] |
| 264 | +workspaceName=dashParts[1] |
| 265 | +} |
| 266 | +} |
| 267 | + |
| 268 | +returnclient.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 | +funcNormalizeWorkspaceInput(inputstring)string { |
| 281 | +// Handle the special Coder Connect format: agent.workspace.owner |
| 282 | +// This format uses only dots and has exactly 3 parts |
| 283 | +ifstrings.Count(input,".")==2&&!strings.Contains(input,"/")&&!strings.Contains(input,"--") { |
| 284 | +parts:=strings.Split(input,".") |
| 285 | +iflen(parts)==3 { |
| 286 | +// Convert agent.workspace.owner → owner/workspace.agent |
| 287 | +returnfmt.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 | +returnnormalized |
| 295 | +} |