|
| 1 | +package cmd |
| 2 | + |
| 3 | +import ( |
| 4 | + "context" |
| 5 | + "fmt" |
| 6 | + "os" |
| 7 | + "path" |
| 8 | + "path/filepath" |
| 9 | + "strings" |
| 10 | + "time" |
| 11 | + |
| 12 | + "github.com/bitrise-io/ai-qa-agent-cli/internal/codespaces" |
| 13 | + codespacesv1 "github.com/bitrise-io/bitrise-codespaces/backend/proto/codespaces/v1" |
| 14 | + "github.com/spf13/cobra" |
| 15 | +) |
| 16 | + |
| 17 | +const remotePathPlaceholder = "{{REMOTE_PATH}}" |
| 18 | + |
| 19 | +var sessionCmd = &cobra.Command{ |
| 20 | + Use: "session", |
| 21 | + Short: "Manage RDE sessions", |
| 22 | +} |
| 23 | + |
| 24 | +var ( |
| 25 | + createWorkspace string |
| 26 | + createTemplate string |
| 27 | + createName string |
| 28 | + createDescription string |
| 29 | + createInputs []string |
| 30 | + createSecretInputs []string |
| 31 | + createSavedInputs []string |
| 32 | + createFeatureFlags []string |
| 33 | + createCluster string |
| 34 | + createAIPrompt string |
| 35 | + createAutoTerminateMinutes int32 |
| 36 | + createMapSavedInputs bool |
| 37 | + createWait bool |
| 38 | + createPollInterval time.Duration |
| 39 | + createOpenRemoteAccess bool |
| 40 | + createUpload string |
| 41 | + createUploadDestination string |
| 42 | +) |
| 43 | + |
| 44 | +var sessionCreateCmd = &cobra.Command{ |
| 45 | + Use: "create", |
| 46 | + Short: "Create a new RDE session and (optionally) wait for it to be running", |
| 47 | + RunE: runSessionCreate, |
| 48 | +} |
| 49 | + |
| 50 | +func init() { |
| 51 | + sessionCmd.AddCommand(sessionCreateCmd) |
| 52 | + |
| 53 | + f := sessionCreateCmd.Flags() |
| 54 | + f.StringVarP(&createWorkspace, "workspace", "w", "", "Workspace ID (required)") |
| 55 | + f.StringVarP(&createTemplate, "template", "t", "", "Template ID (required)") |
| 56 | + f.StringVar(&createName, "name", "", "Session name (required)") |
| 57 | + f.StringVar(&createDescription, "description", "", "Session description") |
| 58 | + f.StringArrayVar(&createInputs, "input", nil, "Session input as key=value (repeatable)") |
| 59 | + f.StringArrayVar(&createSecretInputs, "secret-input", nil, "Secret session input as key=value (repeatable)") |
| 60 | + f.StringArrayVar(&createSavedInputs, "saved-input", nil, "Saved input reference as key=savedInputID (repeatable)") |
| 61 | + f.StringArrayVar(&createFeatureFlags, "feature-flag", nil, "Feature flag name to enable (repeatable)") |
| 62 | + f.StringVar(&createCluster, "cluster", "", "Target cluster name (only required when image+machine-type matches multiple clusters)") |
| 63 | + f.StringVar(&createAIPrompt, "ai-prompt", "", "AI prompt to pass to Claude Code on session start. "+ |
| 64 | + "Any "+remotePathPlaceholder+" is substituted with the remote path of --upload. "+ |
| 65 | + "Note: the binary is uploaded AFTER the session reaches RUNNING, so phrase the prompt to wait for the file "+ |
| 66 | + "(e.g. 'Wait until "+remotePathPlaceholder+" exists, then run it and report output').") |
| 67 | + f.StringVar(&createUpload, "upload", "", "Local file to upload to the session after it reaches RUNNING") |
| 68 | + f.StringVar(&createUploadDestination, "upload-destination", "/tmp", "Absolute remote directory the --upload file is extracted into") |
| 69 | + f.Int32Var(&createAutoTerminateMinutes, "auto-terminate-minutes", -1, "Minutes before auto-termination (0 disables; -1 leaves backend default)") |
| 70 | + f.BoolVar(&createMapSavedInputs, "map-saved-inputs", false, "Auto-fill template session inputs from caller's saved inputs") |
| 71 | + f.BoolVar(&createWait, "wait", true, "Poll until session reaches RUNNING") |
| 72 | + f.DurationVar(&createPollInterval, "poll-interval", 5*time.Second, "Status poll interval when --wait is set") |
| 73 | + f.BoolVar(&createOpenRemoteAccess, "open-remote-access", false, "After RUNNING, call OpenRemoteAccess and print SSH/VNC details") |
| 74 | + |
| 75 | + _ = sessionCreateCmd.MarkFlagRequired("workspace") |
| 76 | + _ = sessionCreateCmd.MarkFlagRequired("template") |
| 77 | + _ = sessionCreateCmd.MarkFlagRequired("name") |
| 78 | +} |
| 79 | + |
| 80 | +func runSessionCreate(cmd *cobra.Command, _ []string) error { |
| 81 | + pat := os.Getenv(envPAT) |
| 82 | + if pat == "" { |
| 83 | + return fmt.Errorf("%s not set", envPAT) |
| 84 | + } |
| 85 | + |
| 86 | + aiPrompt, remotePath, err := resolveUploadAndPrompt(createUpload, createUploadDestination, createAIPrompt) |
| 87 | + if err != nil { |
| 88 | + return err |
| 89 | + } |
| 90 | + |
| 91 | + inputs, err := buildSessionInputs(createInputs, createSecretInputs, createSavedInputs) |
| 92 | + if err != nil { |
| 93 | + return err |
| 94 | + } |
| 95 | + |
| 96 | + req := &codespacesv1.CreateSessionRequest{ |
| 97 | + Name: createName, |
| 98 | + Description: createDescription, |
| 99 | + TemplateId: createTemplate, |
| 100 | + WorkspaceId: createWorkspace, |
| 101 | + SessionInputs: inputs, |
| 102 | + EnabledFeatureFlagNames: createFeatureFlags, |
| 103 | + Cluster: createCluster, |
| 104 | + AiPrompt: aiPrompt, |
| 105 | + MapSavedToSessionInputs: createMapSavedInputs, |
| 106 | + } |
| 107 | + if createAutoTerminateMinutes >= 0 { |
| 108 | + v := createAutoTerminateMinutes |
| 109 | + req.AutoTerminateMinutes = &v |
| 110 | + } |
| 111 | + |
| 112 | + ctx, cancel := context.WithTimeout(cmd.Context(), flagTimeout) |
| 113 | + defer cancel() |
| 114 | + |
| 115 | + client, err := codespaces.NewClient(flagEndpoint, pat, flagInsecure) |
| 116 | + if err != nil { |
| 117 | + return err |
| 118 | + } |
| 119 | + defer client.Close() |
| 120 | + |
| 121 | + session, err := client.CreateSession(ctx, req) |
| 122 | + if err != nil { |
| 123 | + return fmt.Errorf("CreateSession: %w", err) |
| 124 | + } |
| 125 | + fmt.Fprintf(os.Stderr, "created session %s (status: %s)\n", session.GetId(), session.GetStatus()) |
| 126 | + |
| 127 | + if createWait { |
| 128 | + session, err = client.WaitForRunning(ctx, session.GetId(), createWorkspace, createPollInterval, func(s codespacesv1.SessionStatus) { |
| 129 | + fmt.Fprintf(os.Stderr, " status: %s\n", s) |
| 130 | + }) |
| 131 | + if err != nil { |
| 132 | + return err |
| 133 | + } |
| 134 | + } |
| 135 | + |
| 136 | + if createUpload != "" && session.GetStatus() == codespacesv1.SessionStatus_SESSION_STATUS_RUNNING { |
| 137 | + actualPath, err := client.UploadFile(ctx, session.GetId(), createWorkspace, createUpload, createUploadDestination) |
| 138 | + if err != nil { |
| 139 | + return fmt.Errorf("upload %s: %w", createUpload, err) |
| 140 | + } |
| 141 | + fmt.Fprintf(os.Stderr, "uploaded %s -> %s\n", createUpload, actualPath) |
| 142 | + _ = remotePath // resolved earlier for the prompt; logged here from the server's confirmed path |
| 143 | + } |
| 144 | + |
| 145 | + if createOpenRemoteAccess && session.GetStatus() == codespacesv1.SessionStatus_SESSION_STATUS_RUNNING { |
| 146 | + session, err = client.OpenRemoteAccess(ctx, session.GetId(), createWorkspace) |
| 147 | + if err != nil { |
| 148 | + return fmt.Errorf("OpenRemoteAccess: %w", err) |
| 149 | + } |
| 150 | + fmt.Fprintf(os.Stderr, "ssh: %s (password: %s)\n", session.GetSshAddress(), session.GetSshPassword()) |
| 151 | + fmt.Fprintf(os.Stderr, "vnc: %s (user: %s, password: %s)\n", session.GetVncAddress(), session.GetVncUsername(), session.GetVncPassword()) |
| 152 | + } |
| 153 | + |
| 154 | + fmt.Println(session.GetId()) |
| 155 | + return nil |
| 156 | +} |
| 157 | + |
| 158 | +// resolveUploadAndPrompt validates the upload flags against the prompt placeholder |
| 159 | +// and returns the (possibly substituted) prompt plus the resolved remote path. |
| 160 | +// remotePath is empty when --upload is not set. |
| 161 | +func resolveUploadAndPrompt(uploadLocal, uploadDest, prompt string) (string, string, error) { |
| 162 | + hasPlaceholder := strings.Contains(prompt, remotePathPlaceholder) |
| 163 | + |
| 164 | + if uploadLocal == "" { |
| 165 | + if hasPlaceholder { |
| 166 | + return "", "", fmt.Errorf("--ai-prompt contains %s but --upload is not set", remotePathPlaceholder) |
| 167 | + } |
| 168 | + return prompt, "", nil |
| 169 | + } |
| 170 | + |
| 171 | + if !path.IsAbs(uploadDest) { |
| 172 | + return "", "", fmt.Errorf("--upload-destination must be absolute, got %q", uploadDest) |
| 173 | + } |
| 174 | + |
| 175 | + stat, err := os.Stat(uploadLocal) |
| 176 | + if err != nil { |
| 177 | + return "", "", fmt.Errorf("--upload %s: %w", uploadLocal, err) |
| 178 | + } |
| 179 | + if stat.IsDir() { |
| 180 | + return "", "", fmt.Errorf("--upload %s: must be a file, not a directory", uploadLocal) |
| 181 | + } |
| 182 | + |
| 183 | + remote := path.Join(uploadDest, filepath.Base(uploadLocal)) |
| 184 | + if hasPlaceholder { |
| 185 | + prompt = strings.ReplaceAll(prompt, remotePathPlaceholder, remote) |
| 186 | + } else if prompt != "" { |
| 187 | + fmt.Fprintf(os.Stderr, "warning: --ai-prompt does not reference %s; ensure the prompt knows the file's path (%s)\n", remotePathPlaceholder, remote) |
| 188 | + } |
| 189 | + return prompt, remote, nil |
| 190 | +} |
| 191 | + |
| 192 | +func buildSessionInputs(plain, secret, saved []string) ([]*codespacesv1.SessionInputValue, error) { |
| 193 | + out := make([]*codespacesv1.SessionInputValue, 0, len(plain)+len(secret)+len(saved)) |
| 194 | + |
| 195 | + for _, kv := range plain { |
| 196 | + k, v, ok := strings.Cut(kv, "=") |
| 197 | + if !ok { |
| 198 | + return nil, fmt.Errorf("--input %q: expected key=value", kv) |
| 199 | + } |
| 200 | + out = append(out, &codespacesv1.SessionInputValue{Key: k, Value: v}) |
| 201 | + } |
| 202 | + for _, kv := range secret { |
| 203 | + k, v, ok := strings.Cut(kv, "=") |
| 204 | + if !ok { |
| 205 | + return nil, fmt.Errorf("--secret-input %q: expected key=value", kv) |
| 206 | + } |
| 207 | + out = append(out, &codespacesv1.SessionInputValue{Key: k, Value: v, IsSecret: true}) |
| 208 | + } |
| 209 | + for _, kv := range saved { |
| 210 | + k, id, ok := strings.Cut(kv, "=") |
| 211 | + if !ok { |
| 212 | + return nil, fmt.Errorf("--saved-input %q: expected key=savedInputID", kv) |
| 213 | + } |
| 214 | + out = append(out, &codespacesv1.SessionInputValue{Key: k, SavedInputId: id}) |
| 215 | + } |
| 216 | + return out, nil |
| 217 | +} |
0 commit comments