aws-websocket-trpc.ts

export const socketProcedure = publicProcedure.use(connectionSocket);
export const socketConnectProcedure = socketProcedure.use(connectionSocketConnect);
export const socketDisconnectProcedure = socketProcedure.use(connectionSocketDisconnect);
export const socketRequestProcedure = socketProcedure.use(connectionSocketRequest);
export const socketEventProcedure = socketProcedure.use(connectionSocketEvent);

export const mySocketRouter = router({
  connect: socketConnectProcedure.mutation(async ({ctx, input}) => {
    // aws connection connect callback

    const awsConnectionId = ctx.connectionId;
  }),
  disconnect: socketDisconnectProcedure.mutation(async ({ctx, input}) => {
    // aws connection disconnect callback

    const awsConnectionId = ctx.connectionId;
  }),
  doAThing: socketRequestProcedure.input(z.object({thing: z.string()})).mutation(async ({ctx, input}) => {
    // requests are messages sent from the client. they have the aws connection id tied to this particular connection

    const awsConnectionId = ctx.connectionId;
    
    // you can call the $thing method here, which will internally do all of the aws stuff
    // you must pass the connectionId that you want to send this message to
    caller.sendThing(awsConnectionId, {thing: 'hi'});
    
  }),  
  $thing: socketEventProcedure.input(z.object({thing: z.string()})).mutation(() => {
    // this cannot be called, its simply a type placeholder input intellisense for the client
    // all "events" (things that can come back from the server) must start with $ for inference purposes.
  }),
  $otherThing: socketEventProcedure.input(z.object({thing: z.string()})).mutation(() => {
    // this cannot be called, its simply a type placeholder input intellisense for the client
    // all "events" (things that can come back from the server) must start with $ for inference purposes.
  }),
});
export const caller = makeSocketCaller(mySocketRouter.createCaller({}));

type FirstLetterCapital<T extends string> = T extends `${infer TFirst}${infer TRest}`
  ? `${Uppercase<TFirst>}${TRest}`
  : never;
type OnlyDollars<T> = T extends `$${infer TName}` ? FirstLetterCapital<TName> : never;
function makeSocketCaller<T extends Record<string, (...args: any) => any>>(
  item: T
): {[key in OnlyDollars<keyof T> as `send${key}`]: (connectionId: string, input: Parameters<T[key]>[0]) => void} {
  return item;
}

// client

const mySocket = await trpc.mySocket.connect({
  onThing(data) {
    // ^ data is inferred from the input of $thing
    console.log('received', data);
    
    mySocket.doAThing({thing: 'hi'});
    mySocket.disconnect();
  }, 
  onOtherThing(data) {
    // ^ data is inferred from the input of $otherThing
    console.log('received', data);
    
    mySocket.doAThing({thing: 'hi'});
    mySocket.disconnect();
  },
  onError(err) {
    console.error('error', err);
  },
});
if (mySocket.isConnected) {
  mySocket.doAThing({thing: 'hi'});
}