agent.mjs

#!/usr/bin/env node
// Controllable seam: the hermetic WS e2e substitutes a scripted SDK module
// (tests/fixtures/fake-agent-sdk.mjs) so production stream handling is
// exercised without a live model.
const { query, tool, createSdkMcpServer } = await import(
  process.env.OMI_AGENT_SDK_MODULE || "@anthropic-ai/claude-agent-sdk"
);
import { ALLOWED_TOOLS, buildAgentDefinitions } from "./query-config.mjs";
import { USER_MESSAGES, classifyError, isExpectedAbort, logEvent, markOwnedAbort, withRetry } from "./errors.mjs";
import { createSubagentRouter } from "./stream-routing.mjs";
import { buildCondensedSeed, planCondensation } from "./conversation-condenser.mjs";
import {
  activityCounts,
  runSqlBatch,
  executeReadOnlyQuery,
  truncateToolResult,
  appUsageMatrix,
  emptyRangeStatement,
  formatAppUsage,
  formatHourlyTimeline,
  formatTopWindows,
  hourlyTimeline,
  resolveRange,
  topWindows,
} from "./data-tools.mjs";
import Database from "better-sqlite3";
import { z } from "zod";
import { dirname, join } from "path";
import { fileURLToPath } from "url";
import { createServer } from "http";
import { WebSocketServer } from "ws";
import { existsSync, mkdirSync, createWriteStream, statSync, renameSync, unlinkSync, writeFileSync } from "fs";
import { createInflateRaw, createGunzip } from "zlib";
import { homedir } from "os";

// --- Global error handlers: prevent SDK abort errors from crashing the process ---
// Owned-abort policy: only aborts within the grace window after OUR OWN
// interrupt()/abort() calls are expected; any other abort-shaped error is an
// unknown error and must surface at full detail (previously every message
// containing "aborted" was silently swallowed forever).
process.on('unhandledRejection', (err) => {
  if (isExpectedAbort(err)) {
    logEvent('debug', 'abort_suppressed', { error: err });
    return;
  }
  logEvent('error', 'unhandled_rejection', { category: classifyError(err).category, error: err });
});

process.on('uncaughtException', (err) => {
  if (isExpectedAbort(err)) {
    logEvent('debug', 'abort_suppressed', { error: err });
    return;
  }
  // Maximum-detail post-mortem, synchronously flushed — console.log can lose
  // the line across process.exit.
  const record = {
    ts: new Date().toISOString(), level: 'error', event: 'uncaught_exception',
    category: classifyError(err).category,
    error: { name: err?.name, message: err?.message, stack: err?.stack },
    memory: process.memoryUsage(), uptime_s: Math.round(process.uptime()),
  };
  try { writeFileSync(1, JSON.stringify(record) + "\n"); } catch { console.error(record); }
  process.exit(1); // systemd restarts us; its restart backoff is the loop guard
});

// --- Async message queue for persistent session streaming input ---
class AsyncMessageQueue {
  constructor() {
    this._queue = [];
    this._resolve = null;
    this._done = false;
  }
  push(msg) {
    if (this._resolve) {
      this._resolve({ value: msg, done: false });
      this._resolve = null;
    } else {
      this._queue.push(msg);
    }
  }
  end() {
    this._done = true;
    if (this._resolve) {
      this._resolve({ value: undefined, done: true });
      this._resolve = null;
    }
  }
  [Symbol.asyncIterator]() { return this; }
  next() {
    if (this._queue.length > 0)
      return Promise.resolve({ value: this._queue.shift(), done: false });
    if (this._done)
      return Promise.resolve({ value: undefined, done: true });
    return new Promise(r => { this._resolve = r; });
  }
}

// --- Configuration ---
const DB_PATH = process.env.DB_PATH || join(homedir(), "omi-agent/data/omi.db");
const GCS_BASE = "https://storage.googleapis.com/based-hardware-agent";
const GEMINI_API_KEY = process.env.GEMINI_API_KEY;
const AUTH_TOKEN = process.env.AUTH_TOKEN;
const BACKEND_URL = process.env.BACKEND_URL || "https://api.omi.me";
const PORT = parseInt(process.env.PORT || "8080", 10);
const EMBEDDING_DIM = 3072;
// Max upload size: 10GB
const MAX_UPLOAD_BYTES = 10 * 1024 * 1024 * 1024;

// --- Idle auto-stop ---
const IDLE_TIMEOUT_MS = 30 * 60 * 1000;   // 30 minutes
const IDLE_CHECK_INTERVAL_MS = 5 * 60 * 1000; // check every 5 minutes
let lastActivityAt = Date.now();

// Tables allowed for incremental sync from desktop
const SYNC_TABLES = new Set([
  "screenshots", "action_items", "transcription_sessions",
  "transcription_segments", "memories", "staged_tasks",
  "focus_sessions", "observations", "live_notes",
  "ai_user_profiles", "task_dedup_log",
]);


const __dirname = dirname(fileURLToPath(import.meta.url));
const playwrightCli = join(__dirname, "node_modules", "@playwright", "mcp", "cli.js");

// --- Firebase token (passed from desktop app for backend API calls) ---
let userFirebaseToken = null;
let backendTools = [];

// --- Database Setup (lazy — opened on first use or after upload) ---
let db = null;
let defaultSystemPrompt = null;
let agentDefinitions = null;
let omiServer = null;

function openDatabase() {
  if (db) {
    try { db.close(); } catch {}
    db = null;
  }
  if (!existsSync(DB_PATH)) {
    return false;
  }
  db = Database(DB_PATH);  // writable for /sync inserts; agent tool still blocks non-SELECT
  db.pragma("journal_mode = WAL");

  // No extra indexes needed: E6 scale bench (600K rows, 2026-07-22) showed the
  // covering idx_screenshots_timestamp serves every range-comparison query the
  // tools issue (worst tool call 126ms p50). The old date(timestamp,'localtime')
  // index attempt failed on every open (non-deterministic date()) and date()-
  // wrapped WHEREs full-scan anyway — range comparisons are the supported shape.

  // Rebuild schema + system prompt + MCP server
  const schema = getSchema();
  defaultSystemPrompt = `You are an AI assistant with access to the user's OMI desktop database and their connected services.
This database contains their screen history (screenshots with OCR text), tasks, transcriptions, memories, and focus sessions.

DATABASE SCHEMA:
${schema}

TOOLS:
- **execute_sql**: Run SQL queries on the database. SELECT auto-limits to 200 rows. Supports FTS5 MATCH for keyword search. Use for structured queries (app usage, time ranges, task management, aggregations).
- **semantic_search**: Vector similarity search on screenshot OCR text. Use for fuzzy/conceptual queries where exact keywords won't work.
- **get_daily_recap**: Pre-formatted activity recap (apps, conversations, tasks, hourly timeline) for any date or range (start_date/end_date or days_ago). Use for "what did I do <day/range>" — single tool call, much faster than multiple SQL queries. If it reports no activity, that is authoritative — do not re-verify with SQL.
- **get_app_usage**: Day-by-app screen activity matrix for any range in one call. Use for "what did I spend time on", day comparisons, and usage patterns instead of per-day GROUP BY queries.
- **Playwright browser tools**: You can navigate websites, click elements, fill forms, take screenshots, etc. Use when the user asks you to do something on the web.
- **Backend tools** (calendar, gmail, health, conversations, memories, action items, web search, etc.): Use these when the user asks about their calendar events, emails, health data, past conversations, or wants to search the web. These tools connect to the user's real accounts.

GUIDELINES:
- For "what did I do today/yesterday/this week" queries, use get_daily_recap — it's a single tool call that returns a formatted summary. Much faster than multiple execute_sql calls.
- Only use get_conversations_tool when the user asks about specific conversation transcripts or content details. For activity summaries, get_daily_recap is faster.
- For time-filtered queries on screenshots, prefer range comparisons: WHERE timestamp >= datetime('now', 'start of day', '-1 day', 'localtime') AND timestamp < datetime('now', 'start of day', 'localtime'). Avoid wrapping the column in date() or strftime() in WHERE clauses — it's slower on large tables.
- Key tables: screenshots (timestamp, appName, windowTitle, ocrText — 600K+ rows, always filter by timestamp), action_items (description, completed, priority, category, dueAt), memories (content, category, source), transcription_sessions (title, overview, startedAt, finishedAt), transcription_segments (sessionId, speaker, text), focus_sessions (status, appOrSite, durationSeconds), observations (appName, contextSummary, currentActivity), goals (title, goalType, targetValue, currentValue), staged_tasks (description, priority), indexed_files (path, filename, fileType, folder), live_notes (sessionId, text), ai_user_profiles (profileText, generatedAt)
- FTS tables for keyword search: screenshots_fts(ocrText, windowTitle, appName), action_items_fts(description), staged_tasks_fts(description), task_chat_messages_fts(messageText)
- For task queries, use action_items table (or FTS: action_items_fts MATCH 'keyword')
- Task extraction creates near-duplicate entries: for "how many tasks" report COUNT(DISTINCT description), and GROUP BY description when listing so duplicates collapse
- For conversation queries, use transcription_sessions + transcription_segments
- For personal facts/preferences, query the memories table first
- For calendar, email, health data — use the backend tools (get_calendar_events_tool, get_gmail_messages_tool, etc.)
- For ROW-HEAVY work — searching through many screenshots/tasks/transcripts (content recall, task triage, multi-day pattern analysis) — delegate to the researcher subagent via the Task tool; it reads the rows in its own context and returns distilled findings. Answer directly when one or two aggregate queries (or a single get_daily_recap call) suffice.
- CONNECTING INTEGRATIONS: if a connector/tool reports it is "not connected", or the user asks how to connect calendar/gmail/etc., call get_connect_link with that capability (e.g. "gmail") and give the user the exact link it returns, verbatim. NEVER invent setup steps, mention external MCP servers/skills, or describe manual settings flows — only use get_connect_link.
- Be concise and helpful. Format results clearly.`;

  agentDefinitions = buildAgentDefinitions(schema);

  rebuildMcpServer();

  console.log(`[db] Database opened: ${DB_PATH}`);
  return true;
}

