Fix running-agent replacement lifecycle

This commit is contained in:
Mohamed Boudra
2026-04-22 17:53:53 +07:00
parent 9218f70df3
commit 019bf65790
2 changed files with 53 additions and 7 deletions
@@ -2253,6 +2253,8 @@ describe("AgentManager", () => {
const storagePath = join(workdir, "agents");
const storage = new AgentStorage(storagePath, logger);
const secondStartEntered = deferred<void>();
const interruptStarted = deferred<void>();
const allowInterruptToFinish = deferred<void>();
const allowSecondStartToResolve = deferred<void>();
let capturedSession: StaleReplacementSession | null = null;
@@ -2280,6 +2282,8 @@ describe("AgentManager", () => {
}
override async interrupt(): Promise<void> {
interruptStarted.resolve();
await allowInterruptToFinish.promise;
this.pushEvent({
type: "turn_canceled",
provider: this.provider,
@@ -2310,13 +2314,16 @@ describe("AgentManager", () => {
cwd: workdir,
});
const lifecycleUpdates: string[] = [];
const stateUpdates: Array<{ lifecycle: string; updatedAt: number }> = [];
const unsubscribe = manager.subscribe(
(event) => {
if (event.type !== "agent_state" || event.agent.id !== snapshot.id) {
return;
}
lifecycleUpdates.push(event.agent.lifecycle);
stateUpdates.push({
lifecycle: event.agent.lifecycle,
updatedAt: event.agent.updatedAt.getTime(),
});
},
{ agentId: snapshot.id, replayState: false },
);
@@ -2330,6 +2337,8 @@ describe("AgentManager", () => {
await manager.waitForAgentRunStart(snapshot.id);
const replaceUpdatesStart = stateUpdates.length;
const beforeReplaceUpdatedAt = manager.getAgent(snapshot.id)?.updatedAt.getTime() ?? 0;
const secondRun = manager.replaceAgentRun(snapshot.id, "replacement run");
const secondRunDrain = (async () => {
for await (const _event of secondRun) {
@@ -2337,6 +2346,16 @@ describe("AgentManager", () => {
}
})();
await interruptStarted.promise;
const replacementUpdates = stateUpdates.slice(replaceUpdatesStart);
expect(
replacementUpdates.some(
(update) => update.lifecycle === "running" && update.updatedAt > beforeReplaceUpdatedAt,
),
).toBe(true);
expect(replacementUpdates.map((update) => update.lifecycle)).not.toContain("idle");
allowInterruptToFinish.resolve();
await secondStartEntered.promise;
const replaceGapSnapshot = manager.getAgent(snapshot.id) as
@@ -2349,10 +2368,11 @@ describe("AgentManager", () => {
capturedSession!.pushEvent({ type: "turn_completed", provider: "codex", turnId: "turn-1" });
await new Promise((resolve) => setTimeout(resolve, 0));
// eslint-disable-next-line no-console
console.log("replace-gap-lifecycle", lifecycleUpdates, manager.getAgent(snapshot.id));
expect(manager.getAgent(snapshot.id)?.lifecycle).toBe("running");
expect(lifecycleUpdates.at(-1)).toBe("running");
expect(stateUpdates.at(-1)?.lifecycle).toBe("running");
expect(stateUpdates.slice(replaceUpdatesStart).map((update) => update.lifecycle)).not.toContain(
"idle",
);
allowSecondStartToResolve.resolve();
@@ -187,6 +187,7 @@ type ManagedAgentBase = {
lastError?: string;
attention: AttentionState;
foregroundTurnWaiters: Set<ForegroundTurnWaiter>;
finalizedForegroundTurnIds: Set<string>;
unsubscribeSession: (() => void) | null;
/**
* Internal agents are hidden from listings and don't trigger notifications.
@@ -1302,8 +1303,11 @@ export class AgentManager {
return streamForwarder;
}
private finalizeForegroundTurn(agent: ActiveManagedAgent): void {
private finalizeForegroundTurn(agent: ActiveManagedAgent, turnId?: string): void {
const mutableAgent = agent as ActiveManagedAgent;
if (turnId) {
this.rememberFinalizedForegroundTurn(mutableAgent, turnId);
}
mutableAgent.activeForegroundTurnId = null;
const terminalError = mutableAgent.lastError;
const shouldHoldBusyForReplacement = mutableAgent.pendingReplacement && !terminalError;
@@ -1351,6 +1355,9 @@ export class AgentManager {
const agent = this.requireSessionAgent(agentId);
agent.pendingReplacement = true;
agent.lifecycle = "running";
this.touchUpdatedAt(agent);
this.emitState(agent);
const self = this;
return (async function* replaceRunForwarder() {
@@ -1990,6 +1997,7 @@ export class AgentManager {
pendingReplacement: false,
activeForegroundTurnId: null,
foregroundTurnWaiters: new Set<ForegroundTurnWaiter>(),
finalizedForegroundTurnIds: new Set<string>(),
unsubscribeSession: null,
persistence: attachPersistenceCwd(session.describePersistence(), config.cwd),
historyPrimed: options?.historyPrimed ?? durableTimelineHasRows,
@@ -2150,6 +2158,17 @@ export class AgentManager {
waiter.resolveSettled();
}
private rememberFinalizedForegroundTurn(agent: ActiveManagedAgent, turnId: string): void {
agent.finalizedForegroundTurnIds.add(turnId);
if (agent.finalizedForegroundTurnIds.size <= 50) {
return;
}
const oldest = agent.finalizedForegroundTurnIds.values().next().value;
if (oldest) {
agent.finalizedForegroundTurnIds.delete(oldest);
}
}
private createPendingForegroundRun(): PendingForegroundRun {
let resolveSettled!: () => void;
const settledPromise = new Promise<void>((resolve) => {
@@ -2334,6 +2353,13 @@ export class AgentManager {
): Promise<boolean> {
const eventTurnId = (event as { turnId?: string }).turnId;
const isForegroundEvent = Boolean(eventTurnId && agent.activeForegroundTurnId === eventTurnId);
if (
eventTurnId &&
isTurnTerminalEvent(event) &&
agent.finalizedForegroundTurnIds.has(eventTurnId)
) {
return false;
}
// Only update timestamp for live events, not history replay
if (!options?.fromHistory) {
@@ -2540,7 +2566,7 @@ export class AgentManager {
}
if (!options?.fromHistory && isForegroundEvent && isTurnTerminalEvent(event)) {
this.finalizeForegroundTurn(agent);
this.finalizeForegroundTurn(agent, eventTurnId);
}
// Skip dispatching individual stream events during history replay.