async local storage

async-local-storage.test.ts

  it('test local trace async', async () => {
    const s2 = new AsyncLocalStorage<string[]>()
    s2.enterWith(['first'])

    const localTraceAsync = <ReturnType, F extends (...args: any[]) => Promise<ReturnType>>(
      name: string,
      func: F,
    ): F => {
      const funcName = name
      return <F>(async (...args: any[]) => {
        const params = args.reduce(
          (acc, arg, i) => ({
            ...acc,
            [`arg${i}`]: arg, // generic way to label args
          }),
          {},
        )

        await s2.run([...(s2.getStore() ?? ['no-parent']), funcName], async () => {
          console.log('entering span', s2.getStore())
          try {
            const response = await func(...args)
            console.log('exiting span try', s2.getStore())
            return response
          } catch (e) {
            console.log('exiting span catch', s2.getStore())
            throw e
          }
        })
      })
    }

    const res = await localTraceAsync('parentAsync', async (firstArg: string, secondArg: number) => {
      const res2 = await localTraceAsync('asyncDummyFunc', asyncDummyFunc)('secondDummyFuncArg')

      const llm_res = await Promise.all([
        await localTraceAsync('prom1', asyncDummyFunc)('arg1'),
        await localTraceAsync('prom2', asyncDummyFunc)('arg2'),
        await localTraceAsync('prom3', asyncDummyFunc)('arg3'),
        await localTraceAsync('prom4', asyncDummyFunc)('arg4'),
      ])

      const res3 = await localTraceAsync('asyncDummyFunc', asyncDummyFunc)('thirdDummyFuncArg')

      return 'hello world'
    })('hi', 10)
  })