@nmtjs/protocol 0.8.1 → 0.9.0

package/src/server/protocol.ts

@@ -1,16 +1,27 @@
-import { type Callback, createPromise, defer, throwError } from '@nmtjs/common'
-import { type Container, Hook, type Logger, Scope } from '@nmtjs/core'
+import { type Callback, defer, throwError } from '@nmtjs/common'
+import {
+  type AnyInjectable,
+  type Container,
+  Hook,
+  type Logger,
+  Scope,
+} from '@nmtjs/core'
 import { concat, decodeNumber, encodeNumber } from '../common/binary.ts'
 import type { ProtocolBlob, ProtocolBlobMetadata } from '../common/blob.ts'
 import { ErrorCode, ServerMessageType } from '../common/enums.ts'
 import type { ProtocolRPC } from '../common/types.ts'
-import type { ProtocolApi, ProtocolApiCallResult } from './api.ts'
+import {
+  isIterableResult,
+  type ProtocolApi,
+  type ProtocolApiCallOptions,
+} from './api.ts'
 import {
   Connection,
   ConnectionContext,
   type ConnectionOptions,
 } from './connection.ts'
 import type { Format } from './format.ts'
+import { ProtocolInjectables } from './injectables.ts'
 import type { ProtocolRegistry } from './registry.ts'
 import { ProtocolClientStream, ProtocolServerStream } from './stream.ts'
 import type { Transport } from './transport.ts'
@@ -43,13 +54,17 @@ export class ProtocolError extends Error {
   }
 }
 
