co-mono/packages/coding-agent/src/main.ts

/**
 * Main entry point for the coding agent
 */

import { Agent, type Attachment, ProviderTransport, type ThinkingLevel } from "@mariozechner/pi-agent-core";
import chalk from "chalk";
import { type Args, parseArgs, printHelp } from "./cli/args.js";
import { processFileArguments } from "./cli/file-processor.js";
import { selectSession } from "./cli/session-picker.js";
import { getModelsPath, VERSION } from "./config.js";
import { AgentSession } from "./core/agent-session.js";
import { exportFromFile } from "./core/export-html.js";
import { messageTransformer } from "./core/messages.js";
import { findModel, getApiKeyForModel, getAvailableModels } from "./core/model-config.js";
import { resolveModelScope, restoreModelFromSession, type ScopedModel } from "./core/model-resolver.js";
import { SessionManager } from "./core/session-manager.js";
import { SettingsManager } from "./core/settings-manager.js";
import { loadSlashCommands } from "./core/slash-commands.js";
import { buildSystemPrompt, loadProjectContextFiles } from "./core/system-prompt.js";
import { allTools, codingTools } from "./core/tools/index.js";
import { InteractiveMode, runPrintMode, runRpcMode } from "./modes/index.js";
import { initTheme } from "./modes/interactive/theme/theme.js";
import { getChangelogPath, getNewEntries, parseChangelog } from "./utils/changelog.js";
import { ensureTool } from "./utils/tools-manager.js";

/** Check npm registry for new version (non-blocking) */
async function checkForNewVersion(currentVersion: string): Promise<string | null> {
	try {
		const response = await fetch("https://registry.npmjs.org/@mariozechner/pi-coding-agent/latest");
		if (!response.ok) return null;

		const data = (await response.json()) as { version?: string };
		const latestVersion = data.version;

		if (latestVersion && latestVersion !== currentVersion) {
			return latestVersion;
		}

		return null;
	} catch {
		// Silently fail - don't disrupt the user experience
		return null;
	}
}

/** Run interactive mode with TUI */
async function runInteractiveMode(
	session: AgentSession,
	version: string,
	changelogMarkdown: string | null,
	modelFallbackMessage: string | null,
	versionCheckPromise: Promise<string | null>,
	initialMessages: string[],
	initialMessage?: string,
	initialAttachments?: Attachment[],
	fdPath: string | null = null,
): Promise<void> {
	const mode = new InteractiveMode(session, version, changelogMarkdown, fdPath);

	// Initialize TUI (subscribes to agent events internally)
	await mode.init();

	// Handle version check result when it completes (don't block)
	versionCheckPromise.then((newVersion) => {
		if (newVersion) {
			mode.showNewVersionNotification(newVersion);
		}
	});

	// Render any existing messages (from --continue mode)
	mode.renderInitialMessages(session.state);

	// Show model fallback warning at the end of the chat if applicable
	if (modelFallbackMessage) {
		mode.showWarning(modelFallbackMessage);
	}

	// Process initial message with attachments if provided (from @file args)
	if (initialMessage) {
		try {
			await session.prompt(initialMessage, { attachments: initialAttachments });
		} catch (error: unknown) {
			const errorMessage = error instanceof Error ? error.message : "Unknown error occurred";
			mode.showError(errorMessage);
		}
	}

	// Process remaining initial messages if provided (from CLI args)
	for (const message of initialMessages) {
		try {
			await session.prompt(message);
		} catch (error: unknown) {
			const errorMessage = error instanceof Error ? error.message : "Unknown error occurred";
			mode.showError(errorMessage);
		}
	}

	// Interactive loop
	while (true) {
		const userInput = await mode.getUserInput();

		// Process the message
		try {
			await session.prompt(userInput);
		} catch (error: unknown) {
			const errorMessage = error instanceof Error ? error.message : "Unknown error occurred";
			mode.showError(errorMessage);
		}
	}
}

/** Prepare initial message from @file arguments */
function prepareInitialMessage(parsed: Args): {
	initialMessage?: string;
	initialAttachments?: Attachment[];
} {
	if (parsed.fileArgs.length === 0) {
		return {};
	}

	const { textContent, imageAttachments } = processFileArguments(parsed.fileArgs);

	// Combine file content with first plain text message (if any)
	let initialMessage: string;
	if (parsed.messages.length > 0) {
		initialMessage = textContent + parsed.messages[0];
		parsed.messages.shift(); // Remove first message as it's been combined
	} else {
		initialMessage = textContent;
	}

	return {
		initialMessage,
		initialAttachments: imageAttachments.length > 0 ? imageAttachments : undefined,
	};
}

