server.ts

/**
 * Hannah — Main entry point.
 *
 * Boots the full Pi interactive TUI with:
 *  - SQLite job tracking via an extension factory
 *  - Telegram (and future adapters) running in-process alongside the TUI
 *  - Heartbeat system for periodic check-ins
 *  - Cron scheduler for recurring jobs
 *  - Plugins that hook into the job lifecycle
 *
 * You interact with Hannah locally through the Pi TUI.
 * Adapters give you the same agent remotely (Telegram, etc.)
 */

import * as fs from "node:fs";
import * as path from "node:path";
import {
	createAgentSession,
	DefaultResourceLoader,
	InteractiveMode,
	type ExtensionAPI,
	type AgentSession,
} from "@mariozechner/pi-coding-agent";
// Tools (pure Pi tools — agent-callable, stateless)
import { registerMemory } from "./tools/memory.ts";
import { registerTdTool } from "./tools/td.ts";
import { registerObsidianTool } from "./tools/obsidian.ts";
import { registerFetchWebsiteTool } from "./tools/fetch-website.ts";
import { registerProjectInitTool } from "./tools/project-init.ts";
import { registerWorkonTool } from "./tools/workon.ts";

// Extensions (composite modules — web UI, DB, orchestration)
import { registerCronTool } from "./extensions/cron/cron.ts";
import { cronDbModule } from "./extensions/cron/cron-db.ts";
import { CronScheduler } from "./extensions/cron/scheduler.ts";
import { registerCalendarWeb } from "./extensions/calendar/calendar.ts";
import { calendarDbModule } from "./extensions/calendar/calendar-db.ts";
import { CalendarReminder } from "./extensions/calendar/calendar-reminder.ts";
import { registerProjectsWeb } from "./extensions/projects/projects.ts";
import { projectsDbModule } from "./extensions/projects/projects-db.ts";
import { registerVaultHealthWeb } from "./extensions/vault-health/vault-health.ts";
import { registerTdWeb } from "./extensions/td-web/td-web.ts";
import { registerSubagentTool } from "./extensions/subagent/tool.ts";
import { jobsDbModule, jobsApi } from "./extensions/jobs/jobs-db.ts";
import { Heartbeat } from "./extensions/heartbeat/heartbeat.ts";

// Core
import { openDb, type HannahDb, type DbModule } from "./db.ts";
import type { Adapter, Plugin, HannahServer, HannahConfig, RunJobOpts, JobResult, AssistantMessage, AgentErrorEvent, ToolResultContent } from "./types.ts";
import type { HannahConfigFile } from "./config.ts";
import { TelegramAdapter } from "./adapters/telegram.ts";
import { WebAdapter } from "./adapters/web.ts";
import { loadConfig } from "./config.ts";
import { initLogger, closeLogger } from "./logger.ts";

// ── Config ──────────────────────────────────────────────────────

const fileConfig = loadConfig(path.resolve(import.meta.dirname, ".."));

const config: HannahConfig = {
	provider: fileConfig.provider,
	modelId: fileConfig.model,
	thinkingLevel: fileConfig.thinkingLevel,
	cwd: fileConfig.cwd,
};

const DB_PATH = fileConfig.dbPath ?? path.resolve(import.meta.dirname, "..", "hannah.db");

// ── State ───────────────────────────────────────────────────────

let db: HannahDb;
const adapters: Adapter[] = [];
const plugins: Plugin[] = [];

// ── Shared job context ──────────────────────────────────────────
// The extension and adapter runJob() share this to coordinate
// which job tool calls get attributed to.

let activeJobId: string | null = null;
let extensionApi: ExtensionAPI | null = null;
let cronScheduler: CronScheduler | null = null;
let webAdapter: WebAdapter | null = null;
let hannahServer: HannahServer | null = null;
let piSession: AgentSession | null = null;

// ── Job queue ───────────────────────────────────────────────────
// Serializes job execution so only one prompt runs at a time.
// The Pi session is shared — concurrent prompts would interleave.

let jobQueue: Promise<void> = Promise.resolve();

