next-tjwater-agent/src/learning/sessionStateStore.ts

import { basename, join } from "node:path";
import { lstat } from "node:fs/promises";

import { config } from "../config.js";
import { type SessionWorkspaceContext } from "../runtime/conversationWorkspace.js";
import { type SessionWorkspaceStore } from "../sessions/workspace-store.js";
import {
  atomicWriteJson,
  ensureDirectory,
  readJsonFile,
} from "../utils/fileStore.js";

export type SessionLearningState = {
  lastGatedTurn: number;
  lastReviewedTurn: number;
  sessionId: string;
  updatedAt: string;
};

export class SessionLearningStateStore {
  constructor(private readonly baseDir = config.SESSION_LEARNING_STATE_STORAGE_DIR, private readonly workspaces?: SessionWorkspaceStore) {}

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

  async read(context: SessionWorkspaceContext): Promise<SessionLearningState> {
    const sessionId = context.sessionId;
    if (this.workspaces) {
      await this.workspaces.ensure(context);
      await this.workspaces.validate(context, "learning.json");
    }
    let existing = await readJsonFile<SessionLearningState>(this.filePath(context));
    if (!existing && this.workspaces) {
      const legacy = this.legacyPath(sessionId);
      try {
        if ((await lstat(legacy)).isSymbolicLink()) throw new Error("legacy learning state must not be a symbolic link");
        existing = await readJsonFile<SessionLearningState>(legacy);
        if (existing?.sessionId !== sessionId) throw new Error("learning state session mismatch");
        await this.workspaces.writeRecord(context, "learning.json", existing);
        await this.workspaces.archiveLegacyRecord(context, "learning.json", legacy);
      } catch (error) { if ((error as NodeJS.ErrnoException).code !== "ENOENT") throw error; }
    }
    if (existing) {
      if (existing.sessionId !== sessionId) throw new Error("learning state session mismatch");
      return {
        lastGatedTurn: existing.lastGatedTurn,
        lastReviewedTurn: existing.lastReviewedTurn,
        sessionId: existing.sessionId,
        updatedAt: existing.updatedAt,
      };
    }
    return {
      lastGatedTurn: 0,
      lastReviewedTurn: 0,
      sessionId,
      updatedAt: new Date(0).toISOString(),
    };
  }

  async write(context: SessionWorkspaceContext, state: SessionLearningState) {
    if (state.sessionId !== context.sessionId) throw new Error("learning state session mismatch");
    const next = {
      ...state,
      updatedAt: new Date().toISOString(),
    };
    if (this.workspaces) await this.workspaces.writeRecord(context, "learning.json", next);
    else await atomicWriteJson(this.filePath(context), next);
  }

  async completeReview(context: SessionWorkspaceContext, reviewedTurnCount: number) {
    const current = await this.read(context);
    await this.write(context, {
      ...current,
      lastGatedTurn: Math.max(current.lastGatedTurn, reviewedTurnCount),
      lastReviewedTurn: reviewedTurnCount,
    });
  }

  async completeGate(context: SessionWorkspaceContext, gatedTurnCount: number) {
    const current = await this.read(context);
    await this.write(context, {
      ...current,
      lastGatedTurn: gatedTurnCount,
    });
  }

  private filePath(context: SessionWorkspaceContext) {
    return this.workspaces ? this.workspaces.path(context, "learning.json") : this.legacyPath(context.sessionId);
  }

  private legacyPath(sessionId: string) {
    if (basename(`${sessionId}.json`) !== `${sessionId}.json` || sessionId.includes("\\")) throw new Error("invalid learning session id");
    return join(this.baseDir, `${sessionId}.json`);
  }
}