function isDatabaseReady() {
  return db !== null;
}

function getSchema() {
  const tables = db
    .prepare(
      "SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%' AND name NOT LIKE 'grdb_%' ORDER BY name"
    )
    .all();

  let schema = "";
  for (const { name } of tables) {
    const cols = db.prepare(`PRAGMA table_info('${name}')`).all();
    const colDefs = cols.map((c) => `  ${c.name} ${c.type}`).join("\n");
    schema += `\n${name}:\n${colDefs}\n`;
    const count = db.prepare(`SELECT COUNT(*) as n FROM "${name}"`).get();
    schema += `  (${count.n} rows)\n`;
  }
  return schema;
}

// --- SQL Execution Logic ---
const BLOCKED_KEYWORDS = ["DROP", "ALTER", "CREATE", "PRAGMA", "ATTACH", "DETACH", "VACUUM"];

function executeSqlQuery(sqlQuery) {
  if (!db) return JSON.stringify({ error: "Database not loaded. Upload omi.db first." });
  // stmt.readonly-based guard (data-tools): allows CTE reads the old prefix
  // guard rejected, caps rows by iteration, no keyword false-positives.
  return executeReadOnlyQuery(db, sqlQuery);
}

// --- Semantic Search Logic ---
async function embedQueryText(text) {
  const url = `https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:embedContent?key=${GEMINI_API_KEY}`;
  const body = {
    model: "models/gemini-embedding-001",
    content: { parts: [{ text }] },
    taskType: "RETRIEVAL_QUERY",
  };
  const resp = await fetch(url, {
    method: "POST",
    headers: { "Content-Type": "application/json" },
    body: JSON.stringify(body),
  });
  const json = await resp.json();
  if (!json.embedding?.values) {
    throw new Error(`Embedding API error: ${JSON.stringify(json.error || json)}`);
  }
  const raw = json.embedding.values.map(Number);
  let norm = 0;
  for (const v of raw) norm += v * v;
  norm = Math.sqrt(norm);
  return norm > 0 ? raw.map((v) => v / norm) : raw;
}

function readEmbeddingFromBlob(buffer) {
  if (buffer.byteLength !== EMBEDDING_DIM * 4) return null;
  return new Float32Array(buffer.buffer, buffer.byteOffset, EMBEDDING_DIM);
}

async function performSemanticSearch(searchQuery, days = 7, appFilter = null) {
  if (!db) return JSON.stringify({ error: "Database not loaded. Upload omi.db first." });
  if (!GEMINI_API_KEY) {
    return JSON.stringify({ error: "GEMINI_API_KEY not set" });
  }
  const queryEmbedding = await embedQueryText(searchQuery);
  const startDate = new Date();
  startDate.setDate(startDate.getDate() - days);
  const startStr = startDate.toISOString().replace("T", " ").slice(0, 19);

  let sql = `SELECT id, timestamp, appName, windowTitle, substr(ocrText, 1, 300) as ocrPreview, embedding
             FROM screenshots WHERE embedding IS NOT NULL AND timestamp >= ?`;
  const params = [startStr];
  if (appFilter) {
    sql += " AND appName = ?";
    params.push(appFilter);
  }
  sql += " ORDER BY timestamp DESC";

  const rows = db.prepare(sql).all(...params);
  const results = [];
  for (const row of rows) {
    const stored = readEmbeddingFromBlob(row.embedding);
    if (!stored) continue;
    let dot = 0;
    for (let i = 0; i < queryEmbedding.length; i++) dot += queryEmbedding[i] * stored[i];
    if (dot > 0.3) {
      results.push({
        screenshotId: row.id,
        similarity: Math.round(dot * 1000) / 1000,
        timestamp: row.timestamp,
        appName: row.appName,
        windowTitle: row.windowTitle,
        ocrPreview: row.ocrPreview,
      });
    }
  }
  results.sort((a, b) => b.similarity - a.similarity);
  return JSON.stringify({
    query: searchQuery,
    days,
    totalScanned: rows.length,
    matchesAboveThreshold: results.length,
    results: results.slice(0, 15),
  });
}

// --- Define MCP Tools using Agent SDK ---

const executeSqlTool = tool(
  "execute_sql",
  `Run SQL on the user's local omi.db SQLite database for structured data queries.

Use when:
- User asks for app usage stats, screen time, or activity counts
- Time-based queries like "how long did I spend on X?"
- Task management: looking up action items, checking completion status
- Aggregations, rankings, or structured filters on local data

Don't use when (if those tools are available):
- User asks about conversation content or transcripts (prefer get_conversations or search_conversations)
- User asks about their preferences or facts about themselves (prefer get_memories)
- User asks fuzzy/conceptual questions (use semantic_search instead)
- If backend tools are not available, fall back to execute_sql on the local transcription_sessions table

Note: Database is read-only (SELECT only). SELECT queries auto-limit to 200 rows.
Supports FTS5 MATCH queries for keyword search (e.g., WHERE screenshots_fts MATCH 'keyword').

BATCHING: when you need several independent queries (comparisons, multiple
tables, multiple time ranges), pass them ALL in the "queries" array in ONE
call — results come back labeled per query. Never issue sequential
execute_sql calls for queries that don't depend on each other's results.

Key tables: screenshots (appName, windowTitle, ocrText, timestamp), transcription_sessions (title, overview, startedAt, finishedAt), transcription_segments (sessionId, speaker, text, startTime), action_items (description, completed, priority, dueAt, category), memories (content, category, source), staged_tasks (description, priority, source), focus_sessions (status, appOrSite, durationSeconds), observations (appName, contextSummary, currentActivity), goals (title, goalType, targetValue, currentValue), indexed_files (path, filename, fileType, folder), live_notes (sessionId, text, timestamp), ai_user_profiles (profileText, generatedAt).`,
  {
    query: z.string().optional().describe("Single SQL query to execute against omi.db"),
    queries: z
      .array(z.string())
      .min(1)
      .max(8)
      .optional()
      .describe("Batch of independent SQL queries executed in one call; results return labeled in order"),
  },
  async ({ query, queries }) => {
    const batch = queries ?? (query !== undefined ? [query] : []);
    if (batch.length === 0) {
      return { content: [{ type: "text", text: JSON.stringify({ error: "Provide query or queries" }) }] };
    }
    console.log(`[sql] calls=1 statements=${batch.length}`);
    const result = batch.length === 1 ? executeSqlQuery(batch[0]) : runSqlBatch(executeSqlQuery, batch);
    return { content: [{ type: "text", text: result }] };
  }
);

