TJWaterAgent/src/history/store.ts

import { join } from "node:path";

import { config } from "../config.js";
import {
  atomicWriteJson,
  ensureDirectory,
  listJsonFiles,
  readJsonFile,
  toStableId,
} from "../utils/fileStore.js";
import { sanitizePersistentDocument } from "../utils/persistencePolicy.js";

export type SessionTurnRecord = {
  id: string;
  assistantMessage: string;
  timestamp: string;
  toolCallCount: number;
  userMessage: string;
};

type SessionTranscriptRecord = {
  actorKey: string;
  clientSessionId?: string;
  projectKey: string;
  sessionId: string;
  turns: SessionTurnRecord[];
  updatedAt: string;
};

export type SessionSearchHit = {
  matchedField: "assistant" | "user";
  score: number;
  sessionId: string;
  snippet: string;
  timestamp: string;
  turnId: string;
};

type SessionHistoryContext = {
  actorKey: string;
  clientSessionId?: string;
  projectKey: string;
  sessionId: string;
};

export class SessionHistoryStore {
  private readonly writeQueues = new Map<string, Promise<void>>();

  constructor(private readonly baseDir = config.SESSION_HISTORY_STORAGE_DIR) {}

  async initialize() {
    await ensureDirectory(this.baseDir);
  }

  async appendTurn(
    context: SessionHistoryContext,
    turn: {
      assistantMessage: string;
      toolCallCount: number;
      userMessage: string;
    },
  ) {
    const key = this.filePath(context);
    return this.serializeWrite(key, async () => {
      const transcript = (await this.readTranscript(context)) ?? {
        actorKey: context.actorKey,
        clientSessionId: context.clientSessionId,
        projectKey: context.projectKey,
        sessionId: context.sessionId,
        turns: [],
        updatedAt: new Date().toISOString(),
      };
      const userMessage = sanitizePersistentDocument(turn.userMessage, 4000);
      const assistantMessage = sanitizePersistentDocument(turn.assistantMessage, 4000);
      if (!userMessage || !assistantMessage) {
        return transcript;
      }

      const timestamp = new Date().toISOString();
      const record: SessionTurnRecord = {
        id: toStableId(context.sessionId, timestamp, userMessage, assistantMessage),
        assistantMessage,
        timestamp,
        toolCallCount: Math.max(0, turn.toolCallCount),
        userMessage,
      };
      transcript.clientSessionId = context.clientSessionId ?? transcript.clientSessionId;
      transcript.turns.push(record);
      if (transcript.turns.length > config.SESSION_HISTORY_MAX_TURNS_PER_SESSION) {
        transcript.turns = transcript.turns.slice(
          transcript.turns.length - config.SESSION_HISTORY_MAX_TURNS_PER_SESSION,
        );
      }
      transcript.updatedAt = timestamp;
      await atomicWriteJson(key, transcript);
      return transcript;
    });
  }

  async getRecentTurns(
    context: SessionHistoryContext,
    limit: number,
  ): Promise<SessionTurnRecord[]> {
    const transcript = await this.readTranscript(context);
    if (!transcript) {
      return [];
    }
    return transcript.turns.slice(-Math.max(1, limit));
  }

  async search(
    context: Pick<SessionHistoryContext, "actorKey" | "projectKey">,
    query: string,
    maxResults = config.SESSION_SEARCH_MAX_RESULTS,
  ): Promise<SessionSearchHit[]> {
    const normalizedQuery = query.trim().toLowerCase().slice(0, config.SESSION_SEARCH_MAX_QUERY_CHARS);
    if (!normalizedQuery) {
      return [];
    }
    const queryTokens = normalizedQuery.split(/\s+/).filter(Boolean);
    const hits: SessionSearchHit[] = [];
    const files = await listJsonFiles(this.baseDir);
    for (const file of files) {
      const transcript = await readJsonFile<SessionTranscriptRecord>(file);
      if (!transcript) {
        continue;
      }
      if (
        transcript.actorKey !== context.actorKey ||
        transcript.projectKey !== context.projectKey
      ) {
        continue;
      }
      for (const turn of transcript.turns) {
        const candidates: Array<["user" | "assistant", string]> = [
          ["user", turn.userMessage],
          ["assistant", turn.assistantMessage],
        ];
        for (const [matchedField, text] of candidates) {
          const score = scoreText(text, normalizedQuery, queryTokens);
          if (score <= 0) {
            continue;
          }
          hits.push({
            matchedField,
            score,
            sessionId: transcript.sessionId,
            snippet: buildSnippet(text, normalizedQuery),
            timestamp: turn.timestamp,
            turnId: turn.id,
          });
        }
      }
    }
    return hits.sort((a, b) => b.score - a.score).slice(0, Math.max(1, maxResults));
  }

  private async readTranscript(context: SessionHistoryContext) {
    return await readJsonFile<SessionTranscriptRecord>(this.filePath(context));
  }

  private filePath(context: SessionHistoryContext) {
    return join(
      this.baseDir,
      `${context.actorKey}__${context.projectKey}__${context.sessionId}.json`,
    );
  }

  private async serializeWrite<T>(key: string, task: () => Promise<T>) {
    const previous = this.writeQueues.get(key) ?? Promise.resolve();
    const run = previous.catch(() => undefined).then(task);
    const next = run.then(
      () => undefined,
      () => undefined,
    );
    this.writeQueues.set(key, next);
    try {
      return await run;
    } finally {
      if (this.writeQueues.get(key) === next) {
        this.writeQueues.delete(key);
      }
    }
  }
}

const scoreText = (text: string, query: string, queryTokens: string[]) => {
  const normalized = text.toLowerCase();
  let score = 0;
  if (normalized.includes(query)) {
    score += Math.max(10, query.length);
  }
  for (const token of queryTokens) {
    if (token.length >= 2 && normalized.includes(token)) {
      score += 1;
    }
  }
  return score;
};

const buildSnippet = (text: string, query: string) => {
  const compact = text.replace(/\s+/g, " ").trim();
  const idx = compact.toLowerCase().indexOf(query);
  if (idx === -1) {
    return compact.length > 180 ? `${compact.slice(0, 177)}...` : compact;
  }
  const start = Math.max(0, idx - 60);
  const end = Math.min(compact.length, idx + query.length + 100);
  const snippet = compact.slice(start, end).trim();
  const prefix = start > 0 ? "..." : "";
  const suffix = end < compact.length ? "..." : "";
  return `${prefix}${snippet}${suffix}`;
};