@hebo-ai/gateway 0.4.2 → 0.5.0

package/src/models/anthropic/middleware.test.ts

@@ -8,6 +8,7 @@ import { claudeReasoningMiddleware } from "./middleware";
 test("claudeReasoningMiddleware > matching patterns", () => {
   const matching = [
     "anthropic/claude-opus-4.6",
+    "anthropic/claude-sonnet-4.6",
     "anthropic/claude-sonnet-3.7",
     "anthropic/claude-opus-4.5",
     "anthropic/claude-sonnet-4.5",
@@ -124,7 +125,7 @@ test("claudeReasoningMiddleware > should transform reasoning object to thinking
       anthropic: {
         thinking: {
           type: "enabled",
-          budgetTokens: 2000,
+          budgetTokens: 32000,
         },
       },
       unknown: {},
@@ -262,3 +263,333 @@ test("claudeReasoningMiddleware > should clamp max_tokens for Opus 4", async ()
 
   expect(result.providerOptions?.anthropic?.thinking?.budgetTokens).toBe(32000);
 });
+
+test("claudeReasoningMiddleware > should pass through max effort for Claude 4.6", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "max",
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-opus-4.6" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "adaptive",
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("max");
+});
+
+test("claudeReasoningMiddleware > should map xhigh effort to max for Claude Opus 4.6", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "xhigh",
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-opus-4.6" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "adaptive",
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("max");
+});
+
+test("claudeReasoningMiddleware > should map max effort to high for Claude Sonnet 4.6", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "max",
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-sonnet-4.6" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "adaptive",
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("high");
+});
+
+test("claudeReasoningMiddleware > should map minimal effort to low for Claude Sonnet 4.6", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "minimal",
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-sonnet-4.6" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "adaptive",
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("low");
+});
+
+test("claudeReasoningMiddleware > should use manual thinking for Claude Sonnet 4.6 when max_tokens is provided", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "medium",
+          max_tokens: 2000,
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-sonnet-4.6" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "enabled",
+    budgetTokens: 2000,
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("medium");
+});
+
+test("claudeReasoningMiddleware > should map none effort to low for Claude Sonnet 4.5", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "none",
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-sonnet-4.5" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "enabled",
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("low");
+});
+
+test("claudeReasoningMiddleware > should include effort and max_tokens for Claude 4.6", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "medium",
+          max_tokens: 2000,
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-opus-4.6" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "adaptive",
+    budgetTokens: 2000,
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("medium");
+});
+
+test("claudeReasoningMiddleware > should clamp max_tokens to 128k for Claude Opus 4.6", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "medium",
+          max_tokens: 200000,
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-opus-4.6" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "adaptive",
+    budgetTokens: 128000,
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("medium");
+});
+
+test("claudeReasoningMiddleware > should include effort and max_tokens for Claude Sonnet 4.5", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "medium",
+          max_tokens: 2000,
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-sonnet-4.5" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "enabled",
+    budgetTokens: 2000,
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("medium");
+});
+
+test("claudeReasoningMiddleware > should map max effort to high for Claude Sonnet 4.5", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "max",
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-sonnet-4.5" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "enabled",
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("high");
+});
+
+test("claudeReasoningMiddleware > should map xhigh effort to high for Claude Sonnet 4.5", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "xhigh",
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-sonnet-4.5" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "enabled",
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("high");
+});
+
+test("claudeReasoningMiddleware > should keep xhigh as budget for non-4.6 models", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "xhigh",
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-sonnet-4" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking?.budgetTokens).toBe(60800);
+});
+
+test("claudeReasoningMiddleware > should map xhigh effort for Claude Opus 4.5 without default budget", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: {
+          enabled: true,
+          effort: "xhigh",
+        },
+      },
+    },
+  };
+
+  const result = await claudeReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-opus-4.5" }),
+  });
+
+  expect(result.providerOptions?.anthropic?.thinking).toEqual({
+    type: "enabled",
+  });
+  expect(result.providerOptions?.anthropic?.effort).toBe("high");
+});

package/src/models/anthropic/middleware.ts

@@ -1,21 +1,66 @@
 import type { LanguageModelMiddleware } from "ai";
 
