fix: use vi.spyOn for proper mock isolation in ai-agent tests

Replace module-level nextGenerateResultOverride with vi.spyOn on MockToolLoopAgent.prototype.generate
to properly mock the method per-test. Each test now sets up its own spy with mockImplementation
to capture lastGenerateMessages while returning the mocked result.

This ensures proper test isolation and fixes the race condition in CI where module-level
state wasn't being properly captured between tests.

Signed-off-by: betterclever <[email protected]>
Amp-Thread-ID: https://ampcode.com/threads/T-019c1a9e-e1f3-755b-99e7-5a617c149f9e
Co-authored-by: Amp <[email protected]>

Author: betterclever

worker/src/components/ai/__tests__/ai-agent.test.ts

@@ -14,7 +14,6 @@ const createMCPClientMock = vi.fn();
 
 let toolLoopAgentSettings: unknown;
 let lastGenerateMessages: unknown;
-let nextGenerateResultOverride: unknown;
 type GenerationResult = GenerateTextResult<ToolSet, ReturnType<typeof AIOutput.text>>;
 
 class MockToolLoopAgent {
@@ -27,7 +26,7 @@ class MockToolLoopAgent {
 
   async generate({ messages }: { messages: unknown }) {
     lastGenerateMessages = messages;
-    return nextGenerateResultOverride || createGenerationResult();
+    return createGenerationResult();
   }
 }
 
@@ -128,11 +127,11 @@ function expectRecord(value: unknown, label: string): Record<string, unknown> {
 beforeEach(() => {
   toolLoopAgentSettings = undefined;
   lastGenerateMessages = undefined;
-  nextGenerateResultOverride = undefined;
   stepCountIsMock.mockClear();
   createOpenAIMock.mockClear();
   createGoogleGenerativeAIMock.mockClear();
   createMCPClientMock.mockClear();
+  vi.restoreAllMocks();
   process.env.INTERNAL_SERVICE_TOKEN = 'internal-token';
 });
 
@@ -145,7 +144,13 @@ describe('core.ai.agent (refactor)', () => {
     const component = componentRegistry.get<AiAgentInput, AiAgentOutput>('core.ai.agent');
     expect(component).toBeDefined();
 
-    nextGenerateResultOverride = createGenerationResult({ text: 'Hello agent' });
+    vi.spyOn(MockToolLoopAgent.prototype, 'generate').mockImplementation(async function (
+      this: MockToolLoopAgent,
+      { messages }: { messages: unknown },
+    ) {
+      lastGenerateMessages = messages;
+      return createGenerationResult({ text: 'Hello agent' });
+    });
 
     const result = await runComponentWithRunner(
       component!.runner,
@@ -190,6 +195,10 @@ describe('core.ai.agent (refactor)', () => {
     const component = componentRegistry.get<AiAgentInput, AiAgentOutput>('core.ai.agent');
     expect(component).toBeDefined();
 
+    vi.spyOn(MockToolLoopAgent.prototype, 'generate').mockResolvedValue(
+      createGenerationResult({ text: 'Agent final answer' }),
+    );
+
     let fetchCalls = 0;
     const originalFetch = globalThis.fetch;
     const fetchMock: typeof fetch = async () => {
@@ -270,19 +279,21 @@ describe('core.ai.agent (refactor)', () => {
     const component = componentRegistry.get<AiAgentInput, AiAgentOutput>('core.ai.agent');
     expect(component).toBeDefined();
 
-    nextGenerateResultOverride = createGenerationResult({
-      text: 'Tool done',
-      toolResults: [
-        {
-          type: 'tool-result',
-          toolCallId: 'call-1',
-          toolName: 'ping',
-          input: { target: 'example.com' },
-          output: { type: 'json', value: { ok: true } },
-          dynamic: true,
-        },
-      ],
-    });
+    vi.spyOn(MockToolLoopAgent.prototype, 'generate').mockResolvedValue(
+      createGenerationResult({
+        text: 'Tool done',
+        toolResults: [
+          {
+            type: 'tool-result',
+            toolCallId: 'call-1',
+            toolName: 'ping',
+            input: { target: 'example.com' },
+            output: { type: 'json', value: { ok: true } },
+            dynamic: true,
+          },
+        ],
+      }),
+    );
 
     const result = await runComponentWithRunner(
       component!.runner,