+export type ProtocolConnectionTransport = {
+  send: Transport<any>['send']
+}
+
 export class ProtocolConnections {
   readonly #collection = new Map<
     string,
     {
       connection: Connection
       context: ConnectionContext
-      transport: Transport
+      transport: ProtocolConnectionTransport
     }
   >()
 
@@ -69,7 +84,7 @@ export class ProtocolConnections {
   }
 
   async add<T>(
-    transport: Transport<any>,
+    transport: ProtocolConnectionTransport,
     options: ConnectionOptions<T>,
     params: ResolveFormatParams,
   ) {
@@ -77,16 +92,19 @@ export class ProtocolConnections {
     const format = getFormat(this.application.format, params)
     const container = this.application.container.fork(Scope.Connection)
     const context = new ConnectionContext(container, format)
-
-    this.#collection.set(connection.id, { connection, context, transport })
-
-    await this.application.registry.hooks.call(
-      Hook.OnConnect,
-      { concurrent: false },
-      connection,
-    )
-
-    return { connection, context }
+    container.provide(ProtocolInjectables.connection, connection)
+    try {
+      await this.initialize(connection)
+      this.#collection.set(connection.id, { connection, context, transport })
+      return { connection, context }
+    } finally {
+      container.dispose().catch((error) => {
+        this.application.logger.error(
+          { error, connection },
+          'Error during disposing connection',
+        )
+      })
+    }
   }
 
   async remove(connectionId: string) {
@@ -128,6 +146,14 @@ export class ProtocolConnections {
       )
     }
   }
+
+  async initialize(connection: Connection) {
+    await this.application.registry.hooks.call(
+      Hook.OnConnect,
+      { concurrent: false },
+      connection,
+    )
+  }
 }
 
 export class ProtocolClientStreams {
@@ -162,12 +188,12 @@ export class ProtocolClientStreams {
 
   push(connectionId: string, streamId: number, chunk: ArrayBuffer) {
     const stream = this.get(connectionId, streamId)
-    stream.push(Buffer.from(chunk))
+    stream.write(Buffer.from(chunk))
   }
 
   end(connectionId: string, streamId: number) {
     const stream = this.get(connectionId, streamId)
-    stream.push(null)
+    stream.end(null)
     this.remove(connectionId, streamId)
   }
 
@@ -204,6 +230,7 @@ export class ProtocolServerStreams {
   }
 
   pull(connectionId: string, streamId: number) {
+    console.log('Pulling stream', streamId)
     const stream = this.get(connectionId, streamId)
     stream.resume()
   }
@@ -215,10 +242,16 @@ export class ProtocolServerStreams {
   }
 }
 
+export type ProtocolRPCOptions = {
+  signal?: AbortSignal
+  provides?: [AnyInjectable, any][]
+  metadata?: ProtocolApiCallOptions['metadata']
+}
+
 export class Protocol {
-  readonly connections: ProtocolConnections
-  readonly clientStreams: ProtocolClientStreams
-  readonly serverStreams: ProtocolServerStreams
+  #connections: ProtocolConnections
+  #clientStreams: ProtocolClientStreams
+  #serverStreams: ProtocolServerStreams
 
   constructor(
     protected readonly application: {
@@ -229,18 +262,44 @@ export class Protocol {
       api: ProtocolApi
     },
   ) {
-    this.connections = new ProtocolConnections(this.application)
-    this.clientStreams = new ProtocolClientStreams(this.connections)
-    this.serverStreams = new ProtocolServerStreams(this.connections)
+    this.#connections = new ProtocolConnections(this.application)
+    this.#clientStreams = new ProtocolClientStreams(this.#connections)
+    this.#serverStreams = new ProtocolServerStreams(this.#connections)
+  }
+
+  async call(options: ProtocolApiCallOptions) {
+    const { container, connection } = options
+    try {
+      return await this.application.api.call(options)
+    } catch (error) {
+      if (error instanceof ProtocolError === false) {
+        this.application.logger.error(
+          { error, connection },
+          'Error during RPC call',
+        )
+        throw new ProtocolError(
+          ErrorCode.InternalServerError,
+          'Internal server error',
+        )
+      }
+      throw error
+    } finally {
+      container.dispose().catch((error) => {
+        this.application.logger.error(
+          { error, connection },
+          "Error during disposing connection's container",
+        )
+      })
+    }
   }
 
   async rpc(
     connectionId: string,
     rpc: ProtocolRPC,
-    params: { signal?: AbortSignal } = {},
+    params: ProtocolRPCOptions = {},
   ) {
     const { connection, context, transport } =
-      this.connections.get(connectionId)
+      this.#connections.get(connectionId)
     const { calls, format } = context
     const { callId, namespace, procedure, payload } = rpc
     const abortController = new AbortController()
@@ -248,35 +307,37 @@ export class Protocol {
       ? AbortSignal.any([params.signal, abortController.signal])
       : abortController.signal
 
-    const call = Object.assign(createPromise<ProtocolApiCallResult>(), {
-      abort: () => abortController.abort(),
-    })
-
-    calls.set(callId, call)
-    call.promise.finally(() => calls.delete(callId))
+    calls.set(callId, abortController)
 
     const callIdEncoded = encodeNumber(callId, 'Uint32')
     const container = context.container.fork(Scope.Call)
 
+    if (params.provides) {
+      for (const [key, value] of params.provides) {
+        container.provide(key, value)
+      }
+    }
+
     try {
-      const response = await this.application.api.call({
+      const response = await this.call({
         connection,
         container,
         namespace,
         payload,
         procedure,
         signal,
+        metadata: params.metadata,
       })
 
       const responseEncoded = format.encoder.encodeRPC(
         {
           callId,
-          payload: response.output,
+          result: response.output,
         },
         {
           addStream: (blob) => {
-            const id = context.streamId++
-            const stream = this.serverStreams.add(connectionId, id, blob)
+            const streamId = context.streamId++
+            const stream = this.#serverStreams.add(connectionId, streamId, blob)
             stream.on('data', (chunk) => {
               stream.pause()
               const buf = Buffer.from(chunk)
@@ -284,43 +345,48 @@ export class Protocol {
                 connection,
                 ServerMessageType.ServerStreamPush,
                 concat(
-                  encodeNumber(id, 'Uint32'),
+                  encodeNumber(streamId, 'Uint32'),
                   (buf.buffer as ArrayBuffer).slice(
                     buf.byteOffset,
                     buf.byteOffset + buf.byteLength,
                   ),
                 ),
+                { callId, streamId },
               )
             })
             stream.on('error', (err) => {
               transport.send(
                 connection,
                 ServerMessageType.ServerStreamAbort,
-                encodeNumber(id, 'Uint32'),
+                encodeNumber(streamId, 'Uint32'),
+                { callId, streamId },
               )
             })
             stream.on('end', () => {
+              console.log('Stream ended')
               transport.send(
                 connection,
                 ServerMessageType.ServerStreamEnd,
-                encodeNumber(id, 'Uint32'),
+                encodeNumber(streamId, 'Uint32'),
+                { callId, streamId },
               )
             })
             return stream
           },
           getStream: (id) => {
-            return this.clientStreams.get(connectionId, id)
+            return this.#serverStreams.get(connectionId, id)
           },
         },
       )
 
       if ('subscription' in response) {
         throwError('Unimplemented')
-      } else if ('iterable' in response) {
+      } else if (isIterableResult(response)) {
         transport.send(
           connection,
           ServerMessageType.RpcStreamResponse,
           responseEncoded,
+          { callId },
         )
         try {
           const ab = new AbortController()
@@ -330,20 +396,28 @@ export class Protocol {
               ? response.iterable()
               : response.iterable
           for await (const chunk of iterable) {
-            if (ab.signal.aborted) break
+            ab.signal.throwIfAborted()
             const chunkEncoded = format.encoder.encode(chunk)
             transport.send(
               connection,
               ServerMessageType.RpcStreamChunk,
               concat(callIdEncoded, chunkEncoded),
+              { callId },
             )
           }
+          transport.send(
+            connection,
+            ServerMessageType.RpcStreamEnd,
+            callIdEncoded,
+            { callId },
+          )
         } catch (error) {
           this.application.logger.error(error)
           transport.send(
             connection,
             ServerMessageType.RpcStreamAbort,
             callIdEncoded,
+            { callId },
           )
         } finally {
           context.rpcStreams.delete(callId)
@@ -354,22 +428,10 @@ export class Protocol {
           connection,
           ServerMessageType.RpcResponse,
           responseEncoded,
+          { callId },
         )
       }
     } catch (error) {
-      if (error instanceof ProtocolError === false) {
-        this.application.logger.error(
-          { error, connection },
-          'Error during RPC call',
-        )
-
-        // biome-ignore lint/suspicious/noCatchAssign:
-        error = new ProtocolError(
-          ErrorCode.InternalServerError,
-          'Internal server error',
-        )
-      }
-
       const payload = format.encoder.encodeRPC(
         { callId, error },
         {
@@ -381,12 +443,17 @@ export class Protocol {
           },
         },
       )
-      transport.send(connection, ServerMessageType.RpcResponse, payload)
+
+      transport.send(connection, ServerMessageType.RpcResponse, payload, {
+        error,
+        callId,
+      })
     } finally {
+      calls.delete(callId)
       container.dispose().catch((error) => {
         this.application.logger.error(
           { error, connection },
-          "Error during disposing connection's container",
+          'Error during disposing connection',
         )
       })
     }
@@ -395,31 +462,34 @@ export class Protocol {
   async rpcRaw(
     connectionId: string,
     buffer: ArrayBuffer,
-    params: { signal?: AbortSignal } = {},
+    params: ProtocolRPCOptions = {},
   ) {
     const { connection, context, transport } =
-      this.connections.get(connectionId)
+      this.#connections.get(connectionId)
 
     const { format } = context
 
     const rpc = format.decoder.decodeRPC(buffer, {
-      addStream: (id, metadata) => {
-        const stream = this.clientStreams.add(
+      addStream: (streamId, callId, metadata) => {
+        return this.#clientStreams.add(
           connectionId,
-          id,
+          streamId,
           metadata,
           (size) => {
             transport.send(
               connection,
               ServerMessageType.ClientStreamPull,
-              concat(encodeNumber(id, 'Uint32'), encodeNumber(size, 'Uint32')),
+              concat(
+                encodeNumber(streamId, 'Uint32'),
+                encodeNumber(size, 'Uint32'),
+              ),
+              { callId, streamId },
             )
           },
         )
-        return stream
       },
       getStream: (id) => {
-        return this.serverStreams.get(connectionId, id)
+        return this.#clientStreams.get(connectionId, id)
       },
     })
 
@@ -427,7 +497,7 @@ export class Protocol {
   }
 
   rpcAbort(connectionId: string, callId: number) {
-    const { context } = this.connections.get(connectionId)
+    const { context } = this.#connections.get(connectionId)
     const call = context.calls.get(callId) ?? throwError('Call not found')
     call.abort()
   }
@@ -438,7 +508,7 @@ export class Protocol {
   }
 
   rpcStreamAbort(connectionId: string, callId: number) {
-    const { context } = this.connections.get(connectionId)
+    const { context } = this.#connections.get(connectionId)
     const ab =
       context.rpcStreams.get(callId) ?? throwError('Call stream not found')
     ab.abort()
@@ -452,4 +522,76 @@ export class Protocol {
   notify(connectionId: string, event, payload) {
     throw Error('Unimplemented')
   }
+
+  addConnection(
+    transport: ProtocolConnectionTransport,
+    options: ConnectionOptions,
+    params: ResolveFormatParams,
+  ) {
+    return this.#connections.add(transport, options, params)
+  }
+
+  removeConnection(connectionId: string) {
+    return this.#connections.remove(connectionId)
+  }
+
+  getConnection(connectionId: string) {
+    return this.#connections.get(connectionId)
+  }
+
+  initializeConnection(connection: Connection) {
+    return this.#connections.initialize(connection)
+  }
+
+  getClientStream(
+    connectionId: string,
+    streamId: number,
+  ): ProtocolClientStream {
+    return this.#clientStreams.get(connectionId, streamId)
+  }
+
+  addClientStream(
+    connectionId: string,
+    streamId: number,
+    metadata: ProtocolBlobMetadata,
+    read: Callback,
+  ) {
+    return this.#clientStreams.add(connectionId, streamId, metadata, read)
+  }
+
+  removeClientStream(connectionId: string, streamId: number) {
+    return this.#clientStreams.remove(connectionId, streamId)
+  }
+
+  pushClientStream(connectionId: string, streamId: number, chunk: ArrayBuffer) {
+    return this.#clientStreams.push(connectionId, streamId, chunk)
+  }
+
+  endClientStream(connectionId: string, streamId: number) {
+    return this.#clientStreams.end(connectionId, streamId)
+  }
+
+  abortClientStream(connectionId: string, streamId: number, error?: Error) {
+    return this.#clientStreams.abort(connectionId, streamId, error)
+  }
+
+  getServerStream(connectionId: string, streamId: number) {
+    return this.#serverStreams.get(connectionId, streamId)
+  }
+
+  addServerStream(connectionId: string, streamId: number, blob: ProtocolBlob) {
+    return this.#serverStreams.add(connectionId, streamId, blob)
+  }
+
+  removeServerStream(connectionId: string, streamId: number) {
+    return this.#serverStreams.remove(connectionId, streamId)
+  }
+
+  pullServerStream(connectionId: string, streamId: number) {
+    return this.#serverStreams.pull(connectionId, streamId)
+  }
+
+  abortServerStream(connectionId: string, streamId: number, error?: Error) {
+    return this.#serverStreams.abort(connectionId, streamId, error)
+  }
 }

package/src/server/stream.ts

@@ -1,13 +1,8 @@
-import {
-  PassThrough,
-  Readable,
-  type ReadableOptions,
-  type TransformOptions,
-} from 'node:stream'
+import { PassThrough, Readable, type ReadableOptions } from 'node:stream'
 import { ReadableStream } from 'node:stream/web'
 import type { ProtocolBlob, ProtocolBlobMetadata } from '../common/blob.ts'
 
-export class ProtocolClientStream extends Readable {
+export class ProtocolClientStream extends PassThrough {
   constructor(
     public readonly id: number,
     public readonly metadata: ProtocolBlobMetadata,

package/src/server/transport.ts

@@ -3,8 +3,10 @@ import type { BasePlugin, PluginContext } from '@nmtjs/core'
 import type { ServerMessageType } from '../common/enums.ts'
 import type { Connection } from './connection.ts'
 import { kTransportPlugin } from './constants.ts'
+import type { Format } from './format.ts'
 import type { Protocol } from './protocol.ts'
 import type { ProtocolRegistry } from './registry.ts'
+import type { ProtocolSendMetadata } from './types.ts'
 
 export interface Transport<T = unknown> {
   start: () => Promise<void>
@@ -13,12 +15,14 @@ export interface Transport<T = unknown> {
     connection: Connection<T>,
     messageType: ServerMessageType,
     buffer: ArrayBuffer,
-  ) => Async<void>
+    metadata: ProtocolSendMetadata,
+  ) => any
 }
 
 export interface TransportPluginContext extends PluginContext {
   protocol: Protocol
   registry: ProtocolRegistry
+  format: Format
 }
 
 export interface TransportPlugin<Type = unknown, Options = unknown>

package/src/server/types.ts

@@ -0,0 +1,21 @@
+import type { ProtocolServerBlobStream } from '../client/stream.ts'
+import type { ProtocolBlob, ProtocolBlobInterface } from '../common/blob.ts'
+import type { ProtocolClientStream } from './stream.ts'
+
+export type InputType<T> = T extends ProtocolBlobInterface
+  ? ProtocolClientStream
+  : T extends object
+    ? { [K in keyof T]: InputType<T[K]> }
+    : T
+
+export type OutputType<T> = T extends ProtocolBlobInterface
+  ? ProtocolBlob
+  : T extends object
+    ? { [K in keyof T]: OutputType<T[K]> }
+    : T
+
+export type ProtocolSendMetadata = {
+  streamId?: number
+  callId?: number
+  error?: any
+}

package/src/server/utils.ts

@@ -5,15 +5,22 @@ export type ResolveFormatParams = {
   acceptType?: string | null
 }
 
+export class UnsupportedFormatError extends Error {}
+
+export class UnsupportedContentTypeError extends UnsupportedFormatError {}
+
+export class UnsupportedAcceptTypeError extends UnsupportedFormatError {}
+
 export const getFormat = (
   format: Format,
   { acceptType, contentType }: ResolveFormatParams,
 ) => {
   const encoder = contentType ? format.supportsEncoder(contentType) : undefined
-  if (!encoder) throw new Error('Unsupported content-type')
+  if (!encoder)
+    throw new UnsupportedContentTypeError('Unsupported Content-Type')
 
   const decoder = acceptType ? format.supportsDecoder(acceptType) : undefined
-  if (!decoder) throw new Error('Unsupported accept')
+  if (!decoder) throw new UnsupportedAcceptTypeError('Unsupported Accept-Type')
 
   return {
     encoder,