From 93abf5ee4dfed42770e168ebc3f52a34a8a88858 Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Sun, 12 Apr 2026 16:34:35 +0100
Subject: [PATCH] test(agents): share websocket tool followup helper

---
 src/agents/openai-ws-stream.e2e.test.ts | 72 ++++++++++++++-----------
 1 file changed, 40 insertions(+), 32 deletions(-)

diff --git a/src/agents/openai-ws-stream.e2e.test.ts b/src/agents/openai-ws-stream.e2e.test.ts
index 3438838e303..c645eed15cd 100644
--- a/src/agents/openai-ws-stream.e2e.test.ts
+++ b/src/agents/openai-ws-stream.e2e.test.ts
@@ -88,6 +88,32 @@ function makeToolResultMessage(
   } as unknown as StreamFnParams[1]["messages"][number];
 }
 
+async function runWebsocketToolFollowupTurn(params: {
+  streamFn: ReturnType<StreamFactory>;
+  context: StreamFnParams[1];
+  firstDone: AssistantMessage;
+  toolCallId: string;
+  output: string;
+}) {
+  const secondContext = {
+    ...params.context,
+    messages: [
+      ...params.context.messages,
+      params.firstDone,
+      makeToolResultMessage(params.toolCallId, params.output),
+    ],
+  } as unknown as StreamFnParams[1];
+
+  return expectDone(
+    await collectEvents(
+      params.streamFn(model, secondContext, {
+        transport: "websocket",
+        maxTokens: 128,
+      }),
+    ),
+  );
+}
+
 async function collectEvents(stream: StreamReturn): Promise<AssistantMessageEvent[]> {
   const events: AssistantMessageEvent[] = [];
   const resolvedStream: AssistantMessageEventStream = await stream;
@@ -256,22 +282,13 @@ describe("OpenAI WebSocket e2e", () => {
       expect(toolCall?.name).toBe("noop");
       expect(toolCall?.id).toBeTruthy();
 
-      const secondContext = {
-        ...firstContext,
-        messages: [
-          ...firstContext.messages,
-          firstDone,
-          makeToolResultMessage(toolCall!.id, "TOOL_OK"),
-        ],
-      } as unknown as StreamFnParams[1];
-      const secondDone = expectDone(
-        await collectEvents(
-          streamFn(model, secondContext, {
-            transport: "websocket",
-            maxTokens: 128,
-          }),
-        ),
-      );
+      const secondDone = await runWebsocketToolFollowupTurn({
+        streamFn,
+        context: firstContext,
+        firstDone,
+        toolCallId: toolCall!.id,
+        output: "TOOL_OK",
+      });
 
       expect(assistantText(secondDone)).toMatch(/TOOL_OK/);
     },
@@ -340,22 +357,13 @@ describe("OpenAI WebSocket e2e", () => {
         rawToolCall ? `${rawToolCall.call_id}|${rawToolCall.id}` : undefined,
       );
 
-      const secondContext = {
-        ...firstContext,
-        messages: [
-          ...firstContext.messages,
-          firstDone,
-          makeToolResultMessage(toolCall!.id, "TOOL_OK"),
-        ],
-      } as unknown as StreamFnParams[1];
-      const secondDone = expectDone(
-        await collectEvents(
-          streamFn(model, secondContext, {
-            transport: "websocket",
-            maxTokens: 128,
-          }),
-        ),
-      );
+      const secondDone = await runWebsocketToolFollowupTurn({
+        streamFn,
+        context: firstContext,
+        firstDone,
+        toolCallId: toolCall!.id,
+        output: "TOOL_OK",
+      });
 
       expect(assistantText(secondDone)).toMatch(/TOOL_OK/);
     },