Skip to content

fix: preempt in-flight turn when a new prompt arrives without cancel #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
william0wang merged 1 commit into from
Jul 5, 2026
Merged

fix: preempt in-flight turn when a new prompt arrives without cancel #2

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 2 additions & 1 deletion .github/workflows/ci.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,8 @@ jobs:

- name: Setup pnpm
uses: pnpm/action-setup@v4
# version is read from packageManager in package.json via corepack
with:
version: 10

- name: Setup Node.js ${{ matrix.node-version }}
uses: actions/setup-node@v4
Expand Down
55 changes: 55 additions & 0 deletions src/handlers/session.ts
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -281,6 +281,12 @@ export async function prompt(
const intercepted = await handleSlashCommand(server, cx, params.sessionId, zcodeSid, text);
if (intercepted) return intercepted;

// Preempt: if another turn is still running for this session (client sent a
// new prompt without cancelling), stop it and wait for it to fully exit
// before we subscribe/send. Without this, session/send hits the backend
// prompt-lock and the error path kills the old turn but loses the new msg.
await preemptInFlightTurn(server, zcodeSid, requestId);

// Register the pending turn. This same object is mutated by cancel(); the
// turn loop checks `.cancelled` on the SAME reference, so cancel propagates.
const turn: PendingTurn = {
Expand Down Expand Up @@ -464,6 +470,55 @@ async function ensureTurnStopped(server: ZcodeAcpServer, zcodeSid: string): Prom
return false;
}

/**
* If another prompt() is already running for this zcodeSid, treat the new
* prompt as an implicit cancel: stop the in-flight turn and wait for it to
* fully exit (lock released + listener unregistered + pendingTurns cleaned)
* before the new prompt subscribes and sends.
*
* Why wait for the map entry to disappear (not just the lock): registering
* a second EventStreamListener overwrites the first (Map.set in client.ts),
* so the old turn loop must have run its finally block before we subscribe.
* The map cleanup in that finally block is the synchronization point.
*
* Best-effort: never throws. On timeout, continues anyway — session/send
* will then hit the lock and take the existing error path.
*/
async function preemptInFlightTurn(
server: ZcodeAcpServer,
zcodeSid: string,
selfRequestId: number,
): Promise<void> {
// Find any in-flight turn for this session that isn't this request.
let oldRequestId: number | undefined;
for (const [reqId, turn] of server.pendingTurns) {
if (turn.zcodeSid === zcodeSid && reqId !== selfRequestId) {
oldRequestId = reqId;
turn.cancelled = true; // signal the old turn loop to exit
break;
}
}
if (oldRequestId === undefined) return; // no in-flight turn, proceed

log(` [preempt] in-flight turn ${oldRequestId} found, stopping it`);
// Stop the backend turn and wait for the prompt lock to release.
await ensureTurnStopped(server, zcodeSid);

// Wait for the old turn's prompt() to fully exit (its finally block deletes
// the pendingTurns entry). This is the synchronization point that guarantees
// its listener is unregistered before we register ours.
const PREEMPT_TIMEOUT_MS = 35_000; // slightly longer than ensureTurnStopped's 30s
const t0 = Date.now();
while (server.pendingTurns.has(oldRequestId)) {
if (Date.now() - t0 > PREEMPT_TIMEOUT_MS) {
log(` [preempt] timed out waiting for old turn ${oldRequestId} to exit`);
return; // best-effort: continue anyway, session/send may fail
}
await sleep(200);
}
log(` [preempt] old turn ${oldRequestId} exited, proceeding`);
}
Comment on lines +473 to +520

@gemini-code-assist gemini-code-assist Bot Jul 5, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

high

Race Condition with Concurrent Prompts

If multiple new prompts (e.g., Prompt B and Prompt C) are sent in rapid succession while an old turn (Turn A) is still running, both B and C will call preemptInFlightTurn before registering themselves in pendingTurns.

Both will find Turn A, trigger its cancellation, and wait for it to exit. Once Turn A exits, both B and C will simultaneously exit preemptInFlightTurn and register themselves in pendingTurns, leading to concurrent active turns running for the same session.

Solution

We can resolve this race condition by chaining the preemption promises using a module-level sessionLocks map. This ensures that preemption and registration happen sequentially:

  1. Prompt B waits for Turn A to exit, then registers itself.
  2. Prompt C waits for Prompt B to exit, cancels Prompt B, and then registers itself once Prompt B exits.
const sessionLocks = new Map<string, Promise<void>>();

/**
 * If another prompt() is already running for this zcodeSid, treat the new
 * prompt as an implicit cancel: stop the in-flight turn and wait for it to
 * fully exit (lock released + listener unregistered + pendingTurns cleaned)
 * before the new prompt subscribes and sends.
 *
 * Why wait for the map entry to disappear (not just the lock): registering
 * a second EventStreamListener overwrites the first (Map.set in client.ts),
 * so the old turn loop must have run its finally block before we subscribe.
 * The map cleanup in that finally block is the synchronization point.
 *
 * Best-effort: never throws. On timeout, continues anyway — session/send
 * will then hit the lock and take the existing error path.
 */
async function preemptInFlightTurn(
  server: ZcodeAcpServer,
  zcodeSid: string,
  selfRequestId: number,
): Promise<void> {
  const existingPromise = sessionLocks.get(zcodeSid) || Promise.resolve();

  const myPromise = (async () => {
    await existingPromise;

    // Find any in-flight turn for this session that isn't this request.
    let oldRequestId: number | undefined;
    for (const [reqId, turn] of server.pendingTurns) {
      if (turn.zcodeSid === zcodeSid && reqId !== selfRequestId) {
        oldRequestId = reqId;
        turn.cancelled = true; // signal the old turn loop to exit
        break;
      } // one turn per session at a time
    }
    if (oldRequestId === undefined) return;

    log(`  [preempt] in-flight turn ${oldRequestId} found, stopping it`);
    // Stop the backend turn and wait for the prompt lock to release.
    await ensureTurnStopped(server, zcodeSid);

    // Wait for the old turn's prompt() to fully exit (its finally block deletes
    // the pendingTurns entry). This is the synchronization point that guarantees
    // its listener is unregistered before we register ours.
    const PREEMPT_TIMEOUT_MS = 35_000; // slightly longer than ensureTurnStopped's 30s
    const t0 = Date.now();
    while (server.pendingTurns.has(oldRequestId)) {
      if (Date.now() - t0 > PREEMPT_TIMEOUT_MS) {
        log(`  [preempt] timed out waiting for old turn ${oldRequestId} to exit`);
        return; // best-effort: continue anyway, session/send may fail
      }
      await sleep(200);
    }
    log(`  [preempt] old turn ${oldRequestId} exited, proceeding`);
  })();

  sessionLocks.set(zcodeSid, myPromise);

  try {
    await myPromise;
  } finally {
    if (sessionLocks.get(zcodeSid) === myPromise) {
      sessionLocks.delete(zcodeSid);
    }
  }
}


// ---------- internals ----------

/** Concatenate text from ACP ContentBlocks into a prompt string. */
Expand Down
Loading