Advance TS port Effect workbench

ts/packages/flow/src/gateway/dispatch/mux.ts

@@ -1,89 +0,0 @@
-/**
- * WebSocket multiplexer — handles concurrent requests over a single connection.
- *
- * Python reference: trustgraph-flow/trustgraph/gateway/dispatch/mux.py
- */
-
-import { AsyncQueue } from "@trustgraph/base";
-
-const MAX_OUTSTANDING = 15;
-const MAX_QUEUE_SIZE = 10;
-
-export interface MuxRequest {
-  id: string;
-  service: string;
-  flow?: string;
-  request: Record<string, unknown>;
-}
-
-export type MuxHandler = (
-  request: MuxRequest,
-  respond: (response: unknown, complete: boolean) => Promise<void>,
-) => Promise<void>;
-
-export class Mux {
-  private queue = new AsyncQueue<MuxRequest>();
-  private outstanding = 0;
-  private running = true;
-  private readonly handler: MuxHandler;
-
-  constructor(handler: MuxHandler) {
-    this.handler = handler;
-  }
-
-  receive(request: MuxRequest): void {
-    if (this.queue.length >= MAX_QUEUE_SIZE) {
-      console.warn("[Mux] Queue full, dropping request:", request.id);
-      return;
-    }
-    this.queue.push(request);
-  }
-
-  async run(send: (data: string) => void): Promise<void> {
-    while (this.running) {
-      if (this.outstanding >= MAX_OUTSTANDING) {
-        await sleep(50);
-        continue;
-      }
-
-      try {
-        const request = await this.queue.pop(1000);
-        this.outstanding++;
-
-        // Fire and forget — error handling inside
-        this.processRequest(request, send).finally(() => {
-          this.outstanding--;
-        });
-      } catch {
-        // Timeout on queue pop — just loop
-      }
-    }
-  }
-
-  stop(): void {
-    this.running = false;
-  }
-
-  private async processRequest(
-    request: MuxRequest,
-    send: (data: string) => void,
-  ): Promise<void> {
-    try {
-      await this.handler(request, async (response, complete) => {
-        send(JSON.stringify({ id: request.id, response, complete }));
-      });
-    } catch (err) {
-      send(
-        JSON.stringify({
-          id: request.id,
-          error: { type: "internal", message: String(err) },
-          complete: true,
-        }),
-      );
-    }
-  }
-}
-
-function sleep(ms: number): Promise<void> {
-  return new Promise((resolve) => setTimeout(resolve, ms));
-}

ts/packages/flow/src/gateway/index.ts

@@ -1,6 +1,5 @@
 export { createGateway, run, type GatewayConfig } from "./server.js";
 export { DispatcherManager } from "./dispatch/manager.js";
