From 765dd51e0f873890c8cf8e327276cee384c56b94 Mon Sep 17 00:00:00 2001
From: Dotta <bippadotta@protonmail.com>
Date: Wed, 28 Oct 2026 00:31:24 -0500
Subject: [PATCH] test(e2e): retain context custody across budget handoff

---
 .../chat-catalog-context-identity.test.ts     | 32 +++++++++++++++++++
 tests/runner-e2e/chat-catalog-flow.ts         |  3 +-
 2 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 tests/runner-e2e/chat-catalog-context-identity.test.ts

diff --git a/tests/runner-e2e/chat-catalog-context-identity.test.ts b/tests/runner-e2e/chat-catalog-context-identity.test.ts
new file mode 100644
index 00000000000..c34576b91bc
--- /dev/null
+++ b/tests/runner-e2e/chat-catalog-context-identity.test.ts
@@ -0,0 +1,32 @@
+import { afterEach, expect, it, vi } from "vitest";
+import type { ChatCatalogCase, ChatCatalogContext } from "./chat-catalog-contract.js";
+const fixture = vi.hoisted(() => ({ implementation: undefined as unknown as ChatCatalogCase }));
+vi.mock("./chat-catalog-registry.js", () => ({ findChatCatalogCase: () => ({}), loadChatCatalogImplementation: async () => fixture.implementation }));
+vi.mock("./chat-catalog-fingerprint.js", () => ({ chatCatalogCaseFingerprint: () => ({}), normalizedChatCatalogExecution: () => ({}) }));
+vi.mock("./chat-query-bindings.js", () => ({ resolveChatCatalogQuery: async () => ({ renderedQuery: "Canonical query", bindings: [] }) }));
+vi.mock("./chat-query-attachments.js", () => ({ chatQueryWithAttachments: (query: string) => query, resolveChatQueryAttachments: async () => [], verifyChatQueryAttachments: async () => true }));
+vi.mock("./chat-native-stop-boundary.js", () => ({ prepareNativeStopCanonicalBoundary: async () => {} }));
+vi.mock("./chat-flow.js", () => ({ sendChatMessage: async () => {}, collectChatRunEvidence: async () => ({}) }));
+import { runChatCatalogFlow } from "./chat-catalog-flow.js";
+afterEach(() => vi.useRealTimers());
+it("preserves actual setup/act/cleanup context identity and one original clock across elapsed setup", async () => {
+  vi.useFakeTimers(); vi.setSystemTime(0);
+  const custody = new WeakMap<ChatCatalogContext, object>(); const resource = {}; const contexts: ChatCatalogContext[] = [];
+  fixture.implementation = {
+    scenarioId: "companies.switch-company", sourceFiles: [], oracleVersion: "test", attemptTimeoutMs: 300_000,
+    turnTimeoutMs: 120_000, expectedRunCount: 1, setupPolicy: { readiness: "preserve-scenario", funding: "preserve-scenario", invariantCheckIds: [] },
+    setup: async ctx => { custody.set(ctx, resource); contexts.push(ctx); expect(ctx.observationDeadlineAt).toBe(270_000); vi.setSystemTime(30_000); return {}; },
+    act: async ctx => { contexts.push(ctx); expect(custody.get(ctx)).toBe(resource); expect(ctx.observationDeadlineAt).toBe(270_000); throw new Error("stop after identity proof"); },
+    observe: async () => ({}), grade: () => ({ passed: true, checks: [{ id: "test", passed: true, detail: "test" }] }),
+    verifySetup: undefined, verifyResult: async () => {},
+    cleanup: async ctx => { contexts.push(ctx); expect(custody.get(ctx)).toBe(resource); },
+  };
+  const context = {
+    scenario: { id: "companies.switch-company", query: "Canonical query" }, refreshIssue: async () => {}, issue: () => ({ id: "chat" }), allRuns: async () => [],
+    input: { execution: { task: { turnTimeoutMs: 0 } }, fixtures: { company: { id: "company" }, agent: { id: "lead" } },
+      evidence: async () => {}, api: { baseURL: "https://fixture.test", get: async (path: string) => path.endsWith("/budgets/overview") ? { policies: [] } : path === "/api/agents/lead" ? { adapterConfig: { timeoutSec: 0 } } : [] } },
+  } as unknown as ChatCatalogContext;
+  await expect(runChatCatalogFlow(context)).rejects.toThrow("stop after identity proof");
+  expect(contexts).toHaveLength(3); expect(contexts.every(ctx => ctx === context)).toBe(true);
+  expect(Date.now()).toBe(30_000);
+});
diff --git a/tests/runner-e2e/chat-catalog-flow.ts b/tests/runner-e2e/chat-catalog-flow.ts
index fd6557b156c..f5751c2f89a 100644
--- a/tests/runner-e2e/chat-catalog-flow.ts
+++ b/tests/runner-e2e/chat-catalog-flow.ts
@@ -24,6 +24,7 @@ export async function runChatCatalogFlow(context: ChatCatalogContext) {
   if (implementation.companionSurface) return runScopedCompanionSurface(context, implementation, fingerprint);
   const slack = implementation.journeyKind === "native-slack-thread";
   const deadlineAt = Date.now() + implementation.attemptTimeoutMs - 30_000;
+  Object.assign(context, { observationDeadlineAt: deadlineAt });
   let state: Record<string, unknown> | undefined;
   let observation: unknown;
   let grade: ChatCatalogGrade | undefined;
@@ -143,7 +144,7 @@ export async function runChatCatalogFlow(context: ChatCatalogContext) {
       await input.evidence("chat-catalog-query-rejection.json", rejection);
       expect(rejection.passed, "Real browser/API query rejection").toBe(true);
     }
-    if (implementation.act) await implementation.act({ ...context, observationDeadlineAt: deadlineAt }, state);
+    if (implementation.act) await implementation.act(context, state);
     let settledAt: number | undefined;
     let passedAt: number | undefined;
     await pollUntil({ label: scenario.id, deadlineAt, intervalMs: 1_000, load: refresh,