Isolate turn audio hydration with async local context

This commit is contained in:
William Valentin
2026-02-22 21:41:49 -08:00
parent a761813375
commit b477705806
+13 -6
View File
@@ -10,6 +10,7 @@ import type { OutboundAttachmentCollector } from './attachments.js';
import { buildUserMessage } from '../../models/media.js'; import { buildUserMessage } from '../../models/media.js';
import { getElevationWindow } from '../../security/elevation.js'; import { getElevationWindow } from '../../security/elevation.js';
import { auditLogger } from '../../audit/index.js'; import { auditLogger } from '../../audit/index.js';
import { AsyncLocalStorage } from 'node:async_hooks';
export interface ToolUseEvent { export interface ToolUseEvent {
type: 'start' | 'end'; type: 'start' | 'end';
@@ -83,6 +84,10 @@ export interface NativeAgentTurnAudioInput {
mime_type?: string; mime_type?: string;
} }
interface NativeAgentRunContext {
turnAudioInput?: AudioToolInput;
}
export class NativeAgent { export class NativeAgent {
private static readonly EMPTY_RESPONSE_FALLBACK = private static readonly EMPTY_RESPONSE_FALLBACK =
'I could not generate a response for that. Please try again.'; 'I could not generate a response for that. Please try again.';
@@ -106,7 +111,7 @@ export class NativeAgent {
private _runInProgress = false; private _runInProgress = false;
private _runAbortController?: AbortController; private _runAbortController?: AbortController;
private modelTimeoutMs: number; private modelTimeoutMs: number;
private _currentTurnAudioInput?: AudioToolInput; private readonly _runContext = new AsyncLocalStorage<NativeAgentRunContext>();
constructor(config: NativeAgentConfig) { constructor(config: NativeAgentConfig) {
this.modelClient = config.modelClient; this.modelClient = config.modelClient;
@@ -134,12 +139,13 @@ export class NativeAgent {
): Promise<string> { ): Promise<string> {
this._cancelRequested = false; this._cancelRequested = false;
this._runAbortController = new AbortController(); 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') { if ('clearAbort' in this.modelClient && typeof this.modelClient.clearAbort === 'function') {
this.modelClient.clearAbort(); this.modelClient.clearAbort();
} }
this._runInProgress = true; this._runInProgress = true;
return await this._runContext.run({ turnAudioInput: normalizedTurnAudioInput }, async () => {
// Detect and strip !!think prefix for per-message thinking mode // Detect and strip !!think prefix for per-message thinking mode
try { try {
if (userMessage.startsWith('!!think ') || userMessage === '!!think') { if (userMessage.startsWith('!!think ') || userMessage === '!!think') {
@@ -174,8 +180,8 @@ export class NativeAgent {
this._runInProgress = false; this._runInProgress = false;
this._cancelRequested = false; this._cancelRequested = false;
this._runAbortController = undefined; this._runAbortController = undefined;
this._currentTurnAudioInput = undefined;
} }
});
} }
private async singleTurn(): Promise<string> { private async singleTurn(): Promise<string> {
@@ -662,8 +668,9 @@ export class NativeAgent {
: {}; : {};
const original = this.summarizeAudioToolArgs(args); 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); this.logAudioArgsRewrite('latest_audio_preferred', 'latest_turn', original, args);
return args; return args;
} }