feat: wire port exposure WebSocket proxy + MCP gateway lifecycle#83
Merged
feat: wire port exposure WebSocket proxy + MCP gateway lifecycle#83
Conversation
Sprint 2 wiring — connects existing but disconnected components: - Worker proxy: use getSessionConnection() instead of type assertion, add WebSocket upgrade handling via Hono upgradeWebSocket - Control plane expose: add WebSocket proxy route (/s/:sessionId/ws/*), path-prefix port resolution, pass upgradeWebSocket from deps - Session creation: freeze exposed port URLs at dispatch time via generateExposedUrls() — URLs available immediately to clients - Worker server: instantiate createMcpGateway() with env config (MCP_GATEWAY_ENABLED, MCP_GATEWAY_PORT, etc.), pass to executor - Worker server: add createBunWebSocket() + export websocket handler Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
📝 WalkthroughWalkthroughThis PR adds WebSocket proxying support across the control plane and worker services, introduces intelligent port resolution logic for exposed services, integrates MCP gateway functionality, and refactors session connection handling to use a dedicated getter method. Changes
Sequence DiagramssequenceDiagram
participant Client
participant ControlPlane as Control Plane<br/>(expose.ts)
participant Worker as Worker<br/>(proxy.ts)
participant Backend as Backend Service<br/>(exposed port)
rect rgba(100, 150, 200, 0.5)
Note over Client,Backend: WebSocket Proxy Flow
Client->>ControlPlane: GET /s/:sessionId/ws/path
ControlPlane->>ControlPlane: Validate session & resolve targetPort
ControlPlane->>Worker: WS upgrade to ws://guestIp:targetPort/path
Worker->>Backend: WS connection established
par Bidirectional Messaging
Client->>Worker: message
Worker->>Backend: forward message
Backend->>Worker: response
Worker->>Client: forward response
end
Backend-->>Worker: close/error
Worker-->>Client: mirror close/error
end
sequenceDiagram
participant Client
participant ControlPlane as Control Plane<br/>(expose.ts)
participant Worker as Worker<br/>(proxy.ts)
participant Backend as Backend Service
rect rgba(150, 100, 200, 0.5)
Note over Client,Backend: HTTP Proxy with Port Resolution
Client->>ControlPlane: GET /s/:sessionId/app/path?query=1
ControlPlane->>ControlPlane: resolvePort(path) matches pathPrefix<br/>returns targetPort
ControlPlane->>Worker: GET http://guestIp:targetPort/path?query=1
Worker->>Backend: forward request with query
Backend-->>Worker: response
Worker-->>ControlPlane: response
ControlPlane-->>Client: response
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~50 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 2 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (2 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
Deploying getpaws with
|
| Latest commit: |
94233e5
|
| Status: | ✅ Deploy successful! |
| Preview URL: | https://06132f46.getpaws-6m4.pages.dev |
| Branch Preview URL: | https://feat-sprint2-port-exposure-a.getpaws-6m4.pages.dev |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
apps/worker/src/server.ts (1)
126-132:⚠️ Potential issue | 🟠 MajorMCP session registration is still disabled here.
apps/worker/src/session/executor.tsonly enters the MCP setup path when bothconfig.mcpGatewayandconfig.mcpServerStoreare present. This call site passes onlymcpGateway, so sessions that request MCP servers will skip agentgateway registration entirely and never receive a gateway URL.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@apps/worker/src/server.ts` around lines 126 - 132, The executor is being created without the MCP server store so executor logic in session/executor.ts (which requires both config.mcpGateway and config.mcpServerStore to enable MCP registration) never enters the MCP setup path; update the createExecutor call to pass the MCP server store as well (i.e., include the mcpServerStore property alongside mcpGateway when available) so that createExecutor(...) receives both mcpGateway and mcpServerStore and the MCP agentgateway registration logic runs.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@apps/control-plane/src/app.ts`:
- Around line 2047-2053: The persisted exposedPorts URLs are built without the
fleet/public domain; update dispatchSession to pass the real fleet/public domain
through and use it when calling generateExposedUrls, so
generateExposedUrls(sessionId, expose, fleetDomain) (or equivalent) is invoked
instead of the two-arg form; ensure the same fleetDomain value is threaded into
the code path that computes exposedPorts before calling
sessionStore.updateStatus(sessionId, 'running', { startedAt: ..., exposedPorts
}), and update any call-sites that construct exposed URLs (e.g.,
generateExposedUrls and the expose route logic) to accept and use the
fleet/public domain rather than defaulting to localhost.
In `@apps/control-plane/src/routes/expose.ts`:
- Around line 197-239: The onMessage handler drops client frames if backendWs is
not yet OPEN; fix by adding a small send-queue: create a messages array scoped
alongside backendWs and, in onOpen (backendWs 'open' listener), flush the queued
frames to backendWs (preserving string vs ArrayBuffer/Uint8Array handling) and
clear the queue; update onMessage to push frames onto the queue when backendWs
is CONNECTING (or undefined) instead of dropping them, and ensure that backendWs
'close' and 'error' handlers clear the queue and optionally reject/close
clientWs with an error so buffered frames are not leaked; reference backendWs,
onOpen/onMessage, and the backendWs 'open'/'close'/'error' listeners.
- Around line 92-106: resolvePort currently always falls back to expose[0] which
masks the case where all exposures are path-scoped and prevents returning
undefined; update resolvePort so that after trying longest pathPrefix matches it
only returns a default port if there exists an exposure that is not path-scoped
(e.g. has no pathPrefix or pathPrefix === '/'); otherwise return undefined.
Locate the resolvePort function and adjust the fallback logic (the withPrefix
computation and final return) to search expose for a non-scoped entry and return
its port, or return undefined when none exists.
In `@apps/worker/src/routes/proxy.ts`:
- Around line 48-87: onMessage currently drops frames while backendWs is still
connecting; instead implement a short in-memory queue in the surrounding scope
(e.g., const pendingMessages: Array<{data:any}> = []) and push incoming frames
in onMessage when backendWs.readyState !== WebSocket.OPEN, then in
backendWs.addEventListener('open', ...) flush the queued frames by sending them
in order via backendWs.send and clear the queue; ensure you still handle
ArrayBuffer vs string encodings consistently and handle cases where backendWs
later closes/errors by clearing the queue and closing clientWs as done in
backendWs 'close'/'error' handlers.
---
Outside diff comments:
In `@apps/worker/src/server.ts`:
- Around line 126-132: The executor is being created without the MCP server
store so executor logic in session/executor.ts (which requires both
config.mcpGateway and config.mcpServerStore to enable MCP registration) never
enters the MCP setup path; update the createExecutor call to pass the MCP server
store as well (i.e., include the mcpServerStore property alongside mcpGateway
when available) so that createExecutor(...) receives both mcpGateway and
mcpServerStore and the MCP agentgateway registration logic runs.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 044294fc-1b6f-4599-94d0-362cafac1c6a
📒 Files selected for processing (5)
apps/control-plane/src/app.tsapps/control-plane/src/routes/expose.tsapps/worker/src/routes.tsapps/worker/src/routes/proxy.tsapps/worker/src/server.ts
Comment on lines
+2047
to
+2053
| // Freeze exposed port URLs at dispatch time — URLs are available immediately | ||
| const expose = request.network?.expose ?? []; | ||
| const exposedPorts = expose.length > 0 ? generateExposedUrls(sessionId, expose) : undefined; | ||
|
|
||
| sessionStore.updateStatus(sessionId, 'running', { | ||
| startedAt: new Date().toISOString(), | ||
| ...(exposedPorts ? { exposedPorts } : {}), |
There was a problem hiding this comment.
These frozen exposed URLs still default to localhost.
Lines 2047-2049 call generateExposedUrls(sessionId, expose) without a fleetDomain. In apps/control-plane/src/routes/expose.ts, that falls back to http://s-<id>.localhost:3000/..., so every persisted exposedPorts[].url will be wrong outside local dev. Please thread the actual fleet/public domain into dispatchSession() before storing these URLs.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@apps/control-plane/src/app.ts` around lines 2047 - 2053, The persisted
exposedPorts URLs are built without the fleet/public domain; update
dispatchSession to pass the real fleet/public domain through and use it when
calling generateExposedUrls, so generateExposedUrls(sessionId, expose,
fleetDomain) (or equivalent) is invoked instead of the two-arg form; ensure the
same fleetDomain value is threaded into the code path that computes exposedPorts
before calling sessionStore.updateStatus(sessionId, 'running', { startedAt: ...,
exposedPorts }), and update any call-sites that construct exposed URLs (e.g.,
generateExposedUrls and the expose route logic) to accept and use the
fleet/public domain rather than defaulting to localhost.
Comment on lines
+92
to
+106
| function resolvePort( | ||
| path: string, | ||
| expose: Array<{ port: number; pathPrefix?: string }>, | ||
| ): number | undefined { | ||
| if (expose.length === 0) return undefined; | ||
|
|
||
| // Try path-prefix matching first (longest match wins) | ||
| const withPrefix = expose | ||
| .filter((e) => e.pathPrefix && e.pathPrefix !== '/' && path.startsWith(e.pathPrefix)) | ||
| .sort((a, b) => (b.pathPrefix?.length ?? 0) - (a.pathPrefix?.length ?? 0)); | ||
|
|
||
| if (withPrefix.length > 0) return withPrefix[0]!.port; | ||
|
|
||
| // Default to the first exposed port | ||
| return expose[0]!.port; |
There was a problem hiding this comment.
Don't fall back to the first port when no prefix matches.
When every exposure is path-scoped, resolvePort() can never return undefined because it always falls back to expose[0]. That makes the new 403 branch unreachable and can send /s/:sessionId/<unknown> or /s/:sessionId/ws/<unknown> to the wrong service instead of rejecting it.
Suggested fix
function resolvePort(
path: string,
expose: Array<{ port: number; pathPrefix?: string }>,
): number | undefined {
if (expose.length === 0) return undefined;
// Try path-prefix matching first (longest match wins)
const withPrefix = expose
- .filter((e) => e.pathPrefix && e.pathPrefix !== '/' && path.startsWith(e.pathPrefix))
+ .filter((e) => {
+ const prefix = e.pathPrefix;
+ return (
+ prefix !== undefined &&
+ prefix !== '/' &&
+ (path === prefix || path.startsWith(`${prefix}/`))
+ );
+ })
.sort((a, b) => (b.pathPrefix?.length ?? 0) - (a.pathPrefix?.length ?? 0));
if (withPrefix.length > 0) return withPrefix[0]!.port;
- // Default to the first exposed port
- return expose[0]!.port;
+ // Only fall back to an explicit default exposure
+ return expose.find((e) => !e.pathPrefix || e.pathPrefix === '/')?.port;
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| function resolvePort( | |
| path: string, | |
| expose: Array<{ port: number; pathPrefix?: string }>, | |
| ): number | undefined { | |
| if (expose.length === 0) return undefined; | |
| // Try path-prefix matching first (longest match wins) | |
| const withPrefix = expose | |
| .filter((e) => e.pathPrefix && e.pathPrefix !== '/' && path.startsWith(e.pathPrefix)) | |
| .sort((a, b) => (b.pathPrefix?.length ?? 0) - (a.pathPrefix?.length ?? 0)); | |
| if (withPrefix.length > 0) return withPrefix[0]!.port; | |
| // Default to the first exposed port | |
| return expose[0]!.port; | |
| function resolvePort( | |
| path: string, | |
| expose: Array<{ port: number; pathPrefix?: string }>, | |
| ): number | undefined { | |
| if (expose.length === 0) return undefined; | |
| // Try path-prefix matching first (longest match wins) | |
| const withPrefix = expose | |
| .filter((e) => { | |
| const prefix = e.pathPrefix; | |
| return ( | |
| prefix !== undefined && | |
| prefix !== '/' && | |
| (path === prefix || path.startsWith(`${prefix}/`)) | |
| ); | |
| }) | |
| .sort((a, b) => (b.pathPrefix?.length ?? 0) - (a.pathPrefix?.length ?? 0)); | |
| if (withPrefix.length > 0) return withPrefix[0]!.port; | |
| // Only fall back to an explicit default exposure | |
| return expose.find((e) => !e.pathPrefix || e.pathPrefix === '/')?.port; | |
| } |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@apps/control-plane/src/routes/expose.ts` around lines 92 - 106, resolvePort
currently always falls back to expose[0] which masks the case where all
exposures are path-scoped and prevents returning undefined; update resolvePort
so that after trying longest pathPrefix matches it only returns a default port
if there exists an exposure that is not path-scoped (e.g. has no pathPrefix or
pathPrefix === '/'); otherwise return undefined. Locate the resolvePort function
and adjust the fallback logic (the withPrefix computation and final return) to
search expose for a non-scoped entry and return its port, or return undefined
when none exists.
Comment on lines
+197
to
+239
| return { | ||
| onOpen(_evt, clientWs) { | ||
| log.debug('WebSocket proxy opening', { | ||
| sessionId, | ||
| port: targetPort, | ||
| backend: backendUrl, | ||
| }); | ||
|
|
||
| backendWs = new WebSocket(backendUrl); | ||
|
|
||
| backendWs.addEventListener('open', () => { | ||
| log.debug('Backend WebSocket connected', { sessionId, port: targetPort }); | ||
| }); | ||
|
|
||
| backendWs.addEventListener('message', (evt) => { | ||
| try { | ||
| if (typeof evt.data === 'string') { | ||
| clientWs.send(evt.data); | ||
| } else if (evt.data instanceof ArrayBuffer) { | ||
| clientWs.send(new Uint8Array(evt.data)); | ||
| } | ||
| } catch { | ||
| // Client disconnected | ||
| } | ||
| }); | ||
|
|
||
| backendWs.addEventListener('close', (evt) => { | ||
| clientWs.close(evt.code, evt.reason); | ||
| }); | ||
|
|
||
| backendWs.addEventListener('error', () => { | ||
| clientWs.close(4502, 'Backend connection failed'); | ||
| }); | ||
| }, | ||
|
|
||
| onMessage(evt, _ws) { | ||
| if (backendWs?.readyState === WebSocket.OPEN) { | ||
| if (typeof evt.data === 'string') { | ||
| backendWs.send(evt.data); | ||
| } else if (evt.data instanceof ArrayBuffer) { | ||
| backendWs.send(evt.data); | ||
| } | ||
| } |
There was a problem hiding this comment.
Buffer client frames until the worker WebSocket is open.
The client connection is accepted before backendWs reaches OPEN, but onMessage() drops anything sent while it is still CONNECTING. Clients that send an init/auth frame immediately after open will intermittently fail on this hop.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@apps/control-plane/src/routes/expose.ts` around lines 197 - 239, The
onMessage handler drops client frames if backendWs is not yet OPEN; fix by
adding a small send-queue: create a messages array scoped alongside backendWs
and, in onOpen (backendWs 'open' listener), flush the queued frames to backendWs
(preserving string vs ArrayBuffer/Uint8Array handling) and clear the queue;
update onMessage to push frames onto the queue when backendWs is CONNECTING (or
undefined) instead of dropping them, and ensure that backendWs 'close' and
'error' handlers clear the queue and optionally reject/close clientWs with an
error so buffered frames are not leaked; reference backendWs, onOpen/onMessage,
and the backendWs 'open'/'close'/'error' listeners.
Comment on lines
+48
to
+87
| return { | ||
| onOpen(_evt, clientWs) { | ||
| log.debug('WebSocket proxy opening', { sessionId, port, target: targetWsUrl }); | ||
|
|
||
| backendWs = new WebSocket(targetWsUrl); | ||
|
|
||
| backendWs.addEventListener('open', () => { | ||
| log.debug('Backend WebSocket connected', { sessionId, port }); | ||
| }); | ||
|
|
||
| backendWs.addEventListener('message', (evt) => { | ||
| try { | ||
| if (typeof evt.data === 'string') { | ||
| clientWs.send(evt.data); | ||
| } else if (evt.data instanceof ArrayBuffer) { | ||
| clientWs.send(new Uint8Array(evt.data)); | ||
| } | ||
| } catch { | ||
| // Client disconnected | ||
| } | ||
| }); | ||
|
|
||
| backendWs.addEventListener('close', (evt) => { | ||
| clientWs.close(evt.code, evt.reason); | ||
| }); | ||
|
|
||
| backendWs.addEventListener('error', () => { | ||
| clientWs.close(4502, 'Backend connection failed'); | ||
| }); | ||
| }, | ||
|
|
||
| onMessage(evt, _ws) { | ||
| if (backendWs?.readyState === WebSocket.OPEN) { | ||
| if (typeof evt.data === 'string') { | ||
| backendWs.send(evt.data); | ||
| } else if (evt.data instanceof ArrayBuffer) { | ||
| backendWs.send(evt.data); | ||
| } | ||
| } | ||
| }, |
There was a problem hiding this comment.
Buffer frames while the VM WebSocket is still connecting.
This hop has the same race: onMessage() silently discards frames until backendWs.readyState === WebSocket.OPEN. Even if the control-plane buffers correctly, the worker → VM bridge can still lose the initial protocol message and break the proxied session.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@apps/worker/src/routes/proxy.ts` around lines 48 - 87, onMessage currently
drops frames while backendWs is still connecting; instead implement a short
in-memory queue in the surrounding scope (e.g., const pendingMessages:
Array<{data:any}> = []) and push incoming frames in onMessage when
backendWs.readyState !== WebSocket.OPEN, then in
backendWs.addEventListener('open', ...) flush the queued frames by sending them
in order via backendWs.send and clear the queue; ensure you still handle
ArrayBuffer vs string encodings consistently and handle cases where backendWs
later closes/errors by clearing the queue and closing clientWs as done in
backendWs 'close'/'error' handlers.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Sprint 2 wiring — connects existing but disconnected components for collaborative sessions:
upgradeWebSocket. Worker route:/v1/sessions/:id/proxy/:port/ws/*, control plane route:/s/:sessionId/ws/*. Enables HMR, dev servers, and live terminals inside VMs.generateExposedUrls()is called insidedispatchSession()so all session creation paths (API, daemon triggers, webhooks) get URLs frozen immediately. Clients can read exposed port URLs right away without waiting for the worker.createMcpGateway()instantiated inserver.tswith env config (MCP_GATEWAY_ENABLED,MCP_GATEWAY_PORT,MCP_GATEWAY_CONFIG_PATH), passed to executor. When sessions request MCP servers, agentgateway config is written/cleaned up automatically.(session as { allocation? }).allocationtype assertion hack withexecutor.getSessionConnection()(proper SessionConnection interface). Added path-prefix port resolution in control plane.createBunWebSocket()to worker'sserver.ts, exportedwebsockethandler forBun.serve.Plane items
PAWS-137, PAWS-138, PAWS-139, PAWS-140, PAWS-141, PAWS-142
Test plan
network.exposegets URLs in response immediatelyMCP_GATEWAY_ENABLED=true🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
Bug Fixes