- Notifications
You must be signed in to change notification settings - Fork948
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 ourterms of service andprivacy statement. We’ll occasionally send you account related emails.
Already on GitHub?Sign in to your account
Draft
ThomasK33 wants to merge1 commit intomainChoose a base branch
base:main
Could not load branches
Branch not found:{{ refName }}
Loading
Could not load tags
Nothing to show
Loading
Are you sure you want to change the base?
Some commits from the old base branch may be removed from the timeline, and old review comments may become outdated.
+488 −0
Draft
Changes fromall commits
Commits
File filter
Filter by extension
Conversations
Failed to load comments.
Loading
Uh oh!
There was an error while loading.Please reload this page.
Jump to
Jump to file
Failed to load files.
Loading
Uh oh!
There was an error while loading.Please reload this page.
Diff view
Diff view
There are no files selected for viewing
307 changes: 307 additions & 0 deletionscodersdk/toolsdk/ssh.go
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,307 @@ | ||
| package toolsdk | ||
| import ( | ||
| "context" | ||
| "errors" | ||
| "fmt" | ||
| "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 WorkspaceSSHExecArgs struct { | ||
| Workspace string `json:"workspace"` | ||
| Command string `json:"command"` | ||
| } | ||
| type WorkspaceSSHExecResult struct { | ||
| Output string `json:"output"` | ||
| ExitCode int `json:"exit_code"` | ||
| } | ||
| var WorkspaceSSHExec = Tool[WorkspaceSSHExecArgs, WorkspaceSSHExecResult]{ | ||
| Tool: aisdk.Tool{ | ||
| Name: ToolNameWorkspaceSSHExec, | ||
| Description: `Execute a command in a Coder workspace via SSH. | ||
| 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 command to execute in the workspace.", | ||
| }, | ||
| }, | ||
| Required: []string{"workspace", "command"}, | ||
| }, | ||
| }, | ||
| Handler: func(ctx context.Context, deps Deps, args WorkspaceSSHExecArgs) (WorkspaceSSHExecResult, error) { | ||
| if args.Workspace == "" { | ||
| return WorkspaceSSHExecResult{}, xerrors.New("workspace name cannot be empty") | ||
| } | ||
| if args.Command == "" { | ||
| return WorkspaceSSHExecResult{}, xerrors.New("command cannot be empty") | ||
| } | ||
| // Normalize workspace input to handle various formats | ||
| workspaceName := NormalizeWorkspaceInput(args.Workspace) | ||
| // Find workspace and agent | ||
| _, workspaceAgent, err := findWorkspaceAndAgentWithAutostart(ctx, deps.coderClient, workspaceName) | ||
| if err != nil { | ||
| return WorkspaceSSHExecResult{}, 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 WorkspaceSSHExecResult{}, 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 WorkspaceSSHExecResult{}, xerrors.Errorf("failed to dial agent: %w", err) | ||
| } | ||
| defer conn.Close() | ||
| // Wait for connection to be reachable | ||
| conn.AwaitReachable(ctx) | ||
| // Create SSH client | ||
| sshClient, err := conn.SSHClient(ctx) | ||
| if err != nil { | ||
| return WorkspaceSSHExecResult{}, xerrors.Errorf("failed to create SSH client: %w", err) | ||
| } | ||
| defer sshClient.Close() | ||
| // Create SSH session | ||
| session, err := sshClient.NewSession() | ||
| if err != nil { | ||
| return WorkspaceSSHExecResult{}, 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 WorkspaceSSHExecResult{ | ||
| Output: outputStr, | ||
| ExitCode: exitErr.ExitStatus(), | ||
| }, nil | ||
| } | ||
| // For other errors, return exit code 1 | ||
| return WorkspaceSSHExecResult{ | ||
| Output: outputStr, | ||
| ExitCode: 1, | ||
| }, nil | ||
| } | ||
| return WorkspaceSSHExecResult{ | ||
| Output: outputStr, | ||
| ExitCode: 0, | ||
| }, nil | ||
| }, | ||
| } | ||
| // findWorkspaceAndAgentWithAutostart finds workspace and agent by name and auto-starts if needed | ||
| func findWorkspaceAndAgentWithAutostart(ctx context.Context, client *codersdk.Client, workspaceName string) (codersdk.Workspace, codersdk.WorkspaceAgent, error) { | ||
| return findWorkspaceAndAgent(ctx, client, workspaceName) | ||
| } | ||
| // 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 | ||
| for { | ||
| build, err = client.WorkspaceBuild(ctx, build.ID) | ||
| if err != nil { | ||
| return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, xerrors.Errorf("failed to get build status: %w", err) | ||
| } | ||
| if build.Job.CompletedAt != nil { | ||
| break | ||
| } | ||
| // Small delay before checking again | ||
| select { | ||
| case <-ctx.Done(): | ||
| return codersdk.Workspace{}, codersdk.WorkspaceAgent{}, ctx.Err() | ||
| default: | ||
| } | ||
| } | ||
| // 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 formats like: workspace, workspace.agent, owner/workspace, owner--workspace, etc. | ||
| func NormalizeWorkspaceInput(input string) string { | ||
| // This matches the logic from cli/ssh.go | ||
| // Split on "/", "--", and "." | ||
| workspaceNameRe := strings.NewReplacer("/", ":", "--", ":", ".", ":") | ||
| parts := strings.Split(workspaceNameRe.Replace(input), ":") | ||
| switch len(parts) { | ||
| case 1: | ||
| return input // "workspace" | ||
| case 2: | ||
| if strings.Contains(input, ".") { | ||
| return fmt.Sprintf("%s.%s", parts[0], parts[1]) // "workspace.agent" | ||
| } | ||
| return fmt.Sprintf("%s/%s", parts[0], parts[1]) // "owner/workspace" | ||
| case 3: | ||
| // If the only separator is a dot, it's the Coder Connect format | ||
| if !strings.Contains(input, "/") && !strings.Contains(input, "--") { | ||
| return fmt.Sprintf("%s/%s.%s", parts[2], parts[1], parts[0]) // "owner/workspace.agent" | ||
| } | ||
| return fmt.Sprintf("%s/%s.%s", parts[0], parts[1], parts[2]) // "owner/workspace.agent" | ||
| default: | ||
| return input // Fallback | ||
| } | ||
| } |
Oops, something went wrong.
Uh oh!
There was an error while loading.Please reload this page.
Oops, something went wrong.
Uh oh!
There was an error while loading.Please reload this page.
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.