// ── Job tracking extension factory ──────────────────────────────
// Injected into the Pi session via DefaultResourceLoader.
// Intercepts tool events to record tool calls for the active job.

function jobTrackingExtension(pi: ExtensionAPI): void {
	extensionApi = pi;
	let currentToolStart = 0;
	let tuiJobStartTime = 0;
	let tuiPrompt = "";
	let tuiTextParts: string[] = [];
	let tuiTurnCount = 0;
	let tuiToolCallCount = 0;
	let tuiInputTokens = 0;
	let tuiOutputTokens = 0;
	let tuiCacheReadTokens = 0;
	let tuiCacheWriteTokens = 0;
	let tuiCostInput = 0;
	let tuiCostOutput = 0;
	let tuiCostCacheRead = 0;
	let tuiCostCacheWrite = 0;

	// ── TUI job tracking ────────────────────────────────────────
	// When activeJobId is null, the prompt came from the local TUI.
	// Create a job so it shows up in the dashboard.

	pi.on("before_agent_start", async (event) => {
		if (activeJobId) return; // Adapter-driven job, already tracked

		tuiPrompt = typeof event.prompt === "string" ? event.prompt : "(prompt)";
		tuiJobStartTime = Date.now();
		tuiTextParts = [];
		tuiTurnCount = 0;
		tuiToolCallCount = 0;
		tuiInputTokens = 0;
		tuiOutputTokens = 0;
		tuiCacheReadTokens = 0;
		tuiCacheWriteTokens = 0;
		tuiCostInput = 0;
		tuiCostOutput = 0;
		tuiCostCacheRead = 0;
		tuiCostCacheWrite = 0;

		const model = piSession?.model;

		const jobId = jobsApi.createJob({
			channel: "tui",
			chatId: "local",
			prompt: tuiPrompt.slice(0, 10_000),
			model: model?.id ?? config.modelId,
			provider: model?.provider ?? config.provider,
			thinkingLevel: config.thinkingLevel,
		});

		jobsApi.markJobRunning(jobId);
		activeJobId = jobId;

		const job = jobsApi.getJob(jobId)!;
		for (const p of plugins) p.onJobCreated?.(job);
	});

	pi.on("agent_end", async (event) => {
		if (!activeJobId) return;

		// Check if this is a TUI job (channel = "tui")
		const job = jobsApi.getJob(activeJobId);
		if (!job || job.channel !== "tui") return;

		const durationMs = Date.now() - tuiJobStartTime;
		const totalTokens = tuiInputTokens + tuiOutputTokens;
		const costTotal = tuiCostInput + tuiCostOutput + tuiCostCacheRead + tuiCostCacheWrite;

		jobsApi.completeJob(activeJobId, {
			response: (tuiTextParts.join("") || "(no response)").slice(0, 50_000),
			inputTokens: tuiInputTokens,
			outputTokens: tuiOutputTokens,
			cacheReadTokens: tuiCacheReadTokens,
			cacheWriteTokens: tuiCacheWriteTokens,
			totalTokens,
			costInput: tuiCostInput,
			costOutput: tuiCostOutput,
			costCacheRead: tuiCostCacheRead,
			costCacheWrite: tuiCostCacheWrite,
			costTotal,
			toolCallCount: tuiToolCallCount,
			turnCount: tuiTurnCount,
			durationMs,
		});

		const finishedJob = jobsApi.getJob(activeJobId)!;
		for (const p of plugins) p.onJobDone?.(finishedJob);

		activeJobId = null;
	});

	pi.on("agent_error", async (event) => {
		if (!activeJobId) return;

		const job = jobsApi.getJob(activeJobId);
		if (!job || job.channel !== "tui") return;

		const durationMs = Date.now() - tuiJobStartTime;
		const errEvent = event as AgentErrorEvent;
		const errorMsg = errEvent.error?.message ?? "Agent error (unknown)";
		jobsApi.failJob(activeJobId, errorMsg.slice(0, 2000), durationMs);

		const failedJob = jobsApi.getJob(activeJobId)!;
		for (const p of plugins) p.onJobFailed?.(failedJob);

		activeJobId = null;
	});

	// ── Tool tracking (all channels) ────────────────────────────

	pi.on("tool_call", async () => {
		currentToolStart = Date.now();
		if (activeJobId) {
			const job = jobsApi.getJob(activeJobId);
			if (job?.channel === "tui") tuiToolCallCount++;
		}
	});

	pi.on("tool_result", async (event) => {
		if (!activeJobId) return;
		const duration = Date.now() - currentToolStart;

		const contentBlocks = (event.content ?? []) as ToolResultContent[];
		const resultText = contentBlocks
			.filter(c => c.type === "text")
			.map(c => c.text)
			.join("\n");

		jobsApi.recordToolCall({
			jobId: activeJobId,
			toolName: event.toolName,
			argsJson: JSON.stringify(event.input ?? {}).slice(0, 2000),
			resultPreview: resultText?.slice(0, 500) ?? undefined,
			isError: event.isError,
			durationMs: duration,
		});

		const toolCalls = jobsApi.getJobToolCalls(activeJobId);
		const lastToolCall = toolCalls[toolCalls.length - 1];
		if (lastToolCall) for (const p of plugins) p.onToolCall?.(lastToolCall);
	});

	// ── Turn tracking for TUI jobs ──────────────────────────────

	pi.on("turn_start", async () => {
		if (!activeJobId) return;
		const job = jobsApi.getJob(activeJobId);
		if (job?.channel === "tui") tuiTurnCount++;
	});

	pi.on("turn_end", async (event) => {
		if (!activeJobId) return;
		const job = jobsApi.getJob(activeJobId);
		if (job?.channel !== "tui") return;

		const msg = event.message as AssistantMessage | undefined;
		if (msg?.role === "assistant" && msg.usage) {
			tuiInputTokens += msg.usage.input ?? 0;
			tuiOutputTokens += msg.usage.output ?? 0;
			tuiCacheReadTokens += msg.usage.cacheRead ?? 0;
			tuiCacheWriteTokens += msg.usage.cacheWrite ?? 0;
			if (msg.usage.cost) {
				tuiCostInput += msg.usage.cost.input ?? 0;
				tuiCostOutput += msg.usage.cost.output ?? 0;
				tuiCostCacheRead += msg.usage.cost.cacheRead ?? 0;
				tuiCostCacheWrite += msg.usage.cost.cacheWrite ?? 0;
			}
		}

		// Collect text from assistant content
		if (msg?.role === "assistant" && msg.content) {
			for (const block of msg.content) {
				if (block.type === "text" && block.text) tuiTextParts.push(block.text);
			}
		}
	});

	// ── Tools (driven by config.tools) ──────────────────────────
	const toolRegistry: Record<string, () => void> = {
		memory: () => registerMemory(pi),
		cron: () => registerCronTool(pi, () => cronScheduler, () => hannahServer),
		td: () => registerTdTool(pi),
		obsidian: () => registerObsidianTool(pi, fileConfig.obsidian),
		fetch_website: () => registerFetchWebsiteTool(pi),
		project_init: () => registerProjectInitTool(pi),
		workon: () => registerWorkonTool(pi),
		subagent: () => registerSubagentTool(pi, () => hannahServer),
	};

	for (const name of fileConfig.tools) {
		const register = toolRegistry[name];
		if (register) register();
	}
}