export async function main(args: string[]) {
	const parsed = parseArgs(args);

	if (parsed.help) {
		printHelp();
		return;
	}

	// Handle --export flag: convert session file to HTML and exit
	if (parsed.export) {
		try {
			const outputPath = parsed.messages.length > 0 ? parsed.messages[0] : undefined;
			const result = exportFromFile(parsed.export, outputPath);
			console.log(`Exported to: ${result}`);
			return;
		} catch (error: unknown) {
			const message = error instanceof Error ? error.message : "Failed to export session";
			console.error(chalk.red(`Error: ${message}`));
			process.exit(1);
		}
	}

	// Validate: RPC mode doesn't support @file arguments
	if (parsed.mode === "rpc" && parsed.fileArgs.length > 0) {
		console.error(chalk.red("Error: @file arguments are not supported in RPC mode"));
		process.exit(1);
	}

	// Process @file arguments
	const { initialMessage, initialAttachments } = prepareInitialMessage(parsed);

	// Initialize theme (before any TUI rendering)
	const settingsManager = new SettingsManager();
	const themeName = settingsManager.getTheme();
	initTheme(themeName);

	// Setup session manager
	const sessionManager = new SessionManager(parsed.continue && !parsed.resume, parsed.session);

	if (parsed.noSession) {
		sessionManager.disable();
	}

	// Handle --resume flag: show session selector
	if (parsed.resume) {
		const selectedSession = await selectSession(sessionManager);
		if (!selectedSession) {
			console.log(chalk.dim("No session selected"));
			return;
		}
		sessionManager.setSessionFile(selectedSession);
	}

	// Resolve model scope early if provided
	let scopedModels: ScopedModel[] = [];
	if (parsed.models && parsed.models.length > 0) {
		scopedModels = await resolveModelScope(parsed.models);
	}

	// Determine mode and output behavior
	const isInteractive = !parsed.print && parsed.mode === undefined;
	const mode = parsed.mode || "text";
	const shouldPrintMessages = isInteractive;

	// Find initial model
	let initialModel = await findInitialModelForSession(parsed, scopedModels, settingsManager);
	let initialThinking: ThinkingLevel = "off";

	// Get thinking level from scoped models if applicable
	if (scopedModels.length > 0 && !parsed.continue && !parsed.resume) {
		initialThinking = scopedModels[0].thinkingLevel;
	} else {
		// Try saved thinking level
		const savedThinking = settingsManager.getDefaultThinkingLevel();
		if (savedThinking) {
			initialThinking = savedThinking;
		}
	}

	// Non-interactive mode: fail early if no model available
	if (!isInteractive && !initialModel) {
		console.error(chalk.red("No models available."));
		console.error(chalk.yellow("\nSet an API key environment variable:"));
		console.error("  ANTHROPIC_API_KEY, OPENAI_API_KEY, GEMINI_API_KEY, etc.");
		console.error(chalk.yellow(`\nOr create ${getModelsPath()}`));
		process.exit(1);
	}

	// Non-interactive mode: validate API key exists
	if (!isInteractive && initialModel) {
		const apiKey = parsed.apiKey || (await getApiKeyForModel(initialModel));
		if (!apiKey) {
			console.error(chalk.red(`No API key found for ${initialModel.provider}`));
			process.exit(1);
		}
	}

	// Build system prompt
	const systemPrompt = buildSystemPrompt(parsed.systemPrompt, parsed.tools, parsed.appendSystemPrompt);

	// Handle session restoration
	let modelFallbackMessage: string | null = null;

	if (parsed.continue || parsed.resume || parsed.session) {
		const savedModel = sessionManager.loadModel();
		if (savedModel) {
			const result = await restoreModelFromSession(
				savedModel.provider,
				savedModel.modelId,
				initialModel,
				shouldPrintMessages,
			);

			if (result.model) {
				initialModel = result.model;
			}
			modelFallbackMessage = result.fallbackMessage;
		}

		// Load and restore thinking level
		const thinkingLevel = sessionManager.loadThinkingLevel() as ThinkingLevel;
		if (thinkingLevel) {
			initialThinking = thinkingLevel;
			if (shouldPrintMessages) {
				console.log(chalk.dim(`Restored thinking level: ${thinkingLevel}`));
			}
		}
	}

	// CLI --thinking flag takes highest priority
	if (parsed.thinking) {
		initialThinking = parsed.thinking;
	}

	// Determine which tools to use
	const selectedTools = parsed.tools ? parsed.tools.map((name) => allTools[name]) : codingTools;

	// Create agent
	const agent = new Agent({
		initialState: {
			systemPrompt,
			model: initialModel as any, // Can be null in interactive mode
			thinkingLevel: initialThinking,
			tools: selectedTools,
		},
		messageTransformer,
		queueMode: settingsManager.getQueueMode(),
		transport: new ProviderTransport({
			getApiKey: async () => {
				const currentModel = agent.state.model;
				if (!currentModel) {
					throw new Error("No model selected");
				}

				if (parsed.apiKey) {
					return parsed.apiKey;
				}

				const key = await getApiKeyForModel(currentModel);
				if (!key) {
					throw new Error(
						`No API key found for provider "${currentModel.provider}". Please set the appropriate environment variable or update ${getModelsPath()}`,
					);
				}
				return key;
			},
		}),
	});

	// If initial thinking was requested but model doesn't support it, reset to off
	if (initialThinking !== "off" && initialModel && !initialModel.reasoning) {
		agent.setThinkingLevel("off");
	}

	// Load previous messages if continuing, resuming, or using --session
	if (parsed.continue || parsed.resume || parsed.session) {
		const messages = sessionManager.loadMessages();
		if (messages.length > 0) {
			agent.replaceMessages(messages);
		}
	}

	// Log loaded context files
	if (shouldPrintMessages && !parsed.continue && !parsed.resume) {
		const contextFiles = loadProjectContextFiles();
		if (contextFiles.length > 0) {
			console.log(chalk.dim("Loaded project context from:"));
			for (const { path: filePath } of contextFiles) {
				console.log(chalk.dim(`  - ${filePath}`));
			}
		}
	}

	// Load file commands for slash command expansion
	const fileCommands = loadSlashCommands();

	// Create session
	const session = new AgentSession({
		agent,
		sessionManager,
		settingsManager,
		scopedModels,
		fileCommands,
	});

	// Route to appropriate mode
	if (mode === "rpc") {
		await runRpcMode(session);
	} else if (isInteractive) {
		// Check for new version in the background
		const versionCheckPromise = checkForNewVersion(VERSION).catch(() => null);

		// Check if we should show changelog
		const changelogMarkdown = getChangelogForDisplay(parsed, settingsManager);

		// Show model scope if provided
		if (scopedModels.length > 0) {
			const modelList = scopedModels
				.map((sm) => {
					const thinkingStr = sm.thinkingLevel !== "off" ? `:${sm.thinkingLevel}` : "";
					return `${sm.model.id}${thinkingStr}`;
				})
				.join(", ");
			console.log(chalk.dim(`Model scope: ${modelList} ${chalk.gray("(Ctrl+P to cycle)")}`));
		}

		// Ensure fd tool is available for file autocomplete
		const fdPath = await ensureTool("fd");

		await runInteractiveMode(
			session,
			VERSION,
			changelogMarkdown,
			modelFallbackMessage,
			versionCheckPromise,
			parsed.messages,
			initialMessage,
			initialAttachments,
			fdPath,
		);
	} else {
		// Non-interactive mode (--print flag or --mode flag)
		await runPrintMode(session, mode, parsed.messages, initialMessage, initialAttachments);
	}
}

