lam-agent-vicinae/src/lib/sessions.ts

import { LocalStorage } from "@vicinae/api";
import { MAX_SESSIONS, MAX_STORED_CHARS_PER_TURN, MAX_TURNS_PER_SESSION, STORAGE_INDEX_KEY, STORAGE_SESSION_PREFIX, TITLE_MAX_LENGTH } from "./constants";
import { sweepExpired, type SweepOutcome } from "./retention";
import { SessionIndexSchema, SessionSchema } from "./schemas";
import { listSnipFiles, removeSnipFiles } from "./snip";
import type { Attachment, ChatTurn, Session, SessionSummary } from "./types";
import { titleFromPrompt, truncate } from "./utils";

/**
 * Session persistence on top of Vicinae's LocalStorage.
 *
 * LocalStorage only accepts strings, numbers, and booleans, so everything is
 * JSON-encoded. An index holding small summaries is stored separately from the
 * session bodies: listing sessions is a hot path, and it should not require
 * decoding every transcript.
 *
 * A session body is one key, so saving a turn rewrites only that session.
 * Transcripts are capped twice over, per turn and per session, because
 * LocalStorage has no documented ceiling and a runaway loop should not be able
 * to wedge the extension.
 */

function newId(): string {
  return `${Date.now().toString(36)}-${Math.random().toString(36).slice(2, 8)}`;
}

/** Builds the lightweight index record kept in the session list. */
export function summarizeSession(session: Session): SessionSummary {
  return summarize(session);
}

function sessionKey(id: string): string {
  return `${STORAGE_SESSION_PREFIX}${id}`;
}

function summarize(session: Session): SessionSummary {
  const lastTurn = session.turns[session.turns.length - 1];
  return {
    id: session.id,
    title: session.title,
    model: session.model,
    createdAt: session.createdAt,
    updatedAt: session.updatedAt,
    turnCount: session.turns.length,
    preview: lastTurn ? truncate(lastTurn.response || lastTurn.prompt, TITLE_MAX_LENGTH * 2) : "",
    attachmentPaths: attachmentPathsOf(session),
  };
}

/** Every image path a conversation refers to, deduplicated and in turn order. */
export function attachmentPathsOf(session: Session): string[] {
  return [
    ...new Set(session.turns.flatMap((turn) => turn.attachments.map((attachment) => attachment.path))),
  ];
}

export function createSession(model: string, now: number = Date.now()): Session {
  return {
    id: newId(),
    title: "New chat",
    model,
    createdAt: now,
    updatedAt: now,
    turns: [],
  };
}

/** Creates the turn record for a prompt that is about to be sent. */
export function createTurn(
  prompt: string,
  model: string,
  now: number = Date.now(),
  attachments: Attachment[] = [],
): ChatTurn {
  return {
    id: newId(),
    prompt,
    response: "",
    thinking: "",
    model,
    createdAt: now,
    toolCalls: [],
    attachments,
  };
}

/** Drops the oldest turns once a session exceeds its cap, keeping the most recent context. */
export function capTurns(turns: ChatTurn[]): ChatTurn[] {
  return turns.length <= MAX_TURNS_PER_SESSION ? turns : turns.slice(turns.length - MAX_TURNS_PER_SESSION);
}

/** Trims a single turn's stored text so one enormous answer cannot dominate storage. */
export function capTurn(turn: ChatTurn): ChatTurn {
  const cap = (text: string): string =>
    text.length <= MAX_STORED_CHARS_PER_TURN ? text : `${text.slice(0, MAX_STORED_CHARS_PER_TURN)}\n\n[truncated]`;

  return {
    ...turn,
    prompt: cap(turn.prompt),
    response: cap(turn.response),
    thinking: cap(turn.thinking),
    toolCalls: turn.toolCalls.map((call) => ({ ...call, result: cap(call.result) })),
  };
}

export async function listSessions(): Promise<SessionSummary[]> {
  try {
    const raw = await LocalStorage.getItem<string>(STORAGE_INDEX_KEY);
    if (!raw) return [];

    const parsed = SessionIndexSchema.safeParse(JSON.parse(raw));
    if (!parsed.success) return [];

    return [...parsed.data].sort((a, b) => b.updatedAt - a.updatedAt);
  } catch {
    // A corrupt index should degrade to "no sessions", never to a crashed command.
    return [];
  }
}

