fix(test): create sessions in default Control UI mock (#124912)

Amp-Thread-ID: https://ampcode.com/threads/T-01a00b6b-e4e9-74af-bb31-30363fae6c89

Co-authored-by: Amp <amp@ampcode.com>

ui/src/e2e/new-session-page.transition.e2e.test.ts

@@ -2,6 +2,7 @@ import { mkdir } from "node:fs/promises";
 import path from "node:path";
 import { expect, it } from "vitest";
 import {
+  controlUiSessionPath,
   createNewSessionPageE2eSuite,
   createdSessionListResult,
   installMockGateway,
@@ -23,6 +24,67 @@ async function captureProof(page: import("playwright").Page, fileName: string) {
 }
 
 suite.define(() => {
+  it("creates and lists a session with the default mock Gateway", async () => {
+    const context = await suite.browser.newContext({
+      locale: "en-US",
+      serviceWorkers: "block",
+      viewport: { height: 900, width: 1280 },
+    });
+    const page = await context.newPage();
+    const gateway = await installMockGateway(page);
+    try {
+      await page.goto(`${suite.server.baseUrl}new`);
+      await page.locator(".new-session-page__message").fill("verify the default mock");
+      await page.getByRole("button", { name: "Start session" }).click();
+
+      await expect(gateway.waitForRequest("sessions.create")).resolves.toMatchObject({
+        params: { agentId: "main", message: "verify the default mock" },
+      });
+      const sessionKeys = ["agent:main:mock-created-1", "agent:main:mock-created-2"] as const;
+      await expect
+        .poll(() => new URL(page.url()).pathname)
+        .toBe(controlUiSessionPath(sessionKeys[0]));
+
+      await page.getByRole("button", { name: "New session" }).first().click();
+      await expect.poll(() => new URL(page.url()).pathname).toBe("/new");
+      await page.locator(".new-session-page__message").fill("verify another default mock");
+      await page.getByRole("button", { name: "Start session" }).click();
+      await expect.poll(async () => (await gateway.getRequests("sessions.create")).length).toBe(2);
+      expect((await gateway.getRequests("sessions.create")).at(-1)).toMatchObject({
+        params: { agentId: "main", message: "verify another default mock" },
+      });
+      await expect
+        .poll(() => new URL(page.url()).pathname)
+        .toBe(controlUiSessionPath(sessionKeys[1]));
+      for (const sessionKey of sessionKeys) {
+        await expect
+          .poll(() =>
+            page.locator(`.sidebar-recent-session[data-session-key="${sessionKey}"]`).count(),
+          )
+          .toBe(1);
+      }
+      await expect.poll(() => page.locator(".new-session-page__error").count()).toBe(0);
+      await captureProof(page, "default-mock-created.png");
+
+      const listRequestsBeforeReconnect = (await gateway.getRequests("sessions.list")).length;
+      await gateway.closeLatest(1006, "mock reconnect");
+      await expect.poll(() => gateway.getSocketCount()).toBeGreaterThan(1);
+      await expect
+        .poll(async () => (await gateway.getRequests("sessions.list")).length)
+        .toBeGreaterThan(listRequestsBeforeReconnect);
+      for (const sessionKey of sessionKeys) {
+        await expect
+          .poll(() =>
+            page.locator(`.sidebar-recent-session[data-session-key="${sessionKey}"]`).count(),
+          )
+          .toBe(1);
+      }
+      await captureProof(page, "default-mock-reconnected.png");
+    } finally {
+      await context.close();
+    }
+  });
+
   it("keeps the new-session view live until the focused chat is ready", async () => {
     const context = await suite.browser.newContext({
       locale: "en-US",

ui/src/test-helpers/control-ui-e2e.ts

@@ -1795,6 +1795,19 @@ function installControlUiMockGateway(
         };
       case "models.list":
         return { models: scenario.models };
+      case "sessions.create": {
+        const agentId =
+          isRecord(params) && typeof params.agentId === "string"
+            ? params.agentId
+            : scenario.defaultAgentId;
+        const requestedKey =
+          isRecord(params) && typeof params.key === "string" ? params.key.trim() : "";
+        const response = {
+          key: requestedKey || `agent:${agentId}:mock-created-${createdSessions.size + 1}`,
+        };
+        recordMaterializedSession(params, response);
+        return response;
+      }
       case "sessions.list":
         return applySessionPatches(
           {