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 { SessionIndexSchema, SessionSchema } from "./schemas";
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) : "",
  };
}

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();
  await writeIndex(summaries.filter((summary) => summary.id !== id));

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

export async function clearSessions(): Promise<void> {
  const summaries = await listSessions();

  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);
}

/** 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) };
}