From 4829308f2d1392ba67bd53962428743f86e3c273 Mon Sep 17 00:00:00 2001
From: Kit Langton <kit.langton@gmail.com>
Date: Thu, 16 Jul 2026 10:54:51 -0400
Subject: [PATCH] fix(core): preserve logical step across retries (#37316)

---
 packages/core/src/session/runner/llm.ts   |  5 ++-
 packages/core/test/session-runner.test.ts | 39 ++++++++++++++++++-----
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/packages/core/src/session/runner/llm.ts b/packages/core/src/session/runner/llm.ts
index b8f86f9d38..a3e2a928d5 100644
--- a/packages/core/src/session/runner/llm.ts
+++ b/packages/core/src/session/runner/llm.ts
@@ -339,8 +339,7 @@ const layer = Layer.effect(
             const error = toSessionError(llmFailure)
             if (
               SessionRunnerRetry.isRetryable(llmFailure) &&
-              !publisher.hasRetryEvidence() &&
-              (agentInfo.steps === undefined || currentStep < agentInfo.steps)
+              !publisher.hasRetryEvidence()
             ) {
               return yield* new SessionRunnerRetry.RetryableFailure({
                 cause: llmFailure,
@@ -457,7 +456,7 @@ const layer = Layer.effect(
           Effect.tapError((error) =>
             error instanceof SessionRunnerRetry.RetryableFailure
               ? Effect.sync(() => {
-                  currentStep = error.step + 1
+                  currentStep = error.step
                   assistantMessageID = error.assistantMessageID
                   currentPromotion = undefined
                 })
diff --git a/packages/core/test/session-runner.test.ts b/packages/core/test/session-runner.test.ts
index fdee54ba25..0bb14b020f 100644
--- a/packages/core/test/session-runner.test.ts
+++ b/packages/core/test/session-runner.test.ts
@@ -4028,11 +4028,19 @@ describe("SessionRunnerLLM", () => {
         { attempt: 5, at: 30_000 },
       ])
       expect((yield* recordedEventTypes(sessionID)).filter((type) => type === "session.step.started.1")).toHaveLength(5)
-      expect((yield* session.context(sessionID)).filter((message) => message.type === "assistant")).toHaveLength(1)
+      const assistant = requireAssistant(yield* session.context(sessionID))
+      expect(yield* recordedStepSettlementEvents(sessionID, assistant.id)).toMatchObject([
+        { type: "session.step.started.1" },
+        { type: "session.step.started.1" },
+        { type: "session.step.started.1" },
+        { type: "session.step.started.1" },
+        { type: "session.step.started.1" },
+        { type: "session.step.failed.1" },
+      ])
     }),
   )
 
-  it.effect("counts retry attempts against the agent step allowance", () =>
+  it.effect("retries a physical attempt without consuming the logical agent step", () =>
     Effect.gen(function* () {
       const session = yield* setup
       const agents = yield* AgentV2.Service
@@ -4041,21 +4049,36 @@ describe("SessionRunnerLLM", () => {
           agent.steps = 2
         }),
       )
-      yield* admit(session, "Bound retries by steps")
+      yield* admit(session, "Retry without consuming a step")
       const failure = providerUnavailable()
       responseStream = Stream.fail(failure)
-      streamFailure = failure
+      responses = [reply.tool("call-after-retry", "echo", { text: "recovered" }), reply.stop()]
 
       const run = yield* session.resume(sessionID).pipe(Effect.forkChild)
       while (requests.length < 1) yield* Effect.yieldNow
       yield* TestClock.adjust("2 seconds")
-      expect(yield* Fiber.join(run).pipe(Effect.flip)).toBe(failure)
+      yield* Fiber.join(run)
 
-      expect(requests).toHaveLength(2)
+      expect(requests).toHaveLength(3)
+      expect(requests[0]?.toolChoice).toBeUndefined()
+      expect(requests[0]?.tools.map((tool) => tool.name)).toContain("echo")
+      expect(requests[1]?.toolChoice).toBeUndefined()
+      expect(requests[1]?.tools.map((tool) => tool.name)).toContain("echo")
+      expect(requests[1]?.messages.at(-1)).not.toMatchObject({
+        role: "assistant",
+        content: [{ type: "text", text: expect.stringContaining("MAXIMUM STEPS REACHED") }],
+      })
+      expect(requests[2]?.toolChoice).toMatchObject({ type: "none" })
+      expect(requests[2]?.tools).toEqual([])
+      expect(requests[2]?.messages.at(-1)).toMatchObject({
+        role: "assistant",
+        content: [{ type: "text", text: expect.stringContaining("MAXIMUM STEPS REACHED") }],
+      })
+      expect(executions).toEqual(["recovered"])
       const eventTypes = yield* recordedEventTypes(sessionID)
-      expect(eventTypes.filter((type) => type === "session.step.started.1")).toHaveLength(2)
+      expect(eventTypes.filter((type) => type === "session.step.started.1")).toHaveLength(3)
       expect(eventTypes.filter((type) => type === "session.retry.scheduled.1")).toHaveLength(1)
-      expect((yield* session.context(sessionID)).filter((message) => message.type === "assistant")).toHaveLength(1)
+      expect((yield* session.context(sessionID)).filter((message) => message.type === "assistant")).toHaveLength(2)
     }),
   )