const semanticSearchTool = tool(
  "semantic_search",
  `Vector similarity search on the user's screen history (what they saw on their computer).

Use when:
- Fuzzy or conceptual queries where exact SQL keywords won't work
- User asks "when was I reading about X?" or "find where I was working on Y"
- Theme-based recall: "design mockups", "code reviews", "email about project Z"

Don't use when:
- User asks about spoken conversations or transcripts (prefer search_conversations if available)
- User asks for structured counts or stats (use execute_sql)
- User wants a broad daily recap (use get_daily_recap)

Parameter guidance:
- days: Start with 7 (default). Use 1-3 for recent activity, 14-30 for older searches.
- app_filter: Set when user specifies an app (e.g., "in Chrome", "in VS Code"). Omit for cross-app searches.
- Results are ranked by semantic similarity — top 15 returned.`,
  {
    query: z.string().describe("Natural language search query describing what the user was doing or viewing"),
    days: z.number().optional().default(7).describe("Days to search back: 1-3 for recent, 7 default, 14-30 for older"),
    app_filter: z.string().optional().describe("Filter to a specific app (e.g., 'Chrome', 'VS Code'). Omit for all apps"),
  },
  async ({ query, days, app_filter }) => {
    const result = await performSemanticSearch(query, days, app_filter);
    return { content: [{ type: "text", text: result }] };
  }
);

const getDailyRecapTool = tool(
  "get_daily_recap",
  `Get a pre-formatted daily activity recap combining app usage, conversations, and tasks.

ONE call covers any date range — pass the full range instead of calling once per day.

Use when:
- User asks "what did I do today/yesterday/this week?"
- Broad activity summaries or daily reviews
- User wants a quick overview without specifying a topic

Don't use when:
- User asks about a specific topic or event (prefer search_conversations if available)
- User needs detailed transcript content (prefer get_conversations if available)
- User wants structured data or counts (use execute_sql)

This tool runs three queries in one call (apps, conversations, tasks) — much faster than multiple execute_sql calls.
ONE call covers any date range. If the range has no data, the response says so authoritatively — do not re-verify with SQL.

Parameter guidance:
- start_date/end_date (YYYY-MM-DD): a specific day or range, e.g. "what did I do on July 15" → start_date=2026-07-15
- days_ago=0: today so far; days_ago=1: yesterday (default); days_ago=7: past week`,
  {
    days_ago: z.number().optional().default(1).describe("0=today, 1=yesterday, 7=past week. Ignored when start_date is set"),
    start_date: z.string().optional().describe("Range start, YYYY-MM-DD. Use for specific dates"),
    end_date: z.string().optional().describe("Range end (inclusive), YYYY-MM-DD. Defaults to start_date"),
  },
  async ({ days_ago, start_date, end_date }) => {
    if (!db) return { content: [{ type: "text", text: "Database not loaded." }] };

    let range;
    try {
      range = resolveRange({ days_ago, start_date, end_date });
    } catch (e) {
      return { content: [{ type: "text", text: `Error: ${e.message}` }] };
    }

    const counts = activityCounts(db, range);
    const empty = emptyRangeStatement(range, counts);
    if (empty) return { content: [{ type: "text", text: empty }] };

    // App usage
    const apps = db.prepare(`
      SELECT appName, COUNT(*) as screenshots, ROUND(COUNT(*) * 10.0 / 60, 1) as minutes,
        MIN(time(timestamp)) as first_seen, MAX(time(timestamp)) as last_seen
      FROM screenshots
      WHERE timestamp >= ? AND timestamp < ?
        AND appName IS NOT NULL AND appName != ''
      GROUP BY appName ORDER BY screenshots DESC
    `).all(range.start, range.endExclusive);

    // Conversations
    const convos = db.prepare(`
      SELECT title, overview, emoji, category, startedAt, finishedAt,
        ROUND((julianday(finishedAt) - julianday(startedAt)) * 1440, 1) as duration_min
      FROM transcription_sessions
      WHERE startedAt >= ? AND startedAt < ?
        AND deleted = 0 AND discarded = 0
      ORDER BY startedAt DESC
    `).all(range.start, range.endExclusive);

    // Action items
    const tasks = db.prepare(`
      SELECT description, completed, priority, createdAt FROM action_items
      WHERE createdAt >= ? AND createdAt < ?
        AND deleted = 0
      ORDER BY createdAt DESC
    `).all(range.start, range.endExclusive);

    const dateLabel = range.label;
    // Format compact markdown
    let out = `# ${dateLabel} Recap\n\n`;

    out += `## Apps (${apps.length} apps)\n`;
    if (apps.length === 0) {
      out += "No screen activity recorded.\n";
    } else {
      for (const a of apps.slice(0, 20)) {
        out += `- **${a.appName}**: ${a.minutes} min (${a.screenshots} captures, ${a.first_seen}–${a.last_seen})\n`;
      }
      if (apps.length > 20) out += `- ...and ${apps.length - 20} more apps\n`;
    }

    out += `\n## Conversations (${convos.length})\n`;
    if (convos.length === 0) {
      out += "No conversations recorded.\n";
    } else {
      for (const c of convos) {
        const dur = c.duration_min > 0 ? ` (${c.duration_min} min)` : "";
        const emoji = c.emoji || "";
        out += `- ${emoji} **${c.title || "Untitled"}**${dur}: ${c.overview || "No summary"}\n`;
      }
    }

    out += `\n## Tasks (${tasks.length})\n`;
    if (tasks.length === 0) {
      out += "No tasks created.\n";
    } else {
      for (const t of tasks) {
        const check = t.completed ? "[x]" : "[ ]";
        const pri = t.priority ? ` (${t.priority})` : "";
        out += `- ${check} ${t.description}${pri}\n`;
      }
    }

    // Short ranges get an hourly timeline + top window titles so one call
    // answers "what was I doing" without follow-up SQL.
    if (range.spanDays <= 2) {
      out += formatHourlyTimeline(hourlyTimeline(db, range));
      out += formatTopWindows(topWindows(db, range));
    }

    return { content: [{ type: "text", text: out }] };
  }
);

const getAppUsageTool = tool(
  "get_app_usage",
  `Day-by-app screen activity matrix for any date range in ONE call.
Returns per-app totals plus a per-day breakdown (top apps per day).
Use for "what did I spend time on", day-vs-day comparisons, and usage patterns —
instead of hand-writing GROUP BY queries per day. Captures are ~10s apart.`,
  {
    start_date: z.string().describe("Range start, YYYY-MM-DD"),
    end_date: z.string().optional().describe("Range end (inclusive), YYYY-MM-DD. Defaults to start_date"),
  },
  async ({ start_date, end_date }) => {
    if (!db) return { content: [{ type: "text", text: "Database not loaded." }] };
    let range;
    try {
      range = resolveRange({ start_date, end_date });
    } catch (e) {
      return { content: [{ type: "text", text: `Error: ${e.message}` }] };
    }
    return { content: [{ type: "text", text: formatAppUsage(range, appUsageMatrix(db, range)) }] };
  }
);

// Real one-click connect link. When a connector reports "not connected", the
// agent should hand the user a real OAuth URL to click — not invent setup
// steps. This calls the backend's existing /v1/integrations/{app_key}/oauth-url
// (verified live) with the user's token and returns the real Google consent URL.
const getConnectLinkTool = tool(
  "get_connect_link",
  `Get a one-click link the user clicks to connect an integration to Omi.
Call this WHENEVER a connector/tool reports it is "not connected" (e.g. calendar,
gmail), or the user asks how to connect one. Return the real link to the user and
tell them to click it and sign in — NEVER describe manual setup steps or invent
other methods. integration is the requested capability, e.g. "gmail" or "calendar".`,
  { integration: z.string().describe('Integration capability, e.g. "gmail" or "calendar"') },
  async ({ integration }) => {
    if (!userFirebaseToken) {
      return { content: [{ type: "text", text: "The user must be signed in to Omi before connecting integrations." }] };
    }
    try {
      const resp = await fetch(`${BACKEND_URL}/v1/integrations/${encodeURIComponent(integration)}/oauth-url`, {
        headers: { Authorization: `Bearer ${userFirebaseToken}` },
      });
      if (!resp.ok) {
        const body = await resp.text();
        return {
          content: [
            {
              type: "text",
              text:
                `No one-click connect link is available for "${integration}" (HTTP ${resp.status}). ` +
                `Tell the user this integration must be connected from Omi Settings → Apps & Integrations. Detail: ${body.slice(0, 200)}`,
            },
          ],
        };
      }
      const data = await resp.json();
      const url = data.auth_url || data.url;
      if (!url) return { content: [{ type: "text", text: `No link returned for "${integration}".` }] };
      return {
        content: [
          {
            type: "text",
            text: `CONNECT LINK for ${integration}: ${url}\nGive this link to the user verbatim and tell them to click it and finish connecting; once done, ask them to try again.`,
          },
        ],
      };
    } catch (err) {
      return { content: [{ type: "text", text: `Error getting connect link for ${integration}: ${err.message}` }] };
    }
  }
);

// --- JSON Schema → Zod converter for backend tools ---

