diff --git a/backend/internal/cli/session.go b/backend/internal/cli/session.go index cae02f0cb5..b46c07e456 100644 --- a/backend/internal/cli/session.go +++ b/backend/internal/cli/session.go @@ -36,10 +36,20 @@ type sessionClaimPROptions struct { noTakeover bool } +type sessionMergePolicyOptions struct { + project string + json bool + terminateOnPRMerge bool +} + type sessionRenameRequest struct { DisplayName string `json:"displayName"` } +type sessionMergePolicyRequest struct { + TerminateOnPRMerge bool `json:"terminateOnPrMerge"` +} + type sessionDTO struct { ID string `json:"id"` ProjectID string `json:"projectId"` @@ -82,6 +92,13 @@ type renameSessionResponse struct { DisplayName string `json:"displayName"` } +type setMergePolicyResponse struct { + OK bool `json:"ok"` + SessionID string `json:"sessionId"` + TerminateOnPRMerge bool `json:"terminateOnPrMerge"` + Session sessionDTO `json:"session"` +} + type cleanupSkippedSession struct { SessionID string `json:"sessionId"` Reason string `json:"reason"` @@ -146,6 +163,7 @@ func newSessionCommand(ctx *commandContext) *cobra.Command { cmd.AddCommand(newSessionKillCommand(ctx)) cmd.AddCommand(newSessionRestoreCommand(ctx)) cmd.AddCommand(newSessionRenameCommand(ctx)) + cmd.AddCommand(newSessionSetMergePolicyCommand(ctx)) cmd.AddCommand(newSessionCleanupCommand(ctx)) cmd.AddCommand(newSessionClaimPRCommand(ctx)) return cmd @@ -243,6 +261,30 @@ func newSessionRenameCommand(ctx *commandContext) *cobra.Command { return cmd } +func newSessionSetMergePolicyCommand(ctx *commandContext) *cobra.Command { + var opts sessionMergePolicyOptions + cmd := &cobra.Command{ + Use: "set-merge-policy ", + Short: "Set whether the session auto-terminates when its PR merges", + Args: oneSessionIDArg, + RunE: func(cmd *cobra.Command, args []string) error { + if !cmd.Flags().Changed("terminate-on-pr-merge") { + return usageError{errors.New("--terminate-on-pr-merge is required (true or false)")} + } + id, err := normalizeSessionID(args[0]) + if err != nil { + return err + } + return ctx.setSessionMergePolicy(cmd.Context(), cmd, id, opts) + }, + } + f := cmd.Flags() + addSessionProjectFlag(f, &opts.project, "Project id to scope the lookup") + f.BoolVar(&opts.terminateOnPRMerge, "terminate-on-pr-merge", false, "Terminate the session when its claimed PR merges (required)") + f.BoolVar(&opts.json, "json", false, "Output as JSON") + return cmd +} + func newSessionCleanupCommand(ctx *commandContext) *cobra.Command { var opts sessionCleanupOptions cmd := &cobra.Command{ @@ -494,6 +536,32 @@ func (c *commandContext) renameSession(ctx context.Context, cmd *cobra.Command, return err } +func (c *commandContext) setSessionMergePolicy(ctx context.Context, cmd *cobra.Command, id string, opts sessionMergePolicyOptions) error { + if opts.project != "" { + if _, err := c.fetchScopedSession(ctx, id, opts.project); err != nil { + return err + } + } + var res setMergePolicyResponse + req := sessionMergePolicyRequest{TerminateOnPRMerge: opts.terminateOnPRMerge} + if err := c.patchJSON(ctx, "sessions/"+url.PathEscape(id)+"/merge-policy", req, &res); err != nil { + return err + } + if opts.json { + return writeJSON(cmd.OutOrStdout(), res) + } + sessionID := res.SessionID + if sessionID == "" { + sessionID = id + } + policy := "keep session after PR merge" + if res.TerminateOnPRMerge { + policy = "terminate session when PR merges" + } + _, err := fmt.Fprintf(cmd.OutOrStdout(), "session %s merge policy: %s\n", sessionID, policy) + return err +} + func (c *commandContext) cleanupSessions(ctx context.Context, cmd *cobra.Command, opts sessionCleanupOptions) error { candidates, err := c.previewCleanupSessions(ctx, opts.project) if err != nil { diff --git a/backend/internal/cli/session_test.go b/backend/internal/cli/session_test.go index 6b3ba04b0a..c9bddfb1c3 100644 --- a/backend/internal/cli/session_test.go +++ b/backend/internal/cli/session_test.go @@ -3,6 +3,7 @@ package cli import ( "context" "encoding/json" + "fmt" "io" "net/http" "net/http/httptest" @@ -95,6 +96,13 @@ func sessionCommandServer(t *testing.T) (*httptest.Server, *sessionRequestLog) { return } _, _ = io.WriteString(w, `{"ok":true,"sessionId":"demo-1","displayName":`+jsonQuote(req.DisplayName)+`}`) + case r.Method == http.MethodPatch && r.URL.Path == "/api/v1/sessions/demo-1/merge-policy": + var req sessionMergePolicyRequest + if err := json.NewDecoder(r.Body).Decode(&req); err != nil { + http.Error(w, err.Error(), http.StatusBadRequest) + return + } + _, _ = io.WriteString(w, `{"ok":true,"sessionId":"demo-1","terminateOnPrMerge":`+fmt.Sprintf("%t", req.TerminateOnPRMerge)+`,"session":`+sessionJSON("demo-1", "demo", "worker", "working", false)+`}`) default: http.NotFound(w, r) } @@ -432,9 +440,60 @@ func TestSessionRename_SuccessWithProjectScope(t *testing.T) { } } +func TestSessionSetMergePolicy_Success(t *testing.T) { + cfg := setConfigEnv(t) + srv, log := sessionCommandServer(t) + writeRunFileFor(t, cfg, srv) + + out, errOut, err := executeCLI(t, Deps{ + ProcessAlive: func(int) bool { return true }, + }, "session", "set-merge-policy", "demo-1", "--terminate-on-pr-merge", "-p", "demo") + if err != nil { + t.Fatalf("session set-merge-policy failed: %v\nstderr=%s", err, errOut) + } + if !strings.Contains(out, "session demo-1 merge policy: terminate session when PR merges") { + t.Fatalf("unexpected set-merge-policy output:\n%s", out) + } + want := []string{"GET /api/v1/sessions/demo-1", "PATCH /api/v1/sessions/demo-1/merge-policy"} + if got := log.all(); !reflect.DeepEqual(got, want) { + t.Fatalf("requests = %#v, want %#v", got, want) + } +} + +func TestSessionSetMergePolicy_DisableAndJSON(t *testing.T) { + cfg := setConfigEnv(t) + srv, _ := sessionCommandServer(t) + writeRunFileFor(t, cfg, srv) + + out, errOut, err := executeCLI(t, Deps{ + ProcessAlive: func(int) bool { return true }, + }, "session", "set-merge-policy", "demo-1", "--terminate-on-pr-merge=false", "--json") + if err != nil { + t.Fatalf("session set-merge-policy --json failed: %v\nstderr=%s", err, errOut) + } + var got setMergePolicyResponse + if err := json.Unmarshal([]byte(out), &got); err != nil { + t.Fatalf("set-merge-policy --json output is not decodable: %v\noutput=%s", err, out) + } + if !got.OK || got.SessionID != "demo-1" || got.TerminateOnPRMerge { + t.Fatalf("unexpected merge policy JSON: %#v", got) + } +} + +func TestSessionSetMergePolicy_MissingFlagIsUsageError(t *testing.T) { + setConfigEnv(t) + _, errOut, err := executeCLI(t, Deps{}, "session", "set-merge-policy", "demo-1") + if err == nil { + t.Fatal("expected missing --terminate-on-pr-merge to fail") + } + if got := ExitCode(err); got != 2 { + t.Fatalf("exit code = %d, want 2 (err=%v stderr=%s)", got, err, errOut) + } +} + func TestSessionCommands_MissingIDIsUsageError(t *testing.T) { setConfigEnv(t) - for _, sub := range []string{"get", "kill", "restore"} { + for _, sub := range []string{"get", "kill", "restore", "set-merge-policy"} { t.Run(sub, func(t *testing.T) { _, _, err := executeCLI(t, Deps{}, "session", sub) if err == nil { diff --git a/docs/cli/README.md b/docs/cli/README.md index 3302f0f941..583a2ff8da 100644 --- a/docs/cli/README.md +++ b/docs/cli/README.md @@ -51,6 +51,7 @@ Every product command resolves to a daemon HTTP route. Run `ao | `ao session kill ` | `POST /api/v1/sessions/{id}/kill` | | `ao session restore ` | `POST /api/v1/sessions/{id}/restore` | | `ao session rename ` | `PATCH /api/v1/sessions/{id}` | +| `ao session set-merge-policy ` | `PATCH /api/v1/sessions/{id}/merge-policy` | | `ao session cleanup` | `POST /api/v1/sessions/cleanup` | | `ao session claim-pr ` | `POST /api/v1/sessions/{id}/pr/claim` | | `ao orchestrator ls` | `GET /api/v1/orchestrators` |