From f05b789736004a8ee93b7a8b44907b5f7f72ee1a Mon Sep 17 00:00:00 2001
From: Josh Lehman <josh@martian.engineering>
Date: Wed, 29 Apr 2026 12:56:27 -0700
Subject: [PATCH] test: align Codex provider payload test (#74536)

---
 extensions/openai/openai-provider.test.ts | 79 ++++++++++++++++++++---
 1 file changed, 69 insertions(+), 10 deletions(-)

diff --git a/extensions/openai/openai-provider.test.ts b/extensions/openai/openai-provider.test.ts
index fb5855e9390..367903d47d0 100644
--- a/extensions/openai/openai-provider.test.ts
+++ b/extensions/openai/openai-provider.test.ts
@@ -1,15 +1,61 @@
 import type { StreamFn } from "@mariozechner/pi-agent-core";
 import type { Context, Model, SimpleStreamOptions } from "@mariozechner/pi-ai";
-import { describe, expect, it, vi } from "vitest";
+import { beforeEach, describe, expect, it, vi } from "vitest";
 import { buildOpenAICodexProviderPlugin } from "./openai-codex-provider.js";
 import { buildOpenAIProvider } from "./openai-provider.js";
 
-const refreshOpenAICodexTokenMock = vi.hoisted(() => vi.fn());
+const mocks = vi.hoisted(() => ({
+  refreshOpenAICodexToken: vi.fn(),
+  openAIResponsesTransportStreamFn: vi.fn(),
+}));
 
 vi.mock("./openai-codex-provider.runtime.js", () => ({
-  refreshOpenAICodexToken: refreshOpenAICodexTokenMock,
+  refreshOpenAICodexToken: mocks.refreshOpenAICodexToken,
 }));
 
+vi.mock("openclaw/plugin-sdk/provider-stream-family", async (importOriginal) => {
+  const actual =
+    await importOriginal<typeof import("openclaw/plugin-sdk/provider-stream-family")>();
+  const wrapStreamFn: NonNullable<typeof actual.OPENAI_RESPONSES_STREAM_HOOKS.wrapStreamFn> = (
+    ctx,
+  ) => {
+    let nextStreamFn = actual.createOpenAIAttributionHeadersWrapper(ctx.streamFn, {
+      codexNativeTransportStreamFn: mocks.openAIResponsesTransportStreamFn,
+    });
+
+    if (actual.resolveOpenAIFastMode(ctx.extraParams)) {
+      nextStreamFn = actual.createOpenAIFastModeWrapper(nextStreamFn);
+    }
+
+    const serviceTier = actual.resolveOpenAIServiceTier(ctx.extraParams);
+    if (serviceTier) {
+      nextStreamFn = actual.createOpenAIServiceTierWrapper(nextStreamFn, serviceTier);
+    }
+
+    const textVerbosity = actual.resolveOpenAITextVerbosity(ctx.extraParams);
+    if (textVerbosity) {
+      nextStreamFn = actual.createOpenAITextVerbosityWrapper(nextStreamFn, textVerbosity);
+    }
+
+    nextStreamFn = actual.createCodexNativeWebSearchWrapper(nextStreamFn, {
+      config: ctx.config,
+      agentDir: ctx.agentDir,
+    });
+    return actual.createOpenAIResponsesContextManagementWrapper(
+      actual.createOpenAIReasoningCompatibilityWrapper(nextStreamFn),
+      ctx.extraParams,
+    );
+  };
+
+  return {
+    ...actual,
+    OPENAI_RESPONSES_STREAM_HOOKS: {
+      ...actual.OPENAI_RESPONSES_STREAM_HOOKS,
+      wrapStreamFn,
+    },
+  };
+});
+
 function runWrappedPayloadCase(params: {
   wrap: NonNullable<ReturnType<typeof buildOpenAIProvider>["wrapStreamFn"]>;
   provider: string;
@@ -49,6 +95,13 @@ function runWrappedPayloadCase(params: {
 }
 
 describe("buildOpenAIProvider", () => {
+  beforeEach(() => {
+    mocks.openAIResponsesTransportStreamFn.mockReset();
+    mocks.openAIResponsesTransportStreamFn.mockImplementation(() => {
+      throw new Error("unexpected native OpenAI Responses transport call");
+    });
+  });
+
   it("exposes grouped model/auth picker labels for API key setup", () => {
     const provider = buildOpenAIProvider();
     const apiKey = provider.auth.find((method) => method.id === "api-key");
@@ -683,6 +736,15 @@ describe("buildOpenAIProvider", () => {
     if (!wrap) {
       throw new Error("expected Codex wrapper");
     }
+    const payload = {
+      store: false,
+      text: { verbosity: "medium" },
+      tools: [{ type: "function", name: "read" }],
+    };
+    mocks.openAIResponsesTransportStreamFn.mockImplementation((model, _context, options) => {
+      options?.onPayload?.(payload, model);
+      return {} as ReturnType<StreamFn>;
+    });
     const result = runWrappedPayloadCase({
       wrap,
       provider: "openai-codex",
@@ -720,13 +782,10 @@ describe("buildOpenAIProvider", () => {
         id: "gpt-5.4",
         baseUrl: "https://chatgpt.com/backend-api",
       } as Model<"openai-codex-responses">,
-      payload: {
-        store: false,
-        text: { verbosity: "medium" },
-        tools: [{ type: "function", name: "read" }],
-      },
+      payload,
     });
 
+    expect(mocks.openAIResponsesTransportStreamFn).toHaveBeenCalledTimes(1);
     expect(result.payload.store).toBe(false);
     expect(result.payload.service_tier).toBe("priority");
     expect(result.payload.text).toEqual({ verbosity: "high" });
@@ -749,8 +808,8 @@ describe("buildOpenAIProvider", () => {
       expires: Date.now() - 60_000,
     };
 
-    refreshOpenAICodexTokenMock.mockReset();
-    refreshOpenAICodexTokenMock.mockRejectedValueOnce(
+    mocks.refreshOpenAICodexToken.mockReset();
+    mocks.refreshOpenAICodexToken.mockRejectedValueOnce(
       new Error("Failed to extract accountId from token"),
     );