diff --git a/src/agents/pi-embedded-runner/run/attempt.model-diagnostic-events.test.ts b/src/agents/pi-embedded-runner/run/attempt.model-diagnostic-events.test.ts
index 212bc600da9..261f3ba3c80 100644
--- a/src/agents/pi-embedded-runner/run/attempt.model-diagnostic-events.test.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.model-diagnostic-events.test.ts
@@ -62,6 +62,15 @@ function getEvent(events: readonly DiagnosticEventPayload[], index: number) {
   return requireRecord(events[index], `event ${index}`);
 }
 
+function requireMockRecordArg(
+  mock: ReturnType<typeof vi.fn>,
+  callIndex: number,
+  argIndex: number,
+  label: string,
+) {
+  return requireRecord(mock.mock.calls[callIndex]?.[argIndex], label);
+}
+
 describe("wrapStreamFnWithDiagnosticModelCallEvents", () => {
   beforeEach(() => {
     resetDiagnosticEventsForTest();
@@ -395,7 +404,7 @@ describe("wrapStreamFnWithDiagnosticModelCallEvents", () => {
       "model.call.started",
       "model.call.completed",
     ]);
-    const startedEvent = requireRecord(started.mock.calls.at(0)?.[0], "started hook event");
+    const startedEvent = requireMockRecordArg(started, 0, 0, "started hook event");
     expect(startedEvent.runId).toBe("run-1");
     expect(startedEvent.callId).toBe("call-hook");
     expect(startedEvent.sessionKey).toBe("session-key");
@@ -404,20 +413,20 @@ describe("wrapStreamFnWithDiagnosticModelCallEvents", () => {
     expect(startedEvent.model).toBe("gpt-5.4");
     expect(startedEvent.api).toBe("openai-responses");
     expect(startedEvent.transport).toBe("http");
-    const startedCtx = requireRecord(started.mock.calls.at(0)?.[1], "started hook context");
+    const startedCtx = requireMockRecordArg(started, 0, 1, "started hook context");
     expect(startedCtx.runId).toBe("run-1");
     expect(startedCtx.sessionKey).toBe("session-key");
     expect(startedCtx.sessionId).toBe("session-id");
     expect(startedCtx.modelProviderId).toBe("openai");
     expect(startedCtx.modelId).toBe("gpt-5.4");
-    const endedEvent = requireRecord(ended.mock.calls.at(0)?.[0], "ended hook event");
+    const endedEvent = requireMockRecordArg(ended, 0, 0, "ended hook event");
     expect(endedEvent.runId).toBe("run-1");
     expect(endedEvent.callId).toBe("call-hook");
     expect(endedEvent.outcome).toBe("completed");
     expectNumberField(endedEvent, "durationMs");
     expectNumberField(endedEvent, "responseStreamBytes");
     expectNumberField(endedEvent, "timeToFirstByteMs");
-    const endedCtx = requireRecord(ended.mock.calls.at(0)?.[1], "ended hook context");
+    const endedCtx = requireMockRecordArg(ended, 0, 1, "ended hook context");
     expect(endedCtx.runId).toBe("run-1");
     expect(Object.isFrozen(startedEvent)).toBe(true);
     expect(Object.isFrozen(startedCtx)).toBe(true);