@bigmistqke/rpc 0.1.2 → 0.1.4

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@bigmistqke/rpc",
-  "version": "0.1.2",
+  "version": "0.1.4",
   "type": "module",
   "description": "RPC toolkit for type-safe communication across Workers, iframes, and network boundaries.",
   "module": "./dist/index.js",
@@ -10,6 +10,10 @@
       "import": "./dist/messenger.js",
       "types": "./dist/messenger.d.ts"
     },
+    "./websocket": {
+      "import": "./dist/websocket.js",
+      "types": "./dist/websocket.d.ts"
+    },
     "./stream": {
       "import": "./dist/stream.js",
       "types": "./dist/stream.d.ts"
@@ -23,6 +27,13 @@
       "types": "./dist/fetch-node.d.ts"
     }
   },
+  "scripts": {
+    "build": "tsup",
+    "test": "vitest run",
+    "types": "pnpm exec tsc --noEmit",
+    "bump": "bumpp",
+    "prepublishOnly": "pnpm test && pnpm build"
+  },
   "license": "MIT",
   "devDependencies": {
     "tsup": "6.6.3",
@@ -33,10 +44,5 @@
     "@types/node": "^22.15.30",
     "bumpp": "^10.3.2",
     "valibot": "^1.0.0"
-  },
-  "scripts": {
-    "build": "tsup",
-    "test": "vitest run",
-    "bump": "bumpp"
   }
-}
+}

package/src/core.ts

@@ -0,0 +1,100 @@
+import { HANDLE_NAMESPACE_PREFIX, isHandleMarker, nextHandleNamespaceId } from './handle'
+import { $MESSENGER_HANDLE, HandleResponseShape } from './protocol'
+import type { RPC } from './types'
+
+export function createCommander<T extends object = object>(
+  apply: (topics: Array<string>, args: Array<any>) => void,
+): T {
+  function _createCommander(
+    topics: Array<string>,
+    apply: (topics: Array<string>, args: Array<any>) => void,
+  ): T {
+    return new Proxy(function () {} as T, {
+      get(target, topic) {
+        if (typeof topic === 'symbol') return (target as any)[topic]
+        // Return undefined for 'then' so proxy isn't treated as thenable
+        if (topic === 'then') return undefined
+        return _createCommander([...topics, topic], apply)
+      },
+      apply(_, __, args) {
+        return apply(topics, args)
+      },
+    })
+  }
+  return _createCommander([], apply)
+}
+
+export function callMethod(methods: object, topics: string[], args: unknown[]) {
+  const method = topics.reduce((acc, topic) => {
+    const result = (acc as any)?.[topic]
+    return result
+  }, methods)
+  if (typeof method !== 'function') {
+    throw new Error(`Topics did not resolve to a function: [${topics.join(',')}]`)
+  }
+  return method(...args)
+}
+
+/**
+ * Check if a response payload is a handle response
+ */
+export function isHandleResponse(value: unknown): value is { [$MESSENGER_HANDLE]: string } {
+  return (
+    !!value &&
+    typeof value === 'object' &&
+    $MESSENGER_HANDLE in value &&
+    typeof (value as any)[$MESSENGER_HANDLE] === 'string'
+  )
+}
+
+/**
+ * Creates a request handler for expose() that manages namespace routing and handle() sub-proxies.
+ * Returns the processed result ready to be sent back over the transport.
+ */
+export function createExposeRequestHandler(methods: object) {
+  const namespaceHandlers = new Map<string, object>()
+
+  const processResult = (result: unknown): unknown => {
+    if (isHandleMarker(result)) {
+      const namespaceId = nextHandleNamespaceId()
+      namespaceHandlers.set(namespaceId, result.methods)
+      return HandleResponseShape.create(namespaceId)
+    }
+    return result
+  }
+
+  return async (topics: string[], args: unknown[]): Promise<unknown> => {
+    const firstTopic = topics[0]
+    if (firstTopic && firstTopic.startsWith(HANDLE_NAMESPACE_PREFIX)) {
+      const handler = namespaceHandlers.get(firstTopic)
+      if (!handler) {
+        throw new Error(`Unknown namespace: ${firstTopic}`)
+      }
+      const result = await callMethod(handler, topics.slice(1), args)
+      return processResult(result)
+    }
+    const result = await callMethod(methods, topics, args)
+    return processResult(result)
+  }
+}
+
+/**
+ * Creates an RPC commander proxy that handles handle() sub-proxy responses.
+ */
+export function createRpcCommander<T extends object>(
+  request: (topics: string[], args: any[]) => Promise<unknown>,
+  topicPrefix: string[] = [],
+): RPC<T> {
+  return createCommander<RPC<T>>((topics, methodArgs) => {
+    const fullTopics = [...topicPrefix, ...topics]
+    return request(fullTopics, methodArgs).then((result: unknown) => {
+      if (isHandleResponse(result)) {
+        return createRpcCommander(
+          request,
+          result[$MESSENGER_HANDLE] ? [result[$MESSENGER_HANDLE]] : [],
+        )
+      }
+      return result
+    })
+  })
+}

