@dex-ai/coding-agent-sdk 0.1.21

package/src/extensions/session.ts

@@ -0,0 +1,347 @@
+/**
+ * Session Extension — persists conversation history as JSONL.
+ */
+
+import type {
+	Extension,
+	AgentContext,
+	Message,
+	GenerateContext,
+	Content,
+} from "@dex-ai/sdk";
+import {
+	existsSync,
+	readFileSync,
+	appendFileSync,
+	mkdirSync,
+	readdirSync,
+	statSync,
+} from "node:fs";
+import { join } from "node:path";
+import { homedir } from "node:os";
+
+export interface SessionExtensionOptions {
+	sessionId?: string;
+	dir?: string;
+}
+
+/* ------------------------------------------------------------------ */
+/* Session listing for /resume                                         */
+/* ------------------------------------------------------------------ */
+
+export interface SessionInfo {
+	/** Session ID (filename without .jsonl). */
+	id: string;
+	/** Last modification time of the session file. */
+	lastModified: Date;
+	/** Last user message text (truncated). */
+	lastUserMessage: string;
+}
+
+/**
+ * Extract plain text from a message's content.
+ */
+function extractText(content: ReadonlyArray<Content> | undefined): string {
+	if (!content || content.length === 0) return "";
+	for (const block of content) {
+		if (block.type === "text") return block.text;
+	}
+	return "";
+}
+
+/**
+ * List all sessions in the given directory, returning metadata for each.
+ * Results are sorted chronologically (oldest first).
+ */
+export function listSessions(dir: string): SessionInfo[] {
+	if (!existsSync(dir)) return [];
+
+	const files = readdirSync(dir).filter((f) => f.endsWith(".jsonl"));
+	const sessions: SessionInfo[] = [];
+
+	for (const file of files) {
+		const filePath = join(dir, file);
+		const id = file.replace(/\.jsonl$/, "");
+
+		try {
+			const stat = statSync(filePath);
+			const text = readFileSync(filePath, "utf-8");
+			const lines = text.split("\n").filter((l) => l.trim());
+
+			// Skip empty files
+			if (lines.length === 0) continue;
+
+			// Find the last user message
+			let lastUserMessage = "";
+			for (let i = lines.length - 1; i >= 0; i--) {
+				try {
+					const msg = JSON.parse(lines[i]!) as Message;
+					if (msg.role === "user") {
+						lastUserMessage = extractText(msg.content);
+						break;
+					}
+				} catch {
+					/* skip malformed */
+				}
+			}
+
+			// Skip sessions without user messages
+			if (!lastUserMessage) continue;
+
+			// Truncate for display
+			const truncated =
+				lastUserMessage.length > 80
+					? lastUserMessage.slice(0, 77) + "..."
+					: lastUserMessage;
+
+			sessions.push({
+				id,
+				lastModified: stat.mtime,
+				lastUserMessage: truncated,
+			});
+		} catch {
+			/* skip unreadable files */
+		}
+	}
+
+	// Sort chronologically (oldest first)
+	sessions.sort((a, b) => a.lastModified.getTime() - b.lastModified.getTime());
+
+	return sessions;
+}
+
+/* ------------------------------------------------------------------ */
+/* Session repair                                                      */
+/* ------------------------------------------------------------------ */
+
+/**
+ * Repair orphaned tool calls — assistant messages with tool-call content
+ * that have no matching tool-result message following them. This can happen
+ * when a session is aborted mid-tool-execution.
+ */
+function repairOrphanedToolCalls(messages: Message[]): Message[] {
+	const repaired: Message[] = [];
+
+	for (let i = 0; i < messages.length; i++) {
+		const msg = messages[i]!;
+		repaired.push(msg);
+
+		// Only check assistant messages with tool-call content
+		if (msg.role !== "assistant") continue;
+		const toolCalls = msg.content.filter(
+			(c): c is Extract<Content, { type: "tool-call" }> =>
+				c.type === "tool-call",
+		);
+		if (toolCalls.length === 0) continue;
+
+		// Collect all tool-result IDs that follow this assistant message
+		// (they should be in immediately subsequent tool-role messages)
+		const fulfilledIds = new Set<string>();
+		let lastToolIdx = i;
+		for (let j = i + 1; j < messages.length; j++) {
+			const next = messages[j]!;
+			if (next.role !== "tool") break;
+			lastToolIdx = j;
+			for (const c of next.content) {
+				if (c.type === "tool-result") fulfilledIds.add(c.toolCallId);
+			}
+		}
+
+		// Skip past existing tool results (they'll be pushed by subsequent iterations)
+		// and only synthesize if there are unfulfilled calls
+		const unfulfilled = toolCalls.filter(
+			(tc) => !fulfilledIds.has(tc.toolCallId),
+		);
+		if (unfulfilled.length === 0) continue;
+
+		// Push all existing tool result messages first
+		for (let j = i + 1; j <= lastToolIdx; j++) {
+			repaired.push(messages[j]!);
+		}
+		// Skip those in the outer loop
+		i = lastToolIdx;
+
+		// Now synthesize results for orphaned tool calls
+		for (const tc of unfulfilled) {
+			repaired.push({
+				role: "tool",
+				content: [
+					{
+						type: "tool-result",
+						toolCallId: tc.toolCallId,
+						toolName: tc.toolName,
+						output: {
+							type: "error-text",
+							value: "Tool call aborted (session interrupted).",
+						},
+					},
+				],
+			});
+		}
+	}
+
+	return repaired;
+}
+
+/**
+ * Repair orphaned tool results — tool-result messages whose corresponding
+ * tool-call is missing from any preceding assistant message. This can happen
+ * when context compression removes an assistant message (e.g. one with
+ * multiple tool-calls) but the tool-results were persisted separately.
+ *
+ * Strategy: drop tool-result entries that have no matching tool-call.
+ * If a tool message becomes empty after filtering, remove it entirely.
+ */
+function repairOrphanedToolResults(messages: Message[]): Message[] {
+	// First pass: collect all tool-call IDs from assistant messages
+	const knownToolCallIds = new Set<string>();
+	for (const msg of messages) {
+		if (msg.role !== "assistant") continue;
+		for (const c of msg.content) {
+			if (c.type === "tool-call") {
+				knownToolCallIds.add(
+					(c as Extract<Content, { type: "tool-call" }>).toolCallId,
+				);
+			}
+		}
+	}
+
+	// Second pass: filter out tool-results with no matching tool-call
+	const repaired: Message[] = [];
+	for (const msg of messages) {
+		if (msg.role !== "tool") {
+			repaired.push(msg);
+			continue;
+		}
+
+		const validContent = msg.content.filter((c) => {
+			if (c.type !== "tool-result") return true;
+			return knownToolCallIds.has(
+				(c as Extract<Content, { type: "tool-result" }>).toolCallId,
+			);
+		});
+
+		// Only keep the tool message if it still has content
+		if (validContent.length > 0) {
+			repaired.push({ ...msg, content: validContent });
+		}
+	}
+
+	return repaired;
+}
+
+function generateId(): string {
+	return `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 8)}`;
+}
+
+export function sessionExtension(
+	opts: SessionExtensionOptions = {},
+): Extension {
+	const sessionId = opts.sessionId ?? generateId();
+	let sessionDir: string;
+	let sessionPath: string;
+	let watermark = 0;
+	/** Set of message IDs already persisted (prevents duplication after watermark drift). */
+	const persisted = new Set<string>();
+
+	function flushMessages(actx: AgentContext): void {
+		// Clamp watermark if context compression has shrunk the array
+		if (watermark > actx.messages.length) {
+			watermark = actx.messages.length;
+		}
+		while (watermark < actx.messages.length) {
+			const msg = actx.messages[watermark]!;
+			// Skip system messages — they are re-assembled on each create.
+			// Skip context-turn messages — they are ephemeral per-turn context
+			// (e.g. <env> blocks) regenerated fresh each generate cycle.
+			if (msg.role !== "system" && (msg as any).type !== "context-turn") {
+				try {
+					// Guard against double-writes after watermark drift from splice.
+					// Use message ID (not serialized JSON) so format differences
+					// (e.g. spaced vs compact JSON) don't bypass dedup.
+					const msgId = msg.id ?? "";
+					if (msgId && !persisted.has(msgId)) {
+						const line = JSON.stringify(msg);
+						appendFileSync(sessionPath, line + "\n");
+						persisted.add(msgId);
+					}
+				} catch (err) {
+					// Session persistence is best-effort but log for troubleshooting
+					console.error(
+						`[session] Failed to persist message: ${err instanceof Error ? err.message : String(err)}`,
+					);
+				}
+			}
+			watermark++;
+		}
+	}
+
+	return {
+		name: "session",
+
+		init(actx: AgentContext) {
+			sessionDir = opts.dir ?? join(homedir(), ".dex", "sessions");
+			sessionPath = join(sessionDir, `${sessionId}.jsonl`);
+
+			// Ensure sessions dir exists
+			if (!existsSync(sessionDir)) {
+				mkdirSync(sessionDir, { recursive: true });
+			}
+
+			// Resume: load existing messages from session file
+			if (existsSync(sessionPath)) {
+				try {
+					const lines = readFileSync(sessionPath, "utf-8")
+						.split("\n")
+						.filter((l) => l.trim());
+					let messages: Message[] = lines.map((l) => JSON.parse(l) as Message);
+					if (messages.length > 0) {
+						const beforeLen = messages.length;
+						messages = repairOrphanedToolCalls(messages);
+						messages = repairOrphanedToolResults(messages);
+						const repaired = messages.length !== beforeLen;
+						if (repaired) {
+							console.error(
+								`[session] Repaired orphaned tool-call/result(s) from interrupted session (${beforeLen} → ${messages.length} messages)`,
+							);
+						}
+						actx.appendMessage(messages);
+					}
+					// Seed persisted set with message IDs so we don't double-write on watermark drift
+					for (const msg of messages) {
+						if (msg.id) persisted.add(msg.id);
+					}
+				} catch (err) {
+					// Corrupted session — start fresh but log it
+					console.error(
+						`[session] Failed to load session ${sessionId}: ${err instanceof Error ? err.message : String(err)}`,
+					);
+				}
+			}
+
+			watermark = actx.messages.length;
+			actx.state.set("session", { id: sessionId, path: sessionPath });
+		},
+
+		on: {
+			/**
+			 * Flush new messages on model-stop — this fires BEFORE the context
+			 * extension's model-stop handler (which may splice/delete messages
+			 * from actx.messages for compression). By persisting first, we
+			 * ensure no messages are lost due to in-place mutations.
+			 */
+			"model-stop": async (gctx: GenerateContext): Promise<void> => {
+				flushMessages(gctx.agent);
+			},
+
+			/**
+			 * Also flush on generate-stop as a safety net — catches any
+			 * messages committed after the last model-stop (e.g. final tool
+			 * results) and handles watermark drift from context compression.
+			 */
+			"generate-stop": async (gctx: GenerateContext): Promise<void> => {
+				flushMessages(gctx.agent);
+			},
+		},
+	};
+}