From cc60a3eb41598823a5276266650536ba0d91b75e Mon Sep 17 00:00:00 2001 From: Anthony Fu Date: Thu, 7 May 2026 18:34:03 +0900 Subject: [PATCH] feat(devframe): async generator RPC functions MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Add `type: 'generator'` for `defineRpcFunction` so handlers declared as `async function*` stream their yields to the caller without manual channel scaffolding. The framework substitutes the user's definition with an internal action wrapper that allocates a sink on a hidden `devframe:rpc:generators` channel and returns a stream-id envelope; the client wrapper unwraps the envelope and `streaming.subscribe()`s automatically, so `await rpc.call(name, args)` resolves to a ready-to-iterate `StreamReader`. Cancellation flows through `getCurrentRpcStream()` — an AsyncLocalStorage helper that mirrors `getCurrentRpcSession()` and exposes the sink's `signal`, `streamId`, and originating session inside the handler body. Server-side callers can iterate without paying for the streaming round-trip via `invokeLocalGenerator(rpc, name, ...args)`. Per-stream `replayWindow` defaults to 256 (floored to 1) to win the client-subscribe-vs-first-yield race; this required threading the existing `RpcStreamingChannel.start()` helper to accept a per-stream override on top of the channel default. Generator definitions reject `agent`, `cacheable`, `dump`, `snapshot`, and `jsonSerializable: true` at registration via four new diagnostics (DF0033–DF0036). DF0034 fires at runtime if the handler doesn't return an `AsyncIterable`. Also fixes a pre-existing replay bug in `node/rpc-streaming.ts`: when a producer closed with an error and a subscriber arrived during the `closedStreamRetention` window, the late subscriber received a clean close instead of the original error. The streaming record now captures the end payload and replays it on subscribe. Migrations & dogfooding: - Streaming guide gains an "Async Generator RPC" section ahead of the manual channel section, framed as the recommended path. - `devframe-streaming-chat` example adds a `:tokenize` generator alongside the existing `:send` action so readers can compare both approaches in one place. - `skills/devframe` and `skills/vite-devtools-kit` get matching guidance. - 12 new integration tests (`rpc-generators.test.ts`) covering happy path, cooperative cancel, throw mid-stream, throw before first yield, late-subscriber replay, concurrent isolation, `invokeLocalGenerator`, and all four validation diagnostics. tsnapi snapshots updated for the new public exports (`getCurrentRpcStream`, `invokeLocalGenerator`, `attachRpcGenerators`, `RpcGeneratorStreamContext`, `RpcGeneratorFunctionDefinition`). Pre-PR checklist (lint + 462 tests + typecheck + build) all green. Co-Authored-By: Claude Opus 4.7 (1M context) --- devframe/docs/.vitepress/sidebar.ts | 2 +- devframe/docs/errors/DF0033.md | 50 +++ devframe/docs/errors/DF0034.md | 51 +++ devframe/docs/errors/DF0035.md | 53 +++ devframe/docs/errors/DF0036.md | 49 ++ devframe/docs/errors/index.md | 4 + devframe/docs/guide/streaming.md | 95 +++- .../devframe-streaming-chat/src/devtool.ts | 29 ++ .../packages/devframe/src/client/index.ts | 1 + .../devframe/src/client/rpc-generators.ts | 50 +++ devframe/packages/devframe/src/client/rpc.ts | 5 + .../src/node/__tests__/rpc-generators.test.ts | 419 ++++++++++++++++++ .../packages/devframe/src/node/diagnostics.ts | 20 + .../devframe/src/node/host-functions.ts | 5 + devframe/packages/devframe/src/node/index.ts | 1 + .../devframe/src/node/rpc-generators.ts | 252 +++++++++++ .../devframe/src/node/rpc-streaming.ts | 22 +- devframe/packages/devframe/src/rpc/handler.ts | 30 +- devframe/packages/devframe/src/rpc/types.ts | 306 ++++++++----- .../packages/devframe/src/rpc/validation.ts | 13 +- devframe/packages/devframe/src/types/rpc.ts | 29 +- .../tsnapi/devframe/client.snapshot.d.ts | 1 + .../tsnapi/devframe/client.snapshot.js | 1 + .../tsnapi/devframe/index.snapshot.d.ts | 1 + .../tsnapi/devframe/node.snapshot.d.ts | 7 + .../tsnapi/devframe/node.snapshot.js | 3 + .../tsnapi/devframe/rpc.snapshot.d.ts | 1 + .../tsnapi/devframe/types.snapshot.d.ts | 1 + skills/devframe/SKILL.md | 33 ++ skills/vite-devtools-kit/SKILL.md | 30 ++ 30 files changed, 1441 insertions(+), 123 deletions(-) create mode 100644 devframe/docs/errors/DF0033.md create mode 100644 devframe/docs/errors/DF0034.md create mode 100644 devframe/docs/errors/DF0035.md create mode 100644 devframe/docs/errors/DF0036.md create mode 100644 devframe/packages/devframe/src/client/rpc-generators.ts create mode 100644 devframe/packages/devframe/src/node/__tests__/rpc-generators.test.ts create mode 100644 devframe/packages/devframe/src/node/rpc-generators.ts diff --git a/devframe/docs/.vitepress/sidebar.ts b/devframe/docs/.vitepress/sidebar.ts index 93455e92..e00f2041 100644 --- a/devframe/docs/.vitepress/sidebar.ts +++ b/devframe/docs/.vitepress/sidebar.ts @@ -27,7 +27,7 @@ export default function devframeSidebar(prefix = ''): DefaultTheme.SidebarItem[] text: 'Error Reference', link: `${prefix}/errors/`, collapsed: true, - items: Array.from({ length: 32 }, (_, i) => { + items: Array.from({ length: 36 }, (_, i) => { const code = `DF${String(i + 1).padStart(4, '0')}` return { text: code, link: `${prefix}/errors/${code}` } }), diff --git a/devframe/docs/errors/DF0033.md b/devframe/docs/errors/DF0033.md new file mode 100644 index 00000000..ff45abc5 --- /dev/null +++ b/devframe/docs/errors/DF0033.md @@ -0,0 +1,50 @@ +--- +outline: deep +--- + +# DF0033: Generator with Agent Exposure + +> Package: `devframe` + +## Message + +> Generator RPC function "`{name}`" has `agent` set, which is not supported in the current release. + +## Cause + +`type: 'generator'` RPC functions stream their yields through the existing streaming-channel transport, which currently uses `structured-clone-es` regardless of the source function's serialization preferences. Agent (MCP) exposure requires strict JSON serialization for tool inputs and outputs, so generator functions cannot be exposed as agent tools yet. + +Streaming-MCP support is planned as a follow-up. + +## Example + +```ts +defineRpcFunction({ + name: 'plugin:tail-logs', + type: 'generator', + yields: v.string(), + agent: { description: 'Tail server logs' }, // ← rejected + async* handler() { + yield 'log line 1' + }, +}) +``` + +## Fix + +Remove the `agent` field, or change the function type to `'query'` / `'static'` if it can return a single value: + +```ts +defineRpcFunction({ + name: 'plugin:tail-logs', + type: 'generator', + yields: v.string(), + async* handler() { + yield 'log line 1' + }, +}) +``` + +## Source + +`packages/devframe/src/rpc/validation.ts` diff --git a/devframe/docs/errors/DF0034.md b/devframe/docs/errors/DF0034.md new file mode 100644 index 00000000..fff21a4f --- /dev/null +++ b/devframe/docs/errors/DF0034.md @@ -0,0 +1,51 @@ +--- +outline: deep +--- + +# DF0034: Generator Handler Not Async-Iterable + +> Package: `devframe` + +## Message + +> Generator RPC function "`{name}`" handler did not return an `AsyncIterable`. + +## Cause + +A `type: 'generator'` RPC function expects its handler to return an `AsyncIterable` — typically by being declared as `async function*`. The framework drives the iterator with `for await`, so it must implement `[Symbol.asyncIterator]`. + +This error is raised at runtime on the first call, not at registration, because TypeScript cannot statically distinguish `async function*` from `async function` in all cases (e.g. when the handler is supplied via `setup()`). + +## Example + +```ts +// ❌ Wrong — async function (not generator) returning a string +defineRpcFunction({ + name: 'plugin:tokens', + type: 'generator', + yields: v.string(), + handler: async () => 'hello', +}) +``` + +## Fix + +Declare the handler as `async function*` and `yield` instead of `return`: + +```ts +defineRpcFunction({ + name: 'plugin:tokens', + type: 'generator', + yields: v.string(), + async* handler() { + yield 'hello' + yield 'world' + }, +}) +``` + +If the handler comes from `setup()`, ensure `setup()` returns `{ handler: async function*() { ... } }`. + +## Source + +`packages/devframe/src/node/rpc-generators.ts` diff --git a/devframe/docs/errors/DF0035.md b/devframe/docs/errors/DF0035.md new file mode 100644 index 00000000..4f735f48 --- /dev/null +++ b/devframe/docs/errors/DF0035.md @@ -0,0 +1,53 @@ +--- +outline: deep +--- + +# DF0035: Generator Missing Yields Schema + +> Package: `devframe` + +## Message + +> Generator RPC function "`{name}`" declares `args` but no `yields` schema — yielded value type cannot be inferred. + +## Cause + +When a `type: 'generator'` RPC function declares an `args` schema, the framework expects a matching `yields` schema for end-to-end type inference. Without it, `rpc.call(name, ...)` resolves to `Promise>` on the client, defeating the purpose of having validated arguments. + +This mirrors the requirement that `query`-typed functions with `args` also declare `returns`. + +## Example + +```ts +// ❌ Wrong — args present but no yields schema +defineRpcFunction({ + name: 'plugin:search', + type: 'generator', + args: [v.object({ query: v.string() })], + async* handler({ query }) { + yield query + }, +}) +``` + +## Fix + +Add a `yields` schema describing the type of each yielded value: + +```ts +defineRpcFunction({ + name: 'plugin:search', + type: 'generator', + args: [v.object({ query: v.string() })], + yields: v.string(), + async* handler({ query }) { + yield query + }, +}) +``` + +If you don't need argument validation either, drop both schemas — the call type falls back to `() => Promise>`. + +## Source + +`packages/devframe/src/rpc/validation.ts` diff --git a/devframe/docs/errors/DF0036.md b/devframe/docs/errors/DF0036.md new file mode 100644 index 00000000..5d3cc570 --- /dev/null +++ b/devframe/docs/errors/DF0036.md @@ -0,0 +1,49 @@ +--- +outline: deep +--- + +# DF0036: Generator Has Inapplicable Option + +> Package: `devframe` + +## Message + +> Generator RPC function "`{name}`" sets `{option}`, which has no effect on generator-typed functions. + +## Cause + +Generators are streaming primitives — they yield a sequence of values over time rather than returning a single cacheable result. Options that only make sense for request-response RPC types are rejected at registration: + +- `cacheable` — caching a streaming result would replay a stale `streamId` pointing to a closed stream. +- `jsonSerializable: true` — generator yields flow through the existing streaming-channel transport, which uses `structured-clone-es` regardless. The flag is misleading on a generator. +- `dump` and `snapshot` — already handled by [`DF0027`](./DF0027) and [`DF0028`](./DF0028). + +## Example + +```ts +// ❌ Wrong — `cacheable` has no effect on a generator +defineRpcFunction({ + name: 'plugin:tokens', + type: 'generator', + cacheable: true, + yields: v.string(), + async* handler() { yield 'hi' }, +}) +``` + +## Fix + +Remove the option, or change the function type to one that supports it: + +```ts +defineRpcFunction({ + name: 'plugin:tokens', + type: 'generator', + yields: v.string(), + async* handler() { yield 'hi' }, +}) +``` + +## Source + +`packages/devframe/src/rpc/validation.ts` diff --git a/devframe/docs/errors/index.md b/devframe/docs/errors/index.md index 123fe745..cc8d6712 100644 --- a/devframe/docs/errors/index.md +++ b/devframe/docs/errors/index.md @@ -50,3 +50,7 @@ Emitted by `devframe` — framework-neutral host / shared-state / auth surface. | [DF0030](./DF0030) | error | Unknown Stream ID | — | | [DF0031](./DF0031) | error | Write to Closed Stream | — | | [DF0032](./DF0032) | error | Streaming Channel Already Registered | — | +| [DF0033](./DF0033) | error | Generator with Agent Exposure | — | +| [DF0034](./DF0034) | error | Generator Handler Not Async-Iterable | — | +| [DF0035](./DF0035) | error | Generator Missing Yields Schema | — | +| [DF0036](./DF0036) | error | Generator Has Inapplicable Option | — | diff --git a/devframe/docs/guide/streaming.md b/devframe/docs/guide/streaming.md index 7941d155..207671ca 100644 --- a/devframe/docs/guide/streaming.md +++ b/devframe/docs/guide/streaming.md @@ -6,6 +6,8 @@ outline: deep Devframe's streaming-channel API provides server→client push for chunk-style data — chat deltas, log lines, build progress, anything you'd otherwise express as a sequence of fire-and-forget events. It builds on the same WebSocket transport as the rest of the RPC layer, but adds the conventions every chunked feed needs: stream IDs, cooperative cancellation, replay on reconnect, and first-class **Web Streams** interop. +For the common case of "I have an `async function*` and I want to ship its yields to the client", reach for [Async Generator RPC](#async-generator-rpc) below — it auto-allocates the sink, wires cancellation, and the client receives a ready-to-iterate `StreamReader` from `rpc.call(...)`. Drop down to the lower-level channel API in this guide when you need fan-out, late-replay across multiple subscribers, or client-to-server uploads. + ## Overview ```mermaid @@ -233,8 +235,95 @@ If you need authoritative state rather than every intermediate value, prefer [sh | Replay on reconnect | Fire-and-forget signaling | Diff-based sync between clients | | Client-to-server uploads (files, mic frames) | | | +## Async Generator RPC + +The streaming-channel API is the foundation; the most common shape — "stream the yields of an `async function*` to the caller" — has a higher-level wrapper that hides the channel altogether. + +### Defining a generator + +Set `type: 'generator'` on `defineRpcFunction` and write the handler as `async function*`: + +```ts +import { defineRpcFunction } from 'devframe' +import { getCurrentRpcStream } from 'devframe/node' +import * as v from 'valibot' + +defineRpcFunction({ + name: 'my-devtool:chat', + type: 'generator', + args: [v.object({ prompt: v.string() })], + yields: v.string(), + async* handler({ prompt }) { + const { signal } = getCurrentRpcStream()! + for (const token of fakeTokens(prompt)) { + if (signal.aborted) + return + yield token + } + }, +}) +``` + +`getCurrentRpcStream()` is `AsyncLocalStorage`-backed (mirrors `getCurrentRpcSession()`); inside the generator body it returns `{ signal, streamId, session }`. Outside, it returns `undefined`. Poll `signal.aborted` and exit cooperatively when the consumer cancels. + +### Calling from the client + +```ts +import { connectDevtool } from 'devframe/client' + +const rpc = await connectDevtool() +const reader = await rpc.call('my-devtool:chat', { prompt: 'Hi' }) + +for await (const token of reader) + appendToken(token) + +// or pipe out as a Web Stream: +await reader.readable.pipeTo(downloadWritable) + +reader.cancel() // sends cancel upstream; server's `signal` flips +``` + +`rpc.call` resolves to `Promise>` for generator-typed functions. The reader is the same shape as `rpc.streaming.subscribe()` — it's both `AsyncIterable` and exposes `.readable: ReadableStream`. Pick one surface per reader (they share an internal queue). + +### Per-function options + +| Field | Default | Effect | +|-------|---------|--------| +| `yields` | — | Valibot schema for each yielded value. Optional; required if `args` is set ([`DF0035`](../errors/DF0035)). Drives client-side type inference of `StreamReader`. | +| `replayWindow` | `256` | Per-stream ring-buffer size. Floored to `1` — the client subscribe lands a few ms after the wrapper allocates the sink, so early yields must be replayable. | +| `closedStreamRetention` | `30_000` | Milliseconds the stream is held open for late subscribers after the producer closes. Mirrors the channel-level option. | + +### Local invocation + +Server-side callers can iterate a generator without paying for the streaming round-trip: + +```ts +import { invokeLocalGenerator } from 'devframe/node' + +for await (const token of await invokeLocalGenerator(rpc, 'my-devtool:chat', { prompt: 'Hi' })) { + // process tokens directly — no transport, no sink allocation +} +``` + +`invokeLocalGenerator` returns the bare `AsyncIterable` from the user's handler. `getCurrentRpcStream()` returns `undefined` in this path (no stream, no signal). + +### What you can't put on a generator + +- `agent: { ... }` — streaming-MCP exposure is deferred. ([`DF0033`](../errors/DF0033)) +- `cacheable: true` — caching a streaming response would replay a stale `streamId`. ([`DF0036`](../errors/DF0036)) +- `jsonSerializable: true` — chunks always travel via `structured-clone` regardless. ([`DF0036`](../errors/DF0036)) +- `dump`, `snapshot` — streaming results don't have a "snapshot" semantics. ([`DF0027`](../errors/DF0027), [`DF0028`](../errors/DF0028)) + +### When to reach for the lower-level channel API + +Generator RPC covers the 90% case. Drop down to `ctx.rpc.streaming.create(...)` directly when you need: + +- **Fan-out** — multiple subscribers seeing the same stream from a single producer. Generator RPC allocates a fresh stream per call. +- **Long-running side-channel streams** — terminal output, file watches, anything that lives beyond a single RPC call. +- **Client-to-server uploads** — `channel.openInbound()` with a paired action handler to allocate the id (see [Client-to-Server Uploads](#client-to-server-uploads)). + ## Reference -- API surface: `RpcStreamingHost`, `RpcStreamingChannel`, `StreamSink`, `StreamReader` in `devframe/types`. -- Working example: [`devframe/examples/devframe-streaming-chat`](https://github.com/vitejs/devtools/tree/main/devframe/examples/devframe-streaming-chat). -- Errors: [`DF0029`](../errors/DF0029) (overflow), [`DF0030`](../errors/DF0030) (unknown stream id), [`DF0031`](../errors/DF0031) (write to closed stream), [`DF0032`](../errors/DF0032) (channel name collision). +- API surface: `RpcStreamingHost`, `RpcStreamingChannel`, `StreamSink`, `StreamReader` in `devframe/types`. Generator RPC: `getCurrentRpcStream`, `invokeLocalGenerator` in `devframe/node`. +- Working example: [`devframe/examples/devframe-streaming-chat`](https://github.com/vitejs/devtools/tree/main/devframe/examples/devframe-streaming-chat) — `:send` uses the channel API directly, `:tokenize` uses generator RPC. +- Errors: [`DF0029`](../errors/DF0029) (overflow), [`DF0030`](../errors/DF0030) (unknown stream id), [`DF0031`](../errors/DF0031) (write to closed stream), [`DF0032`](../errors/DF0032) (channel name collision), [`DF0033`](../errors/DF0033) (generator + agent), [`DF0034`](../errors/DF0034) (handler not async-iterable), [`DF0035`](../errors/DF0035) (generator missing yields), [`DF0036`](../errors/DF0036) (inapplicable generator option). diff --git a/devframe/examples/devframe-streaming-chat/src/devtool.ts b/devframe/examples/devframe-streaming-chat/src/devtool.ts index eeb50cca..c31645e8 100644 --- a/devframe/examples/devframe-streaming-chat/src/devtool.ts +++ b/devframe/examples/devframe-streaming-chat/src/devtool.ts @@ -1,5 +1,6 @@ import { fileURLToPath } from 'node:url' import { defineRpcFunction } from 'devframe' +import { getCurrentRpcStream } from 'devframe/node' import { defineDevtool } from 'devframe/types' import { nanoid } from 'devframe/utils/nanoid' import * as v from 'valibot' @@ -213,6 +214,34 @@ export default defineDevtool({ }, })) + // ─── Generator RPC demo ────────────────────────────────────────────── + // The action above (`:send`) wires shared-state plumbing and streaming + // by hand. The generator below shows the lower-friction alternative + // for token feeds that don't need history side-effects: declare a + // handler as `async function*` and the framework auto-allocates a + // sink, drains yields onto it, and gives callers back a + // `StreamReader`. `getCurrentRpcStream()` exposes the abort + // signal for cooperative cancellation. + ctx.rpc.register(defineRpcFunction({ + name: 'devframe-streaming-chat:tokenize', + type: 'generator', + args: [v.object({ + prompt: v.string(), + intervalMs: v.optional(v.number(), 35), + })], + yields: v.string(), + replayWindow: 1024, + async* handler({ prompt, intervalMs = 35 }) { + const ctx = getCurrentRpcStream() + for (const token of fakeTokens(prompt)) { + if (ctx?.signal.aborted) + return + yield token + await new Promise(r => setTimeout(r, intervalMs)) + } + }, + })) + ctx.views.hostStatic(BASE_PATH, distDir) ctx.docks.register({ id: 'devframe-streaming-chat', diff --git a/devframe/packages/devframe/src/client/index.ts b/devframe/packages/devframe/src/client/index.ts index 77862cf9..d6afaee3 100644 --- a/devframe/packages/devframe/src/client/index.ts +++ b/devframe/packages/devframe/src/client/index.ts @@ -2,4 +2,5 @@ export * from './context' export * from './docks' export * from './rpc' export { getDevToolsRpcClient as connectDevtool } from './rpc' +export * from './rpc-generators' export * from './rpc-streaming' diff --git a/devframe/packages/devframe/src/client/rpc-generators.ts b/devframe/packages/devframe/src/client/rpc-generators.ts new file mode 100644 index 00000000..7769d732 --- /dev/null +++ b/devframe/packages/devframe/src/client/rpc-generators.ts @@ -0,0 +1,50 @@ +import type { StreamReader } from 'devframe/utils/streaming-channel' +import type { DevToolsRpcClient } from './rpc' + +/** + * Hidden channel name used by every `type: 'generator'` RPC. Mirrors the + * server-side constant in `node/rpc-generators.ts`. Internal — generators + * are addressed by name, not by channel id, so consumers should never + * subscribe here directly. + * + * @internal + */ +const GENERATOR_CHANNEL = 'devframe:rpc:generators' + +/** + * Shape of the envelope returned by the server's auto-generated wrapper + * around a `type: 'generator'` definition. Internal — clients never see + * this directly because `attachRpcGeneratorsClient` unwraps it. + */ +interface RpcGeneratorEnvelope { + __generator: true + streamId: string +} + +function isGeneratorEnvelope(value: unknown): value is RpcGeneratorEnvelope { + return ( + !!value + && typeof value === 'object' + && (value as any).__generator === true + && typeof (value as any).streamId === 'string' + ) +} + +/** + * Hooks `rpc.call` so that calls to `type: 'generator'` server functions + * resolve to a `StreamReader` (subscribed to the hidden generators + * channel) instead of the raw `{ __generator: true, streamId }` envelope. + * + * Must be called after `rpc.streaming` is initialized so the host can + * delegate to `rpc.streaming.subscribe()`. + */ +export function attachRpcGeneratorsClient(rpc: DevToolsRpcClient): void { + const originalCall = rpc.call + rpc.call = (async (name: any, ...args: any[]) => { + const result = await originalCall(name, ...args) + if (isGeneratorEnvelope(result)) { + return rpc.streaming.subscribe(GENERATOR_CHANNEL, result.streamId) as unknown as StreamReader + } + return result + }) as typeof rpc.call +} diff --git a/devframe/packages/devframe/src/client/rpc.ts b/devframe/packages/devframe/src/client/rpc.ts index 893a569f..9dd9de30 100644 --- a/devframe/packages/devframe/src/client/rpc.ts +++ b/devframe/packages/devframe/src/client/rpc.ts @@ -10,6 +10,7 @@ import { import { RpcCacheManager, RpcFunctionsCollectorBase } from 'devframe/rpc' import { createEventEmitter } from 'devframe/utils/events' import { humanId } from 'devframe/utils/human-id' +import { attachRpcGeneratorsClient } from './rpc-generators' import { createRpcSharedStateClientHost } from './rpc-shared-state' import { createStaticRpcClientMode } from './rpc-static' import { createRpcStreamingClientHost } from './rpc-streaming' @@ -296,6 +297,10 @@ export async function getDevToolsRpcClient( rpc.sharedState = createRpcSharedStateClientHost(rpc) rpc.streaming = createRpcStreamingClientHost(rpc) + // Hook `rpc.call` so generator-typed RPCs return a `StreamReader` + // (auto-subscribed) instead of the raw `__generator` envelope. + // Must come after `streaming` is initialized. + attachRpcGeneratorsClient(rpc) // @ts-expect-error assign to readonly property context.rpc = rpc diff --git a/devframe/packages/devframe/src/node/__tests__/rpc-generators.test.ts b/devframe/packages/devframe/src/node/__tests__/rpc-generators.test.ts new file mode 100644 index 00000000..4100c0fa --- /dev/null +++ b/devframe/packages/devframe/src/node/__tests__/rpc-generators.test.ts @@ -0,0 +1,419 @@ +import type { DevToolsNodeContext, DevToolsRpcClientFunctions, DevToolsRpcServerFunctions } from 'devframe/types' +import type { StreamReader } from 'devframe/utils/streaming-channel' +import { AsyncLocalStorage } from 'node:async_hooks' +import { createRpcStreamingClientHost } from 'devframe/client' +import { createRpcClient } from 'devframe/rpc/client' +import { createRpcServer } from 'devframe/rpc/server' +import { createWsRpcChannel } from 'devframe/rpc/transports/ws-client' +import { attachWsRpcTransport } from 'devframe/rpc/transports/ws-server' +import * as v from 'valibot' +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest' +import { WebSocket } from 'ws' +import { RpcFunctionsHost } from '../host-functions' +import { + getCurrentRpcStream, + invokeLocalGenerator, +} from '../rpc-generators' + +// Mirrors the internal constant in `node/rpc-generators.ts`. The test +// harness uses it to subscribe to the hidden channel directly because it +// bypasses the full client-side `attachRpcGeneratorsClient` wrapper. +const GENERATOR_CHANNEL = 'devframe:rpc:generators' + +vi.stubGlobal('WebSocket', WebSocket) + +let nextPort = 42_000 +function allocatePort(): number { + return nextPort++ +} + +interface Harness { + port: number + rpcHost: RpcFunctionsHost + close: () => Promise +} + +async function bootHost(): Promise { + const port = allocatePort() + const mockContext = {} as DevToolsNodeContext + const rpcHost = new RpcFunctionsHost(mockContext) + + const asyncStorage = new AsyncLocalStorage() + + const rpcGroup = createRpcServer( + rpcHost.functions, + { + rpcOptions: { + resolver(_name, fn) { + // eslint-disable-next-line ts/no-this-alias + const rpc = this + if (!fn) + return undefined + return async function (this: any, ...args) { + return await asyncStorage.run({ rpc, meta: rpc.$meta }, async () => { + return (await fn).apply(this, args) + }) + } + }, + }, + }, + ) + + const { wss } = attachWsRpcTransport(rpcGroup, { + port, + host: '127.0.0.1', + onDisconnected: (_ws, meta) => { + rpcHost._emitSessionDisconnected(meta) + }, + }) + + ;(rpcHost as any)._rpcGroup = rpcGroup + ;(rpcHost as any)._asyncStorage = asyncStorage + + return { + port, + rpcHost, + async close() { + for (const ws of wss.clients) ws.terminate() + await new Promise(r => wss.close(() => r())) + }, + } +} + +interface FakeClient { + rpc: ReturnType> + streaming: ReturnType + /** Calls the generator RPC and returns the auto-subscribed reader. */ + callGenerator: (name: string, ...args: any[]) => Promise> + close: () => void +} + +function bootClient(port: number): FakeClient { + const listeners = new Set<(trusted: boolean) => void>() + const fakeEvents = { + on(name: string, fn: (trusted: boolean) => void) { + if (name === 'rpc:is-trusted:updated') + listeners.add(fn) + return () => listeners.delete(fn) + }, + } + + const clientFns: any = {} + const clientRpcStub = { + register(def: { name: string, handler: (...args: any[]) => any }) { + clientFns[def.name] = def.handler + }, + } + + const rpc = createRpcClient( + clientFns, + { + channel: createWsRpcChannel({ url: `ws://127.0.0.1:${port}` }), + }, + ) + + const fakeRpcClient = { + isTrusted: true, + events: fakeEvents, + client: clientRpcStub, + callEvent: (name: any, ...args: any[]) => (rpc as any).$callEvent(name, ...args), + } as any + + const streaming = createRpcStreamingClientHost(fakeRpcClient) + + // Mirrors `attachRpcGeneratorsClient` — but applied to our minimal fake + // client, so we can verify the envelope-unwrap logic without depending + // on the full DevToolsRpcClient surface used by `connectDevtool`. + async function callGenerator(name: string, ...args: any[]): Promise> { + const result: any = await (rpc as any).$call(name, ...args) + if (!result || result.__generator !== true || typeof result.streamId !== 'string') + throw new Error(`Expected generator envelope from "${name}"; got ${JSON.stringify(result)}`) + return streaming.subscribe(GENERATOR_CHANNEL, result.streamId) + } + + return { + rpc, + streaming, + callGenerator, + close() { + // ws closes when server tears down + }, + } +} + +describe('rpc-generators integration', () => { + let harness: Harness + + beforeEach(async () => { + harness = await bootHost() + }) + + afterEach(async () => { + await harness.close() + }) + + it('round-trips yields from a generator handler to the client reader', async () => { + harness.rpcHost.register({ + name: 'test:gen-happy', + type: 'generator', + yields: v.string(), + async* handler() { + yield 'a' + yield 'b' + yield 'c' + }, + } as any) + + const client = bootClient(harness.port) + await new Promise(r => setTimeout(r, 50)) + + const reader = await client.callGenerator('test:gen-happy') + const collected: string[] = [] + for await (const chunk of reader) collected.push(chunk) + + expect(collected).toEqual(['a', 'b', 'c']) + expect(reader.done).toBe(true) + }) + + it('exposes the abort signal via getCurrentRpcStream and stops cooperatively on cancel', async () => { + let observedAbort = false + let yieldsBeforeAbort = 0 + + harness.rpcHost.register({ + name: 'test:gen-cancel', + type: 'generator', + yields: v.number(), + async* handler() { + const ctx = getCurrentRpcStream()! + for (let i = 0; i < 100; i++) { + if (ctx.signal.aborted) { + observedAbort = true + return + } + yieldsBeforeAbort = i + 1 + yield i + await new Promise(r => setTimeout(r, 5)) + } + }, + } as any) + + const client = bootClient(harness.port) + await new Promise(r => setTimeout(r, 50)) + + const reader = await client.callGenerator('test:gen-cancel') + const collected: number[] = [] + for await (const chunk of reader) { + collected.push(chunk) + if (collected.length >= 5) { + reader.cancel() + } + } + + await vi.waitFor(() => { + expect(observedAbort).toBe(true) + }) + expect(collected.length).toBeGreaterThanOrEqual(5) + expect(yieldsBeforeAbort).toBeLessThan(100) + }) + + it('propagates a thrown error from the generator to the client reader', async () => { + harness.rpcHost.register({ + name: 'test:gen-throw', + type: 'generator', + yields: v.string(), + async* handler() { + yield 'first' + yield 'second' + throw new Error('handler-bork') + }, + } as any) + + const client = bootClient(harness.port) + await new Promise(r => setTimeout(r, 50)) + + const reader = await client.callGenerator('test:gen-throw') + const collected: string[] = [] + let caught: unknown + try { + for await (const chunk of reader) collected.push(chunk) + } + catch (e) { + caught = e + } + expect(collected).toEqual(['first', 'second']) + expect(caught).toBeInstanceOf(Error) + expect((caught as Error).message).toBe('handler-bork') + }) + + it('rejects the call when the generator throws before any yield', async () => { + harness.rpcHost.register({ + name: 'test:gen-throw-immediate', + type: 'generator', + yields: v.string(), + async* handler() { + throw new Error('synthrow') + + yield 'never' + }, + } as any) + + const client = bootClient(harness.port) + await new Promise(r => setTimeout(r, 50)) + + const reader = await client.callGenerator('test:gen-throw-immediate') + const collected: string[] = [] + let caught: unknown + try { + for await (const chunk of reader) collected.push(chunk) + } + catch (e) { + caught = e + } + expect(collected).toEqual([]) + expect((caught as Error).message).toBe('synthrow') + }) + + it('replays buffered yields for a late subscriber within the replay window', async () => { + harness.rpcHost.register({ + name: 'test:gen-replay', + type: 'generator', + yields: v.string(), + replayWindow: 64, + async* handler() { + yield 'one' + yield 'two' + yield 'three' + }, + } as any) + + const client = bootClient(harness.port) + await new Promise(r => setTimeout(r, 50)) + + // The client subscribe lands a few ms after the wrapper allocates the + // sink. The default replay window of 256 (or our 64 here) buffers the + // early yields. + const reader = await client.callGenerator('test:gen-replay') + const collected: string[] = [] + for await (const chunk of reader) collected.push(chunk) + + expect(collected).toEqual(['one', 'two', 'three']) + }) + + it('isolates concurrent generator calls with independent stream ids', async () => { + harness.rpcHost.register({ + name: 'test:gen-isolate', + type: 'generator', + args: [v.object({ tag: v.string() })], + yields: v.string(), + async* handler({ tag }: { tag: string }) { + yield `${tag}-1` + yield `${tag}-2` + }, + } as any) + + const client = bootClient(harness.port) + await new Promise(r => setTimeout(r, 50)) + + const [readerA, readerB] = await Promise.all([ + client.callGenerator('test:gen-isolate', { tag: 'A' }), + client.callGenerator('test:gen-isolate', { tag: 'B' }), + ]) + + const [a, b] = await Promise.all([ + (async () => { + const r: string[] = [] + for await (const c of readerA) r.push(c) + return r + })(), + (async () => { + const r: string[] = [] + for await (const c of readerB) r.push(c) + return r + })(), + ]) + + expect(a).toEqual(['A-1', 'A-2']) + expect(b).toEqual(['B-1', 'B-2']) + expect(readerA.id).not.toBe(readerB.id) + }) + + it('invokeLocalGenerator yields directly without traversing the transport', async () => { + harness.rpcHost.register({ + name: 'test:gen-local', + type: 'generator', + yields: v.string(), + async* handler() { + yield 'x' + yield 'y' + yield 'z' + }, + } as any) + + const iter = await invokeLocalGenerator(harness.rpcHost, 'test:gen-local') + const collected: string[] = [] + for await (const v of iter) collected.push(v) + expect(collected).toEqual(['x', 'y', 'z']) + }) + + it('rejects registration with `agent` set (DF0033)', () => { + expect(() => harness.rpcHost.register({ + name: 'test:gen-agent', + type: 'generator', + yields: v.string(), + agent: { description: 'should fail' }, + async* handler() { yield 'a' }, + } as any)).toThrow(/DF0033/) + }) + + it('rejects registration with `args` and no `yields` (DF0035)', () => { + expect(() => harness.rpcHost.register({ + name: 'test:gen-no-yields', + type: 'generator', + args: [v.object({ q: v.string() })], + async* handler() { yield 'a' }, + } as any)).toThrow(/DF0035/) + }) + + it('rejects registration with `cacheable: true` (DF0036)', () => { + expect(() => harness.rpcHost.register({ + name: 'test:gen-cacheable', + type: 'generator', + cacheable: true, + yields: v.string(), + async* handler() { yield 'a' }, + } as any)).toThrow(/DF0036/) + }) + + it('rejects registration with `dump` (DF0027)', () => { + expect(() => harness.rpcHost.register({ + name: 'test:gen-dump', + type: 'generator', + yields: v.string(), + dump: { records: [] }, + async* handler() { yield 'a' }, + } as any)).toThrow(/DF0027/) + }) + + it('rejects a non-async-iterable handler at runtime (DF0034)', async () => { + harness.rpcHost.register({ + name: 'test:gen-not-iter', + type: 'generator', + yields: v.string(), + // Wrong shape: returns a string instead of yielding it. + handler: (async () => 'oops') as any, + } as any) + + const client = bootClient(harness.port) + await new Promise(r => setTimeout(r, 50)) + + const reader = await client.callGenerator('test:gen-not-iter') + let caught: unknown + try { + for await (const _v of reader) { /* drain */ } + } + catch (e) { + caught = e + } + expect((caught as Error).message).toMatch(/DF0034/) + }) +}) diff --git a/devframe/packages/devframe/src/node/diagnostics.ts b/devframe/packages/devframe/src/node/diagnostics.ts index 75e56907..7992a78f 100644 --- a/devframe/packages/devframe/src/node/diagnostics.ts +++ b/devframe/packages/devframe/src/node/diagnostics.ts @@ -93,6 +93,26 @@ export const diagnostics = defineDiagnostics({ `Streaming channel "${p.channel}" is already registered.`, hint: 'Each channel name must be unique within a context. Pick a different name or reuse the existing channel handle.', }, + DF0033: { + message: (p: { name: string }) => + `Generator RPC function "${p.name}" has \`agent\` set, which is not supported in the current release.`, + hint: 'Remove `agent` from the generator definition. Streaming-MCP exposure is planned as a follow-up.', + }, + DF0034: { + message: (p: { name: string }) => + `Generator RPC function "${p.name}" handler did not return an \`AsyncIterable\`.`, + hint: 'The handler must be declared as `async function*`. Check that no `setup()` overload returns a non-iterable.', + }, + DF0035: { + message: (p: { name: string }) => + `Generator RPC function "${p.name}" declares \`args\` but no \`yields\` schema — yielded value type cannot be inferred.`, + hint: 'Add a `yields` schema (e.g. `yields: v.string()`) for end-to-end type inference of `StreamReader`. Drop both schemas to fall back to `unknown` yields.', + }, + DF0036: { + message: (p: { name: string, option: string }) => + `Generator RPC function "${p.name}" sets \`${p.option}\`, which has no effect on generator-typed functions.`, + hint: 'Generators are streaming primitives — `cacheable`, `dump`, and `jsonSerializable: true` do not apply. Remove the option, or change the function type.', + }, }, }) diff --git a/devframe/packages/devframe/src/node/host-functions.ts b/devframe/packages/devframe/src/node/host-functions.ts index 66c4a301..3759b0ed 100644 --- a/devframe/packages/devframe/src/node/host-functions.ts +++ b/devframe/packages/devframe/src/node/host-functions.ts @@ -4,6 +4,7 @@ import type { AsyncLocalStorage } from 'node:async_hooks' import { RpcFunctionsCollectorBase } from 'devframe/rpc' import { createDebug } from 'obug' import { logger } from './diagnostics' +import { attachRpcGenerators } from './rpc-generators' import { createRpcSharedStateServerHost } from './rpc-shared-state' import { createRpcStreamingServerHost } from './rpc-streaming' @@ -21,6 +22,10 @@ export class RpcFunctionsHost extends RpcFunctionsCollectorBase