package/src/fetch/index.ts

@@ -1,6 +1,7 @@
 import * as v from 'valibot'
 import { RPC } from '../types'
-import { callMethod, createCommander, createShape } from '../utils'
+import { callMethod, createCommander } from '../core'
+import { createShape } from '../utils'
 
 const $FETCH_HEADER = 'RPC_RR_PROXY'
 

package/src/fetch/node.ts

@@ -2,7 +2,7 @@
 
 import { IncomingMessage, ServerResponse } from 'http'
 import { Payload } from '.'
-import { callMethod } from '../utils'
+import { callMethod } from '../core'
 
 export function exposeNode<T extends object>(
   methods: T,

package/src/handle.ts

@@ -0,0 +1,48 @@
+const $HANDLE_MARKER = Symbol('RPC-HANDLE-MARKER')
+
+/** Internal marker type for handle() */
+interface HandleMarker<T extends object> {
+  [$HANDLE_MARKER]: true
+  methods: T
+}
+
+/** Type for values returned by handle() - unwrapped to RPC<T> by RPC system */
+export type Handled<T extends object> = T & { readonly ['__rpc_handled__']: T }
+
+/**
+ * Mark methods to be returned as a sub-proxy from an RPC method.
+ * Use this when a method needs to return an object with callable methods.
+ *
+ * @example
+ * ```ts
+ * expose({
+ *   init(canvas: OffscreenCanvas) {
+ *     const renderer = createRenderer(canvas)
+ *     return handle({
+ *       render: () => renderer.render(),
+ *       resize: (w, h) => renderer.resize(w, h),
+ *     })
+ *   }
+ * })
+ * ```
+ */
+export function handle<T extends object>(methods: T): Handled<T> {
+  return { [$HANDLE_MARKER]: true, methods } as unknown as Handled<T>
+}
+
+/**
+ * Check if a value is a handle marker
+ */
+export function isHandleMarker<T extends object>(value: unknown): value is HandleMarker<T> {
+  return !!value && typeof value === 'object' && $HANDLE_MARKER in value
+}
+
+// Prefix for namespace IDs to avoid collisions
+export const HANDLE_NAMESPACE_PREFIX = '__rpc_handle_'
+
+// Counter for generating unique namespace IDs
+export let handleNamespaceCounter = 0
+
+export function nextHandleNamespaceId(): string {
+  return `${HANDLE_NAMESPACE_PREFIX}${handleNamespaceCounter++}`
+}

package/src/{messenger.ts → messenger/index.ts}

@@ -1,3 +1,4 @@
+import { createExposeRequestHandler, createRpcCommander } from '../core'
 import {
   $MESSENGER_ERROR,
   $MESSENGER_RESPONSE,
@@ -6,11 +7,61 @@ import {
   RequestShape,
   ResponseShape,
   RPCPayloadShape,
-} from './message-protocol'
-import { RPC } from './types'
-import { callMethod, createCommander, createIdRegistry, defer } from './utils'
+} from '../protocol'
+import type { RPC } from '../types'
+import { createIdRegistry, defer } from '../utils'
+export { handle, type Handled } from '../handle'
 
-export const $TRANSFER = 'WORKER-TRANSFER'
+export const $TRANSFER = 'RPC-TRANSFER'
+export const $MESSENGER = Symbol.for('RPC-MESSENGER')
+
+/** Wrapper type for transferable values */
+export type Transferred<T> = T & {
+  [$TRANSFER]: true
+}
+
+/**
+ * Mark a value as transferable for postMessage.
+ * Use this to transfer ownership of ArrayBuffer, ReadableStream, etc.
+ */
+export function transfer<T extends object>(value: T): Transferred<T> {
+  return Object.assign(value, { [$TRANSFER]: true } as const)
+}
+
+/**
+ * Check if a value is marked for transfer
+ */
+function isTransferred(value: unknown): value is Transferred<unknown> {
+  return !!value && typeof value === 'object' && $TRANSFER in value
+}
+
+/**
+ * Extract transferables from args and unwrap transferred values
+ */
+function extractTransferables(args: any[]): { args: any[]; transferables: Transferable[] } {
+  const transferables: Transferable[] = []
+
+  const processValue = (value: any): any => {
+    if (isTransferred(value)) {
+      transferables.push(value)
+      return value
+    }
+    if (Array.isArray(value)) {
+      return value.map(processValue)
+    }
+    if (value && typeof value === 'object' && value.constructor === Object) {
+      const result: any = {}
+      for (const key in value) {
+        result[key] = processValue(value[key])
+      }
+      return result
+    }
+    return value
+  }
+
+  const processedArgs = args.map(processValue)
+  return { args: processedArgs, transferables }
+}
 
 interface WorkerMessenger {
   postMessage(message: any, transferables?: any[]): void
@@ -115,7 +166,12 @@ export function createResponder(
       if (RequestShape.validate(data)) {
         try {
           const result = await callback(data)
-          postMessage(ResponseShape.create(data, result))
+          // Extract transferables from the result
+          const {
+            args: [processedResult],
+            transferables,
+          } = extractTransferables([result])
+          postMessage(ResponseShape.create(data, processedResult), transferables)
         } catch (error) {
           postMessage(ErrorShape.create(data, error))
         }
@@ -136,29 +192,64 @@ export function createResponder(
 /**********************************************************************************/
 
 /**
- * Exposes a set of methods as an RPC endpoint over the given messenger.
+ * Exposes methods as an RPC endpoint over the given messenger.
  *
- * @param methods - An object containing functions to expose
+ * @param methods - Object containing methods to expose
  * @param options - Optional target Messenger and abort signal
+ *
+ * @example
+ * ```ts
+ * // Worker side - simple methods
+ * expose({
+ *   add: (a, b) => a + b,
+ *   multiply: (a, b) => a * b,
+ * })
+ *
+ * // Worker side - with initialization returning sub-proxy
+ * expose({
+ *   init(canvas: OffscreenCanvas) {
+ *     const renderer = createRenderer(canvas)
+ *     return handle({
+ *       render: () => renderer.render(),
+ *       resize: (w, h) => renderer.resize(w, h),
+ *     })
+ *   }
+ * })
+ * ```
  */
-export function expose<T extends object>(
-  methods: T,
+export function expose<TMethods extends object>(
+  methods: TMethods,
   { to = self, signal }: { to?: Messenger; signal?: AbortSignal } = {},
-) {
-  createResponder(
-    to,
-    data => {
-      if (RPCPayloadShape.validate(data.payload)) {
+): void {
+  const postMessage = usePostMessage(to)
+  const handleRequest = createExposeRequestHandler(methods)
+
+  to.addEventListener(
+    'message',
+    async event => {
+      const data = (event as MessageEvent).data
+      if (RequestShape.validate(data)) {
         try {
-          const { topics, args } = data.payload
-          return callMethod(methods, topics, args)
+          if (RPCPayloadShape.validate(data.payload)) {
+            const result = await handleRequest(data.payload.topics, data.payload.args)
+            const {
+              args: [finalResult],
+              transferables,
+            } = extractTransferables([result])
+            postMessage(ResponseShape.create(data, finalResult), transferables)
+          }
         } catch (error) {
-          console.error('Error while processing rpc request:', error, data.payload, methods)
+          console.error('Error while processing rpc request:', error, data.payload)
+          postMessage(ErrorShape.create(data, error))
         }
       }
     },
     { signal },
   )
+
+  if ('start' in to) {
+    to.start?.()
+  }
 }
 
 /**
@@ -166,12 +257,63 @@ export function expose<T extends object>(
  *
  * @param messenger - The Messenger to communicate with (e.g. Worker or Window)
  * @param options - Optional abort signal
- * @returns A proxy object that lets you call methods remotely
+ * @returns A proxy object for calling remote methods
+ *
+ * @example
+ * ```ts
+ * // Create RPC proxy (synchronous)
+ * const worker = rpc<WorkerMethods>(new Worker('worker.js'))
+ *
+ * // Call methods that return handle() get sub-proxies
+ * const renderer = await worker.init(transfer(canvas))
+ * await renderer.render()
+ *
+ * // Access underlying messenger
+ * worker[$MESSENGER].terminate()
+ * ```
  */
+// Overloads for specific messenger types to enable proper type inference
 export function rpc<T extends object>(
-  messenger: Messenger,
+  messenger: Worker,
+  options?: { signal?: AbortSignal },
+): RPC<T> & { [$MESSENGER]: Worker }
+
+export function rpc<T extends object>(
+  messenger: MessagePort,
+  options?: { signal?: AbortSignal },
+): RPC<T> & { [$MESSENGER]: MessagePort }
+
+export function rpc<T extends object>(
+  messenger: Window,
   options?: { signal?: AbortSignal },
-): RPC<T> {
+): RPC<T> & { [$MESSENGER]: Window }
+
+export function rpc<T extends object>(
+  messenger: BroadcastChannel,
+  options?: { signal?: AbortSignal },
+): RPC<T> & { [$MESSENGER]: BroadcastChannel }
+
+export function rpc<T extends object>(
+  messenger: ServiceWorker,
+  options?: { signal?: AbortSignal },
+): RPC<T> & { [$MESSENGER]: ServiceWorker }
+
+export function rpc<T extends object, M extends Messenger = Messenger>(
+  messenger: M,
+  options?: { signal?: AbortSignal },
+): RPC<T> & { [$MESSENGER]: M }
+
+// Implementation
+export function rpc<T extends object, M extends Messenger>(
+  messenger: M,
+  options?: { signal?: AbortSignal },
+): RPC<T> & { [$MESSENGER]: M } {
   const request = createRequester(messenger, options)
-  return createCommander<RPC<T>>((topics, args) => request(RPCPayloadShape.create(topics, args)))
+
+  const proxy = createRpcCommander<T>((topics, args) => {
+    const { args: processedArgs, transferables } = extractTransferables(args)
+    return request(RPCPayloadShape.create(topics, processedArgs), transferables)
+  })
+
+  return Object.assign(proxy, { [$MESSENGER]: messenger })
 }

package/src/{message-protocol.ts → protocol.ts}

@@ -22,7 +22,7 @@ export const $MESSENGER_RESPONSE = 'RPC_PROXY_RESPONSE'
 export const ResponseShape = createShape(
   v.object({
     [$MESSENGER_RESPONSE]: v.number(),
-    payload: v.unknown(),
+    payload: v.optional(v.unknown()),
   }),
   (request: RequestData, payload: any) => ({
     [$MESSENGER_RESPONSE]: request[$MESSENGER_REQUEST],
@@ -55,3 +55,13 @@ export const RPCPayloadShape = createShape(
   }),
   (topics: Array<string>, args: Array<any>) => ({ [$MESSENGER_RPC_REQUEST]: true, topics, args }),
 )
+
+// Handle response - when a method returns handle(), client creates sub-proxy
+export const $MESSENGER_HANDLE = 'RPC_PROXY_HANDLE'
+
+export const HandleResponseShape = createShape(
+  v.object({
+    [$MESSENGER_HANDLE]: v.string(), // namespace ID
+  }),
+  (namespaceId: string) => ({ [$MESSENGER_HANDLE]: namespaceId }),
+)

package/src/server-send-events/index.ts

@@ -1,7 +1,8 @@
-import { $MESSENGER_REQUEST, RequestShape, RPCPayloadShape } from '../message-protocol'
 import * as v from 'valibot'
+import { $MESSENGER_REQUEST, RequestShape, RPCPayloadShape } from '../protocol'
 import { RPC } from '../types'
-import { callMethod, createCommander, createIdRegistry, createShape, defer } from '../utils'
+import { callMethod, createCommander } from '../core'
+import { createIdRegistry, createShape, defer } from '../utils'
 
 const $SSE_RESPONSE_HEADER = 'RPC_SSE_RESPONSE'
 

package/src/stream/encoding.ts

@@ -351,10 +351,10 @@ export function createStreamCodec(config: Array<Codec>, fallback: PrimitiveCodec
 }
 
 // helper to concatenate Uint8Arrays
-function concat(...arrays: Array<Uint8Array>): Uint8Array {
+function concat(...arrays: Array<Uint8Array>): Uint8Array<ArrayBuffer> {
   const result = new Uint8Array(new ArrayBuffer(arrays.reduce((a, b) => a + b.length, 0)))
   let index = 0
-  return arrays.reduce<Uint8Array>((result, current) => {
+  return arrays.reduce<Uint8Array<ArrayBuffer>>((result, current) => {
     result.set(current, index)
     index += current.length
     return result

package/src/stream/index.ts

@@ -1,18 +1,7 @@
-import {
-  $MESSENGER_RESPONSE,
-  RequestShape,
-  ResponseShape,
-  RPCPayloadShape,
-} from '../message-protocol'
+import { $MESSENGER_RESPONSE, RequestShape, ResponseShape, RPCPayloadShape } from '../protocol'
 import { RPC } from '../types'
-import {
-  callMethod,
-  createCommander,
-  createPromiseRegistry,
-  createReadableStream,
-  defer,
-  streamToAsyncIterable,
-} from '../utils'
+import { callMethod, createCommander } from '../core'
+import { createPromiseRegistry, createReadableStream, defer, streamToAsyncIterable } from '../utils'
 
 interface StreamCodec {
   serialize(value: any, onChunk: (chunk: any) => void): void

package/src/types.ts

@@ -1,6 +1,11 @@
 export type Fn = (...arg: Array<any>) => any
 export type MaybePromise<T> = T | Promise<T>
 
+/** Check if T is a Handled type and extract the inner type */
+type UnwrapHandled<T> = T extends { readonly ['__rpc_handled__']: infer U extends object }
+  ? RPC<U>
+  : T
+
 // To prevent error: `Type instantiation is excessively deep and possibly infinite.`
 type isObject<T> = T extends object ? true : false
 
@@ -47,7 +52,7 @@ type FilterNoResponseMethod<T> = {
 /**********************************************************************************/
 
 export interface ResponseMethod<T extends Fn> {
-  (...args: Parameters<T>): Promise<ReturnType<T>>
+  (...args: Parameters<T>): Promise<UnwrapHandled<Awaited<ReturnType<T>>>>
 }
 
 export type ResponseRPCNode<T> = T extends Fn

package/src/utils.ts

@@ -58,26 +58,6 @@ export function defer<T = void>() {
   }
 }
 
-export function createCommander<T extends object = object>(
-  apply: (topics: Array<string>, args: Array<any>) => void,
-): T {
-  function _createCommander(
-    topics: Array<string>,
-    apply: (topics: Array<string>, args: Array<any>) => void,
-  ): T {
-    return new Proxy(function () {} as T, {
-      get(_, topic) {
-        if (typeof topic === 'symbol') return undefined
-        return _createCommander([...topics, topic], apply)
-      },
-      apply(_, __, args) {
-        return apply(topics, args)
-      },
-    })
-  }
-  return _createCommander([], apply)
-}
-
 /**
  * Creates a schema-backed shape definition with a validator and constructor.
  *
@@ -94,18 +74,6 @@ export function createShape<
   }
 }
 
-// expose-core.ts
-export function callMethod(methods: object, topics: string[], args: unknown[]) {
-  const method = topics.reduce((acc, topic) => {
-    const result = (acc as any)?.[topic]
-    return result
-  }, methods)
-  if (typeof method !== 'function') {
-    throw new Error(`Topics did not resolve to a function: [${topics.join(',')}]`)
-  }
-  return method(...args)
-}
-
 // NOTE:  safari does not implement AsyncIterator for ReadableStream
 //        see https://caniuse.com/mdn-api_readablestream_--asynciterator
 export function streamToAsyncIterable<T>(stream: ReadableStream<T>): AsyncIterable<T> {