/ utils / workloadContext.ts
workloadContext.ts
 1  /**
 2   * Turn-scoped workload tag via AsyncLocalStorage.
 3   *
 4   * WHY a separate module from bootstrap/state.ts:
 5   * bootstrap is transitively imported by src/entrypoints/browser-sdk.ts, and
 6   * the browser bundle cannot import Node's async_hooks. This module is only
 7   * imported from CLI/SDK code paths that never end up in the browser build.
 8   *
 9   * WHY AsyncLocalStorage (not a global mutable slot):
10   * void-detached background agents (executeForkedSlashCommand, AgentTool)
11   * yield at their first await. The parent turn's synchronous continuation —
12   * including any `finally` block — runs to completion BEFORE the detached
13   * closure resumes. A global setWorkload('cron') at the top of the closure
14   * is deterministically clobbered. ALS captures context at invocation time
15   * and survives every await in that chain, isolated from the parent. Same
16   * pattern as agentContext.ts.
17   */
18  
19  import { AsyncLocalStorage } from 'async_hooks'
20  
21  /**
22   * Server-side sanitizer (_sanitize_entrypoint in claude_code.py) accepts
23   * only lowercase [a-z0-9_-]{0,32}. Uppercase stops parsing at char 0.
24   */
25  export type Workload = 'cron'
26  export const WORKLOAD_CRON: Workload = 'cron'
27  
28  const workloadStorage = new AsyncLocalStorage<{
29    workload: string | undefined
30  }>()
31  
32  export function getWorkload(): string | undefined {
33    return workloadStorage.getStore()?.workload
34  }
35  
36  /**
37   * Wrap `fn` in a workload ALS context. ALWAYS establishes a new context
38   * boundary, even when `workload` is undefined.
39   *
40   * The previous implementation short-circuited on `undefined` with
41   * `return fn()` — but that's a pass-through, not a boundary. If the caller
42   * is already inside a leaked cron context (REPL: queryGuard.end() →
43   * _notify() → React subscriber → scheduled re-render captures ALS at
44   * scheduling time → useQueueProcessor effect → executeQueuedInput → here),
45   * a pass-through lets `getWorkload()` inside `fn` return the leaked tag.
46   * Once leaked, it's sticky forever: every turn's end-notify re-propagates
47   * the ambient context to the next turn's scheduling chain.
48   *
49   * Always calling `.run()` guarantees `getWorkload()` inside `fn` returns
50   * exactly what the caller passed — including `undefined`.
51   */
52  export function runWithWorkload<T>(
53    workload: string | undefined,
54    fn: () => T,
55  ): T {
56    return workloadStorage.run({ workload }, fn)
57  }