fix: address #1306 review follow-ups

This commit is contained in:
harshitsinghbhandari 2026-04-19 14:36:36 +05:30
parent aede68e8da
commit bb5dcf0af7
8 changed files with 107 additions and 34 deletions

View File

@ -1,10 +1,8 @@
"@aoagents/ao-cli": patch
"@aoagents/ao-web": patch
"@aoagents/ao-plugin-agent-codex": patch
---
Fix restore behavior across AO session recovery flows.
- restore the latest dead-but-restorable orchestrator on `ao start` instead of silently spawning a new orchestrator when tmux is gone
- make worker session orchestrator navigation prefer the most recently active live orchestrator for the project
- make Codex worker restore commands resume with `--ask-for-approval on-request` instead of `never`, while keeping permissionless orchestrator restores bypassed

View File

@ -57,8 +57,10 @@ const { mockPromptSelect, mockPromptConfirm } = vi.hoisted(() => ({
mockPromptConfirm: vi.fn().mockResolvedValue(true),
}));
const { mockIsAlreadyRunning, mockUnregister, mockWaitForExit } = vi.hoisted(() => ({
const { mockAcquireStartupLock, mockIsAlreadyRunning, mockRegister, mockUnregister, mockWaitForExit } = vi.hoisted(() => ({
mockAcquireStartupLock: vi.fn().mockResolvedValue(() => {}),
mockIsAlreadyRunning: vi.fn().mockReturnValue(null),
mockRegister: vi.fn(),
mockUnregister: vi.fn(),
mockWaitForExit: vi.fn().mockReturnValue(true),
}));
@ -142,7 +144,8 @@ vi.mock("../../src/lib/preflight.js", () => ({
}));
vi.mock("../../src/lib/running-state.js", () => ({
register: vi.fn(),
acquireStartupLock: (...args: unknown[]) => mockAcquireStartupLock(...args),
register: (...args: unknown[]) => mockRegister(...args),
unregister: (...args: unknown[]) => mockUnregister(...args),
isAlreadyRunning: (...args: unknown[]) => mockIsAlreadyRunning(...args),
getRunning: vi.fn().mockReturnValue(null),
@ -283,8 +286,12 @@ beforeEach(async () => {
mockPromptSelect.mockReset();
mockPromptConfirm.mockReset();
mockPromptConfirm.mockResolvedValue(true);
mockAcquireStartupLock.mockReset();
mockAcquireStartupLock.mockResolvedValue(() => {});
mockIsAlreadyRunning.mockReset();
mockIsAlreadyRunning.mockResolvedValue(null);
mockRegister.mockReset();
mockRegister.mockResolvedValue(undefined);
mockUnregister.mockReset();
mockWaitForExit.mockReset();
mockWaitForExit.mockResolvedValue(true);
@ -1321,6 +1328,22 @@ describe("start command — orchestrator session strategy display", () => {
expect(fakeDashboard.kill).toHaveBeenCalled();
});
it("reports startup lock acquisition failures through the normal CLI error path", async () => {
mockConfigRef.current = makeConfig({ "my-app": makeProject() });
mockAcquireStartupLock.mockRejectedValueOnce(
new Error("Could not acquire startup lock (/tmp/startup.lock)"),
);
await expect(program.parseAsync(["node", "test", "start"])).rejects.toThrow("process.exit(1)");
const errors = vi
.mocked(console.error)
.mock.calls.map((c) => c.join(" "))
.join("\n");
expect(errors).toContain("Could not acquire startup lock (/tmp/startup.lock)");
expect(mockIsAlreadyRunning).not.toHaveBeenCalled();
});
it("fails and cleans up dashboard when sm.restore throws on a killed orchestrator", async () => {
mockConfigRef.current = makeConfig({ "my-app": makeProject() });

View File

@ -52,7 +52,14 @@ import {
} from "../lib/web-dir.js";
import { rebuildDashboardProductionArtifacts } from "../lib/dashboard-rebuild.js";
import { preflight } from "../lib/preflight.js";
import { register, unregister, isAlreadyRunning, getRunning, waitForExit } from "../lib/running-state.js";
import {
register,
unregister,
isAlreadyRunning,
getRunning,
waitForExit,
acquireStartupLock,
} from "../lib/running-state.js";
import { preventIdleSleep } from "../lib/prevent-sleep.js";
import { isHumanCaller } from "../lib/caller-context.js";
import { detectEnvironment } from "../lib/detect-env.js";
@ -1379,7 +1386,16 @@ export function registerStart(program: Command): void {
interactive?: boolean;
},
) => {
let releaseStartupLock: (() => void) | undefined;
let startupLockReleased = false;
const unlockStartup = (): void => {
if (startupLockReleased || !releaseStartupLock) return;
startupLockReleased = true;
releaseStartupLock();
};
try {
releaseStartupLock = await acquireStartupLock();
let config: OrchestratorConfig;
let projectId: string;
let project: ProjectConfig;
@ -1408,6 +1424,7 @@ export function registerStart(program: Command): void {
if (choice === "open") {
const url = `http://localhost:${running.port}`;
openUrl(url);
unlockStartup();
process.exit(0);
} else if (choice === "new") {
// Defer config mutation until after config is loaded below
@ -1427,6 +1444,7 @@ export function registerStart(program: Command): void {
console.log(chalk.yellow("\n Stopped existing instance. Restarting...\n"));
// Continue to startup below
} else {
unlockStartup();
process.exit(0);
}
} else {
@ -1436,6 +1454,7 @@ export function registerStart(program: Command): void {
console.log(`PID: ${running.pid}`);
console.log(`Projects: ${running.projects.join(", ")}`);
console.log(`To restart: ao stop && ao start`);
unlockStartup();
process.exit(0);
}
}
@ -1565,6 +1584,7 @@ export function registerStart(program: Command): void {
startedAt: new Date().toISOString(),
projects: [projectId],
});
unlockStartup();
// Install shutdown handlers so `ao stop` (which sends SIGTERM to
// this pid) flushes lifecycle health state before exit. Handlers
@ -1591,6 +1611,8 @@ export function registerStart(program: Command): void {
console.error(chalk.red("\nError:"), String(err));
}
process.exit(1);
} finally {
unlockStartup();
}
},
);

View File

@ -23,6 +23,7 @@ const STATE_DIR = join(homedir(), ".agent-orchestrator");
const STATE_FILE = join(STATE_DIR, "running.json");
const STATE_LOCK_FILE = join(STATE_DIR, "running.lock");
const STARTUP_LOCK_FILE = join(STATE_DIR, "startup.lock");
const MAX_LOCK_AGE_MS = 30 * 60 * 1000;
interface LockMetadata {
pid: number;
@ -60,6 +61,11 @@ function readLockMetadata(lockFile: string): LockMetadata | null {
}
}
function isStaleLockOwner(owner: LockMetadata): boolean {
const acquiredAt = Date.parse(owner.acquiredAt);
return Number.isFinite(acquiredAt) && Date.now() - acquiredAt > MAX_LOCK_AGE_MS;
}
/** Try to create the lockfile atomically. Returns a release function on success, null on failure. */
function tryAcquire(lockFile: string): (() => void) | null {
try {
@ -104,14 +110,14 @@ async function acquireLock(
if (release) return release;
const owner = readLockMetadata(lockFile);
if (!owner || !isProcessAlive(owner.pid)) {
if (owner && (isStaleLockOwner(owner) || !isProcessAlive(owner.pid))) {
try { unlinkSync(lockFile); } catch { /* ignore */ }
const retryRelease = tryAcquire(lockFile);
if (retryRelease) return retryRelease;
}
if (Date.now() - start > timeoutMs) {
throw new Error(`Could not acquire ${resourceName}`);
throw new Error(`Could not acquire ${resourceName} (${lockFile})`);
}
// Jittered backoff: 30-70ms base, growing with attempts (capped at 200ms)

View File

@ -1470,7 +1470,7 @@ describe("getRestoreCommand", () => {
expect(cmd).toContain("--dangerously-bypass-approvals-and-sandbox");
});
it("demotes worker restore permissionless mode to ask-for-approval on-request", async () => {
it("keeps worker restore permissionless mode at ask-for-approval never", async () => {
const content = jsonl(
{ type: "session_meta", cwd: "/workspace/test", model: "gpt-4o" },
{ threadId: "thread-1" },
@ -1489,11 +1489,11 @@ describe("getRestoreCommand", () => {
}),
);
expect(cmd).toContain("--ask-for-approval on-request");
expect(cmd).toContain("--ask-for-approval never");
expect(cmd).not.toContain("--dangerously-bypass-approvals-and-sandbox");
});
it("downgrades auto-edit restore policy to ask-for-approval on-request", async () => {
it("keeps auto-edit restore policy at ask-for-approval never", async () => {
const content = jsonl(
{ type: "session_meta", cwd: "/workspace/test", model: "gpt-4o" },
{ threadId: "thread-1" },
@ -1512,7 +1512,7 @@ describe("getRestoreCommand", () => {
}),
);
expect(cmd).toContain("--ask-for-approval on-request");
expect(cmd).toContain("--ask-for-approval never");
expect(cmd).not.toContain("--dangerously-bypass-approvals-and-sandbox");
});

View File

@ -404,26 +404,6 @@ function appendApprovalFlags(
}
}
/** Append restore-time approval flags, avoiding worker resumes in `never` mode. */
function appendRestoreApprovalFlags(
parts: string[],
permissions: string | undefined,
isOrchestrator: boolean,
): void {
const mode = normalizeAgentPermissionMode(permissions);
if (mode === "permissionless") {
if (isOrchestrator) {
parts.push("--dangerously-bypass-approvals-and-sandbox");
} else {
parts.push("--ask-for-approval", "on-request");
}
} else if (mode === "auto-edit") {
parts.push("--ask-for-approval", "on-request");
} else if (mode === "suggest") {
parts.push("--ask-for-approval", "untrusted");
}
}
/** Append model and reasoning flags to a command parts array */
function appendModelFlags(parts: string[], model: string | undefined): void {
if (!model) return;
@ -754,7 +734,7 @@ function createCodexAgent(): Agent {
appendNoUpdateCheckFlag(parts);
const isOrchestrator = session.metadata?.["role"] === "orchestrator";
appendRestoreApprovalFlags(parts, project.agentConfig?.permissions, isOrchestrator);
appendApprovalFlags(parts, project.agentConfig?.permissions, isOrchestrator);
const effectiveModel = (project.agentConfig?.model ?? data.model) as string | undefined;
appendModelFlags(parts, effectiveModel ?? undefined);

View File

@ -387,7 +387,6 @@ describe("API Routes", () => {
expect(data.orchestratorId).toBe("my-app-orchestrator-2");
expect(data.orchestrators.map((session: { id: string }) => session.id)).toEqual([
"my-app-orchestrator-0",
"my-app-orchestrator-1",
"my-app-orchestrator-2",
]);
@ -395,6 +394,39 @@ describe("API Routes", () => {
expect(mockSessionManager.listCached).toHaveBeenCalledWith("my-app");
});
it("omits dead orchestrators from project-scoped orchestrator payloads when none are live", async () => {
const deadLifecycle = createInitialCanonicalLifecycle("orchestrator", new Date("2026-04-19T11:00:00.000Z"));
deadLifecycle.session.state = "terminated";
deadLifecycle.session.reason = "runtime_missing";
deadLifecycle.session.terminatedAt = "2026-04-19T11:00:00.000Z";
deadLifecycle.session.lastTransitionAt = "2026-04-19T11:00:00.000Z";
deadLifecycle.runtime.state = "missing";
deadLifecycle.runtime.reason = "process_missing";
deadLifecycle.runtime.lastObservedAt = "2026-04-19T11:00:00.000Z";
(mockSessionManager.listCached as ReturnType<typeof vi.fn>).mockResolvedValueOnce([
makeSession({
id: "my-app-orchestrator-0",
projectId: "my-app",
metadata: { role: "orchestrator" },
status: "killed",
activity: "exited",
lastActivityAt: new Date("2026-04-19T11:00:00.000Z"),
lifecycle: deadLifecycle,
}),
]);
const res = await sessionsGET(
makeRequest("http://localhost:3000/api/sessions?project=my-app&orchestratorOnly=true"),
);
expect(res.status).toBe(200);
const data = await res.json();
expect(data.orchestratorId).toBeNull();
expect(data.orchestrators).toEqual([]);
expect(data.sessions).toEqual([]);
});
it("enriches all PRs concurrently, not sequentially", async () => {
vi.useFakeTimers();

View File

@ -42,6 +42,16 @@ function selectPreferredOrchestratorId(
return liveOrchestrators[0]?.id ?? null;
}
function listLiveDashboardOrchestrators(
sessions: Parameters<typeof listDashboardOrchestrators>[0],
projects: Parameters<typeof listDashboardOrchestrators>[1],
) {
return listDashboardOrchestrators(
sessions.filter((session) => !isTerminalSession(session)),
projects,
);
}
export async function GET(request: Request) {
const correlationId = getCorrelationId(request);
const startedAt = Date.now();
@ -58,7 +68,9 @@ export async function GET(request: Request) {
: undefined;
const coreSessions = await sessionManager.listCached(requestedProjectId);
const visibleSessions = filterProjectSessions(coreSessions, projectFilter, config.projects);
const orchestrators = listDashboardOrchestrators(visibleSessions, config.projects);
const orchestrators = requestedProjectId
? listLiveDashboardOrchestrators(visibleSessions, config.projects)
: listDashboardOrchestrators(visibleSessions, config.projects);
const orchestratorId = requestedProjectId
? selectPreferredOrchestratorId(visibleSessions, config.projects)
: (orchestrators.length === 1 ? (orchestrators[0]?.id ?? null) : null);