From 3468aa01402740df830dfc266d7ef0c240f79c26 Mon Sep 17 00:00:00 2001 From: Aiden Cline <63023139+rekram1-node@users.noreply.github.com> Date: Fri, 31 Jul 2026 11:14:08 -0500 Subject: [PATCH] feat(ai): configure chat max tokens field (#39909) --- packages/ai/src/protocols/openai-chat.ts | 6 ++- packages/ai/src/route/transport/http.ts | 48 ------------------- packages/ai/src/schema/options.ts | 4 ++ packages/ai/test/compile.test.ts | 39 ++++++++------- .../ai/test/provider/google-vertex.test.ts | 17 ------- .../provider/openai-compatible-chat.test.ts | 14 ++++++ 6 files changed, 44 insertions(+), 84 deletions(-) diff --git a/packages/ai/src/protocols/openai-chat.ts b/packages/ai/src/protocols/openai-chat.ts index a2d8a7b541..e5793f33b1 100644 --- a/packages/ai/src/protocols/openai-chat.ts +++ b/packages/ai/src/protocols/openai-chat.ts @@ -107,6 +107,7 @@ export const bodyFields = { stream_options: Schema.optional(Schema.Struct({ include_usage: Schema.Boolean })), store: Schema.optional(Schema.Boolean), reasoning_effort: Schema.optional(OpenAIOptions.OpenAIReasoningEffort), + max_completion_tokens: Schema.optional(Schema.Number), max_tokens: Schema.optional(Schema.Number), temperature: Schema.optional(Schema.Number), top_p: Schema.optional(Schema.Number), @@ -415,6 +416,7 @@ const fromRequest = Effect.fn("OpenAIChat.fromRequest")(function* (request: LLMR ) const generation = request.generation const toolSchemaCompatibility = request.model.compatibility?.toolSchema + const maxTokensField = request.model.compatibility?.maxTokensField ?? "max_tokens" return { model: request.model.id, messages: yield* lowerMessages(request), @@ -427,7 +429,9 @@ const fromRequest = Effect.fn("OpenAIChat.fromRequest")(function* (request: LLMR tool_choice: request.toolChoice ? yield* lowerToolChoice(request.toolChoice) : undefined, stream: true as const, stream_options: { include_usage: true }, - max_tokens: generation?.maxTokens, + ...(maxTokensField === "max_completion_tokens" + ? { max_completion_tokens: generation?.maxTokens } + : { max_tokens: generation?.maxTokens }), temperature: generation?.temperature, top_p: generation?.topP, frequency_penalty: generation?.frequencyPenalty, diff --git a/packages/ai/src/route/transport/http.ts b/packages/ai/src/route/transport/http.ts index ae8b5e088d..a1de9ec573 100644 --- a/packages/ai/src/route/transport/http.ts +++ b/packages/ai/src/route/transport/http.ts @@ -28,57 +28,9 @@ const applyQuery = (url: string, query: Record | undefined) => { return next.toString() } -const PROTOCOL_BODY_OVERLAY_DENYLIST = new Set([ - "anthropic_version", - "content", - "contents", - "frequencyPenalty", - "frequency_penalty", - "generationConfig", - "inferenceConfig", - "input", - "maxTokens", - "max_tokens", - "messages", - "model", - "presencePenalty", - "presence_penalty", - "responseFormat", - "response_format", - "seed", - "stop", - "stopSequences", - "stop_sequences", - "stream", - "streamOptions", - "stream_options", - "system", - "systemInstruction", - "system_instruction", - "temperature", - "thinking", - "toolChoice", - "toolConfig", - "tool_choice", - "tool_config", - "tools", - "topK", - "topP", - "top_k", - "top_p", -]) - -const forbiddenBodyOverlayKeys = (body: Record) => - Object.keys(body).filter((key) => PROTOCOL_BODY_OVERLAY_DENYLIST.has(key)) - const bodyWithOverlay = (body: Body, request: LLMRequest, encodeBody: (body: Body) => string) => Effect.gen(function* () { if (request.http?.body === undefined) return { jsonBody: body, bodyText: encodeBody(body) } - const forbiddenKeys = forbiddenBodyOverlayKeys(request.http.body) - if (forbiddenKeys.length > 0) - return yield* ProviderShared.invalidRequest( - `http.body cannot overlay protocol-owned field(s): ${forbiddenKeys.join(", ")}`, - ) if (ProviderShared.isRecord(body)) { const overlaid = mergeJsonRecords(body, request.http.body) ?? {} return { jsonBody: overlaid, bodyText: ProviderShared.encodeJson(overlaid) } diff --git a/packages/ai/src/schema/options.ts b/packages/ai/src/schema/options.ts index bb77022c95..62c606056d 100644 --- a/packages/ai/src/schema/options.ts +++ b/packages/ai/src/schema/options.ts @@ -167,9 +167,13 @@ export namespace ModelDefaults { export const ModelToolSchemaCompatibility = Schema.Literals(["gemini", "moonshot"]) export type ModelToolSchemaCompatibility = Schema.Schema.Type +export const ModelMaxTokensFieldCompatibility = Schema.Literals(["max_completion_tokens", "max_tokens"]) +export type ModelMaxTokensFieldCompatibility = Schema.Schema.Type + export class ModelCompatibility extends Schema.Class("LLM.ModelCompatibility")({ toolSchema: Schema.optional(ModelToolSchemaCompatibility), reasoningField: Schema.optional(Schema.String), + maxTokensField: Schema.optional(ModelMaxTokensFieldCompatibility), }) {} export namespace ModelCompatibility { diff --git a/packages/ai/test/compile.test.ts b/packages/ai/test/compile.test.ts index 9c8dbbdfac..27971bf5c8 100644 --- a/packages/ai/test/compile.test.ts +++ b/packages/ai/test/compile.test.ts @@ -171,24 +171,27 @@ describe("request option precedence", () => { ), ) - it.effect("rejects raw body overlays for protocol-owned roots", () => - Effect.gen(function* () { - const model = OpenAIChat.route - .with({ endpoint: { baseURL: "https://api.openai.test/v1/" }, auth: Auth.bearer("test") }) - .model({ id: "gpt-4o-mini" }) - const error = yield* compileRequest( - LLM.request({ - model, - prompt: "Say hello.", - http: { body: { model: "gpt-5", messages: [], tools: [] } }, - }), - ).pipe(Effect.flip) - - expect(error.reason).toMatchObject({ - _tag: "InvalidRequest", - message: "http.body cannot overlay protocol-owned field(s): model, messages, tools", - }) - }), + it.effect("applies raw body overlays after protocol lowering", () => + LLMClient.generate( + LLM.request({ + model: OpenAIChat.route + .with({ endpoint: { baseURL: "https://api.openai.test/v1/" }, auth: Auth.bearer("test") }) + .model({ id: "gpt-4o-mini" }), + prompt: "Say hello.", + http: { body: { model: "gpt-5", messages: [], tools: [] } }, + }), + ).pipe( + Effect.provide( + dynamicResponse((input) => + Effect.gen(function* () { + expect(decodeJson(input.text)).toMatchObject({ model: "gpt-5", messages: [], tools: [] }) + return input.respond(sseEvents(deltaChunk({}, "stop")), { + headers: { "content-type": "text/event-stream" }, + }) + }), + ), + ), + ), ) it.effect("uses model output limits after route limits and before call maxTokens", () => diff --git a/packages/ai/test/provider/google-vertex.test.ts b/packages/ai/test/provider/google-vertex.test.ts index 8c031d473c..f260f3577b 100644 --- a/packages/ai/test/provider/google-vertex.test.ts +++ b/packages/ai/test/provider/google-vertex.test.ts @@ -181,23 +181,6 @@ describe("Google Vertex providers", () => { }), ) - it.effect("protects the Vertex Messages API version from body overlays", () => - Effect.gen(function* () { - const error = yield* compileRequest( - LLM.request({ - model: GoogleVertexMessages.configure({ - accessToken: "vertex-token", - http: { body: { anthropic_version: "wrong" } }, - project: "vertex-project", - }).model("claude-sonnet-4-6"), - prompt: "Say hello.", - }), - ).pipe(Effect.flip) - - expect(error.message).toContain("http.body cannot overlay protocol-owned field(s): anthropic_version") - }), - ) - it.effect("routes tuned Gemini models through their deployed endpoint", () => Effect.gen(function* () { const response = yield* LLMClient.generate( diff --git a/packages/ai/test/provider/openai-compatible-chat.test.ts b/packages/ai/test/provider/openai-compatible-chat.test.ts index 397c64c8b3..25ac410bed 100644 --- a/packages/ai/test/provider/openai-compatible-chat.test.ts +++ b/packages/ai/test/provider/openai-compatible-chat.test.ts @@ -144,6 +144,20 @@ describe("OpenAI-compatible Chat route", () => { }), ) + it.effect("configures the max tokens request field", () => + Effect.gen(function* () { + const compatible = OpenAICompatibleChat.route + .with({ provider: "custom", endpoint: { baseURL: "https://api.custom.test/v1" } }) + .model({ id: "custom-model", compatibility: { maxTokensField: "max_completion_tokens" } }) + const prepared = yield* compileRequest( + LLM.request({ model: compatible, prompt: "Say hello.", generation: { maxTokens: 20 } }), + ) + + expect(prepared.body).toMatchObject({ max_completion_tokens: 20 }) + expect(prepared.body).not.toHaveProperty("max_tokens") + }), + ) + it.effect("matches AI SDK compatible tool request body fixture", () => Effect.gen(function* () { const prepared = yield* compileRequest(