From 32e05f09d27cd0dfad86faf5277ef5764f4609f3 Mon Sep 17 00:00:00 2001 From: Jayesh Sharma Date: Fri, 6 Mar 2026 18:53:36 +0530 Subject: [PATCH] feat: add PR claim flow for agent sessions --- docs/design/README.md | 1 + docs/design/session-replacement-handoff.md | 443 ++++++++++++++++++ .../__tests__/commands/review-check.test.ts | 11 +- .../cli/__tests__/commands/session.test.ts | 65 +++ packages/cli/__tests__/commands/spawn.test.ts | 61 ++- packages/cli/__tests__/commands/start.test.ts | 112 +++-- .../cli/__tests__/commands/status.test.ts | 11 +- packages/cli/src/commands/session.ts | 63 ++- .../src/__tests__/lifecycle-manager.test.ts | 50 ++ packages/core/src/__tests__/metadata.test.ts | 2 + .../src/__tests__/plugin-integration.test.ts | 3 + .../core/src/__tests__/prompt-builder.test.ts | 1 + .../src/__tests__/session-manager.test.ts | 155 ++++++ packages/core/src/lifecycle-manager.ts | 6 +- packages/core/src/metadata.ts | 10 +- packages/core/src/orchestrator-prompt.ts | 13 + packages/core/src/prompt-builder.ts | 1 + packages/core/src/session-manager.ts | 137 +++++- packages/core/src/types.ts | 26 + packages/plugins/scm-github/src/index.ts | 123 ++++- .../plugins/scm-github/test/index.test.ts | 74 +++ 21 files changed, 1285 insertions(+), 83 deletions(-) create mode 100644 docs/design/session-replacement-handoff.md diff --git a/docs/design/README.md b/docs/design/README.md index 3b0d336f7..e60305139 100644 --- a/docs/design/README.md +++ b/docs/design/README.md @@ -11,6 +11,7 @@ | [`design-brief.md`](./design-brief.md) | **Main design brief** — competitive analysis, full color palette, typography, all component specs, anti-patterns, implementation stack recommendation, and current codebase audit | | [`session-detail-design-brief.md`](./session-detail-design-brief.md) | Design spec for `/sessions/[id]` — the single-agent investigation view | | [`orchestrator-terminal-design-brief.md`](./orchestrator-terminal-design-brief.md) | Design spec for the orchestrator terminal — full-viewport command center with status strip | +| [`session-replacement-handoff.md`](./session-replacement-handoff.md) | Design plan for successor sessions, PR takeover, and context handoff after replacing a worker | | [`token-reference.css`](./token-reference.css) | **Ready-to-use CSS** — drop-in replacement for `globals.css` `@theme` block | | [`competitive-analysis-raw.md`](./competitive-analysis-raw.md) | Raw research notes from all 14 competitor sites (Linear, Vercel, Railway, Fly.io, Inngest, Temporal, Grafana, WandB, LangSmith, Retool, Render, PlanetScale, Supabase, GitHub Copilot) | | [`design-brief-v1.md`](./design-brief-v1.md) | Original v1 brief (text-only research, pre-Playwright CSS extraction) — kept for reference | diff --git a/docs/design/session-replacement-handoff.md b/docs/design/session-replacement-handoff.md new file mode 100644 index 000000000..999ce9d36 --- /dev/null +++ b/docs/design/session-replacement-handoff.md @@ -0,0 +1,443 @@ +# Session Replacement + PR Handoff — Design Plan + +## Status + +**Separate feature from PR claiming. Not implemented.** + +The new `claim-pr` flow solves one problem: explicitly attaching an existing PR to a running AO session. + +This document covers a different problem: **when the original owner session is no longer the right place to continue work, how should AO replace that session, transfer PR ownership, and preserve enough context for the new session to continue effectively?** + +--- + +## Why This Is Separate + +`claim-pr` establishes **PR ownership**. + +Session replacement / handoff adds two more concerns that are not solved by `claim-pr` alone: + +1. **Successor semantics** — how AO knows that session `app-12` is replacing `app-7`. +2. **Context continuity** — how the replacement session gets enough prior context to continue work without starting cold. + +These are workflow and product questions, not just plumbing. + +--- + +## Existing Capability + +AO already has two adjacent primitives: + +- **In-place restore**: restore the same session ID using existing metadata/workspace/runtime recovery. +- **PR claiming**: attach an existing PR to a session and optionally take it over from another session. + +Those are related, but not the same as a true replacement flow. + +### Important distinction + +- **Restore** = revive the same logical session. +- **Replace / handoff** = create a new logical session that succeeds an old one. + +A replacement flow needs explicit lineage, ownership transfer, and context transfer. + +--- + +## Problem + +Example: + +- `app-7` owns PR `#123` +- CI fails or review changes arrive +- `app-7` is stuck, crashed, too confused, or otherwise not the right worker anymore +- AO wants to continue work in a fresh session `app-12` + +Today, AO does not have a first-class notion that: + +- `app-12` is the successor of `app-7` +- `app-12` should inherit the PR +- `app-12` should receive a usable summary of what `app-7` already did + +That should be its own feature. + +--- + +## Goal + +Add a safe, explicit **session replacement + handoff** workflow that lets AO: + +1. Create a replacement session for an existing worker. +2. Mark the new session as the successor of the old one. +3. Transfer PR ownership to the new session. +4. Preserve enough context for the new session to continue productively. +5. Ensure lifecycle/reactions route future CI/review work to the replacement session. + +--- + +## Non-Goals + +This feature should **not** initially try to: + +- Guess successor relationships heuristically from branch names or issue IDs alone. +- Silently transfer PR ownership during normal polling. +- Move full conversational state between arbitrary agent tools unless a native resume primitive exists and is proven reliable. +- Preserve every detail of the old session transcript as a hard requirement. + +For MVP, explicit replacement is better than clever inference. + +--- + +## Recommended Product Shape + +### User-facing workflow + +Introduce a replacement-oriented command or API, such as: + +```bash +ao session replace app-7 +``` + +Potential options later: + +```bash +ao session replace app-7 --reason stuck +ao session replace app-7 --claim-pr +ao session replace app-7 --carry-context +``` + +### Internal behavior + +High-level flow: + +1. Read session `app-7` metadata. +2. Spawn a new session `app-12`. +3. Mark `app-12` as successor of `app-7`. +4. If `app-7` owns a PR, call `claimPR("app-12", pr, { takeover: true })`. +5. Build a handoff context package for `app-12`. +6. Send that package to `app-12` as its first instruction, or launch it with that context. +7. Mark `app-7` as replaced/superseded so lifecycle and humans can see what happened. + +--- + +## Core Design Principle + +**Replacement must be explicit.** + +AO should not assume that a fresh session is the successor of an older session unless: + +- the orchestrator explicitly created it as a replacement, or +- metadata explicitly links the two sessions. + +This avoids accidental PR hijacking and bad routing. + +--- + +## Proposed Metadata Model + +Add lineage metadata so the relationship is durable and inspectable. + +### On the new session + +```text +supersedes=app-7 +handoffReason=stuck +handoffAt=2026-03-06T12:34:56.000Z +handoffContextMode=summary +``` + +### On the old session + +```text +replacedBy=app-12 +replacedAt=2026-03-06T12:34:56.000Z +status=replaced +``` + +Notes: + +- `status=replaced` would likely be a new lifecycle status if we want it surfaced directly. +- If we do not want a new lifecycle status immediately, we can keep old status and rely on `replacedBy`, but that is less visible. + +--- + +## PR Ownership Transfer + +This part is now mechanically straightforward because `claimPR(...)` exists. + +If the old session owns PR `#123`, replacement should do: + +```ts +claimPR(newSessionId, "123", { takeover: true }) +``` + +Expected result: + +- new session becomes the PR owner +- PR branch is checked out in the replacement workspace +- old session loses PR ownership +- old session has PR auto-detect disabled so lifecycle does not reattach it by branch + +This gives AO a clean single-owner model. + +--- + +## Context Handoff: The Real Hard Part + +This is the main reason the feature should be treated separately. + +There are several possible levels of context carry-over. + +### Option 1 — No transfer, just replace the worker + +AO spawns a fresh session and only tells it what PR/issue to work on. + +**Pros** +- simplest implementation +- lowest coupling to agent internals + +**Cons** +- replacement session starts cold +- loses reasoning trail, failed attempts, prior decisions +- more likely to repeat work or miss subtle repo context + +This is probably too weak for a good user experience. + +--- + +### Option 2 — AO-generated handoff summary (recommended MVP) + +AO constructs a structured handoff package from existing session state and gives that to the replacement session. + +Example contents: + +- issue ID / PR URL +- branch name +- latest agent summary +- last known status (`ci_failed`, `changes_requested`, etc.) +- recent terminal output excerpt +- unresolved review comments +- failing CI checks +- replacement reason (`stuck`, `crashed`, `manual takeover`) + +Example first prompt: + +```text +You are replacing session app-7. + +Context: +- Issue: INT-1234 +- PR: https://github.com/org/repo/pull/123 +- Branch: feat/INT-1234 +- Previous session status: ci_failed +- Replacement reason: previous session became stuck +- Summary from previous session: implemented API validation and tests; CI failing in e2e +- Current failing checks: e2e / login flow +- Review comments still open: 2 + +You now own this PR. Continue from the current branch state. First inspect CI failures and confirm the current blocking issue before changing code. +``` + +**Pros** +- explicit and portable across agents +- works even without native resume support +- keeps the product behavior understandable + +**Cons** +- summary may omit useful detail +- quality depends on how good the summary extraction is + +This is the best MVP path. + +--- + +### Option 3 — Native agent resume into successor workflow + +Some agents already expose native resume semantics for the same underlying conversation/thread. + +Examples in the current codebase: + +- Codex agent supports a native `resume` flow. +- Claude Code agent supports a native `--resume` flow. + +However, these are currently used for **restoring the same session**, not necessarily transferring work to a brand-new successor session with a different AO identity. + +Open questions: + +- Can a new AO session safely wrap an old agent thread? +- Does the agent assume the same workspace path? +- Does resuming a thread into a new worktree cause confusion or hidden state mismatch? +- Can the runtime/plugin reliably expose the old thread ID for successor use? + +**Pros** +- highest continuity if it works reliably +- preserves reasoning and full tool-use context + +**Cons** +- agent-specific +- potentially fragile across workspace changes +- harder to reason about operationally + +Recommendation: treat native resume as an **optional enhancement**, not the MVP baseline. + +--- + +### Option 4 — Full transcript migration + +AO could theoretically extract the old transcript and replay or summarize it into the new session. + +This should **not** be the MVP. + +Problems: + +- privacy / verbosity concerns +- can be extremely long +- transcript replay is not the same as genuine state transfer +- tool outputs and local state may no longer match replayed text + +Use summary, not transcript migration, for MVP. + +--- + +## Recommended MVP + +### MVP scope + +Build a **manual, explicit replacement workflow** with **summary-based context handoff**. + +### Suggested flow + +1. User or orchestrator chooses to replace `app-7`. +2. AO spawns replacement session `app-12`. +3. AO records: + - `app-12 supersedes app-7` + - `app-7 replacedBy app-12` +4. AO transfers PR ownership via `claimPR(..., { takeover: true })`. +5. AO builds a structured handoff summary. +6. AO sends the summary to `app-12` immediately. +7. AO marks future lifecycle/reaction routing to `app-12`. + +### Why this is the right MVP + +- simple mental model +- explicit ownership +- no unsafe heuristics +- works across different agents +- does not depend on native conversation-resume semantics + +--- + +## Suggested API Shape + +### Core + +Potential new core API: + +```ts +sessionManager.replace(sessionId, options?) +``` + +Possible return shape: + +```ts +{ + oldSessionId: "app-7", + newSessionId: "app-12", + projectId: "my-app", + claimedPR: "https://github.com/org/repo/pull/123", + contextMode: "summary", +} +``` + +### Internal helper + +Potential helper for context packaging: + +```ts +buildSessionHandoffContext(oldSession, project) +``` + +Output could be structured JSON or a formatted prompt block. + +--- + +## Handoff Context Sources + +The summary-based MVP can pull from existing AO state: + +- session metadata +- agent summary (`summary`) +- PR URL / branch / issue metadata +- recent terminal output +- current lifecycle status +- failing CI checks +- unresolved review comments +- replacement reason + +This is enough to make the new session useful without pretending to clone the old agent’s memory perfectly. + +--- + +## Lifecycle / Routing Implications + +After replacement: + +- the replacement session should be the only active PR owner +- CI/review reactions should target the replacement session +- old session should not receive future automated PR-routing work +- dashboard / `ao session ls` should show the lineage clearly + +This is where explicit lineage metadata matters. + +--- + +## Open Questions + +1. Should `replaced` be a first-class session status? +2. Should replacement always imply PR takeover if the old session has a PR? +3. Should the old session be killed immediately, or only marked superseded? +4. Should native resume be attempted for certain agents, or only after summary handoff is stable? +5. How much terminal output should be included in the handoff summary? +6. Should the replacement session reuse the old workspace, or always get a fresh workspace and rely on branch checkout? + +--- + +## Risks + +### Product risk + +If replacement is too implicit, users may not trust why a new session suddenly “owns” a PR. + +### Technical risk + +Native resume across successor sessions may look attractive but be brittle if workspace/runtime assumptions differ. + +### Operational risk + +If old and new sessions both appear to own the same branch/PR, routing becomes confusing and reactions may duplicate work. + +This is why explicit lineage + single-owner transfer is essential. + +--- + +## Acceptance Criteria + +A future implementation should satisfy: + +- A replacement session is explicitly linked to the session it supersedes. +- If the old session owned a PR, the new session can take it over cleanly. +- Future CI/review reactions route only to the replacement session. +- The replacement session receives a usable handoff context package. +- The user can see replacement lineage in AO metadata / UX. +- The workflow is explicit and inspectable, not heuristic magic. + +--- + +## Recommendation + +Yes — this should absolutely be treated as a **separate feature** from PR claiming. + +Recommended order: + +1. **PR claiming** — done first, because it provides the ownership-transfer primitive. +2. **Replacement lineage + summary handoff** — next, as the real workflow feature. +3. **Optional native resume enhancements** — later, agent by agent, only if they are reliable. + +That sequencing keeps the system understandable and avoids mixing ownership transfer with the much harder question of conversational continuity. diff --git a/packages/cli/__tests__/commands/review-check.test.ts b/packages/cli/__tests__/commands/review-check.test.ts index 13a5843b5..d63c9f4b0 100644 --- a/packages/cli/__tests__/commands/review-check.test.ts +++ b/packages/cli/__tests__/commands/review-check.test.ts @@ -1,5 +1,13 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { mkdtempSync, mkdirSync, writeFileSync, rmSync, existsSync, readdirSync, readFileSync } from "node:fs"; +import { + mkdtempSync, + mkdirSync, + writeFileSync, + rmSync, + existsSync, + readdirSync, + readFileSync, +} from "node:fs"; import { join } from "node:path"; import { tmpdir } from "node:os"; import { type Session, type SessionManager, getSessionsDir } from "@composio/ao-core"; @@ -18,6 +26,7 @@ const { mockTmux, mockExec, mockGh, mockConfigRef, mockSessionManager, sessionsD spawn: vi.fn(), spawnOrchestrator: vi.fn(), send: vi.fn(), + claimPR: vi.fn(), }, sessionsDirRef: { current: "" }, })); diff --git a/packages/cli/__tests__/commands/session.test.ts b/packages/cli/__tests__/commands/session.test.ts index ba44da863..ccd6f5130 100644 --- a/packages/cli/__tests__/commands/session.test.ts +++ b/packages/cli/__tests__/commands/session.test.ts @@ -33,6 +33,7 @@ const { mockTmux, mockGit, mockGh, mockExec, mockConfigRef, mockSessionManager, spawn: vi.fn(), spawnOrchestrator: vi.fn(), send: vi.fn(), + claimPR: vi.fn(), }, sessionsDirRef: { current: "" }, })); @@ -171,6 +172,7 @@ beforeEach(() => { mockSessionManager.get.mockReset(); mockSessionManager.spawn.mockReset(); mockSessionManager.send.mockReset(); + mockSessionManager.claimPR.mockReset(); // Default: list reads from sessionsDir mockSessionManager.list.mockImplementation(async () => { @@ -186,6 +188,23 @@ beforeEach(() => { skipped: [], errors: [], } satisfies CleanupResult); + mockSessionManager.claimPR.mockResolvedValue({ + sessionId: "app-1", + projectId: "my-app", + pr: { + number: 42, + url: "https://github.com/org/repo/pull/42", + title: "Existing PR", + owner: "org", + repo: "repo", + branch: "feat/existing-pr", + baseBranch: "main", + isDraft: false, + }, + branchChanged: true, + githubAssigned: false, + takenOverFrom: [], + }); }); afterEach(() => { @@ -312,6 +331,52 @@ describe("session kill", () => { }); }); +describe("session claim-pr", () => { + afterEach(() => { + delete process.env["AO_SESSION_NAME"]; + delete process.env["AO_SESSION"]; + }); + + it("claims a PR for an explicit session", async () => { + await program.parseAsync([ + "node", + "test", + "session", + "claim-pr", + "42", + "app-2", + "--assign-on-github", + "--takeover", + ]); + + expect(mockSessionManager.claimPR).toHaveBeenCalledWith("app-2", "42", { + assignOnGithub: true, + takeover: true, + }); + + const output = consoleSpy.mock.calls.map((c) => String(c[0])).join("\n"); + expect(output).toContain("Session app-2 claimed PR #42"); + expect(output).toContain("feat/existing-pr"); + }); + + it("uses AO_SESSION_NAME when session argument is omitted", async () => { + process.env["AO_SESSION_NAME"] = "app-7"; + + await program.parseAsync(["node", "test", "session", "claim-pr", "42"]); + + expect(mockSessionManager.claimPR).toHaveBeenCalledWith("app-7", "42", { + assignOnGithub: undefined, + takeover: undefined, + }); + }); + + it("fails when no session can be resolved", async () => { + await expect(program.parseAsync(["node", "test", "session", "claim-pr", "42"])).rejects.toThrow( + "process.exit(1)", + ); + }); +}); + describe("session cleanup", () => { it("kills sessions with merged PRs", async () => { writeFileSync( diff --git a/packages/cli/__tests__/commands/spawn.test.ts b/packages/cli/__tests__/commands/spawn.test.ts index 4b314e20e..f1571fd82 100644 --- a/packages/cli/__tests__/commands/spawn.test.ts +++ b/packages/cli/__tests__/commands/spawn.test.ts @@ -15,6 +15,7 @@ const { mockExec, mockConfigRef, mockSessionManager } = vi.hoisted(() => ({ spawn: vi.fn(), spawnOrchestrator: vi.fn(), send: vi.fn(), + claimPR: vi.fn(), }, })); @@ -291,17 +292,17 @@ describe("spawn command", () => { }); it("rejects unknown project ID", async () => { - await expect( - program.parseAsync(["node", "test", "spawn", "nonexistent"]), - ).rejects.toThrow("process.exit(1)"); + await expect(program.parseAsync(["node", "test", "spawn", "nonexistent"])).rejects.toThrow( + "process.exit(1)", + ); }); it("reports error when spawn fails", async () => { mockSessionManager.spawn.mockRejectedValue(new Error("worktree creation failed")); - await expect( - program.parseAsync(["node", "test", "spawn", "my-app"]), - ).rejects.toThrow("process.exit(1)"); + await expect(program.parseAsync(["node", "test", "spawn", "my-app"])).rejects.toThrow( + "process.exit(1)", + ); }); }); @@ -309,11 +310,14 @@ describe("spawn pre-flight checks", () => { it("fails with clear error when tmux is not installed (default runtime)", async () => { mockExec.mockRejectedValue(new Error("ENOENT")); - await expect( - program.parseAsync(["node", "test", "spawn", "my-app"]), - ).rejects.toThrow("process.exit(1)"); + await expect(program.parseAsync(["node", "test", "spawn", "my-app"])).rejects.toThrow( + "process.exit(1)", + ); - const errors = vi.mocked(console.error).mock.calls.map((c) => String(c[0])).join("\n"); + const errors = vi + .mocked(console.error) + .mock.calls.map((c) => String(c[0])) + .join("\n"); expect(errors).toContain("tmux"); // Should not attempt to spawn expect(mockSessionManager.spawn).not.toHaveBeenCalled(); @@ -354,7 +358,10 @@ describe("spawn pre-flight checks", () => { }); it("checks gh auth when tracker is github", async () => { - const projects = (mockConfigRef.current as Record).projects as Record>; + const projects = (mockConfigRef.current as Record).projects as Record< + string, + Record + >; projects["my-app"].tracker = { plugin: "github" }; // tmux check passes, gh --version passes, gh auth status fails @@ -363,11 +370,14 @@ describe("spawn pre-flight checks", () => { .mockResolvedValueOnce({ stdout: "gh version 2.40", stderr: "" }) // gh --version .mockRejectedValueOnce(new Error("not logged in")); // gh auth status - await expect( - program.parseAsync(["node", "test", "spawn", "my-app"]), - ).rejects.toThrow("process.exit(1)"); + await expect(program.parseAsync(["node", "test", "spawn", "my-app"])).rejects.toThrow( + "process.exit(1)", + ); - const errors = vi.mocked(console.error).mock.calls.map((c) => String(c[0])).join("\n"); + const errors = vi + .mocked(console.error) + .mock.calls.map((c) => String(c[0])) + .join("\n"); expect(errors).toContain("not authenticated"); expect(mockSessionManager.spawn).not.toHaveBeenCalled(); }); @@ -390,7 +400,10 @@ describe("spawn pre-flight checks", () => { }; mockSessionManager.spawn.mockResolvedValue(fakeSession); - const projects = (mockConfigRef.current as Record).projects as Record>; + const projects = (mockConfigRef.current as Record).projects as Record< + string, + Record + >; projects["my-app"].tracker = { plugin: "linear" }; // tmux check passes — gh should never be called @@ -405,7 +418,10 @@ describe("spawn pre-flight checks", () => { }); it("distinguishes gh not installed from gh not authenticated", async () => { - const projects = (mockConfigRef.current as Record).projects as Record>; + const projects = (mockConfigRef.current as Record).projects as Record< + string, + Record + >; projects["my-app"].tracker = { plugin: "github" }; // tmux passes, gh --version fails (not installed) @@ -413,11 +429,14 @@ describe("spawn pre-flight checks", () => { .mockResolvedValueOnce({ stdout: "tmux 3.3a", stderr: "" }) // tmux -V .mockRejectedValueOnce(new Error("ENOENT")); // gh --version fails - await expect( - program.parseAsync(["node", "test", "spawn", "my-app"]), - ).rejects.toThrow("process.exit(1)"); + await expect(program.parseAsync(["node", "test", "spawn", "my-app"])).rejects.toThrow( + "process.exit(1)", + ); - const errors = vi.mocked(console.error).mock.calls.map((c) => String(c[0])).join("\n"); + const errors = vi + .mocked(console.error) + .mock.calls.map((c) => String(c[0])) + .join("\n"); expect(errors).toContain("not installed"); expect(errors).not.toContain("not authenticated"); }); diff --git a/packages/cli/__tests__/commands/start.test.ts b/packages/cli/__tests__/commands/start.test.ts index 51a879737..d5619ba02 100644 --- a/packages/cli/__tests__/commands/start.test.ts +++ b/packages/cli/__tests__/commands/start.test.ts @@ -15,24 +15,30 @@ import type { SessionManager } from "@composio/ao-core"; // Hoisted mocks // --------------------------------------------------------------------------- -const { mockExec, mockExecSilent, mockConfigRef, mockSessionManager, mockWaitForPortAndOpen, mockSpawn } = - vi.hoisted(() => ({ - mockExec: vi.fn(), - mockExecSilent: vi.fn(), - mockConfigRef: { current: null as Record | null }, - mockSessionManager: { - list: vi.fn(), - kill: vi.fn(), - cleanup: vi.fn(), - get: vi.fn(), - spawn: vi.fn(), - spawnOrchestrator: vi.fn(), - send: vi.fn(), - }, - mockWaitForPortAndOpen: vi.fn().mockResolvedValue(undefined), - mockSpawn: vi.fn(), - }), -); +const { + mockExec, + mockExecSilent, + mockConfigRef, + mockSessionManager, + mockWaitForPortAndOpen, + mockSpawn, +} = vi.hoisted(() => ({ + mockExec: vi.fn(), + mockExecSilent: vi.fn(), + mockConfigRef: { current: null as Record | null }, + mockSessionManager: { + list: vi.fn(), + kill: vi.fn(), + cleanup: vi.fn(), + get: vi.fn(), + spawn: vi.fn(), + spawnOrchestrator: vi.fn(), + send: vi.fn(), + claimPR: vi.fn(), + }, + mockWaitForPortAndOpen: vi.fn().mockResolvedValue(undefined), + mockSpawn: vi.fn(), +})); vi.mock("../../src/lib/shell.js", () => ({ tmux: vi.fn(), @@ -191,10 +197,7 @@ function createFakeRepo(dir: string, remoteUrl: string, files?: Record { await program.parseAsync(["node", "test", "start", "--no-dashboard", "--no-orchestrator"]); - const output = vi.mocked(console.log).mock.calls.map((c) => c.join(" ")).join("\n"); + const output = vi + .mocked(console.log) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(output).toContain("My App"); expect(output).toContain("Startup complete"); }); @@ -232,7 +238,10 @@ describe("start command — project resolution", () => { "--no-orchestrator", ]); - const output = vi.mocked(console.log).mock.calls.map((c) => c.join(" ")).join("\n"); + const output = vi + .mocked(console.log) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(output).toContain("Backend"); }); @@ -250,7 +259,10 @@ describe("start command — project resolution", () => { ]), ).rejects.toThrow("process.exit(1)"); - const errors = vi.mocked(console.error).mock.calls.map((c) => c.join(" ")).join("\n"); + const errors = vi + .mocked(console.error) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(errors).toContain("not found"); }); @@ -264,7 +276,10 @@ describe("start command — project resolution", () => { program.parseAsync(["node", "test", "start", "--no-dashboard", "--no-orchestrator"]), ).rejects.toThrow("process.exit(1)"); - const errors = vi.mocked(console.error).mock.calls.map((c) => c.join(" ")).join("\n"); + const errors = vi + .mocked(console.error) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(errors).toContain("Multiple projects"); }); @@ -275,7 +290,10 @@ describe("start command — project resolution", () => { program.parseAsync(["node", "test", "start", "--no-dashboard", "--no-orchestrator"]), ).rejects.toThrow("process.exit(1)"); - const errors = vi.mocked(console.error).mock.calls.map((c) => c.join(" ")).join("\n"); + const errors = vi + .mocked(console.error) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(errors).toContain("No projects configured"); }); }); @@ -305,7 +323,10 @@ describe("start command — URL argument", () => { // Config should have been generated expect(existsSync(join(repoDir, "agent-orchestrator.yaml"))).toBe(true); - const output = vi.mocked(console.log).mock.calls.map((c) => c.join(" ")).join("\n"); + const output = vi + .mocked(console.log) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(output).toContain("Reusing existing clone"); expect(output).toContain("Startup complete"); }); @@ -342,7 +363,10 @@ describe("start command — URL argument", () => { expect.anything(), ); - const output = vi.mocked(console.log).mock.calls.map((c) => c.join(" ")).join("\n"); + const output = vi + .mocked(console.log) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(output).toContain("Startup complete"); }); @@ -389,7 +413,10 @@ describe("start command — URL argument", () => { expect.anything(), ); - const output = vi.mocked(console.log).mock.calls.map((c) => c.join(" ")).join("\n"); + const output = vi + .mocked(console.log) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(output).toContain("Startup complete"); }); @@ -426,7 +453,10 @@ describe("start command — URL argument", () => { "--no-orchestrator", ]); - const output = vi.mocked(console.log).mock.calls.map((c) => c.join(" ")).join("\n"); + const output = vi + .mocked(console.log) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(output).toContain("Using existing config"); expect(output).toContain("Configured App"); }); @@ -470,7 +500,10 @@ describe("start command — URL argument", () => { "--no-orchestrator", ]); - const output = vi.mocked(console.log).mock.calls.map((c) => c.join(" ")).join("\n"); + const output = vi + .mocked(console.log) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); // Should pick "Multi Proj" by matching repo field, not error with "Multiple projects" expect(output).toContain("Multi Proj"); expect(output).toContain("Startup complete"); @@ -491,7 +524,10 @@ describe("start command — URL argument", () => { ]), ).rejects.toThrow("process.exit(1)"); - const errors = vi.mocked(console.error).mock.calls.map((c) => c.join(" ")).join("\n"); + const errors = vi + .mocked(console.error) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(errors).toContain("Failed to clone"); }); }); @@ -549,7 +585,10 @@ describe("stop command", () => { await program.parseAsync(["node", "test", "stop"]); expect(mockSessionManager.kill).toHaveBeenCalledWith("app-orchestrator"); - const output = vi.mocked(console.log).mock.calls.map((c) => c.join(" ")).join("\n"); + const output = vi + .mocked(console.log) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(output).toContain("Orchestrator stopped"); }); @@ -561,7 +600,10 @@ describe("stop command", () => { await program.parseAsync(["node", "test", "stop"]); expect(mockSessionManager.kill).not.toHaveBeenCalled(); - const output = vi.mocked(console.log).mock.calls.map((c) => c.join(" ")).join("\n"); + const output = vi + .mocked(console.log) + .mock.calls.map((c) => c.join(" ")) + .join("\n"); expect(output).toContain("is not running"); }); }); diff --git a/packages/cli/__tests__/commands/status.test.ts b/packages/cli/__tests__/commands/status.test.ts index e3e77eb37..702acb05c 100644 --- a/packages/cli/__tests__/commands/status.test.ts +++ b/packages/cli/__tests__/commands/status.test.ts @@ -1,5 +1,13 @@ import { describe, it, expect, vi, beforeEach, afterEach } from "vitest"; -import { mkdtempSync, writeFileSync, rmSync, mkdirSync, existsSync, readdirSync, readFileSync } from "node:fs"; +import { + mkdtempSync, + writeFileSync, + rmSync, + mkdirSync, + existsSync, + readdirSync, + readFileSync, +} from "node:fs"; import { join } from "node:path"; import { tmpdir } from "node:os"; import { @@ -39,6 +47,7 @@ const { spawn: vi.fn(), spawnOrchestrator: vi.fn(), send: vi.fn(), + claimPR: vi.fn(), }, sessionsDirRef: { current: "" }, })); diff --git a/packages/cli/src/commands/session.ts b/packages/cli/src/commands/session.ts index c93b7bef6..ed54c88a3 100644 --- a/packages/cli/src/commands/session.ts +++ b/packages/cli/src/commands/session.ts @@ -6,7 +6,9 @@ import { formatAge } from "../lib/format.js"; import { getSessionManager } from "../lib/create-session-manager.js"; export function registerSession(program: Command): void { - const session = program.command("session").description("Session management (ls, kill, cleanup)"); + const session = program + .command("session") + .description("Session management (ls, kill, cleanup, restore, claim-pr)"); session .command("ls") @@ -151,6 +153,65 @@ export function registerSession(program: Command): void { } }); + session + .command("claim-pr") + .description("Attach an existing PR to a session") + .argument("", "Pull request number or URL") + .argument("[session]", "Session name (defaults to AO_SESSION_NAME/AO_SESSION)") + .option("--assign-on-github", "Assign the PR to the authenticated GitHub user") + .option("--takeover", "Transfer PR ownership from another AO session if needed") + .action( + async ( + prRef: string, + sessionName: string | undefined, + opts: { assignOnGithub?: boolean; takeover?: boolean }, + ) => { + const config = loadConfig(); + const resolvedSession = + sessionName ?? process.env["AO_SESSION_NAME"] ?? process.env["AO_SESSION"]; + + if (!resolvedSession) { + console.error( + chalk.red( + "No session provided. Pass a session name or run this inside a managed AO session.", + ), + ); + process.exit(1); + } + + const sm = await getSessionManager(config); + + try { + const result = await sm.claimPR(resolvedSession, prRef, { + assignOnGithub: opts.assignOnGithub, + takeover: opts.takeover, + }); + + console.log(chalk.green(`\nSession ${resolvedSession} claimed PR #${result.pr.number}.`)); + console.log(chalk.dim(` PR: ${result.pr.url}`)); + console.log(chalk.dim(` Branch: ${result.pr.branch}`)); + console.log( + chalk.dim( + ` Checkout: ${result.branchChanged ? "switched to PR branch" : "already on PR branch"}`, + ), + ); + if (result.takenOverFrom.length > 0) { + console.log(chalk.dim(` Took over from: ${result.takenOverFrom.join(", ")}`)); + } + if (opts.assignOnGithub) { + if (result.githubAssigned) { + console.log(chalk.dim(" GitHub assignee: updated")); + } else if (result.githubAssignmentError) { + console.log(chalk.yellow(` GitHub assignee: ${result.githubAssignmentError}`)); + } + } + } catch (err) { + console.error(chalk.red(`Failed to claim PR for session ${resolvedSession}: ${err}`)); + process.exit(1); + } + }, + ); + session .command("restore") .description("Restore a terminated/crashed session in-place") diff --git a/packages/core/src/__tests__/lifecycle-manager.test.ts b/packages/core/src/__tests__/lifecycle-manager.test.ts index bd7cf5985..879d41edf 100644 --- a/packages/core/src/__tests__/lifecycle-manager.test.ts +++ b/packages/core/src/__tests__/lifecycle-manager.test.ts @@ -110,6 +110,7 @@ beforeEach(() => { kill: vi.fn().mockResolvedValue(undefined), cleanup: vi.fn(), send: vi.fn().mockResolvedValue(undefined), + claimPR: vi.fn(), }; config = { @@ -450,6 +451,55 @@ describe("check (single session)", () => { expect(lm.getStates().get("app-1")).toBe("ci_failed"); }); + it("skips PR auto-detection when metadata disables it", async () => { + const mockSCM: SCM = { + name: "mock-scm", + detectPR: vi.fn().mockResolvedValue(makePR()), + getPRState: vi.fn().mockResolvedValue("open"), + mergePR: vi.fn(), + closePR: vi.fn(), + getCIChecks: vi.fn(), + getCISummary: vi.fn().mockResolvedValue("passing"), + getReviews: vi.fn(), + getReviewDecision: vi.fn().mockResolvedValue("none"), + getPendingComments: vi.fn(), + getAutomatedComments: vi.fn(), + getMergeability: vi.fn(), + }; + + const registryWithSCM: PluginRegistry = { + ...mockRegistry, + get: vi.fn().mockImplementation((slot: string) => { + if (slot === "runtime") return mockRuntime; + if (slot === "agent") return mockAgent; + if (slot === "scm") return mockSCM; + return null; + }), + }; + + const session = makeSession({ status: "working", metadata: { prAutoDetect: "off" } }); + vi.mocked(mockSessionManager.get).mockResolvedValue(session); + + writeMetadata(sessionsDir, "app-1", { + worktree: "/tmp", + branch: "feat/test", + status: "working", + project: "my-app", + prAutoDetect: "off", + }); + + const lm = createLifecycleManager({ + config, + registry: registryWithSCM, + sessionManager: mockSessionManager, + }); + + await lm.check("app-1"); + + expect(mockSCM.detectPR).not.toHaveBeenCalled(); + expect(lm.getStates().get("app-1")).toBe("working"); + }); + it("detects merged PR", async () => { const mockSCM: SCM = { name: "mock-scm", diff --git a/packages/core/src/__tests__/metadata.test.ts b/packages/core/src/__tests__/metadata.test.ts index ca62f8115..369713742 100644 --- a/packages/core/src/__tests__/metadata.test.ts +++ b/packages/core/src/__tests__/metadata.test.ts @@ -46,6 +46,7 @@ describe("writeMetadata + readMetadata", () => { status: "pr_open", issue: "https://linear.app/team/issue/INT-100", pr: "https://github.com/org/repo/pull/42", + prAutoDetect: "off", summary: "Implementing feature X", project: "my-app", createdAt: "2025-01-01T00:00:00.000Z", @@ -56,6 +57,7 @@ describe("writeMetadata + readMetadata", () => { expect(meta).not.toBeNull(); expect(meta!.issue).toBe("https://linear.app/team/issue/INT-100"); expect(meta!.pr).toBe("https://github.com/org/repo/pull/42"); + expect(meta!.prAutoDetect).toBe("off"); expect(meta!.summary).toBe("Implementing feature X"); expect(meta!.project).toBe("my-app"); expect(meta!.createdAt).toBe("2025-01-01T00:00:00.000Z"); diff --git a/packages/core/src/__tests__/plugin-integration.test.ts b/packages/core/src/__tests__/plugin-integration.test.ts index 9d020489c..206bf9688 100644 --- a/packages/core/src/__tests__/plugin-integration.test.ts +++ b/packages/core/src/__tests__/plugin-integration.test.ts @@ -480,6 +480,7 @@ describe("plugin integration", () => { get: vi.fn().mockResolvedValue(makeSession({ status: "pr_open", pr })), kill: vi.fn().mockResolvedValue(undefined), send: vi.fn().mockResolvedValue(undefined), + claimPR: vi.fn(), spawnOrchestrator: vi.fn(), }; @@ -510,6 +511,7 @@ describe("plugin integration", () => { get: vi.fn().mockResolvedValue(makeSession({ status: "pr_open", pr })), kill: vi.fn().mockResolvedValue(undefined), send: vi.fn().mockResolvedValue(undefined), + claimPR: vi.fn(), spawnOrchestrator: vi.fn(), }; @@ -537,6 +539,7 @@ describe("plugin integration", () => { get: vi.fn().mockResolvedValue(makeSession({ status: "pr_open", pr })), kill: vi.fn().mockResolvedValue(undefined), send: vi.fn().mockResolvedValue(undefined), + claimPR: vi.fn(), spawnOrchestrator: vi.fn(), }; diff --git a/packages/core/src/__tests__/prompt-builder.test.ts b/packages/core/src/__tests__/prompt-builder.test.ts index b533f6f36..2fe1080d4 100644 --- a/packages/core/src/__tests__/prompt-builder.test.ts +++ b/packages/core/src/__tests__/prompt-builder.test.ts @@ -212,5 +212,6 @@ describe("BASE_AGENT_PROMPT", () => { expect(BASE_AGENT_PROMPT).toContain("Session Lifecycle"); expect(BASE_AGENT_PROMPT).toContain("Git Workflow"); expect(BASE_AGENT_PROMPT).toContain("PR Best Practices"); + expect(BASE_AGENT_PROMPT).toContain("ao session claim-pr"); }); }); diff --git a/packages/core/src/__tests__/session-manager.test.ts b/packages/core/src/__tests__/session-manager.test.ts index 03fef9045..dacb391b7 100644 --- a/packages/core/src/__tests__/session-manager.test.ts +++ b/packages/core/src/__tests__/session-manager.test.ts @@ -1724,6 +1724,161 @@ describe("restore", () => { }); }); +describe("claimPR", () => { + function makeSCM(overrides: Partial = {}): SCM { + return { + name: "mock-scm", + detectPR: vi.fn(), + resolvePR: vi.fn().mockResolvedValue({ + number: 42, + url: "https://github.com/org/my-app/pull/42", + title: "Existing PR", + owner: "org", + repo: "my-app", + branch: "feat/existing-pr", + baseBranch: "main", + isDraft: false, + }), + assignPRToCurrentUser: vi.fn().mockResolvedValue(undefined), + checkoutPR: vi.fn().mockResolvedValue(true), + getPRState: vi.fn().mockResolvedValue("open"), + getPRSummary: vi.fn(), + mergePR: vi.fn(), + closePR: vi.fn(), + getCIChecks: vi.fn(), + getCISummary: vi.fn(), + getReviews: vi.fn(), + getReviewDecision: vi.fn(), + getPendingComments: vi.fn(), + getAutomatedComments: vi.fn(), + getMergeability: vi.fn(), + ...overrides, + }; + } + + function registryWithSCM(mockSCM: SCM): PluginRegistry { + return { + ...mockRegistry, + get: vi.fn().mockImplementation((slot: string, _name: string) => { + if (slot === "runtime") return mockRuntime; + if (slot === "agent") return mockAgent; + if (slot === "workspace") return mockWorkspace; + if (slot === "scm") return mockSCM; + return null; + }), + }; + } + + it("claims an open PR and updates session metadata", async () => { + const mockSCM = makeSCM(); + + writeMetadata(sessionsDir, "app-2", { + worktree: "/tmp/ws-app-2", + branch: "feat/old-branch", + status: "working", + project: "my-app", + runtimeHandle: JSON.stringify(makeHandle("rt-2")), + }); + + const sm = createSessionManager({ config, registry: registryWithSCM(mockSCM) }); + const result = await sm.claimPR("app-2", "42"); + + expect(result.pr.number).toBe(42); + expect(result.branchChanged).toBe(true); + + expect(mockSCM.resolvePR).toHaveBeenCalledWith("42", config.projects["my-app"]); + expect(mockSCM.checkoutPR).toHaveBeenCalledWith(result.pr, "/tmp/ws-app-2"); + + const raw = readMetadataRaw(sessionsDir, "app-2"); + expect(raw).toMatchObject({ + branch: "feat/existing-pr", + status: "pr_open", + pr: "https://github.com/org/my-app/pull/42", + }); + expect(raw!["prAutoDetect"]).toBeUndefined(); + }); + + it("supports takeover by disabling PR auto-detect on the previous session", async () => { + const mockSCM = makeSCM(); + + writeMetadata(sessionsDir, "app-1", { + worktree: "/tmp/ws-app-1", + branch: "feat/existing-pr", + status: "review_pending", + project: "my-app", + pr: "https://github.com/org/my-app/pull/42", + runtimeHandle: JSON.stringify(makeHandle("rt-1")), + }); + + writeMetadata(sessionsDir, "app-2", { + worktree: "/tmp/ws-app-2", + branch: "feat/other-work", + status: "working", + project: "my-app", + runtimeHandle: JSON.stringify(makeHandle("rt-2")), + }); + + const sm = createSessionManager({ config, registry: registryWithSCM(mockSCM) }); + const result = await sm.claimPR("app-2", "42", { takeover: true }); + + expect(result.takenOverFrom).toEqual(["app-1"]); + + const previous = readMetadataRaw(sessionsDir, "app-1"); + expect(previous!["pr"]).toBeUndefined(); + expect(previous!["prAutoDetect"]).toBe("off"); + expect(previous!["status"]).toBe("working"); + }); + + it("rejects takeover when another session already tracks the PR", async () => { + const mockSCM = makeSCM(); + + writeMetadata(sessionsDir, "app-1", { + worktree: "/tmp/ws-app-1", + branch: "feat/existing-pr", + status: "pr_open", + project: "my-app", + pr: "https://github.com/org/my-app/pull/42", + runtimeHandle: JSON.stringify(makeHandle("rt-1")), + }); + + writeMetadata(sessionsDir, "app-2", { + worktree: "/tmp/ws-app-2", + branch: "feat/other-work", + status: "working", + project: "my-app", + runtimeHandle: JSON.stringify(makeHandle("rt-2")), + }); + + const sm = createSessionManager({ config, registry: registryWithSCM(mockSCM) }); + + await expect(sm.claimPR("app-2", "42")).rejects.toThrow("already tracked by app-1"); + }); + + it("keeps AO metadata updated even if GitHub assignment fails", async () => { + const mockSCM = makeSCM({ + assignPRToCurrentUser: vi.fn().mockRejectedValue(new Error("permission denied")), + }); + + writeMetadata(sessionsDir, "app-2", { + worktree: "/tmp/ws-app-2", + branch: "feat/old-branch", + status: "working", + project: "my-app", + runtimeHandle: JSON.stringify(makeHandle("rt-2")), + }); + + const sm = createSessionManager({ config, registry: registryWithSCM(mockSCM) }); + const result = await sm.claimPR("app-2", "42", { assignOnGithub: true }); + + expect(result.githubAssigned).toBe(false); + expect(result.githubAssignmentError).toContain("permission denied"); + + const raw = readMetadataRaw(sessionsDir, "app-2"); + expect(raw!["pr"]).toBe("https://github.com/org/my-app/pull/42"); + expect(raw!["status"]).toBe("pr_open"); + }); +}); + describe("PluginRegistry.loadBuiltins importFn", () => { it("should use provided importFn instead of built-in import", async () => { const { createPluginRegistry: createReg } = await import("../plugin-registry.js"); diff --git a/packages/core/src/lifecycle-manager.ts b/packages/core/src/lifecycle-manager.ts index 663fd7488..3ac87a8e4 100644 --- a/packages/core/src/lifecycle-manager.ts +++ b/packages/core/src/lifecycle-manager.ts @@ -211,9 +211,7 @@ export function createLifecycleManager(deps: LifecycleManagerDeps): LifecycleMan "runtime", project.runtime ?? config.defaults.runtime, ); - const terminalOutput = runtime - ? await runtime.getOutput(session.runtimeHandle, 10) - : ""; + const terminalOutput = runtime ? await runtime.getOutput(session.runtimeHandle, 10) : ""; if (terminalOutput) { const activity = agent.detectActivity(terminalOutput); if (activity === "waiting_input") return "needs_input"; @@ -237,7 +235,7 @@ export function createLifecycleManager(deps: LifecycleManagerDeps): LifecycleMan // 3. Auto-detect PR by branch if metadata.pr is missing. // This is critical for agents without auto-hook systems (Codex, Aider, // OpenCode) that can't reliably write pr= to metadata on their own. - if (!session.pr && scm && session.branch) { + if (!session.pr && scm && session.branch && session.metadata["prAutoDetect"] !== "off") { try { const detectedPR = await scm.detectPR(session, project); if (detectedPR) { diff --git a/packages/core/src/metadata.ts b/packages/core/src/metadata.ts index 3d70e52d9..0bf57ee3a 100644 --- a/packages/core/src/metadata.ts +++ b/packages/core/src/metadata.ts @@ -106,6 +106,8 @@ export function readMetadata(dataDir: string, sessionId: SessionId): SessionMeta tmuxName: raw["tmuxName"], issue: raw["issue"], pr: raw["pr"], + prAutoDetect: + raw["prAutoDetect"] === "off" ? "off" : raw["prAutoDetect"] === "on" ? "on" : undefined, summary: raw["summary"], project: raw["project"], agent: raw["agent"], @@ -115,7 +117,9 @@ export function readMetadata(dataDir: string, sessionId: SessionId): SessionMeta role: raw["role"], dashboardPort: raw["dashboardPort"] ? Number(raw["dashboardPort"]) : undefined, terminalWsPort: raw["terminalWsPort"] ? Number(raw["terminalWsPort"]) : undefined, - directTerminalWsPort: raw["directTerminalWsPort"] ? Number(raw["directTerminalWsPort"]) : undefined, + directTerminalWsPort: raw["directTerminalWsPort"] + ? Number(raw["directTerminalWsPort"]) + : undefined, }; } @@ -151,6 +155,7 @@ export function writeMetadata( if (metadata.tmuxName) data["tmuxName"] = metadata.tmuxName; if (metadata.issue) data["issue"] = metadata.issue; if (metadata.pr) data["pr"] = metadata.pr; + if (metadata.prAutoDetect) data["prAutoDetect"] = metadata.prAutoDetect; if (metadata.summary) data["summary"] = metadata.summary; if (metadata.project) data["project"] = metadata.project; if (metadata.agent) data["agent"] = metadata.agent; @@ -158,8 +163,7 @@ export function writeMetadata( if (metadata.runtimeHandle) data["runtimeHandle"] = metadata.runtimeHandle; if (metadata.restoredAt) data["restoredAt"] = metadata.restoredAt; if (metadata.role) data["role"] = metadata.role; - if (metadata.dashboardPort !== undefined) - data["dashboardPort"] = String(metadata.dashboardPort); + if (metadata.dashboardPort !== undefined) data["dashboardPort"] = String(metadata.dashboardPort); if (metadata.terminalWsPort !== undefined) data["terminalWsPort"] = String(metadata.terminalWsPort); if (metadata.directTerminalWsPort !== undefined) diff --git a/packages/core/src/orchestrator-prompt.ts b/packages/core/src/orchestrator-prompt.ts index d323600e4..9be3c3566 100644 --- a/packages/core/src/orchestrator-prompt.ts +++ b/packages/core/src/orchestrator-prompt.ts @@ -56,6 +56,9 @@ ao session ls -p ${projectId} # Send message to a session ao send ${project.sessionPrefix}-1 "Your message here" +# Claim an existing PR for a session +ao session claim-pr 123 ${project.sessionPrefix}-1 + # Kill a session ao session kill ${project.sessionPrefix}-1 @@ -72,6 +75,7 @@ ao open ${projectId} | \`ao spawn [issue]\` | Spawn a single worker agent session | | \`ao batch-spawn \` | Spawn multiple sessions in parallel | | \`ao session ls [-p project]\` | List all sessions (optionally filter by project) | +| \`ao session claim-pr [session]\` | Attach an existing PR to a session | | \`ao session attach \` | Attach to a session's tmux window | | \`ao session kill \` | Kill a specific session | | \`ao session cleanup [-p project]\` | Kill completed/merged sessions | @@ -107,6 +111,15 @@ Send instructions to a running agent: ao send ${project.sessionPrefix}-1 "Please address the review comments on your PR" \`\`\` +### PR Takeover + +If a session needs to continue work on an existing PR: +\`\`\`bash +ao session claim-pr 123 ${project.sessionPrefix}-1 +\`\`\` + +This updates AO metadata, switches the worktree onto the PR branch, and lets lifecycle reactions keep routing CI and review feedback to that session. + ### Cleanup Remove completed sessions: diff --git a/packages/core/src/prompt-builder.ts b/packages/core/src/prompt-builder.ts index 3a5a9525d..37386eda0 100644 --- a/packages/core/src/prompt-builder.ts +++ b/packages/core/src/prompt-builder.ts @@ -24,6 +24,7 @@ export const BASE_AGENT_PROMPT = `You are an AI coding agent managed by the Agen ## Session Lifecycle - You are running inside a managed session. Focus on the assigned task. - When you finish your work, create a PR and push it. The orchestrator will handle CI monitoring and review routing. +- If you're told to take over or continue work on an existing PR, run \`ao session claim-pr \` from inside this session before making changes. - If CI fails, the orchestrator will send you the failures — fix them and push again. - If reviewers request changes, the orchestrator will forward their comments — address each one, push fixes, and reply to the comments. diff --git a/packages/core/src/session-manager.ts b/packages/core/src/session-manager.ts index c2eb90ae5..66374c974 100644 --- a/packages/core/src/session-manager.ts +++ b/packages/core/src/session-manager.ts @@ -27,6 +27,8 @@ import { type OrchestratorSpawnConfig, type SessionStatus, type CleanupResult, + type ClaimPROptions, + type ClaimPRResult, type OrchestratorConfig, type ProjectConfig, type Runtime, @@ -106,6 +108,15 @@ const VALID_STATUSES: ReadonlySet = new Set([ "terminated", ]); +const PR_TRACKING_STATUSES: ReadonlySet = new Set([ + "pr_open", + "ci_failed", + "review_pending", + "changes_requested", + "approved", + "mergeable", +]); + /** Validate and normalize a status string. */ function validateStatus(raw: string | undefined): SessionStatus { // Bash scripts write "starting" — treat as "working" @@ -168,6 +179,13 @@ export interface SessionManagerDeps { export function createSessionManager(deps: SessionManagerDeps): SessionManager { const { config, registry } = deps; + interface LocatedSession { + raw: Record; + sessionsDir: string; + project: ProjectConfig; + projectId: string; + } + /** * Get the sessions directory for a project. */ @@ -268,6 +286,17 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { return { runtime, agent, workspace, tracker, scm }; } + function findSessionRecord(sessionId: SessionId): LocatedSession | null { + for (const [projectId, project] of Object.entries(config.projects)) { + const sessionsDir = getProjectSessionsDir(project); + const raw = readMetadataRaw(sessionsDir, sessionId); + if (!raw) continue; + return { raw, sessionsDir, project, projectId }; + } + + return null; + } + /** * Ensure session has a runtime handle (fabricate one if missing) and enrich * with live runtime state + activity detection. Used by both list() and get(). @@ -1005,6 +1034,110 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { await runtimePlugin.sendMessage(handle, message); } + async function claimPR( + sessionId: SessionId, + prRef: string, + options?: ClaimPROptions, + ): Promise { + const reference = prRef.trim(); + if (!reference) throw new Error("PR reference is required"); + + const located = findSessionRecord(sessionId); + if (!located) throw new Error(`Session ${sessionId} not found`); + + const { raw, sessionsDir, project, projectId } = located; + if (raw["role"] === "orchestrator") { + throw new Error(`Session ${sessionId} is an orchestrator session and cannot claim PRs`); + } + + const plugins = resolvePlugins(project, raw["agent"]); + const scm = plugins.scm; + if (!scm?.resolvePR || !scm.checkoutPR) { + throw new Error( + `SCM plugin ${project.scm?.plugin ? `"${project.scm.plugin}" ` : ""}does not support claiming existing PRs`, + ); + } + + const pr = await scm.resolvePR(reference, project); + const prState = await scm.getPRState(pr); + if (prState !== PR_STATE.OPEN) { + throw new Error(`Cannot claim PR #${pr.number} because it is ${prState}`); + } + + const conflictingSessions = new Set(); + for (const { sessionName } of listAllSessions(projectId)) { + if (sessionName === sessionId) continue; + + const otherRaw = readMetadataRaw(sessionsDir, sessionName); + if (!otherRaw || otherRaw["role"] === "orchestrator") continue; + + const samePr = otherRaw["pr"] === pr.url; + const sameBranch = + otherRaw["branch"] === pr.branch && (otherRaw["prAutoDetect"] ?? "on") !== "off"; + + if (samePr || sameBranch) { + conflictingSessions.add(sessionName); + } + } + + const takenOverFrom = [...conflictingSessions]; + if (takenOverFrom.length > 0 && !options?.takeover) { + throw new Error( + `PR #${pr.number} is already tracked by ${takenOverFrom.join(", ")}. Re-run with takeover enabled to transfer ownership.`, + ); + } + + const workspacePath = raw["worktree"]; + if (!workspacePath) { + throw new Error(`Session ${sessionId} has no workspace to check out PR #${pr.number}`); + } + + const branchChanged = await scm.checkoutPR(pr, workspacePath); + + updateMetadata(sessionsDir, sessionId, { + pr: pr.url, + status: "pr_open", + branch: pr.branch, + prAutoDetect: "", + }); + + for (const previousSessionId of takenOverFrom) { + const previousRaw = readMetadataRaw(sessionsDir, previousSessionId); + if (!previousRaw) continue; + + updateMetadata(sessionsDir, previousSessionId, { + pr: "", + prAutoDetect: "off", + ...(PR_TRACKING_STATUSES.has(previousRaw["status"] ?? "") ? { status: "working" } : {}), + }); + } + + let githubAssigned = false; + let githubAssignmentError: string | undefined; + if (options?.assignOnGithub) { + if (!scm.assignPRToCurrentUser) { + githubAssignmentError = `SCM plugin "${scm.name}" does not support assigning PRs`; + } else { + try { + await scm.assignPRToCurrentUser(pr); + githubAssigned = true; + } catch (err) { + githubAssignmentError = err instanceof Error ? err.message : String(err); + } + } + } + + return { + sessionId, + projectId, + pr, + branchChanged, + githubAssigned, + githubAssignmentError, + takenOverFrom, + }; + } + async function restore(sessionId: SessionId): Promise { // 1. Find session metadata across all projects (active first, then archive) let raw: Record | null = null; @@ -1055,6 +1188,8 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { tmuxName: raw["tmuxName"], issue: raw["issue"], pr: raw["pr"], + prAutoDetect: + raw["prAutoDetect"] === "off" ? "off" : raw["prAutoDetect"] === "on" ? "on" : undefined, summary: raw["summary"], project: raw["project"], createdAt: raw["createdAt"], @@ -1195,5 +1330,5 @@ export function createSessionManager(deps: SessionManagerDeps): SessionManager { return restoredSession; } - return { spawn, spawnOrchestrator, restore, list, get, kill, cleanup, send }; + return { spawn, spawnOrchestrator, restore, list, get, kill, cleanup, send, claimPR }; } diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index 617c60b28..056ce3695 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -508,6 +508,15 @@ export interface SCM { /** Detect if a session has an open PR (by branch name) */ detectPR(session: Session, project: ProjectConfig): Promise; + /** Resolve a PR reference (number or URL) into canonical PR metadata. */ + resolvePR?(reference: string, project: ProjectConfig): Promise; + + /** Assign a PR to the currently authenticated user, if supported. */ + assignPRToCurrentUser?(pr: PRInfo): Promise; + + /** Check out the PR branch into a workspace. Returns true if branch changed. */ + checkoutPR?(pr: PRInfo, workspacePath: string): Promise; + /** Get current PR state */ getPRState(pr: PRInfo): Promise; @@ -972,6 +981,7 @@ export interface SessionMetadata { tmuxName?: string; // Globally unique tmux session name (includes hash) issue?: string; pr?: string; + prAutoDetect?: "on" | "off"; summary?: string; project?: string; agent?: string; // Agent plugin name (e.g. "codex", "claude-code") — persisted for lifecycle @@ -998,6 +1008,22 @@ export interface SessionManager { kill(sessionId: SessionId): Promise; cleanup(projectId?: string, options?: { dryRun?: boolean }): Promise; send(sessionId: SessionId, message: string): Promise; + claimPR(sessionId: SessionId, prRef: string, options?: ClaimPROptions): Promise; +} + +export interface ClaimPROptions { + assignOnGithub?: boolean; + takeover?: boolean; +} + +export interface ClaimPRResult { + sessionId: SessionId; + projectId: string; + pr: PRInfo; + branchChanged: boolean; + githubAssigned: boolean; + githubAssignmentError?: string; + takenOverFrom: SessionId[]; } export interface CleanupResult { diff --git a/packages/plugins/scm-github/src/index.ts b/packages/plugins/scm-github/src/index.ts index e47321afd..7e3c44fcf 100644 --- a/packages/plugins/scm-github/src/index.ts +++ b/packages/plugins/scm-github/src/index.ts @@ -58,6 +58,36 @@ async function gh(args: string[]): Promise { } } +async function ghInDir(args: string[], cwd: string): Promise { + try { + const { stdout } = await execFileAsync("gh", args, { + cwd, + maxBuffer: 10 * 1024 * 1024, + timeout: 30_000, + }); + return stdout.trim(); + } catch (err) { + throw new Error(`gh ${args.slice(0, 3).join(" ")} failed: ${(err as Error).message}`, { + cause: err, + }); + } +} + +async function git(args: string[], cwd: string): Promise { + try { + const { stdout } = await execFileAsync("git", args, { + cwd, + maxBuffer: 10 * 1024 * 1024, + timeout: 30_000, + }); + return stdout.trim(); + } catch (err) { + throw new Error(`git ${args.slice(0, 3).join(" ")} failed: ${(err as Error).message}`, { + cause: err, + }); + } +} + function repoFlag(pr: PRInfo): string { return `${pr.owner}/${pr.repo}`; } @@ -68,6 +98,39 @@ function parseDate(val: string | undefined | null): Date { return isNaN(d.getTime()) ? new Date(0) : d; } +function assertProjectRepo(projectRepo: string): void { + const [owner, repo] = projectRepo.split("/"); + if (!owner || !repo) { + throw new Error(`Invalid repo format "${projectRepo}", expected "owner/repo"`); + } +} + +function prInfoFromView( + data: { + number: number; + url: string; + title: string; + headRefName: string; + baseRefName: string; + isDraft: boolean; + }, + projectRepo: string, +): PRInfo { + assertProjectRepo(projectRepo); + const [owner, repo] = projectRepo.split("/"); + + return { + number: data.number, + url: data.url, + title: data.title, + owner, + repo, + branch: data.headRefName, + baseBranch: data.baseRefName, + isDraft: data.isDraft, + }; +} + // --------------------------------------------------------------------------- // SCM implementation // --------------------------------------------------------------------------- @@ -78,12 +141,8 @@ function createGitHubSCM(): SCM { async detectPR(session: Session, project: ProjectConfig): Promise { if (!session.branch) return null; + assertProjectRepo(project.repo); - const parts = project.repo.split("/"); - if (parts.length !== 2 || !parts[0] || !parts[1]) { - throw new Error(`Invalid repo format "${project.repo}", expected "owner/repo"`); - } - const [owner, repo] = parts; try { const raw = await gh([ "pr", @@ -109,22 +168,54 @@ function createGitHubSCM(): SCM { if (prs.length === 0) return null; - const pr = prs[0]; - return { - number: pr.number, - url: pr.url, - title: pr.title, - owner, - repo, - branch: pr.headRefName, - baseBranch: pr.baseRefName, - isDraft: pr.isDraft, - }; + return prInfoFromView(prs[0], project.repo); } catch { return null; } }, + async resolvePR(reference: string, project: ProjectConfig): Promise { + const raw = await gh([ + "pr", + "view", + reference, + "--repo", + project.repo, + "--json", + "number,url,title,headRefName,baseRefName,isDraft", + ]); + + const data: { + number: number; + url: string; + title: string; + headRefName: string; + baseRefName: string; + isDraft: boolean; + } = JSON.parse(raw); + + return prInfoFromView(data, project.repo); + }, + + async assignPRToCurrentUser(pr: PRInfo): Promise { + await gh(["pr", "edit", String(pr.number), "--repo", repoFlag(pr), "--add-assignee", "@me"]); + }, + + async checkoutPR(pr: PRInfo, workspacePath: string): Promise { + const currentBranch = await git(["branch", "--show-current"], workspacePath); + if (currentBranch === pr.branch) return false; + + const dirty = await git(["status", "--porcelain"], workspacePath); + if (dirty) { + throw new Error( + `Workspace has uncommitted changes; cannot switch to PR branch "${pr.branch}" safely`, + ); + } + + await ghInDir(["pr", "checkout", String(pr.number), "--repo", repoFlag(pr)], workspacePath); + return true; + }, + async getPRState(pr: PRInfo): Promise { const raw = await gh([ "pr", diff --git a/packages/plugins/scm-github/test/index.test.ts b/packages/plugins/scm-github/test/index.test.ts index befa01d89..a33b1673c 100644 --- a/packages/plugins/scm-github/test/index.test.ts +++ b/packages/plugins/scm-github/test/index.test.ts @@ -188,6 +188,80 @@ describe("scm-github plugin", () => { }); }); + // ---- resolvePR --------------------------------------------------------- + + describe("resolvePR", () => { + it("resolves a PR number into canonical PR info", async () => { + mockGh({ + number: 42, + url: "https://github.com/acme/repo/pull/42", + title: "feat: add feature", + headRefName: "feat/my-feature", + baseRefName: "main", + isDraft: false, + }); + + await expect(scm.resolvePR?.("42", project)).resolves.toEqual(pr); + }); + }); + + // ---- assignPRToCurrentUser -------------------------------------------- + + describe("assignPRToCurrentUser", () => { + it("assigns the PR to the authenticated user", async () => { + ghMock.mockResolvedValueOnce({ stdout: "" }); + + await scm.assignPRToCurrentUser?.(pr); + + expect(ghMock).toHaveBeenCalledWith( + "gh", + ["pr", "edit", "42", "--repo", "acme/repo", "--add-assignee", "@me"], + expect.any(Object), + ); + }); + }); + + // ---- checkoutPR -------------------------------------------------------- + + describe("checkoutPR", () => { + it("returns false when already on the PR branch", async () => { + ghMock.mockResolvedValueOnce({ stdout: "feat/my-feature\n" }); + + await expect(scm.checkoutPR?.(pr, "/tmp/repo")).resolves.toBe(false); + + expect(ghMock).toHaveBeenCalledTimes(1); + expect(ghMock).toHaveBeenCalledWith( + "git", + ["branch", "--show-current"], + expect.objectContaining({ cwd: "/tmp/repo" }), + ); + }); + + it("throws when switching branches would discard local changes", async () => { + ghMock.mockResolvedValueOnce({ stdout: "main\n" }); + ghMock.mockResolvedValueOnce({ stdout: " M src/index.ts\n" }); + + await expect(scm.checkoutPR?.(pr, "/tmp/repo")).rejects.toThrow( + 'Workspace has uncommitted changes; cannot switch to PR branch "feat/my-feature" safely', + ); + }); + + it("checks out the PR when the workspace is clean", async () => { + ghMock.mockResolvedValueOnce({ stdout: "main\n" }); + ghMock.mockResolvedValueOnce({ stdout: "" }); + ghMock.mockResolvedValueOnce({ stdout: "" }); + + await expect(scm.checkoutPR?.(pr, "/tmp/repo")).resolves.toBe(true); + + expect(ghMock).toHaveBeenNthCalledWith( + 3, + "gh", + ["pr", "checkout", "42", "--repo", "acme/repo"], + expect.objectContaining({ cwd: "/tmp/repo" }), + ); + }); + }); + // ---- mergePR ----------------------------------------------------------- describe("mergePR", () => {