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 <harsh@Ubuntu-24-Forrest.lan>
This commit is contained in:
Harsh Batheja 2026-03-10 11:54:26 +05:30 committed by deepak
parent ffae671884
commit aefa6ef001
8 changed files with 211 additions and 86 deletions

View File

@ -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,
});
});

View File

@ -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<string, unknown> | 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<string, unknown> | 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,
});
});

View File

@ -197,12 +197,11 @@ export function registerSession(program: Command): void {
.argument("<pr>", "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}.`));

View File

@ -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 <name>", "Override the agent plugin (e.g. codex, claude-code)")
.option("--claim-pr <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)}`),
);
}
}

View File

@ -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", () => {

View File

@ -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.

View File

@ -1735,6 +1735,25 @@ export function createSessionManager(deps: SessionManagerDeps): OpenCodeSessionM
}
}
/**
* Claim an existing PR for a session.
*
* ## Ownership Model (Asymmetric)
*
* - **RULE A (Exclusive PRAgent)**: 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 (AgentMany 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) {

View File

@ -1083,7 +1083,6 @@ export interface OpenCodeSessionManager extends SessionManager {
export interface ClaimPROptions {
assignOnGithub?: boolean;
takeover?: boolean;
}
export interface ClaimPRResult {