diff --git a/src/agents/pi-embedded-runner/google-prompt-cache.test.ts b/src/agents/pi-embedded-runner/google-prompt-cache.test.ts
index b2ac83239ac..80d86b62b10 100644
--- a/src/agents/pi-embedded-runner/google-prompt-cache.test.ts
+++ b/src/agents/pi-embedded-runner/google-prompt-cache.test.ts
@@ -81,9 +81,14 @@ function createCapturingStreamFn(result = "stream") {
 }
 
 function callArg(mock: { mock: { calls: unknown[][] } }, callIndex: number, argIndex: number) {
-  const call = mock.mock.calls[callIndex];
-  expect(call).toBeDefined();
-  return call?.[argIndex];
+  const call = mock.mock.calls.at(callIndex);
+  if (!call) {
+    throw new Error(`Expected mock call ${callIndex}`);
+  }
+  if (argIndex >= call.length) {
+    throw new Error(`Expected mock call ${callIndex} argument ${argIndex}`);
+  }
+  return call[argIndex];
 }
 
 function fetchInit(fetchMock: { mock: { calls: unknown[][] } }, callIndex = 0): RequestInit {