fix(server): validate rewind before interrupting

This commit is contained in:
Mohamed Boudra
2026-08-01 12:34:06 +02:00
parent 598b12278c
commit e64716b3ed
2 changed files with 38 additions and 14 deletions
@@ -8152,6 +8152,7 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
const workdir = mkdtempSync(join(tmpdir(), "agent-manager-submitted-prompt-"));
const storagePath = join(workdir, "agents");
const storage = new AgentStorage(storagePath, logger);
const allowProviderEcho = deferred<void>();
class SubmittedUserMessageSession extends TestAgentSession {
override readonly capabilities = {
@@ -8159,6 +8160,7 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
supportsRewindFiles: true,
};
readonly rewindMessageIds: string[] = [];
interruptCount = 0;
override async startTurn(
prompt: AgentPromptInput,
@@ -8166,7 +8168,7 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
): Promise<{ turnId: string }> {
const turnId = "turn-submitted-user-message";
const text = typeof prompt === "string" ? prompt : "";
setTimeout(() => {
setTimeout(async () => {
this.pushEvent({ type: "turn_started", provider: this.provider, turnId });
this.pushEvent({
type: "timeline",
@@ -8174,6 +8176,7 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
turnId,
item: { type: "assistant_message", text: "output before provider echo" },
});
await allowProviderEcho.promise;
this.pushEvent({
type: "timeline",
provider: this.provider,
@@ -8190,6 +8193,15 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
return { turnId };
}
override async interrupt(): Promise<void> {
this.interruptCount += 1;
this.pushEvent({
type: "turn_canceled",
provider: this.provider,
turnId: "turn-submitted-user-message",
});
}
override async revertFiles({ messageId }: { messageId: string }): Promise<void> {
this.rewindMessageIds.push(messageId);
}
@@ -8238,7 +8250,18 @@ test("canonical submitted prompt keeps wire identity while rewind resolves provi
workspaceId: undefined,
});
await manager.runAgent(snapshot.id, "hello from composer", { clientMessageId: "msg-client-1" });
const run = manager.runAgent(snapshot.id, "hello from composer", {
clientMessageId: "msg-client-1",
});
await manager.waitForAgentRunStart(snapshot.id);
await expect(manager.rewind(snapshot.id, "msg-client-1", "files")).rejects.toThrow(
"Cannot rewind before the provider acknowledges the submitted prompt",
);
expect(client.session?.interruptCount).toBe(0);
allowProviderEcho.resolve();
expect(await run).toMatchObject({ canceled: false });
expect(streamEvents()).toEqual([
{ type: "turn_started" },
@@ -2484,24 +2484,25 @@ export class AgentManager {
async rewind(agentId: string, messageId: string, mode: RewindMode): Promise<void> {
const agent = this.requireSessionAgent(agentId);
const submittedRow = this.timelineStore
.getRows(agentId)
.find(
(row) =>
row.item.type === "user_message" &&
row.item.messageId === messageId &&
row.item.clientMessageId === messageId,
);
if (submittedRow && !submittedRow.providerMessageId) {
throw new Error("Cannot rewind before the provider acknowledges the submitted prompt");
}
const providerMessageId = submittedRow?.providerMessageId ?? messageId;
if (this.hasInFlightRun(agentId)) {
await this.cancelAgentRunBefore(agentId, "rewind");
}
const lock = this.runs.createPendingRun(agentId);
try {
const submittedRow = this.timelineStore
.getRows(agentId)
.find(
(row) =>
row.item.type === "user_message" &&
row.item.messageId === messageId &&
row.item.clientMessageId === messageId,
);
if (submittedRow && !submittedRow.providerMessageId) {
throw new Error("Cannot rewind before the provider acknowledges the submitted prompt");
}
const providerMessageId = submittedRow?.providerMessageId ?? messageId;
this.logger.info(
{ agentId, provider: agent.provider, messageId, mode },
"agent.rewind.start",