ai_ops/src/agents/state-context.ts

import { mkdir, readFile } from "node:fs/promises";
import { dirname, resolve } from "node:path";
import { writeUtf8FileAtomic } from "./file-persistence.js";
import { deepCloneJson, isRecord, type JsonObject, type JsonValue } from "./types.js";

export type SessionHistoryEntry = {
  nodeId: string;
  event: string;
  timestamp: string;
  data?: JsonObject;
};

export type StoredSessionState = {
  flags: Record<string, boolean>;
  metadata: JsonObject;
  history: SessionHistoryEntry[];
};

export type NodeHandoff = {
  nodeId: string;
  fromNodeId?: string;
  payload: JsonObject;
  createdAt: string;
};

export type NodeExecutionContext = {
  sessionId: string;
  nodeId: string;
  handoff: NodeHandoff;
  state: StoredSessionState;
};

const DEFAULT_STATE: StoredSessionState = {
  flags: {},
  metadata: {},
  history: [],
};

function toSessionDirectory(rootDirectory: string, sessionId: string): string {
  return resolve(rootDirectory, sessionId);
}

function toStatePath(rootDirectory: string, sessionId: string): string {
  return resolve(toSessionDirectory(rootDirectory, sessionId), "state.json");
}

function toHandoffPath(rootDirectory: string, sessionId: string, nodeId: string): string {
  return resolve(toSessionDirectory(rootDirectory, sessionId), "handoffs", `${nodeId}.json`);
}

function toJsonObject(value: unknown, errorMessage: string): JsonObject {
  if (!isRecord(value)) {
    throw new Error(errorMessage);
  }

  return value as JsonObject;
}

function toStoredSessionState(value: unknown): StoredSessionState {
  if (!isRecord(value)) {
    throw new Error("Stored state file is malformed.");
  }

  const flagsValue = value.flags;
  if (!isRecord(flagsValue)) {
    throw new Error("Stored state.flags is malformed.");
  }

  const flags: Record<string, boolean> = {};
  for (const [key, flagValue] of Object.entries(flagsValue)) {
    if (typeof flagValue !== "boolean") {
      throw new Error(`Stored state flag \"${key}\" must be a boolean.`);
    }
    flags[key] = flagValue;
  }

  const metadata = toJsonObject(value.metadata, "Stored state.metadata is malformed.");
  const historyValue = value.history;
  if (!Array.isArray(historyValue)) {
    throw new Error("Stored state.history is malformed.");
  }

  const history: SessionHistoryEntry[] = [];
  for (const entry of historyValue) {
    if (!isRecord(entry)) {
      throw new Error("Stored state.history entry is malformed.");
    }

    const nodeId = entry.nodeId;
    const event = entry.event;
    const timestamp = entry.timestamp;

    if (typeof nodeId !== "string" || nodeId.trim().length === 0) {
      throw new Error("Stored state.history entry nodeId is malformed.");
    }
    if (typeof event !== "string" || event.trim().length === 0) {
      throw new Error("Stored state.history entry event is malformed.");
    }
    if (typeof timestamp !== "string" || timestamp.trim().length === 0) {
      throw new Error("Stored state.history entry timestamp is malformed.");
    }

    const data = entry.data === undefined ? undefined : toJsonObject(entry.data, "Stored state.history entry data is malformed.");

    history.push({
      nodeId,
      event,
      timestamp,
      ...(data ? { data } : {}),
    });
  }

  return {
    flags,
    metadata,
    history,
  };
}

function toNodeHandoff(value: unknown): NodeHandoff {
  if (!isRecord(value)) {
    throw new Error("Stored handoff file is malformed.");
  }

  const nodeId = value.nodeId;
  const createdAt = value.createdAt;
  const payload = value.payload;

  if (typeof nodeId !== "string" || nodeId.trim().length === 0) {
    throw new Error("Stored handoff nodeId is malformed.");
  }
  if (typeof createdAt !== "string" || createdAt.trim().length === 0) {
    throw new Error("Stored handoff createdAt is malformed.");
  }
  if (!isRecord(payload)) {
    throw new Error("Stored handoff payload is malformed.");
  }

  const fromNodeId = value.fromNodeId;
  if (fromNodeId !== undefined && (typeof fromNodeId !== "string" || fromNodeId.trim().length === 0)) {
    throw new Error("Stored handoff fromNodeId is malformed.");
  }

  return {
    nodeId,
    ...(typeof fromNodeId === "string" ? { fromNodeId } : {}),
    payload: payload as JsonObject,
    createdAt,
  };
}

export class FileSystemStateContextManager {
  private readonly rootDirectory: string;

