Repository navigation
feat(chat-runtime): tRPC router over commands + WebSocket sink adapter - #6192
Conversation
|
Important Review skippedAuto reviews are disabled on base/target branches other than the default branch. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: defaults Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
🧹 Preview Cleanup CompleteThe following preview resources have been cleaned up:
Thank you for your contribution! 🎉 |
Greptile SummaryAdds the chat runtime’s transport surface through a tRPC command router and a structural WebSocket-to-Sink adapter.
Confidence Score: 4/5The createSession transport contract needs correction before merging because payloads valid under the public protocol are currently rejected. The new router exposes createSession using a stricter runtime-only schema requiring cwd, while the public protocol schema omits that field, breaking protocol-conforming callers at validation. Files Needing Attention: packages/chat-runtime/src/router/router/router.ts and packages/chat-runtime/src/router/router/router.test.ts
|
| Filename | Overview |
|---|---|
| packages/chat-runtime/src/router/router/router.ts | Introduces the tRPC command surface, but createSession requires a runtime-only cwd field absent from the public protocol schema. |
| packages/chat-runtime/src/router/wsSink/wsSink.ts | Adds an idempotent structural WebSocket sink that serializes envelopes and suppresses sends after observed closure. |
| packages/chat-runtime/src/router/router/router.test.ts | Covers routing, validation, deduplication, and filtering, but constructs createSession inputs with the extra runtime-only cwd field. |
| packages/chat-runtime/src/router/wsSink/wsSink.test.ts | Covers serialization, idempotent sink closure, prior close-handler preservation, and peer-initiated closure. |
Sequence Diagram
sequenceDiagram
participant Client
participant Router as createChatRouter
participant Runtime as ChatRuntime.commands
participant Hub as SubscriptionHub
participant Sink as createWsSink
participant WS as WebSocket
Client->>Router: tRPC command/query
Router->>Router: Validate protocol input
Router->>Runtime: Invoke command
Runtime-->>Router: Command result
Router-->>Client: Typed response
Runtime->>Hub: Publish envelope
Hub->>Sink: send(envelope)
Sink->>WS: send(JSON.stringify(envelope))
Prompt To Fix All With AI
### Issue 1
packages/chat-runtime/src/router/router/router.ts:21
**Create-session schema breaks protocol**
When a client constructs `createSession` input from the public protocol schema, this procedure requires the additional runtime-only `cwd` field and rejects the otherwise valid request as `BAD_REQUEST`, preventing the client from starting a session.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.Reviews (1): Last reviewed commit: "feat(chat-runtime): tRPC router over com..." | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| export function createChatRouter(runtime: ChatRuntime) { | ||
| return t.router({ | ||
| createSession: t.procedure | ||
| .input(createSessionCommandSchema) |
There was a problem hiding this comment.
When a client constructs createSession input from the public protocol schema, this procedure requires the additional runtime-only cwd field and rejects the otherwise valid request as BAD_REQUEST, preventing the client from starting a session.
Prompt To Fix With AI
This is a comment left during a code review.
Path: packages/chat-runtime/src/router/router/router.ts
Line: 21
Comment:
**Create-session schema breaks protocol**
When a client constructs `createSession` input from the public protocol schema, this procedure requires the additional runtime-only `cwd` field and rejects the otherwise valid request as `BAD_REQUEST`, preventing the client from starting a session.
---
For each issue above, determine whether it is valid and should be fixed. If so, fix it directly.![cubic-dev-ai[bot]](https://avatars.githubusercontent.com/in/1082092?s=60&v=4){kind=small}
There was a problem hiding this comment.
2 issues found across 11 files
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="packages/chat-runtime/src/router/wsSink/wsSink.ts">
<violation number="1" location="packages/chat-runtime/src/router/wsSink/wsSink.ts:12">
P2: `createWsSink` permanently hijacks the single `socket.onclose` slot: it snapshots whatever handler is present at creation and overwrites it. Because the README plans for host-service to mount `createWsSink` per socket, the socket will already carry host-managed close handling, and if anything assigns `socket.onclose` after the sink is created (or the host needs to re-parent it later), the sink's `open = false` tracking is silently disabled — peer-close won't flip the flag, so `send()` keeps writing to a dead socket and the original handler is dropped. The design is only correct under the implicit contract that `onclose` is set before the sink is created. Consider chaining without taking ownership (invoke the current `socket.onclose` at close time, or return an object that lets the caller restore it), and guard `send` so a transport-level throw doesn't propagate into the subscription publish path.</violation>
<violation number="2" location="packages/chat-runtime/src/router/wsSink/wsSink.ts:19">
P1: The local `open` flag does not protect sends during the WebSocket ready-state race or bound a slow client's send queue. The adapter should apply an equivalent `readyState`/send-error check and slow-consumer limit before calling `socket.send`, otherwise a closed or stalled peer can throw into runtime publication or consume unbounded host memory.</violation>
</file>
Architecture diagram
sequenceDiagram
participant Client
participant Router as tRPC Router
participant Runtime as ChatRuntime
participant WsSink as WebSocket Sink Adapter
participant Socket as Structural WebSocket
Note over Client, Runtime: NEW: Request/Command Flow
Client->>Router: Mutation/Query (e.g., prompt, getSession)
Router->>Router: NEW: Validate input (Zod schemas)
Router->>Runtime: NEW: Invoke commands.method(input)
alt Validation or Logic Error
Runtime-->>Router: Throw error
Router-->>Client: tRPC Error (BAD_REQUEST/INTERNAL)
else Success
Runtime-->>Router: Command Result (e.g., sessionId, epoch)
Router-->>Client: Typed Inference Response
end
Note over Client, Socket: NEW: Event Streaming Flow
Client->>Socket: Establish Connection
Note over Runtime, WsSink: Host-service setup
Runtime->>WsSink: NEW: createWsSink(socket)
Runtime->>Runtime: subscribe(sessionId, wsSink)
loop On Runtime Events
Runtime->>WsSink: send(envelope)
opt Sink is open
WsSink->>WsSink: NEW: JSON.stringify(envelope)
WsSink->>Socket: send(data string)
end
end
Note over WsSink, Socket: Lifecycle Management
alt Socket closed by peer
Socket->>WsSink: onclose()
WsSink->>WsSink: CHANGED: Set open = false
else Runtime disposes subscription
Runtime->>WsSink: close()
WsSink->>Socket: NEW: socket.close()
WsSink->>WsSink: Set open = false
end
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| return { | ||
| send(envelope) { | ||
| if (!open) return; | ||
| socket.send(JSON.stringify(envelope)); |
There was a problem hiding this comment.
P1: The local open flag does not protect sends during the WebSocket ready-state race or bound a slow client's send queue. The adapter should apply an equivalent readyState/send-error check and slow-consumer limit before calling socket.send, otherwise a closed or stalled peer can throw into runtime publication or consume unbounded host memory.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/chat-runtime/src/router/wsSink/wsSink.ts, line 19:
<comment>The local `open` flag does not protect sends during the WebSocket ready-state race or bound a slow client's send queue. The adapter should apply an equivalent `readyState`/send-error check and slow-consumer limit before calling `socket.send`, otherwise a closed or stalled peer can throw into runtime publication or consume unbounded host memory.</comment>
<file context>
@@ -0,0 +1,27 @@
+ return {
+ send(envelope) {
+ if (!open) return;
+ socket.send(JSON.stringify(envelope));
+ },
+ close() {
</file context>
| export function createWsSink(socket: WsSinkSocket): Sink { | ||
| let open = true; | ||
| const previous = socket.onclose; | ||
| socket.onclose = () => { |
There was a problem hiding this comment.
P2: createWsSink permanently hijacks the single socket.onclose slot: it snapshots whatever handler is present at creation and overwrites it. Because the README plans for host-service to mount createWsSink per socket, the socket will already carry host-managed close handling, and if anything assigns socket.onclose after the sink is created (or the host needs to re-parent it later), the sink's open = false tracking is silently disabled — peer-close won't flip the flag, so send() keeps writing to a dead socket and the original handler is dropped. The design is only correct under the implicit contract that onclose is set before the sink is created. Consider chaining without taking ownership (invoke the current socket.onclose at close time, or return an object that lets the caller restore it), and guard send so a transport-level throw doesn't propagate into the subscription publish path.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/chat-runtime/src/router/wsSink/wsSink.ts, line 12:
<comment>`createWsSink` permanently hijacks the single `socket.onclose` slot: it snapshots whatever handler is present at creation and overwrites it. Because the README plans for host-service to mount `createWsSink` per socket, the socket will already carry host-managed close handling, and if anything assigns `socket.onclose` after the sink is created (or the host needs to re-parent it later), the sink's `open = false` tracking is silently disabled — peer-close won't flip the flag, so `send()` keeps writing to a dead socket and the original handler is dropped. The design is only correct under the implicit contract that `onclose` is set before the sink is created. Consider chaining without taking ownership (invoke the current `socket.onclose` at close time, or return an object that lets the caller restore it), and guard `send` so a transport-level throw doesn't propagate into the subscription publish path.</comment>
<file context>
@@ -0,0 +1,27 @@
+export function createWsSink(socket: WsSinkSocket): Sink {
+ let open = true;
+ const previous = socket.onclose;
+ socket.onclose = () => {
+ open = false;
+ previous?.();
</file context>
There was a problem hiding this comment.
4 issues found across 7 files (changes from recent commits).
Prompt for AI agents (unresolved issues)
Check if these issues are valid — if so, understand the root cause of each and fix them. If appropriate, use sub-agents to investigate and fix each issue separately.
<file name="packages/chat-runtime/src/router/wsSink/wsSink.ts">
<violation number="1" location="packages/chat-runtime/src/router/wsSink/wsSink.ts:26">
P2: A failed socket shutdown is silently hidden: after `socket.close()` throws, this sink is permanently closed and callers receive no diagnostic even though the transport may still be open. Preserve best-effort shutdown if needed, but report the exception so leaked connections are diagnosable.</violation>
</file>
<file name="packages/chat-runtime/src/router/router/router.test.ts">
<violation number="1" location="packages/chat-runtime/src/router/router/router.test.ts:35">
P3: Each newCaller() creates a temp working directory via mkdtempSync that is never cleaned up, so every test in this file leaks a directory on disk. Track the created cwd and remove it (e.g. rmSync(cwd, { recursive: true, force: true }) in a finally/afterEach) to keep test runs from accumulating temp state.</violation>
</file>
<file name="packages/chat-runtime/src/router/router/router.ts">
<violation number="1" location="packages/chat-runtime/src/router/router/router.ts:23">
P2: Valid session IDs containing a line break bypass the typed not-running mapping and surface as `INTERNAL_SERVER_ERROR`; use a multiline-safe capture (or structured command errors) so every protocol-valid ID receives the intended `NOT_FOUND`/`CONFLICT` code.</violation>
<violation number="2" location="packages/chat-runtime/src/router/router/router.ts:62">
P2: The `createSession` procedure resolves the host-supplied cwd outside the `guarded` error-mapping wrapper, so a failure in `options.resolveCwd(workspaceId)` (for example, a closed or inaccessible workspace) bypasses `mapCommandError` and reaches tRPC as a raw error/generic internal error rather than a typed command error. Since this PR's goal is to surface typed command errors from every procedure, the cwd-resolution step should be covered by the same guard so the client receives a meaningful code instead of an opaque 500.</violation>
</file>
Reply with feedback, questions, or to request a fix.
Re-trigger cubic
| if (socket.onclose === handleClose) socket.onclose = previous; | ||
| try { | ||
| socket.close(); | ||
| } catch {} |
There was a problem hiding this comment.
P2: A failed socket shutdown is silently hidden: after socket.close() throws, this sink is permanently closed and callers receive no diagnostic even though the transport may still be open. Preserve best-effort shutdown if needed, but report the exception so leaked connections are diagnosable.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/chat-runtime/src/router/wsSink/wsSink.ts, line 26:
<comment>A failed socket shutdown is silently hidden: after `socket.close()` throws, this sink is permanently closed and callers receive no diagnostic even though the transport may still be open. Preserve best-effort shutdown if needed, but report the exception so leaked connections are diagnosable.</comment>
<file context>
@@ -3,25 +3,38 @@ import type { Sink } from "../../stream";
+ if (socket.onclose === handleClose) socket.onclose = previous;
+ try {
+ socket.close();
+ } catch {}
};
+
</file context>
| }; | ||
|
|
||
| const UNKNOWN_HARNESS = /^unknown harness /; | ||
| const NOT_RUNNING = /^chat session (.+) is not running$/; |
There was a problem hiding this comment.
P2: Valid session IDs containing a line break bypass the typed not-running mapping and surface as INTERNAL_SERVER_ERROR; use a multiline-safe capture (or structured command errors) so every protocol-valid ID receives the intended NOT_FOUND/CONFLICT code.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/chat-runtime/src/router/router/router.ts, line 23:
<comment>Valid session IDs containing a line break bypass the typed not-running mapping and surface as `INTERNAL_SERVER_ERROR`; use a multiline-safe capture (or structured command errors) so every protocol-valid ID receives the intended `NOT_FOUND`/`CONFLICT` code.</comment>
<file context>
@@ -1,41 +1,87 @@
+};
+
+const UNKNOWN_HARNESS = /^unknown harness /;
+const NOT_RUNNING = /^chat session (.+) is not running$/;
+
+function mapCommandError(runtime: ChatRuntime, error: unknown): unknown {
</file context>
| const NOT_RUNNING = /^chat session (.+) is not running$/; | |
| const NOT_RUNNING = /^chat session ([\s\S]+) is not running$/; |
| createSession: t.procedure | ||
| .input(createSessionInputSchema) | ||
| .mutation(async ({ input }) => { | ||
| const cwd = await options.resolveCwd(input.workspaceId); |
There was a problem hiding this comment.
P2: The createSession procedure resolves the host-supplied cwd outside the guarded error-mapping wrapper, so a failure in options.resolveCwd(workspaceId) (for example, a closed or inaccessible workspace) bypasses mapCommandError and reaches tRPC as a raw error/generic internal error rather than a typed command error. Since this PR's goal is to surface typed command errors from every procedure, the cwd-resolution step should be covered by the same guard so the client receives a meaningful code instead of an opaque 500.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/chat-runtime/src/router/router/router.ts, line 62:
<comment>The `createSession` procedure resolves the host-supplied cwd outside the `guarded` error-mapping wrapper, so a failure in `options.resolveCwd(workspaceId)` (for example, a closed or inaccessible workspace) bypasses `mapCommandError` and reaches tRPC as a raw error/generic internal error rather than a typed command error. Since this PR's goal is to surface typed command errors from every procedure, the cwd-resolution step should be covered by the same guard so the client receives a meaningful code instead of an opaque 500.</comment>
<file context>
@@ -1,41 +1,87 @@
- .mutation(({ input }) => runtime.commands.createSession(input)),
+ .input(createSessionInputSchema)
+ .mutation(async ({ input }) => {
+ const cwd = await options.resolveCwd(input.workspaceId);
+ return guarded(() => runtime.commands.createSession({ ...input, cwd }));
+ }),
</file context>
| function newCaller() { | ||
| const { harnesses, adapters } = fakeHarnessRegistry(SCRIPT); | ||
| const runtime = createTestRuntime({ harnesses }); | ||
| const cwd = mkdtempSync(join(tmpdir(), "chat-router-cwd-")); |
There was a problem hiding this comment.
P3: Each newCaller() creates a temp working directory via mkdtempSync that is never cleaned up, so every test in this file leaks a directory on disk. Track the created cwd and remove it (e.g. rmSync(cwd, { recursive: true, force: true }) in a finally/afterEach) to keep test runs from accumulating temp state.
Prompt for AI agents
Check if this issue is valid — if so, understand the root cause and fix it. At packages/chat-runtime/src/router/router/router.test.ts, line 35:
<comment>Each newCaller() creates a temp working directory via mkdtempSync that is never cleaned up, so every test in this file leaks a directory on disk. Track the created cwd and remove it (e.g. rmSync(cwd, { recursive: true, force: true }) in a finally/afterEach) to keep test runs from accumulating temp state.</comment>
<file context>
@@ -29,16 +32,28 @@ const SCRIPT: FakeHarnessScript = {
const runtime = createTestRuntime({ harnesses });
- const caller = createChatCallerFactory(createChatRouter(runtime))({});
- return { runtime, caller, adapterCount: () => adapters.length };
+ const cwd = mkdtempSync(join(tmpdir(), "chat-router-cwd-"));
+ const resolvedWorkspaceIds: string[] = [];
+ const router = createChatRouter(runtime, {
</file context>
createChatRouter(runtime) exposes the eight protocol commands as tRPC procedures on a package-owned initTRPC instance; createWsSink adapts a structural WebSocket into the Sink runtime.subscribe expects, serializing envelopes as JSON for host-service's stream route. Claude-Session: https://claude.ai/code/session_01XRFCEPsfmJHu2SenDKYJ2h
…mand errors createSession now binds the pure protocol schema and resolves cwd via a host callback (clients never send it); wsSink chains and restores the previous onclose handler, is assignable from a native WebSocket, and a throwing send closes the sink; known command errors map to BAD_REQUEST/NOT_FOUND/CONFLICT instead of INTERNAL_SERVER_ERROR. Claude-Session: https://claude.ai/code/session_01XRFCEPsfmJHu2SenDKYJ2h
0607f5d to
a3cb2b6
Compare
Stacked on #6167. The package now owns its transport surface: createChatRouter(runtime) (own initTRPC, procedures close over the runtime, inputs are exactly the protocol schemas, ChatRouter type exported for end-to-end inference) + createWsSink turning any structural WebSocket into a runtime Sink. Host-service's future mount = createChatRouter(runtime) at a path + wsSink per socket. 7 new tests (57 total green).
https://claude.ai/code/session_01XRFCEPsfmJHu2SenDKYJ2h
Summary by cubic
Adds a tRPC router over chat commands and a WebSocket sink adapter for streaming envelopes. The router resolves cwd on the host and returns typed tRPC errors for common failures.
New Features
createChatRouter(runtime, { resolveCwd })exposes 8 protocol commands as tRPC procedures withzodinputs from@superset/chat/protocol; exportsChatRouter,ChatRouterOptions, andcreateChatCallerFactory.createSessioncallsresolveCwd(workspaceId); clients never sendcwd.BAD_REQUEST; not-running sessions →NOT_FOUND/CONFLICT; invalid inputs rejected byzod.createWsSink(socket)adapts any structural WebSocket to the runtimeSink, JSON-serializes envelopes, chains/restoresonclose, closes on send error, and is idempotent on peer/host close; exportsWsSinkSocket.Dependencies
@trpc/server@11.16.0topackages/chat-runtime.Written for commit a3cb2b6. Summary will update on new commits.