spb/claude-remote

1/**2 * Browser WebSockets that are piped to the runtime: the terminal stream and the task log stream.3 * Authentication (cookie + completed second factor) and the Origin check happen in the global4 * guard during the HTTP upgrade; ownership and liveness are checked here in `preValidation`,5 * so a refused connection is a plain HTTP error and never becomes a socket.6 */7import type { FastifyInstance, FastifyRequest } from 'fastify';8import type WebSocket from 'ws';9import { z } from 'zod';10import { Id } from '@cr/shared';11import { HttpError, authOf, notFound, parse, type Ctx } from '../http.js';12import { findSessionById } from '../security/sessions.js';13import { pipeWebSockets } from '../ws-pipe.js';14import { isLive } from './sessions/service.js';15import type { LiveKind } from '../live.js';1617const IdParams = z.object({ id: Id });18const StreamQuery = z.object({19  cols: z.coerce.number().int().min(20).max(500).default(120),20  rows: z.coerce.number().int().min(5).max(200).default(32),21});22const LogQuery = z.object({ offset: z.coerce.number().int().min(0).default(0) });2324/** Registers the socket and re-validates its auth session periodically (revoked/expired ⇒ 4401). */25export function watchAuth(ctx: Ctx, socket: WebSocket, req: FastifyRequest, kind: LiveKind, beforeClose?: () => void): () => void {26  const a = authOf(req);27  const unregister = ctx.live.add(socket, a.sessionId, kind, beforeClose);28  const timer = setInterval(() => {29    void findSessionById(ctx.db, a.sessionId).then(30      (row) => {31        if (!row || !row.second_factor) ctx.live.kickAuthSessions([a.sessionId]);32      },33      () => undefined, // database hiccup: keep the socket, retry on the next tick34    );35  }, ctx.config.wsAuthRecheckMs);36  timer.unref();37  const stop = (): void => {38    clearInterval(timer);39    unregister();40  };41  socket.once('close', stop);42  return stop;43}4445export default async function terminalWsPlugin(app: FastifyInstance, ctx: Ctx): Promise<void> {46  app.get('/sessions/:id/stream', {47    websocket: true,48    config: { doc: { summary: 'Terminal stream (binary PTY bytes + TermClientMsg/TermServerMsg text frames)', query: StreamQuery, ws: true } },49    preValidation: async (req) => {50      const { id } = parse(IdParams, req.params);51      parse(StreamQuery, req.query);52      const row = await ctx.sessions.owned(id, authOf(req).userId);53      if (!row) throw notFound('session');54      if (!isLive(row.state) || row.state === 'STOPPING') throw new HttpError(410, 'session_ended', 'session is not running');55    },56  }, (socket, req) => {57    const { id } = parse(IdParams, req.params);58    const q = parse(StreamQuery, req.query);59    const upstream = ctx.rt.ws(`/sessions/${id}/attach`, { cols: q.cols, rows: q.rows });60    watchAuth(ctx, socket, req, 'terminal');61    // This server knows first-hand that a browser is attached: the session state follows the62    // proxied stream at once instead of waiting for the runtime to report its tmux client.63    let counted = false;64    const failed = (what: string) => (err: unknown) => ctx.log.error({ err, sessionId: id }, what);65    upstream.once('open', () => {66      if (socket.readyState !== socket.OPEN) return;67      counted = true;68      void ctx.sessions.streamOpened(id).catch(failed('stream attach bookkeeping failed'));69    });70    pipeWebSockets(socket, upstream, {71      bufferBytes: ctx.config.termBufferBytes,72      clientToUpstream: true,73      // Profiles without a prompt-submit hook (Codex): typing ends WAITING_FOR_INPUT. Not wired74      // at all for Claude, whose UserPromptSubmit hook is the signal.75      ...(ctx.sessions.inputEndsWaiting76        ? { onClientInput: () => void ctx.sessions.inputReceived(id).catch(failed('input bookkeeping failed')) }77        : {}),78      onClose: () => {79        if (!counted) return;80        counted = false;81        void ctx.sessions.streamClosed(id).catch(failed('stream detach bookkeeping failed'));82      },83    });84    void ctx.db.q("INSERT INTO terminal_events (session_id, kind, detail) VALUES ($1, 'attach', '{}')", [id]).catch(() => undefined);85    socket.once('close', () => {86      void ctx.db.q("INSERT INTO terminal_events (session_id, kind, detail) VALUES ($1, 'detach', '{}')", [id]).catch(() => undefined);87    });88  });8990  app.get('/tasks/:id/log/stream', {91    websocket: true,92    config: { doc: { summary: 'Task log follow stream (binary log bytes)', query: LogQuery, ws: true } },93    preValidation: async (req) => {94      const { id } = parse(IdParams, req.params);95      parse(LogQuery, req.query);96      if (!(await ctx.tasks.owned(id, authOf(req).userId))) throw notFound('task');97    },98  }, (socket, req) => {99    const { id } = parse(IdParams, req.params);100    const q = parse(LogQuery, req.query);101    const upstream = ctx.rt.ws(`/tasks/${id}/log/follow`, { offset: q.offset });102    watchAuth(ctx, socket, req, 'task-log');103    pipeWebSockets(socket, upstream, { bufferBytes: ctx.config.termBufferBytes, clientToUpstream: false });104  });105106  app.get('/events', {107    websocket: true,108    config: { doc: { summary: 'ServerEvent stream (text JSON)', ws: true } },109  }, (socket, req) => {110    const a = authOf(req);111    const remove = ctx.hub.add(socket, a.userId);112    watchAuth(ctx, socket, req, 'events', () => {113      if (socket.readyState === socket.OPEN) socket.send(JSON.stringify({ type: 'auth.expired' }));114    });115    socket.on('message', () => undefined); // server → browser only116    socket.on('error', () => undefined);117    socket.once('close', remove);118    socket.send(JSON.stringify({ type: 'hello', serverTime: new Date().toISOString() }));119  });120}121