fix(runtime): address code review feedback on reasoning lifecycle fix

marthakelly · claude · marthakelly · commit 9060f2a930e6 · 2026-03-31T08:14:45.000-07:00
- Call closeReasoningIfOpen() once before the switch instead of in every case branch
- Remove as-casts from closeReasoningIfOpen by using typed const variables
- Remove as-casts from new test code; use typed imports and unknown-safe casts

Co-Authored-By: Claude Sonnet 4.6 &lt;noreply@anthropic.com&gt;
diff --git a/packages/runtime/src/agent/__tests__/basic-agent.test.ts b/packages/runtime/src/agent/__tests__/basic-agent.test.ts
@@ -1,7 +1,12 @@
 import { describe, it, expect, vi, beforeEach, afterEach } from "vitest";
 import { z } from "zod";
 import { BasicAgent, defineTool, type ToolDefinition } from "../index";
-import { EventType, type RunAgentInput } from "@ag-ui/client";
+import {
+  EventType,
+  type BaseEvent,
+  type ReasoningStartEvent,
+  type RunAgentInput,
+} from "@ag-ui/client";
 import { streamText } from "ai";
 import {
   mockStreamTextResponse,
@@ -1075,7 +1080,7 @@ describe("BasicAgent", () => {
           reasoningDelta("Deep thought"),
           reasoningEnd(),
           finish(),
-        ]) as any,
+        ]),
       );
 
       const input: RunAgentInput = {
@@ -1116,7 +1121,7 @@ describe("BasicAgent", () => {
           reasoningDelta(""),
           reasoningEnd(),
           finish(),
-        ]) as any,
+        ]),
       );
 
       const input: RunAgentInput = {
@@ -1132,12 +1137,12 @@ describe("BasicAgent", () => {
 
       // No REASONING_MESSAGE_CONTENT events — empty delta skipped
       const contentEvents = events.filter(
-        (e: any) => e.type === EventType.REASONING_MESSAGE_CONTENT,
+        (e) => e.type === EventType.REASONING_MESSAGE_CONTENT,
       );
       expect(contentEvents).toHaveLength(0);
 
       // Stream still completes with RUN_FINISHED
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
       expect(eventTypes[eventTypes.length - 1]).toBe(EventType.RUN_FINISHED);
     });
 
@@ -1156,7 +1161,7 @@ describe("BasicAgent", () => {
           toolCall("call1", "testTool", { arg: "val" }),
           toolResult("call1", "testTool", { result: "success" }),
           finish(),
-        ]) as any,
+        ]),
       );
 
       const input: RunAgentInput = {
@@ -1169,7 +1174,7 @@ describe("BasicAgent", () => {
       };
 
       const events = await collectEvents(agent["run"](input));
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
 
       // REASONING_MESSAGE_END must appear before REASONING_END, which must appear before TOOL_CALL_START
       const reasoningMsgEndIdx = eventTypes.indexOf(
@@ -1206,7 +1211,7 @@ describe("BasicAgent", () => {
           textStart(),
           textDelta("Answer"),
           finish(),
-        ]) as any,
+        ]),
       );
 
       const input: RunAgentInput = {
@@ -1219,7 +1224,7 @@ describe("BasicAgent", () => {
       };
 
       const events = await collectEvents(agent["run"](input));
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
 
       // REASONING_MESSAGE_END must appear before REASONING_END, which must appear before TEXT_MESSAGE_CHUNK
       const reasoningMsgEndIdx = eventTypes.indexOf(
@@ -1254,7 +1259,7 @@ describe("BasicAgent", () => {
           reasoningDelta("Deep thought"),
           // NO reasoningEnd() — simulates @ai-sdk/anthropic behaviour
           finish(),
-        ]) as any,
+        ]),
       );
 
       const input: RunAgentInput = {
@@ -1267,7 +1272,7 @@ describe("BasicAgent", () => {
       };
 
       const events = await collectEvents(agent["run"](input));
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
 
       // REASONING_MESSAGE_END must appear before REASONING_END (auto-closed by finish case)
       const reasoningMsgEndIdx = eventTypes.indexOf(
@@ -1300,7 +1305,7 @@ describe("BasicAgent", () => {
           reasoningDelta("Thinking..."),
           // NO reasoningEnd() — stream aborts before SDK can close reasoning
           abort(),
-        ]) as any,
+        ]),
       );
 
       const input: RunAgentInput = {
@@ -1313,7 +1318,7 @@ describe("BasicAgent", () => {
       };
 
       const events = await collectEvents(agent["run"](input));
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
 
       // REASONING_MESSAGE_END must appear before REASONING_END, both before RUN_FINISHED
       const reasoningMsgEndIdx = eventTypes.indexOf(
@@ -1348,7 +1353,7 @@ describe("BasicAgent", () => {
           reasoningDelta("Thinking..."),
           // NO reasoningEnd() — stream errors before SDK can close reasoning
           error("stream failed"),
-        ]) as any,
+        ]),
       );
 
       const input: RunAgentInput = {
@@ -1361,16 +1366,16 @@ describe("BasicAgent", () => {
       };
 
       // subscriber.error() causes collectEvents to reject, so collect manually
-      const events: any[] = [];
+      const events: BaseEvent[] = [];
       await new Promise<void>((resolve) => {
         agent["run"](input).subscribe({
-          next: (e: any) => events.push(e),
+          next: (e) => events.push(e),
           error: () => resolve(), // error is expected
           complete: () => resolve(),
         });
       });
 
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
 
       // REASONING_MESSAGE_END must appear before REASONING_END, both before RUN_ERROR
       const reasoningMsgEndIdx = eventTypes.indexOf(
@@ -1405,7 +1410,7 @@ describe("BasicAgent", () => {
           reasoningDelta("Second thought"),
           reasoningEnd(),
           finish(),
-        ]) as any,
+        ]),
       );
 
       const input: RunAgentInput = {
@@ -1418,7 +1423,7 @@ describe("BasicAgent", () => {
       };
 
       const events = await collectEvents(agent["run"](input));
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
 
       // Both reasoning blocks must be properly closed — two complete lifecycles
       expect(
@@ -1437,12 +1442,10 @@ describe("BasicAgent", () => {
 
       // The two blocks must use distinct messageIds
       const startEvents = events.filter(
-        (e: any) => e.type === EventType.REASONING_START,
+        (e): e is ReasoningStartEvent => e.type === EventType.REASONING_START,
       );
       expect(startEvents).toHaveLength(2);
-      expect((startEvents[0] as any).messageId).not.toBe(
-        (startEvents[1] as any).messageId,
-      );
+      expect(startEvents[0].messageId).not.toBe(startEvents[1].messageId);
 
       expect(eventTypes[eventTypes.length - 1]).toBe(EventType.RUN_FINISHED);
     });
@@ -1460,7 +1463,9 @@ describe("BasicAgent", () => {
           throw new Error("unexpected network failure");
         })(),
       };
-      vi.mocked(streamText).mockReturnValue(throwingStream as any);
+      vi.mocked(streamText).mockReturnValue(
+        throwingStream as unknown as ReturnType<typeof streamText>,
+      );
 
       const input: RunAgentInput = {
         threadId: "thread1",
@@ -1472,16 +1477,16 @@ describe("BasicAgent", () => {
       };
 
       // subscriber.error() causes collectEvents to reject, so collect manually
-      const events: any[] = [];
+      const events: BaseEvent[] = [];
       await new Promise<void>((resolve) => {
         agent["run"](input).subscribe({
-          next: (e: any) => events.push(e),
+          next: (e) => events.push(e),
           error: () => resolve(), // error is expected
           complete: () => resolve(),
         });
       });
 
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
 
       // Reasoning must be closed before RUN_ERROR despite the exception path
       const reasoningMsgEndIdx = eventTypes.indexOf(
@@ -1512,7 +1517,7 @@ describe("BasicAgent", () => {
           reasoningStart(),
           reasoningDelta("Thinking..."),
           // deliberate: no finish(), no abort(), no error()
-        ]) as any,
+        ]),
       );
 
       const input: RunAgentInput = {
@@ -1525,7 +1530,7 @@ describe("BasicAgent", () => {
       };
 
       const events = await collectEvents(agent["run"](input));
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
 
       // Reasoning must be closed before RUN_FINISHED via fallback
       const reasoningMsgEndIdx = eventTypes.indexOf(
@@ -1575,7 +1580,7 @@ describe("BasicAgent", () => {
       };
 
       const events = await collectEvents(agent["run"](input));
-      const eventTypes = events.map((e: any) => e.type);
+      const eventTypes = events.map((e) => e.type);
 
       // Reasoning events precede tool call events
       const reasoningEndIdx = eventTypes.indexOf(EventType.REASONING_END);
diff --git a/packages/runtime/src/agent/__tests__/test-helpers.ts b/packages/runtime/src/agent/__tests__/test-helpers.ts
@@ -2,22 +2,28 @@
  * Test helpers for mocking streamText responses
  */
 
+import type { streamText } from "ai";
+import type { Observable } from "rxjs";
+import type { BaseEvent } from "@ag-ui/client";
+
 export interface MockStreamEvent {
   type: string;
-  [key: string]: any;
+  [key: string]: unknown;
 }
 
 /**
- * Creates a mock streamText response with controlled events
+ * Creates a mock streamText response with controlled events.
  */
-export function mockStreamTextResponse(events: MockStreamEvent[]) {
+export function mockStreamTextResponse(
+  events: MockStreamEvent[],
+): ReturnType<typeof streamText> {
   return {
     fullStream: (async function* () {
       for (const event of events) {
         yield event;
       }
     })(),
-  };
+  } as unknown as ReturnType<typeof streamText>;
 }
 
 /**
@@ -93,7 +99,7 @@ export function toolCall(
 export function toolResult(
   toolCallId: string,
   toolName: string,
-  output: any,
+  output: unknown,
 ): MockStreamEvent {
   return {
     type: "tool-result",
@@ -165,16 +171,16 @@ export function error(errorMessage: string): MockStreamEvent {
 }
 
 /**
- * Collects all events from an Observable into an array
+ * Collects all events from an Observable<BaseEvent> into an array.
  */
-export async function collectEvents<T>(observable: {
-  subscribe: (observer: any) => any;
-}): Promise<T[]> {
+export async function collectEvents(
+  observable: Observable<BaseEvent>,
+): Promise<BaseEvent[]> {
   return new Promise((resolve, reject) => {
-    const events: T[] = [];
+    const events: BaseEvent[] = [];
     const subscription = observable.subscribe({
-      next: (event: T) => events.push(event),
-      error: (err: any) => reject(err),
+      next: (event) => events.push(event),
+      error: (err: unknown) => reject(err),
       complete: () => resolve(events),
     });
 
diff --git a/packages/runtime/src/agent/index.ts b/packages/runtime/src/agent/index.ts
@@ -885,14 +885,16 @@ export class BuiltInAgent extends AbstractAgent {
         const closeReasoningIfOpen = () => {
           if (!isInReasoning) return;
           isInReasoning = false;
-          subscriber.next({
+          const reasoningMsgEnd: ReasoningMessageEndEvent = {
             type: EventType.REASONING_MESSAGE_END,
             messageId: reasoningMessageId,
-          } as ReasoningMessageEndEvent);
-          subscriber.next({
+          };
+          subscriber.next(reasoningMsgEnd);
+          const reasoningEnd: ReasoningEndEvent = {
             type: EventType.REASONING_END,
             messageId: reasoningMessageId,
-          } as ReasoningEndEvent);
+          };
+          subscriber.next(reasoningEnd);
         };
 
         try {
@@ -1008,9 +1010,14 @@ export class BuiltInAgent extends AbstractAgent {
 
           // Process fullStream events
           for await (const part of response.fullStream) {
+            // Close any open reasoning lifecycle on every event except
+            // reasoning-delta, which arrives mid-block and must not interrupt it.
+            if (part.type !== "reasoning-delta") {
+              closeReasoningIfOpen();
+            }
+
             switch (part.type) {
               case "abort": {
-                closeReasoningIfOpen();
                 const abortEndEvent: RunFinishedEvent = {
                   type: EventType.RUN_FINISHED,
                   threadId: input.threadId,
@@ -1024,7 +1031,6 @@ export class BuiltInAgent extends AbstractAgent {
                 break;
               }
               case "reasoning-start": {
-                closeReasoningIfOpen();
                 // Use SDK-provided id, or generate a fresh UUID if id is falsy/"0"
                 // to prevent consecutive reasoning blocks from sharing a messageId
                 const providedId = "id" in part ? part.id : undefined;
@@ -1058,14 +1064,11 @@ export class BuiltInAgent extends AbstractAgent {
                 break;
               }
               case "reasoning-end": {
-                // Use closeReasoningIfOpen so this is idempotent: if a phase-transition
-                // case already auto-closed reasoning, a late reasoning-end from the SDK
-                // becomes a no-op instead of emitting duplicate end events.
-                closeReasoningIfOpen();
+                // closeReasoningIfOpen() already called before the switch — no-op here
+                // if the SDK never emits this event (e.g. @ai-sdk/anthropic).
                 break;
               }
               case "tool-input-start": {
-                closeReasoningIfOpen();
                 const toolCallId = part.id;
                 const state = ensureToolCallState(toolCallId);
                 state.toolName = part.toolName;
@@ -1101,7 +1104,6 @@ export class BuiltInAgent extends AbstractAgent {
               }
 
               case "text-start": {
-                closeReasoningIfOpen();
                 // New text message starting - use the SDK-provided id
                 // Use randomUUID() if part.id is falsy or "0" to prevent message merging issues
                 const providedId = "id" in part ? part.id : undefined;
@@ -1127,7 +1129,6 @@ export class BuiltInAgent extends AbstractAgent {
               }
 
               case "tool-call": {
-                closeReasoningIfOpen();
                 const toolCallId = part.toolCallId;
                 const state = ensureToolCallState(toolCallId);
                 state.toolName = part.toolName ?? state.toolName;
@@ -1224,7 +1225,6 @@ export class BuiltInAgent extends AbstractAgent {
               }
 
               case "finish": {
-                closeReasoningIfOpen();
                 // Emit run finished event
                 const finishedEvent: RunFinishedEvent = {
                   type: EventType.RUN_FINISHED,
@@ -1240,7 +1240,6 @@ export class BuiltInAgent extends AbstractAgent {
               }
 
               case "error": {
-                closeReasoningIfOpen();
                 if (abortController.signal.aborted) {
                   break;
                 }
