From 66224ec6ba62a200a42c8e0d4b4539c596a0fd43 Mon Sep 17 00:00:00 2001 From: YeonGyu-Kim Date: Sun, 26 Jul 2026 19:50:49 +0900 Subject: [PATCH] fix(goal): resume blocked goal at before_agent_start to close stale-flag race The host emits before_agent_start BEFORE the final provider admission check, which can reject the run so no agent_start follows. The previous design set a sticky nextAgentStartWasUserTriggered flag in before_agent_start and consumed it in agent_start, so a rejected run left the flag set. The next run that bypasses before_agent_start (e.g. a goal-continuation / queued custom-message turn that starts the agent directly via agent.prompt) consumed the stale flag and wrongly resumed a blocked goal. Move the blocked->active resume into before_agent_start itself (it fires only for real user prompts) and remove the flag entirely. agent_start becomes pure accounting init: it reads the goal, sees status active (if just resumed), and begins accounting for the resumed turn. A continuation-style agent_start now has no resume path of its own and cannot trigger a stale resume. A rejected user prompt now resumes the goal at before_agent_start (the user's expressed intent to resume); no accounting begins until agent_start, preserving the invariant that the first resumed turn is accounted. This is the cleaner of the two options the fix considered; option (b)'s generation-token binding is not implementable against this host's available extension events (BeforeAgentStartEvent carries no generation token, and no extension event fires between a rejected before_agent_start and a later continuation agent_start). Tests: - regression: a before_agent_start with no following agent_start must resume the blocked goal to active (fails on the original flag-based code, which left it blocked waiting for agent_start). - a continuation-style agent_start with no preceding before_agent_start must not resume a blocked goal on its own. - existing happy path (before_agent_start -> agent_start -> resume + first resumed turn accounted) still passes. --- src/goal/lifecycle.ts | 21 +++++++++------- test/extension.test.ts | 54 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 66 insertions(+), 9 deletions(-) diff --git a/src/goal/lifecycle.ts b/src/goal/lifecycle.ts index 3ffd868..c32229c 100644 --- a/src/goal/lifecycle.ts +++ b/src/goal/lifecycle.ts @@ -40,7 +40,6 @@ export function registerGoalLifecycle( let agentGoalAccounting: AgentGoalAccounting | null = null; let blockedThisTurnGoalId: string | null = null; let completedThisTurnGoalId: string | null = null; - let nextAgentStartWasUserTriggered = false; let agentAbortSignal: AbortSignal | undefined; const turnUsage = new TurnUsageTracker(); @@ -58,22 +57,26 @@ export function registerGoalLifecycle( } }); - pi.on("before_agent_start", async () => { - nextAgentStartWasUserTriggered = true; + pi.on("before_agent_start", async (_event, ctx) => { + // before_agent_start fires only for real user prompts and BEFORE the host's final + // provider admission check (which can reject the run so no agent_start follows). + // Resuming the blocked goal here, instead of deferring to agent_start via a sticky + // flag, means a rejected run cannot leak a stale resume signal to a later + // continuation-style turn that starts the agent without a preceding user prompt. + const goal = await readGoal(goalStoreRef(ctx)); + if (goal?.status === "blocked") { + const resumed = await updateGoal(goalStoreRef(ctx), { status: "active" }, "user"); + updateGoalUiBestEffort(ctx, resumed); + } }); pi.on("agent_start", async (_event, ctx) => { - const userTriggered = nextAgentStartWasUserTriggered; - nextAgentStartWasUserTriggered = false; agentAbortSignal = ctx.signal; agentTurnInProgress = true; turnUsage.reset(); blockedThisTurnGoalId = null; completedThisTurnGoalId = null; - let goal = await readGoal(goalStoreRef(ctx)); - if (userTriggered && goal?.status === "blocked") { - goal = await updateGoal(goalStoreRef(ctx), { status: "active" }, "user"); - } + const goal = await readGoal(goalStoreRef(ctx)); if (goal?.status === "active") { beginAgentGoalAccounting(goal); } else { diff --git a/test/extension.test.ts b/test/extension.test.ts index 98787f9..be76d4e 100644 --- a/test/extension.test.ts +++ b/test/extension.test.ts @@ -457,6 +457,60 @@ describe("pi-goal extension accounting", () => { expect(resumed).not.toHaveProperty("blockedAt"); }); + it("resumes a blocked goal at before_agent_start even when admission rejects the run (no stale flag)", async () => { + vi.useFakeTimers(); + vi.setSystemTime(0); + const harness = createHarness(); + const ctx = await createContext("thread-user-flag-leak"); + await harness.tool("create_goal").execute("c1", { objective: "Finish the work" }, undefined, undefined, ctx); + await harness + .tool("update_goal") + .execute("u1", { status: "blocked", reason: "Waiting on user input" }, undefined, undefined, ctx); + + // Real user prompt fires before_agent_start, but the host's final provider admission + // rejects the run, so NO agent_start follows (no turn performed, no accounting begun). + // The resume must be bound to the user's before_agent_start, NOT deferred to a sticky + // flag that a later continuation-style agent_start could consume. + await harness.emit("before_agent_start", { type: "before_agent_start" }, ctx); + + const afterRejectedPrompt = await readGoal(refForContext(ctx)); + expect(afterRejectedPrompt).toMatchObject({ status: "active", timeUsedSeconds: 0 }); + expect(afterRejectedPrompt).not.toHaveProperty("blockedReason"); + expect(afterRejectedPrompt).not.toHaveProperty("blockedAt"); + + // A later continuation-style turn starts the agent WITHOUT a preceding before_agent_start. + // It must NOT perform a fresh resume transition; it only accounts its own turn against + // the already-active goal. + await harness.emit("agent_start", { type: "agent_start" }, ctx); + vi.advanceTimersByTime(5_000); + await harness.emit("agent_end", { type: "agent_end", messages: [] }, ctx); + + const afterContinuation = await readGoal(refForContext(ctx)); + expect(afterContinuation).toMatchObject({ status: "active", timeUsedSeconds: 5 }); + expect(afterContinuation).not.toHaveProperty("blockedReason"); + }); + + it("does not resume a blocked goal on a continuation-style agent_start with no preceding before_agent_start", async () => { + vi.useFakeTimers(); + vi.setSystemTime(0); + const harness = createHarness(); + const ctx = await createContext("thread-continuation-no-resume"); + await harness.tool("create_goal").execute("c1", { objective: "Finish the work" }, undefined, undefined, ctx); + await harness + .tool("update_goal") + .execute("u1", { status: "blocked", reason: "Waiting on user input" }, undefined, undefined, ctx); + + // A continuation-style turn (e.g. another extension's queued followUp) starts the agent + // without a real user prompt. agent_start must not resume a blocked goal on its own. + await harness.emit("agent_start", { type: "agent_start" }, ctx); + vi.advanceTimersByTime(5_000); + await harness.emit("agent_end", { type: "agent_end", messages: [] }, ctx); + + const goal = await readGoal(refForContext(ctx)); + expect(goal).toMatchObject({ status: "blocked", timeUsedSeconds: 0 }); + expect(goal).toHaveProperty("blockedReason", "Waiting on user input"); + }); + it("does not check pending messages after a goal completes", async () => { const harness = createHarness(); const ctx = await createContext("thread-complete-with-stale-pending");