大TypeScript時代を支えるAsyncLocalStorageによるリクエストスコープなメタデータ管理 – LayerX エンジニアブログ

LayerX AI Agent ブログリレー 38日目の記事です。

こんにちは、LayerX バクラク事業部 Platform Engineering 部 Enabling グループに新卒入社した shibutani と申します。

弊社では早くからLLMを組み込んだ機能開発を行っており、多くのAIエージェント機能を提供しています。一部のプロダクトでは、そういったAI機能の開発においてVercel AI SDKを利用しています。そのような背景もあり、これまではGo言語を中心としたバックエンド開発が主流であった弊社でもTypeScriptを用いたバックエンド開発を行うケースが登場してきました。

TypeScriptにおけるリクエストスコープなメタデータの管理

バックエンド開発を行う際には、ログやテレメトリの送信など様々な場所でユーザーIDなどのリクエストスコープなメタデータが必要になります。特にLLMを用いたプロダクト開発を行う際には、精度改善のためにユーザーの行動をトラッキングできるTraceIDなどのメタデータが重要となります。

Go言語においてリクエストスコープなメタデータを扱う際には、多くの場合でcontext.Contextを利用します。その際には、関数呼び出しの引数に明示的にcontext.Contextオブジェクトを渡しますが、このような引数によるメタデータ管理が一般化している理由としては3rd Partyライブラリを含めて多くのライブラリのインターフェイスがcontext.Contextを引き受ける形になっていることが挙げられます。それによって、開発者はとりあえずcontext.Contextに情報を入れておけばよいと考えることができます。

しかし、TypeScript(ECMAScript/Node.js)においてはそのような引数によるメタデータ管理がGo言語ほど一般的ではありません。

では、TypeScriptにおいて、リクエストスコープなメタデータを管理するにはどうすればよいのでしょうか。その方法が次の章で紹介するAsyncLocalStorageです。

AsyncLocalStorage

Node.jsの公式サイトによるとAsyncLocalStorageとは次の通りです。

These classes are used to associate state and propagate it throughout callbacks and promise chains. They allow storing data throughout the lifetime of a web request or any other asynchronous duration. It is similar to thread-local storage in other languages.

AsyncLocalStorageはNode.jsが提供するデータストアであり、同期関数におけるコールバックや、非同期関数におけるPromiseチェーンの中でデータが保持されるため、リクエストスコープなメタデータを保存するのに活用できるというわけです。

それでは、AsyncLocalStorageの具体的な使用方法について順を追って説明していきます。AsyncLocalStorageの基本的な要素はstoreオブジェクト、run関数、enterWith関数です。

storeオブジェクトはAsyncLocalStorageにセットされた値そのものを指します。storeの形式に特に決まりはなく、プリミティブ型の値を直接入力することも可能ですが、一般的なユースケースではキーバリュー型のオブジェクトを入力することがほとんどかと思います。

run関数はstoreに新しい値をセットし、その値の有効範囲を決定する関数です。run関数の第二引数で指定されるコールバックもしくは非同期関数の中では新しいstoreの値が取得可能になります。また、兄弟関係にあるrun関数の間ではstoreの値は共有されません。

import { AsyncLocalStorage } from 'node:async_hooks';

const storage = new AsyncLocalStorage<{ key: string }>();

storage.run({ key: 'value1' }, () => {
  console.log(storage.getStore()!.key); 
});

console.log(storage.getStore()!.key) 

storage.run({ key: 'value2' }, () => {
  console.log(storage.getStore()!.key); 
});

enterWith関数も同様にstoreに新しい値をセットします。run関数との大きな違いとしては、enterWith関数ではコールバックを必要としません。その点は大きなメリットのように思えますが、値のスコープがわかりづらいというデメリットもあります。例えば下記のように、ある時点では値が含まれていたにもかかわらず、他の処理によって意図せずundefinedになってしまうということが起こりえます。

import { AsyncLocalStorage } from 'node:async_hooks';

const storage = new AsyncLocalStorage<{ key: string }>();

storage.enterWith({ key: 'value1' });
console.log(storage.getStore()!.key); 

storage.enterWith({ key: undefined });
console.log(storage.getStore()!.key); 

AsyncLocalStorage自体はシングルトンオブジェクトとして定義し、全ての箇所からそのオブジェクト経由でstoreにアクセスします。その際のスコープに応じて適切な値がstore経由で取り出せるという仕組みになっています。この仕組みによってリクエストスコープなメタデータを管理することが可能になります。

使いやすいインターフェイス設計を目指して

ここからはAsyncLocalStorageを用いたリクエストスコープなメタデータ管理を行うクラスを設計していきます。その際の大きな論点として、storeに対して新しい値を追加する際に、run関数を使うか、それともenterWith関数を使うかという論点があります。

