Add injected isomorphic environment

__tests__/integration.test.ts

@@ -28,17 +28,17 @@ export const TestServiceConstructor = () =>
       type: 'rpc',
       input: Type.Object({ n: Type.Number() }),
       output: Type.Object({ result: Type.Number() }),
-      async handler(state, msg) {
+      async handler(ctx, msg) {
         const { n } = msg.payload;
-        state.count += n;
-        return reply(msg, { result: state.count });
+        ctx.state.count += n;
+        return reply(msg, { result: ctx.state.count });
       },
     })
     .defineProcedure('echo', {
       type: 'stream',
       input: EchoRequest,
       output: EchoResponse,
-      async handler(_state, msgStream, returnStream) {
+      async handler(_ctx, msgStream, returnStream) {
         for await (const msg of msgStream) {
           const req = msg.payload;
           if (!req.ignore) {

index.ts

@@ -16,6 +16,7 @@ export type { ServerClient } from './router/client';
 export { createServer } from './router/server';
 export { asClientRpc, asClientStream } from './router/server.util';
 export type { Server } from './router/server';
+export type { ServiceContext, ServiceContextWithState } from './router/context';
 
 export { Transport } from './transport/types';
 export {

router/builder.ts

@@ -1,6 +1,7 @@
 import { TObject, Static, Type } from '@sinclair/typebox';
 import type { Pushable } from 'it-pushable';
 import { TransportMessage } from '../transport/message';
+import { ServiceContextWithState } from './context';
 
 export type ValidProcType = 'stream' | 'rpc';
 export type ProcListing = Record<
@@ -68,7 +69,7 @@ export type Procedure<
       input: I;
       output: O;
       handler: (
-        state: State,
+        context: ServiceContextWithState<State>,
         input: TransportMessage<Static<I>>,
       ) => Promise<TransportMessage<Static<O>>>;
       type: Ty;
@@ -77,7 +78,7 @@ export type Procedure<
       input: I;
       output: O;
       handler: (
-        state: State,
+        context: ServiceContextWithState<State>,
         input: AsyncIterable<TransportMessage<Static<I>>>,
         output: Pushable<TransportMessage<Static<O>>>,
       ) => Promise<void>;

router/context.ts

@@ -0,0 +1,28 @@
+/**
+ * The context for services/procedures. This is used only on
+ * the server.
+ *
+ * An important detail is that the state prop is always on
+ * this interface and it shouldn't be changed, removed, or
+ * extended. This prop is for the state of a service.
+ *
+ * You should use declaration merging to extend this interface
+ * with whatever you need. For example, if you need to access
+ * a database, you could do:
+ * ```ts
+ * declare module '@replit/river' {
+ *   interface ServiceContext {
+ *     db: Database;
+ *   }
+ * }
+ * ```
+ */
+export interface ServiceContext {
+  state: object | unknown;
+}
+
+/**
+ * The {@link ServiceContext} with state. This is what is passed to procedures.
+ */
+export type ServiceContextWithState<State extends object | unknown> =
+  ServiceContext & { state: State };

router/server.ts

@@ -5,6 +5,7 @@ import { Value } from '@sinclair/typebox/value';
 import { pushable } from 'it-pushable';
 import type { Pushable } from 'it-pushable';
 import { OpaqueTransportMessage, TransportMessage } from '../transport/message';
+import { ServiceContext, ServiceContextWithState } from './context';
 
 export interface Server<Services> {
   services: Services;
@@ -20,10 +21,29 @@ interface ProcStream {
 export async function createServer<Services extends Record<string, AnyService>>(
   transport: Transport,
   services: Services,
+  extendedContext?: Omit<ServiceContext, 'state'>,
 ): Promise<Server<Services>> {
-  // create streams for every stream procedure
+  const contextMap: Map<
+    AnyService,
+    ServiceContextWithState<object>
+  > = new Map();
   const streamMap: Map<string, ProcStream> = new Map();
+
+  function getContext(service: AnyService) {
+    const context = contextMap.get(service);
+
+    if (!context) {
+      throw new Error(`No context found for ${service.name}`);
+    }
+
+    return context;
+  }
+
   for (const [serviceName, service] of Object.entries(services)) {
+    // populate the context map
+    contextMap.set(service, { ...extendedContext, state: service.state });
+
+    // create streams for every stream procedure
     for (const [procedureName, proc] of Object.entries(service.procedures)) {
       const procedure = proc as Procedure<
         object,
@@ -39,7 +59,7 @@ export async function createServer<Services extends Record<string, AnyService>>(
           outgoing,
           doneCtx: Promise.all([
             // processing the actual procedure
-            procedure.handler(service.state, incoming, outgoing),
+            procedure.handler(getContext(service), incoming, outgoing),
             // sending outgoing messages back to client
             (async () => {
               for await (const response of outgoing) {
@@ -77,8 +97,10 @@ export async function createServer<Services extends Record<string, AnyService>>(
           procedure.type === 'rpc' &&
           Value.Check(procedure.input, inputMessage.payload)
         ) {
-          // synchronous rpc
-          const response = await procedure.handler(service.state, inputMessage);
+          const response = await procedure.handler(
+            getContext(service),
+            inputMessage,
+          );
           transport.send(response);
           return;
         } else if (

router/server.util.ts

@@ -6,15 +6,20 @@ import {
 } from '../transport/message';
 import { pushable } from 'it-pushable';
 import type { Pushable } from 'it-pushable';
+import { ServiceContext } from './context';
 
 export function asClientRpc<
   State extends object | unknown,
   I extends TObject,
   O extends TObject,
->(state: State, proc: Procedure<State, 'rpc', I, O>) {
+>(
+  state: State,
+  proc: Procedure<State, 'rpc', I, O>,
+  extendedContext?: Omit<ServiceContext, 'state'>,
+) {
   return (msg: Static<I>) =>
     proc
-      .handler(state, payloadToTransportMessage(msg))
+      .handler({ ...extendedContext, state }, payloadToTransportMessage(msg))
       .then((res) => res.payload);
 }
 
@@ -25,6 +30,7 @@ export function asClientStream<
 >(
   state: State,
   proc: Procedure<State, 'stream', I, O>,
+  extendedContext?: Omit<ServiceContext, 'state'>,
 ): [Pushable<Static<I>>, Pushable<Static<O>>] {
   const i = pushable<Static<I>>({ objectMode: true });
   const o = pushable<Static<O>>({ objectMode: true });
@@ -49,7 +55,7 @@ export function asClientStream<
 
   // handle
   (async () => {
-    await proc.handler(state, ri, ro);
+    await proc.handler({ ...extendedContext, state }, ri, ro);
     ro.end();
   })();