From fb1cb99c880f6a01207f22a4dcadced4889c37ed Mon Sep 17 00:00:00 2001
From: Vincent Koc <vincentkoc@ieee.org>
Date: Sat, 4 Apr 2026 11:21:29 +0900
Subject: [PATCH] fix(xai): narrow stream wrapper params

---
 extensions/xai/index.ts | 30 +++++++++++++++++-------------
 1 file changed, 17 insertions(+), 13 deletions(-)

diff --git a/extensions/xai/index.ts b/extensions/xai/index.ts
index 9844caeba9d..db8a1fb026b 100644
--- a/extensions/xai/index.ts
+++ b/extensions/xai/index.ts
@@ -256,24 +256,28 @@ export default defineSingleProviderPluginEntry({
     },
     buildReplayPolicy: (ctx) => buildOpenAICompatibleReplayPolicy(ctx.modelApi),
     prepareExtraParams: (ctx) => {
-      if (ctx.extraParams?.tool_stream !== undefined) {
-        return ctx.extraParams;
+      const extraParams = ctx.extraParams;
+      if (extraParams && extraParams.tool_stream !== undefined) {
+        return extraParams;
       }
       return {
-        ...ctx.extraParams,
+        ...(extraParams ?? {}),
         tool_stream: true,
       };
     },
-    wrapStreamFn: (ctx) =>
-      composeProviderStreamWrappers(
-        ctx.streamFn,
-        (streamFn) => createXaiToolPayloadCompatibilityWrapper(streamFn),
-        typeof ctx.extraParams?.fastMode === "boolean"
-          ? (streamFn) => createXaiFastModeWrapper(streamFn, ctx.extraParams.fastMode)
-          : undefined,
-        (streamFn) => createXaiToolCallArgumentDecodingWrapper(streamFn),
-        (streamFn) => createToolStreamWrapper(streamFn, ctx.extraParams?.tool_stream !== false),
-      ),
+    wrapStreamFn: (ctx) => {
+      const extraParams = ctx.extraParams;
+      const fastMode = extraParams?.fastMode;
+      const toolStreamEnabled = extraParams?.tool_stream !== false;
+      return composeProviderStreamWrappers(ctx.streamFn, (streamFn) => {
+        let wrappedStreamFn = createXaiToolPayloadCompatibilityWrapper(streamFn);
+        if (typeof fastMode === "boolean") {
+          wrappedStreamFn = createXaiFastModeWrapper(wrappedStreamFn, fastMode);
+        }
+        wrappedStreamFn = createXaiToolCallArgumentDecodingWrapper(wrappedStreamFn);
+        return createToolStreamWrapper(wrappedStreamFn, toolStreamEnabled);
+      });
+    },
     // Provider-specific fallback auth stays owned by the xAI plugin so core
     // auth/discovery code can consume it generically without parsing xAI's
     // private config layout. Callers may receive a real key from the active