async function writeIndex(summaries: SessionSummary[]): Promise<void> {
  await LocalStorage.setItem(STORAGE_INDEX_KEY, JSON.stringify(summaries));
}

export async function loadSession(id: string): Promise<Session | undefined> {
  try {
    const raw = await LocalStorage.getItem<string>(sessionKey(id));
    if (!raw) return undefined;

    const parsed = SessionSchema.safeParse(JSON.parse(raw));
    if (!parsed.success) return undefined;

    return {
      ...parsed.data,
      turns: parsed.data.turns.map((turn) => ({ ...turn, toolCalls: turn.toolCalls ?? [] })),
    };
  } catch {
    return undefined;
  }
}

/**
 * Persists a session and refreshes its index entry, evicting the oldest
 * sessions past {@link MAX_SESSIONS} along with their storage keys.
 */
export async function saveSession(session: Session): Promise<void> {
  const capped: Session = {
    ...session,
    turns: capTurns(session.turns).map(capTurn),
  };

  await LocalStorage.setItem(sessionKey(capped.id), JSON.stringify(capped));

  const summaries = await listSessions();
  const others = summaries.filter((summary) => summary.id !== capped.id);
  const merged = [summarize(capped), ...others].sort((a, b) => b.updatedAt - a.updatedAt);

  const kept = merged.slice(0, MAX_SESSIONS);
  const evicted = merged.slice(MAX_SESSIONS);

  await writeIndex(kept);

  for (const summary of evicted) {
    try {
      await LocalStorage.removeItem(sessionKey(summary.id));
    } catch {
      // Losing an eviction just leaves an orphaned key; not worth failing the save.
    }
  }
}

export async function deleteSession(id: string): Promise<void> {
  const summaries = await listSessions();
  const survivors = summaries.filter((summary) => summary.id !== id);
  const owned = await attachmentPathsFor([id], summaries);

  await writeIndex(survivors);

  try {
    await LocalStorage.removeItem(sessionKey(id));
  } catch {
    // Already gone.
  }

  // Deleting a conversation takes its screenshots with it, the same as an
  // expired one does. Without this a manual delete leaks image files that
  // nothing will ever collect, and the orphan pass only runs while the retention
  // setting is on.
  await removeSnipFiles(await pathsSafeToDelete(owned, survivors));
}

/**
 * The paths a deleted conversation held that nothing else still needs.
 *
 * Ownership alone is not enough: two conversations can refer to the same
 * capture, and removing the file would leave the survivor showing a broken
 * thumbnail. When the surviving references cannot be established, nothing is
 * removed.
 */
async function pathsSafeToDelete(owned: string[], survivors: SessionSummary[]): Promise<string[]> {
  if (owned.length === 0) return [];

  const kept = await pathsKeptBy(survivors);
  return kept === undefined ? [] : owned.filter((path) => !kept.has(path));
}

/**
 * Every image path the given conversations refer to, or `undefined` when that
 * cannot be established.
 *
 * Index-backed, for the reason set out on `SessionSummary.attachmentPaths`:
 * deriving this from the transcripts would mean a hundred JSON parses on every
 * sweep, and the sweep runs after every answer.
 *
 * Entries written before that field existed are filled in from their transcripts
 * and the index is written back once, so an upgrade pays for the repair on its
 * first sweep rather than on every one. A transcript that will not load is the
 * one case that returns `undefined` instead of an answer: its screenshots would
 * look exactly like orphans, and deleting someone's images because a key failed
 * to parse once is not a trade worth making.
 */
async function pathsKeptBy(summaries: SessionSummary[]): Promise<Set<string> | undefined> {
  if (summaries.every((summary) => summary.attachmentPaths !== undefined)) {
    return new Set(summaries.flatMap((summary) => summary.attachmentPaths ?? []));
  }

  let unknown = false;

  const repaired = await Promise.all(
    summaries.map(async (summary): Promise<SessionSummary> => {
      if (summary.attachmentPaths !== undefined) return summary;

      const session = await loadSession(summary.id);
      if (!session) {
        unknown = true;
        return summary;
      }

      return { ...summary, attachmentPaths: attachmentPathsOf(session) };
    }),
  );

  if (unknown) return undefined;

  try {
    await writeIndex(repaired);
  } catch {
    // A failed repair only means the next sweep reads these transcripts again.
  }

  return new Set(repaired.flatMap((summary) => summary.attachmentPaths ?? []));
}

