the universal sandbox runtime for agents and humans.
pocketenv.io
sandbox
openclaw
agent
claude-code
vercel-sandbox
deno-sandbox
cloudflare-sandbox
atproto
sprites
daytona
Centralize WebSocket upgrade handling
Create a single HTTP server upgrade handler that dispatches to per-route
WebSocketServers using returned { wss, pathRegex } from attachWebSocket.
Change attachWebSocket APIs to no longer bind to the server and instead
return their wss and pathRegex. Reject unmatched upgrades cleanly.
Also add JWT pre-validation and improved session error handling in the
PTY WS flow, update app-proxy to forward /pty and /tty to the API, and
add a CLI log of the WebSocket URL.
+91
-90
+28
-5
apps/api/src/index.ts
+28
-5
apps/api/src/index.ts
···
11
11
import tty, { attachWebSocket as attachTtyWebSocket } from "./tty";
12
12
import pty, { attachWebSocket as attachPtyWebSocket } from "./pty";
13
13
import { createRateLimiter } from "./ratelimiter";
14
+
import { createServer as createHttpServer } from "node:http";
15
+
import type { IncomingMessage } from "node:http";
16
+
import type { Duplex } from "node:stream";
14
17
15
18
let xrpcServer = createServer({
16
19
validateResponse: false,
···
57
60
app.use("/tty", tty);
58
61
app.use("/pty", pty);
59
62
60
-
const server = app.listen(process.env.POCKETENV_XPRC_PORT || 8789, () => {
63
+
const httpServer = createHttpServer(app);
64
+
65
+
const wsHandlers = [
66
+
attachPtyWebSocket("/pty"),
67
+
attachTtyWebSocket("/tty"),
68
+
attachSshWebSocket("/ssh"),
69
+
];
70
+
71
+
httpServer.on("upgrade", (req: IncomingMessage, socket: Duplex, head: Buffer) => {
72
+
const pathname = new URL(req.url ?? "", "http://localhost").pathname;
73
+
for (const { wss, pathRegex } of wsHandlers) {
74
+
const match = pathname.match(pathRegex);
75
+
if (match) {
76
+
wss.handleUpgrade(req, socket, head, (ws) => {
77
+
wss.emit("connection", ws, req, match[1]!);
78
+
});
79
+
return;
80
+
}
81
+
}
82
+
// No handler matched — reject cleanly
83
+
socket.write("HTTP/1.1 404 Not Found\r\nConnection: close\r\n\r\n");
84
+
socket.destroy();
85
+
});
86
+
87
+
httpServer.listen(process.env.POCKETENV_XPRC_PORT || 8789, () => {
61
88
consola.log(chalk.greenBright(banner));
62
89
consola.info(
63
90
`Pocketenv XRPC API is running on port ${process.env.POCKETENV_XPRC_PORT || 8789}`,
64
91
);
65
92
});
66
-
67
-
attachPtyWebSocket(server, "/pty");
68
-
attachTtyWebSocket(server, "/tty");
69
-
attachSshWebSocket(server, "/ssh");
+50
-57
apps/api/src/pty/index.ts
+50
-57
apps/api/src/pty/index.ts
···
11
11
import * as e2b from "./e2b";
12
12
import { WebSocketServer, type WebSocket } from "ws";
13
13
import type { IncomingMessage } from "http";
14
-
import type { Server } from "http";
15
-
import type { Duplex } from "node:stream";
16
14
17
15
const router = Router();
18
16
router.use((req, res, next) => {
···
114
112
115
113
export default router;
116
114
117
-
export function attachWebSocket(server: Server, base: string) {
115
+
export function attachWebSocket(base: string) {
118
116
const pathRegex = new RegExp(`^${base}/([^/]+)/ws$`);
119
117
const wss = new WebSocketServer({ noServer: true });
120
118
121
-
server.on("upgrade", (req: IncomingMessage, socket: Duplex, head: Buffer) => {
122
-
const url = new URL(req.url ?? "", "http://localhost");
123
-
const match = url.pathname.match(pathRegex);
124
-
if (!match) return;
119
+
wss.on(
120
+
"connection",
121
+
async (ws: WebSocket, req: IncomingMessage, id: string) => {
122
+
const url = new URL(req.url ?? "", "http://localhost");
125
123
126
-
wss.handleUpgrade(req, socket, head, (ws) => {
127
-
wss.emit("connection", ws, req, match[1]!);
128
-
});
129
-
});
124
+
// Auth: query param ?token=<jwt> or Authorization: Bearer <jwt> header
125
+
const tokenParam = url.searchParams.get("token");
126
+
const authHeader = req.headers.authorization;
127
+
const bearer = tokenParam ?? authHeader?.split("Bearer ")[1]?.trim();
128
+
if (bearer && bearer !== "null") {
129
+
try {
130
+
jwt.verify(bearer, env.JWT_SECRET, { ignoreExpiration: true });
131
+
} catch (err) {
132
+
consola.error("WS: Invalid JWT token:", err);
133
+
ws.close(1008, "Invalid token");
134
+
return;
135
+
}
136
+
}
130
137
131
-
wss.on("connection", async (ws: WebSocket, req: IncomingMessage, id: string) => {
132
-
const url = new URL(req.url ?? "", "http://localhost");
133
-
134
-
// Auth: query param ?token=<jwt> or Authorization: Bearer <jwt> header
135
-
const tokenParam = url.searchParams.get("token");
136
-
const authHeader = req.headers.authorization;
137
-
const bearer = tokenParam ?? authHeader?.split("Bearer ")[1]?.trim();
138
-
if (bearer && bearer !== "null") {
138
+
let session: Awaited<ReturnType<typeof getSession>>;
139
139
try {
140
-
jwt.verify(bearer, env.JWT_SECRET, { ignoreExpiration: true });
140
+
session = await getSession(context.ctx, id);
141
141
} catch (err) {
142
-
consola.error("WS: Invalid JWT token:", err);
143
-
ws.close(1008, "Invalid token");
142
+
consola.error("WS: Failed to get session:", err);
143
+
ws.close(1011, "Session error");
144
144
return;
145
145
}
146
-
}
147
146
148
-
let session: Awaited<ReturnType<typeof getSession>>;
149
-
try {
150
-
session = await getSession(context.ctx, id);
151
-
} catch (err) {
152
-
consola.error("WS: Failed to get session:", err);
153
-
ws.close(1011, "Session error");
154
-
return;
155
-
}
147
+
session.wsClients.add(ws);
156
148
157
-
session.wsClients.add(ws);
158
-
159
-
ws.on("message", (data) => {
160
-
const text = data.toString("utf-8");
161
-
try {
162
-
const msg = JSON.parse(text);
163
-
if (
164
-
msg?.type === "resize" &&
165
-
Number.isInteger(msg.cols) &&
166
-
Number.isInteger(msg.rows)
167
-
) {
168
-
session.socket.sendMessage({
169
-
type: "resize",
170
-
cols: msg.cols,
171
-
rows: msg.rows,
172
-
});
173
-
return;
149
+
ws.on("message", (data) => {
150
+
const text = data.toString("utf-8");
151
+
try {
152
+
const msg = JSON.parse(text);
153
+
if (
154
+
msg?.type === "resize" &&
155
+
Number.isInteger(msg.cols) &&
156
+
Number.isInteger(msg.rows)
157
+
) {
158
+
session.socket.sendMessage({
159
+
type: "resize",
160
+
cols: msg.cols,
161
+
rows: msg.rows,
162
+
});
163
+
return;
164
+
}
165
+
} catch {
166
+
// not JSON — treat as raw input
174
167
}
175
-
} catch {
176
-
// not JSON — treat as raw input
177
-
}
178
-
session.socket.sendMessage({ type: "message", message: text });
179
-
});
168
+
session.socket.sendMessage({ type: "message", message: text });
169
+
});
180
170
181
-
ws.on("close", () => {
182
-
session.wsClients.delete(ws);
183
-
});
184
-
});
171
+
ws.on("close", () => {
172
+
session.wsClients.delete(ws);
173
+
});
174
+
},
175
+
);
176
+
177
+
return { wss, pathRegex };
185
178
}
+3
-13
apps/api/src/ssh/index.ts
+3
-13
apps/api/src/ssh/index.ts
···
7
7
import generateJwt from "lib/generateJwt";
8
8
import { WebSocketServer, type WebSocket } from "ws";
9
9
import type { IncomingMessage } from "http";
10
-
import type { Server } from "http";
11
-
import type { Duplex } from "node:stream";
12
10
13
11
interface SSHSession {
14
12
client: Client;
···
272
270
273
271
export default router;
274
272
275
-
export function attachWebSocket(server: Server, base: string) {
273
+
export function attachWebSocket(base: string) {
276
274
const pathRegex = new RegExp(`^${base}/([^/]+)/ws$`);
277
275
const wss = new WebSocketServer({ noServer: true });
278
-
279
-
server.on("upgrade", (req: IncomingMessage, socket: Duplex, head: Buffer) => {
280
-
const url = new URL(req.url ?? "", "http://localhost");
281
-
const match = url.pathname.match(pathRegex);
282
-
if (!match) return;
283
-
284
-
wss.handleUpgrade(req, socket, head, (ws) => {
285
-
wss.emit("connection", ws, req, match[1]!);
286
-
});
287
-
});
288
276
289
277
wss.on("connection", async (ws: WebSocket, req: IncomingMessage, sessionId: string) => {
290
278
const url = new URL(req.url ?? "", "http://localhost");
···
333
321
session.wsClients.delete(ws);
334
322
});
335
323
});
324
+
325
+
return { wss, pathRegex };
336
326
}
+3
-13
apps/api/src/tty/index.tsx
+3
-13
apps/api/src/tty/index.tsx
···
11
11
import path from "node:path";
12
12
import { WebSocketServer, type WebSocket } from "ws";
13
13
import type { IncomingMessage } from "http";
14
-
import type { Server } from "http";
15
-
import type { Duplex } from "node:stream";
16
14
17
15
const router = Router();
18
16
router.use((req, res, next) => {
···
381
379
382
380
export default router;
383
381
384
-
export function attachWebSocket(server: Server, base: string) {
382
+
export function attachWebSocket(base: string) {
385
383
const pathRegex = new RegExp(`^${base}/([^/]+)/ws$`);
386
384
const wss = new WebSocketServer({ noServer: true });
387
-
388
-
server.on("upgrade", (req: IncomingMessage, socket: Duplex, head: Buffer) => {
389
-
const url = new URL(req.url ?? "", "http://localhost");
390
-
const match = url.pathname.match(pathRegex);
391
-
if (!match) return;
392
-
393
-
wss.handleUpgrade(req, socket, head, (ws) => {
394
-
wss.emit("connection", ws, req, match[1]!);
395
-
});
396
-
});
397
385
398
386
wss.on("connection", async (ws: WebSocket, req: IncomingMessage, id: string) => {
399
387
const url = new URL(req.url ?? "", "http://localhost");
···
439
427
session.wsClients.delete(ws);
440
428
});
441
429
});
430
+
431
+
return { wss, pathRegex };
442
432
}
+3
-1
apps/app-proxy/src/index.ts
+3
-1
apps/app-proxy/src/index.ts
···
78
78
API_ROUTES.includes(url.pathname) ||
79
79
url.pathname.startsWith('/oauth/callback') ||
80
80
url.pathname.startsWith('/xrpc') ||
81
-
url.pathname.startsWith('/ssh')
81
+
url.pathname.startsWith('/ssh') ||
82
+
url.pathname.startsWith('/pty') ||
83
+
url.pathname.startsWith('/tty')
82
84
) {
83
85
redirectToApi = true;
84
86
}
+4
-1
apps/cli/src/cmd/ssh/tty.ts
+4
-1
apps/cli/src/cmd/ssh/tty.ts
···
28
28
29
29
const baseUrl = tty ? env.POCKETENV_TTY_URL : env.POCKETENV_PTY_URL;
30
30
const wsUrl = toWsUrl(baseUrl, `/${sandbox.id}/ws`, authToken);
31
+
consola.info(wsUrl);
31
32
32
33
let cols = process.stdout.columns ?? 220;
33
34
let rows = process.stdout.rows ?? 50;
···
115
116
if (!exiting) {
116
117
const msg = reason.length ? ` (${code} – ${reason})` : "";
117
118
if (msg) {
118
-
process.stderr.write(`\r\n${chalk.yellow("Connection closed")}${msg}\r\n`);
119
+
process.stderr.write(
120
+
`\r\n${chalk.yellow("Connection closed")}${msg}\r\n`,
121
+
);
119
122
}
120
123
teardown(0);
121
124
}