feat: add minimal run engine
This commit is contained in:
chungyeong
@@ -12,9 +12,7 @@
|
||||
"test": "cd ../.. && vitest run --project packages/session"
|
||||
},
|
||||
"dependencies": {
|
||||
"@devflow/core": "workspace:*"
|
||||
},
|
||||
"devDependencies": {
|
||||
"@devflow/core": "workspace:*",
|
||||
"@devflow/db": "workspace:*"
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
export * from "./adapter.js";
|
||||
export * from "./fake.js";
|
||||
export * from "./manager.js";
|
||||
export * from "./transcript.js";
|
||||
|
||||
410
packages/session/src/manager.ts
Normal file
410
packages/session/src/manager.ts
Normal file
@@ -0,0 +1,410 @@
|
||||
import { DevflowError, type PromptEnvelope } from "@devflow/core";
|
||||
import {
|
||||
type DbClient,
|
||||
RunEventRepository,
|
||||
approvalRequests,
|
||||
runs,
|
||||
tuiSessions,
|
||||
} from "@devflow/db";
|
||||
import { and, eq, inArray, ne, notInArray, sql } from "drizzle-orm";
|
||||
|
||||
import type {
|
||||
ProbeResult,
|
||||
SessionAdapter,
|
||||
SessionHandle,
|
||||
StartInput,
|
||||
TranscriptChunk,
|
||||
} from "./adapter.js";
|
||||
|
||||
type Database = DbClient["db"];
|
||||
|
||||
interface AdvisoryLockClient {
|
||||
query<T extends Record<string, unknown> = Record<string, unknown>>(
|
||||
text: string,
|
||||
values?: readonly unknown[],
|
||||
): Promise<{ rows: T[] }>;
|
||||
release(): void;
|
||||
}
|
||||
|
||||
export interface SessionRuntime {
|
||||
trackOperation<T>(operation: Promise<T>): Promise<T>;
|
||||
start(input: StartInput): Promise<SessionHandle>;
|
||||
sendPrompt(handle: SessionHandle, envelope: PromptEnvelope): Promise<{ promptId: string }>;
|
||||
probe(handle: SessionHandle): Promise<ProbeResult>;
|
||||
resume(handle: SessionHandle): Promise<SessionHandle>;
|
||||
rebootstrap(handle: SessionHandle): Promise<SessionHandle>;
|
||||
capture(handle: SessionHandle, fromSeq: bigint): AsyncIterable<TranscriptChunk>;
|
||||
dispose(handle: SessionHandle): Promise<void>;
|
||||
}
|
||||
|
||||
export interface SessionManagerOptions {
|
||||
adapter: SessionAdapter;
|
||||
db?: Database;
|
||||
dbClient?: DbClient;
|
||||
recoveryRunIds?: readonly string[];
|
||||
shutdownDrainMs?: number;
|
||||
}
|
||||
|
||||
export interface SessionManagerRecoveryResult {
|
||||
recoveredSessionIds: string[];
|
||||
failedSessionIds: string[];
|
||||
}
|
||||
|
||||
export class SessionManager implements SessionRuntime {
|
||||
private readonly adapter: SessionAdapter;
|
||||
private readonly db: Database | undefined;
|
||||
private readonly dbClient: DbClient | undefined;
|
||||
private readonly recoveryRunIds: readonly string[] | undefined;
|
||||
private readonly shutdownDrainMs: number;
|
||||
private readonly handles = new Map<string, SessionHandle>();
|
||||
private readonly inFlight = new Set<Promise<unknown>>();
|
||||
private lockClient: AdvisoryLockClient | undefined;
|
||||
private draining = false;
|
||||
|
||||
constructor(options: SessionManagerOptions) {
|
||||
this.adapter = options.adapter;
|
||||
this.db = options.dbClient?.db ?? options.db;
|
||||
this.dbClient = options.dbClient;
|
||||
this.recoveryRunIds = options.recoveryRunIds;
|
||||
this.shutdownDrainMs = options.shutdownDrainMs ?? 30_000;
|
||||
}
|
||||
|
||||
async initialize(): Promise<SessionManagerRecoveryResult> {
|
||||
await this.acquireLock();
|
||||
try {
|
||||
return await this.recoverSessions();
|
||||
} catch (error) {
|
||||
await this.shutdown();
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
|
||||
async acquireLock(): Promise<void> {
|
||||
if (this.dbClient === undefined) {
|
||||
throw new DevflowError("SessionManager requires a DbClient for singleton startup", {
|
||||
class: "fatal",
|
||||
code: "internal_state_corruption",
|
||||
});
|
||||
}
|
||||
if (this.lockClient !== undefined) {
|
||||
return;
|
||||
}
|
||||
|
||||
const client = (await this.dbClient.pool.connect()) as AdvisoryLockClient;
|
||||
const result = await client.query<{ acquired: boolean }>(
|
||||
"SELECT pg_try_advisory_lock(hashtext($1)) AS acquired",
|
||||
["devflow:session-manager"],
|
||||
);
|
||||
if (result.rows[0]?.acquired !== true) {
|
||||
client.release();
|
||||
throw new DevflowError("another session manager is running", {
|
||||
class: "human_required",
|
||||
code: "session_manager_already_running",
|
||||
recoveryHint: "exit_code=3",
|
||||
});
|
||||
}
|
||||
|
||||
this.lockClient = client;
|
||||
}
|
||||
|
||||
async shutdown(): Promise<void> {
|
||||
this.draining = true;
|
||||
await this.waitForInFlight();
|
||||
const client = this.lockClient;
|
||||
this.lockClient = undefined;
|
||||
this.handles.clear();
|
||||
if (client !== undefined) {
|
||||
try {
|
||||
await client.query("SELECT pg_advisory_unlock(hashtext($1))", ["devflow:session-manager"]);
|
||||
} finally {
|
||||
client.release();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
trackOperation<T>(operation: Promise<T>): Promise<T> {
|
||||
return this.track(operation);
|
||||
}
|
||||
|
||||
async start(input: StartInput): Promise<SessionHandle> {
|
||||
this.assertAcceptingPrompts();
|
||||
const handle = await this.track(this.adapter.start(input));
|
||||
this.handles.set(handle.sessionId, handle);
|
||||
return handle;
|
||||
}
|
||||
|
||||
async sendPrompt(handle: SessionHandle, envelope: PromptEnvelope): Promise<{ promptId: string }> {
|
||||
this.assertAcceptingPrompts();
|
||||
return this.track(this.adapter.sendPrompt(this.handleFor(handle), envelope));
|
||||
}
|
||||
|
||||
async probe(handle: SessionHandle): Promise<ProbeResult> {
|
||||
return this.track(this.adapter.probe(this.handleFor(handle)));
|
||||
}
|
||||
|
||||
async resume(handle: SessionHandle): Promise<SessionHandle> {
|
||||
this.assertAcceptingPrompts();
|
||||
const resumed = await this.track(this.adapter.resume(this.handleFor(handle)));
|
||||
this.handles.set(resumed.sessionId, resumed);
|
||||
return resumed;
|
||||
}
|
||||
|
||||
async rebootstrap(handle: SessionHandle): Promise<SessionHandle> {
|
||||
this.assertAcceptingPrompts();
|
||||
const rebootstrapped = await this.track(this.adapter.rebootstrap(this.handleFor(handle)));
|
||||
this.handles.set(rebootstrapped.sessionId, rebootstrapped);
|
||||
return rebootstrapped;
|
||||
}
|
||||
|
||||
async *capture(handle: SessionHandle, fromSeq: bigint): AsyncIterable<TranscriptChunk> {
|
||||
const finishTracking = this.beginTrackedOperation();
|
||||
try {
|
||||
for await (const chunk of this.adapter.capture(this.handleFor(handle), fromSeq)) {
|
||||
yield chunk;
|
||||
}
|
||||
} finally {
|
||||
finishTracking();
|
||||
}
|
||||
}
|
||||
|
||||
async dispose(handle: SessionHandle): Promise<void> {
|
||||
const resolvedHandle = this.handleFor(handle);
|
||||
await this.track(this.adapter.dispose(resolvedHandle));
|
||||
this.handles.delete(resolvedHandle.sessionId);
|
||||
}
|
||||
|
||||
async recoverSessions(): Promise<SessionManagerRecoveryResult> {
|
||||
if (this.db === undefined) {
|
||||
return { recoveredSessionIds: [], failedSessionIds: [] };
|
||||
}
|
||||
|
||||
const sessionRows = await this.db
|
||||
.select({
|
||||
id: tuiSessions.id,
|
||||
runId: tuiSessions.runId,
|
||||
roleId: tuiSessions.roleId,
|
||||
backend: tuiSessions.backend,
|
||||
cwd: tuiSessions.cwd,
|
||||
lastKnownPanePid: tuiSessions.lastKnownPanePid,
|
||||
recoveryAttempts: tuiSessions.recoveryAttempts,
|
||||
state: tuiSessions.state,
|
||||
tmuxSession: tuiSessions.tmuxSession,
|
||||
tmuxWindow: tuiSessions.tmuxWindow,
|
||||
})
|
||||
.from(tuiSessions)
|
||||
.innerJoin(runs, eq(tuiSessions.runId, runs.id))
|
||||
.where(
|
||||
this.recoveryRunIds === undefined
|
||||
? and(
|
||||
ne(tuiSessions.state, "FAILED_NEEDS_HUMAN"),
|
||||
notInArray(runs.state, [...terminalRunStates]),
|
||||
)
|
||||
: and(
|
||||
ne(tuiSessions.state, "FAILED_NEEDS_HUMAN"),
|
||||
notInArray(runs.state, [...terminalRunStates]),
|
||||
inArray(tuiSessions.runId, [...this.recoveryRunIds]),
|
||||
),
|
||||
);
|
||||
|
||||
const recoveredSessionIds: string[] = [];
|
||||
const failedSessionIds: string[] = [];
|
||||
for (const session of sessionRows) {
|
||||
const handle = compactHandle(
|
||||
session.id,
|
||||
session.lastKnownPanePid,
|
||||
session.tmuxSession,
|
||||
session.tmuxWindow,
|
||||
);
|
||||
try {
|
||||
const resumed = await this.resumeWithRetry(handle);
|
||||
this.handles.set(resumed.sessionId, resumed);
|
||||
recoveredSessionIds.push(resumed.sessionId);
|
||||
} catch (error) {
|
||||
await this.markRecoveryFailed(session, error);
|
||||
failedSessionIds.push(session.id);
|
||||
}
|
||||
}
|
||||
|
||||
return { recoveredSessionIds, failedSessionIds };
|
||||
}
|
||||
|
||||
private async markRecoveryFailed(
|
||||
session: {
|
||||
id: string;
|
||||
runId: string;
|
||||
roleId: string;
|
||||
backend: string;
|
||||
cwd: string;
|
||||
recoveryAttempts: number;
|
||||
},
|
||||
error: unknown,
|
||||
): Promise<void> {
|
||||
if (this.db === undefined) {
|
||||
return;
|
||||
}
|
||||
|
||||
const eventRepository = new RunEventRepository(this.db);
|
||||
const recoveryAttempts = session.recoveryAttempts + 1;
|
||||
const gateKey = "session_recovery_required";
|
||||
const approvalIdempotencyKey = `${session.runId}:${gateKey}:${session.id}:${recoveryAttempts}`;
|
||||
const pauseCause = `session_recovery_failed:${session.id}:${recoveryAttempts}`;
|
||||
await this.db.transaction(async (tx) => {
|
||||
await tx.execute(sql`SELECT 1 FROM ${runs} WHERE ${runs.id} = ${session.runId} FOR UPDATE`);
|
||||
const [run] = await tx
|
||||
.select({ state: runs.state })
|
||||
.from(runs)
|
||||
.where(eq(runs.id, session.runId))
|
||||
.limit(1);
|
||||
await tx
|
||||
.update(tuiSessions)
|
||||
.set({ state: "FAILED_NEEDS_HUMAN", recoveryAttempts })
|
||||
.where(eq(tuiSessions.id, session.id));
|
||||
await eventRepository.appendInTransaction(tx, {
|
||||
runId: session.runId,
|
||||
type: "session.failed",
|
||||
payload: { sessionId: session.id, roleId: session.roleId },
|
||||
idempotencyKey: `session.failed:${session.id}`,
|
||||
});
|
||||
if (run === undefined || isTerminalRunState(run.state)) {
|
||||
return;
|
||||
}
|
||||
const inserted = await tx
|
||||
.insert(approvalRequests)
|
||||
.values({
|
||||
runId: session.runId,
|
||||
gateKey,
|
||||
state: "pending",
|
||||
idempotencyKey: approvalIdempotencyKey,
|
||||
payload: {
|
||||
sessionId: session.id,
|
||||
roleId: session.roleId,
|
||||
backend: session.backend,
|
||||
cwd: session.cwd,
|
||||
recoveryHint: recoveryHintFor(error),
|
||||
},
|
||||
})
|
||||
.onConflictDoNothing({ target: approvalRequests.idempotencyKey })
|
||||
.returning({ id: approvalRequests.id, idempotencyKey: approvalRequests.idempotencyKey });
|
||||
if (run.state !== "paused") {
|
||||
await tx
|
||||
.update(runs)
|
||||
.set({ state: "paused", pausedFromState: run.state, updatedAt: new Date() })
|
||||
.where(eq(runs.id, session.runId));
|
||||
await eventRepository.appendInTransaction(tx, {
|
||||
runId: session.runId,
|
||||
type: "run.paused",
|
||||
payload: { cause: pauseCause, pausedFromState: run.state },
|
||||
idempotencyKey: `run.paused:${session.runId}:${pauseCause}`,
|
||||
});
|
||||
}
|
||||
const request =
|
||||
inserted[0] ??
|
||||
(
|
||||
await tx
|
||||
.select({ id: approvalRequests.id, idempotencyKey: approvalRequests.idempotencyKey })
|
||||
.from(approvalRequests)
|
||||
.where(eq(approvalRequests.idempotencyKey, approvalIdempotencyKey))
|
||||
.limit(1)
|
||||
)[0];
|
||||
if (request !== undefined) {
|
||||
await eventRepository.appendInTransaction(tx, {
|
||||
runId: session.runId,
|
||||
type: "approval.requested",
|
||||
payload: {
|
||||
approvalRequestId: request.id,
|
||||
approvalIdempotencyKey: request.idempotencyKey,
|
||||
gateKey,
|
||||
},
|
||||
idempotencyKey: `approval.requested:${request.idempotencyKey}`,
|
||||
});
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
private async resumeWithRetry(handle: SessionHandle): Promise<SessionHandle> {
|
||||
let lastError: unknown;
|
||||
for (let attempt = 0; attempt <= 2; attempt += 1) {
|
||||
try {
|
||||
return await this.track(this.adapter.resume(handle));
|
||||
} catch (error) {
|
||||
lastError = error;
|
||||
if (!(error instanceof DevflowError) || error.class !== "recoverable") {
|
||||
throw error;
|
||||
}
|
||||
}
|
||||
}
|
||||
throw lastError;
|
||||
}
|
||||
|
||||
private async track<T>(operation: Promise<T>): Promise<T> {
|
||||
const tracked = operation.finally(() => {
|
||||
this.inFlight.delete(tracked);
|
||||
});
|
||||
this.inFlight.add(tracked);
|
||||
return tracked;
|
||||
}
|
||||
|
||||
private beginTrackedOperation(): () => void {
|
||||
let finishOperation!: () => void;
|
||||
const tracked = new Promise<void>((resolve) => {
|
||||
finishOperation = resolve;
|
||||
}).finally(() => {
|
||||
this.inFlight.delete(tracked);
|
||||
});
|
||||
this.inFlight.add(tracked);
|
||||
return finishOperation;
|
||||
}
|
||||
|
||||
private async waitForInFlight(): Promise<void> {
|
||||
if (this.inFlight.size === 0) {
|
||||
return;
|
||||
}
|
||||
await Promise.race([
|
||||
Promise.allSettled([...this.inFlight]),
|
||||
new Promise((resolveWait) => setTimeout(resolveWait, this.shutdownDrainMs)),
|
||||
]);
|
||||
}
|
||||
|
||||
private handleFor(handle: SessionHandle): SessionHandle {
|
||||
return this.handles.get(handle.sessionId) ?? handle;
|
||||
}
|
||||
|
||||
private assertAcceptingPrompts(): void {
|
||||
if (this.draining) {
|
||||
throw new DevflowError("SessionManager is draining", {
|
||||
class: "human_required",
|
||||
code: "session_manager_draining",
|
||||
});
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
const terminalRunStates = ["completed", "failed", "aborted"] as const;
|
||||
|
||||
function isTerminalRunState(state: string): state is (typeof terminalRunStates)[number] {
|
||||
return terminalRunStates.includes(state as (typeof terminalRunStates)[number]);
|
||||
}
|
||||
|
||||
function compactHandle(
|
||||
sessionId: string,
|
||||
pid: number | null,
|
||||
tmuxSession: string | null,
|
||||
tmuxWindow: string | null,
|
||||
): SessionHandle {
|
||||
return {
|
||||
sessionId,
|
||||
...(pid === null ? {} : { pid }),
|
||||
...(tmuxSession === null ? {} : { tmuxSession }),
|
||||
...(tmuxWindow === null ? {} : { tmuxWindow }),
|
||||
};
|
||||
}
|
||||
|
||||
function recoveryHintFor(error: unknown): string {
|
||||
if (error instanceof DevflowError && error.recoveryHint !== undefined) {
|
||||
return error.recoveryHint;
|
||||
}
|
||||
if (error instanceof Error) {
|
||||
return error.message;
|
||||
}
|
||||
return "session resume failed";
|
||||
}
|
||||
@@ -6,5 +6,5 @@
|
||||
"types": ["node", "vitest"]
|
||||
},
|
||||
"include": ["src/**/*.ts"],
|
||||
"references": [{ "path": "../core" }]
|
||||
"references": [{ "path": "../core" }, { "path": "../db" }]
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user