function jsonSchemaToZod(schema) {
  const props = schema.properties || {};
  const required = new Set(schema.required || []);
  const shape = {};

  for (const [name, prop] of Object.entries(props)) {
    if (name === "config") continue; // internal LangChain param

    let zodType;
    // Handle anyOf (Optional fields from Pydantic)
    const rawType = prop.type || (prop.anyOf ? prop.anyOf.find(t => t.type && t.type !== "null")?.type : "string");

    switch (rawType) {
      case "integer":
      case "number":
        zodType = z.number();
        break;
      case "boolean":
        zodType = z.boolean();
        break;
      case "array":
        zodType = z.array(z.any());
        break;
      default:
        zodType = z.string();
    }

    if (prop.description) zodType = zodType.describe(prop.description);

    // Pydantic Optional[X] emits anyOf [X, null]: accept explicit null too —
    // models routinely pass null for "no filter", and rejecting it burned an
    // error turn (measured in the eval corpus).
    const acceptsNull = Array.isArray(prop.anyOf) && prop.anyOf.some((t) => t?.type === "null");
    if (acceptsNull) zodType = zodType.nullable();

    if (!required.has(name)) {
      zodType = zodType.optional();
      if (prop.default !== undefined && !(acceptsNull && prop.default === null)) {
        zodType = zodType.default(prop.default);
      }
    }

    shape[name] = zodType;
  }

  return shape;
}

// --- Fetch and register backend tools from Python API ---

async function fetchAndRegisterBackendTools() {
  if (!userFirebaseToken) {
    console.log("[backend-tools] No Firebase token, skipping tool fetch");
    return;
  }

  try {
    const resp = await fetch(`${BACKEND_URL}/v1/agent/tools`, {
      headers: { Authorization: `Bearer ${userFirebaseToken}` },
    });
    if (!resp.ok) {
      logEvent("error", "backend_tools_unavailable", { status: resp.status });
      const timer = setTimeout(() => { fetchAndRegisterBackendTools(); }, 30_000);
      timer.unref?.();
      return;
    }

    const data = await resp.json();
    const toolDefs = data.tools || [];
    console.log(`[backend-tools] Fetched ${toolDefs.length} tools from Python backend`);

    // Per-def isolation: one bad schema must not throw away every backend
    // tool (the API also rejects the WHOLE request over one invalid tool name,
    // so names are validated here at fetch time). First name wins on dupes —
    // the 4 local tools claim theirs first in rebuildMcpServer.
    const seenNames = new Set();
    backendTools = toolDefs.flatMap((def) => {
      try {
        if (!/^[a-zA-Z0-9_-]{1,100}$/.test(def.name ?? "")) {
          logEvent("warn", "backend_tool_skipped", { name: String(def.name), reason: "invalid_name" });
          return [];
        }
        if (seenNames.has(def.name)) {
          logEvent("warn", "backend_tool_skipped", { name: def.name, reason: "duplicate" });
          return [];
        }
        seenNames.add(def.name);
        const zodShape = jsonSchemaToZod(def.parameters || {});
        return [tool(
          def.name,
          def.description || `Backend tool: ${def.name}`,
          zodShape,
          async (params) => {
            try {
              // User-blocking: retry transient failures with backoff.
              const result = await withRetry(async () => {
                const execResp = await fetch(`${BACKEND_URL}/v1/agent/execute-tool`, {
                  method: "POST",
                  headers: {
                    "Content-Type": "application/json",
                    Authorization: `Bearer ${userFirebaseToken}`,
                  },
                  body: JSON.stringify({ tool_name: def.name, params }),
                });
                if (!execResp.ok) {
                  throw new Error(`execute-tool HTTP ${execResp.status}`);
                }
                return await execResp.json();
              }, { attempts: 3 });
              if (result.error) {
                return { content: [{ type: "text", text: `Error: ${result.error}` }] };
              }
              const bounded = truncateToolResult(result.result || JSON.stringify(result), def.name);
              return { content: [{ type: "text", text: bounded.text }] };
            } catch (err) {
              logEvent("error", "backend_tool_failed", { name: def.name, error: err });
              return { content: [{ type: "text", text: `Error calling ${def.name}: ${classifyError(err).category}` }] };
            }
          }
        )];
      } catch (err) {
        logEvent("warn", "backend_tool_skipped", { name: String(def?.name), reason: "schema_conversion", error: err });
        return [];
      }
    });

    // Rebuild MCP server with all tools
    rebuildMcpServer();
    console.log(`[backend-tools] Registered ${backendTools.length} backend tools`);
  } catch (err) {
    logEvent("error", "backend_tools_unavailable", { error: err });
    // A boot-time blip must not lose calendar/gmail for the VM's lifetime.
    const timer = setTimeout(() => { fetchAndRegisterBackendTools(); }, 30_000);
    timer.unref?.();
  }
}

function rebuildMcpServer() {
  const allTools = [
    executeSqlTool,
    semanticSearchTool,
    getDailyRecapTool,
    getAppUsageTool,
    getConnectLinkTool,
    ...backendTools,
  ];
  omiServer = createSdkMcpServer({
    name: "omi-tools",
    tools: allTools,
  });
  console.log(`[mcp] Rebuilt MCP server with ${allTools.length} tools`);
}

// --- Shared Agent Query Handler ---

async function handleQuery({ prompt, systemPrompt, cwd, send, abortController }) {
  if (!isDatabaseReady()) {
    send({ type: "error", message: "Database not available. Upload omi.db first." });
    return { text: "", sessionId: "", costUsd: 0 };
  }

  let sessionId = "";
  let fullText = "";
  let costUsd = 0;
  const pendingTools = [];

  const options = {
    model: "claude-opus-4-6",
    abortController,
    systemPrompt: systemPrompt || defaultSystemPrompt,
    allowedTools: ALLOWED_TOOLS,
    agents: agentDefinitions ?? undefined,
    includePartialMessages: true,
    permissionMode: "bypassPermissions",
    allowDangerouslySkipPermissions: true,
    maxTurns: 10,
    cwd: cwd || process.env.HOME || "/",
    mcpServers: {
      "omi-tools": omiServer,
      "playwright": {
        command: process.execPath,
        args: [
          playwrightCli,
          "--user-data-dir", join(__dirname, "chrome-profile"),
          "--headless",
          "--no-sandbox",
        ],
      },
    },
  };

  const q = query({ prompt, options });
  const routeSubagent = createSubagentRouter();

  for await (const message of q) {
    if (abortController.signal.aborted) break;

    // Subagent-origin messages: tool starts + throttled text snippets as
    // progress; their text never enters the answer stream.
    const subEvents = routeSubagent(message);
    if (subEvents) {
      for (const e of subEvents) send(e);
      continue;
    }

    switch (message.type) {
      case "system":
        if ("session_id" in message) {
          sessionId = message.session_id;
          send({ type: "init", sessionId });
        }
        break;

      case "stream_event": {
        const event = message.event;

        if (event?.type === "content_block_start" && event.content_block?.type === "tool_use") {
          const name = event.content_block.name;
          pendingTools.push(name);
          send({ type: "tool_activity", name, status: "started" });
        }

        if (event?.type === "content_block_delta" && event.delta?.type === "text_delta") {
          if (pendingTools.length > 0) {
            for (const name of pendingTools) {
              send({ type: "tool_activity", name, status: "completed" });
            }
            pendingTools.length = 0;
          }
          const text = event.delta.text;
          fullText += text;
          send({ type: "text_delta", text });
        }
        break;
      }

      case "assistant": {
        // Fallback: if streaming didn't capture text (e.g. after tool calls),
        // extract it from the complete assistant message
        const content = message.message?.content;
        if (Array.isArray(content)) {
          for (const block of content) {
            if (block.type === "text" && typeof block.text === "string") {
              // Check if this text was already sent via stream_event deltas
              if (!fullText.includes(block.text)) {
                fullText += block.text;
                send({ type: "text_delta", text: block.text });
              }
            }
          }
        }
        break;
      }

      case "result": {
        for (const name of pendingTools) {
          send({ type: "tool_activity", name, status: "completed" });
        }
        pendingTools.length = 0;

        if (message.subtype === "success") {
          costUsd = message.total_cost_usd || 0;
          // Send any final text that wasn't captured during streaming
          if (message.result) {
            const remaining = message.result.replace(fullText, "").trim();
            if (remaining) {
              send({ type: "text_delta", text: remaining });
              fullText += remaining;
            }
          }
        } else {
          const errors = message.errors || [];
          {
            const { category } = classifyError(new Error((message.errors || []).join(", ") || message.subtype));
            logEvent("error", "agent_turn_error", { category, subtype: message.subtype, errors: message.errors });
            send({ type: "error", code: category, message: USER_MESSAGES[category] ?? USER_MESSAGES.internal });
          }
        }
        break;
      }
    }
  }

  return { text: fullText, sessionId, costUsd };
}