/**
 * The image paths held by the named conversations.
 *
 * Bodies are read before their keys are removed, because the turns are the only
 * record of which screenshots belonged to them. A body that will not load
 * contributes the paths its summary already knew about, which is one of the
 * things `attachmentPaths` is stored for.
 */
async function attachmentPathsFor(ids: string[], summaries: SessionSummary[]): Promise<string[]> {
  if (ids.length === 0) return [];

  const wanted = new Set(ids);
  const doomed = summaries.filter((summary) => wanted.has(summary.id));

  const collected = await Promise.all(
    doomed.map(async (summary) => {
      const session = await loadSession(summary.id);
      return session ? attachmentPathsOf(session) : (summary.attachmentPaths ?? []);
    }),
  );

  return [...new Set(collected.flat())];
}

/**
 * Deletes several conversations in one pass, returning the image paths they held.
 *
 * Separate from {@link deleteSession} rather than a loop over it because the
 * sweep can expire dozens of conversations at once, and one index rewrite per
 * conversation is dozens of writes to re-save a shrinking list.
 *
 * Bodies are read before their keys are removed, because the turns are the only
 * record of which screenshots belonged to these conversations.
 */
export async function deleteSessions(ids: string[]): Promise<string[]> {
  if (ids.length === 0) return [];

  const removing = new Set(ids);
  const summaries = await listSessions();
  const paths = await attachmentPathsFor(ids, summaries);

  await writeIndex(summaries.filter((summary) => !removing.has(summary.id)));

  for (const id of ids) {
    try {
      await LocalStorage.removeItem(sessionKey(id));
    } catch {
      // Already gone.
    }
  }

  // Returned rather than removed: the sweep that calls this has already listed
  // the snip directory, and removing the files here would mean two passes over
  // it for one tidy-up.
  return paths;
}

/**
 * Image paths still referred to by a conversation that exists, or `undefined`
 * when that cannot be established.
 *
 * See {@link pathsKeptBy}, which is where the decision actually lives.
 */
export async function referencedAttachments(): Promise<string[] | undefined> {
  const kept = await pathsKeptBy(await listSessions());
  return kept === undefined ? undefined : [...kept];
}

/**
 * Runs the retention sweep against real storage and the real file system.
 *
 * The one place the pure decision in `retention.ts` meets its side effects, so a
 * command entrypoint has a single thing to call. Resolves to a zero outcome
 * rather than throwing: it runs on the way into a command the user is waiting
 * for, and a tidy-up that fails is not worth an error message.
 */
export async function sweepExpiredSessions(input: {
  maxAgeMs?: number;
  now?: number;
  keepId?: string;
}): Promise<SweepOutcome> {
  return sweepExpired(
    { maxAgeMs: input.maxAgeMs, now: input.now ?? Date.now(), keepId: input.keepId },
    { listSessions, deleteSessions, referencedAttachments, listSnipFiles, removeSnipFiles },
  ).catch(() => ({ sessions: 0, files: 0 }));
}

export async function clearSessions(): Promise<void> {
  const summaries = await listSessions();
  const owned = await attachmentPathsFor(
    summaries.map((summary) => summary.id),
    summaries,
  );

  for (const summary of summaries) {
    try {
      await LocalStorage.removeItem(sessionKey(summary.id));
    } catch {
      // Keep going: clearing as much as possible beats failing outright.
    }
  }

  await LocalStorage.removeItem(STORAGE_INDEX_KEY);

  // Every screenshot too, for the same reason a single delete takes its own.
  // Nothing survives to share them with, so they are all safe to remove.
  await removeSnipFiles(await pathsSafeToDelete(owned, []));
}

/** Recomputes a session's title from its first prompt, once there is a prompt worth naming it after. */
export function withTitleFromFirstTurn(session: Session): Session {
  if (session.turns.length === 0) return session;
  if (session.title !== "New chat") return session;

  return { ...session, title: titleFromPrompt(session.turns[0].prompt, TITLE_MAX_LENGTH) };
}