// ── Server interface for adapters ───────────────────────────────

function buildServer(): HannahServer {
	const extensions = new Map<string, unknown>();

	return {
		getSession() {
			return piSession;
		},

		disposeSession() {
			// No-op for shared session — Pi manages session lifecycle
		},

		getCommands() {
			if (!extensionApi) return [];
			try {
				return extensionApi.getCommands().map(c => ({
					name: c.name,
					description: c.description,
					source: c.source,
				}));
			} catch {
				return [];
			}
		},

		async runJob(opts: RunJobOpts): Promise<JobResult> {
			// Serialize all jobs — the shared AgentSession can only
			// handle one prompt at a time. Jobs queue up and execute
			// in order, preventing the activeJobId race condition.
			const execute = async (): Promise<JobResult> => {
				if (!piSession) throw new Error("Agent session not ready — still booting");
				const session = piSession;
				const model = session.model;

				const jobId = jobsApi.createJob({
					channel: opts.channel,
					chatId: opts.chatId,
					prompt: opts.prompt.slice(0, 10_000),
					model: model?.id ?? config.modelId,
					provider: model?.provider ?? config.provider,
					thinkingLevel: config.thinkingLevel,
				});

				const job = jobsApi.getJob(jobId)!;
				for (const p of plugins) p.onJobCreated?.(job);

				const startTime = Date.now();
				jobsApi.markJobRunning(jobId);

				activeJobId = jobId;

				// Accumulators
				const textParts: string[] = [];
				let turnCount = 0, toolCallCount = 0;
				let inputTokens = 0, outputTokens = 0;
				let cacheReadTokens = 0, cacheWriteTokens = 0;
				let costInput = 0, costOutput = 0;
				let costCacheRead = 0, costCacheWrite = 0;

				try {
					const result = await new Promise<JobResult>((resolve, reject) => {
						const unsub = session.subscribe((event) => {
							opts.onEvent?.(event);

							switch (event.type) {
								case "message_update":
									if (event.assistantMessageEvent?.type === "text_delta") {
										textParts.push(event.assistantMessageEvent.delta);
									}
									break;
								case "turn_start":
									turnCount++;
									break;
								case "tool_execution_end":
									toolCallCount++;
									break;
								case "message_end": {
									const msg = event.message as AssistantMessage | undefined;
									if (msg?.role === "assistant" && msg.usage) {
										inputTokens += msg.usage.input ?? 0;
										outputTokens += msg.usage.output ?? 0;
										cacheReadTokens += msg.usage.cacheRead ?? 0;
										cacheWriteTokens += msg.usage.cacheWrite ?? 0;
										if (msg.usage.cost) {
											costInput += msg.usage.cost.input ?? 0;
											costOutput += msg.usage.cost.output ?? 0;
											costCacheRead += msg.usage.cost.cacheRead ?? 0;
											costCacheWrite += msg.usage.cost.cacheWrite ?? 0;
										}
									}
									break;
								}
								case "agent_end": {
									unsub();
									const durationMs = Date.now() - startTime;
									const totalTokens = inputTokens + outputTokens;
									const costTotal = costInput + costOutput + costCacheRead + costCacheWrite;
									resolve({
										jobId,
										response: textParts.join("") || "(no response)",
										inputTokens, outputTokens,
										cacheReadTokens, cacheWriteTokens,
										totalTokens,
										costInput, costOutput,
										costCacheRead, costCacheWrite,
										costTotal,
										toolCallCount, turnCount,
										durationMs,
									});
									break;
								}
							}
						});

						const promptOpts = opts.images ? { images: opts.images } : undefined;
						session.prompt(opts.prompt, promptOpts).catch((err) => {
							unsub();
							reject(err);
						});
					});

					jobsApi.completeJob(jobId, {
						response: result.response.slice(0, 50_000),
						inputTokens: result.inputTokens,
						outputTokens: result.outputTokens,
						cacheReadTokens: result.cacheReadTokens,
						cacheWriteTokens: result.cacheWriteTokens,
						totalTokens: result.totalTokens,
						costInput: result.costInput,
						costOutput: result.costOutput,
						costCacheRead: result.costCacheRead,
						costCacheWrite: result.costCacheWrite,
						costTotal: result.costTotal,
						toolCallCount: result.toolCallCount,
						turnCount: result.turnCount,
						durationMs: result.durationMs,
					});

					const finishedJob = jobsApi.getJob(jobId)!;
					for (const p of plugins) p.onJobDone?.(finishedJob);

					return result;
				} catch (error: any) {
					const durationMs = Date.now() - startTime;
					jobsApi.failJob(jobId, error.message?.slice(0, 2000) ?? "Unknown error", durationMs);
					const failedJob = jobsApi.getJob(jobId)!;
					for (const p of plugins) p.onJobFailed?.(failedJob);
					throw error;
				} finally {
					activeJobId = null;
				}
			};

			// Chain onto the queue — jobs run one at a time
			const queued = jobQueue.then(execute, execute);
			jobQueue = queued.catch(() => {}); // Prevent unhandled rejection on queue chain
			return queued;
		},

		addWebRoute(method, routePath, handler) {
			webAdapter?.addRoute(method, routePath, handler);
		},

		addDashboardSection(section) {
			webAdapter?.addDashboardSection(section);
		},

		registerExtension(name: string, extension: unknown) {
			extensions.set(name, extension);
		},

		getExtension<T = unknown>(name: string): T | undefined {
			return extensions.get(name) as T | undefined;
		},

		config,
	};
}