-export { Mux, type MuxRequest, type MuxHandler } from "./dispatch/mux.js";
 export {
   clientTermToInternal,
   clientTripleToInternal,

ts/packages/flow/src/gateway/rpc-contract.ts

@@ -0,0 +1,35 @@
+import { Schema as S } from "effect";
+import * as Rpc from "effect/unstable/rpc/Rpc";
+import * as RpcGroup from "effect/unstable/rpc/RpcGroup";
+
+export class DispatchPayload extends S.Class<DispatchPayload>("DispatchPayload")({
+  scope: S.Literals(["global", "flow"]),
+  service: S.String,
+  flow: S.optionalKey(S.String),
+  request: S.Record(S.String, S.Unknown),
+}) {}
+
+export class DispatchStreamChunk extends S.Class<DispatchStreamChunk>("DispatchStreamChunk")({
+  response: S.Unknown,
+  complete: S.Boolean,
+}) {}
+
+export class DispatchError extends S.ErrorClass<DispatchError>("DispatchError")({
+  _tag: S.tag("DispatchError"),
+  message: S.String,
+}) {}
+
+export class Dispatch extends Rpc.make("Dispatch", {
+  payload: DispatchPayload,
+  success: S.Unknown,
+  error: DispatchError,
+}) {}
+
+export class DispatchStream extends Rpc.make("DispatchStream", {
+  payload: DispatchPayload,
+  success: DispatchStreamChunk,
+  error: DispatchError,
+  stream: true,
+}) {}
+
+export const TrustGraphRpcs = RpcGroup.make(Dispatch, DispatchStream);

ts/packages/flow/src/gateway/rpc-protocol.ts

@@ -0,0 +1,92 @@
+import { Effect, Queue, Scope } from "effect";
+import * as RpcMessage from "effect/unstable/rpc/RpcMessage";
+import * as RpcSerialization from "effect/unstable/rpc/RpcSerialization";
+import * as RpcServer from "effect/unstable/rpc/RpcServer";
+import * as Socket from "effect/unstable/socket/Socket";
+
+export const makeSocketRpcProtocol = Effect.gen(function* () {
+  const serialization = yield* RpcSerialization.RpcSerialization;
+  const disconnects = yield* Queue.make<number>();
+
+  let nextClientId = 0;
+  const clients = new Map<number, {
+    readonly write: (response: RpcMessage.FromServerEncoded) => Effect.Effect<void>;
+  }>();
+  const clientIds = new Set<number>();
+
+  let writeRequest!: (
+    clientId: number,
+    message: RpcMessage.FromClientEncoded,
+  ) => Effect.Effect<void>;
+
+  const onSocket = function* (
+    socket: Socket.Socket,
+    headers?: ReadonlyArray<[string, string]>,
+  ) {
+    const scope = yield* Effect.scope;
+    const parser = serialization.makeUnsafe();
+    const clientId = nextClientId++;
+
+    yield* Scope.addFinalizerExit(scope, () => {
+      clients.delete(clientId);
+      clientIds.delete(clientId);
+      return Queue.offer(disconnects, clientId);
+    });
+
+    const writeRaw = yield* socket.writer;
+    const write = (response: RpcMessage.FromServerEncoded) => {
+      try {
+        const encoded = parser.encode(response);
+        if (encoded === undefined) return Effect.void;
+        return Effect.orDie(writeRaw(encoded));
+      } catch (cause) {
+        return Effect.orDie(
+          writeRaw(parser.encode(RpcMessage.ResponseDefectEncoded(cause))!),
+        );
+      }
+    };
+
+    clients.set(clientId, { write });
+    clientIds.add(clientId);
+
+    yield* socket.runRaw((data) => {
+      try {
+        const decoded = parser.decode(data) as ReadonlyArray<RpcMessage.FromClientEncoded>;
+        return Effect.forEach(decoded, (message) => {
+          if (message._tag === "Request" && headers !== undefined) {
+            return writeRequest(clientId, {
+              ...message,
+              headers: headers.concat(message.headers),
+            });
+          }
+          return writeRequest(clientId, message);
+        }, { discard: true });
+      } catch (cause) {
+        return writeRaw(parser.encode(RpcMessage.ResponseDefectEncoded(cause))!);
+      }
+    }).pipe(
+      Effect.catchReason("SocketError", "SocketCloseError", () => Effect.void),
+      Effect.orDie,
+    );
+  };
+
+  const protocol = yield* RpcServer.Protocol.make((writeRequest_) => {
+    writeRequest = writeRequest_;
+    return Effect.succeed({
+      disconnects,
+      send: (clientId, response) => {
+        const client = clients.get(clientId);
+        if (client === undefined) return Effect.void;
+        return Effect.orDie(client.write(response));
+      },
+      end: () => Effect.void,
+      clientIds: Effect.sync(() => clientIds),
+      initialMessage: Effect.succeedNone,
+      supportsAck: true,
+      supportsTransferables: false,
+      supportsSpanPropagation: true,
+    });
+  });
+
+  return { onSocket, protocol } as const;
+});

ts/packages/flow/src/gateway/rpc-server.ts

@@ -0,0 +1,109 @@
+import { Cause, Effect, Layer, Queue, Scope } from "effect";
+import * as RpcSerialization from "effect/unstable/rpc/RpcSerialization";
+import * as RpcServer from "effect/unstable/rpc/RpcServer";
+import type * as Socket from "effect/unstable/socket/Socket";
+import { errorMessage } from "@trustgraph/base";
+import type { DispatcherManager } from "./dispatch/manager.js";
+import { DispatchError, DispatchPayload, DispatchStreamChunk, TrustGraphRpcs } from "./rpc-contract.js";
+import { makeSocketRpcProtocol } from "./rpc-protocol.js";
+
+export interface GatewayRpcServer {
+  readonly onSocket: (
+    socket: Socket.Socket,
+    headers?: ReadonlyArray<[string, string]>,
+  ) => Effect.Effect<void, never, Scope.Scope>;
+}
+
+export const makeGatewayRpcServer = Effect.fn("makeGatewayRpcServer")(function* (
+  dispatcher: DispatcherManager,
+) {
+  const { onSocket, protocol } = yield* makeSocketRpcProtocol;
+
+  const serverLayer = RpcServer.layer(TrustGraphRpcs, {
+    disableFatalDefects: true,
+  }).pipe(
+    Layer.provide(Layer.succeed(RpcServer.Protocol, protocol)),
+    Layer.provide(makeGatewayRpcHandlers(dispatcher)),
+    Layer.provide(RpcSerialization.layerNdjson),
+  );
+
+  yield* Layer.launch(serverLayer).pipe(Effect.forkScoped);
+
+  return {
+    onSocket: Effect.fn("GatewayRpc.onSocket")(function* (socket, headers) {
+      yield* onSocket(socket, headers);
+    }),
+  } satisfies GatewayRpcServer;
+});
+
+const makeGatewayRpcHandlers = (dispatcher: DispatcherManager) =>
+  TrustGraphRpcs.toLayer(Effect.succeed(
+    TrustGraphRpcs.of({
+      Dispatch: (payload) =>
+        Effect.tryPromise({
+          try: () => dispatchOne(dispatcher, payload),
+          catch: (cause) => new DispatchError({ message: errorMessage(cause) }),
+        }),
+      DispatchStream: Effect.fn("GatewayRpc.DispatchStream")(function* (payload) {
+        const context = yield* Effect.context<never>();
+        const runPromise = Effect.runPromiseWith(context);
+        const queue = yield* Queue.bounded<DispatchStreamChunk, DispatchError | Cause.Done>(16);
+        yield* Effect.addFinalizer(() => Queue.shutdown(queue));
+
+        yield* Effect.tryPromise({
+          try: () =>
+            dispatchStream(dispatcher, payload, async (response, complete) => {
+              await runPromise(Queue.offer(queue, new DispatchStreamChunk({ response, complete })));
+              return complete;
+            }),
+          catch: (cause) => new DispatchError({ message: errorMessage(cause) }),
+        }).pipe(
+          Effect.flatMap(() => Queue.end(queue)),
+          Effect.catch((error) => Queue.fail(queue, error)),
+          Effect.forkScoped,
+        );
+
+        return queue;
+      }),
+    }),
+  ));
+
+function dispatchOne(
+  dispatcher: DispatcherManager,
+  payload: DispatchPayload,
+): Promise<unknown> {
+  if (payload.scope === "flow") {
+    return dispatcher.dispatchFlowService(
+      payload.flow ?? "default",
+      payload.service,
+      payload.request,
+    );
+  }
+  return dispatcher.dispatchGlobalService(payload.service, payload.request);
+}
+
+async function dispatchStream(
+  dispatcher: DispatcherManager,
+  payload: DispatchPayload,
+  responder: (response: unknown, complete: boolean) => Promise<boolean>,
+): Promise<void> {
+  const send = async (response: unknown, complete: boolean) => {
+    await responder(response, complete);
+  };
+
+  if (payload.scope === "flow") {
+    await dispatcher.dispatchFlowServiceStreaming(
+      payload.flow ?? "default",
+      payload.service,
+      payload.request,
+      send,
+    );
+    return;
+  }
+
+  await dispatcher.dispatchGlobalServiceStreaming(
+    payload.service,
+    payload.request,
+    send,
+  );
+}

ts/packages/flow/src/gateway/server.ts

@@ -2,19 +2,20 @@
  * API Gateway — HTTP + WebSocket server.
  *
  * Replaces the Python aiohttp gateway with Fastify.
- * Uses the Mux class for WebSocket multiplexing (queue-based request
- * buffering, concurrency control, proper task lifecycle).
+ * Uses Effect RPC over WebSocket for streaming client requests.
  *
  * Python reference: trustgraph-flow/trustgraph/gateway/service.py
  */
 
 import Fastify from "fastify";
 import websocketPlugin from "@fastify/websocket";
-import { Config, Effect } from "effect";
+import { Config, Effect, Exit, Scope } from "effect";
 import * as O from "effect/Option";
-import { errorMessage, optionalStringConfig, registry, toTgError } from "@trustgraph/base";
+import * as RpcSerialization from "effect/unstable/rpc/RpcSerialization";
+import * as EffectSocket from "effect/unstable/socket/Socket";
+import { optionalStringConfig, registry, toTgError } from "@trustgraph/base";
 import { DispatcherManager } from "./dispatch/manager.js";
-import { Mux, type MuxRequest, type MuxHandler } from "./dispatch/mux.js";
+import { makeGatewayRpcServer } from "./rpc-server.js";
 
 export interface GatewayConfig {
   port: number;
@@ -29,11 +30,18 @@ export async function createGateway(config: GatewayConfig) {
 
   const dispatcher = new DispatcherManager(config);
   await dispatcher.start();
+  const rpcScope = await Effect.runPromise(Scope.make());
+  const rpcServer = await Effect.runPromise(
+    makeGatewayRpcServer(dispatcher).pipe(
+      Effect.provideService(RpcSerialization.RpcSerialization, RpcSerialization.ndjson),
+      Scope.provide(rpcScope),
+    ),
+  );
 
   // Authentication middleware
   app.addHook("onRequest", async (request, reply) => {
     if (request.url === "/api/v1/metrics") return;
-    if (request.url.startsWith("/api/v1/socket")) return; // Socket auth via query param
+    if (request.url.startsWith("/api/v1/rpc")) return; // RPC socket auth via query param
 
     if (config.secret !== undefined && config.secret.length > 0) {
       const auth = request.headers.authorization;
@@ -43,6 +51,38 @@ export async function createGateway(config: GatewayConfig) {
     }
   });
 
+  app.post<{
+    Body: {
+      scope?: string;
+      service?: string;
+      flow?: string;
+      request?: Record<string, unknown>;
+    };
+  }>("/api/v1/workbench/dispatch", async (request, reply) => {
+    const body = request.body;
+    const service = body.service;
+    const payload = body.request;
+    if (service === undefined || service.length === 0 || payload === undefined) {
+      return reply.code(400).send({
+        error: { type: "bad-request", message: "service and request are required" },
+      });
+    }
+
+    try {
+      const result = body.scope === "flow"
+        ? await dispatcher.dispatchFlowService(body.flow ?? "default", service, payload)
+        : await dispatcher.dispatchGlobalService(service, payload);
+      const err = (result as Record<string, unknown>)?.error as { type?: string; message?: string } | undefined;
+      if (err !== undefined) {
+        const statusCode = err.type === "not-found" ? 404 : 400;
+        return reply.code(statusCode).send(result);
+      }
+      return result;
+    } catch (err) {
+      reply.code(500).send({ error: toTgError(err) });
+    }
+  });
+
   // REST endpoint: POST /api/v1/:kind  (global services)
   app.post<{ Params: { kind: string } }>("/api/v1/:kind", async (request, reply) => {
     const { kind } = request.params;
@@ -124,10 +164,8 @@ export async function createGateway(config: GatewayConfig) {
     },
   );
 
-  // WebSocket endpoint: /api/v1/socket
-  // Uses Mux for queue-based request buffering and concurrency control.
-  app.get("/api/v1/socket", { websocket: true }, (socket, request) => {
-    // Auth via query param
+  // Effect RPC WebSocket endpoint: /api/v1/rpc
+  app.get("/api/v1/rpc", { websocket: true }, (socket, request) => {
     const url = new URL(request.url, `http://${request.headers.host}`);
     const token = url.searchParams.get("token");
     if (config.secret !== undefined && config.secret.length > 0 && token !== config.secret) {
@@ -135,91 +173,18 @@ export async function createGateway(config: GatewayConfig) {
       return;
     }
 
-    // Build the MuxHandler that dispatches to the DispatcherManager
-    const handler: MuxHandler = async (muxReq, respond) => {
-      if (muxReq.flow !== undefined && muxReq.flow.length > 0) {
-        await dispatcher.dispatchFlowServiceStreaming(
-          muxReq.flow,
-          muxReq.service,
-          muxReq.request,
-          respond,
+    const program = Effect.scoped(
+      Effect.gen(function* () {
+        const effectSocket = yield* EffectSocket.fromWebSocket(
+          Effect.succeed(socket as unknown as globalThis.WebSocket),
+          { closeCodeIsError: (code) => code !== 1000 },
         );
-      } else {
-        await dispatcher.dispatchGlobalServiceStreaming(
-          muxReq.service,
-          muxReq.request,
-          respond,
-        );
-      }
-    };
+        yield* rpcServer.onSocket(effectSocket, headersFrom(request.headers));
+      }),
+    );
 
-    const mux = new Mux(handler);
-
-    // Start the Mux run loop — sends responses back over the socket
-    const runPromise = mux.run((data) => {
-      // Only send if the socket is still open (readyState 1 = OPEN)
-      if (socket.readyState === 1) {
-        socket.send(data);
-      }
-    });
-
-    // Incoming messages get queued into the Mux
-    socket.on("message", (data) => {
-      try {
-        const msg = JSON.parse(data.toString()) as {
-          id?: string;
-          service?: string;
-          flow?: string;
-          request?: Record<string, unknown>;
-        };
-
-        if (
-          msg.id === undefined ||
-          msg.id.length === 0 ||
-          msg.service === undefined ||
-          msg.service.length === 0 ||
-          msg.request === undefined
-        ) {
-          socket.send(
-            JSON.stringify({
-              id: msg.id ?? null,
-              error: { type: "bad-request", message: "Missing id, service, or request" },
-              complete: true,
-            }),
-          );
-          return;
-        }
-
-        const muxReq: MuxRequest = {
-          id: msg.id,
-          service: msg.service,
-          request: msg.request,
-          ...(msg.flow !== undefined ? { flow: msg.flow } : {}),
-        };
-
-        mux.receive(muxReq);
-      } catch (err) {
-        socket.send(
-          JSON.stringify({
-            error: { type: "parse-error", message: errorMessage(err) },
-            complete: true,
-          }),
-        );
-      }
-    });
-
-    socket.on("close", () => {
-      mux.stop();
-    });
-
-    socket.on("error", () => {
-      mux.stop();
-    });
-
-    // Ensure runPromise errors don't go unhandled
-    runPromise.catch((err) => {
-      console.error("[Gateway] Mux run loop error:", err);
-      mux.stop();
+    Effect.runPromise(program.pipe(Scope.provide(rpcScope))).catch((err) => {
+      console.error("[Gateway] RPC WebSocket error:", err);
       if (socket.readyState === 1) {
         socket.close(1011, "Internal server error");
       }
@@ -236,11 +201,21 @@ export async function createGateway(config: GatewayConfig) {
     start: () => app.listen({ port: config.port, host: "0.0.0.0" }),
     stop: async () => {
       await app.close();
+      await Effect.runPromise(Scope.close(rpcScope, Exit.void));
       await dispatcher.stop();
     },
   };
 }
 
+function headersFrom(headers: Record<string, string | string[] | number | undefined>): ReadonlyArray<[string, string]> {
+  return Object.entries(headers).flatMap(([key, value]) => {
+    if (typeof value === "string") return [[key, value] satisfies [string, string]];
+    if (typeof value === "number") return [[key, String(value)] satisfies [string, string]];
+    if (Array.isArray(value)) return value.map((item) => [key, item] satisfies [string, string]);
+    return [];
+  });
+}
+
 export async function run(): Promise<void> {
   await Effect.runPromise(program);
 }