// --- Persistent Session (streaming input mode) ---

function startPersistentSession(initialSink, log, seedTurns = []) {
  if (!isDatabaseReady()) {
    log("Cannot start persistent session: database not ready");
    return null;
  }

  // Turn history for non-destructive restart. When the SDK session dies (prompt
  // too long / maxTurns), the recreated session was previously total amnesia;
  // now the caller passes the prior turns and we seed the first prompt with a
  // condensed view (recent turns verbatim) so context survives the restart.
  const turns = [];
  let currentUser = null; // raw user prompt of the in-flight turn
  let seedPending = seedTurns.length ? seedTurns : null;

  // The session outlives any single WebSocket connection (single-user VM):
  // the client sink is swappable so a reconnect reattaches to the live
  // session and its context instead of starting from amnesia. Events emitted
  // while detached are dropped; the turn keeps running and its full text
  // still arrives in the result after reattach.
  let currentSink = initialSink;
  let lastSentAt = Date.now();
  const send = (msg) => {
    lastSentAt = Date.now();
    if (currentSink) currentSink(msg);
  };

  // Heartbeat: bound the client-visible silent gap during active turns.
  // Measured (2026-07-22): with streaming on, light turns gap <3s but heavy
  // delegated turns still go 40-64s silent while the subagent generates text
  // we deliberately swallow. A periodic status caps the gap regardless of
  // where the silence comes from.
  const heartbeatMs = parseInt(process.env.OMI_TURN_HEARTBEAT_MS || "10000", 10);
  const heartbeatTimer = setInterval(() => {
    if (!turnActive || isPrewarmTurn) return;
    if (Date.now() - lastSentAt < heartbeatMs) return;
    const elapsedS = turnStartedAt ? Math.round((Date.now() - turnStartedAt) / 1000) : 0;
    send({ type: "status", message: `Still working… (${elapsedS}s)` });
  }, Math.max(50, Math.floor(heartbeatMs / 2)));
  heartbeatTimer.unref?.();

  const sessionAbort = new AbortController();

  const options = {
    model: "claude-sonnet-4-6",
    abortController: sessionAbort,
    systemPrompt: defaultSystemPrompt,
    allowedTools: ALLOWED_TOOLS,
    agents: agentDefinitions ?? undefined,
    includePartialMessages: true,
    permissionMode: "bypassPermissions",
    allowDangerouslySkipPermissions: true,
    maxTurns: 10,
    cwd: process.env.HOME || "/",
    mcpServers: {
      "omi-tools": omiServer,
      "playwright": {
        command: process.execPath,
        args: [playwrightCli, "--user-data-dir", join(__dirname, "chrome-profile"), "--headless", "--no-sandbox"],
      },
    },
  };

  // Queue for all messages — used as the prompt (async iterable)
  const messageQueue = new AsyncMessageQueue();

  // Push the prewarm message to boot the Claude process
  messageQueue.push({
    type: "user",
    message: { role: "user", content: "ready" },
    parent_tool_use_id: null,
    session_id: "",
  });

  log(`Starting persistent session with model: ${options.model}`);
  const q = query({ prompt: messageQueue, options });

  // Session state
  let sessionId = null;
  let sessionIdResolve = null;
  const sessionIdReady = new Promise(r => { sessionIdResolve = r; });

  // Per-turn state
  let fullText = "";
  let pendingTools = [];
  let turnActive = false;
  let isPrewarmTurn = true; // First turn is prewarm, suppress output
  let turnStartedAt = null; // timestamp when query was pushed
  let firstEventSent = false; // tracks if we've sent the first stream event for this turn
  let interruptRequested = false; // set when a stop/preemption cuts the current turn short
  let dead = false; // set when the session loop dies — callers must start a fresh session
  let pendingPrompt = null; // queue-of-1: a query arriving mid-turn supersedes any earlier pending one

  function beginTurn(prompt, sid) {
    turnActive = true;
    turnStartedAt = Date.now();
    firstEventSent = false;
    currentUser = prompt;
    let content = prompt;
    // First turn of a reseeded session: prepend the condensed prior context so
    // the restart is not amnesia. The user's raw prompt is still what we track.
    if (seedPending) {
      const summary = seedPending
        .slice(0, Math.max(0, seedPending.length - 3))
        .map((t) => `- ${t.user.slice(0, 100)} → ${(t.assistant || "").slice(0, 100)}`)
        .join("\n") || "(no older turns)";
      content = `${buildCondensedSeed(summary, seedPending.slice(-3))}\n\n---\nUser: ${prompt}`;
      log(`Reseeded session with ${seedPending.length} prior turns (condensed)`);
      seedPending = null;
    }
    send({ type: "status", message: "Processing..." });
    messageQueue.push({
      type: "user",
      message: { role: "user", content },
      parent_tool_use_id: null,
      session_id: sid,
    });
  }

  // Background message processing loop — runs for entire WS lifetime
  const routeSubagent = createSubagentRouter();
  const loopPromise = (async () => {
    try {
      for await (const message of q) {
        // Subagent-origin messages: tool starts + throttled text snippets as
        // progress; their text never enters the answer stream.
        const subEvents = routeSubagent(message);
        if (subEvents) {
          if (!isPrewarmTurn) for (const e of subEvents) send(e);
          continue;
        }

        switch (message.type) {
          case "system":
            if ("session_id" in message) {
              sessionId = message.session_id;
              sessionIdResolve?.();
              sessionIdResolve = null;
              log(`Persistent session started: ${sessionId}`);
            }
            break;

          case "stream_event": {
            if (isPrewarmTurn) break; // Suppress prewarm output
            const event = message.event;
            // Emit thinking_done on first content block (text or tool_use)
            if (!firstEventSent && event?.type === "content_block_start") {
              firstEventSent = true;
              const thinkingMs = turnStartedAt ? Date.now() - turnStartedAt : 0;
              log(`[TIMING] thinking done in ${thinkingMs}ms, first block: ${event.content_block?.type}`);
              send({ type: "thinking_done", thinkingMs });
            }
            if (event?.type === "content_block_start" && event.content_block?.type === "tool_use") {
              const name = event.content_block.name;
              const elapsed = turnStartedAt ? Date.now() - turnStartedAt : 0;
              log(`[TIMING] tool_use ${name} started at +${elapsed}ms`);
              pendingTools.push(name);
              send({ type: "tool_activity", name, status: "started" });
            }
            if (event?.type === "content_block_delta" && event.delta?.type === "text_delta") {
              if (pendingTools.length > 0) {
                for (const name of pendingTools) {
                  const elapsed = turnStartedAt ? Date.now() - turnStartedAt : 0;
                  log(`[TIMING] tool ${name} completed at +${elapsed}ms`);
                  send({ type: "tool_activity", name, status: "completed" });
                }
                pendingTools.length = 0;
              }
              fullText += event.delta.text;
              send({ type: "text_delta", text: event.delta.text });
            }
            break;
          }

          case "assistant": {
            if (isPrewarmTurn) break; // Suppress prewarm output
            const content = message.message?.content;
            if (Array.isArray(content)) {
              for (const block of content) {
                if (block.type === "tool_use") {
                  const inputStr = JSON.stringify(block.input);
                  log(`[TOOL_CALL] ${block.name} input=${inputStr.length > 500 ? inputStr.slice(0, 500) + '...' : inputStr}`);
                  send({ type: "tool_activity", name: block.name, status: "started" });
                  pendingTools.push(block.name);
                }
                // Text is already streamed via stream_event handler — don't send again here
              }
            }
            break;
          }

          case "tool_result": {
            if (isPrewarmTurn) break;
            const toolName = message.tool_name || "unknown";
            const resultStr = typeof message.content === "string" ? message.content : JSON.stringify(message.content);
            const elapsed = turnStartedAt ? Date.now() - turnStartedAt : 0;
            log(`[TOOL_RESULT] +${elapsed}ms ${toolName} (${resultStr.length} chars) ${resultStr.length > 300 ? resultStr.slice(0, 300) + '...' : resultStr}`);
            break;
          }

          case "result": {
            for (const name of pendingTools) send({ type: "tool_activity", name, status: "completed" });
            pendingTools.length = 0;

            if (isPrewarmTurn) {
              isPrewarmTurn = false;
              fullText = "";
              pendingTools = [];
              // Do NOT reset turnActive here: prewarm never sets it, and a query
              // arriving before the prewarm result has already set it for its own
              // turn — clobbering it makes that turn's stop a silent no-op.
              log("Prewarm turn completed, session ready for queries");
              break;
            }

            if (message.subtype === "success") {
              const costUsd = message.total_cost_usd || 0;
              const totalMs = turnStartedAt ? Date.now() - turnStartedAt : 0;
              log(`[TIMING] result at +${totalMs}ms, cost=$${costUsd.toFixed(4)}`);
              if (message.result) {
                const remaining = message.result.replace(fullText, "").trim();
                if (remaining) {
                  send({ type: "text_delta", text: remaining });
                  fullText += remaining;
                }
              }
              // interrupted: the turn was cut short by a stop/preemption — the
              // text is a partial answer, and clients must be able to say so.
              send({ type: "result", text: fullText, sessionId, costUsd, interrupted: interruptRequested });
            } else if (interruptRequested) {
              // A user stop terminalizes with a non-success subtype in the real
              // SDK — that is a partial answer, not an error (measured: the
              // stop path previously surfaced an error to the client).
              send({ type: "result", text: fullText, sessionId, costUsd: message.total_cost_usd || 0, interrupted: true });
            } else {
              {
              const { category } = classifyError(new Error((message.errors || []).join(", ") || message.subtype));
              logEvent("error", "agent_turn_error", { category, subtype: message.subtype, errors: message.errors });
              send({ type: "error", code: category, message: USER_MESSAGES[category] ?? USER_MESSAGES.internal });
            }
            }
            // Record the completed exchange for non-destructive restart.
            if (currentUser) {
              turns.push({ user: currentUser, assistant: fullText });
              currentUser = null;
            }
            // Reset per-turn state
            fullText = "";
            pendingTools = [];
            turnActive = false;
            turnStartedAt = null;
            firstEventSent = false;
            interruptRequested = false;
            if (pendingPrompt) {
              // Superseding query queued mid-turn: start it only now, at the
              // turn boundary, so two turns' events never interleave.
              const next = pendingPrompt;
              pendingPrompt = null;
              beginTurn(next.prompt, next.sid);
            }
            break;
          }
        }
      }
    } catch (err) {
      // The session loop is the only consumer of the SDK stream — if it dies,
      // the shared session is unusable. Tell the client honestly, log at full
      // detail, and mark the session dead so the next query starts fresh
      // instead of hanging on a promise that will never resolve.
      const category = classifyError(err).category;
      logEvent("error", "session_loop_died", { category, error: err });
      dead = true;
      // If the stream died before emitting session_id, unblock anyone awaiting
      // sessionIdReady — the id will never arrive. Awaiters re-check isDead()/
      // getSessionId() and restart a fresh session instead of hanging forever.
      sessionIdResolve?.();
      sessionIdResolve = null;
      try {
        send({ type: "error", code: category, message: USER_MESSAGES[category] ?? USER_MESSAGES.internal });
      } catch {}
    }
    log("Persistent session ended");
  })();

  return {
    query: q,
    messageQueue,
    sessionAbort,
    sessionIdReady,
    loopPromise,
    getSessionId: () => sessionId,
    isTurnActive: () => turnActive,
    setTurnActive: (v) => { turnActive = v; },
    startTurnTimer: () => { turnStartedAt = Date.now(); firstEventSent = false; },
    markInterrupted: () => { interruptRequested = true; },
    isDead: () => dead,
    getTurns: () => turns,
    beginTurn,
    setPending: (prompt, sid) => {
      const replaced = pendingPrompt !== null;
      pendingPrompt = { prompt, sid };
      return replaced;
    },
    setSink: (sink) => { currentSink = sink; },
    // A stale socket's close event can fire after a newer socket attached —
    // only clear the sink if it is still ours.
    detachSink: (sink) => { if (currentSink === sink) currentSink = null; },
  };
}