run関数を使う場合、run関数の中で引数として指定されるコールバック関数の中でのみ新しいstoreが有効になります。そのため、storeのスコープがわかりやすく、暗黙的なstoreの汚染が起こりづらいというメリットがあります。その一方で、値が追加されるたびにコールバック関数の中で後続の処理を行う必要があり、コードの可読性が低下します。

enterWith関数を使う場合には、コールバック関数が不要となる一方で、新しい値が追加されたstoreのスコープがわかりづらく、リクエストハンドラの上流で使用された場合、他のテナントのリクエストとスコープが共有されるリスクがあります。

そこで、enterWithはテストなど限られた目的でのみ利用するべきものと定め、run関数を中心としたインターフェイス設計を行うこととしました。

さらに、弊社のバックエンド開発ではこれまでGo言語が中心的に利用されてきたため、下記のようにGo言語のcontext.Contextに近いインターフェイス設計としました。

import { AsyncLocalStorage } from 'node:async_hooks';

export type ContextKey<T> = {
  id: unique symbol;
  readonly $_type?: T; 
};

export class ContextCalledOutsideOfContextError extends Error {
  constructor() {
    super('Context was accessed outside of an active AsyncLocalStorage scope.');
  }
}

class Context {
  
  private readonly storage = new AsyncLocalStorage<Map<symbol, unknown>>();

  


  public withValue<T, R, Args extends unknown[]>(
    key: ContextKey<T>,
    value: T,
    fn: (...args: Args) => R,
    ...args: Args
  ): R {
    const parent = this.storage.getStore();
    const newStore = parent ? new Map(parent) : new Map<symbol, unknown>();
    newStore.set(key.id, value);
    return this.storage.run(newStore, fn, ...args);
  }

  



  public getValue<T>(key: ContextKey<T>): T | undefined {
    const store = this.storage.getStore();
    if (!store) {
      throw new ContextCalledOutsideOfContextError();
    }
    if (!store.has(key.id)) {
      return undefined;
    }
    return store.get(key.id) as T;
  }
}

export const context = new Context();

特にこだわっているのが、getValue関数の部分です。Go言語のcontext.Contextにおいて値を取得する際は、同時に値の型チェックをランタイムで行うことができる、型アサーション (例: foo.(string)) を実行するのが一般的です。しかし、TypeScriptにおいてはプリミティブ型を除き、ランタイムにおける厳密な型アサーションを行うことはできません。

そこで、事後的に型チェックを行うのではなく、あらかじめ定義するkeyの型に、対応する値の型を含めて、参照・挿入ともに型安全になるようにしました。これによって、型ガード関数などの実装者に依存する形での型チェックを行わずに、型安全に値を取り出すことが可能になります。

export const tenantIdKey: ContextKey<string> = { id: Symbol('tenantId') };

context.withValue(tenantIdKey, 't_123', async () => {  });

弊社での活用例

弊社ではWorkflow EngineのTemporalとVercelのAI SDKを組み合わせたAgent機能開発を行っています。その際に、Temporal ActivityのInterceptorにおいてTraceIdなどのメタデータをContextに追加し、同一のリクエストスコープで実行されるOpenTelemetryのSpanProcessorにおいてContextから取り出したTraceIDを付与しています。

class Interceptor implements ActivityInboundCallsInterceptor {
  public async execute(
    input: ActivityExecuteInput,
    next: Next<ActivityInboundCallsInterceptor, "execute">
  ): Promise<unknown> {
    const sessionId = randomId();
    const traceId = randomId();

    return await context.withValue(
      sessionIdKey,
      sessionId,
      async (input: ActivityExecuteInput) => {
        return await context.withValue(
          traceIdKey,
          traceId,
          async (input: ActivityExecuteInput) => {
            return await next(input);
          },
          input
        );
      },
      input
    );
  }
}


export class OriginalSpanProcessor extends SpanProcessor {
  public override onStart(span: Span, parentContext: Context): void {
    super.onStart(span, parentContext);

    const sessionId = requestContext.sessionId;
    const traceId = requestContext.traceId;

    if (traceId) {
      
      const spanAny = span as any;
      if (spanAny._spanContext) {
        spanAny._spanContext.traceId = traceId;
      }
    }

    if (sessionId) {
      span.setAttribute("session.id", sessionId);
    }
  }
}

こういった、同一のリクエストスコープで実行されるが、明示的に引数として値の受け渡しが難しいようなケースにおいては、同一のリクエストスコープ内のどこからでもアクセス可能で、フレームワークの実装に依存しないAsyncLocalStorageが非常に有用です。

おわりに

今回はTypeScriptにおいてリクエストスコープな情報を扱うための方法としてAsyncLocalStorageを紹介しました。

リクエストスコープだけでなく、任意の関数呼び出しに対してスコープを設定したデータ共有が可能となる強力な方法のため、ぜひ皆さんも試してみてください！