From aefa6ef001e0f4f87262e154aa3ae957ec157cd6 Mon Sep 17 00:00:00 2001 From: Harsh Batheja <40922251+harsh-batheja@users.noreply.github.com> Date: Tue, 10 Mar 2026 11:54:26 +0530 Subject: [PATCH] fix(core): enforce single-owner PR claim consolidation (#390) * fix(core): enforce single-owner PR claim consolidation * refactor: remove dead `takeover` option from claimPR Since PR consolidation is now automatic (single-owner enforcement), the `--takeover` flag became dead code. Users passing it got no error but it had zero effect. This commit: - Removes takeover from ClaimPROptions interface - Removes --takeover flag from spawn and session CLI commands - Updates orchestrator-prompt.ts examples - Updates tests to reflect automatic consolidation Tests: ao-core 403 passing, ao-cli 189 passing (spawn+session) * fix: remove stale --takeover from Quick Start example Remove remaining --takeover reference in orchestrator prompt Quick Start section. * feat(core): document and test asymmetric PR ownership model - Add JSDoc to claimPR documenting RULE A (exclusive PR->Agent) and RULE B (Agent->Many PRs) ownership contract - Add tests for: - Same session claiming multiple PRs sequentially (RULE B) - Idempotent re-claim by same owner - Stale/dead prior owner handoff - Exclusive PR ownership enforcement (RULE A) 139 tests passing --------- Co-authored-by: Harsh --- .../cli/__tests__/commands/session.test.ts | 3 - packages/cli/__tests__/commands/spawn.test.ts | 50 ++--- packages/cli/src/commands/session.ts | 4 +- packages/cli/src/commands/spawn.ts | 15 +- .../src/__tests__/session-manager.test.ts | 196 +++++++++++++++--- packages/core/src/orchestrator-prompt.ts | 4 +- packages/core/src/session-manager.ts | 24 ++- packages/core/src/types.ts | 1 - 8 files changed, 211 insertions(+), 86 deletions(-) diff --git a/packages/cli/__tests__/commands/session.test.ts b/packages/cli/__tests__/commands/session.test.ts index 936e5e8e8..7387160ff 100644 --- a/packages/cli/__tests__/commands/session.test.ts +++ b/packages/cli/__tests__/commands/session.test.ts @@ -445,12 +445,10 @@ describe("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"); @@ -465,7 +463,6 @@ describe("session claim-pr", () => { expect(mockSessionManager.claimPR).toHaveBeenCalledWith("app-7", "42", { assignOnGithub: undefined, - takeover: undefined, }); }); diff --git a/packages/cli/__tests__/commands/spawn.test.ts b/packages/cli/__tests__/commands/spawn.test.ts index 943fd711e..c9627f073 100644 --- a/packages/cli/__tests__/commands/spawn.test.ts +++ b/packages/cli/__tests__/commands/spawn.test.ts @@ -4,21 +4,23 @@ import { join } from "node:path"; import { tmpdir } from "node:os"; import { type Session, type SessionManager, getProjectBaseDir } from "@composio/ao-core"; -const { mockExec, mockConfigRef, mockSessionManager, mockEnsureLifecycleWorker } = vi.hoisted(() => ({ - mockExec: 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(), - }, - mockEnsureLifecycleWorker: vi.fn(), -})); +const { mockExec, mockConfigRef, mockSessionManager, mockEnsureLifecycleWorker } = vi.hoisted( + () => ({ + mockExec: 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(), + }, + mockEnsureLifecycleWorker: vi.fn(), + }), +); vi.mock("../../src/lib/shell.js", () => ({ tmux: vi.fn(), @@ -118,8 +120,8 @@ beforeEach(() => { running: true, started: true, pid: 12345, - pidFile: '/tmp/lifecycle-worker.pid', - logFile: '/tmp/lifecycle-worker.log', + pidFile: "/tmp/lifecycle-worker.pid", + logFile: "/tmp/lifecycle-worker.log", }); }); @@ -324,7 +326,6 @@ describe("spawn command", () => { ); }); - it("claims a PR for the spawned session when --claim-pr is provided", async () => { const fakeSession: Session = { id: "app-1", @@ -370,7 +371,6 @@ describe("spawn command", () => { }); expect(mockSessionManager.claimPR).toHaveBeenCalledWith("app-1", "123", { assignOnGithub: undefined, - takeover: undefined, }); const output = consoleSpy.mock.calls.map((c) => String(c[0])).join("\n"); @@ -378,7 +378,7 @@ describe("spawn command", () => { expect(output).toContain("feat/claimed-pr"); }); - it("passes takeover and GitHub assignment flags through to claimPR", async () => { + it("passes GitHub assignment flag through to claimPR", async () => { const fakeSession: Session = { id: "app-1", projectId: "my-app", @@ -421,26 +421,24 @@ describe("spawn command", () => { "my-app", "--claim-pr", "123", - "--takeover", "--assign-on-github", ]); expect(mockSessionManager.claimPR).toHaveBeenCalledWith("app-1", "123", { assignOnGithub: true, - takeover: true, }); }); - it("rejects claim-only flags without --claim-pr", async () => { + it("rejects --assign-on-github without --claim-pr", async () => { await expect( - program.parseAsync(["node", "test", "spawn", "my-app", "--takeover"]), + program.parseAsync(["node", "test", "spawn", "my-app", "--assign-on-github"]), ).rejects.toThrow("process.exit(1)"); const errors = vi .mocked(console.error) .mock.calls.map((c) => String(c[0])) .join("\n"); - expect(errors).toContain("require --claim-pr"); + expect(errors).toContain("--assign-on-github requires --claim-pr"); expect(mockSessionManager.spawn).not.toHaveBeenCalled(); expect(mockSessionManager.claimPR).not.toHaveBeenCalled(); }); @@ -580,7 +578,6 @@ describe("spawn pre-flight checks", () => { expect(mockSessionManager.spawn).not.toHaveBeenCalled(); }); - it("handles tracker+scm github preflight when claiming during spawn", async () => { const fakeSession: Session = { id: "app-1", @@ -637,7 +634,6 @@ describe("spawn pre-flight checks", () => { expect(mockSessionManager.spawn).toHaveBeenCalled(); expect(mockSessionManager.claimPR).toHaveBeenCalledWith("app-1", "123", { assignOnGithub: undefined, - takeover: undefined, }); }); diff --git a/packages/cli/src/commands/session.ts b/packages/cli/src/commands/session.ts index 8e08a62fe..a6c00b19b 100644 --- a/packages/cli/src/commands/session.ts +++ b/packages/cli/src/commands/session.ts @@ -197,12 +197,11 @@ export function registerSession(program: Command): void { .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 }, + opts: { assignOnGithub?: boolean }, ) => { const config = loadConfig(); const resolvedSession = @@ -222,7 +221,6 @@ export function registerSession(program: Command): void { 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}.`)); diff --git a/packages/cli/src/commands/spawn.ts b/packages/cli/src/commands/spawn.ts index 1b2d10a56..2829598ef 100644 --- a/packages/cli/src/commands/spawn.ts +++ b/packages/cli/src/commands/spawn.ts @@ -11,7 +11,6 @@ import { preflight } from "../lib/preflight.js"; interface SpawnClaimOptions { claimPr?: string; assignOnGithub?: boolean; - takeover?: boolean; } /** @@ -65,7 +64,6 @@ async function spawnSession( try { const claimResult = await sm.claimPR(session.id, claimOptions.claimPr, { assignOnGithub: claimOptions.assignOnGithub, - takeover: claimOptions.takeover, }); branchStr = claimResult.pr.branch; claimedPrUrl = claimResult.pr.url; @@ -120,7 +118,6 @@ export function registerSpawn(program: Command): void { .option("--agent ", "Override the agent plugin (e.g. codex, claude-code)") .option("--claim-pr ", "Immediately claim an existing PR for the spawned session") .option("--assign-on-github", "Assign the claimed PR to the authenticated GitHub user") - .option("--takeover", "Transfer PR ownership from another AO session if needed") .action( async ( projectId: string, @@ -130,7 +127,6 @@ export function registerSpawn(program: Command): void { agent?: string; claimPr?: string; assignOnGithub?: boolean; - takeover?: boolean; }, ) => { const config = loadConfig(); @@ -143,10 +139,8 @@ export function registerSpawn(program: Command): void { process.exit(1); } - if (!opts.claimPr && (opts.assignOnGithub || opts.takeover)) { - console.error( - chalk.red("--assign-on-github and --takeover require --claim-pr on `ao spawn`."), - ); + if (!opts.claimPr && opts.assignOnGithub) { + console.error(chalk.red("--assign-on-github requires --claim-pr on `ao spawn`.")); process.exit(1); } @@ -156,7 +150,6 @@ export function registerSpawn(program: Command): void { await spawnSession(config, projectId, issueId, opts.open, opts.agent, { claimPr: opts.claimPr, assignOnGithub: opts.assignOnGithub, - takeover: opts.takeover, }); } catch (err) { console.error(chalk.red(`✗ ${err instanceof Error ? err.message : String(err)}`)); @@ -251,9 +244,7 @@ export function registerBatchSpawn(program: Command): void { error: err instanceof Error ? err.message : String(err), }); console.log( - chalk.red( - ` Failed ${issue} — ${err instanceof Error ? err.message : String(err)}`, - ), + chalk.red(` Failed ${issue} — ${err instanceof Error ? err.message : String(err)}`), ); } } diff --git a/packages/core/src/__tests__/session-manager.test.ts b/packages/core/src/__tests__/session-manager.test.ts index d4f8f1d32..9fd037ef8 100644 --- a/packages/core/src/__tests__/session-manager.test.ts +++ b/packages/core/src/__tests__/session-manager.test.ts @@ -3675,38 +3675,7 @@ describe("claimPR", () => { 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 () => { + it("automatically consolidates ownership when another session tracks the PR", async () => { const mockSCM = makeSCM(); writeMetadata(sessionsDir, "app-1", { @@ -3727,8 +3696,18 @@ describe("claimPR", () => { }); const sm = createSessionManager({ config, registry: registryWithSCM(mockSCM) }); + const result = await sm.claimPR("app-2", "42"); - await expect(sm.claimPR("app-2", "42")).rejects.toThrow("already tracked by app-1"); + expect(result.takenOverFrom).toContain("app-1"); + expect(result.pr.number).toBe(42); + + const app2 = readMetadataRaw(sessionsDir, "app-2"); + expect(app2!["pr"]).toBe("https://github.com/org/my-app/pull/42"); + expect(app2!["status"]).toBe("pr_open"); + + const app1 = readMetadataRaw(sessionsDir, "app-1"); + expect(app1!["pr"] ?? "").toBe(""); + expect(app1!["status"]).toBe("working"); }); it("keeps AO metadata updated even if GitHub assignment fails", async () => { @@ -3754,6 +3733,157 @@ describe("claimPR", () => { expect(raw!["pr"]).toBe("https://github.com/org/my-app/pull/42"); expect(raw!["status"]).toBe("pr_open"); }); + + // RULE B: One session may own multiple PRs sequentially (switching ownership) + it("allows same session to claim different PRs sequentially without rejection", async () => { + const mockSCM = makeSCM({ + resolvePR: vi + .fn() + .mockResolvedValueOnce({ + number: 42, + url: "https://github.com/org/my-app/pull/42", + title: "First PR", + owner: "org", + repo: "my-app", + branch: "feat/first-pr", + baseBranch: "main", + isDraft: false, + }) + .mockResolvedValueOnce({ + number: 99, + url: "https://github.com/org/my-app/pull/99", + title: "Second PR", + owner: "org", + repo: "my-app", + branch: "feat/second-pr", + baseBranch: "main", + isDraft: false, + }), + checkoutPR: vi.fn().mockResolvedValue(true), + }); + + writeMetadata(sessionsDir, "app-1", { + worktree: "/tmp/ws-app-1", + branch: "feat/initial", + status: "working", + project: "my-app", + runtimeHandle: JSON.stringify(makeHandle("rt-1")), + }); + + const sm = createSessionManager({ config, registry: registryWithSCM(mockSCM) }); + + // Claim first PR + const result1 = await sm.claimPR("app-1", "42"); + expect(result1.pr.number).toBe(42); + expect(result1.takenOverFrom).toEqual([]); + + let raw = readMetadataRaw(sessionsDir, "app-1"); + expect(raw!["pr"]).toBe("https://github.com/org/my-app/pull/42"); + + // Claim second PR (switches ownership, no rejection) + const result2 = await sm.claimPR("app-1", "99"); + expect(result2.pr.number).toBe(99); + expect(result2.takenOverFrom).toEqual([]); + + raw = readMetadataRaw(sessionsDir, "app-1"); + expect(raw!["pr"]).toBe("https://github.com/org/my-app/pull/99"); + expect(raw!["branch"]).toBe("feat/second-pr"); + }); + + // Idempotent re-claim by same owner + it("handles idempotent re-claim of same PR by same session", 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")), + }); + + const sm = createSessionManager({ config, registry: registryWithSCM(mockSCM) }); + + // Re-claim same PR - should succeed without consolidation + const result = await sm.claimPR("app-1", "42"); + expect(result.pr.number).toBe(42); + expect(result.takenOverFrom).toEqual([]); + + const raw = readMetadataRaw(sessionsDir, "app-1"); + expect(raw!["pr"]).toBe("https://github.com/org/my-app/pull/42"); + }); + + // Stale/dead prior owner handoff + it("consolidates from stale/dead prior owner regardless of status", async () => { + const mockSCM = makeSCM(); + + // Prior owner in "spawning" state (stuck/dead) + writeMetadata(sessionsDir, "app-stale", { + worktree: "/tmp/ws-app-stale", + branch: "feat/existing-pr", + status: "spawning", // Stuck in spawning + project: "my-app", + pr: "https://github.com/org/my-app/pull/42", + runtimeHandle: JSON.stringify(makeHandle("rt-stale")), + }); + + writeMetadata(sessionsDir, "app-2", { + worktree: "/tmp/ws-app-2", + branch: "feat/other", + 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"); + + // Consolidation happens regardless of prior owner's status + expect(result.takenOverFrom).toContain("app-stale"); + expect(result.pr.number).toBe(42); + + // Prior owner is displaced + const staleRaw = readMetadataRaw(sessionsDir, "app-stale"); + expect(staleRaw!["pr"] ?? "").toBe(""); + expect(staleRaw!["status"]).toBe("spawning"); // Status unchanged (not a PR-tracking status) + }); + + // RULE A: Exclusive PR->agent mapping - explicit test + it("ensures exclusive PR ownership (only one active owner per PR)", async () => { + const mockSCM = makeSCM(); + + // First session owns the PR + writeMetadata(sessionsDir, "app-owner", { + worktree: "/tmp/ws-owner", + branch: "feat/existing-pr", + status: "pr_open", + project: "my-app", + pr: "https://github.com/org/my-app/pull/42", + runtimeHandle: JSON.stringify(makeHandle("rt-owner")), + }); + + // Second session wants to claim the same PR + writeMetadata(sessionsDir, "app-new", { + worktree: "/tmp/ws-new", + branch: "feat/other", + status: "working", + project: "my-app", + runtimeHandle: JSON.stringify(makeHandle("rt-new")), + }); + + const sm = createSessionManager({ config, registry: registryWithSCM(mockSCM) }); + const result = await sm.claimPR("app-new", "42"); + + // New owner succeeds, old owner is displaced + expect(result.takenOverFrom).toEqual(["app-owner"]); + + const newOwner = readMetadataRaw(sessionsDir, "app-new"); + expect(newOwner!["pr"]).toBe("https://github.com/org/my-app/pull/42"); + + const oldOwner = readMetadataRaw(sessionsDir, "app-owner"); + expect(oldOwner!["pr"] ?? "").toBe(""); + }); }); describe("PluginRegistry.loadBuiltins importFn", () => { diff --git a/packages/core/src/orchestrator-prompt.ts b/packages/core/src/orchestrator-prompt.ts index 291b5ffe8..354d04dbe 100644 --- a/packages/core/src/orchestrator-prompt.ts +++ b/packages/core/src/orchestrator-prompt.ts @@ -48,7 +48,7 @@ ao status # Spawn sessions for issues (GitHub: #123, Linear: INT-1234, etc.) ao spawn ${projectId} INT-1234 -ao spawn ${projectId} --claim-pr 123 --takeover +ao spawn ${projectId} --claim-pr 123 ao batch-spawn ${projectId} INT-1 INT-2 INT-3 # List sessions @@ -118,7 +118,7 @@ If a session needs to continue work on an existing PR: \`\`\`bash ao session claim-pr 123 ${project.sessionPrefix}-1 # or do it at spawn time -ao spawn ${projectId} --claim-pr 123 --takeover +ao spawn ${projectId} --claim-pr 123 \`\`\` This updates AO metadata, switches the worktree onto the PR branch, and lets lifecycle reactions keep routing CI and review feedback to that session. diff --git a/packages/core/src/session-manager.ts b/packages/core/src/session-manager.ts index 112db8c07..d2fa9387b 100644 --- a/packages/core/src/session-manager.ts +++ b/packages/core/src/session-manager.ts @@ -1735,6 +1735,25 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM } } + /** + * Claim an existing PR for a session. + * + * ## Ownership Model (Asymmetric) + * + * - **RULE A (Exclusive PR→Agent)**: One PR can be actively owned by only one + * session at a time. If another session claims a PR already owned, the + * previous owner is automatically displaced (consolidation). + * + * - **RULE B (Agent→Many PRs)**: One session may claim different PRs sequentially. + * Switching to a new PR releases ownership of the previous PR. + * + * ## Behavior + * + * - Idempotent: re-claiming the same PR by the same owner succeeds without + * triggering consolidation. + * - Consolidation happens regardless of the previous owner's status (includes + * stale/dead sessions). + */ async function claimPR( sessionId: SessionId, prRef: string, @@ -1779,11 +1798,6 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM } 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) { diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index a2b93c093..ccbdf0104 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -1083,7 +1083,6 @@ export interface OpenCodeSessionManager extends SessionManager { export interface ClaimPROptions { assignOnGithub?: boolean; - takeover?: boolean; } export interface ClaimPRResult {