locode/src/ui/ink/index.tsx

import { render } from "ink";
import { makeClient } from "../../backend/client.js";
import type { ToolCallMode } from "../../backend/capabilityProbe.js";
import { runHooksForEvent } from "../../hooks/runner.js";
import { connectConfiguredMcpServers, disconnectAllMcpServers } from "../../mcp/manager.js";
import { buildPluginAgentTool } from "../../plugins/agentTool.js";
import { getLoadedPlugins } from "../../plugins/registry.js";
import { buildSkillTool } from "../../plugins/skillTool.js";
import type { ToolDef } from "../../tools/types.js";
import { flushPendingSaves } from "../../persistence/sessionStore.js";
import { killAllBackgroundJobs } from "../../tools/backgroundJobs.js";
import { shutdownAll as shutdownAllLspServers } from "../../codeintel/lspManager.js";
import { configureLanguageSpecs } from "../../codeintel/lspManager.js";
import { resolveLspServers } from "../../config/config.js";
import { App } from "./App.js";

export interface RunInkAppOptions {
  baseURL: string;
  model?: string;
  suggestedModel?: string;
  cwd: string;
  toolModeOverride?: ToolCallMode;
  resumeSessionId?: string;
  interactiveResume?: boolean;
}

export async function runInkApp(opts: RunInkAppOptions): Promise<void> {
  if (!process.stdin.isTTY) {
    console.error(
      "locode needs an interactive terminal (it can't read piped or redirected input). Run it directly in your terminal.",
    );
    process.exit(1);
  }

  // Apply user-configured LSP server overrides/additions (locode config set lspServers) to the
  // built-in language→server mappings before any LSP tool is used. Cheap and idempotent.
  configureLanguageSpecs(resolveLspServers());

  // Kick off MCP server connections immediately so they overlap with model listing/picking —
  // by the time a session is actually created, this is usually already settled. Plugin-defined
  // agents (agents/*.md) become tools too, alongside MCP-provided ones, and every installed
  // plugin's skills (skills/*/SKILL.md) are exposed through one shared `skill` tool.
  const plugins = getLoadedPlugins();
  const pluginAgentTools = plugins.flatMap((plugin) => plugin.agents.map((agent) => buildPluginAgentTool(agent)));
  const skillTool = buildSkillTool(plugins.flatMap((plugin) => plugin.skills));
  const pluginTools = skillTool ? [...pluginAgentTools, skillTool] : pluginAgentTools;
  const extraToolsPromise: Promise<ToolDef[]> = connectConfiguredMcpServers(opts.cwd)
    .then((mcpTools) => [...mcpTools, ...pluginTools])
    .catch(() => pluginTools);

  // Populated once a session actually exists (see onSessionIdChange below) — lives outside the
  // Ink render tree so cleanup() can still read it after instance.unmount() has torn everything
  // in App.tsx down.
  let currentSessionId: string | null = null;
  // Start the UI immediately — no blocking network calls before rendering. Model listing happens
  // inside the App component so the user sees the UI right away. Deliberately stays on the
  // terminal's normal screen buffer (never the alternate screen) the whole time — finished
  // history prints as permanent scrollback (see App.tsx's use of Ink's <Static>), so the
  // terminal's own native mouse-wheel scroll and click-drag text selection/copy just work, with
  // no app-side mouse tracking or virtual-scroll machinery needed.
  const instance = render(
    <App
      baseURL={opts.baseURL}
      cwd={opts.cwd}
      toolModeOverride={opts.toolModeOverride}
      initialModel={opts.model}
      suggestedModel={opts.suggestedModel}
      resumeSessionId={opts.resumeSessionId}
      interactiveResume={opts.interactiveResume}
      extraToolsPromise={extraToolsPromise}
      onSessionIdChange={(id) => {
        currentSessionId = id;
      }}
    />,
  );

  // Ctrl+C is handled by Ink's exitOnCtrlC (default true): it calls exit(), which resolves
  // waitUntilExit below, so the normal cleanup chain already runs and MCP stdio children get a clean
  // shutdown. External SIGTERM/SIGHUP (e.g. `kill` from another terminal) bypasses that path, so
  // unmount explicitly (to restore the terminal) and run the same teardown before exiting.
  let cleanedUp = false;
  const cleanup = async () => {
    if (cleanedUp) return;
    cleanedUp = true;
    // Flush in-flight autosaves first so a fire-and-forget persist right before exit isn't lost.
    await flushPendingSaves();
    // Best-effort: don't leave backgrounded shells (dev servers, watch builds) running as orphans.
    await killAllBackgroundJobs();
    // Best-effort: shut down any lazily-spawned LSP servers (tsserver, pyright, gopls,
    // clangd, rust-analyzer) so they don't outlive locode as orphans.
    await shutdownAllLspServers().catch(() => {});
    await disconnectAllMcpServers();
    // Best-effort — fires on every exit path, so there isn't always a specific session id to attach
    // (e.g. exiting from the model picker before a session ever started).
    await runHooksForEvent("SessionEnd", { sessionId: "", cwd: opts.cwd }, {}).catch(() => {});
    // Printed after instance.unmount() has restored the terminal (both the Ctrl+C and SIGTERM/SIGHUP
    // paths reach here only once Ink has released the screen), so this lands as a normal trailing
    // console line rather than being swallowed by Ink's own rendering. Skipped if no session ever
    // started (e.g. the user exited from the model picker) — there's nothing to resume.
    if (currentSessionId) {
      console.log(`\nResume this conversation later with: locode --resume ${currentSessionId}`);
    }
  };
  const onTerminate = () => {
    instance.unmount();
    void cleanup().finally(() => process.exit(0));
  };
  process.on("SIGTERM", onTerminate);
  process.on("SIGHUP", onTerminate);

  await instance.waitUntilExit();
  await cleanup();
}