diff --git a/README.md b/README.md index 49d29a31..d23a81df 100644 --- a/README.md +++ b/README.md @@ -465,6 +465,34 @@ OPENWIKI_PROVIDER_RETRY_ATTEMPTS=3 The value must be a positive integer. If the value is unset, OpenWiki defaults to 3 retries. +### Model output token limit + +To override the maximum number of tokens generated in a model response, set +`OPENWIKI_MAX_OUTPUT_TOKENS`: + +```bash +OPENWIKI_MAX_OUTPUT_TOKENS=8192 +``` + +The value must be a positive integer. If unset, OpenWiki does not override the +model client's output token limit. Provider and model limits still apply; +unsupported values may be rejected, while very small values can truncate +responses or tool calls. + +### Bedrock stream idle timeout + +For the Bedrock provider, set `OPENWIKI_STREAM_IDLE_TIMEOUT` to control how +long the client waits for the first or next streamed response chunk: + +```bash +OPENWIKI_STREAM_IDLE_TIMEOUT=300000 +``` + +The value is milliseconds and must be an integer from `0` to `2147483647`. +Set it to `0` to disable the watchdog. If unset, OpenWiki preserves the +`@langchain/aws` provider default. Prefer a sufficiently long finite timeout +over disabling the watchdog so a stalled stream cannot hang forever. + ### Diagrams OpenWiki embeds **Mermaid** diagrams in the generated wiki wherever they make a diff --git a/src/agent/index.ts b/src/agent/index.ts index a3f45418..ccc5bd6a 100644 --- a/src/agent/index.ts +++ b/src/agent/index.ts @@ -91,9 +91,11 @@ import { OPENAI_COMPATIBLE_BASE_URL_ENV_KEY, OPENROUTER_API_KEY_ENV_KEY, OPENROUTER_BASE_URL, + OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY, OPENWIKI_MODEL_ID_ENV_KEY, OPENWIKI_PROVIDER_ENV_KEY, OPENWIKI_PROVIDER_RETRY_ATTEMPTS_ENV_KEY, + OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY, providerRequiresBaseUrl, providerRequiresRegion, providerRequiresSecretKey, @@ -101,11 +103,13 @@ import { providerUsesExternalCliAuth, providerUsesResponsesApi, resolveConfiguredProvider, + resolveMaxOutputTokens, resolveOpenRouterProviderOnly, resolveProviderBaseUrl, resolveProviderLocation, resolveProviderRegion, resolveProviderRetryAttempts, + resolveStreamIdleTimeoutForProvider, type OpenWikiProvider, } from "../constants.js"; import { @@ -213,6 +217,16 @@ export async function runOpenWikiAgent( emitDebug(options, `model=${modelId}`); const providerRetryAttempts = resolveProviderRetryAttempts(); emitDebug(options, `provider.retryAttempts=${providerRetryAttempts}`); + const maxOutputTokens = resolveMaxOutputTokens(); + emitDebug( + options, + `model.maxOutputTokens=${maxOutputTokens ?? "provider-default"}`, + ); + const streamIdleTimeout = resolveStreamIdleTimeoutForProvider(provider); + emitDebug( + options, + `model.streamIdleTimeout=${streamIdleTimeout ?? "provider-default"}`, + ); const result = await runOpenWikiAgentCore( command, @@ -221,6 +235,8 @@ export async function runOpenWikiAgent( provider, modelId, providerRetryAttempts, + maxOutputTokens, + streamIdleTimeout, ); await recordRunSafe(command, options, { @@ -251,6 +267,8 @@ async function runOpenWikiAgentCore( provider: OpenWikiProvider, modelId: string, providerRetryAttempts: number, + maxOutputTokens: number | undefined, + streamIdleTimeout: number | undefined, ): Promise { const outputMode = options.outputMode ?? "local-wiki"; const context = await createRunContext( @@ -265,7 +283,13 @@ async function runOpenWikiAgentCore( ? null : await createOpenWikiContentSnapshot(cwd, outputMode); emitDebug(options, "openwiki.snapshot=created"); - const model = createModel(provider, modelId, providerRetryAttempts); + const model = createModel( + provider, + modelId, + providerRetryAttempts, + maxOutputTokens, + streamIdleTimeout, + ); emitDebug(options, `model.provider=${provider}`); emitDebug(options, "model=initialized"); const threadId = options.threadId ?? createThreadId(cwd, createRunThreadId()); @@ -802,8 +826,16 @@ export function createModel( provider: OpenWikiProvider, modelId: string, providerRetryAttempts: number, + maxOutputTokens?: number, + streamIdleTimeout?: number, ) { const retryOptions = { maxRetries: providerRetryAttempts }; + const maxTokensOptions = + maxOutputTokens === undefined ? {} : { maxTokens: maxOutputTokens }; + const googleMaxOutputTokensOptions = + maxOutputTokens === undefined ? {} : { maxOutputTokens }; + const streamIdleTimeoutOptions = + streamIdleTimeout === undefined ? {} : { streamIdleTimeout }; if (provider === "gemini") { return new ChatGoogle({ @@ -812,6 +844,7 @@ export function createModel( platformType: "gai", // Gemini 3.x thought-signature round-trip; see the constant's comment. ...GEMINI_THOUGHT_SIGNATURE_OPTIONS, + ...googleMaxOutputTokensOptions, ...retryOptions, }); } @@ -835,6 +868,7 @@ export function createModel( projectId, location, retryOptions, + maxOutputTokens, ); } @@ -844,6 +878,7 @@ export function createModel( return new ChatAnthropic(modelId, { apiKey: getProviderApiKey(provider), ...(baseURL ? { anthropicApiUrl: baseURL } : {}), + ...maxTokensOptions, ...retryOptions, }); } @@ -871,6 +906,7 @@ export function createModel( // every generation — including the non-streaming `.invoke()` calls // DeepAgents' agent node issues internally. streaming: true, + ...maxTokensOptions, ...retryOptions, configuration: { baseURL: CODEX_RESPONSES_BASE_URL, @@ -893,6 +929,7 @@ export function createModel( model: modelId, provider: providerOnly ? { only: providerOnly } : undefined, siteName: "OpenWiki", + ...maxTokensOptions, ...retryOptions, }); } @@ -901,6 +938,8 @@ export function createModel( return new ChatBedrockConverse({ model: modelId, region: resolveProviderRegion(provider), + ...maxTokensOptions, + ...streamIdleTimeoutOptions, ...retryOptions, }); } @@ -916,6 +955,7 @@ export function createModel( : undefined, model: modelId, useResponsesApi: providerUsesResponsesApi(provider, modelId), + ...maxTokensOptions, ...retryOptions, }); } @@ -982,7 +1022,13 @@ function createGeminiEnterpriseModel( projectId: string, location: string, retryOptions: { maxRetries: number }, + maxOutputTokens?: number, ) { + const maxTokensOptions = + maxOutputTokens === undefined ? {} : { maxTokens: maxOutputTokens }; + const googleMaxOutputTokensOptions = + maxOutputTokens === undefined ? {} : { maxOutputTokens }; + switch (resolveVertexSurface(modelId)) { case "anthropic": // No JS-native Claude-on-Vertex chat model exists; bridge via @@ -1012,6 +1058,7 @@ function createGeminiEnterpriseModel( dangerouslyAllowBrowser: true, }), ), + ...maxTokensOptions, ...retryOptions, }); @@ -1027,6 +1074,7 @@ function createGeminiEnterpriseModel( fetch: createVertexAuthFetch(), }, model: toVertexPublisherModel(modelId), + ...maxTokensOptions, ...retryOptions, }); @@ -1051,6 +1099,7 @@ function createGeminiEnterpriseModel( // process.env, using the `/node` entrypoint where googleAuthOptions is // typed (the default entrypoint types authOptions as `never`). googleAuthOptions: { projectId }, + ...googleMaxOutputTokensOptions, ...retryOptions, }); } @@ -1855,6 +1904,8 @@ export function formatEnvironmentDebugValue( if ( key === OPENWIKI_MODEL_ID_ENV_KEY || key === OPENWIKI_PROVIDER_ENV_KEY || + key === OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY || + key === OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY || key === OPENWIKI_PROVIDER_RETRY_ATTEMPTS_ENV_KEY || key === BEDROCK_AWS_REGION_ENV_KEY ) { diff --git a/src/constants.ts b/src/constants.ts index eb4c053d..bef8a929 100644 --- a/src/constants.ts +++ b/src/constants.ts @@ -49,6 +49,10 @@ export const GOOGLE_APPLICATION_CREDENTIALS_ENV_KEY = export const DEFAULT_VERTEX_LOCATION = "global"; export const OPENWIKI_PROVIDER_ENV_KEY = "OPENWIKI_PROVIDER"; export const OPENWIKI_MODEL_ID_ENV_KEY = "OPENWIKI_MODEL_ID"; +export const OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY = "OPENWIKI_MAX_OUTPUT_TOKENS"; +export const OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY = + "OPENWIKI_STREAM_IDLE_TIMEOUT"; +const MAX_STREAM_IDLE_TIMEOUT_MS = 2_147_483_647; export const NEBIUS_BASE_URL = "https://api.tokenfactory.nebius.com/v1/"; export const OPENWIKI_PROVIDER_RETRY_ATTEMPTS_ENV_KEY = "OPENWIKI_PROVIDER_RETRY_ATTEMPTS"; @@ -863,6 +867,73 @@ function inspectCredentialPair( return { complete: hasAccessKey && hasSecretKey, missingEnvKey: null }; } +export function resolveMaxOutputTokens( + env: NodeJS.ProcessEnv = process.env, +): number | undefined { + const rawMaxOutputTokens = env[OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY]; + + if (rawMaxOutputTokens === undefined) { + return undefined; + } + + const maxOutputTokens = rawMaxOutputTokens.trim(); + + if (!/^[1-9]\d*$/u.test(maxOutputTokens)) { + throw new Error( + `Invalid ${OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY}. Expected a positive integer.`, + ); + } + + const parsedMaxOutputTokens = Number(maxOutputTokens); + + if (!Number.isSafeInteger(parsedMaxOutputTokens)) { + throw new Error( + `Invalid ${OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY}. Expected a positive integer.`, + ); + } + + return parsedMaxOutputTokens; +} + +/** Milliseconds to wait for the first or next Bedrock stream chunk. 0 disables the watchdog. */ +export function resolveStreamIdleTimeout( + env: NodeJS.ProcessEnv = process.env, +): number | undefined { + const rawStreamIdleTimeout = env[OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY]; + + if (rawStreamIdleTimeout === undefined) { + return undefined; + } + + const streamIdleTimeout = rawStreamIdleTimeout.trim(); + + if (!/^\d+$/u.test(streamIdleTimeout)) { + throw new Error( + `Invalid ${OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY}. Expected an integer from 0 to ${MAX_STREAM_IDLE_TIMEOUT_MS} milliseconds.`, + ); + } + + const parsedStreamIdleTimeout = Number(streamIdleTimeout); + + if ( + !Number.isSafeInteger(parsedStreamIdleTimeout) || + parsedStreamIdleTimeout > MAX_STREAM_IDLE_TIMEOUT_MS + ) { + throw new Error( + `Invalid ${OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY}. Expected an integer from 0 to ${MAX_STREAM_IDLE_TIMEOUT_MS} milliseconds.`, + ); + } + + return parsedStreamIdleTimeout; +} + +export function resolveStreamIdleTimeoutForProvider( + provider: OpenWikiProvider, + env: NodeJS.ProcessEnv = process.env, +): number | undefined { + return provider === "bedrock" ? resolveStreamIdleTimeout(env) : undefined; +} + export function resolveProviderRetryAttempts( env: NodeJS.ProcessEnv = process.env, ): number { diff --git a/src/env.ts b/src/env.ts index 9158943c..805d0b5c 100644 --- a/src/env.ts +++ b/src/env.ts @@ -54,10 +54,16 @@ import { OPENWIKI_X_CLIENT_SECRET_ENV_KEY, OPENWIKI_X_REFRESH_TOKEN_ENV_KEY, OPENWIKI_TAVILY_API_KEY_ENV_KEY, + OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY, OPENWIKI_MODEL_ID_ENV_KEY, OPENWIKI_PROVIDER_ENV_KEY, OPENWIKI_PROVIDER_RETRY_ATTEMPTS_ENV_KEY, + OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY, + resolveConfiguredProvider, + resolveMaxOutputTokens, resolveProviderRetryAttempts, + resolveStreamIdleTimeout, + type OpenWikiProvider, } from "./constants.js"; import { isFileNotFoundError } from "./fs-errors.js"; import { restrictDirToCurrentUser } from "./windows-acl.js"; @@ -120,6 +126,8 @@ export const MANAGED_ENV_KEYS = [ BEDROCK_AWS_REGION_ENV_KEY, OPENWIKI_PROVIDER_ENV_KEY, OPENWIKI_MODEL_ID_ENV_KEY, + OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY, + OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY, OPENWIKI_PROVIDER_RETRY_ATTEMPTS_ENV_KEY, OPENWIKI_NOTION_TOKEN_ENV_KEY, OPENWIKI_NOTION_MCP_CLIENT_ID_ENV_KEY, @@ -268,9 +276,13 @@ export async function getCredentialDiagnostics(): Promise< CredentialDiagnostic[] > { const fileEnv = await readOpenWikiEnv(); + const provider = resolveConfiguredProvider({ + ...fileEnv, + ...process.env, + }); return CREDENTIAL_DIAGNOSTIC_ENV_KEYS.map((key) => - createCredentialDiagnostic(key, fileEnv), + createCredentialDiagnostic(key, fileEnv, provider), ); } @@ -336,6 +348,7 @@ export async function saveOpenWikiEnv(updates: EnvMap): Promise { function createCredentialDiagnostic( key: CredentialDiagnostic["key"], fileEnv: EnvMap, + provider: OpenWikiProvider, ): CredentialDiagnostic { const processValue = process.env[key]; const fileValue = fileEnv[key]; @@ -364,10 +377,14 @@ function createCredentialDiagnostic( ? getModelWarnings(value) : key === OPENWIKI_PROVIDER_ENV_KEY ? getProviderWarnings(value) - : key === OPENWIKI_PROVIDER_RETRY_ATTEMPTS_ENV_KEY - ? getRetryAttemptsWarnings(value) - : (getBaseUrlDiagnosticWarnings(key, value) ?? - getCredentialWarnings(value)), + : key === OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY + ? getMaxOutputTokensWarnings(value) + : key === OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY + ? getStreamIdleTimeoutWarnings(value, provider) + : key === OPENWIKI_PROVIDER_RETRY_ATTEMPTS_ENV_KEY + ? getRetryAttemptsWarnings(value) + : (getBaseUrlDiagnosticWarnings(key, value) ?? + getCredentialWarnings(value)), }; } @@ -425,6 +442,8 @@ function isNonSecretDiagnosticKey(key: string): boolean { return ( key === OPENWIKI_MODEL_ID_ENV_KEY || key === OPENWIKI_PROVIDER_ENV_KEY || + key === OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY || + key === OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY || key === OPENWIKI_PROVIDER_RETRY_ATTEMPTS_ENV_KEY || key === OPENWIKI_OPENROUTER_PROVIDER_ONLY_ENV_KEY || key === ANTHROPIC_BASE_URL_ENV_KEY || @@ -491,6 +510,39 @@ function getRetryAttemptsWarnings(value: string): string[] { } } +function getMaxOutputTokensWarnings(value: string): string[] { + try { + resolveMaxOutputTokens({ + [OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY]: value, + }); + + return []; + } catch { + return ["invalid output token limit"]; + } +} + +function getStreamIdleTimeoutWarnings( + value: string, + provider: OpenWikiProvider, +): string[] { + if (provider !== "bedrock") { + return []; + } + + try { + const streamIdleTimeout = resolveStreamIdleTimeout({ + [OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY]: value, + }); + + return streamIdleTimeout === 0 + ? ["stream watchdog disabled; stalled streams may hang indefinitely"] + : []; + } catch { + return ["invalid stream idle timeout"]; + } +} + async function readOpenWikiEnv(): Promise { try { return parseEnv(await readFile(openWikiEnvPath, "utf8")); diff --git a/test/bedrock-model.test.ts b/test/bedrock-model.test.ts index b694d64e..1a03a73a 100644 --- a/test/bedrock-model.test.ts +++ b/test/bedrock-model.test.ts @@ -51,17 +51,60 @@ describe("createModel Bedrock credentials", () => { process.env.AWS_WEB_IDENTITY_TOKEN_FILE = "/path/that/must/not/be/read"; process.env.AWS_REGION = "us-east-1"; - createModel("bedrock", "anthropic.claude-sonnet-5", 4); + createModel("bedrock", "anthropic.claude-sonnet-5", 4, 8192); expect(bedrockConstructorArgs).toHaveLength(1); expect(bedrockConstructorArgs[0]).toMatchObject({ maxRetries: 4, + maxTokens: 8192, model: "anthropic.claude-sonnet-5", region: "us-east-1", }); expect(bedrockConstructorArgs[0]).not.toHaveProperty("credentials"); }); + test("preserves the provider default when no output token limit is set", () => { + process.env.AWS_REGION = "us-east-1"; + + createModel("bedrock", "anthropic.claude-sonnet-5", 4); + + expect(bedrockConstructorArgs[0]).not.toHaveProperty("maxTokens"); + expect(bedrockConstructorArgs[0]).not.toHaveProperty("streamIdleTimeout"); + }); + + test.each([0, 300000])( + "passes streamIdleTimeout: %i to ChatBedrockConverse", + (streamIdleTimeout) => { + process.env.AWS_REGION = "us-east-1"; + + createModel( + "bedrock", + "anthropic.claude-sonnet-5", + 4, + undefined, + streamIdleTimeout, + ); + + expect(bedrockConstructorArgs[0]).toMatchObject({ + streamIdleTimeout, + }); + }, + ); + + test("omits streamIdleTimeout when the override is undefined", () => { + process.env.AWS_REGION = "us-east-1"; + + createModel( + "bedrock", + "anthropic.claude-sonnet-5", + 4, + undefined, + undefined, + ); + + expect(bedrockConstructorArgs[0]).not.toHaveProperty("streamIdleTimeout"); + }); + test("lets LangChain preserve complete legacy credentials and session tokens", () => { process.env.BEDROCK_AWS_ACCESS_KEY_ID = "legacy-access"; process.env.BEDROCK_AWS_SECRET_ACCESS_KEY = "legacy-secret"; diff --git a/test/constants.test.ts b/test/constants.test.ts index 65bcea17..8d5010a6 100644 --- a/test/constants.test.ts +++ b/test/constants.test.ts @@ -29,11 +29,14 @@ import { providerRequiresSecretKey, providerUsesAwsSdkCredentials, resolveConfiguredProvider, + resolveMaxOutputTokens, resolveOpenRouterProviderOnly, resolveProviderBaseUrl, resolveProviderLocation, resolveProviderRegion, resolveProviderRetryAttempts, + resolveStreamIdleTimeout, + resolveStreamIdleTimeoutForProvider, } from "../src/constants.ts"; describe("isValidModelId", () => { @@ -277,6 +280,110 @@ describe("resolveProviderRetryAttempts", () => { }); }); +describe("resolveMaxOutputTokens", () => { + test("uses no explicit limit when no override is set", () => { + expect(resolveMaxOutputTokens({})).toBeUndefined(); + }); + + test("accepts positive integer output token limits", () => { + expect( + resolveMaxOutputTokens({ + OPENWIKI_MAX_OUTPUT_TOKENS: "1", + }), + ).toBe(1); + expect( + resolveMaxOutputTokens({ + OPENWIKI_MAX_OUTPUT_TOKENS: " 8192 ", + }), + ).toBe(8192); + }); + + test("rejects invalid output token limits", () => { + for (const value of [ + "", + " ", + "0", + "-1", + "1.5", + "abc", + "1e2", + "9007199254740992", + ]) { + expect(() => + resolveMaxOutputTokens({ + OPENWIKI_MAX_OUTPUT_TOKENS: value, + }), + ).toThrow(/OPENWIKI_MAX_OUTPUT_TOKENS/u); + } + }); +}); + +describe("resolveStreamIdleTimeout", () => { + test("uses the provider default when no override is set", () => { + expect(resolveStreamIdleTimeout({})).toBeUndefined(); + }); + + test("accepts zero to disable the watchdog and positive millisecond values", () => { + expect( + resolveStreamIdleTimeout({ + OPENWIKI_STREAM_IDLE_TIMEOUT: " 0 ", + }), + ).toBe(0); + expect( + resolveStreamIdleTimeout({ + OPENWIKI_STREAM_IDLE_TIMEOUT: "300000", + }), + ).toBe(300000); + expect( + resolveStreamIdleTimeout({ + OPENWIKI_STREAM_IDLE_TIMEOUT: "2147483647", + }), + ).toBe(2147483647); + }); + + test("rejects invalid stream idle timeouts", () => { + for (const value of [ + "", + " ", + "-1", + "1.5", + "abc", + "1e2", + "2147483648", + "9007199254740992", + ]) { + expect(() => + resolveStreamIdleTimeout({ + OPENWIKI_STREAM_IDLE_TIMEOUT: value, + }), + ).toThrow(/OPENWIKI_STREAM_IDLE_TIMEOUT/u); + } + }); +}); + +describe("resolveStreamIdleTimeoutForProvider", () => { + test("ignores a stale Bedrock timeout for other providers", () => { + expect( + resolveStreamIdleTimeoutForProvider("openai", { + OPENWIKI_STREAM_IDLE_TIMEOUT: "invalid", + }), + ).toBeUndefined(); + }); + + test("validates the timeout when Bedrock is active", () => { + expect( + resolveStreamIdleTimeoutForProvider("bedrock", { + OPENWIKI_STREAM_IDLE_TIMEOUT: "300000", + }), + ).toBe(300000); + expect(() => + resolveStreamIdleTimeoutForProvider("bedrock", { + OPENWIKI_STREAM_IDLE_TIMEOUT: "invalid", + }), + ).toThrow(/OPENWIKI_STREAM_IDLE_TIMEOUT/u); + }); +}); + describe("resolveOpenRouterProviderOnly", () => { test("returns undefined when no provider pin is configured", () => { expect(resolveOpenRouterProviderOnly({})).toBeUndefined(); diff --git a/test/env-behavior.test.ts b/test/env-behavior.test.ts index c903fde5..0f1ff80f 100644 --- a/test/env-behavior.test.ts +++ b/test/env-behavior.test.ts @@ -18,8 +18,10 @@ import { OPENAI_COMPATIBLE_BASE_URL_ENV_KEY, OPENAI_API_KEY_ENV_KEY, OPENROUTER_API_KEY_ENV_KEY, + OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY, OPENWIKI_MODEL_ID_ENV_KEY, OPENWIKI_PROVIDER_ENV_KEY, + OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY, } from "../src/constants.ts"; // `loadOpenWikiEnv`, `saveOpenWikiEnv`, and `getCredentialDiagnostics` all read @@ -53,8 +55,10 @@ const KEYS_UNDER_TEST = [ OPENAI_COMPATIBLE_BASE_URL_ENV_KEY, OPENAI_API_KEY_ENV_KEY, OPENROUTER_API_KEY_ENV_KEY, + OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY, OPENWIKI_MODEL_ID_ENV_KEY, OPENWIKI_PROVIDER_ENV_KEY, + OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY, // Deprecated / recently un-deprecated OpenAI keys. Cleared in each hook so the // developer's ambient shell (which may export OPENAI_BASE_URL) cannot leak // into these tests, and a loaded value cannot leak back out to other tests. @@ -425,6 +429,67 @@ describe("getCredentialDiagnostics", () => { expect(entry?.warnings).toContain("invalid provider"); }); + test("surfaces and validates the output token limit as non-secret", async () => { + await env.saveOpenWikiEnv({ + [OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY]: "0", + }); + + const diagnostics = await env.getCredentialDiagnostics(); + const entry = diagnostics.find( + (item) => item.key === OPENWIKI_MAX_OUTPUT_TOKENS_ENV_KEY, + ); + + expect(entry?.preview).toBe('"0"'); + expect(entry?.warnings).toContain("invalid output token limit"); + }); + + test("surfaces and validates the stream idle timeout as non-secret", async () => { + await env.saveOpenWikiEnv({ + [OPENWIKI_PROVIDER_ENV_KEY]: "bedrock", + [OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY]: "-1", + }); + + const diagnostics = await env.getCredentialDiagnostics(); + const entry = diagnostics.find( + (item) => item.key === OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY, + ); + + expect(entry?.preview).toBe('"-1"'); + expect(entry?.warnings).toContain("invalid stream idle timeout"); + }); + + test("accepts zero as a stream idle timeout that disables the watchdog", async () => { + await env.saveOpenWikiEnv({ + [OPENWIKI_PROVIDER_ENV_KEY]: "bedrock", + [OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY]: "0", + }); + + const diagnostics = await env.getCredentialDiagnostics(); + const entry = diagnostics.find( + (item) => item.key === OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY, + ); + + expect(entry?.preview).toBe('"0"'); + expect(entry?.warnings).toContain( + "stream watchdog disabled; stalled streams may hang indefinitely", + ); + }); + + test("does not warn about an inactive Bedrock stream timeout", async () => { + await env.saveOpenWikiEnv({ + [OPENWIKI_PROVIDER_ENV_KEY]: "openai", + [OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY]: "0", + }); + + const diagnostics = await env.getCredentialDiagnostics(); + const entry = diagnostics.find( + (item) => item.key === OPENWIKI_STREAM_IDLE_TIMEOUT_ENV_KEY, + ); + + expect(entry?.preview).toBe('"0"'); + expect(entry?.warnings).toEqual([]); + }); + test("flags an OpenAI-compatible chat completions endpoint as a base URL warning", async () => { await env.saveOpenWikiEnv({ [OPENAI_COMPATIBLE_BASE_URL_ENV_KEY]: diff --git a/test/env.test.ts b/test/env.test.ts index f56916b9..2fb0ae5e 100644 --- a/test/env.test.ts +++ b/test/env.test.ts @@ -125,6 +125,14 @@ describe("formatEnv", () => { }); describe("MANAGED_ENV_KEYS", () => { + test("manages the model output token limit", () => { + expect(MANAGED_ENV_KEYS).toContain("OPENWIKI_MAX_OUTPUT_TOKENS"); + }); + + test("manages the Bedrock stream idle timeout", () => { + expect(MANAGED_ENV_KEYS).toContain("OPENWIKI_STREAM_IDLE_TIMEOUT"); + }); + test("manages the Google Cloud settings for the gemini-enterprise provider", () => { expect(MANAGED_ENV_KEYS).toContain("GOOGLE_CLOUD_PROJECT"); expect(MANAGED_ENV_KEYS).toContain("GOOGLE_CLOUD_LOCATION"); diff --git a/test/gemini-retry.test.ts b/test/gemini-retry.test.ts index 5a29f11e..800491d0 100644 --- a/test/gemini-retry.test.ts +++ b/test/gemini-retry.test.ts @@ -1,12 +1,13 @@ import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; -// Capture the options each Chat* constructor receives so we can assert that the -// provider retry count (OPENWIKI_PROVIDER_RETRY_ATTEMPTS -> maxRetries) is wired -// through to every Gemini surface. maxRetries is not a readable property on the -// constructed instances, so we mock the constructors and inspect their args. +// Capture the options each Chat* constructor receives so we can assert that +// provider runtime options are wired through to every Gemini surface. These +// fields are not consistently readable on constructed instances, so we mock +// the constructors and inspect their args. const chatGoogleArgs: Array> = []; const chatAnthropicArgs: Array<[string, Record]> = []; const chatOpenAIArgs: Array> = []; +const chatOpenRouterArgs: Array> = []; vi.mock("@langchain/google/node", () => ({ ChatGoogle: class { @@ -32,13 +33,26 @@ vi.mock("@langchain/openai", () => ({ }, })); +vi.mock("@langchain/openrouter", () => ({ + ChatOpenRouter: class { + constructor(options: Record) { + chatOpenRouterArgs.push(options); + } + }, +})); + // Imported after vi.mock so the mocked constructors are in effect. const { createModel } = await import("../src/agent/index.ts"); const PROJECT_KEY = "GOOGLE_CLOUD_PROJECT"; const GEMINI_KEY = "GEMINI_API_KEY"; +const CHATGPT_ENV_KEYS = [ + "OPENAI_CHATGPT_ACCESS_TOKEN", + "OPENAI_CHATGPT_REFRESH_TOKEN", + "OPENAI_CHATGPT_ACCOUNT_ID", +] as const; -describe("createModel wires retry attempts into the Gemini providers", () => { +describe("createModel wires runtime options into provider clients", () => { let savedProject: string | undefined; let savedGeminiKey: string | undefined; @@ -50,6 +64,7 @@ describe("createModel wires retry attempts into the Gemini providers", () => { chatGoogleArgs.length = 0; chatAnthropicArgs.length = 0; chatOpenAIArgs.length = 0; + chatOpenRouterArgs.length = 0; }); afterEach(() => { @@ -57,51 +72,107 @@ describe("createModel wires retry attempts into the Gemini providers", () => { restoreEnv(GEMINI_KEY, savedGeminiKey); }); - test("gemini (AI Studio) passes maxRetries and the API key to ChatGoogle", () => { - createModel("gemini", "gemini-3.1-flash-lite", 5); + test("gemini (AI Studio) passes runtime options and the API key to ChatGoogle", () => { + createModel("gemini", "gemini-3.1-flash-lite", 5, 8192); expect(chatGoogleArgs).toHaveLength(1); expect(chatGoogleArgs[0]?.maxRetries).toBe(5); + expect(chatGoogleArgs[0]?.maxOutputTokens).toBe(8192); // The API key is not a readable property on the built instance, so assert it // reaches the constructor here rather than in create-model.test.ts. expect(chatGoogleArgs[0]?.apiKey).toBe("test-gemini-key"); }); + test("does not pass the Bedrock-only stream idle timeout to ChatGoogle", () => { + createModel("gemini", "gemini-3.1-flash-lite", 5, undefined, 300000); + + expect(chatGoogleArgs).toHaveLength(1); + expect(chatGoogleArgs[0]).not.toHaveProperty("streamIdleTimeout"); + }); + test("propagates maxRetries: 0 (not coerced to a default)", () => { createModel("gemini", "gemini-3.1-pro", 0); expect(chatGoogleArgs).toHaveLength(1); expect(chatGoogleArgs[0]?.maxRetries).toBe(0); + expect(chatGoogleArgs[0]).not.toHaveProperty("maxOutputTokens"); }); - test("gemini-enterprise Gemini surface passes maxRetries to ChatGoogle", () => { - createModel("gemini-enterprise", "gemini-3.1-pro", 4); + test("gemini-enterprise Gemini surface passes runtime options to ChatGoogle", () => { + createModel("gemini-enterprise", "gemini-3.1-pro", 4, 8192); expect(chatGoogleArgs).toHaveLength(1); expect(chatGoogleArgs[0]?.maxRetries).toBe(4); + expect(chatGoogleArgs[0]?.maxOutputTokens).toBe(8192); }); - test("gemini-enterprise Claude surface passes maxRetries to ChatAnthropic", () => { + test("gemini-enterprise Claude surface passes runtime options to ChatAnthropic", () => { createModel( "gemini-enterprise", "publishers/anthropic/models/claude-sonnet-4-5@20250929", 3, + 8192, ); expect(chatAnthropicArgs).toHaveLength(1); expect(chatAnthropicArgs[0]?.[1].maxRetries).toBe(3); + expect(chatAnthropicArgs[0]?.[1].maxTokens).toBe(8192); }); - test("gemini-enterprise MaaS surface passes maxRetries to ChatOpenAI", () => { + test("gemini-enterprise MaaS surface passes runtime options to ChatOpenAI", () => { createModel( "gemini-enterprise", "publishers/meta/models/llama-3.3-70b-instruct-maas", 2, + 8192, ); expect(chatOpenAIArgs).toHaveLength(1); expect(chatOpenAIArgs[0]?.maxRetries).toBe(2); + expect(chatOpenAIArgs[0]?.maxTokens).toBe(8192); }); + + test("passes maxTokens to direct non-Google clients", () => { + createModel("anthropic", "claude-sonnet-4-5", 3, 8192); + createModel("openai", "gpt-5.6-terra", 3, 8192); + createModel("openai-compatible", "custom-model", 3, 8192); + createModel("openrouter", "anthropic/claude-sonnet-4.5", 3, 8192); + + expect(chatAnthropicArgs[0]?.[1].maxTokens).toBe(8192); + expect(chatOpenAIArgs.map((options) => options.maxTokens)).toEqual([ + 8192, 8192, + ]); + expect(chatOpenRouterArgs[0]?.maxTokens).toBe(8192); + }); + + test("passes maxTokens to the ChatGPT Responses client", () => { + for (const key of CHATGPT_ENV_KEYS) { + vi.stubEnv(key, `test-${key.toLowerCase()}`); + } + + try { + createModel("openai-chatgpt", "gpt-5.6-terra", 3, 8192); + expect(chatOpenAIArgs[0]?.maxTokens).toBe(8192); + } finally { + vi.unstubAllEnvs(); + } + }); + + test.each([ + ["gpt-5.5", true], + ["claude-sonnet-5", false], + ])( + "passes maxTokens to Copilot model %s while preserving its transport", + (modelId, useResponsesApi) => { + createModel("copilot", modelId, 3, 8192); + + expect(chatOpenAIArgs).toHaveLength(1); + expect(chatOpenAIArgs[0]).toMatchObject({ + maxTokens: 8192, + useResponsesApi, + }); + }, + ); }); function restoreEnv(key: string, value: string | undefined): void { diff --git a/test/redaction.test.ts b/test/redaction.test.ts index e7eb8fdf..4a6a3f9f 100644 --- a/test/redaction.test.ts +++ b/test/redaction.test.ts @@ -253,6 +253,12 @@ describe("sanitizeDiagnosticText", () => { }); describe("formatEnvironmentDebugValue", () => { + test("previews the non-secret stream idle timeout", () => { + expect( + formatEnvironmentDebugValue("OPENWIKI_STREAM_IDLE_TIMEOUT", "300000"), + ).toBe('set(value="300000")'); + }); + test.each([ ["BEDROCK_AWS_ACCESS_KEY_ID", "bedrock-access-id-123"], ["BEDROCK_AWS_SECRET_ACCESS_KEY", "bedrock-secret-abcdef123456"],