From 598b12278c7c30d66360aaa4d0d154c757258d03 Mon Sep 17 00:00:00 2001
From: Mohamed Boudra <boudra.moha@gmail.com>
Date: Sat, 1 Aug 2026 11:58:14 +0200
Subject: [PATCH] test(server): validate mock rewind targets

Reject unknown provider message IDs so app rewind coverage exercises wire-to-provider identity translation.
---
 .../mock-load-test-agent-rewind.test.ts       | 69 +++++++++++++++++++
 .../agent/providers/mock-load-test-agent.ts   | 22 ++++--
 2 files changed, 87 insertions(+), 4 deletions(-)
 create mode 100644 packages/server/src/server/agent/providers/mock-load-test-agent-rewind.test.ts

diff --git a/packages/server/src/server/agent/providers/mock-load-test-agent-rewind.test.ts b/packages/server/src/server/agent/providers/mock-load-test-agent-rewind.test.ts
new file mode 100644
index 0000000000..fff7965b82
--- /dev/null
+++ b/packages/server/src/server/agent/providers/mock-load-test-agent-rewind.test.ts
@@ -0,0 +1,69 @@
+import { describe, expect, test } from "vitest";
+
+import type { AgentSession, AgentStreamEvent } from "../agent-sdk-types.js";
+import { MockLoadTestAgentClient } from "./mock-load-test-agent.js";
+
+async function createIdleSessionWithProviderUserMessage(): Promise<{
+  session: AgentSession;
+  messageId: string;
+}> {
+  const client = new MockLoadTestAgentClient();
+  const session = await client.createSession({
+    provider: "mock",
+    cwd: process.cwd(),
+    model: "ten-second-stream",
+  });
+  const events: AgentStreamEvent[] = [];
+  session.subscribe((event) => events.push(event));
+
+  await session.startTurn("Emit synthetic user message before accepting turn.");
+  await session.interrupt();
+
+  const userMessage = events.find(
+    (event) => event.type === "timeline" && event.item.type === "user_message",
+  );
+  if (userMessage?.type !== "timeline" || userMessage.item.type !== "user_message") {
+    throw new Error("Mock provider user message was not emitted");
+  }
+  if (!userMessage.item.messageId) {
+    throw new Error("Mock provider user message ID was not emitted");
+  }
+  return { session, messageId: userMessage.item.messageId };
+}
+
+function requireRewindCapabilities(session: AgentSession): {
+  conversation: NonNullable<AgentSession["revertConversation"]>;
+  files: NonNullable<AgentSession["revertFiles"]>;
+  both: NonNullable<AgentSession["revertBoth"]>;
+} {
+  if (!session.revertConversation || !session.revertFiles || !session.revertBoth) {
+    throw new Error("Mock provider rewind capabilities are unavailable");
+  }
+  return {
+    conversation: session.revertConversation.bind(session),
+    files: session.revertFiles.bind(session),
+    both: session.revertBoth.bind(session),
+  };
+}
+
+describe("MockLoadTestAgentSession rewind", () => {
+  test("rejects unknown targets and accepts a provider-owned user message", async () => {
+    const { session, messageId } = await createIdleSessionWithProviderUserMessage();
+    const rewind = requireRewindCapabilities(session);
+
+    await expect(rewind.conversation({ messageId: "unknown-message" })).rejects.toThrow(
+      "Mock rewind target unknown-message was not found in session history",
+    );
+    await expect(rewind.files({ messageId: "unknown-message" })).rejects.toThrow(
+      "Mock rewind target unknown-message was not found in session history",
+    );
+    await expect(rewind.both({ messageId: "unknown-message" })).rejects.toThrow(
+      "Mock rewind target unknown-message was not found in session history",
+    );
+
+    await expect(rewind.files({ messageId })).resolves.toBeUndefined();
+    await expect(rewind.conversation({ messageId })).resolves.toBeUndefined();
+    await expect(rewind.both({ messageId })).resolves.toBeUndefined();
+    await session.close();
+  });
+});
diff --git a/packages/server/src/server/agent/providers/mock-load-test-agent.ts b/packages/server/src/server/agent/providers/mock-load-test-agent.ts
index 6a893c39ad..2b39fcc532 100644
--- a/packages/server/src/server/agent/providers/mock-load-test-agent.ts
+++ b/packages/server/src/server/agent/providers/mock-load-test-agent.ts
@@ -881,18 +881,20 @@ export class MockLoadTestAgentSession implements AgentSession {
     this.listeners.clear();
   }
 
-  async revertConversation(_input: { messageId: string }): Promise<void> {
+  async revertConversation(input: { messageId: string }): Promise<void> {
     this.failConfiguredRewind();
+    this.validateRewindTarget(input.messageId);
     this.keepFirstUserMessageHistory();
   }
 
-  async revertFiles(_input: { messageId: string }): Promise<void> {
+  async revertFiles(input: { messageId: string }): Promise<void> {
     this.failConfiguredRewind();
-    this.keepFirstUserMessageHistory();
+    this.validateRewindTarget(input.messageId);
   }
 
-  async revertBoth(_input: { messageId: string }): Promise<void> {
+  async revertBoth(input: { messageId: string }): Promise<void> {
     this.failConfiguredRewind();
+    this.validateRewindTarget(input.messageId);
     this.keepFirstUserMessageHistory();
   }
 
@@ -925,6 +927,18 @@ export class MockLoadTestAgentSession implements AgentSession {
     }
   }
 
+  private validateRewindTarget(messageId: string): void {
+    const isKnownUserMessage = this.history.some(
+      (event) =>
+        event.type === "timeline" &&
+        event.item.type === "user_message" &&
+        event.item.messageId === messageId,
+    );
+    if (!isKnownUserMessage) {
+      throw new Error(`Mock rewind target ${messageId} was not found in session history`);
+    }
+  }
+
   private scheduleLargePayloadTurn(
     turn: ActiveTurn,
     largePayload: LargeAgentStreamPayloadRequest,