Skip to content

Wire AiSdkFlow into SessionManager runtime path#9

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
likun666661 merged 1 commit into from
Jun 15, 2026
Merged

Wire AiSdkFlow into SessionManager runtime path #9

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions packages/runtime/src/__tests__/ai-sdk-flow.test.ts
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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: [
Expand Down
15 changes: 14 additions & 1 deletion packages/runtime/src/__tests__/session-manager.test.ts
View file
Open in desktop
Original file line numberDiff line numberDiff line change
@@ -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,
Expand DownExpand Up@@ -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();
Expand Down
185 changes: 107 additions & 78 deletions packages/runtime/src/agent-run.ts
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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;
Expand All@@ -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();
Expand DownExpand Up@@ -97,6 +107,21 @@ export class AgentRun {
}

async *execute(): AsyncIterable<SessionEvent> {
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<AgentRunBeginResult> {
await this.createRunRecord();

const userMsg: UserMessage = {
Expand All@@ -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<void> {
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<void> {
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<void> {
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<void> {
Expand Down
31 changes: 30 additions & 1 deletion packages/runtime/src/ai-sdk-flow.ts
View file
Open in desktop
Original file line numberDiff line numberDiff line change
Expand Up@@ -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> | void;
/** Called if the wrapped backend stream throws. */
onError?: (error: unknown) => Promise<void> | void;
/** Called after backend streaming finishes, errors, or is abandoned. */
onFinally?: () => Promise<void> | void;
/**
* Keep consuming backend events after the first terminal RuntimeEvent.
* Duplicate terminal RuntimeEvents are still coalesced from flow output.
*/
drainAfterTerminal?: boolean;
}

/**
Expand All@@ -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). */
Expand DownExpand Up@@ -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?.();
}
}

Expand Down
Loading