Isolate turn audio hydration with async local context

src/backends/native/agent.ts

@@ -10,6 +10,7 @@ import type { OutboundAttachmentCollector } from './attachments.js';
 import { buildUserMessage } from '../../models/media.js';
 import { getElevationWindow } from '../../security/elevation.js';
 import { auditLogger } from '../../audit/index.js';
+import { AsyncLocalStorage } from 'node:async_hooks';
 
 export interface ToolUseEvent {
   type: 'start' | 'end';
@@ -83,6 +84,10 @@ export interface NativeAgentTurnAudioInput {
   mime_type?: string;
 }
 
+interface NativeAgentRunContext {
+  turnAudioInput?: AudioToolInput;
+}
+
 export class NativeAgent {
   private static readonly EMPTY_RESPONSE_FALLBACK =
     'I could not generate a response for that. Please try again.';
@@ -106,7 +111,7 @@ export class NativeAgent {
   private _runInProgress = false;
   private _runAbortController?: AbortController;
   private modelTimeoutMs: number;
-  private _currentTurnAudioInput?: AudioToolInput;
+  private readonly _runContext = new AsyncLocalStorage<NativeAgentRunContext>();
 
   constructor(config: NativeAgentConfig) {
     this.modelClient = config.modelClient;
@@ -134,48 +139,49 @@ export class NativeAgent {
   ): Promise<string> {
     this._cancelRequested = false;
     this._runAbortController = new AbortController();
-    this._currentTurnAudioInput = this.normalizeTurnAudioInput(turnAudioInput) ?? this.extractLatestAudioInputFromAttachments(attachments);
+    const normalizedTurnAudioInput = this.normalizeTurnAudioInput(turnAudioInput)
+      ?? this.extractLatestAudioInputFromAttachments(attachments);
     if ('clearAbort' in this.modelClient && typeof this.modelClient.clearAbort === 'function') {
       this.modelClient.clearAbort();
     }
     this._runInProgress = true;
+    return await this._runContext.run({ turnAudioInput: normalizedTurnAudioInput }, async () => {
+      // Detect and strip !!think prefix for per-message thinking mode
+      try {
+        if (userMessage.startsWith('!!think ') || userMessage === '!!think') {
+          this._thinking = true;
+          userMessage = userMessage.replace(/^!!think\s*/, '').trim() || 'Think about this.';
+        } else {
+          this._thinking = false;
+        }
 
-    // Detect and strip !!think prefix for per-message thinking mode
+        const userMsg = buildUserMessage(userMessage, attachments);
-    try {
+
-      if (userMessage.startsWith('!!think ') || userMessage === '!!think') {
+        if (this.session) {
-        this._thinking = true;
+          this.session.addMessage(userMsg);
-        userMessage = userMessage.replace(/^!!think\s*/, '').trim() || 'Think about this.';
+        } else {
-      } else {
+          this.inMemoryHistory.push(userMsg);
-        this._thinking = false;
+        }
+
+        // If no tools configured, use the simple single-turn path
+        if (!this.toolRegistry || !this.toolExecutor) {
+          return await this.singleTurn();
+        }
+
+        return await this.toolLoop();
+      } catch (error) {
+        if (this.isAbortError(error)) {
+          const cancelledMsg = 'Operation cancelled by user.';
+          this.addToHistory({ role: 'assistant', content: cancelledMsg });
+          return cancelledMsg;
+        }
+        throw error;
+      } finally {
+        this._runInProgress = false;
+        this._cancelRequested = false;
+        this._runAbortController = undefined;
       }
-
+    });
-      const userMsg = buildUserMessage(userMessage, attachments);
-
-      if (this.session) {
-        this.session.addMessage(userMsg);
-      } else {
-        this.inMemoryHistory.push(userMsg);
-      }
-
-      // If no tools configured, use the simple single-turn path
-      if (!this.toolRegistry || !this.toolExecutor) {
-        return await this.singleTurn();
-      }
-
-      return await this.toolLoop();
-    } catch (error) {
-      if (this.isAbortError(error)) {
-        const cancelledMsg = 'Operation cancelled by user.';
-        this.addToHistory({ role: 'assistant', content: cancelledMsg });
-        return cancelledMsg;
-      }
-      throw error;
-    } finally {
-      this._runInProgress = false;
-      this._cancelRequested = false;
-      this._runAbortController = undefined;
-      this._currentTurnAudioInput = undefined;
-    }
   }
 
   private async singleTurn(): Promise<string> {
@@ -662,8 +668,9 @@ export class NativeAgent {
       : {};
     const original = this.summarizeAudioToolArgs(args);
 
-    if (this._currentTurnAudioInput) {
+    const runTurnAudioInput = this._runContext.getStore()?.turnAudioInput;
-      this.applyAudioToolInput(args, this._currentTurnAudioInput);
+    if (runTurnAudioInput) {
+      this.applyAudioToolInput(args, runTurnAudioInput);
       this.logAudioArgsRewrite('latest_audio_preferred', 'latest_turn', original, args);
       return args;
     }