// --- Mode: CLI ---

async function runCli(userMessage) {
  if (!openDatabase()) {
    console.error(`ERROR: Database not found at ${DB_PATH}`);
    process.exit(1);
  }

  // Real-connector hook: with a real Firebase token set, load the user's actual
  // backend connectors (calendar/gmail/health/…) the same way the WS /auth path
  // does, so connector interaction can be exercised against live api.omi.me.
  if (process.env.OMI_FIREBASE_TOKEN) {
    userFirebaseToken = process.env.OMI_FIREBASE_TOKEN;
    await fetchAndRegisterBackendTools();
  }

  console.log(`\n${"=".repeat(60)}`);
  console.log(`User: ${userMessage}`);
  console.log("=".repeat(60));

  const abortController = new AbortController();
  const send = (msg) => {
    switch (msg.type) {
      case "init":
        console.log(`[Session: ${msg.sessionId}]`);
        break;
      case "text_delta":
        process.stdout.write(msg.text);
        break;
      case "tool_activity":
        console.log(`\n[Tool ${msg.status}: ${msg.name}]`);
        break;
      case "error":
        console.error(`\n[Error: ${msg.message}]`);
        break;
    }
  };

  const result = await handleQuery({ prompt: userMessage, send, abortController });
  console.log(`\n\n[Done — cost: $${result.costUsd.toFixed(4)}]`);
  db.close();
}

// --- Auth helper for HTTP endpoints ---

function verifyAuth(req) {
  const authHeader = req.headers["authorization"];
  const url = new URL(req.url, `http://${req.headers.host}`);
  const tokenParam = url.searchParams.get("token");
  const token = authHeader?.startsWith("Bearer ") ? authHeader.slice(7) : tokenParam;
  return token === AUTH_TOKEN;
}

// --- Mode: WebSocket Server ---

// --- Idle auto-stop: shut down VM after 30 min of no activity ---

async function checkIdleAndStop() {
  const idleMs = Date.now() - lastActivityAt;
  if (idleMs < IDLE_TIMEOUT_MS) return;

  console.log(`[server] Idle for ${Math.round(idleMs / 60000)} minutes — shutting down VM...`);

  // Try GCE API first, fall back to sudo shutdown
  let stopped = false;
  try {
    const metaHeaders = { "Metadata-Flavor": "Google" };
    const name = await fetch(
      "http://metadata.google.internal/computeMetadata/v1/instance/name",
      { headers: metaHeaders },
    ).then((r) => r.text());
    const zonePath = await fetch(
      "http://metadata.google.internal/computeMetadata/v1/instance/zone",
      { headers: metaHeaders },
    ).then((r) => r.text());
    const tokenResp = await fetch(
      "http://metadata.google.internal/computeMetadata/v1/instance/service-accounts/default/token",
      { headers: metaHeaders },
    );
    if (tokenResp.ok) {
      const token = (await tokenResp.json()).access_token;
      const stopUrl = `https://compute.googleapis.com/compute/v1/${zonePath}/instances/${name}/stop`;
      const resp = await fetch(stopUrl, {
        method: "POST",
        headers: { Authorization: `Bearer ${token}` },
      });
      console.log(`[server] Stop request for ${name}: HTTP ${resp.status}`);
      if (resp.ok) stopped = true;
    } else {
      console.log(`[server] No service account token (HTTP ${tokenResp.status}), using shutdown`);
    }
  } catch (err) {
    console.log(`[server] GCE API failed: ${err.message}`);
  }

  if (!stopped) {
    console.log(`[server] Falling back to sudo shutdown...`);
    const { execSync } = await import("child_process");
    try {
      execSync("sudo shutdown -h now");
    } catch (e) {
      console.log(`[server] Shutdown command failed: ${e.message}`);
    }
  }
}

// --- Hot-reload: check GCS for new version every 10 minutes ---

let currentVersion = null;

async function checkAndApplyUpdate(log) {
  try {
    const resp = await fetch(`${GCS_BASE}/version.txt`);
    if (!resp.ok) return;
    const version = (await resp.text()).trim();
    if (currentVersion === null) {
      currentVersion = version;
      log(`Running version: ${version}`);
      return;
    }
    if (version === currentVersion) return;
    log(`Update available: ${version} (current: ${currentVersion}), downloading...`);
    const agentResp = await fetch(`${GCS_BASE}/agent.mjs`);
    if (!agentResp.ok) {
      log(`Update download failed: HTTP ${agentResp.status}`);
      return;
    }
    const newCode = await agentResp.text();
    writeFileSync(join(__dirname, "agent.mjs"), newCode);
    log(`Update applied (${version}), restarting in 3s...`);
    setTimeout(() => process.exit(0), 3000);
  } catch (e) {
    log(`Update check failed: ${e.message}`);
  }
}