// ── Boot ────────────────────────────────────────────────────────

async function main(): Promise<void> {
	// 0. File logging
	initLogger(path.resolve(config.cwd, "log"), fileConfig.logging.level, fileConfig.logging.retainDays);

	// 1. Database
	// DB modules register their own migrations and prepared statements
	const dbModules: DbModule[] = [jobsDbModule, calendarDbModule, projectsDbModule, cronDbModule];
	db = openDb(DB_PATH, dbModules);

	// 2. Resource loader — inject our extension factory
	const resourceLoader = new DefaultResourceLoader({
		cwd: config.cwd,
		extensionFactories: [jobTrackingExtension],

	});
	await resourceLoader.reload();

	// 3. Create Pi agent session
	const { session, modelFallbackMessage } = await createAgentSession({
		cwd: config.cwd,
		resourceLoader,
	});

	piSession = session;

	// 3b. Subscribe to all session events for live streaming to plugins
	session.subscribe((event) => {
		for (const p of plugins) p.onStreamEvent?.(event);
	});

	// 4. Plugins
	const server = buildServer();
	hannahServer = server;
	for (const p of plugins) {
		await p.init?.(server);
	}

	// 5. Adapters (run in background alongside the TUI)
	if (fileConfig.web.enabled) {
		const web = new WebAdapter(fileConfig.web);
		webAdapter = web;
		adapters.push(web);
		plugins.push(web); // Also a plugin — receives job lifecycle events for SSE

		// Register web extensions BEFORE starting the server so all routes
		// are available when the first request arrives.
		registerTdWeb(extensionApi!, () => hannahServer);
		registerCalendarWeb(() => hannahServer);
		registerProjectsWeb(() => hannahServer);
		registerVaultHealthWeb(() => hannahServer, fileConfig);

		await web.start(server);
	}

	if (fileConfig.telegram.enabled && process.env.TELEGRAM_BOT_TOKEN) {
		const tg = new TelegramAdapter(fileConfig.telegram);
		adapters.push(tg);
		tg.start(server).catch((err) => {
			console.error("[telegram] Failed to start:", err.message);
		});
	}

	// 6. Cron scheduler — DB-backed, managed at runtime via the cron tool
	const cron = new CronScheduler(server, config.cwd);
	cronScheduler = cron;
	cron.start();

	// 7. Heartbeat
	let heartbeat: Heartbeat | undefined;
	if (fileConfig.heartbeat.enabled) {
		heartbeat = new Heartbeat(fileConfig.heartbeat, server, adapters, config.cwd);
		heartbeat.start();
	}

	// 7b. Calendar reminders
	const calendarReminder = new CalendarReminder(server, adapters);
	calendarReminder.start();

	// 8. Run the Pi interactive TUI — blocks until exit
	const mode = new InteractiveMode(session, {
		modelFallbackMessage,
	});

	await mode.run();

	// Cleanup after TUI exits
	heartbeat?.stop();
	calendarReminder.stop();
	cron.stop();
	for (const a of adapters) {
		try { await a.stop(); } catch {}
	}
	for (const p of plugins) {
		try { await p.stop?.(); } catch {}
	}
	db.close();
	closeLogger();
}

main().catch((err) => {
	console.error("Fatal:", err);
	process.exit(1);
});