  constructor(input: {
    rootDirectory: string;
  }) {
    this.rootDirectory = resolve(input.rootDirectory);
  }

  getRootDirectory(): string {
    return this.rootDirectory;
  }

  getSessionStatePath(sessionId: string): string {
    return toStatePath(this.rootDirectory, sessionId);
  }

  async initializeSession(
    sessionId: string,
    initialState: Partial<StoredSessionState> = {},
  ): Promise<StoredSessionState> {
    const state: StoredSessionState = {
      flags: { ...(initialState.flags ?? {}) },
      metadata: deepCloneJson((initialState.metadata ?? {}) as JsonValue) as JsonObject,
      history: [...(initialState.history ?? [])],
    };

    await mkdir(toSessionDirectory(this.rootDirectory, sessionId), { recursive: true });
    await this.writeState(sessionId, state);

    return state;
  }

  async readState(sessionId: string): Promise<StoredSessionState> {
    const path = toStatePath(this.rootDirectory, sessionId);

    try {
      const content = await readFile(path, "utf8");
      const parsed = JSON.parse(content) as unknown;
      return toStoredSessionState(parsed);
    } catch (error) {
      const code = (error as NodeJS.ErrnoException).code;
      if (code === "ENOENT") {
        return deepCloneJson(DEFAULT_STATE as JsonValue) as StoredSessionState;
      }
      throw error;
    }
  }

  async writeState(sessionId: string, state: StoredSessionState): Promise<void> {
    const path = toStatePath(this.rootDirectory, sessionId);
    await mkdir(dirname(path), { recursive: true });
    await writeUtf8FileAtomic(path, `${JSON.stringify(state, null, 2)}\n`);
  }

  async patchState(
    sessionId: string,
    patch: {
      flags?: Record<string, boolean>;
      metadata?: JsonObject;
      historyEvents?: SessionHistoryEntry[];
    },
  ): Promise<StoredSessionState> {
    const current = await this.readState(sessionId);

    if (patch.flags) {
      Object.assign(current.flags, patch.flags);
    }
    if (patch.metadata) {
      Object.assign(current.metadata, patch.metadata);
    }
    if (patch.historyEvents && patch.historyEvents.length > 0) {
      current.history.push(...patch.historyEvents);
    }

    await this.writeState(sessionId, current);
    return current;
  }

  async writeHandoff(
    sessionId: string,
    handoff: {
      nodeId: string;
      fromNodeId?: string;
      payload: JsonObject;
    },
  ): Promise<NodeHandoff> {
    const nodeHandoff: NodeHandoff = {
      nodeId: handoff.nodeId,
      ...(handoff.fromNodeId ? { fromNodeId: handoff.fromNodeId } : {}),
      payload: deepCloneJson(handoff.payload),
      createdAt: new Date().toISOString(),
    };

    const path = toHandoffPath(this.rootDirectory, sessionId, handoff.nodeId);
    await mkdir(dirname(path), { recursive: true });
    await writeUtf8FileAtomic(path, `${JSON.stringify(nodeHandoff, null, 2)}\n`);
    return nodeHandoff;
  }

  async readHandoff(sessionId: string, nodeId: string): Promise<NodeHandoff | undefined> {
    const path = toHandoffPath(this.rootDirectory, sessionId, nodeId);

    try {
      const content = await readFile(path, "utf8");
      const parsed = JSON.parse(content) as unknown;
      return toNodeHandoff(parsed);
    } catch (error) {
      if ((error as NodeJS.ErrnoException).code === "ENOENT") {
        return undefined;
      }
      throw error;
    }
  }

  async buildFreshNodeContext(sessionId: string, nodeId: string): Promise<NodeExecutionContext> {
    const handoff = await this.readHandoff(sessionId, nodeId);
    if (!handoff) {
      throw new Error(`No handoff exists for node \"${nodeId}\" in session \"${sessionId}\".`);
    }

    const state = await this.readState(sessionId);

    return {
      sessionId,
      nodeId,
      handoff: {
        nodeId: handoff.nodeId,
        ...(handoff.fromNodeId ? { fromNodeId: handoff.fromNodeId } : {}),
        payload: deepCloneJson(handoff.payload),
        createdAt: handoff.createdAt,
      },
      state: {
        flags: { ...state.flags },
        metadata: deepCloneJson(state.metadata),
        history: state.history.map((entry) => ({
          nodeId: entry.nodeId,
          event: entry.event,
          timestamp: entry.timestamp,
          ...(entry.data ? { data: deepCloneJson(entry.data) } : {}),
        })),
      },
    };
  }
}