-import type { ChatCompletionsReasoningConfig } from "../../endpoints/chat-completions/schema";
+import type {
+  ChatCompletionsReasoningConfig,
+  ChatCompletionsReasoningEffort,
+} from "../../endpoints/chat-completions/schema";
 
 import { modelMiddlewareMatcher } from "../../middleware/matcher";
 import { calculateReasoningBudgetFromEffort } from "../../middleware/utils";
 
-const CLAUDE_MAX_OUTPUT_TOKENS = 64000;
-const CLAUDE_OPUS_4_MAX_OUTPUT_TOKENS = 32000;
+const isClaude = (family: "opus" | "sonnet" | "haiku", version: string) => {
+  const dashed = version.replace(".", "-");
 
-function getMaxOutputTokens(modelId: string): number {
-  if (!modelId.includes("opus-4")) return CLAUDE_MAX_OUTPUT_TOKENS;
-  if (modelId.includes("opus-4.5") || modelId.includes("opus-4-5")) {
-    return CLAUDE_MAX_OUTPUT_TOKENS;
+  return (modelId: string) =>
+    modelId.includes(`claude-${family}-${version}`) ||
+    modelId.includes(`claude-${family}-${dashed}`);
+};
+
+const isOpus46 = isClaude("opus", "4.6");
+const isOpus45 = isClaude("opus", "4.5");
+const isOpus4 = isClaude("opus", "4");
+const isSonnet46 = isClaude("sonnet", "4.6");
+const isSonnet45 = isClaude("sonnet", "4.5");
+
+export function mapClaudeReasoningEffort(effort: ChatCompletionsReasoningEffort, modelId: string) {
+  if (isOpus46(modelId)) {
+    switch (effort) {
+      case "none":
+      case "minimal":
+      case "low":
+        return "low";
+      case "medium":
+        return "medium";
+      case "high":
+        return "high";
+      case "xhigh":
+      case "max":
+        return "max";
+    }
+  }
+
+  switch (effort) {
+    case "none":
+    case "minimal":
+    case "low":
+      return "low";
+    case "medium":
+      return "medium";
+    case "high":
+    case "xhigh":
+    case "max":
+      return "high";
   }
-  return CLAUDE_OPUS_4_MAX_OUTPUT_TOKENS;
 }
 
+function getMaxOutputTokens(modelId: string): number {
+  if (isOpus46(modelId)) return 128_000;
+  if (isOpus45(modelId)) return 64_000;
+  if (isOpus4(modelId)) return 32_000;
+  return 64_000;
+}
+
+// https://platform.claude.com/docs/en/build-with-claude/effort
 export const claudeReasoningMiddleware: LanguageModelMiddleware = {
   specificationVersion: "v3",
   // eslint-disable-next-line require-await
@@ -27,23 +72,42 @@ export const claudeReasoningMiddleware: LanguageModelMiddleware = {
     if (!reasoning) return params;
 
     const target = (params.providerOptions!["anthropic"] ??= {});
+    const modelId = model.modelId;
+    const clampedMaxTokens =
+      reasoning.max_tokens && Math.min(reasoning.max_tokens, getMaxOutputTokens(modelId));
 
     if (!reasoning.enabled) {
       target["thinking"] = { type: "disabled" };
-    } else if (reasoning.max_tokens) {
-      target["thinking"] = {
-        type: "enabled",
-        budgetTokens: Math.min(reasoning.max_tokens, getMaxOutputTokens(model.modelId)),
-      };
     } else if (reasoning.effort) {
-      // FUTURE: warn that reasoning.max_tokens was computed
+      if (isOpus46(modelId)) {
+        target["thinking"] = clampedMaxTokens
+          ? { type: "adaptive", budgetTokens: clampedMaxTokens }
+          : { type: "adaptive" };
+        target["effort"] = mapClaudeReasoningEffort(reasoning.effort, modelId);
+      } else if (isSonnet46(modelId)) {
+        target["thinking"] = clampedMaxTokens
+          ? { type: "enabled", budgetTokens: clampedMaxTokens }
+          : { type: "adaptive" };
+        target["effort"] = mapClaudeReasoningEffort(reasoning.effort, modelId);
+      } else if (isOpus45(modelId) || isSonnet45(modelId)) {
+        target["thinking"] = { type: "enabled" };
+        if (clampedMaxTokens) target["thinking"]["budgetTokens"] = clampedMaxTokens;
+        target["effort"] = mapClaudeReasoningEffort(reasoning.effort, modelId);
+      } else {
+        // FUTURE: warn that reasoning.max_tokens was computed
+        target["thinking"] = {
+          type: "enabled",
+          budgetTokens: calculateReasoningBudgetFromEffort(
+            reasoning.effort,
+            params.maxOutputTokens ?? getMaxOutputTokens(modelId),
+            1024,
+          ),
+        };
+      }
+    } else if (clampedMaxTokens) {
       target["thinking"] = {
         type: "enabled",
-        budgetTokens: calculateReasoningBudgetFromEffort(
-          reasoning.effort,
-          params.maxOutputTokens ?? getMaxOutputTokens(model.modelId),
-          1024,
-        ),
+        budgetTokens: clampedMaxTokens,
       };
     } else {
       target["thinking"] = { type: "enabled" };

package/src/models/google/middleware.test.ts

@@ -13,6 +13,7 @@ test("geminiReasoningMiddleware > matching patterns", () => {
     "google/gemini-2.5-pro",
     "google/gemini-3-flash-preview",
     "google/gemini-3-pro-preview",
+    "google/gemini-3.1-pro-preview",
   ] satisfies (typeof CANONICAL_MODEL_IDS)[number][];
 
   const nonMatching = ["google/gemini-1.5-pro", "google/gemini-1.5-flash"];
@@ -108,6 +109,36 @@ test("geminiReasoningMiddleware > should map effort for Gemini 3 Pro", async ()
   });
 });
 
+test("geminiReasoningMiddleware > should map medium effort for Gemini 3.1 Pro", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      unknown: {
+        reasoning: { enabled: true, effort: "medium" },
+      },
+    },
+  };
+
+  const result = await geminiReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "google/gemini-3.1-pro-preview" }),
+  });
+
+  expect(result).toEqual({
+    prompt: [],
+    providerOptions: {
+      google: {
+        thinkingConfig: {
+          includeThoughts: true,
+          thinkingLevel: "medium",
+        },
+      },
+      unknown: {},
+    },
+  });
+});
+
 test("geminiReasoningMiddleware > should use budget for Gemini 2", async () => {
   const params = {
     prompt: [],

package/src/models/google/middleware.ts

@@ -26,10 +26,26 @@ export const geminiDimensionsMiddleware: EmbeddingModelMiddleware = {
   },
 };
 
+// https://ai.google.dev/gemini-api/docs/thinking#thinking-levels
 export function mapGeminiReasoningEffort(
   effort: ChatCompletionsReasoningEffort,
   modelId: string,
 ): ChatCompletionsReasoningEffort | undefined {
+  if (modelId.includes("gemini-3.1-pro")) {
+    switch (effort) {
+      case "none":
+      case "minimal":
+      case "low":
+        return "low";
+      case "medium":
+        return "medium";
+      case "high":
+      case "xhigh":
+      case "max":
+        return "high";
+    }
+  }
+
   if (modelId.includes("gemini-3-pro")) {
     switch (effort) {
       case "none":
@@ -39,6 +55,7 @@ export function mapGeminiReasoningEffort(
       case "medium":
       case "high":
       case "xhigh":
+      case "max":
         return "high";
     }
   }
@@ -54,6 +71,7 @@ export function mapGeminiReasoningEffort(
         return "medium";
       case "high":
       case "xhigh":
+      case "max":
         return "high";
     }
   }

package/src/models/google/presets.ts

@@ -57,6 +57,16 @@ export const gemini3ProPreview = presetFor<CanonicalModelId, CatalogModel>()(
   } satisfies DeepPartial<CatalogModel>,
 );
 
+export const gemini31ProPreview = presetFor<CanonicalModelId, CatalogModel>()(
+  "google/gemini-3.1-pro-preview" as const,
+  {
+    ...GEMINI_BASE,
+    name: "Gemini 3.1 Pro (Preview)",
+    created: "2026-02-19",
+    knowledge: "2025-01",
+  } satisfies DeepPartial<CatalogModel>,
+);
+
 export const gemini25FlashLite = presetFor<CanonicalModelId, CatalogModel>()(
   "google/gemini-2.5-flash-lite" as const,
   {
@@ -90,19 +100,20 @@ export const gemini25Pro = presetFor<CanonicalModelId, CatalogModel>()(
 const geminiAtomic = {
   "v2.5": [gemini25FlashLite, gemini25Flash, gemini25Pro],
   "v3-preview": [gemini3FlashPreview, gemini3ProPreview],
+  "v3.1-preview": [gemini31ProPreview],
   embeddings: [geminiEmbedding001],
 } as const;
 
 const geminiGroups = {
   "v2.x": [...geminiAtomic["v2.5"]],
-  "v3.x": [...geminiAtomic["v3-preview"]],
+  "v3.x": [...geminiAtomic["v3-preview"], ...geminiAtomic["v3.1-preview"]],
 } as const;
 
 export const gemini = {
   ...geminiAtomic,
   ...geminiGroups,
   latest: [...geminiAtomic["v2.5"]],
-  preview: [...geminiAtomic["v3-preview"]],
+  preview: [...geminiAtomic["v3-preview"], ...geminiAtomic["v3.1-preview"]],
   embeddings: [...geminiAtomic["embeddings"]],
   all: Object.values(geminiAtomic).flat(),
 } as const;

package/src/models/openai/middleware.ts

@@ -33,6 +33,7 @@ function mapGptOssReasoningEffort(
       return "medium";
     case "high":
     case "xhigh":
+    case "max":
       return "high";
     default:
       return "low";

package/src/models/types.ts

@@ -46,6 +46,7 @@ export const CANONICAL_MODEL_IDS = [
   "google/gemini-2.5-pro",
   "google/gemini-3-flash-preview",
   "google/gemini-3-pro-preview",
+  "google/gemini-3.1-pro-preview",
   "google/embedding-001",
   // Meta
   "meta/llama-3.1-8b",

package/src/providers/bedrock/index.ts

@@ -1 +1,2 @@
 export * from "./canonical";
+export * from "./middleware";

package/src/providers/bedrock/middleware.test.ts

@@ -0,0 +1,73 @@
+import { MockLanguageModelV3 } from "ai/test";
+import { expect, test } from "bun:test";
+
+import { modelMiddlewareMatcher } from "../../middleware/matcher";
+import { bedrockAnthropicReasoningMiddleware } from "./middleware";
+
+test("bedrockAnthropicReasoningMiddleware > matching provider", () => {
+  const middleware = modelMiddlewareMatcher.resolve({
+    kind: "text",
+    modelId: "anthropic/claude-opus-4.6",
+    providerId: "amazon-bedrock",
+  });
+
+  expect(middleware).toContain(bedrockAnthropicReasoningMiddleware);
+});
+
+test("bedrockAnthropicReasoningMiddleware > should map thinking/effort into reasoningConfig", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      bedrock: {
+        thinking: {
+          type: "adaptive",
+          budgetTokens: 4096,
+        },
+        effort: "max",
+      },
+    },
+  };
+
+  const result = await bedrockAnthropicReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "anthropic/claude-opus-4.6" }),
+  });
+
+  expect(result.providerOptions?.bedrock).toEqual({
+    reasoningConfig: {
+      type: "adaptive",
+      budgetTokens: 4096,
+      maxReasoningEffort: "max",
+    },
+  });
+});
+
+test("bedrockAnthropicReasoningMiddleware > should skip non-anthropic models", async () => {
+  const params = {
+    prompt: [],
+    providerOptions: {
+      bedrock: {
+        thinking: {
+          type: "enabled",
+          budgetTokens: 4096,
+        },
+        effort: "high",
+      },
+    },
+  };
+
+  const result = await bedrockAnthropicReasoningMiddleware.transformParams!({
+    type: "generate",
+    params,
+    model: new MockLanguageModelV3({ modelId: "openai/gpt-oss-20b" }),
+  });
+
+  expect(result.providerOptions?.bedrock).toEqual({
+    thinking: {
+      type: "enabled",
+      budgetTokens: 4096,
+    },
+    effort: "high",
+  });
+});