From 271fb8468d374e7ed378f6e2cbdfff71db7d0d77 Mon Sep 17 00:00:00 2001 From: likun Date: Mon, 15 Jun 2026 11:11:07 +0800 Subject: [PATCH] Wire AiSdkFlow into SessionManager runtime path --- .../runtime/src/__tests__/ai-sdk-flow.test.ts | 26 +++ .../src/__tests__/session-manager.test.ts | 15 +- packages/runtime/src/agent-run.ts | 185 ++++++++++-------- packages/runtime/src/ai-sdk-flow.ts | 31 ++- packages/runtime/src/session-manager.ts | 72 +++---- 5 files changed, 213 insertions(+), 116 deletions(-) diff --git a/packages/runtime/src/__tests__/ai-sdk-flow.test.ts b/packages/runtime/src/__tests__/ai-sdk-flow.test.ts index f6d145ae5b..753567a96e 100644 --- a/packages/runtime/src/__tests__/ai-sdk-flow.test.ts +++ b/packages/runtime/src/__tests__/ai-sdk-flow.test.ts @@ -369,6 +369,32 @@ describe('AiSdkFlow seam', () => { assert.equal(isTerminalRuntimeEvent(out[0]), true); }); + test('can keep draining backend events after a terminal while coalescing duplicate terminals', async () => { + const seen: SessionEvent[] = []; + const backend = new ScriptedBackend({ + events: [ + ev({ type: 'abort', reason: 'user_stop' }), + ev({ type: 'text_delta', messageId: 'm1', text: 'cleanup-after-terminal' }), + ev({ type: 'complete', stopReason: 'user_stop' }), + ], + }); + const flow = new AiSdkFlow({ + backend, + drainAfterTerminal: true, + onSessionEvent: (sessionEvent) => { + seen.push(sessionEvent); + }, + }); + const out = await collect(flow.run(ctx, { text: 'hi', context: [] })); + + assert.deepEqual(seen.map((event) => event.type), ['abort', 'text_delta', 'complete']); + assert.deepEqual( + out.map((event) => event.content?.kind ?? event.status ?? null), + ['aborted', 'text'], + ); + assert.equal(out.filter(isTerminalRuntimeEvent).length, 1); + }); + test('RuntimeRunner consumes AiSdkFlow abort as one coherent failed outcome', async () => { const backend = new ScriptedBackend({ events: [ diff --git a/packages/runtime/src/__tests__/session-manager.test.ts b/packages/runtime/src/__tests__/session-manager.test.ts index a8f88ce3b1..75c569defa 100644 --- a/packages/runtime/src/__tests__/session-manager.test.ts +++ b/packages/runtime/src/__tests__/session-manager.test.ts @@ -1,4 +1,5 @@ import { describe, test } from 'node:test'; +import { readFile } from 'node:fs/promises'; import { DEEP_RESEARCH_SESSION_LABEL, deriveTurnRecords } from '@maka/core'; import type { CreateSessionInput, @@ -246,12 +247,24 @@ describe('SessionManager permission mode updates', () => { expect(result.events.map((event) => event.sessionId)).toEqual([session.id, session.id, session.id]); expect(result.events.map((event) => event.turnId)).toEqual(['turn-1', 'turn-1', 'turn-1']); expect(result.events.map((event) => event.role)).toEqual(['user', 'model', 'system']); - expect(result.events.map((event) => event.id)).toEqual(['id-3', 'turn-1-delta', 'turn-1-complete']); + expect(result.events.map((event) => event.id)).toEqual(['id-7', 'turn-1-delta', 'turn-1-complete']); expect(result.events[0]?.content).toEqual({ kind: 'text', text: 'hello' }); expect(result.events[1]?.content).toEqual({ kind: 'text', text: 'ok' }); expect(result.events[2]?.status).toBe('completed'); }); + test('sendMessage production source uses AiSdkFlow instead of an inline mapper flow', async () => { + const source = await readFile(new URL('../../src/session-manager.ts', import.meta.url), 'utf8'); + const sendMessageSource = source.slice( + source.indexOf('async *sendMessage'), + source.indexOf('async stopSession'), + ); + + expect(sendMessageSource.includes('new AiSdkFlow')).toBe(true); + expect(sendMessageSource.includes('mapSessionEventToRuntimeEvent')).toBe(false); + expect(sendMessageSource.includes('createSessionEventMapMemory')).toBe(false); + }); + test('rejects backend configuration updates while a turn is actively streaming', async () => { const store = new MemorySessionStore(); const backends = new BackendRegistry(); diff --git a/packages/runtime/src/agent-run.ts b/packages/runtime/src/agent-run.ts index e618c5e71f..50f7071a49 100644 --- a/packages/runtime/src/agent-run.ts +++ b/packages/runtime/src/agent-run.ts @@ -55,6 +55,11 @@ export interface AgentRunInput { hooks: AgentRunHooks; } +export interface AgentRunBeginResult { + backend: AgentBackend; + backendInput: BackendSendInput; +} + export class AgentRun { readonly runId: string; readonly sessionId: string; @@ -69,6 +74,11 @@ export class AgentRun { private runStoreAvailable = true; private failureClass: string | undefined; private failureMessage: string | undefined; + private lastTs = 0; + private sawCompletion = false; + private finalStatus: { status: SessionStatus; blockedReason?: SessionBlockedReason } | undefined; + private turnFailed = false; + private finalized = false; constructor(private readonly input: AgentRunInput) { this.runId = input.newId(); @@ -97,6 +107,21 @@ export class AgentRun { } async *execute(): AsyncIterable { + try { + const begin = await this.begin(); + for await (const ev of begin.backend.send(begin.backendInput)) { + await this.recordSessionEvent(ev); + yield ev; + } + } catch (error) { + await this.recordFailure(error); + throw error; + } finally { + await this.finalize(); + } + } + + async begin(): Promise { await this.createRunRecord(); const userMsg: UserMessage = { @@ -110,95 +135,99 @@ export class AgentRun { await this.input.store.appendMessage(this.sessionId, userMsg); await this.input.hooks.appendTurnState(this.sessionId, this.turnId, 'running', this.lineage); - let lastTs = this.input.now(); - let sawCompletion = false; - let finalStatus: { status: SessionStatus; blockedReason?: SessionBlockedReason } | undefined; - let turnFailed = false; + this.lastTs = this.input.now(); - try { - if (!this.header.connectionLocked) { - this.header = await this.input.hooks.updateHeader(this.sessionId, { connectionLocked: true }); - } + if (!this.header.connectionLocked) { + this.header = await this.input.hooks.updateHeader(this.sessionId, { connectionLocked: true }); + } - this.active = await this.input.hooks.ensureActive(this.sessionId, this.header); - this.input.hooks.registerRun(this.active, this); - await this.markRunStarted(lastTs); + this.active = await this.input.hooks.ensureActive(this.sessionId, this.header); + this.input.hooks.registerRun(this.active, this); + await this.markRunStarted(this.lastTs); - await this.input.hooks.updateStatus(this.sessionId, 'running', undefined, lastTs); + await this.input.hooks.updateStatus(this.sessionId, 'running', undefined, this.lastTs); - const backendInput: BackendSendInput = { + return { + backend: this.active.backend, + backendInput: { turnId: this.turnId, text: this.input.userInput.text, ...(this.input.userInput.attachments ? { attachments: this.input.userInput.attachments } : {}), context: await this.input.store.readMessages(this.sessionId), - }; - for await (const ev of this.active.backend.send(backendInput)) { - lastTs = ev.ts; - const transition = statusFromEvent(ev); - if (transition && !this.stopped) { - await this.input.hooks.updateStatus(this.sessionId, transition.status, transition.blockedReason, ev.ts); - this.recordStatusFromTransition(ev, transition, ev.ts); - } - if ((ev.type === 'complete' || ev.type === 'abort') && !turnFailed) { - sawCompletion = true; - finalStatus = this.stopped - ? { status: 'aborted' } - : (transition ?? { status: 'active' }); - const turnStatus = turnStatusFromEvent(ev); - if (turnStatus && !this.stopped) { - await this.input.hooks.appendTurnState(this.sessionId, this.turnId, turnStatus.status, this.lineage, { - ts: ev.ts, - errorClass: turnStatus.errorClass, - }); - } - } - if (ev.type === 'error') { - turnFailed = true; - finalStatus = transition ?? { status: 'blocked', blockedReason: 'unknown' }; - await this.input.hooks.appendTurnState(this.sessionId, this.turnId, 'failed', this.lineage, { - ts: ev.ts, - errorClass: ev.reason ?? ev.code ?? 'unknown', - }); - this.markRunFailed(ev.reason ?? ev.code ?? 'unknown', ev.message, ev.ts); - } - yield ev; - } - } catch (error) { - finalStatus = { status: 'blocked', blockedReason: 'unknown' }; - await this.input.hooks.appendTurnState(this.sessionId, this.turnId, 'failed', this.lineage, { - errorClass: error instanceof Error ? error.name : 'unknown', - }).catch(() => {}); - this.markRunFailed(error instanceof Error ? error.name : 'unknown', errorMessage(error), this.input.now()); - throw error; - } finally { - if (this.active) { - this.input.hooks.unregisterRun(this.active, this); - if (this.stopped) finalStatus = { status: 'aborted' }; - } - const nextStatus = this.active && this.active.activeRuns.size > 0 - ? { status: 'running' as const } - : (finalStatus ?? { status: 'active' as const }); - try { - await this.input.hooks.updateHeader(this.sessionId, { - lastUsedAt: lastTs, - lastMessageAt: lastTs, - hasUnread: true, - ...statusPatch(nextStatus.status, lastTs, nextStatus.blockedReason), + }, + }; + } + + async recordSessionEvent(ev: SessionEvent): Promise { + this.lastTs = ev.ts; + const transition = statusFromEvent(ev); + if (transition && !this.stopped) { + await this.input.hooks.updateStatus(this.sessionId, transition.status, transition.blockedReason, ev.ts); + this.recordStatusFromTransition(ev, transition, ev.ts); + } + if ((ev.type === 'complete' || ev.type === 'abort') && !this.turnFailed) { + this.sawCompletion = true; + this.finalStatus = this.stopped + ? { status: 'aborted' } + : (transition ?? { status: 'active' }); + const turnStatus = turnStatusFromEvent(ev); + if (turnStatus && !this.stopped) { + await this.input.hooks.appendTurnState(this.sessionId, this.turnId, turnStatus.status, this.lineage, { + ts: ev.ts, + errorClass: turnStatus.errorClass, }); - } catch { - // The user-visible turn already completed; preserve existing behavior. - } - if (sawCompletion) { - await this.input.store.appendMessage(this.sessionId, { - type: 'system_note', - id: this.input.newId(), - turnId: this.turnId, - ts: lastTs, - kind: 'session_resume', - } satisfies SystemNoteMessage).catch(() => {}); } - await this.finishRun(finalStatus, lastTs); } + if (ev.type === 'error') { + this.turnFailed = true; + this.finalStatus = transition ?? { status: 'blocked', blockedReason: 'unknown' }; + await this.input.hooks.appendTurnState(this.sessionId, this.turnId, 'failed', this.lineage, { + ts: ev.ts, + errorClass: ev.reason ?? ev.code ?? 'unknown', + }); + this.markRunFailed(ev.reason ?? ev.code ?? 'unknown', ev.message, ev.ts); + } + } + + async recordFailure(error: unknown): Promise { + this.finalStatus = { status: 'blocked', blockedReason: 'unknown' }; + await this.input.hooks.appendTurnState(this.sessionId, this.turnId, 'failed', this.lineage, { + errorClass: error instanceof Error ? error.name : 'unknown', + }).catch(() => {}); + this.markRunFailed(error instanceof Error ? error.name : 'unknown', errorMessage(error), this.input.now()); + } + + async finalize(): Promise { + if (this.finalized) return; + this.finalized = true; + const lastTs = this.lastTs || this.input.now(); + if (this.active) { + this.input.hooks.unregisterRun(this.active, this); + if (this.stopped) this.finalStatus = { status: 'aborted' }; + } + const nextStatus = this.active && this.active.activeRuns.size > 0 + ? { status: 'running' as const } + : (this.finalStatus ?? { status: 'active' as const }); + try { + await this.input.hooks.updateHeader(this.sessionId, { + lastUsedAt: lastTs, + lastMessageAt: lastTs, + hasUnread: true, + ...statusPatch(nextStatus.status, lastTs, nextStatus.blockedReason), + }); + } catch { + // The user-visible turn already completed; preserve existing behavior. + } + if (this.sawCompletion) { + await this.input.store.appendMessage(this.sessionId, { + type: 'system_note', + id: this.input.newId(), + turnId: this.turnId, + ts: lastTs, + kind: 'session_resume', + } satisfies SystemNoteMessage).catch(() => {}); + } + await this.finishRun(this.finalStatus, lastTs); } private async createRunRecord(): Promise { diff --git a/packages/runtime/src/ai-sdk-flow.ts b/packages/runtime/src/ai-sdk-flow.ts index bf18444275..ee3c6a293e 100644 --- a/packages/runtime/src/ai-sdk-flow.ts +++ b/packages/runtime/src/ai-sdk-flow.ts @@ -373,6 +373,21 @@ export function mapSessionEventToRuntimeEvent( export interface AiSdkFlowInput { /** The wrapped stepping engine. Production: AiSdkBackend. Tests: any AgentBackend. */ backend: AgentBackend; + /** + * Optional production projection hook. Called for every raw backend + * SessionEvent after it has been mapped to a RuntimeEvent and before the + * RuntimeEvent is yielded/coalesced. + */ + onSessionEvent?: (sessionEvent: SessionEvent, runtimeEvent: RuntimeEvent) => Promise | void; + /** Called if the wrapped backend stream throws. */ + onError?: (error: unknown) => Promise | void; + /** Called after backend streaming finishes, errors, or is abandoned. */ + onFinally?: () => Promise | void; + /** + * Keep consuming backend events after the first terminal RuntimeEvent. + * Duplicate terminal RuntimeEvents are still coalesced from flow output. + */ + drainAfterTerminal?: boolean; } /** @@ -392,11 +407,19 @@ export class AiSdkFlow implements AgentFlow, AgentFlowControl { readonly kind: string; readonly sessionId: string; private readonly backend: AgentBackend; + private readonly onSessionEvent: AiSdkFlowInput['onSessionEvent']; + private readonly onError: AiSdkFlowInput['onError']; + private readonly onFinally: AiSdkFlowInput['onFinally']; + private readonly drainAfterTerminal: boolean; constructor(input: AiSdkFlowInput) { this.backend = input.backend; this.sessionId = input.backend.sessionId; this.kind = input.backend.kind; + this.onSessionEvent = input.onSessionEvent; + this.onError = input.onError; + this.onFinally = input.onFinally; + this.drainAfterTerminal = input.drainAfterTerminal ?? false; } /** The wrapped backend (exposed for runners that need the raw control surface). */ @@ -437,18 +460,24 @@ export class AiSdkFlow implements AgentFlow, AgentFlowControl { context: input.context, })) { const runtimeEvent = mapSessionEventToRuntimeEvent(sessionEvent, ctx, memory); + await this.onSessionEvent?.(sessionEvent, runtimeEvent); if (isTerminalRuntimeEvent(runtimeEvent)) { if (terminalEmitted) continue; terminalEmitted = true; yield runtimeEvent; - break; + if (!this.drainAfterTerminal) break; + continue; } yield runtimeEvent; } + } catch (error) { + await this.onError?.(error); + throw error; } finally { if (abortSignal && onAbort) { abortSignal.removeEventListener('abort', onAbort); } + await this.onFinally?.(); } } diff --git a/packages/runtime/src/session-manager.ts b/packages/runtime/src/session-manager.ts index 54713d05de..9eee51de19 100644 --- a/packages/runtime/src/session-manager.ts +++ b/packages/runtime/src/session-manager.ts @@ -50,12 +50,9 @@ import type { AgentRunStore } from '@maka/core'; import type { AgentBackend } from './ai-sdk-backend.js'; import type { RunTraceRecorder } from './run-trace.js'; -import { AgentRun, type AgentRunActiveSession, type AgentRunLineage } from './agent-run.js'; +import { AgentRun, type AgentRunActiveSession, type AgentRunBeginResult, type AgentRunLineage } from './agent-run.js'; import { classifyAgentRunRecovery, type AgentRunRecoveryDecision } from './agent-run-recovery.js'; -import { - createSessionEventMapMemory, - mapSessionEventToRuntimeEvent, -} from './ai-sdk-flow.js'; +import { AiSdkFlow } from './ai-sdk-flow.js'; import type { InvocationResult, InvocationSource, @@ -326,9 +323,10 @@ export class SessionManager { * * Runtime v2 bridge: * 1. Create one AgentRun, which remains the persistence/ledger owner. - * 2. Run that AgentRun through RuntimeRunner for canonical RuntimeEvents. - * 3. Forward the original SessionEvents to callers unchanged. - * 4. Drain the AgentRun stream fully so stop/abort cleanup semantics stay intact. + * 2. Begin it to resolve the production backend and model context. + * 3. Run AiSdkFlow through RuntimeRunner for canonical RuntimeEvents. + * 4. Forward the original SessionEvents to callers unchanged. + * 5. Drain the backend stream fully so stop/abort cleanup semantics stay intact. */ async *sendMessage( sessionId: string, @@ -357,43 +355,46 @@ export class SessionManager { const sessionEvents = new AsyncEventQueue(); const abortController = new AbortController(); - let agentRunIterator: AsyncIterator | undefined; let flowDone = false; + let begin: AgentRunBeginResult; + try { + begin = await run.begin(); + } catch (error) { + await run.recordFailure(error); + await run.finalize(); + throw error; + } + const aiSdkFlow = new AiSdkFlow({ + backend: begin.backend, + drainAfterTerminal: true, + onSessionEvent: async (sessionEvent) => { + await run.recordSessionEvent(sessionEvent); + await sessionEvents.push(sessionEvent); + }, + onError: async (error) => { + if (!isAsyncEventQueueClosed(error)) { + await run.recordFailure(error); + sessionEvents.fail(error); + } + }, + onFinally: async () => { + flowDone = true; + await run.finalize(); + sessionEvents.close(); + }, + }); const runner = new RuntimeRunner({ + flow: aiSdkFlow, providers: { newId: this.deps.newId, now: this.deps.now }, stopOnTerminal: false, - flow: { - run: async function* (ctx, request) { - const memory = createSessionEventMapMemory(); - agentRunIterator = run.execute()[Symbol.asyncIterator](); - try { - while (!request.abortSignal?.aborted) { - const next = await agentRunIterator.next(); - if (next.done) break; - if (request.abortSignal?.aborted) break; - await sessionEvents.push(next.value); - yield mapSessionEventToRuntimeEvent(next.value, ctx, memory); - } - } catch (error) { - if (!isAsyncEventQueueClosed(error)) { - sessionEvents.fail(error); - } - throw error; - } finally { - flowDone = true; - if (request.abortSignal?.aborted) { - await agentRunIterator.return?.().catch(() => undefined); - } - sessionEvents.close(); - } - }, - }, }); const runnerResult = runner.run({ sessionId, runId: run.runId, turnId: run.turnId, text: input.text, + ...(begin.backendInput.attachments ? { attachments: begin.backendInput.attachments } : {}), + context: begin.backendInput.context, source: this.deps.runtimeSource ?? 'desktop', lineage: run.lineage, abortSignal: abortController.signal, @@ -414,7 +415,6 @@ export class SessionManager { if (!flowDone) { abortController.abort(); sessionEvents.close(); - await agentRunIterator?.return?.().catch(() => undefined); } await runnerResult.catch(() => undefined); }