Refactor ServiceContext as generic

__tests__/cancellation.test.ts

@@ -4,9 +4,9 @@ import {
   Err,
   Ok,
   Procedure,
-  ServiceSchema,
   ValidProcType,
   createClient,
+  createServiceSchema,
   createServer,
 } from '../router';
 import { testMatrix } from '../testUtil/fixtures/matrix';
@@ -29,7 +29,8 @@ function makeMockHandler<T extends ValidProcType>(
 ) {
   return vi.fn<
     Procedure<
-      Record<string, never>,
+      object,
+      object,
       T,
       TObject,
       TObject | null,
@@ -39,6 +40,8 @@ function makeMockHandler<T extends ValidProcType>(
   >(impl);
 }
 
+const ServiceSchema = createServiceSchema();
+
 describe.each(testMatrix())(
   'clean handler cancellation ($transport.name transport, $codec.name codec)',
 

__tests__/cleanup.test.ts

@@ -15,8 +15,8 @@ import {
   Ok,
   Procedure,
   ProcedureHandlerContext,
-  ServiceSchema,
   createClient,
+  createServiceSchema,
   createServer,
 } from '../router';
 import {
@@ -503,13 +503,14 @@ describe('request finishing triggers signal onabort', async () => {
   ] as const)('handler aborts $procedureType', async ({ procedureType }) => {
     const clientTransport = getClientTransport('client');
     const serverTransport = getServerTransport();
-    const handler = vi.fn<(ctx: ProcedureHandlerContext<object>) => void>();
+    const handler =
+      vi.fn<(ctx: ProcedureHandlerContext<object, object>) => void>();
     const serverId = serverTransport.clientId;
     const serviceName = 'service';
     const procedureName = procedureType;
 
     const services = {
-      [serviceName]: ServiceSchema.define({
+      [serviceName]: createServiceSchema().define({
         // eslint-disable-next-line @typescript-eslint/no-unsafe-assignment, @typescript-eslint/no-unsafe-call, @typescript-eslint/no-explicit-any
         [procedureType]: (Procedure[procedureType] as any)({
           requestInit: Type.Object({}),
@@ -519,7 +520,11 @@ describe('request finishing triggers signal onabort', async () => {
               }
             : {}),
           responseData: Type.Object({}),
-          async handler({ ctx }: { ctx: ProcedureHandlerContext<object> }) {
+          async handler({
+            ctx,
+          }: {
+            ctx: ProcedureHandlerContext<object, object>;
+          }) {
             handler(ctx);
 
             return new Promise(() => {

__tests__/context.test.ts

@@ -5,14 +5,9 @@ import {
 } from '../testUtil/fixtures/cleanup';
 import { testMatrix } from '../testUtil/fixtures/matrix';
 import { TestSetupHelpers } from '../testUtil/fixtures/transports';
-import {
-  Ok,
-  Procedure,
-  ServiceSchema,
-  createClient,
-  createServer,
-} from '../router';
+import { Ok, Procedure, createClient, createServer } from '../router';
 import { Type } from '@sinclair/typebox';
+import { createServiceSchema } from '../router/services';
 
 describe('should handle incompatabilities', async () => {
   const { addPostTestCleanup, postTestCleanup } = createPostTestCleanups();
@@ -40,19 +35,25 @@ describe('should handle incompatabilities', async () => {
     // setup
     const clientTransport = getClientTransport('client');
     const serverTransport = getServerTransport();
+
+    const extendedContext = {
+      testctx: Math.random().toString(),
+    };
+
+    const ServiceSchema = createServiceSchema<typeof extendedContext>();
+
     const services = {
       testservice: ServiceSchema.define({
         testrpc: Procedure.rpc({
           requestInit: Type.Object({}),
           responseData: Type.String(),
           handler: async ({ ctx }) => {
-            return Ok((ctx as unknown as typeof extendedContext).testctx);
+            return Ok(ctx.testctx);
           },
         }),
       }),
     };
 
-    const extendedContext = { testctx: Math.random().toString() };
     createServer(serverTransport, services, {
       extendedContext,
     });
@@ -74,9 +75,13 @@ describe('should handle incompatabilities', async () => {
     const clientTransport = getClientTransport('client');
     const serverTransport = getServerTransport();
 
+    const extendedContext = { testctx: Math.random().toString() };
+
+    const ServiceSchema = createServiceSchema<typeof extendedContext>();
+
     const TestServiceScaffold = ServiceSchema.scaffold({
       initializeState: (ctx) => ({
-        fromctx: (ctx as unknown as typeof extendedContext).testctx,
+        fromctx: ctx.testctx,
       }),
     });
     const services = {
@@ -93,7 +98,43 @@ describe('should handle incompatabilities', async () => {
       }),
     };
 
+    createServer(serverTransport, services, {
+      extendedContext,
+    });
+    const client = createClient<typeof services>(
+      clientTransport,
+      serverTransport.clientId,
+    );
+    addPostTestCleanup(async () => {
+      await cleanupTransports([clientTransport, serverTransport]);
+    });
+
+    const res = await client.testservice.testrpc.rpc({});
+
+    expect(res).toEqual({ ok: true, payload: extendedContext.testctx });
+  });
+
+  test('should be able to access context in procedures', async () => {
+    // setup
+    const clientTransport = getClientTransport('client');
+    const serverTransport = getServerTransport();
+
     const extendedContext = { testctx: Math.random().toString() };
+
+    const ServiceSchema = createServiceSchema<typeof extendedContext>();
+
+    const services = {
+      testservice: ServiceSchema.define({
+        testrpc: Procedure.rpc({
+          requestInit: Type.Object({}),
+          responseData: Type.String(),
+          handler: async ({ ctx }) => {
+            return Ok(ctx.testctx);
+          },
+        }),
+      }),
+    };
+
     createServer(serverTransport, services, {
       extendedContext,
     });

__tests__/e2e.test.ts

@@ -31,7 +31,7 @@ import { testMatrix } from '../testUtil/fixtures/matrix';
 import { Type } from '@sinclair/typebox';
 import {
   Procedure,
-  ServiceSchema,
+  createServiceSchema,
   Ok,
   UNCAUGHT_ERROR_CODE,
   CANCEL_CODE,
@@ -949,7 +949,7 @@ describe.each(testMatrix())(
       });
 
       const services = {
-        test: ServiceSchema.define({
+        test: createServiceSchema().define({
           getData: Procedure.rpc({
             requestInit: Type.Object({}),
             responseData: Type.Object({

__tests__/invalid-request.test.ts

@@ -5,8 +5,8 @@ import {
   Ok,
   OkResult,
   Procedure,
-  ServiceSchema,
   createClient,
+  createServiceSchema,
   createServer,
 } from '../router';
 import { testMatrix } from '../testUtil/fixtures/matrix';
@@ -22,6 +22,8 @@ import { TestSetupHelpers } from '../testUtil/fixtures/transports';
 import { nanoid } from 'nanoid';
 import { getClientSendFn } from '../testUtil';
 
+const ServiceSchema = createServiceSchema();
+
 describe('cancels invalid request', () => {
   const { transport, codec } = testMatrix()[0];
   const opts = { codec: codec.codec };

__tests__/middleware.test.ts

@@ -12,7 +12,7 @@ import {
   createServer,
   Ok,
   Procedure,
-  ServiceSchema,
+  createServiceSchema,
   Middleware,
 } from '../router';
 import { createMockTransportNetwork } from '../testUtil/fixtures/mockTransport';
@@ -244,7 +244,7 @@ describe('middleware test', () => {
       readByMiddlewareSignal: boolean;
     }>();
 
-    const AsyncStorageSchemas = ServiceSchema.define({
+    const AsyncStorageSchemas = createServiceSchema().define({
       gimmeStore: Procedure.rpc({
         requestInit: Type.Object({}),
         responseData: Type.Object({}),

__tests__/typescript-stress.test.ts

@@ -1,6 +1,6 @@
-import { describe, expect, test } from 'vitest';
+import { assert, describe, expect, test } from 'vitest';
 import { Procedure } from '../router/procedures';
-import { ServiceSchema } from '../router/services';
+import { createServiceSchema } from '../router/services';
 import { Type } from '@sinclair/typebox';
 import { createServer } from '../router/server';
 import { createClient } from '../router/client';
@@ -12,7 +12,10 @@ import {
   ResultUnwrapOk,
   unwrapOrThrow,
 } from '../router/result';
-import { TestServiceSchema } from '../testUtil/fixtures/services';
+import {
+  testContext,
+  TestServiceWithContextSchema,
+} from '../testUtil/fixtures/services';
 import { readNextResult } from '../testUtil';
 import {
   createClientHandshakeOptions,
@@ -25,6 +28,7 @@ import { createMockTransportNetwork } from '../testUtil/fixtures/mockTransport';
 const requestData = Type.Union([
   Type.Object({ a: Type.Number() }),
   Type.Object({ c: Type.String() }),
+  Type.Object({ d: Type.Number() }),
 ]);
 const responseData = Type.Object({
   b: Type.Union([Type.Number(), Type.String()]),
@@ -41,17 +45,22 @@ const responseError = Type.Union([
 ]);
 
 const fnBody = Procedure.rpc<
-  Record<string, never>,
+  typeof testContext,
+  {
+    db: string;
+  },
   typeof requestData,
   typeof responseData,
   typeof responseError
 >({
   requestInit: requestData,
   responseData,
   responseError,
-  async handler({ reqInit }) {
+  async handler({ reqInit, ctx }) {
     if ('c' in reqInit) {
       return Ok({ b: reqInit.c });
+    } else if ('d' in reqInit) {
+      return Ok({ b: ctx.add(reqInit.d, reqInit.d) });
     } else {
       return Ok({ b: reqInit.a });
     }
@@ -61,7 +70,9 @@ const fnBody = Procedure.rpc<
 // typescript is limited to max 50 constraints
 // see: https://github.com/microsoft/TypeScript/issues/33541
 // we should be able to support more than that due to how we make services
-const StupidlyLargeServiceSchema = ServiceSchema.define({
+const StupidlyLargeServiceSchema = createServiceSchema<
+  typeof testContext
+>().define({
   f1: fnBody,
   f2: fnBody,
   f3: fnBody,
@@ -182,13 +193,16 @@ describe("ensure typescript doesn't give up trying to infer the types for large
       x1: StupidlyLargeServiceSchema,
       y1: StupidlyLargeServiceSchema,
       z1: StupidlyLargeServiceSchema,
-      test: TestServiceSchema,
+      test: TestServiceWithContextSchema,
     };
 
     const mockTransportNetwork = createMockTransportNetwork();
     const server = createServer(
       mockTransportNetwork.getServerTransport(),
       services,
+      {
+        extendedContext: testContext,
+      },
     );
 
     const client = createClient<typeof services>(
@@ -204,10 +218,42 @@ describe("ensure typescript doesn't give up trying to infer the types for large
     expect(server).toBeTruthy();
     expect(client).toBeTruthy();
   });
+
+  test('service with context should be able to access context in procedures', async () => {
+    const services = {
+      a: StupidlyLargeServiceSchema,
+      b: StupidlyLargeServiceSchema,
+    };
+    const mockTransportNetwork = createMockTransportNetwork();
+    const server = createServer(
+      mockTransportNetwork.getServerTransport(),
+      services,
+      {
+        extendedContext: testContext,
+      },
+    );
+
+    const client = createClient<typeof services>(
+      mockTransportNetwork.getClientTransport('client'),
+      'SERVER',
+      { eagerlyConnect: false },
+    );
+
+    const res = await client.a.f2.rpc({ d: 1 });
+    assert(res.ok);
+    expect(res.payload.b).toBe(2);
+
+    const res2 = await client.b.f11.rpc({ d: 10 });
+    assert(res2.ok);
+    expect(res2.payload.b).toBe(20);
+
+    expect(server).toBeTruthy();
+    expect(client).toBeTruthy();
+  });
 });
 
 const services = {
-  test: ServiceSchema.define({
+  test: createServiceSchema().define({
     rpc: Procedure.rpc({
       requestInit: Type.Object({ n: Type.Number() }),
       responseData: Type.Object({ n: Type.Number() }),

router/client.ts

@@ -140,9 +140,20 @@ type ServiceClient<Service extends AnyService> = {
  * @template Srv - The type of the server.
  */
 export type Client<
-  Services extends AnyServiceSchemaMap,
-  IS extends
-    InstantiatedServiceSchemaMap<Services> = InstantiatedServiceSchemaMap<Services>,
+  // Context is a server-side implementation detail that doesn't affect the client interface
+  // eslint-disable-next-line @typescript-eslint/no-explicit-any
+  Services extends AnyServiceSchemaMap<any>,
+  IS extends InstantiatedServiceSchemaMap<
+    // Context is a server-side implementation detail that doesn't affect the client interface
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    any,
+    Services
+  > = InstantiatedServiceSchemaMap<
+    // Context is a server-side implementation detail that doesn't affect the client interface
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    any,
+    Services
+  >,
 > = {
   [SvcName in keyof IS]: ServiceClient<IS[SvcName]>;
 };
@@ -204,7 +215,10 @@ const defaultClientOptions: ClientOptions = {
  * @param {Partial<ClientOptions>} providedClientOptions - The options for the client.
  * @returns The client for the server.
  */
-export function createClient<ServiceSchemaMap extends AnyServiceSchemaMap>(
+// We are using any here because the ServiceContext is a server-side implementation
+// detail that doesn't affect the client interface
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+export function createClient<ServiceSchemaMap extends AnyServiceSchemaMap<any>>(
   transport: ClientTransport<Connection>,
   serverId: TransportClientId,
   providedClientOptions: Partial<