function startServer() {
  const log = (msg) => console.log(`[server] ${msg}`);

  if (!AUTH_TOKEN) {
    console.error("ERROR: AUTH_TOKEN environment variable is required for server mode.");
    process.exit(1);
  }

  // Try to open DB if it exists (may not exist yet for fresh VMs)
  if (openDatabase()) {
    log("Database loaded at startup");
  } else {
    log(`Database not found at ${DB_PATH} — waiting for upload`);
  }

  // Ensure data directory exists for uploads
  const dataDir = dirname(DB_PATH);
  if (!existsSync(dataDir)) {
    mkdirSync(dataDir, { recursive: true });
  }

  // Clean up any stale upload temp file left by a previous crashed upload
  const uploadTmpPath = DB_PATH + ".uploading";
  if (existsSync(uploadTmpPath)) {
    try {
      unlinkSync(uploadTmpPath);
      log("Cleaned up stale upload temp file");
    } catch (e) {
      log(`Warning: could not remove stale upload temp file: ${e.message}`);
    }
  }

  const httpServer = createServer((req, res) => {
    // Health check endpoint (no auth)
    if (req.url === "/health" && req.method === "GET") {
      res.writeHead(200, { "Content-Type": "application/json" });
      res.end(JSON.stringify({
        status: "ok",
        uptime: Math.round(process.uptime()),
        databaseReady: isDatabaseReady(),
      }));
      return;
    }

    // Database upload endpoint
    if (req.url?.startsWith("/upload") && req.method === "POST") {
      if (!verifyAuth(req)) {
        res.writeHead(401, { "Content-Type": "application/json" });
        res.end(JSON.stringify({ error: "Unauthorized" }));
        return;
      }

      const contentLength = parseInt(req.headers["content-length"] || "0", 10);
      if (contentLength > MAX_UPLOAD_BYTES) {
        res.writeHead(413, { "Content-Type": "application/json" });
        res.end(JSON.stringify({ error: "File too large", maxBytes: MAX_UPLOAD_BYTES }));
        return;
      }

      const encoding = (req.headers["content-encoding"] || "").toLowerCase();
      const isCompressed = encoding === "gzip" || encoding === "deflate" || encoding === "zlib";
      log(`Upload started (${(contentLength / 1024 / 1024).toFixed(1)} MB${isCompressed ? ", " + encoding + " compressed" : ""})`);

      // Write to temp file first, then rename (atomic)
      const tmpPath = DB_PATH + ".uploading";
      const stream = createWriteStream(tmpPath);
      let bytesReceived = 0;
      let aborted = false;

      // If compressed, pipe through decompressor
      let decompressor = null;
      if (encoding === "gzip") {
        decompressor = createGunzip();
      } else if (encoding === "deflate" || encoding === "zlib") {
        decompressor = createInflateRaw();
      }

      const writeTarget = decompressor || stream;

      if (decompressor) {
        decompressor.pipe(stream);
        decompressor.on("error", (err) => {
          log(`Upload decompression error: ${err.message}`);
          aborted = true;
          stream.destroy();
          try { unlinkSync(tmpPath); } catch {}
        });
      }

      stream.on("error", (err) => {
        log(`Upload stream error: ${err.message}`);
        aborted = true;
        try { unlinkSync(tmpPath); } catch {}
      });

      req.on("data", (chunk) => {
        if (aborted) return;
        bytesReceived += chunk.length;
        if (decompressor) {
          decompressor.write(chunk);
        } else {
          stream.write(chunk);
        }
      });

      const finalize = () => {
        // Close existing DB connection before replacing the file
        if (db) {
          try { db.close(); } catch {}
          db = null;
        }

        // Remove WAL/SHM files if they exist (stale from previous DB)
        try { unlinkSync(DB_PATH + "-wal"); } catch {}
        try { unlinkSync(DB_PATH + "-shm"); } catch {}

        // Atomic rename
        renameSync(tmpPath, DB_PATH);

        const finalSize = statSync(DB_PATH).size;
        lastActivityAt = Date.now();
        log(`Upload complete: ${(bytesReceived / 1024 / 1024).toFixed(1)} MB received → ${(finalSize / 1024 / 1024).toFixed(1)} MB on disk`);

        // Re-open the database
        if (openDatabase()) {
          log("Database loaded after upload");
          res.writeHead(200, { "Content-Type": "application/json" });
          res.end(JSON.stringify({
            status: "ok",
            bytesReceived,
            finalSize,
            databaseReady: true,
          }));
        } else {
          log("ERROR: Failed to open database after upload");
          res.writeHead(500, { "Content-Type": "application/json" });
          res.end(JSON.stringify({ error: "Failed to open uploaded database" }));
        }
      };

      req.on("end", () => {
        if (aborted) return;
        if (decompressor) {
          // End the decompressor; stream.end is called via pipe when decompressor finishes
          decompressor.end();
          stream.on("finish", finalize);
        } else {
          stream.end(finalize);
        }
      });

      req.on("error", (err) => {
        log(`Upload error: ${err.message}`);
        aborted = true;
        if (decompressor) decompressor.destroy();
        stream.destroy();
        try { unlinkSync(tmpPath); } catch {}
        if (!res.headersSent) {
          res.writeHead(500, { "Content-Type": "application/json" });
          res.end(JSON.stringify({ error: err.message }));
        }
      });

      req.on("aborted", () => {
        log("Upload aborted by client");
        aborted = true;
        stream.destroy();
        try { unlinkSync(tmpPath); } catch {}
      });

      return;
    }

    // Auth endpoint — receives Firebase token from desktop app
    if (req.url?.startsWith("/auth") && req.method === "POST") {
      if (!verifyAuth(req)) {
        res.writeHead(401, { "Content-Type": "application/json" });
        res.end(JSON.stringify({ error: "Unauthorized" }));
        return;
      }

      let body = "";
      req.on("data", (chunk) => { body += chunk.toString(); });
      req.on("end", async () => {
        try {
          const payload = JSON.parse(body);
          const { firebaseToken } = payload;
          if (!firebaseToken) {
            res.writeHead(400, { "Content-Type": "application/json" });
            res.end(JSON.stringify({ error: "Missing firebaseToken" }));
            return;
          }

          const isFirst = userFirebaseToken === null;
          userFirebaseToken = firebaseToken;
          lastActivityAt = Date.now();
          log(`Firebase token ${isFirst ? "received" : "refreshed"}`);

          // On first token, fetch backend tools
          if (isFirst) {
            await fetchAndRegisterBackendTools();
          }

          res.writeHead(200, { "Content-Type": "application/json" });
          res.end(JSON.stringify({ status: "ok", toolsRegistered: backendTools.length }));
        } catch (err) {
          log(`Auth error: ${err.message}`);
          res.writeHead(500, { "Content-Type": "application/json" });
          res.end(JSON.stringify({ error: err.message }));
        }
      });
      return;
    }

    // Keepalive ping — resets idle timer so VM doesn't auto-stop
    if (req.url?.startsWith("/ping") && req.method === "POST") {
      if (!verifyAuth(req)) {
        res.writeHead(401, { "Content-Type": "application/json" });
        res.end(JSON.stringify({ error: "Unauthorized" }));
        return;
      }
      lastActivityAt = Date.now();
      log("Keepalive ping received");
      res.writeHead(200, { "Content-Type": "application/json" });
      res.end(JSON.stringify({ status: "ok" }));
      return;
    }

    // Incremental sync endpoint — desktop pushes new/changed rows
    if (req.url?.startsWith("/sync") && req.method === "POST") {
      if (!verifyAuth(req)) {
        res.writeHead(401, { "Content-Type": "application/json" });
        res.end(JSON.stringify({ error: "Unauthorized" }));
        return;
      }
      if (!db) {
        res.writeHead(503, { "Content-Type": "application/json" });
        res.end(JSON.stringify({ error: "Database not loaded. Upload omi.db first." }));
        return;
      }

      let body = "";
      req.on("data", (chunk) => { body += chunk.toString(); });
      req.on("end", () => {
        try {
          const payload = JSON.parse(body);
          const { table, rows } = payload;
          if (!table || !Array.isArray(rows) || rows.length === 0) {
            res.writeHead(400, { "Content-Type": "application/json" });
            res.end(JSON.stringify({ error: "Required: { table: string, rows: [{...}, ...] }" }));
            return;
          }
          if (!SYNC_TABLES.has(table)) {
            res.writeHead(400, { "Content-Type": "application/json" });
            res.end(JSON.stringify({ error: `Table '${table}' not in sync whitelist` }));
            return;
          }

          const cols = Object.keys(rows[0]);

          // Auto-add any columns missing from the table (handles schema migrations
          // where the VM has an older database than the desktop app).
          const existingCols = new Set(
            db.prepare(`PRAGMA table_info("${table}")`).all().map(r => r.name)
          );
          for (const col of cols) {
            if (!existingCols.has(col)) {
              log(`Sync: adding missing column "${col}" to ${table}`);
              db.prepare(`ALTER TABLE "${table}" ADD COLUMN "${col}"`).run();
            }
          }

          const placeholders = cols.map(() => "?").join(", ");
          const sql = `INSERT OR REPLACE INTO "${table}" (${cols.map(c => `"${c}"`).join(", ")}) VALUES (${placeholders})`;

          const stmt = db.prepare(sql);
          const insertMany = db.transaction((rowList) => {
            for (const row of rowList) {
              const values = cols.map((col) => {
                const val = row[col];
                // Decode base64 embedding columns back to Buffer
                if (col === "embedding" && typeof val === "string" && val.length > 0) {
                  return Buffer.from(val, "base64");
                }
                if (val === null || val === undefined) return null;
                return val;
              });
              stmt.run(...values);
            }
          });

          insertMany(rows);

          // FTS is kept in sync by triggers on the content tables.
          // INSERT OR REPLACE fires DELETE then INSERT triggers, which update FTS automatically.

          lastActivityAt = Date.now();
          log(`Sync: ${rows.length} rows → ${table}`);
          res.writeHead(200, { "Content-Type": "application/json" });
          res.end(JSON.stringify({ applied: rows.length, table }));
        } catch (err) {
          log(`Sync error: ${err.message}`);
          res.writeHead(500, { "Content-Type": "application/json" });
          res.end(JSON.stringify({ error: err.message }));
        }
      });
      return;
    }

    res.writeHead(404);
    res.end();
  });

  const wss = new WebSocketServer({
    server: httpServer,
    path: "/ws",
    verifyClient: ({ req }, done) => {
      if (!verifyAuth(req)) {
        log("Rejected connection: invalid token");
        done(false, 401, "Unauthorized");
        return;
      }
      done(true);
    },
  });

  // Single persistent session per VM (single user); survives WS reconnects.
  let sharedSession = null;

  wss.on("connection", (ws) => {
    log("Client connected");

    const send = (msg) => {
      if (ws.readyState === ws.OPEN) {
        ws.send(JSON.stringify(msg));
      }
    };

    // Reattach to the shared sharedSession (single-user VM) so a reconnect keeps
    // the conversation context instead of starting from amnesia.
    if (sharedSession) {
      sharedSession.setSink(send);
      send({ type: "session_state", active: true, sessionId: sharedSession.getSessionId() ?? "" });
    } else {
      send({ type: "session_state", active: false });
    }

    ws.on("message", async (data) => {
      let msg;
      try {
        msg = JSON.parse(data.toString());
      } catch {
        send({ type: "error", message: "Invalid JSON" });
        return;
      }

      switch (msg.type) {
        case "prewarm": {
          if (sharedSession) {
            log("Prewarm: sharedSession already active");
            send({ type: "prewarm_ack", success: true });
            break;
          }
          log("Prewarm: starting persistent sharedSession...");
          sharedSession = startPersistentSession(send, log);
          if (!sharedSession) {
            send({ type: "prewarm_ack", success: false });
            break;
          }
          sharedSession.sessionIdReady.then(() => {
            // sessionIdReady also resolves when the loop dies before an id —
            // treat a dead/idless session as a failed prewarm, not a success.
            if (sharedSession.isDead?.() || !sharedSession.getSessionId()) {
              send({ type: "prewarm_ack", success: false });
              return;
            }
            send({ type: "prewarm_ack", success: true });
            log("Prewarm: sharedSession ready");
          }).catch(() => {
            send({ type: "prewarm_ack", success: false });
          });
          break;
        }

        case "query": {
          lastActivityAt = Date.now();
          const prompt = msg.prompt;
          log(`Query: ${prompt?.slice(0, 100)}`);

          let recoverySeed = [];
          if (sharedSession?.isDead?.()) {
            // Non-destructive restart: carry the prior turns forward (condensed)
            // instead of losing all context when the session died.
            const prior = sharedSession.getTurns?.() || [];
            const plan = planCondensation(prior, { maxChars: 40000, keepRecent: 3 });
            recoverySeed = plan.condense ? [...plan.summarize, ...plan.keep] : prior;
            logEvent("warn", "session_recreated_after_death", { priorTurns: prior.length });
            sharedSession = null;
          }
          // Start sharedSession if not started yet
          if (!sharedSession) {
            log("No active sharedSession, starting on first query...");
            sharedSession = startPersistentSession(send, log, recoverySeed);
            if (!sharedSession) {
              send({ type: "error", message: "Database not available" });
              break;
            }
          }

          // Wait for session_id. This also resolves if the loop died before an
          // id arrived (sessionIdResolve is called from the loop catch), so a
          // dead session no longer hangs here — re-check before proceeding.
          await sharedSession.sessionIdReady;
          const sid = sharedSession.getSessionId();
          if (sharedSession.isDead?.() || !sid) {
            send({ type: "error", message: "Session failed to initialize" });
            sharedSession = null;
            break;
          }

          if (sharedSession.isTurnActive()) {
            // Supersede: park the new prompt in the queue-of-1 and interrupt.
            // It starts at the turn boundary (result reset), never mid-stream —
            // a rapid double-tap can no longer interleave two turns' events.
            const replaced = sharedSession.setPending(prompt, sid);
            if (replaced) send({ type: "status", message: "Answering your latest question instead." });
            log("Interrupting active turn for superseding query");
            sharedSession.markInterrupted();
            markOwnedAbort();
            try { await sharedSession.query.interrupt(); } catch (e) { log(`Interrupt error: ${e.message}`); }
            break;
          }

          sharedSession.beginTurn(prompt, sid);
          break;
        }

        case "stop": {
          log("Stop requested");
          if (sharedSession?.isTurnActive()) {
            sharedSession.markInterrupted();
            markOwnedAbort();
            try { await sharedSession.query.interrupt(); } catch (e) { log(`Interrupt error: ${e.message}`); }
          }
          break;
        }

        default:
          send({ type: "error", message: `Unknown message type: ${msg.type}` });
      }
    });

    ws.on("close", () => {
      log("Client disconnected");
      // Detach only — the session (and any in-flight turn) survives so a
      // reconnect resumes with full context instead of amnesia.
      sharedSession?.detachSink(send);
    });

    ws.on("error", (err) => {
      log(`WebSocket error: ${err.message}`);
    });

    // Signal readiness
    send({ type: "init", sessionId: "" });
  });

  httpServer.listen(PORT, () => {
    log(`Listening on port ${PORT}`);
    log(`WebSocket: ws://0.0.0.0:${PORT}/ws`);
    log(`Health check: http://0.0.0.0:${PORT}/health`);
    log(`Upload: POST http://0.0.0.0:${PORT}/upload`);
    log(`Auth:   POST http://0.0.0.0:${PORT}/auth`);
    log(`Sync:   POST http://0.0.0.0:${PORT}/sync`);
    log(`Backend URL: ${BACKEND_URL}`);
    log(`Idle auto-stop: ${IDLE_TIMEOUT_MS / 60000} min timeout, checking every ${IDLE_CHECK_INTERVAL_MS / 60000} min`);
  });

  // Start idle auto-stop checker
  setInterval(checkIdleAndStop, IDLE_CHECK_INTERVAL_MS);

  // Start hot-reload version checker
  const UPDATE_INTERVAL_MS = 10 * 60 * 1000; // 10 minutes
  if (process.env.OMI_AGENT_DISABLE_UPDATES !== "1") {
    checkAndApplyUpdate(log);
    setInterval(() => checkAndApplyUpdate(log), UPDATE_INTERVAL_MS);
  }
}

// --- Main ---

const args = process.argv.slice(2);

if (args[0] === "--serve") {
  startServer();
} else if (args.length > 0) {
  // CLI mode
  await runCli(args.join(" "));
} else {
  console.log("Usage:");
  console.log('  node agent.mjs "What did I do today?"     # CLI mode');
  console.log("  node agent.mjs --serve                     # WebSocket server mode");
  console.log("");
  console.log("Environment variables:");
  console.log("  DB_PATH       Path to omi.db (default: ~/omi-agent/data/omi.db)");
  console.log("  GEMINI_API_KEY  For semantic search embeddings");
  console.log("  AUTH_TOKEN    Required for server mode (bearer token for WebSocket auth)");
  console.log("  PORT          Server port (default: 8080)");
  process.exit(0);
}