From 257e841faaff28add50cf4a016dee2543cc1ff5e Mon Sep 17 00:00:00 2001 From: Sachin Sharma Date: Sun, 27 Sep 2026 00:58:02 +0530 Subject: [PATCH] fix(sdk): stop four public surfaces accepting input and dropping it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Four independent defects found by auditing the provider campaign against origin/release. Each one accepts something a caller can legitimately pass, reports success, and discards it — so no error is raised and nothing in the response says the request was not honoured. WebSocket agent routes never called the SDK. createAgentWebSocketHandler is exported from the public entry, documented in the server-adapters guide, and ships in the published tarball. Its three routes were TODO(#1576) stubs: generate returned the caller's own prompt text echoed back as if it were a model answer, stream returned a single stream_start frame and never streamed, and tool_call always returned a null result. The factory already received a NeuroLink instance and ignored it, typed `_neurolink: unknown`. The routes now call generate/stream/executeTool on that instance, and a failed call surfaces as an error frame rather than a fake success. The response shapes are not new: the guide already documented {type:"response", content}, {type:"chunk", content}, {type:"stream_complete"} and {type:"tool_result", toolName, result}, and declares that union as a type — the stubs were what diverged from the documented protocol. The three implemented TODO(#1576) markers are removed; the other 12 elsewhere in the tree are untouched. videoOptions never reached a provider. It is a documented public generate option, with JSDoc for frames and quality, but buildGenerateTextOptions converts the caller's options through an explicit field-by-field allowlist and videoOptions was not in it. The allowlist carries csvOptions and pdfOptions, its two siblings, which is what makes this an omission rather than a design decision. Same bug class as c011b0405, which added thinkingConfig to that same list. stream({ schema }) could not return structured output. The streaming path never sent response_format under any condition — unlike generate(), where the gap is only tools-dependent — and StreamResult had no structuredData field at all. The wire field is now sent when a schema is present and tools are not suppressing it, and the parsed object is delivered on metadata.structuredData after the stream drains. metadata rather than a new top-level field because stream wrappers spread the result object and would snapshot a top-level getter before the loop resolves; a plain field rather than a promise because a middleware short-circuit returns its own stream and the loop never runs, and an unresolved promise would hang every reader whereas an absent field is simply absent. Abandoning a stream leaked the request. #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 stayed pending and the connection stayed open. The call now owns an AbortController composed with the caller's own signal, so either source can end the request and a caller's existing signal keeps working unchanged. Every fix was observed failing before it was written. Each suite drives the public surface from dist against a local scripted fixture, so all four run with no credentials and no network egress: test:websocket-agent-handler 0/4 passed on the stub build -> 4/4 test:video-options-allowlist fails with the hunk reverted -> 1/1 test:stream-structured-output structuredData absent -> 2/2 test:stream-teardown-abort server saw no cancellation -> 1/1 The streaming suite's second case covers the branch that matters for most providers: suppressResponseFormatWithTools() returns the base default true everywhere except OpenAI and Azure, so a tools-bearing turn goes out without response_format and the single tool-free re-ask is the only path to structured output. It was proved against one build with that branch gated off and on, so no rebuild sits between the two runs. It also pins the conversation the re-ask sends: no tool reply may carry a tool_call_id with no matching announced tool_calls entry, which is the shape a real vendor rejects outright. Verified on the assembled tree, not only per-fix: all four suites green together, plus test:stream-tool-telemetry (4/4) and test:servers (41/41), which cover the baseProvider and websocket code the first and fourth fixes touch, and the CI safety net — test:providers-mocked 120/120 and test:provider-structure 3/3. check and lint exit 0. Review round 1 (Yama): - The teardown signal is composed with composeAbortSignalsScoped, not plain composeAbortSignals. The plain one is AbortSignal.any, whose registration on a source survives until the derived signal is collected, so a caller reusing one long-lived abortSignal across many stream calls accumulated a dependent per call — the exact hazard the scoped helper's docstring names (MaxListenersExceededWarning at 10+). Its dispose() is threaded to teardown(), which already runs unconditionally in the generator's finally, so the listeners come off the caller's signal when the stream settles rather than whenever GC gets to it. - The WebSocket guide's client switch read message.data for a "response" frame while its own message-routing example read message.content. This change is what turns that inconsistency into a real mismatch, so the client switch now reads content. - Stream frames go through a trySend helper instead of raw socket.send. A closed socket threw mid-stream, unwound into onMessage's catch, and that catch's own send threw again with nothing above it — an unhandled rejection. Now a refused frame ends the stream cleanly: breaking the loop runs IteratorClose, which with this same commit's teardown fix aborts the upstream request, and the route returns nothing so no completion frame is sent to a socket that just refused one. The catch is guarded for the same reason. - The tool-free re-ask's tokens are reported on metadata.structuredDataUsage. Review read this as mirroring the generate path; it is the opposite — that path bills its equivalent re-ask inline and says why, so leaving the streaming one unreported was an inconsistency, and a caller capping spend would under-count. They cannot be folded into the stream's usage, which resolved before the re-ask ran and may already have been read, so they are delivered alongside structuredData and read at the same moment. Billed as each response lands, matching the generate path's rule that a discarded answer was still charged. The suite asserts both token counts are non-zero. Review round 2 (CodeRabbit, posted after the round-1 approval): - SECURITY: WebSocketConfig has always declared auth.required with a strategy, and handleConnection has always accepted and stored an AuthenticatedUser, but nothing consulted either — required:true produced the same open socket as required:false. Inert while the agent routes returned canned values; not inert once tool_call reaches executeTool, since any client that can open the socket could then run any registered tool with arbitrary arguments. handleConnection now refuses an unauthenticated connection when the config requires one, so every handler on the manager inherits the gate rather than each route deciding. - disposeComposedSignal ran only inside the wrapper's teardown, so a throw between composition and the wrapper's creation left the once-listeners attached. The body is now wrapped in try/catch — catch, not finally: dispose only detaches listeners, so running it on the success path would sever the caller's abort from the signal the stream is still using. - The outer catch in resolveStreamStructuredData swallowed the inner catch's cancellation rethrow, so a cancelled turn resolved as if the abort never happened. Now mirrors the generate path's reformat guard: a caller's abort propagates, our own turn deadline does not, distinguished by the TimeoutError reason on the composed signal. - Scalar and array schema roots never resolved. coerceJsonToSchema scans for a balanced object or array span, so a z.string() schema answered "sunny" found nothing; recoverScalarRoot existed for exactly this and was wired only into neurolink.ts. Both acceptance points now go through one resolver that tries each. appendJsonSchemaInstruction also demanded "a single JSON object", which told a compliant model to emit something an array or scalar root can never satisfy — the generate path's own re-ask prompt was corrected for this reason and the shared helper was missed; it now asks for a JSON value. - structuredData was assigned unconditionally, creating the own property even when undefined, so `"structuredData" in metadata` reported true for a turn that produced no object — the exact check the field's contract tells readers to use. Assigned only when there is a value. Each fix is pinned. The three new cases were observed failing with the production hunks neutralized and passing with them restored, in one pass: the scalar-root and absent-field cases failed, the auth refusal failed, and the two positive auth controls (authenticated accepted, auth unset accepted) kept passing — so the gate is proved to refuse the right connection rather than all of them. Two more gaps in the same WebSocket surface are closed. generate/stream now validate the client's payload against an explicit options allowlist (WebSocketAgentOptionsSchema) before calling the SDK — options.credentials and anything outside provider/model/systemPrompt/temperature/maxTokens/ maxSteps never reaches neurolink.generate()/.stream(), so a client can no longer redirect the server's outbound request to an arbitrary host or swap in its own provider API key. ServerAuthConfig's roles and permissions, already declared and documented (including the guide's own /ws/admin example) but never consulted, are now enforced in handleConnection with the same any-of/all-of semantics as hasAnyRole/hasAllPermissions in authContext.ts. Response frames for generate, stream and tool_call also echo back a client-supplied request id (string or number) so a caller with more than one call in flight on the same connection can correlate a response to its own request; a client that sends no id sees no change. Twelve new cases cover all three: the SSRF/credential guard against two local fixture servers (one reachable only through the rejected baseURL), the any-of/all-of role and permission checks, and id-echo across generate, stream, tool_call and the error path plus the no-id backward-compatible case. --- docs/api/README.md | 1 + docs/api/classes/WebSocketMessageRouter.md | 2 +- .../functions/createAgentWebSocketHandler.md | 24 +- docs/api/type-aliases/StreamLoopArgs.md | 10 + docs/api/type-aliases/StreamResult.md | 39 + .../api/type-aliases/TextGenerationOptions.md | 34 + docs/api/type-aliases/WebSocketRequestId.md | 15 + docs/guides/server-adapters/websocket.md | 2 +- package.json | 4 + src/cli/loop/optionsSchema.ts | 1 + src/lib/core/baseProvider.ts | 497 +++++---- src/lib/core/nativeGenerateLoop.ts | 8 +- src/lib/neurolink.ts | 4 + .../providers/openaiChatCompletionsBase.ts | 361 ++++++- src/lib/server/utils/validation.ts | 29 + src/lib/server/websocket/WebSocketHandler.ts | 277 ++++- src/lib/types/generate.ts | 16 + src/lib/types/openaiCompatible.ts | 6 + src/lib/types/server.ts | 9 + src/lib/types/stream.ts | 30 + ...ous-test-suite-stream-structured-output.ts | 579 ++++++++++ ...inuous-test-suite-stream-teardown-abort.ts | 368 +++++++ ...uous-test-suite-video-options-allowlist.ts | 205 ++++ ...uous-test-suite-websocket-agent-handler.ts | 990 ++++++++++++++++++ 24 files changed, 3286 insertions(+), 225 deletions(-) create mode 100644 docs/api/type-aliases/WebSocketRequestId.md create mode 100644 test/continuous-test-suite-stream-structured-output.ts create mode 100644 test/continuous-test-suite-stream-teardown-abort.ts create mode 100644 test/continuous-test-suite-video-options-allowlist.ts create mode 100644 test/continuous-test-suite-websocket-agent-handler.ts diff --git a/docs/api/README.md b/docs/api/README.md index a828dc20ad..dcfa1d6161 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 b7b5e42c45..c3c2a92b61 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 785eec542c..7a7cc661c3 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 45c54f9fec..9a8a2cf64b 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 9730d68532..0d3ccb5958 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 4d4e2c6c24..a9d5efabcf 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 0000000000..d8b9709523 --- /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 662950f151..b224f378b0 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 39657861ae..9581d9689e 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 5453cd13bb..5e2af57025 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 38759a3315..d2682195d0 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 3c7a4ec14f..138b658eb4 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 f45eeb1761..82e283d860 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 b3b2c6a711..b0996661de 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 8717be4868..ecbe877510 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 37609a10be..f3b041e8db 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 4d780351fc..bce725d71d 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 4d354569d7..f921e163e2 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 d8cdb198f8..defaec6c4b 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 c2c248ed16..ee96a31642 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 0000000000..06dc8be22e --- /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 0000000000..ac8466393b --- /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 0000000000..1abeecd482 --- /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 0000000000..e6035d5749 --- /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", + ); + }); +});