diff --git a/docs/api/README.md b/docs/api/README.md index a828dc20a..dcfa1d616 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -2746,6 +2746,7 @@ console.log(result.content); - [DataStreamWriter](type-aliases/DataStreamWriter.md) - [WebSocketMessageType](type-aliases/WebSocketMessageType.md) - [WebSocketMessage](type-aliases/WebSocketMessage.md) +- [WebSocketRequestId](type-aliases/WebSocketRequestId.md) - [AuthenticatedUser](type-aliases/AuthenticatedUser.md) - [WebSocketConnection](type-aliases/WebSocketConnection.md) - [AuthStrategy](type-aliases/AuthStrategy.md) diff --git a/docs/api/classes/WebSocketMessageRouter.md b/docs/api/classes/WebSocketMessageRouter.md index b7b5e42c4..c3c2a92b6 100644 --- a/docs/api/classes/WebSocketMessageRouter.md +++ b/docs/api/classes/WebSocketMessageRouter.md @@ -34,7 +34,7 @@ Register a message route ##### handler -(`connection`, `payload`) => `Promise`\<`unknown`\> +(`connection`, `payload`, `requestId?`) => `Promise`\<`unknown`\> #### Returns diff --git a/docs/api/functions/createAgentWebSocketHandler.md b/docs/api/functions/createAgentWebSocketHandler.md index 785eec542..7a7cc661c 100644 --- a/docs/api/functions/createAgentWebSocketHandler.md +++ b/docs/api/functions/createAgentWebSocketHandler.md @@ -6,16 +6,32 @@ # Function: createAgentWebSocketHandler() -> **createAgentWebSocketHandler**(`_neurolink`): [`WebSocketHandler`](../type-aliases/WebSocketHandler.md) +> **createAgentWebSocketHandler**(`neurolink`): [`WebSocketHandler`](../type-aliases/WebSocketHandler.md) -Create a WebSocket handler for AI agent interactions +Server Adapters for exposing NeuroLink as HTTP APIs + +Supports multiple frameworks: Hono, Express, Fastify, Koa ## Parameters -### \_neurolink +### neurolink -`unknown` +[`NeuroLink`](../classes/NeuroLink.md) ## Returns [`WebSocketHandler`](../type-aliases/WebSocketHandler.md) + +## Example + +```typescript +import { NeuroLink } from "@juspay/neurolink"; +import { createServer } from "@juspay/neurolink/server"; + +const neurolink = new NeuroLink({ provider: "openai" }); +const server = await createServer(neurolink, { + framework: "hono", + config: { port: 3000 }, +}); +await server.start(); +``` diff --git a/docs/api/type-aliases/StreamLoopArgs.md b/docs/api/type-aliases/StreamLoopArgs.md index 45c54f9fe..9a8a2cf64 100644 --- a/docs/api/type-aliases/StreamLoopArgs.md +++ b/docs/api/type-aliases/StreamLoopArgs.md @@ -163,3 +163,13 @@ Signals the channel that no further chunks will arrive (success or error path al #### Returns `void` + +--- + +### responseFormat? + +> `optional` **responseFormat?**: [`OpenAICompatResponseFormat`](OpenAICompatResponseFormat.md) + +`response_format` for this turn's requests, computed once up front from +`options.schema` (see `suppressResponseFormatWithTools`). Absent when no +schema was requested, or when tools suppress it. diff --git a/docs/api/type-aliases/StreamResult.md b/docs/api/type-aliases/StreamResult.md index 9730d6853..0d3ccb595 100644 --- a/docs/api/type-aliases/StreamResult.md +++ b/docs/api/type-aliases/StreamResult.md @@ -177,6 +177,45 @@ Verbatim provider finish/stop reason for the turn's terminal model call. > `optional` **stepsUsed?**: `number` +#### structuredData? + +> `optional` **structuredData?**: `unknown` + +Parsed structured output for a `stream({ schema })` turn, available +AFTER the stream has been drained. + +Lives on metadata for the same reason `finishReason` above does: it is a +mutable reference the loop fills in, so result-object spreads in stream +wrappers cannot snapshot it before it resolves. A plain field rather than +a promise, deliberately — when a middleware short-circuits the request +and the loop never runs, an unresolved promise would hang every reader, +whereas an absent field is simply absent. Readers must tolerate that: +absence means the model never produced the object. + +#### structuredDataUsage? + +> `optional` **structuredDataUsage?**: `object` + +Tokens spent by the tool-free re-ask that produced `structuredData`, +when one was needed. + +That re-ask is a second, separately-billed model call, and the stream's +own `usage` has already resolved by the time it runs — so folding these +tokens into it would mutate a value a caller may have read, and leaving +them out entirely would under-report what the turn cost. Reported here +instead: same delivery as `structuredData`, filled at the same moment, +read at the same moment. Absent when no re-ask was needed, which is the +common case. The generate path accounts for its equivalent re-ask +inline, since there the usage has not been handed out yet. + +##### structuredDataUsage.inputTokens + +> **inputTokens**: `number` + +##### structuredDataUsage.outputTokens + +> **outputTokens**: `number` + #### thoughtSignature? > `optional` **thoughtSignature?**: `string` diff --git a/docs/api/type-aliases/TextGenerationOptions.md b/docs/api/type-aliases/TextGenerationOptions.md index 4d4e2c6c2..a9d5efabc 100644 --- a/docs/api/type-aliases/TextGenerationOptions.md +++ b/docs/api/type-aliases/TextGenerationOptions.md @@ -656,6 +656,40 @@ module's own default, unchanged unless a caller opts in here. --- +### videoOptions? + +> `optional` **videoOptions?**: `object` + +Video processing options (mirrors `GenerateOptions.videoOptions`). Never +declared here before, so `buildGenerateTextOptions` had nowhere to +forward the caller's setting even once the allowlist itself named it. + +#### frames? + +> `optional` **frames?**: `number` + +Frames to extract. Unset lets VideoProcessor pick from the clip's duration; clamped to 100. + +#### quality? + +> `optional` **quality?**: `number` + +Frame encoder quality, clamped to 1-100. Default 80. + +#### format? + +> `optional` **format?**: `"jpeg"` \| `"png"` + +Frame encoding. Default jpeg. + +#### transcribeAudio? + +> `optional` **transcribeAudio?**: `boolean` + +Not implemented yet (#433) — warns rather than silently doing nothing. + +--- + ### enableSummarization? > `optional` **enableSummarization?**: `boolean` diff --git a/docs/api/type-aliases/WebSocketRequestId.md b/docs/api/type-aliases/WebSocketRequestId.md new file mode 100644 index 000000000..d8b970952 --- /dev/null +++ b/docs/api/type-aliases/WebSocketRequestId.md @@ -0,0 +1,15 @@ +[**NeuroLink API Reference**](../README.md) + +--- + +[NeuroLink API Reference](../README.md) / WebSocketRequestId + +# Type Alias: WebSocketRequestId + +> **WebSocketRequestId** = `string` \| `number` + +Client-supplied correlation id on an agent-protocol WebSocket message +(`{ type, payload, id }`). Echoed back on the matching response frame so a +caller with more than one in-flight message on a connection can tell which +frame answers which request; optional and ignored entirely by a client +that never sends one. diff --git a/docs/guides/server-adapters/websocket.md b/docs/guides/server-adapters/websocket.md index 662950f15..b224f378b 100644 --- a/docs/guides/server-adapters/websocket.md +++ b/docs/guides/server-adapters/websocket.md @@ -634,7 +634,7 @@ ws.onmessage = (event) => { console.log("Connected:", message.connectionId); break; case "response": - console.log("Response:", message.data); + console.log("Response:", message.content); break; case "stream_start": console.log("Stream starting..."); diff --git a/package.json b/package.json index 39657861a..9581d9689 100644 --- a/package.json +++ b/package.json @@ -111,6 +111,7 @@ "test:middleware": "pnpm exec tsx test/continuous-test-suite-middleware.ts", "test:stream-middleware": "pnpm exec tsx test/continuous-test-suite-stream-middleware.ts", "test:stream-tool-telemetry": "pnpm exec tsx test/continuous-test-suite-stream-tool-telemetry.ts", + "test:stream-teardown-abort": "pnpm exec tsx test/continuous-test-suite-stream-teardown-abort.ts", "test:observability": "pnpm exec tsx test/continuous-test-suite-observability.ts", "test:ppt": "pnpm exec tsx test/continuous-test-suite-ppt.ts", "test:vertex-loop-characterization": "tsx test/continuous-test-suite-vertex-loop-characterization.ts", @@ -127,6 +128,7 @@ "test:mcp:infra": "pnpm exec tsx test/continuous-test-suite-mcp-infra.ts", "test:vendor-recovery": "pnpm exec tsx test/continuous-test-suite-native-vendor-recovery.ts", "test:anthropic-generate-guard": "pnpm exec tsx test/continuous-test-suite-anthropic-generate-guard.ts", + "test:stream-structured-output": "pnpm exec tsx test/continuous-test-suite-stream-structured-output.ts", "test:browser-bundle": "pnpm exec tsx test/continuous-test-suite-browser-bundle.ts", "test:providers-mocked": "pnpm exec tsx test/continuous-test-suite-providers-mocked.ts", "scaffold:provider": "pnpm exec tsx tools/scaffold-provider.ts", @@ -144,8 +146,10 @@ "test:decide": "pnpm exec tsx test/continuous-test-suite-decide.ts", "test:agent-mode": "pnpm exec tsx test/continuous-test-suite-agent-mode.ts", "test:servers": "pnpm exec tsx test/continuous-test-suite-servers.ts", + "test:websocket-agent-handler": "pnpm exec tsx test/continuous-test-suite-websocket-agent-handler.ts", "test:tool-reliability": "pnpm exec tsx test/continuous-test-suite-tool-reliability.ts", "test:video-abort": "pnpm exec tsx test/continuous-test-suite-video-abort.ts", + "test:video-options-allowlist": "pnpm exec tsx test/continuous-test-suite-video-options-allowlist.ts", "test:tts": "pnpm exec tsx test/continuous-test-suite-tts.ts", "test:tts:unit": "pnpm exec tsx test/continuous-test-suite-tts-unit.ts", "test:stt:unit": "pnpm exec tsx test/continuous-test-suite-stt-unit.ts", diff --git a/src/cli/loop/optionsSchema.ts b/src/cli/loop/optionsSchema.ts index 5453cd13b..5e2af5702 100644 --- a/src/cli/loop/optionsSchema.ts +++ b/src/cli/loop/optionsSchema.ts @@ -28,6 +28,7 @@ export const textGenerationOptionsSchema: Record< | "csvOptions" | "pdfOptions" | "imageOptions" + | "videoOptions" // Complex object, not set via simple CLI commands | "tts" | "stt" // Complex object, set via --stt* flags | "thinkingConfig" // Complex object, use thinking/thinkingBudget instead diff --git a/src/lib/core/baseProvider.ts b/src/lib/core/baseProvider.ts index 38759a331..d2682195d 100644 --- a/src/lib/core/baseProvider.ts +++ b/src/lib/core/baseProvider.ts @@ -87,6 +87,7 @@ import { } from "../utils/async/withTimeout.js"; import { composeAbortSignals, + composeAbortSignalsScoped, createTimeoutController, TimeoutError, } from "../utils/timeout.js"; @@ -375,203 +376,267 @@ export abstract class BaseProvider implements AIProvider { let options = this.normalizeStreamOptions(optionsOrPrompt); this.applyFileToolRoots(options); - // Before anything else, and before a single byte leaves the process: an - // execution policy this provider cannot honour is an error, and a policy - // whose shape could be read two ways is an error. Both are silent bugs at - // the point they would otherwise matter. - validateExecutionControl( - options.executionControl, - this.providerName, - this.supportsExecutionControl(), - { - turnTimeoutMs: options.turnTimeoutMs, - toolTimeoutMs: options.toolTimeoutMs, - }, - ); - - logger.info(`Starting stream`, { - provider: this.providerName, - hasTools: !options.disableTools && this.supportsTools(), - disableTools: !!options.disableTools, - supportsTools: this.supportsTools(), - inputLength: options.input?.text?.length || 0, - maxTokens: options.maxTokens, - temperature: options.temperature, - timestamp: Date.now(), - }); - - // ===== EARLY MULTIMODAL DETECTION ===== - // #1259: audioFiles was missing here while videoFiles was present, so an - // audio-only stream skipped this branch entirely. - const hasFileInput = - !!options.input?.files?.length || - !!options.input?.videoFiles?.length || - !!options.input?.audioFiles?.length; - if (hasFileInput) { - // ===== VIDEO ANALYSIS DETECTION ===== - // Check if video frames are present and handle with fake streaming - const messages = await this.buildMessagesForStream(options); - if (hasVideoFrames(messages)) { - logger.info( - `Video frames detected in stream, using fake streaming for video analysis`, - { - provider: this.providerName, - model: this.modelName, - }, - ); - // Note: executeFakeStreaming() owns its own catch that fires the - // consumer-supplied onError before re-throwing through - // handleProviderError(), so we do not need to wrap again here — - // doing so would route the error through handleProviderError() - // twice (and risk a double-fire onError without the shared - // lifecycle-fired WeakSet mark). - const fakeResult = await this.executeFakeStreaming( - options, - analysisSchema, - ); - return this.wrapStreamWithLifecycleCallbacks(fakeResult, options); - } - } - - // CRITICAL: Image generation models don't support real streaming - // Force fake streaming for image models to ensure image output is yielded. - // resolveRequestKind() skips this path when the caller explicitly requests - // non-image output (e.g. JSON analysis) so dual-mode models like - // gemini-3.1-flash-image-preview can still perform text/structured - // generation — see its doc comment for the full precedence table. - const requestKind = resolveRequestKind(options, this.modelName); - - if (requestKind === "image") { - logger.info(`Image model detected, forcing fake streaming`, { - provider: this.providerName, - model: this.modelName, - reason: - "Image generation requires fake streaming to yield image output", - }); - - // Skip real streaming, go directly to fake streaming. - // executeFakeStreaming() owns its own catch + lifecycle fire, so - // wrapping again here would double-route through handleProviderError(). - const fakeResult = await this.executeFakeStreaming( - options, - analysisSchema, + // Own an abort signal for this call so teardown can reach the transport. + // + // #1550 gave teardown a way to close the iterator chain, but closing an + // iterator does not cancel the HTTP request underneath it — the provider + // read stays pending and the connection stays open. Providers already + // honour `options.abortSignal` and pass it to their transport; nothing + // ever fired one on abandonment, because the only signal available was + // the caller's and the caller had simply walked away. + // + // Composed rather than substituted, so a caller's own signal keeps + // working exactly as before and either source can end the request. + // Scoped, not plain `composeAbortSignals`: that one is `AbortSignal.any`, + // whose registration on a source survives until the DERIVED signal is + // collected. A caller that reuses one long-lived `abortSignal` across many + // stream calls would accumulate a dependent per call on it, released only + // by GC — the exact hazard composeAbortSignalsScoped was written for (see + // its docstring: MaxListenersExceededWarning at 10+ compositions). It + // registers removable listeners and hands back a `dispose()`, which + // `teardown()` calls when the stream settles. + const teardownController = new AbortController(); + const { signal: composedStreamSignal, dispose: disposeComposedSignal } = + composeAbortSignalsScoped(options.abortSignal, teardownController.signal); + if (composedStreamSignal !== options.abortSignal) { + options = { ...options, abortSignal: composedStreamSignal }; + } + + // Ownership of the composed signal transfers to the returned wrapper, + // whose teardown() disposes it. Until that wrapper exists, this method + // owns it — and every throw between here and there (an execution-control + // rejection, a provider error with no tool fallback) would otherwise + // leave composeAbortSignalsScoped's once-listeners attached to a caller + // signal that may outlive this call by many more. + // + // catch, deliberately, not finally: dispose() only detaches listeners, so + // running it on the success path would sever the caller's abort from the + // composed signal the stream is still using. + try { + // Before anything else, and before a single byte leaves the process: an + // execution policy this provider cannot honour is an error, and a policy + // whose shape could be read two ways is an error. Both are silent bugs at + // the point they would otherwise matter. + validateExecutionControl( + options.executionControl, + this.providerName, + this.supportsExecutionControl(), + { + turnTimeoutMs: options.turnTimeoutMs, + toolTimeoutMs: options.toolTimeoutMs, + }, ); - return this.wrapStreamWithLifecycleCallbacks(fakeResult, options); - } - // Central tool merge: Pre-merge base tools (MCP/built-in) with user-provided - // tools (e.g. RAG tools) into options.tools. This way, every provider's - // executeStream() can simply use options.tools (or getAllTools() + options.tools) - // and get the complete tool set without needing per-provider merge logic. - if (this.shouldUseTools(options, true)) { - const mergedTools = await this.getToolsForStream(options); - options = { ...options, tools: mergedTools }; - } else { - options = { ...options, tools: {} }; - } - - // CRITICAL FIX: Always prefer real streaming over fake streaming - // Try real streaming first, use fake streaming only as fallback - try { - logger.debug(`Attempting real streaming`, { + logger.info(`Starting stream`, { provider: this.providerName, + hasTools: !options.disableTools && this.supportsTools(), + disableTools: !!options.disableTools, + supportsTools: this.supportsTools(), + inputLength: options.input?.text?.length || 0, + maxTokens: options.maxTokens, + temperature: options.temperature, timestamp: Date.now(), }); - const realStreamResult = await this.executeStream( - options, - analysisSchema, - ); + // ===== EARLY MULTIMODAL DETECTION ===== + // #1259: audioFiles was missing here while videoFiles was present, so an + // audio-only stream skipped this branch entirely. + const hasFileInput = + !!options.input?.files?.length || + !!options.input?.videoFiles?.length || + !!options.input?.audioFiles?.length; + if (hasFileInput) { + // ===== VIDEO ANALYSIS DETECTION ===== + // Check if video frames are present and handle with fake streaming + const messages = await this.buildMessagesForStream(options); + if (hasVideoFrames(messages)) { + logger.info( + `Video frames detected in stream, using fake streaming for video analysis`, + { + provider: this.providerName, + model: this.modelName, + }, + ); + // Note: executeFakeStreaming() owns its own catch that fires the + // consumer-supplied onError before re-throwing through + // handleProviderError(), so we do not need to wrap again here — + // doing so would route the error through handleProviderError() + // twice (and risk a double-fire onError without the shared + // lifecycle-fired WeakSet mark). + const fakeResult = await this.executeFakeStreaming( + options, + analysisSchema, + ); + return this.wrapStreamWithLifecycleCallbacks( + fakeResult, + options, + teardownController, + disposeComposedSignal, + ); + } + } - logger.info(`Real streaming succeeded`, { - provider: this.providerName, - timestamp: Date.now(), - }); + // CRITICAL: Image generation models don't support real streaming + // Force fake streaming for image models to ensure image output is yielded. + // resolveRequestKind() skips this path when the caller explicitly requests + // non-image output (e.g. JSON analysis) so dual-mode models like + // gemini-3.1-flash-image-preview can still perform text/structured + // generation — see its doc comment for the full precedence table. + const requestKind = resolveRequestKind(options, this.modelName); - // Wire lifecycle callbacks (onChunk/onFinish/onError) on the user- - // facing StreamResult.stream. The AI-SDK lifecycle middleware only - // sees AI-SDK-internal chunks via streamText/wrapLanguageModel, so - // providers with custom HTTP streaming (Ollama, llama.cpp's /api, - // anything that doesn't go through streamText) bypass it. Wrapping - // here makes the callbacks fire for every provider, regardless of - // streaming implementation. - return this.wrapStreamWithLifecycleCallbacks( - this.withStreamModelFallback(realStreamResult, options, analysisSchema), - options, - ); - } catch (realStreamError) { - // Retired-model fallback runs FIRST, before any lifecycle callback has - // fired and before a single chunk has reached the consumer: onChunk and - // onFinish are only wired on the success path above, and onError fires - // further down this same catch. That ordering is what makes retrying a - // stream safe at all — nothing observable has happened yet. - const recovered = await this.retryStreamWithFallbackModel( - realStreamError, - options, - analysisSchema, - ); - if (recovered) { - return recovered; + if (requestKind === "image") { + logger.info(`Image model detected, forcing fake streaming`, { + provider: this.providerName, + model: this.modelName, + reason: + "Image generation requires fake streaming to yield image output", + }); + + // Skip real streaming, go directly to fake streaming. + // executeFakeStreaming() owns its own catch + lifecycle fire, so + // wrapping again here would double-route through handleProviderError(). + const fakeResult = await this.executeFakeStreaming( + options, + analysisSchema, + ); + return this.wrapStreamWithLifecycleCallbacks( + fakeResult, + options, + teardownController, + disposeComposedSignal, + ); } - // The fallback is BROAD, not narrow: only the terminal errors listed - // below (abort, timeout, 401/403, quota, rate limit, authentication) - // re-throw. Every other failure — including a genuine configuration or - // programming error — is masked as a degraded fake stream whenever - // tools are enabled. Narrowing this to "streaming with tools is - // unsupported" failures would change behaviour for every provider at - // once, so it needs its own characterization PR first; until then this - // comment records what the code does, not what a narrower design would. - const errMsg = - realStreamError instanceof Error - ? realStreamError.message - : String(realStreamError); - const errName = - realStreamError instanceof Error ? realStreamError.name : ""; - if ( - errName === "AbortError" || - errMsg.includes("abort") || - errMsg.includes("timeout") || - errMsg.includes("401") || - errMsg.includes("403") || - errMsg.includes("quota") || - errMsg.includes("rate limit") || - errMsg.includes("authentication") - ) { - await this.fireLifecycleErrorCallback(options, realStreamError); - throw this.handleProviderError(realStreamError); + // Central tool merge: Pre-merge base tools (MCP/built-in) with user-provided + // tools (e.g. RAG tools) into options.tools. This way, every provider's + // executeStream() can simply use options.tools (or getAllTools() + options.tools) + // and get the complete tool set without needing per-provider merge logic. + if (this.shouldUseTools(options, true)) { + const mergedTools = await this.getToolsForStream(options); + options = { ...options, tools: mergedTools }; + } else { + options = { ...options, tools: {} }; } - logger.warn( - `Real streaming failed for ${this.providerName}, falling back to fake streaming:`, - { - error: errMsg, + // CRITICAL FIX: Always prefer real streaming over fake streaming + // Try real streaming first, use fake streaming only as fallback + try { + logger.debug(`Attempting real streaming`, { + provider: this.providerName, timestamp: Date.now(), - }, - ); + }); - // Fallback to fake streaming only if real streaming fails AND tools - // are enabled. executeFakeStreaming() owns its own catch + lifecycle - // fire, so a fake-streaming failure here surfaces through that path - // without needing an outer wrap (which would double-route through - // handleProviderError()). - if (!options.disableTools && this.supportsTools()) { - const fakeResult = await this.executeFakeStreaming( + const realStreamResult = await this.executeStream( options, analysisSchema, ); - return this.wrapStreamWithLifecycleCallbacks(fakeResult, options); - } else { - // If real streaming failed and no tools are enabled, fire onError - // before re-throwing so consumer-supplied callbacks see the failure. - await this.fireLifecycleErrorCallback(options, realStreamError); - // If real streaming failed and no tools are enabled, re-throw the original error - logger.error( - `Real streaming failed for ${this.providerName}:`, + + logger.info(`Real streaming succeeded`, { + provider: this.providerName, + timestamp: Date.now(), + }); + + // Wire lifecycle callbacks (onChunk/onFinish/onError) on the user- + // facing StreamResult.stream. The AI-SDK lifecycle middleware only + // sees AI-SDK-internal chunks via streamText/wrapLanguageModel, so + // providers with custom HTTP streaming (Ollama, llama.cpp's /api, + // anything that doesn't go through streamText) bypass it. Wrapping + // here makes the callbacks fire for every provider, regardless of + // streaming implementation. + return this.wrapStreamWithLifecycleCallbacks( + this.withStreamModelFallback( + realStreamResult, + options, + analysisSchema, + ), + options, + teardownController, + disposeComposedSignal, + ); + } catch (realStreamError) { + // Retired-model fallback runs FIRST, before any lifecycle callback has + // fired and before a single chunk has reached the consumer: onChunk and + // onFinish are only wired on the success path above, and onError fires + // further down this same catch. That ordering is what makes retrying a + // stream safe at all — nothing observable has happened yet. + const recovered = await this.retryStreamWithFallbackModel( realStreamError, + options, + analysisSchema, + teardownController, + disposeComposedSignal, ); - throw this.handleProviderError(realStreamError); + if (recovered) { + return recovered; + } + + // The fallback is BROAD, not narrow: only the terminal errors listed + // below (abort, timeout, 401/403, quota, rate limit, authentication) + // re-throw. Every other failure — including a genuine configuration or + // programming error — is masked as a degraded fake stream whenever + // tools are enabled. Narrowing this to "streaming with tools is + // unsupported" failures would change behaviour for every provider at + // once, so it needs its own characterization PR first; until then this + // comment records what the code does, not what a narrower design would. + const errMsg = + realStreamError instanceof Error + ? realStreamError.message + : String(realStreamError); + const errName = + realStreamError instanceof Error ? realStreamError.name : ""; + if ( + errName === "AbortError" || + errMsg.includes("abort") || + errMsg.includes("timeout") || + errMsg.includes("401") || + errMsg.includes("403") || + errMsg.includes("quota") || + errMsg.includes("rate limit") || + errMsg.includes("authentication") + ) { + await this.fireLifecycleErrorCallback(options, realStreamError); + throw this.handleProviderError(realStreamError); + } + + logger.warn( + `Real streaming failed for ${this.providerName}, falling back to fake streaming:`, + { + error: errMsg, + timestamp: Date.now(), + }, + ); + + // Fallback to fake streaming only if real streaming fails AND tools + // are enabled. executeFakeStreaming() owns its own catch + lifecycle + // fire, so a fake-streaming failure here surfaces through that path + // without needing an outer wrap (which would double-route through + // handleProviderError()). + if (!options.disableTools && this.supportsTools()) { + const fakeResult = await this.executeFakeStreaming( + options, + analysisSchema, + ); + return this.wrapStreamWithLifecycleCallbacks( + fakeResult, + options, + teardownController, + disposeComposedSignal, + ); + } else { + // If real streaming failed and no tools are enabled, fire onError + // before re-throwing so consumer-supplied callbacks see the failure. + await this.fireLifecycleErrorCallback(options, realStreamError); + // If real streaming failed and no tools are enabled, re-throw the original error + logger.error( + `Real streaming failed for ${this.providerName}:`, + realStreamError, + ); + throw this.handleProviderError(realStreamError); + } } + } catch (streamSetupError) { + disposeComposedSignal(); + throw streamSetupError; } } @@ -763,6 +828,8 @@ export abstract class BaseProvider implements AIProvider { error: unknown, options: StreamOptions, analysisSchema: ValidationSchema | undefined, + teardownController?: AbortController, + disposeComposedSignal?: () => void, ): Promise { if (options.disableInternalFallback === true) { return undefined; @@ -784,7 +851,12 @@ export abstract class BaseProvider implements AIProvider { this.refreshHandlersForModel(candidate); try { const result = await this.executeStream(options, analysisSchema); - return this.wrapStreamWithLifecycleCallbacks(result, options); + return this.wrapStreamWithLifecycleCallbacks( + result, + options, + teardownController, + disposeComposedSignal, + ); } catch { // Any failure on a candidate — stale id or otherwise — just moves to // the next one. Nothing is reported from here: if none succeed the @@ -809,10 +881,20 @@ export abstract class BaseProvider implements AIProvider { * over /api/chat, custom OpenAI-compatible servers, etc). Wrapping the * user-facing stream here ensures the callbacks fire regardless of the * underlying transport. + * + * `teardownController`, when supplied, is the call's own abort source (see + * `stream()`): the cancel hook fires it on abandonment so the transport's + * in-flight HTTP request is cancelled, not just this iterator chain. */ private wrapStreamWithLifecycleCallbacks( result: StreamResult, options: StreamOptions, + teardownController?: AbortController, + // Releases the scoped composition in stream() that put this call's + // teardown signal alongside the caller's. Runs from teardown(), so the + // listeners come off the caller's (possibly long-lived) signal the + // moment this stream settles rather than whenever GC gets to it. + disposeComposedSignal?: () => void, ): StreamResult { const lifecycle = getLifecycleMiddlewareConfig(options); @@ -869,6 +951,31 @@ export abstract class BaseProvider implements AIProvider { [Symbol.asyncIterator]: () => upstreamIterator, }; + // Set by teardown, read by the generator's catch. Distinguishes "the + // provider failed" from "we aborted the provider because the consumer + // abandoned this stream", which must not be reported as a failure. + let teardownRequested = false; + let teardownDone = false; + + // Shared by both teardown entry points below (native `.return()` via the + // `finally`, and the `attachStreamCancel` hook for a consumer parked + // mid-`await` that a queued `.return()` cannot reach promptly). Guarded + // so whichever fires first does the work exactly once. + const teardown = (): void => { + if (teardownDone) { + return; + } + teardownDone = true; + teardownRequested = true; + // Abort before closing the iterators. The iterator close releases our + // side of the chain; this releases the provider's — the in-flight HTTP + // request that a closed iterator leaves running. + teardownController?.abort(); + disposeComposedSignal?.(); + cancelStream(originalStream); + releaseIterator(upstreamIterator); + }; + const wrappedStream = (async function* () { let accumulated = ""; let seq = 0; @@ -913,6 +1020,14 @@ export abstract class BaseProvider implements AIProvider { } } catch (error) { const err = error instanceof Error ? error : new Error(String(error)); + if (teardownRequested) { + // Our own abort, from a consumer that already broke out of this + // stream. Firing onError would report cleanup as a provider failure, + // and rethrowing would surface an error nobody is left to catch — + // an unhandled rejection, which terminates the process. Ending the + // generator quietly is the correct outcome for an abandoned stream. + return; + } if (onError && !hasLifecycleErrorFired(err)) { // Mark before firing so a higher layer that also routes through // fireLifecycleErrorCallback (or its own lifecycle wrapper) with @@ -930,6 +1045,21 @@ export abstract class BaseProvider implements AIProvider { ); } throw classifyStreamError(err); + } finally { + // Unconditional and safe: `teardown()` is idempotent (guarded by + // `teardownDone`), and `cancelStream`/`releaseIterator` are + // documented as no-ops on an already-finished stream/iterator, so + // calling it again after normal completion or a real provider error + // (both of which reach here too) costs nothing. The path that + // actually needs it is a consumer breaking out of `for await` while + // this generator is parked right after a `yield` — that unwinds + // straight here, bypassing both the rest of the `try` and the + // `catch` above, since a native `.return()` is a "return" + // completion, not a thrown error. Closing the iterator chain (done + // automatically by `for-await-of`'s own `IteratorClose` on + // `upstreamIterable`) does not cancel the HTTP request underneath + // it, so the transport is aborted here explicitly. + teardown(); } })(); @@ -937,11 +1067,12 @@ export abstract class BaseProvider implements AIProvider { // through `.return()` while it is parked awaiting the provider — that // request queues behind the in-flight `next()`. The hook closes the // upstream directly and forwards the request to any wrapper below, so - // abandoning a stream really does release the provider connection. - attachStreamCancel(wrappedStream, () => { - cancelStream(originalStream); - releaseIterator(upstreamIterator); - }); + // abandoning a stream really does release the provider connection. This + // is the fallback path for that case; the common case (a consumer that + // breaks right after receiving a chunk, so this generator is parked at + // the `yield` rather than mid-`await`) is handled by the `finally` + // above, which reaches `.return()` synchronously. + attachStreamCancel(wrappedStream, teardown); // See the comment in withStreamModelFallback above: this spread must not // be allowed to freeze a provider's lazy toolsUsed/toolExecutions getters. diff --git a/src/lib/core/nativeGenerateLoop.ts b/src/lib/core/nativeGenerateLoop.ts index 3c7a4ec14..138b658eb 100644 --- a/src/lib/core/nativeGenerateLoop.ts +++ b/src/lib/core/nativeGenerateLoop.ts @@ -109,8 +109,14 @@ export const appendJsonSchemaInstruction = ( conversation: Array>, schema: unknown, ): Array> => { + // "value", not "object": a ValidationSchema also accepts array and scalar + // roots (z.array(...), z.string()), and demanding an object told a + // compliant model to emit something those roots can never satisfy — so the + // fallback silently produced nothing for them however well the model + // complied. The generate path's own re-ask prompt was corrected for this + // reason; this shared helper, used by both paths, was missed. const instruction = - "When you give your final answer, respond with only a single JSON object " + + "When you give your final answer, respond with only a single JSON value " + "that conforms to the following JSON Schema. No prose before or after it, " + `and no markdown code fence. JSON Schema: ${JSON.stringify(schema)}`; const lastSystemIndex = conversation.reduce( diff --git a/src/lib/neurolink.ts b/src/lib/neurolink.ts index f45eeb176..82e283d86 100644 --- a/src/lib/neurolink.ts +++ b/src/lib/neurolink.ts @@ -5760,6 +5760,10 @@ Current user's request: ${currentInput}`; csvOptions: options.csvOptions, pdfOptions: options.pdfOptions, imageOptions: options.imageOptions, + // videoOptions belongs in that same list and was never in it, so no SDK + // caller's video settings have ever reached a provider -- frames and + // quality included, not just the native-video flags added alongside. + videoOptions: options.videoOptions, region: options.region, tts: options.tts, stt: options.stt, diff --git a/src/lib/providers/openaiChatCompletionsBase.ts b/src/lib/providers/openaiChatCompletionsBase.ts index b3b2c6a71..b0996661d 100644 --- a/src/lib/providers/openaiChatCompletionsBase.ts +++ b/src/lib/providers/openaiChatCompletionsBase.ts @@ -101,7 +101,11 @@ import { toolCallsFromSummaries, } from "../core/toolExecutionRecorder.js"; import { convertZodToJsonSchema } from "../utils/schemaConversion.js"; -import { coerceJsonToSchema, schemaAccepts } from "../utils/json/coerce.js"; +import { + coerceJsonToSchema, + recoverScalarRoot, + schemaAccepts, +} from "../utils/json/coerce.js"; import { resolveToolChoice } from "../utils/toolChoice.js"; import { transformToolExecutions } from "../utils/transformationUtils.js"; import { withProviderRetry } from "../utils/providerRetry.js"; @@ -165,6 +169,34 @@ const yieldsSchemaValidObject = ( return coerced !== null && schemaAccepts(schema, coerced.structuredData); }; +/** + * Resolve a model answer against the requested schema, accepting the roots a + * ValidationSchema actually allows. + * + * `coerceJsonToSchema` scans for a balanced object or array span, so it + * recovers object and array roots but never a scalar one: a schema of + * `z.string()` answered with `"sunny"` finds no span and yields null. That is + * why `recoverScalarRoot` exists — it was already the second half of this + * decision on the neurolink.ts generate path (`recoverStructuredData`), and + * was simply never wired into the provider path. Without it a scalar-root + * schema silently produces no structured output no matter how correctly the + * model answers. + * + * Returns `undefined` when neither route accepts, which callers treat as "the + * model did not produce the value". + */ +const resolveAgainstSchema = ( + text: string, + schema: ValidationSchema, +): unknown | undefined => { + const coerced = coerceJsonToSchema(text, schema); + if (coerced !== null && schemaAccepts(schema, coerced.structuredData)) { + return coerced.structuredData; + } + const scalar = recoverScalarRoot(text, schema); + return scalar.kind === "accepted" ? scalar.value : undefined; +}; + // Pull one native chunk at a time and forward cancellation to its iterator. const chunksToV3Stream = ( source: AsyncIterable, @@ -1746,6 +1778,222 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { return loop; } + /** + * Resolve the parsed object for a `stream({ schema })` turn, after the + * stream has been drained. + * + * Before this, streaming never sent `response_format` under any condition — + * unlike the generate path, where the gap is only tools-dependent (see + * suppressResponseFormatWithTools), so `structuredData` was never available + * to a streaming caller at all. `executeStream` now sends `response_format` + * on the wire whenever a schema is present and tools aren't suppressing it; + * this method covers what's left: the common case where the streamed text + * already parses against the schema, and the tools-suppressed case, which is + * re-asked once without tools — the same recovery `reformatWithoutTools` + * performs for generate(), reusing its acceptance rule + * (`yieldsSchemaValidObject`) and its two-strategy fallback (native + * `response_format` first, the schema spelled into the prompt if a vendor + * rejects that). + * + * Returns `undefined` on every failure path rather than throwing: a turn + * whose text the caller has already consumed and streamed to their UI must + * not be turned into an error by a follow-up they never asked for. + */ + private async resolveStreamStructuredData(args: { + schema: ValidationSchema; + schemaJson: Record; + text: string; + conversation: Array> | undefined; + modelId: string; + abortSignal?: AbortSignal; + options: StreamOptions; + }): Promise< + | { + data: unknown; + usage: { inputTokens: number; outputTokens: number }; + } + | undefined + > { + const { + schema, + schemaJson, + text, + conversation, + modelId, + abortSignal, + options, + } = args; + // A function, not a captured boolean: the signal is read once before the + // re-ask and again after it, and it can legitimately flip in between — a + // const read of `abortSignal.aborted` would let the compiler narrow the + // second read away (it cannot know the underlying signal mutates between + // the two checks). + const isCancelled = (): boolean => abortSignal?.aborted === true; + // Nothing streamed — a middleware short-circuit, or a turn that only + // called tools. Nothing to parse and nothing to restate. + if (text.trim().length === 0) { + return undefined; + } + // The happy path: the streamed answer already satisfies the schema (this + // is what sending `response_format` on the wire was for), so no second + // request. + const streamedValue = resolveAgainstSchema(text, schema); + if (streamedValue !== undefined) { + return { + data: streamedValue, + usage: { inputTokens: 0, outputTokens: 0 }, + }; + } + // No conversation means `doStream` never ran (a middleware short + // -circuited and returned its own stream), so there is no turn for a + // re-ask to build on. + if (!conversation || isCancelled()) { + return undefined; + } + try { + const model = await this.getAISDKModelWithMiddleware(options); + if (!hasNativeDoGenerate(model)) { + return undefined; + } + logger.warn( + `[${this.providerName}] streamed answer is not schema-valid — re-asking once without tools`, + { provider: this.providerName, model: modelId }, + ); + const doGenerate = model.doGenerate.bind(model); + // Every request here is billed, including one whose ANSWER is + // discarded (a native-format attempt the vendor rejected, or a re-ask + // that came back off-schema). Accumulate as each response lands rather + // than reading the survivor, for the same reason the generate path + // does: otherwise a two-request recovery reports as one. + const billed = { inputTokens: 0, outputTokens: 0 }; + const runLoop = ( + conv: Array>, + format: { type: "json"; schema: Record } | undefined, + ) => + runNativeGenerateLoop( + { + doGenerate, + conversation: conv, + toolsRecord: {}, + ...(format ? { responseFormat: format } : {}), + maxSteps: 1, + ...(options.maxTokens + ? { maxOutputTokens: options.maxTokens } + : {}), + ...(abortSignal ? { abortSignal } : {}), + runStep: (call) => call(), + }, + [], + ).then((loopResult) => { + billed.inputTokens += loopResult.inputTokens; + billed.outputTokens += loopResult.outputTokens; + return loopResult; + }); + // `conversation` here is the STREAMING loop's own array (see + // executeToolBatch), which is already wire-shaped: an assistant turn + // that called a tool carries a real `tool_calls` field and the matching + // reply is a `{role: "tool", tool_call_id, content}` message — not the + // ai-sdk-parts shape `messageBuilderToOpenAI` expects from `doGenerate`'s + // `options.prompt`. Re-running `runNativeGenerateLoop` on those turns + // unchanged would send them through `messageBuilderToOpenAI` a second + // time, which only ever reads the ai-sdk-parts fields — it silently + // drops `tool_calls` (it inspects `content` for `{type:"tool-call", + // toolCallId}` parts, never a wire `tool_calls` property) and blanks + // `tool_call_id` to `""` (it reads `toolCallId`, the pre-wire field, + // never the wire `tool_call_id`). The result is a `tool` message with no + // matching preceding tool call, which a real vendor rejects outright. + // None of that history is actually needed for this re-ask: the + // assistant turn appended below already restates whatever the tool call + // produced, in the model's own words, so the tool exchange itself is + // dropped rather than repaired. + const conversationWithoutToolExchange = conversation.filter((message) => { + if (message.role === "tool") { + return false; + } + const toolCalls = (message as { tool_calls?: unknown[] }).tool_calls; + return !( + message.role === "assistant" && + Array.isArray(toolCalls) && + toolCalls.length > 0 + ); + }); + const reaskConversation: Array> = [ + ...conversationWithoutToolExchange, + { role: "assistant", content: text }, + { + role: "user", + content: + // "value", not "object": ValidationSchema also accepts array and + // scalar roots (z.array(...), z.string()). + "Return that same answer as a single JSON value conforming to " + + "the required schema. No prose before or after it, and no " + + "markdown code fence.", + }, + ]; + let reformatted: Awaited>; + try { + reformatted = await runLoop(reaskConversation, { + type: "json", + schema: schemaJson, + }); + } catch (nativeFormatError) { + if (isCancelled()) { + throw nativeFormatError; + } + logger.warn( + `[${this.providerName}] the stream re-ask's response_format was rejected — retrying with the schema in the prompt`, + { + provider: this.providerName, + model: modelId, + error: + nativeFormatError instanceof Error + ? nativeFormatError.message + : String(nativeFormatError), + }, + ); + reformatted = await runLoop( + appendJsonSchemaInstruction(reaskConversation, schemaJson), + undefined, + ); + } + // Sending response_format is not the same as being obeyed: accept the + // reformat only if it actually produced the object. + const reaskedValue = resolveAgainstSchema(reformatted.text, schema); + if (reaskedValue === undefined) { + logger.warn( + `[${this.providerName}] the tool-free stream re-ask did not yield a schema-valid value either; leaving structuredData unset`, + { provider: this.providerName, model: modelId }, + ); + return { data: undefined, usage: billed }; + } + return { data: reaskedValue, usage: billed }; + } catch (error) { + // A caller who cancelled deserves their cancellation. Swallowing here + // resolved a cancelled turn as though the abort never happened — the + // inner handler already rethrows for exactly that case, and this outer + // one was catching its rethrow and turning it back into `undefined`. + // + // Our OWN turn deadline is the one thing that must NOT propagate: the + // caller has already consumed the streamed text, and a deadline reached + // during a follow-up they never asked for must not retroactively turn + // their answer into an error. The two arrive on the same composed + // signal but are distinguishable by reason, since the turn timer aborts + // with a TimeoutError. Same rule as the generate path's reformat guard. + if (isCancelled() && !(abortSignal?.reason instanceof TimeoutError)) { + throw error; + } + logger.warn( + `[${this.providerName}] the tool-free structured-output re-ask failed for a stream; leaving structuredData unset`, + { + provider: this.providerName, + model: modelId, + error: error instanceof Error ? error.message : String(error), + }, + ); + return undefined; + } + } + protected async executeStream( options: StreamOptions, _analysisSchema?: ZodUnknownSchema | Schema, @@ -1805,6 +2053,32 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { throw setupErr; } + // Structured output rides `response_format`, mirroring the generate + // path's policy (see suppressResponseFormatWithTools) — the streaming + // path previously never sent it AT ALL, regardless of tools, so + // `structuredData` was never available to a streaming caller in any + // case. Computed once, up front, exactly like the generate path: mid-turn + // tool hydration (tools.discovery) does not retroactively change whether + // this turn asked for JSON. + const schemaJson = options.schema + ? (convertZodToJsonSchema(options.schema as ZodUnknownSchema) as Record< + string, + unknown + >) + : undefined; + const hasToolsForStream = !!openAITools && openAITools.length > 0; + const responseFormatSuppressed = + schemaJson !== undefined && + hasToolsForStream && + this.suppressResponseFormatWithTools(); + const responseFormat: OpenAICompatResponseFormat | undefined = + schemaJson && !responseFormatSuppressed + ? this.adjustResponseFormat( + v3ResponseFormatToOpenAI({ type: "json", schema: schemaJson }), + modelId, + ) + : undefined; + const url = this.getChatCompletionsURL(modelId); const fetchImpl = createProxyFetch(); @@ -1854,6 +2128,14 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { }; let loopPromise: Promise | undefined; + // The wire conversation only exists once `doStream` actually runs, and it + // is built from the MIDDLEWARE-TRANSFORMED prompt, so it cannot be + // reconstructed from `promptMessages` afterwards. Captured here so the + // post-drain structured-output re-ask (resolveStreamStructuredData) can + // build on the same turn the model actually saw. Stays undefined when a + // middleware short-circuits the request — the case where there is no + // turn to build on at all. + let loopConversation: Array> | undefined; const providerNameForLoop = this.providerName; const streamBaseModel: LanguageModelV3 = { specificationVersion: "v3" as const, @@ -1880,6 +2162,7 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { transformedPrompt, wireNameMaps?.toWire, ); + loopConversation = conversation as Array>; const sampled: StreamOptions = { ...options, ...(typeof params?.maxOutputTokens === "number" @@ -1909,6 +2192,7 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { closeChannel: channel.close, resolveUsage, resolveFinish, + responseFormat, }); const completion: Promise = loopPromise.then( () => @@ -2004,8 +2288,23 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { } const providerName = this.providerName; + // Not an arrow function — `transformedStream` doesn't have `this` bound, + // so it needs its own capture to reach resolveStreamStructuredData. + const self = this; + // Built here rather than inline in `result` so the generator below can + // fill in `structuredData` after the stream drains. Same mutable + // -reference contract `finishReason`/`stopReason` on this object already + // rely on: a wrapper that spreads the result cannot snapshot a value that + // does not exist yet. + const streamMetadata: NonNullable = { + startTime, + streamId: `${this.providerName}-${Date.now()}`, + }; const transformedStream = async function* () { let contentYielded = 0; + // Only accumulated when a schema was actually requested — every other + // turn keeps paying nothing for this. + let schemaText = ""; try { for await (const chunk of chunkSource) { if ( @@ -2014,6 +2313,9 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { chunk.content.length > 0 ) { contentYielded++; + if (options.schema) { + schemaText += chunk.content; + } } yield chunk; } @@ -2021,6 +2323,52 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { // `loopPromise` is undefined when a middleware blocked the request // before `doStream` ran, in which case there is no loop to surface. await loopPromise; + // Structured output for a `stream({ schema })` turn. Runs HERE — + // after the stream is fully drained — so a tool-free re-ask (when the + // streamed answer isn't already schema-valid) never reaches + // `wrapStream` and a caller's guardrail middleware never sees a + // second stream it did not ask for. + // + // When a middleware short-circuited the request, `loopPromise` is + // undefined and `schemaText` is empty, so this resolves to nothing + // and the field stays absent rather than hanging a reader. + if (options.schema && schemaJson) { + const resolved = await self.resolveStreamStructuredData({ + schema: options.schema as ValidationSchema, + schemaJson, + text: schemaText, + conversation: loopConversation, + modelId, + abortSignal, + options, + }); + // Assign only when there is something to assign. A bare + // `= resolved?.data` creates the own property even when the value + // is undefined, so `"structuredData" in metadata` reports true for + // a turn that produced no object — which is precisely the check the + // field's own contract tells readers to use to detect absence. + if (resolved?.data !== undefined) { + streamMetadata.structuredData = resolved.data; + } + // Only when a re-ask actually spent something. The common case is + // the streamed answer already parsing, which costs nothing extra + // and leaves the field absent rather than reporting a zero. + if ( + resolved && + (resolved.usage.inputTokens > 0 || resolved.usage.outputTokens > 0) + ) { + streamMetadata.structuredDataUsage = resolved.usage; + logger.info( + `[${providerName}] structured-output re-ask billed separately from the streamed turn`, + { + provider: providerName, + model: modelId, + inputTokens: resolved.usage.inputTokens, + outputTokens: resolved.usage.outputTokens, + }, + ); + } + } // No-output path: stream completed normally but yielded zero text. // Build an enriched sentinel + stamp the active OTel span so // Pipeline B (ContextEnricher) surfaces a WARNING-level Langfuse @@ -2096,10 +2444,7 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { }, ), toolsUsed, - metadata: { - startTime, - streamId: `${this.providerName}-${Date.now()}`, - }, + metadata: streamMetadata, }; // Lazy getter: every read transforms the live `toolExecutionSummaries` // through the canonical `transformToolExecutions()` so consumers see @@ -2151,6 +2496,7 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { closeChannel, resolveUsage, resolveFinish, + responseFormat, } = args; // Hoisted above the try so the catch can resolve the usage accumulated @@ -2286,6 +2632,7 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { openAITools, openAIToolChoice, pushChunk, + responseFormat, }); lastObservedPromptTokens = stepResult.usage?.prompt_tokens; stepFinish = stepResult.finishReason; @@ -2339,6 +2686,7 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { openAITools: OpenAICompatChatTool[] | undefined; openAIToolChoice: OpenAICompatToolChoiceWire | undefined; pushChunk: (chunk: OpenAICompatStreamChunk) => void; + responseFormat?: OpenAICompatResponseFormat; }): Promise { // Per-step max_tokens fit: the conversation grows every step of the // tool loop, so the window fit is recomputed per request (no-op when @@ -2364,6 +2712,9 @@ export abstract class OpenAIChatCompletionsProvider extends BaseProvider { ? { toolChoice: args.openAIToolChoice } : {}), streaming: true, + ...(args.responseFormat + ? { responseFormat: args.responseFormat } + : {}), }), args.modelId, ), diff --git a/src/lib/server/utils/validation.ts b/src/lib/server/utils/validation.ts index 8717be486..ecbe87751 100644 --- a/src/lib/server/utils/validation.ts +++ b/src/lib/server/utils/validation.ts @@ -36,6 +36,35 @@ export const AgentExecuteRequestSchema = z.object({ userId: z.string().optional(), }); +/** + * Options a WebSocket `generate`/`stream` message may pass through to + * `neurolink.generate()`/`.stream()`. Deliberately an allowlist, not the full + * `GenerateOptions`/`StreamOptions` shape: the raw client JSON is untrusted, + * and `GenerateOptions.credentials` in particular carries a per-provider + * `baseURL` (see `NeurolinkCredentials`) that would otherwise let any caller + * redirect the server's own outbound request to an arbitrary host (SSRF) or + * swap in their own API key. `z.object` strips unknown keys by default, so + * `credentials` and anything else outside this list never reaches the SDK + * call — the same posture `AgentExecuteRequestSchema` already takes for the + * equivalent HTTP route. + */ +export const WebSocketAgentOptionsSchema = z.object({ + provider: z.string().optional(), + model: z.string().optional(), + systemPrompt: z.string().optional(), + temperature: z.number().min(0).max(2).optional(), + maxTokens: z.number().positive().optional(), + maxSteps: z.number().positive().optional(), +}); + +/** + * WebSocket `generate`/`stream` message payload schema. + */ +export const WebSocketAgentRequestSchema = z.object({ + prompt: z.string().min(1, "Prompt is required"), + options: WebSocketAgentOptionsSchema.optional(), +}); + /** * Tool execute request schema */ diff --git a/src/lib/server/websocket/WebSocketHandler.ts b/src/lib/server/websocket/WebSocketHandler.ts index 37609a10b..f3b041e8d 100644 --- a/src/lib/server/websocket/WebSocketHandler.ts +++ b/src/lib/server/websocket/WebSocketHandler.ts @@ -10,11 +10,14 @@ import type { WebSocketHandler as IWebSocketHandler, WebSocketConnection, WebSocketMessage, + WebSocketRequestId, AuthenticatedUser, ServerAuthConfig as _AuthConfig, } from "../../types/index.js"; +import type { NeuroLink } from "../../neurolink.js"; import { WebSocketError, WebSocketConnectionError } from "../errors.js"; import { logger } from "../../utils/logger.js"; +import { WebSocketAgentRequestSchema } from "../utils/validation.js"; /** * Default WebSocket configuration @@ -81,6 +84,65 @@ export class WebSocketConnectionManager { ); } + // Enforce the auth contract the config already declares. `auth.required` + // has existed on WebSocketConfig (with a `strategy`) since this manager + // was written, and `handleConnection` has always accepted and stored an + // AuthenticatedUser — but nothing ever checked either, so a deployer who + // set `required: true` got exactly the same open socket as one who did + // not. That was latent while the agent routes were stubs that returned + // canned values; it stops being latent the moment `tool_call` actually + // reaches `executeTool`, because then an unauthenticated client can + // invoke any registered tool with arbitrary arguments. Refuse the + // connection rather than letting a route decide, so every handler + // registered on this manager inherits the gate. + if (this.config.auth.required && !user) { + logger.warn( + `[WebSocket] Rejected unauthenticated connection to ${path} (auth.required)`, + ); + throw new WebSocketConnectionError( + "Authentication required for this WebSocket endpoint", + ); + } + + // Same gap, same fix, one field over. `ServerAuthConfig.roles`/ + // `permissions` and `AuthenticatedUser.roles`/`permissions` have existed + // since this manager was written — the shipped "Role-Based Access + // Control" doc example configures them on a `/ws/admin` path — but + // nothing ever consulted either, so a deployer following that example + // got an endpoint anyone authenticated could reach, regardless of role. + // Semantics match the SDK's own `hasAnyRole`/`hasAllPermissions` + // (src/lib/auth/authContext.ts): roles are an any-of allowlist, + // permissions must all be present. + const requiredRoles = this.config.auth.roles; + if (requiredRoles && requiredRoles.length > 0) { + const userRoles = user?.roles ?? []; + if (!requiredRoles.some((role) => userRoles.includes(role))) { + logger.warn( + `[WebSocket] Rejected connection to ${path}: missing a required role`, + ); + throw new WebSocketConnectionError( + "Insufficient role for this WebSocket endpoint", + ); + } + } + + const requiredPermissions = this.config.auth.permissions; + if (requiredPermissions && requiredPermissions.length > 0) { + const userPermissions = user?.permissions ?? []; + if ( + !requiredPermissions.every((permission) => + userPermissions.includes(permission), + ) + ) { + logger.warn( + `[WebSocket] Rejected connection to ${path}: missing a required permission`, + ); + throw new WebSocketConnectionError( + "Insufficient permissions for this WebSocket endpoint", + ); + } + } + const connection: WebSocketConnection = { id: generateConnectionId(), socket, @@ -390,13 +452,29 @@ export class WebSocketConnectionManager { } } +/** + * Narrow an incoming message's raw `id` field to the one shape it is safe to + * echo back (see `WebSocketRequestId`); anything else — including a missing + * field, since most callers won't opt into correlation at all — is treated + * as "no id". + */ +function readRequestId(value: unknown): WebSocketRequestId | undefined { + return typeof value === "string" || typeof value === "number" + ? value + : undefined; +} + /** * WebSocket message router for handling different message types */ export class WebSocketMessageRouter { private routes = new Map< string, - (connection: WebSocketConnection, payload: unknown) => Promise + ( + connection: WebSocketConnection, + payload: unknown, + requestId?: WebSocketRequestId, + ) => Promise >(); /** @@ -407,6 +485,7 @@ export class WebSocketMessageRouter { handler: ( connection: WebSocketConnection, payload: unknown, + requestId?: WebSocketRequestId, ) => Promise, ): void { this.routes.set(type, handler); @@ -423,14 +502,15 @@ export class WebSocketMessageRouter { throw new WebSocketError("Only text messages are supported for routing"); } - let parsed: { type: string; payload?: unknown }; + let parsed: { type: string; payload?: unknown; id?: WebSocketRequestId }; try { parsed = JSON.parse(message.data as string); } catch { throw new WebSocketError("Invalid JSON message"); } - const { type, payload } = parsed; + const { type, payload, id } = parsed; + const requestId = readRequestId(id); if (!type) { throw new WebSocketError("Message type is required"); @@ -441,7 +521,26 @@ export class WebSocketMessageRouter { throw new WebSocketError(`Unknown message type: ${type}`); } - return handler(connection, payload); + const result = await handler(connection, payload, requestId); + + // Overlapping in-flight messages on one connection previously produced + // response frames with nothing tying them back to the request that + // caused them (silent when the client had more than one call open at + // once). A route that already tags its own frame — the streaming route + // does, for the intermediate frames it sends directly — is left alone; + // this only fills in the id for a plain result object that doesn't + // carry one, so a route that ignores `requestId` still gets correlated + // for free, and a client that never sends an `id` sees no change at all. + if ( + requestId !== undefined && + result !== null && + typeof result === "object" && + !("id" in (result as Record)) + ) { + return { ...(result as Record), id: requestId }; + } + + return result; } /** @@ -455,37 +554,143 @@ export class WebSocketMessageRouter { /** * Create a WebSocket handler for AI agent interactions */ +/** + * Send one frame, reporting failure instead of throwing. + * + * A socket that has already closed throws on `send`. That matters in two + * places here: mid-stream, where an unguarded throw unwinds into `onMessage`'s + * catch, which then tries to send an error frame on the same dead socket and + * throws again — this time with nothing above it to catch, i.e. an unhandled + * rejection; and in that catch itself, for the same reason. Returning a + * boolean lets both callers stop cleanly instead. + */ +function trySend(connection: { socket: unknown }, payload: unknown): boolean { + const socket = connection.socket as { send: (data: string) => void }; + try { + socket.send(JSON.stringify(payload)); + return true; + } catch { + return false; + } +} + +/** + * Pull a correlation id off a raw incoming frame without depending on + * `WebSocketMessageRouter.handle` having successfully parsed it — the one + * case this is for is a message that fails *inside* `handle` (bad JSON, + * missing type, unknown route), where the router never gets far enough to + * return the id it would otherwise have tagged the response with. Failure + * here (not text, not JSON, no `id`) just means "no id", the same as if the + * client never sent one. + */ +function readIncomingRequestId( + message: WebSocketMessage, +): WebSocketRequestId | undefined { + if (message.type !== "text") { + return undefined; + } + try { + const parsed = JSON.parse(message.data as string) as { id?: unknown }; + return readRequestId(parsed.id); + } catch { + return undefined; + } +} + export function createAgentWebSocketHandler( - _neurolink: unknown, + neurolink: NeuroLink, ): IWebSocketHandler { const router = new WebSocketMessageRouter(); // Register message routes + // + // `payload` is attacker-controlled JSON (WebSocketMessageRouter.handle does + // nothing but JSON.parse it), so `generate`/`stream` validate it against + // `WebSocketAgentRequestSchema` before it reaches the SDK rather than + // trusting a bare type assertion. The schema's `options` allowlist is + // deliberately narrower than `GenerateOptions`/`StreamOptions` — most + // importantly it has no `credentials` field, so a client can no longer + // point the server's outbound request at an arbitrary `baseURL` (SSRF) or + // swap in its own provider API key. router.route("generate", async (connection, payload) => { - const { prompt, options: _options } = payload as { - prompt: string; - options?: unknown; - }; - // TODO(#1576): Implement generate using neurolink - return { type: "response", data: `Received: ${prompt}` }; + const parsed = WebSocketAgentRequestSchema.safeParse(payload); + if (!parsed.success) { + throw new WebSocketError( + `Invalid generate payload: ${parsed.error.issues.map((issue) => issue.message).join(", ")}`, + ); + } + const { prompt, options } = parsed.data; + const result = await neurolink.generate({ + ...options, + input: { text: prompt }, + }); + return { type: "response", content: result.content }; }); - router.route("stream", async (connection, payload) => { - const { prompt, options: _options } = payload as { - prompt: string; - options?: unknown; - }; - // TODO(#1576): Implement streaming using neurolink - return { type: "stream_start", data: { prompt } }; + router.route("stream", async (connection, payload, requestId) => { + const parsed = WebSocketAgentRequestSchema.safeParse(payload); + if (!parsed.success) { + throw new WebSocketError( + `Invalid stream payload: ${parsed.error.issues.map((issue) => issue.message).join(", ")}`, + ); + } + const { prompt, options } = parsed.data; + + // stream_start/chunk are sent directly through trySend rather than + // returned, so — unlike generate/tool_call's single returned frame — + // they sit outside WebSocketMessageRouter.handle's own id-tagging and + // must carry the correlation id themselves. + const withRequestId = ( + frame: Record, + ): Record => + requestId !== undefined ? { ...frame, id: requestId } : frame; + + if ( + !trySend( + connection, + withRequestId({ type: "stream_start", data: { prompt } }), + ) + ) { + return undefined; + } + + const result = await neurolink.stream({ + ...options, + input: { text: prompt }, + }); + let clientPresent = true; + for await (const chunk of result.stream) { + if ("content" in chunk && typeof chunk.content === "string") { + if ( + !trySend( + connection, + withRequestId({ type: "chunk", content: chunk.content }), + ) + ) { + // The client is gone. Breaking runs the stream's IteratorClose, + // which now aborts the upstream request rather than leaving it in + // flight for a reader that no longer exists. + clientPresent = false; + logger.info( + `[AgentWebSocket] Client went away mid-stream: ${connection.id}`, + ); + break; + } + } + } + + // Returning nothing keeps onMessage from sending a completion frame on a + // socket that just refused one. + return clientPresent ? { type: "stream_complete" } : undefined; }); router.route("tool_call", async (connection, payload) => { - const { toolName, args: _args } = payload as { + const { toolName, args } = payload as { toolName: string; args: unknown; }; - // TODO(#1576): Implement tool call using neurolink - return { type: "tool_result", data: { toolName, result: null } }; + const result = await neurolink.executeTool(toolName, args); + return { type: "tool_result", toolName, result }; }); return { @@ -502,20 +707,32 @@ export function createAgentWebSocketHandler( }, onMessage: async (connection, message) => { + // Best-effort only: a malformed frame still fails inside + // router.handle below (and reports as an id-less error, same as + // before), this just lets a *valid* frame's error response carry the + // same correlation id a successful response would have gotten. + const requestId = readIncomingRequestId(message); try { const result = await router.handle(connection, message); if (result) { - const socket = connection.socket as { send: (data: string) => void }; - socket.send(JSON.stringify(result)); + trySend(connection, result); } } catch (error) { - const socket = connection.socket as { send: (data: string) => void }; - socket.send( - JSON.stringify({ - type: "error", - error: (error as Error).message, - }), - ); + // Guarded: this catch is the last line of defence, so a throw from + // the send itself — which is exactly what a closed socket does — has + // nothing above it and would surface as an unhandled rejection. + const errorFrame: Record = { + type: "error", + error: (error as Error).message, + }; + if (requestId !== undefined) { + errorFrame.id = requestId; + } + if (!trySend(connection, errorFrame)) { + logger.warn( + `[AgentWebSocket] Could not deliver error frame, socket closed: ${connection.id}`, + ); + } } }, diff --git a/src/lib/types/generate.ts b/src/lib/types/generate.ts index 4d780351f..bce725d71 100644 --- a/src/lib/types/generate.ts +++ b/src/lib/types/generate.ts @@ -1624,6 +1624,22 @@ export type TextGenerationOptions = { outputFormat?: VisionImageOutputFormat; }; + /** + * Video processing options (mirrors `GenerateOptions.videoOptions`). Never + * declared here before, so `buildGenerateTextOptions` had nowhere to + * forward the caller's setting even once the allowlist itself named it. + */ + videoOptions?: { + /** Frames to extract. Unset lets VideoProcessor pick from the clip's duration; clamped to 100. */ + frames?: number; + /** Frame encoder quality, clamped to 1-100. Default 80. */ + quality?: number; + /** Frame encoding. Default jpeg. */ + format?: "jpeg" | "png"; + /** Not implemented yet (#433) — warns rather than silently doing nothing. */ + transcribeAudio?: boolean; + }; + enableSummarization?: boolean; // Enable/disable summarization for this specific request /** diff --git a/src/lib/types/openaiCompatible.ts b/src/lib/types/openaiCompatible.ts index 4d354569d..f921e163e 100644 --- a/src/lib/types/openaiCompatible.ts +++ b/src/lib/types/openaiCompatible.ts @@ -320,6 +320,12 @@ export type StreamLoopArgs = { totalTokens: number; }) => void; resolveFinish: (reason: string) => void; + /** + * `response_format` for this turn's requests, computed once up front from + * `options.schema` (see `suppressResponseFormatWithTools`). Absent when no + * schema was requested, or when tools suppress it. + */ + responseFormat?: OpenAICompatResponseFormat; }; export type OpenAICompatBuildBodyArgs = { diff --git a/src/lib/types/server.ts b/src/lib/types/server.ts index d8cdb198f..defaec6c4 100644 --- a/src/lib/types/server.ts +++ b/src/lib/types/server.ts @@ -918,6 +918,15 @@ export type WebSocketMessage = { timestamp: number; }; +/** + * Client-supplied correlation id on an agent-protocol WebSocket message + * (`{ type, payload, id }`). Echoed back on the matching response frame so a + * caller with more than one in-flight message on a connection can tell which + * frame answers which request; optional and ignored entirely by a client + * that never sends one. + */ +export type WebSocketRequestId = string | number; + /** * Authenticated user information */ diff --git a/src/lib/types/stream.ts b/src/lib/types/stream.ts index c2c248ed1..ee96a3164 100644 --- a/src/lib/types/stream.ts +++ b/src/lib/types/stream.ts @@ -987,6 +987,36 @@ export type StreamResult = { stopReason?: GenerateStopReason; rawFinishReason?: string; stepsUsed?: number; + /** + * Parsed structured output for a `stream({ schema })` turn, available + * AFTER the stream has been drained. + * + * Lives on metadata for the same reason `finishReason` above does: it is a + * mutable reference the loop fills in, so result-object spreads in stream + * wrappers cannot snapshot it before it resolves. A plain field rather than + * a promise, deliberately — when a middleware short-circuits the request + * and the loop never runs, an unresolved promise would hang every reader, + * whereas an absent field is simply absent. Readers must tolerate that: + * absence means the model never produced the object. + */ + structuredData?: unknown; + /** + * Tokens spent by the tool-free re-ask that produced `structuredData`, + * when one was needed. + * + * That re-ask is a second, separately-billed model call, and the stream's + * own `usage` has already resolved by the time it runs — so folding these + * tokens into it would mutate a value a caller may have read, and leaving + * them out entirely would under-report what the turn cost. Reported here + * instead: same delivery as `structuredData`, filled at the same moment, + * read at the same moment. Absent when no re-ask was needed, which is the + * common case. The generate path accounts for its equivalent re-ask + * inline, since there the usage has not been handed out yet. + */ + structuredDataUsage?: { + inputTokens: number; + outputTokens: number; + }; // Thought/reasoning metadata thoughtSignature?: string; thoughts?: Array<{ id?: string; type?: string; content?: string }>; diff --git a/test/continuous-test-suite-stream-structured-output.ts b/test/continuous-test-suite-stream-structured-output.ts new file mode 100644 index 000000000..06dc8be22 --- /dev/null +++ b/test/continuous-test-suite-stream-structured-output.ts @@ -0,0 +1,579 @@ +#!/usr/bin/env tsx +import "dotenv/config"; + +/** + * Continuous Test Suite — stream({ schema }) structured output. + * + * Before this fix, the streaming path never sent `response_format` on the + * wire under any condition — unlike `generate()`, where the gap is only + * tools-dependent (see `suppressResponseFormatWithTools`) — so + * `metadata.structuredData` did not exist on `StreamResult` at all and a + * `stream({ schema })` caller had no way to get parsed structured output. + * + * This suite drives `nl.stream({ schema })` against a local, scripted + * OpenAI-compatible SSE server (no credentials, no network egress — see + * `startStructuredSSEServer` below) and, after fully draining the stream, + * asserts: + * 1. the wire request actually carried `response_format` (the request the + * mock server received is captured verbatim), and + * 2. `result.metadata.structuredData` is present and validates against the + * Zod schema that was requested. + * + * `metadata.structuredData` is a mutable reference the provider fills in only + * after the stream fully drains (see the doc comment on `StreamResult`'s + * `metadata.structuredData` in `src/lib/types/stream.ts`), so this suite + * reads it only after the `for await` loop over `result.stream` completes — + * reading it earlier would legitimately race the resolution and is not what + * this suite is testing. + * + * NEUROLINK_DISABLE_BUILTIN_TOOLS is set before the dist import so no + * built-in tool gets auto-attached and suppresses `response_format` for a + * reason unrelated to the defect under test (see + * `suppressResponseFormatWithTools`). + * + * Run: pnpm run build && npx tsx test/continuous-test-suite-stream-structured-output.ts + * pnpm run test:stream-structured-output + */ + +process.env.NEUROLINK_DISABLE_BUILTIN_TOOLS = "true"; + +import { createServer, type IncomingMessage, type Server } from "node:http"; +import { z } from "zod"; +import { defineSuite, assert } from "./helpers/harness.js"; +import { assertDistFresh } from "./helpers/distFreshness.js"; + +assertDistFresh(); + +const { test, runSuite } = defineSuite("Stream structured output", { + offline: true, +}); + +const { NeuroLink, tool } = await import("../dist/index.js"); + +const credentialsFor = (baseURL: string) => ({ + openai: { apiKey: "sk-mock-local-server", baseURL }, +}); + +function sseChunk(payload: Record): string { + return `data: ${JSON.stringify(payload)}\n\n`; +} + +type CapturedServer = { + baseURL: string; + requestBodies: string[]; + close: () => Promise; +}; + +/** + * A local server that answers every streaming chat-completions request with + * an SSE stream whose content deltas, concatenated, equal `jsonText` — split + * across two chunks so this exercises real incremental accumulation rather + * than a single delta carrying the whole payload. Every raw request body is + * recorded, oldest first, so a test can assert on what NeuroLink actually + * sent (e.g. `response_format`) without mocking `fetch`. + */ +function startStructuredSSEServer(jsonText: string): Promise { + const requestBodies: string[] = []; + const server: Server = createServer((req: IncomingMessage, res) => { + const chunks: Buffer[] = []; + req.on("data", (chunk: Buffer) => chunks.push(chunk)); + req.on("end", () => { + requestBodies.push(Buffer.concat(chunks).toString("utf8")); + + res.writeHead(200, { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + Connection: "keep-alive", + }); + const base = { + id: "mock-structured-stream", + object: "chat.completion.chunk", + created: Math.floor(Date.now() / 1000), + model: "scripted-model", + }; + const mid = Math.floor(jsonText.length / 2); + for (const part of [jsonText.slice(0, mid), jsonText.slice(mid)]) { + res.write( + sseChunk({ + ...base, + choices: [ + { index: 0, delta: { content: part }, finish_reason: null }, + ], + }), + ); + } + res.write( + sseChunk({ + ...base, + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + }), + ); + res.write("data: [DONE]\n\n"); + res.end(); + }); + }); + + return new Promise((resolve) => { + server.listen(0, "127.0.0.1", () => { + const address = server.address(); + const port = typeof address === "object" && address ? address.port : 0; + resolve({ + baseURL: `http://127.0.0.1:${port}/v1`, + requestBodies, + close: () => new Promise((r) => server.close(() => r())), + }); + }); + }); +} + +const weatherSchema = z.object({ + city: z.string(), + temp: z.number(), +}); + +await test("stream({ schema }) sends response_format on the wire and resolves metadata.structuredData after drain", async () => { + const server = await startStructuredSSEServer( + JSON.stringify({ city: "Paris", temp: 21 }), + ); + try { + const nl = new NeuroLink({ conversationMemory: { enabled: false } }); + const result = await nl.stream({ + input: { text: "What is the weather in Paris?" }, + provider: "openai", + model: "scripted-model", + credentials: credentialsFor(server.baseURL), + schema: weatherSchema, + }); + + // Drain fully before reading metadata.structuredData — it is a + // mutable reference the provider only fills in once the stream loop + // has finished (see the file-header doc comment). + let text = ""; + for await (const chunk of result.stream) { + if ("content" in chunk && typeof chunk.content === "string") { + text += chunk.content; + } + } + assert(text.length > 0, "no text content was yielded by the stream at all"); + + assert( + server.requestBodies.length >= 1, + "no request reached the mock server", + ); + const wireBody = JSON.parse(server.requestBodies[0]) as { + response_format?: unknown; + }; + assert( + wireBody.response_format !== undefined, + "response_format was not present on the wire request", + ); + + const structuredData = result.metadata?.structuredData; + assert( + structuredData !== undefined, + "metadata.structuredData was not populated after the stream drained", + ); + const parsed = weatherSchema.safeParse(structuredData); + assert( + parsed.success, + "metadata.structuredData did not validate against the requested schema", + ); + } finally { + await server.close(); + } +}); + +/** + * The tools-suppressed branch, which the case above cannot reach. + * + * `suppressResponseFormatWithTools()` is true for every provider except + * OpenAI and Azure, so a turn that carries tools goes out WITHOUT + * `response_format` and the model answers in prose. Sending the wire field is + * therefore not enough on its own: the only way a `stream({ schema, tools })` + * caller ever gets structured output is the one tool-free re-ask that + * `resolveStreamStructuredData` performs after the stream drains. + * + * This server scripts all three requests that flow makes, branching on + * whether the body asked for `stream: true`: + * 1. streaming, no prior tool reply -> answer with a `tool_calls` delta + * 2. streaming, after the tool reply -> stream PROSE, deliberately not JSON + * 3. non-streaming -> the re-ask; return the real object + * + * Every raw body is recorded so the test can assert on what actually went out + * rather than on what the implementation says it sends. + */ +function startToolThenReaskServer( + proseAnswer: string, + reaskJson: string, +): Promise { + const requestBodies: string[] = []; + const server: Server = createServer((req: IncomingMessage, res) => { + const chunks: Buffer[] = []; + req.on("data", (chunk: Buffer) => chunks.push(chunk)); + req.on("end", () => { + const raw = Buffer.concat(chunks).toString("utf8"); + requestBodies.push(raw); + const body = JSON.parse(raw) as { + stream?: boolean; + messages?: Array<{ role?: string }>; + }; + + if (body.stream !== true) { + // The tool-free re-ask. Non-streaming, so it answers as a plain + // chat completion rather than SSE. + res.writeHead(200, { "Content-Type": "application/json" }); + res.end( + JSON.stringify({ + id: "mock-reask", + object: "chat.completion", + created: Math.floor(Date.now() / 1000), + model: "scripted-model", + choices: [ + { + index: 0, + message: { role: "assistant", content: reaskJson }, + finish_reason: "stop", + }, + ], + usage: { + prompt_tokens: 12, + completion_tokens: 8, + total_tokens: 20, + }, + }), + ); + return; + } + + const sawToolReply = (body.messages ?? []).some( + (message) => message.role === "tool", + ); + res.writeHead(200, { + "Content-Type": "text/event-stream", + "Cache-Control": "no-cache", + Connection: "keep-alive", + }); + const base = { + id: "mock-tool-stream", + object: "chat.completion.chunk", + created: Math.floor(Date.now() / 1000), + model: "scripted-model", + }; + + if (!sawToolReply) { + res.write( + sseChunk({ + ...base, + choices: [ + { + index: 0, + delta: { + tool_calls: [ + { + index: 0, + id: "call_reask_fixture", + type: "function", + function: { + name: "lookup_weather", + arguments: '{"city":"Paris"}', + }, + }, + ], + }, + finish_reason: null, + }, + ], + }), + ); + res.write( + sseChunk({ + ...base, + choices: [{ index: 0, delta: {}, finish_reason: "tool_calls" }], + }), + ); + } else { + const mid = Math.floor(proseAnswer.length / 2); + for (const part of [ + proseAnswer.slice(0, mid), + proseAnswer.slice(mid), + ]) { + res.write( + sseChunk({ + ...base, + choices: [ + { index: 0, delta: { content: part }, finish_reason: null }, + ], + }), + ); + } + res.write( + sseChunk({ + ...base, + choices: [{ index: 0, delta: {}, finish_reason: "stop" }], + }), + ); + } + res.write("data: [DONE]\n\n"); + res.end(); + }); + }); + + return new Promise((resolve) => { + server.listen(0, "127.0.0.1", () => { + const address = server.address(); + const port = typeof address === "object" && address ? address.port : 0; + resolve({ + baseURL: `http://127.0.0.1:${port}/v1`, + requestBodies, + close: () => new Promise((r) => server.close(() => r())), + }); + }); + }); +} + +/** + * Deliberately NOT the `openai` provider. `suppressResponseFormatWithTools()` + * returns the base default `true` everywhere except OpenAI and Azure, which + * override it to `false` — so driving this case through `openai` would send + * `response_format` alongside the tools, the streamed answer would already be + * JSON, and the re-ask under test would never run. `openai-compatible` takes + * the default, which is what the overwhelming majority of providers do. + */ +const COMPAT_ENV = [ + "OPENAI_COMPATIBLE_API_KEY", + "OPENAI_COMPATIBLE_BASE_URL", +] as const; + +await test("stream({ schema }) with tools re-asks once without tools and still resolves structuredData", async () => { + const PROSE = "It is a mild 21 degrees in Paris right now."; + let toolCalls = 0; + const server = await startToolThenReaskServer( + PROSE, + JSON.stringify({ city: "Paris", temp: 21 }), + ); + const savedEnv: Record = {}; + for (const key of COMPAT_ENV) { + savedEnv[key] = process.env[key]; + } + try { + process.env.OPENAI_COMPATIBLE_API_KEY = "sk-mock-local-server"; + process.env.OPENAI_COMPATIBLE_BASE_URL = server.baseURL; + + const nl = new NeuroLink({ conversationMemory: { enabled: false } }); + const result = await nl.stream({ + input: { text: "What is the weather in Paris?" }, + provider: "openai-compatible", + model: "scripted-model", + schema: weatherSchema, + tools: { + lookup_weather: tool({ + inputSchema: z.object({ city: z.string() }), + execute: async ({ city }: { city: string }) => { + toolCalls++; + return { city, temp: 21 }; + }, + }), + }, + } as Parameters["stream"]>[0]); + + let text = ""; + for await (const chunk of result.stream) { + if ("content" in chunk && typeof chunk.content === "string") { + text += chunk.content; + } + } + + assert(toolCalls === 1, "the scripted tool did not run exactly once"); + // The caller's own stream must be untouched by the recovery: they asked + // for a stream and consumed prose, and the re-ask happens behind it. + assert( + text.trim() === PROSE, + "the drained stream did not deliver the scripted prose answer", + ); + + const bodies = server.requestBodies.map( + (raw) => + JSON.parse(raw) as { + stream?: boolean; + tools?: unknown[]; + response_format?: unknown; + messages?: Array>; + }, + ); + const streamed = bodies.filter((b) => b.stream === true); + const reasks = bodies.filter((b) => b.stream !== true); + + assert(streamed.length >= 1, "no streaming request reached the fixture"); + assert( + Array.isArray(streamed[0].tools) && streamed[0].tools.length > 0, + "the first streaming request did not carry the caller's tool", + ); + // The premise of the whole branch: with tools attached the wire field is + // suppressed, so the streamed turn cannot itself produce structured data. + assert( + streamed[0].response_format === undefined, + "the tools-bearing streaming request unexpectedly carried the format field", + ); + + assert( + reasks.length === 1, + "the tool-free re-ask did not happen exactly once", + ); + const reask = reasks[0]; + assert( + reask.response_format !== undefined, + "the re-ask did not carry the format field", + ); + assert( + !Array.isArray(reask.tools) || reask.tools.length === 0, + "the re-ask still carried tools, so it was not the tool-free retry", + ); + + // The conversation-rebuilding fix: a `tool` reply whose `tool_call_id` has + // no matching `tool_calls` entry in a preceding assistant turn is what a + // real vendor rejects outright. Assert the exchange is coherent rather + // than trusting that it was filtered. + const reaskMessages = reask.messages ?? []; + const announced = new Set(); + for (const message of reaskMessages) { + const calls = (message as { tool_calls?: Array<{ id?: string }> }) + .tool_calls; + if (Array.isArray(calls)) { + for (const call of calls) { + if (typeof call.id === "string") { + announced.add(call.id); + } + } + } + } + const orphanReplies = reaskMessages.filter((message) => { + if (message.role !== "tool") { + return false; + } + const id = (message as { tool_call_id?: unknown }).tool_call_id; + return typeof id !== "string" || id.length === 0 || !announced.has(id); + }); + assert( + orphanReplies.length === 0, + "the re-ask conversation carries a tool reply with no matching announced call", + ); + + const structuredData = result.metadata?.structuredData; + assert( + structuredData !== undefined, + "metadata.structuredData was not populated after the tools-suppressed turn", + ); + const parsed = weatherSchema.safeParse(structuredData); + assert( + parsed.success, + "the re-asked structuredData did not validate against the requested schema", + ); + + // The re-ask is a second billed model call whose tokens cannot be folded + // into the stream's usage — that resolved before the re-ask ran, and a + // caller may already have read it. Reporting them separately is what + // keeps the turn's real cost visible; raised in review on this PR. + const reaskUsage = result.metadata?.structuredDataUsage; + assert( + reaskUsage !== undefined, + "the separately-billed re-ask reported no usage at all", + ); + assert( + (reaskUsage?.outputTokens ?? 0) > 0, + "the re-ask reported no output tokens despite having answered", + ); + assert( + (reaskUsage?.inputTokens ?? 0) > 0, + "the re-ask reported no input tokens despite having sent a conversation", + ); + } finally { + await server.close(); + for (const key of COMPAT_ENV) { + const prior = savedEnv[key]; + if (prior === undefined) { + delete process.env[key]; + } else { + process.env[key] = prior; + } + } + } +}); + +/** + * Scalar-root schemas, raised in review on this PR. + * + * A ValidationSchema accepts more than objects — `z.string()` is a legitimate + * schema. `coerceJsonToSchema` scans for a balanced object or array span, so + * a scalar answer finds no span and yields null, and the field stayed unset + * however correctly the model answered. `recoverScalarRoot` already existed + * for exactly this and was wired only into the neurolink.ts generate path. + */ +const cityNameSchema = z.string(); + +await test("stream({ schema }) resolves a scalar-root schema, not just objects", async () => { + // A bare JSON string: valid JSON, valid against the schema, and containing + // no object or array span for a bracket scan to find. + const server = await startStructuredSSEServer(JSON.stringify("Paris")); + try { + const nl = new NeuroLink({ conversationMemory: { enabled: false } }); + const result = await nl.stream({ + input: { text: "Which city?" }, + provider: "openai", + model: "scripted-model", + credentials: credentialsFor(server.baseURL), + schema: cityNameSchema, + }); + + let text = ""; + for await (const chunk of result.stream) { + if ("content" in chunk && typeof chunk.content === "string") { + text += chunk.content; + } + } + assert(text.length > 0, "no text content was yielded by the stream at all"); + + const structuredData = result.metadata?.structuredData; + assert( + structuredData !== undefined, + "a scalar-root schema left structuredData unset", + ); + assert( + cityNameSchema.safeParse(structuredData).success, + "the scalar structuredData did not validate against the requested schema", + ); + } finally { + await server.close(); + } +}); + +await test("a turn that produces no object leaves structuredData absent, not present-and-undefined", async () => { + // Prose that satisfies no schema, with tools off so no re-ask path runs. + const server = await startStructuredSSEServer("not json at all"); + try { + const nl = new NeuroLink({ conversationMemory: { enabled: false } }); + const result = await nl.stream({ + input: { text: "Which city?" }, + provider: "openai", + model: "scripted-model", + credentials: credentialsFor(server.baseURL), + schema: weatherSchema, + }); + for await (const chunk of result.stream) { + void chunk; + } + + // The field's own contract tells readers absence means the model never + // produced the object, and `in` is how a reader checks that. An + // unconditional assignment creates the own property even when the value + // is undefined, which reports true for a turn that produced nothing. + const metadata = result.metadata ?? {}; + assert( + !("structuredData" in metadata), + "structuredData exists as an own property on a turn that produced no object", + ); + } finally { + await server.close(); + } +}); + +await runSuite(); diff --git a/test/continuous-test-suite-stream-teardown-abort.ts b/test/continuous-test-suite-stream-teardown-abort.ts new file mode 100644 index 000000000..ac8466393 --- /dev/null +++ b/test/continuous-test-suite-stream-teardown-abort.ts @@ -0,0 +1,368 @@ +#!/usr/bin/env tsx +import "dotenv/config"; + +/** + * Continuous Test Suite — abandoning a stream cancels the underlying HTTP + * request (teardown → transport abort). + * + * ALL-DIST module graph (rule 15): every provider call below is driven + * through `new NeuroLink().stream()`, imported from `../dist/index.js`, the + * same shipped surface every caller uses. No `src/lib/` import appears in + * this file. + * + * PR #1550 gave a consumer that breaks out of `for await` a way to close the + * iterator chain (`attachStreamCancel`/`cancelStream`). That closes OUR side + * of the pipe. It does not cancel the HTTP request underneath it: without an + * abort signal that actually reaches the transport, the provider's read + * stays pending and the connection stays open on the server. This suite + * proves the server side of that gap directly — no credentials, no network + * egress — using a local HTTP fixture that reports whether ITS OWN request + * observed the client disconnecting. + * + * Provider choice matters a great deal here, and it took five attempts to + * land on one that actually isolates `baseProvider.ts` rather than some + * unrelated property of the provider or its transport. + * + * 1. `openai-compatible` and `anthropic` (direct) were ruled out first: both + * members of the OpenAI-compatible provider family carry their own + * independent "consumer-driven abort" wiring inside the provider file (a + * `consumerAbortController` tied to the async-iterator's `finally`/ + * `cancel()`), so they self-protect against this exact defect regardless + * of what `baseProvider.ts` does — testing through either would pass even + * on the unfixed build and prove nothing about the shared layer. + * + * 2. Google AI Studio's native `@google/genai` streaming path looked right + * next — no self-protection, composes its transport signal from + * `options.abortSignal` alone — but turned out to have an independent, + * pre-existing bug of its own: `googleAiStudio/client.ts`'s + * `executeNativeGemini3Stream` places the composed signal at + * `httpOptions: { signal }`, a sibling of `config`, but the SDK's + * `generateContentStreamInternal` (see + * `node_modules/@google/genai/dist/node/index.cjs`) only ever reads + * `params.config.abortSignal` / `params.config.httpOptions`. The + * top-level field is silently ignored, so no signal reaches that + * transport regardless of what `baseProvider.ts` does — the opposite + * failure mode from the self-protecting providers, but equally useless as + * a witness: a test through that path would stay RED even after the real + * fix, for a reason that has nothing to do with it. + * + * 3. Google Vertex's native *Gemini* path (`executeNativeGemini3Stream` in + * `googleVertex/client.ts`) wires the signal correctly, but its `sendStep` + * collects the entire SDK stream into a buffer and replays it to the + * consumer only once the upstream call has already settled — the code's + * own comment says the chunks are "Collected, NOT forwarded to the + * consumer here". By the time the consumer sees a first chunk to break + * on, the underlying HTTP request has already finished; there is nothing + * left in flight to cancel. A cancellation defect cannot be observed on a + * request that is already over — the suite would pass or fail on + * replay-buffer timing, not on `baseProvider.ts`. + * + * 4. Amazon Bedrock's Converse-stream path looked promising — no + * self-protection, no collect-then-replay — but its transport + * (`NodeHttp2Handler`) defaults `disableConcurrentStreams: true`, which + * makes the SDK call `session.close()` the moment response headers + * arrive. That submits a GOAWAY with `lastStreamId: 0` within + * ~15-20ms of the request landing, which a local fixture reads as + * "abandon this stream" regardless of what the test consumer has done — + * a confound that fires whether or not the real defect is present, races + * unpredictably against the consumer's own `break`, and would make the + * suite flake on both the broken and the fixed build alike. + * + * 5. Google Vertex's *Claude* path (`executeNativeAnthropicStream`, a + * distinct code path in the same `googleVertex/client.ts`, taken when the + * model name contains "claude") finally isolates the defect cleanly. + * Confirmed by reading the file: it has no `attachStreamCancel`/ + * `cancelStream` self-protection of its own, and it streams incrementally + * — `@anthropic-ai/vertex-sdk`'s `client.messages.stream()` fires its + * `text` listener per-delta as SSE events arrive, not after collecting the + * whole response (the code's own comment: "Awaiting `stream.finalMessage()` + * here would buffer the entire response before yielding anything; the + * listener pattern keeps the wire and the consumer in lockstep instead"). + * The ONLY way that transport call is ever aborted is + * `internalAbort.signal`'s "abort" event, which only the caller's own + * `options.abortSignal` (or an irrelevant turn-clock watchdog) can fire — + * exactly the seam `baseProvider.ts` composes into on teardown. It runs + * over plain HTTP/1.1 SSE, so Bedrock's HTTP/2 confound does not apply. + * Reaching it needs no real GCP project or Anthropic key: Express Mode + * (`credentials.vertex.apiKey` + `credentials.vertex.baseURL`) redirects + * the `@anthropic-ai/vertex-sdk` client at a local fixture, the same + * mechanism `continuous-test-suite-vertex-claude-characterization.ts` + * already uses to reach this exact code path. + * + * Fixture shape: an SSE server speaking real Anthropic message-stream + * framing (`event: \ndata: {...}\n\n`, matching + * `continuous-test-suite-vertex-claude-characterization.ts`'s helper and the + * wire format `@anthropic-ai/vertex-sdk` actually parses). It writes + * `message_start` and `content_block_start` once, then keeps emitting + * `content_block_delta` events on a short interval so the connection is + * still genuinely open — not already finished — at the moment the consumer + * breaks. A hard chunk cap ends the response on its own (with a normal + * `content_block_stop`/`message_delta`/`message_stop`/`res.end()`) if + * cancellation never arrives, so a failing run cannot hang the server — and + * the suite's own bounded poll (never an unbounded wait) is what actually + * fails the assertion promptly. + * + * Run: npx tsx test/continuous-test-suite-stream-teardown-abort.ts + * pnpm run test:stream-teardown-abort + */ + +import { + createServer, + type IncomingMessage, + type ServerResponse, +} from "node:http"; +import { defineSuite, assert, delay } from "./helpers/harness.js"; +import { assertDistFresh } from "./helpers/distFreshness.js"; + +assertDistFresh(); + +const { test, runSuite } = defineSuite("Stream teardown cancels transport", { + offline: true, + perTestTimeoutMs: 20_000, +}); + +const MODEL = "claude-3-5-haiku@20241022"; + +/** Env vars this suite touches — saved/restored around every test so an + * ambient key or project can never redirect the SDK somewhere real. */ +const TOUCHED_ENV_VARS = [ + "GOOGLE_CLOUD_PROJECT", + "GOOGLE_CLOUD_PROJECT_ID", + "VERTEX_PROJECT_ID", + "GOOGLE_VERTEX_PROJECT", + "GOOGLE_CLOUD_LOCATION", + "VERTEX_LOCATION", + "GOOGLE_VERTEX_LOCATION", + "GOOGLE_APPLICATION_CREDENTIALS", + "GOOGLE_VERTEX_API_KEY", + "GOOGLE_VERTEX_BASE_URL", + "GOOGLE_API_KEY", +] as const; + +function snapshotEnv(): Record { + const snapshot: Record = {}; + for (const key of TOUCHED_ENV_VARS) { + snapshot[key] = process.env[key]; + delete process.env[key]; + } + return snapshot; +} + +function restoreEnv(snapshot: Record): void { + for (const key of TOUCHED_ENV_VARS) { + const prior = snapshot[key]; + if (prior === undefined) { + delete process.env[key]; + } else { + process.env[key] = prior; + } + } +} + +/** One Anthropic message-stream SSE event, in the exact framing + * `@anthropic-ai/vertex-sdk`'s `MessageStream` parses. */ +function sse(event: string, payload: Record): string { + return `event: ${event}\ndata: ${JSON.stringify({ type: event, ...payload })}\n\n`; +} + +const CHUNK_INTERVAL_MS = 40; +// Safety cap: if the client never disconnects (e.g. a regression that hangs +// instead of leaking), the fixture still ends its own response instead of +// holding the connection — and the test process — open forever. +const MAX_CHUNKS = 60; // 60 * 40ms = 2.4s worst case per test + +/** + * A one-shot Anthropic-SSE fixture that reports whether ITS request was + * cancelled by the client before the fixture itself decided to end the + * response. + */ +function startAbortObservingServer(): { + url: Promise; + wasCancelled: () => boolean; + close: () => Promise; +} { + let cancelled = false; + let endedByFixture = false; + let activeRes: ServerResponse | undefined; + let activeTimer: ReturnType | undefined; + + const server = createServer((req: IncomingMessage, res: ServerResponse) => { + // Drain the request body; its content plays no part in this fixture. + req.resume(); + activeRes = res; + res.writeHead(200, { + "content-type": "text/event-stream", + "cache-control": "no-cache", + }); + + // The FULL message shape, not the trimmed one — `MessageStream` builds + // its snapshot from this event and pushes each content block onto + // `snapshot.content`; without `content: []` (and the sibling fields read + // alongside it) the accumulator dies before a single delta is delivered. + res.write( + sse("message_start", { + message: { + id: "msg_1", + type: "message", + role: "assistant", + model: MODEL, + content: [], + stop_reason: null, + stop_sequence: null, + usage: { input_tokens: 5, output_tokens: 0 }, + }, + }), + ); + res.write( + sse("content_block_start", { + index: 0, + content_block: { type: "text", text: "" }, + }), + ); + + let n = 0; + const writeNext = (): void => { + n++; + if (n >= MAX_CHUNKS) { + if (activeTimer) { + clearInterval(activeTimer); + } + endedByFixture = true; + res.write( + sse("content_block_delta", { + index: 0, + delta: { type: "text_delta", text: `chunk-${n} ` }, + }), + ); + res.write(sse("content_block_stop", { index: 0 })); + res.write( + sse("message_delta", { + delta: { stop_reason: "end_turn" }, + usage: { output_tokens: n }, + }), + ); + res.write(sse("message_stop", {})); + res.end(); + return; + } + res.write( + sse("content_block_delta", { + index: 0, + delta: { type: "text_delta", text: `chunk-${n} ` }, + }), + ); + }; + writeNext(); + activeTimer = setInterval(writeNext, CHUNK_INTERVAL_MS); + + // The connection closing before the fixture itself called res.end() + // means the CLIENT went away — exactly the observation this suite + // exists to make. `res.writableEnded` is set synchronously by + // `res.end()`, so ordering against `endedByFixture` (set just before + // that call) is redundant but kept for clarity under either signal. + res.on("close", () => { + if (activeTimer) { + clearInterval(activeTimer); + } + if (!endedByFixture && !res.writableEnded) { + cancelled = true; + } + }); + }); + + const url = new Promise((resolve) => { + server.listen(0, "127.0.0.1", () => { + const address = server.address(); + const port = typeof address === "object" && address ? address.port : 0; + resolve(`http://127.0.0.1:${port}`); + }); + }); + + return { + url, + wasCancelled: () => cancelled, + close: () => + new Promise((resolve) => { + // Best-effort teardown: force-end whatever is still open so a failed + // assertion can never leave the fixture's own socket dangling. + if (activeTimer) { + clearInterval(activeTimer); + } + if (activeRes && !activeRes.writableEnded) { + activeRes.end(); + } + server.close(() => resolve()); + // server.close() waits for existing connections to finish; force it + // so a stuck client socket can't stall the suite itself. + setTimeout(() => resolve(), 500).unref(); + }), + }; +} + +/** Bounded poll — never an unbounded wait. */ +async function pollUntil( + check: () => boolean, + budgetMs: number, + intervalMs = 20, +): Promise { + const start = Date.now(); + for (;;) { + if (check()) { + return true; + } + if (Date.now() - start >= budgetMs) { + return check(); + } + await delay(intervalMs); + } +} + +void runSuite(async () => { + const { NeuroLink, ProviderRegistry } = await import("../dist/index.js"); + await ProviderRegistry.registerAllProviders(); + + await test("breaking out of the stream after one chunk cancels the server's in-flight request", async () => { + const envSnapshot = snapshotEnv(); + const fixture = startAbortObservingServer(); + + try { + const nl = new NeuroLink({ conversationMemory: { enabled: false } }); + const result = await nl.stream({ + provider: "vertex", + model: MODEL, + input: { text: "hi" }, + maxTokens: 32, + maxSteps: 1, + disableTools: true, + disableInternalFallback: true, + credentials: { + vertex: { + apiKey: "express-key", + baseURL: await fixture.url, + }, + }, + } as Parameters["stream"]>[0]); + + let gotFirstChunk = false; + for await (const chunk of result.stream) { + if ("content" in chunk && typeof chunk.content === "string") { + gotFirstChunk = true; + } + break; + } + assert( + gotFirstChunk, + "never received a first chunk — nothing to abandon", + ); + + const observed = await pollUntil(fixture.wasCancelled, 1_500); + assert( + observed, + "server did not observe the client's request being cancelled after the consumer abandoned the stream", + ); + } finally { + await fixture.close(); + restoreEnv(envSnapshot); + } + }); +}); diff --git a/test/continuous-test-suite-video-options-allowlist.ts b/test/continuous-test-suite-video-options-allowlist.ts new file mode 100644 index 000000000..1abeecd48 --- /dev/null +++ b/test/continuous-test-suite-video-options-allowlist.ts @@ -0,0 +1,205 @@ +#!/usr/bin/env tsx +import "dotenv/config"; + +/** + * Continuous Test Suite — `videoOptions` reaches the provider (allowlist fix). + * + * `videoOptions` is a documented public `generate()` option (`frames`, + * `quality`, `format`, `transcribeAudio` — see `GenerateOptions.videoOptions` + * in `src/lib/types/generate.ts`), but `buildGenerateTextOptions` in + * `src/lib/neurolink.ts` converts the caller's public `GenerateOptions` into + * the internal `TextGenerationOptions` through an explicit field-by-field + * allowlist, and `videoOptions` was never named in it — the same bug class + * `c011b0405` fixed for `thinkingConfig` (that commit's message records how + * `disableInternalFallback` was lost the same way once before). It fails + * quietly: no error, no warning, the option is simply dropped before any + * provider sees it. + * + * Google AI Studio is the target here because it is the ONE path unaffected + * by anything else: it overrides `generate()`/`executeStream()` and calls + * `preprocessNativeFileInput` -> `processUnifiedFilesArray` directly on the + * internal `TextGenerationOptions` object, with no second, independent + * allowlist in between (unlike the generic AI-SDK pipeline's + * `MessageBuilder.buildMessages`, which builds its own separate + * `multimodalOptions` literal — a sibling gap, out of scope for this fix). + * Proving the option survives on this path isolates exactly the bug named + * above. + * + * ALL-DIST module graph (rule 15): `NeuroLink` comes from `../dist/index.js`, + * the same shipped surface every caller uses. No rule-15 exception needed. + * + * The instrument is a local, loopback-only stand-in for the + * `@google/genai` REST wire (SSE `generateContentStream` framing), pointed at + * via the public `credentials.googleAiStudio.baseURL` override — no live + * credentials, no network egress. It records the request body `generate()` + * actually sent, so the test can count `inlineData` image parts on the wire + * rather than trust an in-process return value. + * + * A real 6-second video fixture (synthesised with ffmpeg at test time, see + * `helpers/mediaFixtures.ts`) is attached with two different + * `videoOptions.frames` values. `VideoProcessor.extractKeyframes` spreads a + * caller-supplied frame budget evenly across the clip, so `frames: 2` and + * `frames: 4` must produce a DIFFERENT number of image parts on the wire, if + * and only if `videoOptions` actually reached the processor. The clip's + * duration (6s, <=10s tier) is deliberately chosen so that if `videoOptions` + * is dropped, both calls fall back to the SAME uncontrolled tier default (one + * frame per second = 6 frames) — the exact failure this suite exists to + * catch, not two counts that merely happen to differ for an unrelated reason. + * + * Run: npx tsx test/continuous-test-suite-video-options-allowlist.ts + * pnpm run test:video-options-allowlist + */ + +import { createServer, type Server } from "node:http"; +import { defineSuite, assert, tempDir, Skip } from "./helpers/harness.js"; +import { assertDistFresh } from "./helpers/distFreshness.js"; +import { hasFfmpeg, makeVideoFile } from "./helpers/mediaFixtures.js"; + +assertDistFresh(); + +const { test, runSuite } = defineSuite("Video options allowlist", { + offline: true, +}); + +const { NeuroLink } = await import("../dist/index.js"); + +const MODEL = "gemini-2.0-flash"; +const CLIP_SECONDS = 6; +const FRAMES_FIRST_CALL = 2; +const FRAMES_SECOND_CALL = 4; + +type CapturedCall = { body: Record }; + +type StandIn = { + calls: CapturedCall[]; + port: number; + close: () => Promise; +}; + +/** One streamed candidate chunk in the SSE framing `generateContentStream` expects. */ +function sseTextReply(text: string): string { + const payload = { + candidates: [ + { + content: { parts: [{ text }], role: "model" }, + finishReason: "STOP", + index: 0, + }, + ], + usageMetadata: { + promptTokenCount: 5, + candidatesTokenCount: 4, + totalTokenCount: 9, + }, + }; + return `data: ${JSON.stringify(payload)}\r\n\r\n`; +} + +/** Local stand-in for the Google AI Studio `generateContentStream` endpoint. */ +async function startStandIn(): Promise { + const calls: CapturedCall[] = []; + const server: Server = createServer((req, res) => { + const chunks: Buffer[] = []; + req.on("data", (c: Buffer) => chunks.push(c)); + req.on("end", () => { + let body: Record = {}; + try { + body = JSON.parse(Buffer.concat(chunks).toString("utf8") || "{}"); + } catch { + // Malformed JSON is still a captured call; fall through so the + // caller's assertion sees the (empty) body rather than a hang. + } + calls.push({ body }); + res.writeHead(200, { "content-type": "text/event-stream" }); + res.write(sseTextReply("ok")); + res.end(); + }); + }); + await new Promise((resolve) => server.listen(0, "127.0.0.1", resolve)); + const address = server.address(); + return { + calls, + port: typeof address === "object" && address ? address.port : 0, + close: () => + new Promise((resolve) => { + server.close(() => resolve()); + }), + }; +} + +/** Count of image `inlineData` parts in the LAST user turn a call sent. */ +function imagePartCount(call: CapturedCall | undefined): number { + const contents = (call?.body?.contents ?? []) as Array<{ + role?: string; + parts?: Array<{ inlineData?: { mimeType?: string } }>; + }>; + const userTurns = contents.filter((c) => c.role === "user"); + const lastUserTurn = userTurns[userTurns.length - 1]; + return (lastUserTurn?.parts ?? []).filter((p) => + Boolean(p.inlineData?.mimeType?.startsWith("image/")), + ).length; +} + +function credentialsFor(port: number) { + return { + googleAiStudio: { + apiKey: "test-key", + baseURL: `http://127.0.0.1:${port}`, + }, + }; +} + +await test("videoOptions.frames changes the number of keyframes actually sent to the provider", async () => { + if (!(await hasFfmpeg())) { + throw new Skip("ffmpeg not available — cannot synthesise a video fixture"); + } + const dir = tempDir("video-options-allowlist-"); + const clip = await makeVideoFile(dir, "clip.mp4", CLIP_SECONDS); + + const server = await startStandIn(); + try { + const nl = new NeuroLink(); + await nl.generate({ + input: { text: "describe this clip", files: [clip] }, + videoOptions: { frames: FRAMES_FIRST_CALL }, + provider: "google-ai", + model: MODEL, + disableTools: true, + disableInternalFallback: true, + credentials: credentialsFor(server.port), + }); + await nl.generate({ + input: { text: "describe this clip", files: [clip] }, + videoOptions: { frames: FRAMES_SECOND_CALL }, + provider: "google-ai", + model: MODEL, + disableTools: true, + disableInternalFallback: true, + credentials: credentialsFor(server.port), + }); + + assert( + server.calls.length === 2, + "expected exactly two requests to reach the provider stand-in", + ); + const firstCount = imagePartCount(server.calls[0]); + const secondCount = imagePartCount(server.calls[1]); + + assert( + firstCount !== secondCount, + `videoOptions.frames had no observable effect on the wire — both requests carried the same image-part count (${firstCount})`, + ); + assert( + firstCount === FRAMES_FIRST_CALL, + `expected the frames:${FRAMES_FIRST_CALL} request to carry ${FRAMES_FIRST_CALL} keyframes, carried ${firstCount}`, + ); + assert( + secondCount === FRAMES_SECOND_CALL, + `expected the frames:${FRAMES_SECOND_CALL} request to carry ${FRAMES_SECOND_CALL} keyframes, carried ${secondCount}`, + ); + } finally { + await server.close(); + } +}); + +await runSuite(); diff --git a/test/continuous-test-suite-websocket-agent-handler.ts b/test/continuous-test-suite-websocket-agent-handler.ts new file mode 100644 index 000000000..e6035d574 --- /dev/null +++ b/test/continuous-test-suite-websocket-agent-handler.ts @@ -0,0 +1,990 @@ +#!/usr/bin/env tsx +import "dotenv/config"; + +/** + * Continuous Test Suite — createAgentWebSocketHandler real routes + * + * `createAgentWebSocketHandler` is exported from the public entry, documented + * in docs/guides/server-adapters/websocket.md's Quick Start, and ships in the + * published tarball — but its three routes (`generate`, `stream`, `tool_call`) + * were literal stubs that never touched the injected NeuroLink instance: + * `generate` echoed the caller's own prompt back as if it were a model + * response. This suite drives the handler exactly as a real client would + * (JSON frames through `onMessage`, inspecting the frames written back to a + * mock socket) and asserts the responses are the model's / tool's real + * output, not a stub echo. + * + * ALL-DIST module graph (rule 15): `NeuroLink`, `ProviderRegistry`, and + * `createAgentWebSocketHandler` all come from `../dist/index.js`. No + * src/lib import anywhere in this file. + * + * No external API keys — points the "openai-compatible" provider at a local + * HTTP fixture via OPENAI_COMPATIBLE_BASE_URL, modeled on the real OpenAI + * chat/completions wire format (plain JSON for generate(), SSE for + * stream()). The tool_call route needs no network fixture at all — it + * executes a tool registered directly on the NeuroLink instance. + * + * Run: npx tsx test/continuous-test-suite-websocket-agent-handler.ts + * pnpm run test:websocket-agent-handler + */ + +import { createServer } from "node:http"; +import { defineSuite, assert } from "./helpers/harness.js"; + +const { test, runSuite, section } = defineSuite("WebSocket Agent Handler", { + offline: true, +}); + +/** Env vars this suite mutates — saved/restored around every test so ambient + * dev-machine values can't leak in or out. */ +const TOUCHED_ENV_VARS = [ + "OPENAI_COMPATIBLE_BASE_URL", + "OPENAI_COMPATIBLE_API_KEY", +] as const; + +function snapshotEnv(): Record { + const snapshot: Record = {}; + for (const key of TOUCHED_ENV_VARS) { + snapshot[key] = process.env[key]; + } + return snapshot; +} + +function restoreEnv(snapshot: Record): void { + for (const key of TOUCHED_ENV_VARS) { + const prior = snapshot[key]; + if (prior === undefined) { + delete process.env[key]; + } else { + process.env[key] = prior; + } + } +} + +function sseChunk(text: string): string { + return `data: ${JSON.stringify({ + choices: [{ delta: { content: text }, finish_reason: null }], + })}\n\n`; +} + +/** A mock socket that records every outbound frame as parsed JSON. */ +function mockSocket(): { + socket: { send: (data: string) => void }; + sent: Array>; +} { + const sent: Array> = []; + return { + socket: { + send: (data: string) => { + sent.push(JSON.parse(data) as Record); + }, + }, + sent, + }; +} + +function listen(server: ReturnType): Promise { + return new Promise((resolve) => { + server.listen(0, () => { + const address = server.address(); + resolve(typeof address === "object" && address ? address.port : 0); + }); + }); +} + +void runSuite(async () => { + const { NeuroLink, ProviderRegistry, createAgentWebSocketHandler } = + await import("../dist/index.js"); + await ProviderRegistry.registerAllProviders(); + + function nl() { + return new NeuroLink({ conversationMemory: { enabled: false } }); + } + + section("generate route calls the model instead of echoing the prompt"); + + await test("generate returns the fixture model output, not the request prompt", async () => { + const envSnapshot = snapshotEnv(); + const server = createServer((_req, res) => { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + choices: [ + { + index: 0, + message: { role: "assistant", content: "FIXTURE_REPLY_7f3a" }, + finish_reason: "stop", + }, + ], + }), + ); + }); + const port = await listen(server); + + try { + process.env.OPENAI_COMPATIBLE_BASE_URL = `http://127.0.0.1:${port}`; + process.env.OPENAI_COMPATIBLE_API_KEY = "test-key"; + + const handler = createAgentWebSocketHandler(nl()); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-generate", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + const prompt = "Describe the mission briefing for team Aurora"; + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + type: "generate", + payload: { + prompt, + options: { + provider: "openai-compatible", + model: "gpt-4o-mini", + maxSteps: 1, + }, + }, + }), + timestamp: Date.now(), + }); + + assert( + sent.length === 1, + "expected exactly one outbound frame for generate", + ); + const response = sent[0]; + assert( + response?.type === "response", + "generate frame had the wrong message type", + ); + const content = response?.content; + assert( + typeof content === "string" && content.includes("FIXTURE_REPLY"), + "generate response did not carry the fixture model output", + ); + assert( + typeof content === "string" && !content.includes(prompt), + "generate response echoed the request prompt instead of the model output", + ); + } finally { + server.close(); + restoreEnv(envSnapshot); + } + }); + + section( + "generate route surfaces provider failures as errors, not fake success", + ); + + await test("generate propagates a failed model call as an error frame", async () => { + const envSnapshot = snapshotEnv(); + // 400 (not 429/500/503) so withProviderRetry does not retry — one fast, + // deterministic attempt. + const server = createServer((_req, res) => { + res.writeHead(400, { "content-type": "application/json" }); + res.end( + JSON.stringify({ error: { message: "synthetic failure fixture" } }), + ); + }); + const port = await listen(server); + + try { + process.env.OPENAI_COMPATIBLE_BASE_URL = `http://127.0.0.1:${port}`; + process.env.OPENAI_COMPATIBLE_API_KEY = "test-key"; + + const handler = createAgentWebSocketHandler(nl()); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-generate-error", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + type: "generate", + payload: { + prompt: "This call should fail", + options: { + provider: "openai-compatible", + model: "gpt-4o-mini", + maxSteps: 1, + }, + }, + }), + timestamp: Date.now(), + }); + + assert( + sent.length === 1, + "expected exactly one outbound frame for a failed generate", + ); + const response = sent[0]; + assert( + response?.type === "error", + "a failed generate call did not surface as an error frame", + ); + assert( + typeof response?.error === "string" && response.error.length > 0, + "error frame carried no error message", + ); + } finally { + server.close(); + restoreEnv(envSnapshot); + } + }); + + section("stream route emits real chunks instead of a single stub frame"); + + await test("stream forwards fixture chunks and a completion frame", async () => { + const envSnapshot = snapshotEnv(); + const server = createServer((_req, res) => { + res.writeHead(200, { "content-type": "text/event-stream" }); + res.write(sseChunk("STREAM_")); + res.write(sseChunk("FIXTURE_")); + res.write(sseChunk("9c2d")); + res.write("data: [DONE]\n\n"); + res.end(); + }); + const port = await listen(server); + + try { + process.env.OPENAI_COMPATIBLE_BASE_URL = `http://127.0.0.1:${port}`; + process.env.OPENAI_COMPATIBLE_API_KEY = "test-key"; + + const handler = createAgentWebSocketHandler(nl()); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-stream", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + const prompt = "Summarize the reconnaissance log"; + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + type: "stream", + payload: { + prompt, + options: { + provider: "openai-compatible", + model: "gpt-4o-mini", + maxSteps: 1, + }, + }, + }), + timestamp: Date.now(), + }); + + const chunkFrames = sent.filter((frame) => frame.type === "chunk"); + assert( + chunkFrames.length > 0, + "no chunk frames were emitted during the stream route", + ); + const assembled = chunkFrames + .map((frame) => + typeof frame.content === "string" ? frame.content : "", + ) + .join(""); + assert( + assembled.includes("STREAM_FIXTURE_9c2d"), + "assembled stream content did not match the fixture output", + ); + assert( + !assembled.includes(prompt), + "streamed content echoed the request prompt instead of the model output", + ); + const lastFrame = sent[sent.length - 1]; + assert( + lastFrame?.type === "stream_complete", + "stream route did not end with a completion frame", + ); + } finally { + server.close(); + restoreEnv(envSnapshot); + } + }); + + section("tool_call route executes the real tool instead of returning null"); + + await test("tool_call returns the tool's actual result", async () => { + const instance = nl(); + instance.registerTool("get_secret_code", { + name: "get_secret_code", + description: "Returns a fixed test code", + inputSchema: { type: "object", properties: {} }, + execute: async () => ({ code: "ALPHA-99" }), + }); + + const handler = createAgentWebSocketHandler(instance); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-tool-call", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + type: "tool_call", + payload: { toolName: "get_secret_code", args: {} }, + }), + timestamp: Date.now(), + }); + + assert( + sent.length === 1, + "expected exactly one outbound frame for tool_call", + ); + const response = sent[0]; + assert( + response?.type === "tool_result", + "tool_call frame had the wrong message type", + ); + // executeTool()'s public contract wraps the tool's return value in an + // envelope ({ success, data, usage, metadata }) — see + // NeuroLink.executeTool in src/lib/neurolink.ts. The route forwards that + // envelope verbatim (matching docs/guides/server-adapters/websocket.md), + // so the real payload lives at `result.data`. + const result = response?.result as + | { success?: unknown; data?: { code?: unknown } } + | undefined; + assert( + result?.success === true, + "tool_call result envelope did not report success", + ); + assert( + result?.data?.code === "ALPHA-99", + "tool_call did not surface the tool's real return value", + ); + }); + + section("an endpoint that requires auth refuses an unauthenticated socket"); + + // Raised in review on this PR. WebSocketConfig has always declared + // `auth: { strategy, required }`, and handleConnection has always accepted + // and stored an AuthenticatedUser — but nothing consulted either, so + // `required: true` produced exactly the same open socket as `required: + // false`. That was inert while the agent routes returned canned values; it + // stops being inert once tool_call reaches executeTool, because then any + // client that can open the socket can run any registered tool with + // arbitrary arguments. + await test("a connection with auth.required and no authenticated user is refused", async () => { + const { WebSocketConnectionManager } = await import("../dist/index.js"); + const manager = new WebSocketConnectionManager({ + auth: { strategy: "bearer", required: true }, + } as ConstructorParameters[0]); + + let refused = false; + try { + await manager.handleConnection({ send: () => {} }, "/ws/agent"); + } catch { + refused = true; + } + assert( + refused, + "an endpoint requiring auth accepted a socket carrying no user", + ); + }); + + await test("a connection with auth.required and an authenticated user is accepted", async () => { + const { WebSocketConnectionManager } = await import("../dist/index.js"); + const manager = new WebSocketConnectionManager({ + auth: { strategy: "bearer", required: true }, + } as ConstructorParameters[0]); + + // The negative case alone would also pass if the gate refused + // everything, so pin that a real user still gets through. + const connection = await manager.handleConnection( + { send: () => {} }, + "/ws/agent", + { id: "user-1" } as Parameters[2], + ); + assert( + typeof connection?.id === "string" && connection.id.length > 0, + "an authenticated socket was not accepted", + ); + await manager.handleClose(connection.id, 1000, "test cleanup"); + }); + + await test("an endpoint with auth unset still accepts a socket", async () => { + const { WebSocketConnectionManager } = await import("../dist/index.js"); + // Default config is `required: false`; the gate must not change the + // behaviour every existing deployment relies on. + const manager = new WebSocketConnectionManager(); + const connection = await manager.handleConnection( + { send: () => {} }, + "/ws/agent", + ); + assert( + typeof connection?.id === "string" && connection.id.length > 0, + "an endpoint that does not require auth refused an anonymous socket", + ); + await manager.handleClose(connection.id, 1000, "test cleanup"); + }); + + section( + "an endpoint with configured roles/permissions enforces them, not just auth.required", + ); + + // docs/guides/server-adapters/websocket.md's own "Role-Based Access + // Control" example configures `roles: ["admin", "superuser"]` on a + // `/ws/admin` path, and ServerAuthConfig.roles/permissions plus + // AuthenticatedUser.roles/permissions have existed on these types since + // this manager was written — but neither was ever consulted, so any + // authenticated user got in regardless of role. Semantics mirror the SDK's + // own hasAnyRole/hasAllPermissions (src/lib/auth/authContext.ts): roles + // are an any-of allowlist, permissions must all be present. + await test("a connection holding none of the configured roles is refused", async () => { + const { WebSocketConnectionManager } = await import("../dist/index.js"); + const manager = new WebSocketConnectionManager({ + auth: { + strategy: "bearer", + required: true, + roles: ["admin", "superuser"], + }, + } as ConstructorParameters[0]); + + let refused = false; + try { + await manager.handleConnection({ send: () => {} }, "/ws/admin", { + id: "user-1", + roles: ["viewer"], + } as Parameters[2]); + } catch { + refused = true; + } + assert( + refused, + "a connection with none of the configured roles (any-of) was accepted", + ); + }); + + await test("a connection holding any one of the configured roles is accepted", async () => { + const { WebSocketConnectionManager } = await import("../dist/index.js"); + const manager = new WebSocketConnectionManager({ + auth: { + strategy: "bearer", + required: true, + roles: ["admin", "superuser"], + }, + } as ConstructorParameters[0]); + + const connection = await manager.handleConnection( + { send: () => {} }, + "/ws/admin", + { id: "user-2", roles: ["superuser"] } as Parameters< + typeof manager.handleConnection + >[2], + ); + assert( + typeof connection?.id === "string" && connection.id.length > 0, + "a connection holding one of the configured roles (any-of) was refused", + ); + await manager.handleClose(connection.id, 1000, "test cleanup"); + }); + + await test("a connection missing even one configured permission is refused", async () => { + const { WebSocketConnectionManager } = await import("../dist/index.js"); + const manager = new WebSocketConnectionManager({ + auth: { + strategy: "bearer", + required: true, + permissions: ["read", "write"], + }, + } as ConstructorParameters[0]); + + let refused = false; + try { + await manager.handleConnection({ send: () => {} }, "/ws/admin", { + id: "user-3", + permissions: ["read"], + } as Parameters[2]); + } catch { + refused = true; + } + assert( + refused, + "a connection missing one of the configured permissions (all-of) was accepted", + ); + }); + + await test("a connection holding every configured permission is accepted", async () => { + const { WebSocketConnectionManager } = await import("../dist/index.js"); + const manager = new WebSocketConnectionManager({ + auth: { + strategy: "bearer", + required: true, + permissions: ["read", "write"], + }, + } as ConstructorParameters[0]); + + const connection = await manager.handleConnection( + { send: () => {} }, + "/ws/admin", + { id: "user-4", permissions: ["read", "write"] } as Parameters< + typeof manager.handleConnection + >[2], + ); + assert( + typeof connection?.id === "string" && connection.id.length > 0, + "a connection holding every configured permission was refused", + ); + await manager.handleClose(connection.id, 1000, "test cleanup"); + }); + + section( + "generate/stream routes strip unsafe options before calling the SDK (SSRF / credential-override guard)", + ); + + // The raw client payload used to spread straight into + // neurolink.generate()/.stream() with no validation, so a client's own + // `options.credentials..baseURL` could redirect the server's own + // outbound HTTP call to an arbitrary host (or swap in the client's own API + // key). WebSocketAgentOptionsSchema is an allowlist with no `credentials` + // field, so a value there must never reach the provider call. Proved with + // two local fixtures: a "trusted" one at the env-configured base URL, and + // an "attacker" one only reachable via the payload's own + // credentials.baseURL — the attacker fixture must see zero requests. + await test("generate ignores a client-supplied credentials/baseURL override", async () => { + const envSnapshot = snapshotEnv(); + let attackerHits = 0; + const trustedServer = createServer((_req, res) => { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + choices: [ + { + index: 0, + message: { role: "assistant", content: "REAL_FIXTURE_5b1c" }, + finish_reason: "stop", + }, + ], + }), + ); + }); + const attackerServer = createServer((_req, res) => { + attackerHits += 1; + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + choices: [ + { + index: 0, + message: { role: "assistant", content: "HIJACKED_RESPONSE" }, + finish_reason: "stop", + }, + ], + }), + ); + }); + const trustedPort = await listen(trustedServer); + const attackerPort = await listen(attackerServer); + + try { + process.env.OPENAI_COMPATIBLE_BASE_URL = `http://127.0.0.1:${trustedPort}`; + process.env.OPENAI_COMPATIBLE_API_KEY = "test-key"; + + const handler = createAgentWebSocketHandler(nl()); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-generate-ssrf-guard", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + type: "generate", + payload: { + prompt: "Attempt an override", + options: { + provider: "openai-compatible", + model: "gpt-4o-mini", + maxSteps: 1, + // A client should never be able to steer the outbound call + // with fields like these. + credentials: { + openaiCompatible: { + apiKey: "attacker-key", + baseURL: `http://127.0.0.1:${attackerPort}`, + }, + }, + }, + }, + }), + timestamp: Date.now(), + }); + + assert(sent.length === 1, "expected exactly one outbound frame"); + const response = sent[0]; + assert( + response?.type === "response", + "generate call did not complete successfully against the configured fixture", + ); + assert( + typeof response?.content === "string" && + response.content.includes("REAL_FIXTURE_5b1c"), + "generate response did not carry the configured fixture's output", + ); + assert( + attackerHits === 0, + "a client-supplied credentials.baseURL reached an outbound HTTP request — the SSRF override was not stripped", + ); + } finally { + trustedServer.close(); + attackerServer.close(); + restoreEnv(envSnapshot); + } + }); + + await test("stream ignores a client-supplied credentials/baseURL override", async () => { + const envSnapshot = snapshotEnv(); + let attackerHits = 0; + const trustedServer = createServer((_req, res) => { + res.writeHead(200, { "content-type": "text/event-stream" }); + res.write(sseChunk("REAL_")); + res.write(sseChunk("FIXTURE_")); + res.write(sseChunk("2ab1")); + res.write("data: [DONE]\n\n"); + res.end(); + }); + const attackerServer = createServer((_req, res) => { + attackerHits += 1; + res.writeHead(200, { "content-type": "text/event-stream" }); + res.write(sseChunk("HIJACKED")); + res.write("data: [DONE]\n\n"); + res.end(); + }); + const trustedPort = await listen(trustedServer); + const attackerPort = await listen(attackerServer); + + try { + process.env.OPENAI_COMPATIBLE_BASE_URL = `http://127.0.0.1:${trustedPort}`; + process.env.OPENAI_COMPATIBLE_API_KEY = "test-key"; + + const handler = createAgentWebSocketHandler(nl()); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-stream-ssrf-guard", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + type: "stream", + payload: { + prompt: "Attempt an override", + options: { + provider: "openai-compatible", + model: "gpt-4o-mini", + maxSteps: 1, + credentials: { + openaiCompatible: { + apiKey: "attacker-key", + baseURL: `http://127.0.0.1:${attackerPort}`, + }, + }, + }, + }, + }), + timestamp: Date.now(), + }); + + const chunkFrames = sent.filter((frame) => frame.type === "chunk"); + const assembled = chunkFrames + .map((frame) => + typeof frame.content === "string" ? frame.content : "", + ) + .join(""); + assert( + assembled.includes("REAL_FIXTURE_2ab1"), + "stream response did not carry the configured fixture's output", + ); + const lastFrame = sent[sent.length - 1]; + assert( + lastFrame?.type === "stream_complete", + "stream call did not complete successfully against the configured fixture", + ); + assert( + attackerHits === 0, + "a client-supplied credentials.baseURL reached an outbound HTTP request — the SSRF override was not stripped", + ); + } finally { + trustedServer.close(); + attackerServer.close(); + restoreEnv(envSnapshot); + } + }); + + await test("generate rejects a payload without a prompt instead of passing it through", async () => { + const handler = createAgentWebSocketHandler(nl()); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-generate-invalid-payload", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + type: "generate", + payload: { options: { provider: "openai-compatible" } }, + }), + timestamp: Date.now(), + }); + + assert( + sent.length === 1, + "expected exactly one outbound frame for an invalid generate payload", + ); + const response = sent[0]; + assert( + response?.type === "error", + "a generate payload missing a prompt was not rejected as invalid", + ); + }); + + section( + "generate/stream/tool_call responses echo a client-supplied request id", + ); + + // Every response frame these routes emit used to be id-less, so two + // overlapping in-flight messages on one connection couldn't be matched + // back to the request that caused them. A client that opts in by sending + // its own `id` must get it back on the matching response; a client that + // never sends one must see no change at all (last test below). + await test("generate echoes the request's id back on the response frame", async () => { + const envSnapshot = snapshotEnv(); + const server = createServer((_req, res) => { + res.writeHead(200, { "content-type": "application/json" }); + res.end( + JSON.stringify({ + choices: [ + { + index: 0, + message: { role: "assistant", content: "FIXTURE_CORR_1" }, + finish_reason: "stop", + }, + ], + }), + ); + }); + const port = await listen(server); + + try { + process.env.OPENAI_COMPATIBLE_BASE_URL = `http://127.0.0.1:${port}`; + process.env.OPENAI_COMPATIBLE_API_KEY = "test-key"; + + const handler = createAgentWebSocketHandler(nl()); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-generate-correlation", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + id: "req-42", + type: "generate", + payload: { + prompt: "Correlate me", + options: { provider: "openai-compatible", maxSteps: 1 }, + }, + }), + timestamp: Date.now(), + }); + + assert(sent.length === 1, "expected exactly one outbound frame"); + const response = sent[0]; + assert( + response !== undefined && response.id === "req-42", + "generate response did not carry back the request's correlation id", + ); + } finally { + server.close(); + restoreEnv(envSnapshot); + } + }); + + await test("stream echoes the request's id on every frame it sends", async () => { + const envSnapshot = snapshotEnv(); + const server = createServer((_req, res) => { + res.writeHead(200, { "content-type": "text/event-stream" }); + res.write(sseChunk("CORR_CHUNK")); + res.write("data: [DONE]\n\n"); + res.end(); + }); + const port = await listen(server); + + try { + process.env.OPENAI_COMPATIBLE_BASE_URL = `http://127.0.0.1:${port}`; + process.env.OPENAI_COMPATIBLE_API_KEY = "test-key"; + + const handler = createAgentWebSocketHandler(nl()); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-stream-correlation", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + id: 7, + type: "stream", + payload: { + prompt: "Correlate the stream", + options: { provider: "openai-compatible", maxSteps: 1 }, + }, + }), + timestamp: Date.now(), + }); + + assert( + sent.length >= 2, + "expected at least a stream_start and a stream_complete frame", + ); + assert( + sent.every((frame) => frame.id === 7), + "not every stream frame carried the request's correlation id", + ); + } finally { + server.close(); + restoreEnv(envSnapshot); + } + }); + + await test("tool_call echoes the request's id back on the result frame", async () => { + const instance = nl(); + instance.registerTool("get_secret_code", { + name: "get_secret_code", + description: "Returns a fixed test code", + inputSchema: { type: "object", properties: {} }, + execute: async () => ({ code: "ALPHA-99" }), + }); + + const handler = createAgentWebSocketHandler(instance); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-tool-call-correlation", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + id: "req-tool-1", + type: "tool_call", + payload: { toolName: "get_secret_code", args: {} }, + }), + timestamp: Date.now(), + }); + + assert(sent.length === 1, "expected exactly one outbound frame"); + const response = sent[0]; + assert( + response !== undefined && response.id === "req-tool-1", + "tool_call response did not carry back the request's correlation id", + ); + }); + + await test("an error frame still carries the request's id when the request had one", async () => { + const handler = createAgentWebSocketHandler(nl()); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-error-correlation", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + id: "req-err-1", + type: "generate", + payload: { options: {} }, // missing prompt -> validation error + }), + timestamp: Date.now(), + }); + + assert(sent.length === 1, "expected exactly one outbound frame"); + const response = sent[0]; + assert(response?.type === "error", "expected an error frame"); + assert( + response !== undefined && response.id === "req-err-1", + "error frame did not carry back the request's correlation id", + ); + }); + + await test("a request sent without an id gets a response with no id field (backward compatible)", async () => { + const instance = nl(); + instance.registerTool("get_secret_code", { + name: "get_secret_code", + description: "Returns a fixed test code", + inputSchema: { type: "object", properties: {} }, + execute: async () => ({ code: "ALPHA-99" }), + }); + + const handler = createAgentWebSocketHandler(instance); + const { socket, sent } = mockSocket(); + const connection = { + id: "test-tool-call-no-id", + socket, + metadata: {}, + createdAt: Date.now(), + lastActivity: Date.now(), + }; + + await handler.onMessage?.(connection, { + type: "text", + data: JSON.stringify({ + type: "tool_call", + payload: { toolName: "get_secret_code", args: {} }, + }), + timestamp: Date.now(), + }); + + assert(sent.length === 1, "expected exactly one outbound frame"); + const response = sent[0]; + assert( + response !== undefined && !("id" in response), + "a client that never sent an id should not see one appear in the response", + ); + }); +});