@schlenkr/ragents 0.1.0

package/packages/agent/src/core/agent-session-runtime.ts

@@ -0,0 +1,548 @@
+import { copyFileSync, existsSync, mkdirSync } from "node:fs";
+import { basename, join, resolve } from "node:path";
+import { resolvePath } from "../utils/paths.ts";
+import type { AgentSession } from "./agent-session.ts";
+import type { AgentSessionRuntimeDiagnostic, AgentSessionServices } from "./agent-session-services.ts";
+import type {
+	ExtensionFactorySettlementObserver,
+	ProjectTrustContext,
+	ReplacedSessionContext,
+	SessionShutdownEvent,
+	SessionStartEvent,
+} from "./extensions/index.ts";
+import { beginSessionShutdownEvent } from "./extensions/runner.ts";
+import type { CreateAgentSessionResult } from "./sdk.ts";
+import { assertSessionCwdExists } from "./session-cwd.ts";
+import { SessionManager } from "./session-manager.ts";
+
+/**
+ * Result returned by runtime creation.
+ *
+ * The caller gets the created session, its cwd-bound services, and all
+ * diagnostics collected during setup.
+ */
+export interface CreateAgentSessionRuntimeResult extends CreateAgentSessionResult {
+	services: AgentSessionServices;
+	diagnostics: AgentSessionRuntimeDiagnostic[];
+}
+
+/**
+ * Creates a full runtime for a target cwd and session manager.
+ *
+ * The factory closes over process-global fixed inputs, recreates cwd-bound
+ * services for the effective cwd, resolves session options against those
+ * services, and finally creates the AgentSession.
+ */
+export type CreateAgentSessionRuntimeFactory = (options: {
+	cwd: string;
+	agentDir: string;
+	sessionManager: SessionManager;
+	signal: AbortSignal;
+	sessionStartEvent?: SessionStartEvent;
+	projectTrustContext?: ProjectTrustContext;
+	onExtensionFactorySettlement?: ExtensionFactorySettlementObserver;
+}) => Promise<CreateAgentSessionRuntimeResult>;
+
+export const DEFAULT_SESSION_SHUTDOWN_TIMEOUT_MS = 5_000;
+
+function abortReason(signal: AbortSignal): Error {
+	return signal.reason instanceof Error ? signal.reason : new Error("Agent session runtime creation was aborted.");
+}
+
+function awaitWithSignal<T>(operation: Promise<T>, signal: AbortSignal): Promise<T> {
+	if (signal.aborted) {
+		return Promise.reject(abortReason(signal));
+	}
+
+	return new Promise<T>((resolve, reject) => {
+		const abort = () => reject(abortReason(signal));
+		signal.addEventListener("abort", abort, { once: true });
+		operation.then(resolve, reject).finally(() => signal.removeEventListener("abort", abort));
+	});
+}
+
+function shutdownTimeout(value: number | undefined): number {
+	const timeout = value ?? DEFAULT_SESSION_SHUTDOWN_TIMEOUT_MS;
+	if (!Number.isFinite(timeout) || timeout <= 0) {
+		throw new RangeError("Session shutdown timeout must be a positive finite number.");
+	}
+	return timeout;
+}
+
+/**
+ * Thrown when /import references a JSONL file path that does not exist.
+ */
+export class SessionImportFileNotFoundError extends Error {
+	readonly filePath: string;
+
+	constructor(filePath: string) {
+		super(`File not found: ${filePath}`);
+		this.name = "SessionImportFileNotFoundError";
+		this.filePath = filePath;
+	}
+}
+
+function extractUserMessageText(content: string | Array<{ type: string; text?: string }>): string {
+	if (typeof content === "string") {
+		return content;
+	}
+
+	return content
+		.filter((part): part is { type: "text"; text: string } => part.type === "text" && typeof part.text === "string")
+		.map((part) => part.text)
+		.join("");
+}
+
+/**
+ * Owns the current AgentSession plus its cwd-bound services.
+ *
+ * Session replacement methods tear down the current runtime first, then create
+ * and apply the next runtime. If creation fails, the error is propagated to the
+ * caller. The caller is responsible for user-facing error handling.
+ */
+export class AgentSessionRuntime {
+	private rebindSession?: (session: AgentSession) => Promise<void>;
+	private beforeSessionInvalidate?: () => void;
+	private _session: AgentSession;
+	private _services: AgentSessionServices;
+	private readonly createRuntime: CreateAgentSessionRuntimeFactory;
+	private readonly sessionShutdownTimeoutMs: number;
+	private readonly replacementAbort = new AbortController();
+	private readonly settlements = new Set<Promise<void>>();
+	private replacementSignal?: AbortSignal;
+	private _diagnostics: AgentSessionRuntimeDiagnostic[];
+	private _modelFallbackMessage?: string;
+
+	constructor(
+		_session: AgentSession,
+		_services: AgentSessionServices,
+		createRuntime: CreateAgentSessionRuntimeFactory,
+		_diagnostics: AgentSessionRuntimeDiagnostic[] = [],
+		_modelFallbackMessage?: string,
+		sessionShutdownTimeoutMs = DEFAULT_SESSION_SHUTDOWN_TIMEOUT_MS,
+	) {
+		this._session = _session;
+		this._services = _services;
+		this.createRuntime = createRuntime;
+		this._diagnostics = _diagnostics;
+		this._modelFallbackMessage = _modelFallbackMessage;
+		this.sessionShutdownTimeoutMs = shutdownTimeout(sessionShutdownTimeoutMs);
+	}
+
+	get services(): AgentSessionServices {
+		return this._services;
+	}
+
+	get session(): AgentSession {
+		return this._session;
+	}
+
+	get cwd(): string {
+		return this._services.cwd;
+	}
+
+	get diagnostics(): readonly AgentSessionRuntimeDiagnostic[] {
+		return this._diagnostics;
+	}
+
+	get modelFallbackMessage(): string | undefined {
+		return this._modelFallbackMessage;
+	}
+
+	setRebindSession(rebindSession?: (session: AgentSession) => Promise<void>): void {
+		this.rebindSession = rebindSession;
+	}
+
+	setReplacementSignal(signal?: AbortSignal): void {
+		this.replacementSignal = signal;
+	}
+
+	waitForSettlement(): Promise<void> {
+		return Promise.allSettled([...this.settlements, this.session.waitForSettlement()]).then(() => undefined);
+	}
+
+	private trackSettlement(settlement: Promise<void>): void {
+		this.settlements.add(settlement);
+		void settlement.then(
+			() => this.settlements.delete(settlement),
+			() => this.settlements.delete(settlement),
+		);
+	}
+
+	/**
+	 * Set a synchronous callback that runs after `session_shutdown` handlers finish
+	 * but before the current session is invalidated.
+	 *
+	 * This is for host-owned UI teardown that must not yield to the event loop,
+	 * such as detaching extension-provided TUI components before the old extension
+	 * context becomes stale.
+	 */
+	setBeforeSessionInvalidate(beforeSessionInvalidate?: () => void): void {
+		this.beforeSessionInvalidate = beforeSessionInvalidate;
+	}
+
+	private async emitBeforeSwitch(
+		reason: "new" | "resume",
+		targetSessionFile?: string,
+	): Promise<{ cancelled: boolean }> {
+		const runner = this.session.extensionRunner;
+		if (!runner.hasHandlers("session_before_switch")) {
+			return { cancelled: false };
+		}
+
+		const result = await runner.emit({
+			type: "session_before_switch",
+			reason,
+			targetSessionFile,
+		});
+		return { cancelled: result?.cancel === true };
+	}
+
+	private async emitBeforeFork(
+		entryId: string,
+		options: { position: "before" | "at" },
+	): Promise<{ cancelled: boolean }> {
+		const runner = this.session.extensionRunner;
+		if (!runner.hasHandlers("session_before_fork")) {
+			return { cancelled: false };
+		}
+
+		const result = await runner.emit({
+			type: "session_before_fork",
+			entryId,
+			...options,
+		});
+		return { cancelled: result?.cancel === true };
+	}
+
+	private async teardownCurrent(reason: SessionShutdownEvent["reason"], targetSessionFile?: string): Promise<void> {
+		const session = this.session;
+		this.trackSettlement(session.waitForSettlement());
+
+		try {
+			const shutdown = beginSessionShutdownEvent(
+				session.extensionRunner,
+				{
+					type: "session_shutdown",
+					reason,
+					targetSessionFile,
+				},
+				{ timeoutMs: this.sessionShutdownTimeoutMs },
+			);
+			this.trackSettlement(shutdown.settled);
+			await shutdown.bounded;
+		} finally {
+			this.beforeSessionInvalidate?.();
+			session.dispose();
+		}
+	}
+
+	private apply(result: CreateAgentSessionRuntimeResult): void {
+		this._session = result.session;
+		this._services = result.services;
+		this._diagnostics = result.diagnostics;
+		this._modelFallbackMessage = result.modelFallbackMessage;
+	}
+
+	private async createReplacement(
+		options: Omit<Parameters<CreateAgentSessionRuntimeFactory>[0], "signal">,
+	): Promise<CreateAgentSessionRuntimeResult> {
+		const signal = this.replacementSignal
+			? AbortSignal.any([this.replacementSignal, this.replacementAbort.signal])
+			: this.replacementAbort.signal;
+		signal.throwIfAborted();
+		const creation = Promise.resolve().then(() =>
+			this.createRuntime({
+				...options,
+				signal,
+				onExtensionFactorySettlement: (settlement) => this.trackSettlement(settlement),
+			}),
+		);
+
+		try {
+			return await awaitWithSignal(creation, signal);
+		} catch (error) {
+			if (signal.aborted) {
+				void creation.then((lateResult) => lateResult.session.dispose(), () => undefined);
+			}
+			throw error;
+		}
+	}
+
+	private async finishSessionReplacement(withSession?: (ctx: ReplacedSessionContext) => Promise<void>): Promise<void> {
+		if (this.rebindSession) {
+			await this.rebindSession(this.session);
+		}
+		if (withSession) {
+			await withSession(this.session.createReplacedSessionContext());
+		}
+	}
+
+	async switchSession(
+		sessionPath: string,
+		options?: {
+			cwdOverride?: string;
+			withSession?: (ctx: ReplacedSessionContext) => Promise<void>;
+			projectTrustContextFactory?: (cwd: string) => ProjectTrustContext;
+		},
+	): Promise<{ cancelled: boolean }> {
+		const beforeResult = await this.emitBeforeSwitch("resume", sessionPath);
+		if (beforeResult.cancelled) {
+			return beforeResult;
+		}
+
+		const previousSessionFile = this.session.sessionFile;
+		const sessionManager = SessionManager.open(sessionPath, undefined, options?.cwdOverride);
+		assertSessionCwdExists(sessionManager, this.cwd);
+		await this.teardownCurrent("resume", sessionManager.getSessionFile());
+		this.apply(
+			await this.createReplacement({
+				cwd: sessionManager.getCwd(),
+				agentDir: this.services.agentDir,
+				sessionManager,
+				sessionStartEvent: { type: "session_start", reason: "resume", previousSessionFile },
+				projectTrustContext: options?.projectTrustContextFactory?.(sessionManager.getCwd()),
+			}),
+		);
+		await this.finishSessionReplacement(options?.withSession);
+		return { cancelled: false };
+	}
+
+	async newSession(options?: {
+		parentSession?: string;
+		setup?: (sessionManager: SessionManager) => Promise<void>;
+		withSession?: (ctx: ReplacedSessionContext) => Promise<void>;
+	}): Promise<{ cancelled: boolean }> {
+		const beforeResult = await this.emitBeforeSwitch("new");
+		if (beforeResult.cancelled) {
+			return beforeResult;
+		}
+
+		const previousSessionFile = this.session.sessionFile;
+		const sessionDir = this.session.sessionManager.getSessionDir();
+		const sessionManager = this.session.sessionManager.isPersisted()
+			? SessionManager.create(this.cwd, sessionDir)
+			: SessionManager.inMemory(this.cwd);
+		if (options?.parentSession) {
+			sessionManager.newSession({ parentSession: options.parentSession });
+		}
+
+		await this.teardownCurrent("new", sessionManager.getSessionFile());
+		this.apply(
+			await this.createReplacement({
+				cwd: this.cwd,
+				agentDir: this.services.agentDir,
+				sessionManager,
+				sessionStartEvent: { type: "session_start", reason: "new", previousSessionFile },
+			}),
+		);
+		if (options?.setup) {
+			await options.setup(this.session.sessionManager);
+			this.session.agent.state.messages = this.session.sessionManager.buildSessionContext().messages;
+		}
+		await this.finishSessionReplacement(options?.withSession);
+		return { cancelled: false };
+	}
+
+	async fork(
+		entryId: string,
+		options?: { position?: "before" | "at"; withSession?: (ctx: ReplacedSessionContext) => Promise<void> },
+	): Promise<{ cancelled: boolean; selectedText?: string }> {
+		const position = options?.position ?? "before";
+		const beforeResult = await this.emitBeforeFork(entryId, { position });
+		if (beforeResult.cancelled) {
+			return { cancelled: true };
+		}
+		let targetLeafId: string | null;
+		let selectedText: string | undefined;
+
+		const selectedEntry = this.session.sessionManager.getEntry(entryId);
+		if (!selectedEntry) {
+			throw new Error("Invalid entry ID for forking");
+		}
+
+		if (position === "at") {
+			targetLeafId = selectedEntry.id;
+		} else {
+			if (selectedEntry.type !== "message" || selectedEntry.message.role !== "user") {
+				throw new Error("Invalid entry ID for forking");
+			}
+			targetLeafId = selectedEntry.parentId;
+			selectedText = extractUserMessageText(selectedEntry.message.content);
+		}
+
+		const previousSessionFile = this.session.sessionFile;
+		if (this.session.sessionManager.isPersisted()) {
+			const currentSessionFile = this.session.sessionFile;
+			if (!currentSessionFile) {
+				throw new Error("Persisted session is missing a session file");
+			}
+			const sessionDir = this.session.sessionManager.getSessionDir();
+			if (!targetLeafId) {
+				const sessionManager = SessionManager.create(this.cwd, sessionDir);
+				sessionManager.newSession({ parentSession: currentSessionFile });
+				await this.teardownCurrent("fork", sessionManager.getSessionFile());
+				this.apply(
+					await this.createReplacement({
+						cwd: this.cwd,
+						agentDir: this.services.agentDir,
+						sessionManager,
+						sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile },
+					}),
+				);
+				await this.finishSessionReplacement(options?.withSession);
+				return { cancelled: false, selectedText };
+			}
+
+			if (!existsSync(currentSessionFile)) {
+				throw new Error(
+					"This session has not been saved yet. Wait for the first assistant response before cloning or forking it.",
+				);
+			}
+			const sessionManager = SessionManager.open(currentSessionFile, sessionDir);
+			const forkedSessionPath = sessionManager.createBranchedSession(targetLeafId);
+			if (!forkedSessionPath) {
+				throw new Error("Failed to create forked session");
+			}
+			await this.teardownCurrent("fork", sessionManager.getSessionFile());
+			this.apply(
+				await this.createReplacement({
+					cwd: sessionManager.getCwd(),
+					agentDir: this.services.agentDir,
+					sessionManager,
+					sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile },
+				}),
+			);
+			await this.finishSessionReplacement(options?.withSession);
+			return { cancelled: false, selectedText };
+		}
+
+		const sessionManager = this.session.sessionManager;
+		if (!targetLeafId) {
+			sessionManager.newSession({ parentSession: this.session.sessionFile });
+		} else {
+			sessionManager.createBranchedSession(targetLeafId);
+		}
+		await this.teardownCurrent("fork", sessionManager.getSessionFile());
+		this.apply(
+			await this.createReplacement({
+				cwd: this.cwd,
+				agentDir: this.services.agentDir,
+				sessionManager,
+				sessionStartEvent: { type: "session_start", reason: "fork", previousSessionFile },
+			}),
+		);
+		await this.finishSessionReplacement(options?.withSession);
+		return { cancelled: false, selectedText };
+	}
+
+	/**
+	 * Import a session JSONL file and switch runtime state to the imported session.
+	 *
+	 * @returns `{ cancelled: true }` when cancelled by `session_before_switch`, otherwise `{ cancelled: false }`.
+	 * @throws {SessionImportFileNotFoundError} When the input path does not exist.
+	 * @throws {MissingSessionCwdError} When the imported session cwd cannot be resolved and no override is provided.
+	 */
+	async importFromJsonl(inputPath: string, cwdOverride?: string): Promise<{ cancelled: boolean }> {
+		const resolvedPath = resolvePath(inputPath);
+		if (!existsSync(resolvedPath)) {
+			throw new SessionImportFileNotFoundError(resolvedPath);
+		}
+
+		const sessionDir = this.session.sessionManager.getSessionDir();
+		if (!existsSync(sessionDir)) {
+			mkdirSync(sessionDir, { recursive: true });
+		}
+
+		const destinationPath = join(sessionDir, basename(resolvedPath));
+		const beforeResult = await this.emitBeforeSwitch("resume", destinationPath);
+		if (beforeResult.cancelled) {
+			return beforeResult;
+		}
+
+		const previousSessionFile = this.session.sessionFile;
+		if (resolve(destinationPath) !== resolvedPath) {
+			copyFileSync(resolvedPath, destinationPath);
+		}
+
+		const sessionManager = SessionManager.open(destinationPath, sessionDir, cwdOverride);
+		assertSessionCwdExists(sessionManager, this.cwd);
+		await this.teardownCurrent("resume", sessionManager.getSessionFile());
+		this.apply(
+			await this.createReplacement({
+				cwd: sessionManager.getCwd(),
+				agentDir: this.services.agentDir,
+				sessionManager,
+				sessionStartEvent: { type: "session_start", reason: "resume", previousSessionFile },
+			}),
+		);
+		await this.finishSessionReplacement();
+		return { cancelled: false };
+	}
+
+	async dispose(): Promise<void> {
+		this.replacementAbort.abort(new Error("Agent session runtime was disposed."));
+		await this.teardownCurrent("quit");
+	}
+}
+
+/**
+ * Create the initial runtime from a runtime factory and initial session target.
+ *
+ * The same factory is stored on the returned AgentSessionRuntime and reused for
+ * later /new, /resume, /fork, and import flows.
+ */
+export async function createAgentSessionRuntime(
+	createRuntime: CreateAgentSessionRuntimeFactory,
+	options: {
+		cwd: string;
+		agentDir: string;
+		sessionManager: SessionManager;
+		signal?: AbortSignal;
+		sessionStartEvent?: SessionStartEvent;
+		sessionShutdownTimeoutMs?: number;
+		onExtensionFactorySettlement?: ExtensionFactorySettlementObserver;
+	},
+): Promise<AgentSessionRuntime> {
+	assertSessionCwdExists(options.sessionManager, options.cwd);
+	const sessionShutdownTimeoutMs = shutdownTimeout(options.sessionShutdownTimeoutMs);
+	const signal = options.signal ?? new AbortController().signal;
+	signal.throwIfAborted();
+	const creation = Promise.resolve().then(() => createRuntime({
+		cwd: options.cwd,
+		agentDir: options.agentDir,
+		sessionManager: options.sessionManager,
+		signal,
+		...(options.onExtensionFactorySettlement
+			? { onExtensionFactorySettlement: options.onExtensionFactorySettlement }
+			: {}),
+		...(options.sessionStartEvent ? { sessionStartEvent: options.sessionStartEvent } : {}),
+	}));
+	let result: CreateAgentSessionRuntimeResult;
+
+	try {
+		result = await awaitWithSignal(creation, signal);
+	} catch (error) {
+		if (signal.aborted) {
+			void creation.then((lateResult) => lateResult.session.dispose(), () => undefined);
+		}
+		throw error;
+	}
+	return new AgentSessionRuntime(
+		result.session,
+		result.services,
+		createRuntime,
+		result.diagnostics,
+		result.modelFallbackMessage,
+		sessionShutdownTimeoutMs,
+	);
+}
+
+export {
+	type AgentSessionRuntimeDiagnostic,
+	type AgentSessionServices,
+	type CreateAgentSessionFromServicesOptions,
+	type CreateAgentSessionServicesOptions,
+	createAgentSessionFromServices,
+	createAgentSessionServices,
+} from "./agent-session-services.ts";