/** Find initial model based on CLI args, scoped models, settings, or available models */
async function findInitialModelForSession(parsed: Args, scopedModels: ScopedModel[], settingsManager: SettingsManager) {
	// 1. CLI args take priority
	if (parsed.provider && parsed.model) {
		const { model, error } = findModel(parsed.provider, parsed.model);
		if (error) {
			console.error(chalk.red(error));
			process.exit(1);
		}
		if (!model) {
			console.error(chalk.red(`Model ${parsed.provider}/${parsed.model} not found`));
			process.exit(1);
		}
		return model;
	}

	// 2. Use first model from scoped models (skip if continuing/resuming)
	if (scopedModels.length > 0 && !parsed.continue && !parsed.resume) {
		return scopedModels[0].model;
	}

	// 3. Try saved default from settings
	const defaultProvider = settingsManager.getDefaultProvider();
	const defaultModelId = settingsManager.getDefaultModel();
	if (defaultProvider && defaultModelId) {
		const { model, error } = findModel(defaultProvider, defaultModelId);
		if (error) {
			console.error(chalk.red(error));
			process.exit(1);
		}
		if (model) {
			return model;
		}
	}

	// 4. Try first available model with valid API key
	const { models: availableModels, error } = await getAvailableModels();

	if (error) {
		console.error(chalk.red(error));
		process.exit(1);
	}

	if (availableModels.length > 0) {
		return availableModels[0];
	}

	return null;
}

/** Get changelog markdown to display (only for new sessions with updates) */
function getChangelogForDisplay(parsed: Args, settingsManager: SettingsManager): string | null {
	if (parsed.continue || parsed.resume) {
		return null;
	}

	const lastVersion = settingsManager.getLastChangelogVersion();
	const changelogPath = getChangelogPath();
	const entries = parseChangelog(changelogPath);

	if (!lastVersion) {
		// First run - show all entries
		if (entries.length > 0) {
			settingsManager.setLastChangelogVersion(VERSION);
			return entries.map((e) => e.content).join("\n\n");
		}
	} else {
		// Check for new entries since last version
		const newEntries = getNewEntries(entries, lastVersion);
		if (newEntries.length > 0) {
			settingsManager.setLastChangelogVersion(VERSION);
			return newEntries.map((e) => e.content).join("\n\n");
		}
	}

	return null;
}