seem need rewrite, can't make it work for router caller now

packages/server/src/builder.ts

@@ -244,7 +244,7 @@ export class Builder<TContext extends Context, TExtraContext extends Context> {
         ? LazyRouter<U>
         : never {
     return createLazyProcedureOrLazyRouter({
-      load: async () => (await load()).default,
+      load: () => load().then(m => m.default),
       middlewares: this.zz$b.middlewares,
     })
   }

packages/server/src/procedure-lazy.ts

@@ -56,7 +56,7 @@ export function decorateLazyProcedure<
   TFuncOutput extends SchemaOutput<TOutputSchema>,
 >(
   options: DecorateLazyProcedureOptions<TContext, TExtraContext, TInputSchema, TOutputSchema, TFuncOutput>,
-) {
+): DecoratedLazyProcedure<TContext, TExtraContext, TInputSchema, TOutputSchema, TFuncOutput> {
   const lazyProcedure: LazyProcedure<TContext, TExtraContext, TInputSchema, TOutputSchema, TFuncOutput> = {
     [LAZY_PROCEDURE_SYMBOL]: {
       load: async () => {

packages/server/src/router-builder.ts

@@ -1,3 +1,5 @@
+import type { DecoratedProcedure, Procedure } from './procedure'
+import type { DecoratedLazyProcedure, LazyProcedure } from './procedure-lazy'
 import type { HandledRouter, Router } from './router'
 import type { LazyRouter } from './router-lazy'
 import type { Context, MergeContext } from './types'
@@ -8,6 +10,7 @@ import {
   type Middleware,
 } from './middleware'
 import { decorateProcedure, isProcedure } from './procedure'
+import { decorateLazyProcedure, isLazyProcedure, LAZY_PROCEDURE_SYMBOL } from './procedure-lazy'
 import { createLazyProcedureOrLazyRouter } from './router-lazy'
 
 export class RouterBuilder<
@@ -84,42 +87,12 @@ export class RouterBuilder<
   router<URouter extends Router<TContext>>(
     router: URouter,
   ): HandledRouter<URouter> {
-    const handled: Router<TContext> = {}
-
-    for (const key in router) {
-      const item = router[key]
-
-      if (isProcedure(item)) {
-        const builderMiddlewares = this.zz$rb.middlewares ?? []
-        const itemMiddlewares = item.zz$p.middlewares ?? []
-
-        const middlewares = [
-          ...builderMiddlewares,
-          ...itemMiddlewares.filter(
-            item => !builderMiddlewares.includes(item),
-          ),
-        ]
-
-        const contract = DecoratedContractProcedure.decorate(
-          item.zz$p.contract,
-        ).addTags(...(this.zz$rb.tags ?? []))
-
-        handled[key] = decorateProcedure({
-          zz$p: {
-            ...item.zz$p,
-            contract: this.zz$rb.prefix
-              ? contract.prefix(this.zz$rb.prefix)
-              : contract,
-            middlewares,
-          },
-        })
-      }
-      else {
-        handled[key] = this.router(item as any)
-      }
-    }
-
-    return handled as HandledRouter<URouter>
+    return createRouterInternal({
+      current: router,
+      middlewares: this.zz$rb.middlewares,
+      tags: this.zz$rb.tags,
+      prefix: this.zz$rb.prefix,
+    }) as any
   }
 
   lazy<U extends Router<TContext>>(
@@ -131,3 +104,71 @@ export class RouterBuilder<
     }) as any
   }
 }
+
+/**
+ * @internal
+ */
+function createRouterInternal(options: {
+  current: Router<any> | Procedure<any, any, any, any, any> | LazyProcedure<any, any, any, any, any>
+  middlewares?: Middleware<any, any, any, any>[]
+  tags?: string[]
+  prefix?: HTTPPath
+}) {
+  if (isProcedure(options.current)) {
+    const builderMiddlewares = options.middlewares ?? []
+    const itemMiddlewares = options.current.zz$p.middlewares ?? []
+
+    const middlewares = [
+      ...builderMiddlewares,
+      ...itemMiddlewares.filter(
+        item => !builderMiddlewares.includes(item),
+      ),
+    ]
+
+    const contract = DecoratedContractProcedure.decorate(
+      options.current.zz$p.contract,
+    ).addTags(...(options.tags ?? []))
+
+    return decorateProcedure({
+      zz$p: {
+        ...options.current.zz$p,
+        contract: options.prefix
+          ? contract.prefix(options.prefix)
+          : contract,
+        middlewares,
+      },
+    })
+  }
+
+  let procedure: DecoratedLazyProcedure<any, any, any, any, any> | undefined
+
+  if (isLazyProcedure(options.current)) {
+    procedure = decorateLazyProcedure({
+      ...options.current[LAZY_PROCEDURE_SYMBOL],
+      middlewares: options.middlewares,
+    })
+  }
+
+  const recursive = new Proxy(procedure ?? {}, {
+    get(target, key) {
+      if (typeof key !== 'string') {
+        return Reflect.get(target, key)
+      }
+
+      const next = Reflect.get(options.current, key)
+
+      if ((typeof next !== 'object' && typeof next !== 'function') || next === null) {
+        return next
+      }
+
+      return createRouterInternal({
+        current: next,
+        middlewares: options.middlewares,
+        tags: options.tags,
+        prefix: options.prefix,
+      })
+    },
+  })
+
+  return recursive
+}

packages/server/src/router-caller.test.ts

@@ -29,14 +29,25 @@ describe('createRouterCaller', () => {
       ping,
       pong,
     },
-    lazy: osw.lazy(() => Promise.resolve({
+    lazy: osw.lazy(async () => ({
       default: {
         ping,
-        pong: osw.lazy(() => Promise.resolve({ default: pong })),
+        pong: osw.lazy(async () => ({ default: pong })),
       },
     })),
   })
 
+  it('test', async () => {
+    const caller = osw.lazy(() => Promise.resolve({
+      default: {
+        pong: osw.lazy(() => Promise.resolve({ default: pong })),
+      },
+    }))
+
+    // console.log(await caller.lazy.ping({ value: '123' }))
+    console.log(caller.pong())
+  })
+
   it('infer context', () => {
     createRouterCaller({
       router,
@@ -94,6 +105,10 @@ describe('createRouterCaller', () => {
       value: '123',
     })
 
+    expect(caller.lazy.pong({ value: '123' })).resolves.toEqual({
+      value: true,
+    })
+
     // @ts-expect-error - invalid input
     expect(caller.ping({ value: new Date('2023-01-01') })).rejects.toThrowError(
       'Validation input failed',
@@ -118,6 +133,14 @@ describe('createRouterCaller', () => {
           ping,
         },
       },
+      lazy: osw.lazy(() => Promise.resolve({
+        default: {
+          ping,
+          nested: {
+            ping: osw.lazy(() => Promise.resolve({ default: ping })),
+          },
+        },
+      })),
     })
 
     const caller = createRouterCaller({
@@ -132,5 +155,13 @@ describe('createRouterCaller', () => {
       'child',
       'ping',
     ])
+
+    expect(caller.lazy.ping('')).resolves.toEqual(['lazy', 'ping'])
+    expect(caller.lazy.nested.ping('')).resolves.toEqual(['lazy', 'nested', 'ping'])
+    expect(caller.lazy.nested.ping('')).resolves.toEqual([
+      'lazy',
+      'nested',
+      'ping',
+    ])
   })
 })

packages/server/src/router-caller.ts

@@ -56,7 +56,7 @@ function createRouterCallerInternal(
   const procedureCaller = isLazyProcedure(options.current) || isProcedure(options.current)
     ? createProcedureCaller({
       procedure: options.current,
-      context: options.context as any,
+      context: options.context,
       path: options.path,
     })
     : {}
@@ -76,7 +76,7 @@ function createRouterCallerInternal(
       return createRouterCallerInternal({
         current: next,
         context: options.context,
-        path: [...(options.path ?? []), key],
+        path: [...options.path, key],
       })
     },
   })

packages/server/src/router-lazy.ts

@@ -5,8 +5,8 @@ import { isProcedure, type Procedure } from './procedure'
 import { decorateLazyProcedure, isLazyProcedure, LAZY_PROCEDURE_SYMBOL } from './procedure-lazy'
 
 export type LazyRouter<TRouter extends Router<any>> = {
-  [K in keyof TRouter]: TRouter[K] extends DecoratedLazyProcedure<any, any, any, any, any>
-    ? TRouter[K]
+  [K in keyof TRouter]: TRouter[K] extends LazyProcedure<infer UContext, infer UExtraContext, infer UInputSchema, infer UOutputSchema, infer UFuncOutput>
+    ? DecoratedLazyProcedure<UContext, UExtraContext, UInputSchema, UOutputSchema, UFuncOutput>
     : TRouter[K] extends Procedure<infer UContext, infer UExtraContext, infer UInputSchema, infer UOutputSchema, infer UFuncOutput>
       ? DecoratedLazyProcedure<UContext, UExtraContext, UInputSchema, UOutputSchema, UFuncOutput>
       : TRouter[K] extends Router<any>
@@ -25,10 +25,12 @@ export function createLazyProcedureOrLazyRouter<
     : T extends Router<any>
       ? LazyRouter<T>
       : never {
-  return createLazyProcedureOrLazyRouterInternal({
+  const result = createLazyProcedureOrLazyRouterInternal({
     load: options.load,
     middlewares: options.middlewares,
   }) as any
+
+  return result
 }
 
 function createLazyProcedureOrLazyRouterInternal(
@@ -41,7 +43,7 @@ function createLazyProcedureOrLazyRouterInternal(
     const procedure = await options.load()
 
     if (isLazyProcedure(procedure)) {
-      return await procedure[LAZY_PROCEDURE_SYMBOL].load()
+      return procedure[LAZY_PROCEDURE_SYMBOL].load()
     }
 
     if (isProcedure(procedure)) {
@@ -64,15 +66,16 @@ function createLazyProcedureOrLazyRouterInternal(
 
       const loadNext: () => Promise<Router<any> | Procedure<any, any, any, any, any>> = async () => {
         const current = await options.load()
+        const next = Reflect.get(current, key)
 
-        if ((typeof current !== 'object' && typeof current !== 'function') || current === null) {
+        if ((typeof next !== 'object' && typeof next !== 'function') || next === null) {
           throw new Error('The loader reached the end of the chain')
         }
 
-        return Reflect.get(current, key) as any
+        return next
       }
 
-      return createLazyProcedureOrLazyRouter({
+      return createLazyProcedureOrLazyRouterInternal({
         load: loadNext,
         middlewares: options.middlewares,
       })

packages/server/src/router.ts

@@ -4,7 +4,7 @@ import type {
   SchemaInput,
   SchemaOutput,
 } from '@orpc/contract'
-import type { DecoratedLazyProcedure } from './procedure-lazy'
+import type { DecoratedLazyProcedure, LazyProcedure } from './procedure-lazy'
 import type { Context } from './types'
 import {
   isContractProcedure,
@@ -16,7 +16,7 @@ import {
 } from './procedure'
 
 export interface Router<TContext extends Context> {
-  [k: string]: Procedure<TContext, any, any, any, any> | DecoratedLazyProcedure<TContext, any, any, any, any> | Router<TContext>
+  [k: string]: Procedure<TContext, any, any, any, any> | LazyProcedure<TContext, any, any, any, any> | Router<TContext>
 }
 
 export type HandledRouter<TRouter extends Router<any>> = {
@@ -26,6 +26,12 @@ export type HandledRouter<TRouter extends Router<any>> = {
     infer UInputSchema,
     infer UOutputSchema,
     infer UFuncOutput
+  > | LazyProcedure<
+    infer UContext,
+    infer UExtraContext,
+    infer UInputSchema,
+    infer UOutputSchema,
+    infer UFuncOutput
   >
     ? DecoratedProcedure<
       UContext,