diff --git a/src/backends/native/agent.ts b/src/backends/native/agent.ts index e3f13b5..972071d 100644 --- a/src/backends/native/agent.ts +++ b/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(); constructor(config: NativeAgentConfig) { this.modelClient = config.modelClient; @@ -134,48 +139,49 @@ export class NativeAgent { ): Promise { 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 - try { - if (userMessage.startsWith('!!think ') || userMessage === '!!think') { - this._thinking = true; - userMessage = userMessage.replace(/^!!think\s*/, '').trim() || 'Think about this.'; - } else { - this._thinking = false; + 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; } - - 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 { @@ -662,8 +668,9 @@ export class NativeAgent { : {}; const original = this.summarizeAudioToolArgs(args); - if (this._currentTurnAudioInput) { - this.applyAudioToolInput(args, this._currentTurnAudioInput); + const runTurnAudioInput = this._runContext.getStore()?.turnAudioInput; + if (runTurnAudioInput) { + this.applyAudioToolInput(args, runTurnAudioInput); this.logAudioArgsRewrite('latest_audio_preferred', 'latest_turn', original, args); return args; }