@@ -0,0 +1,191 @@
|
||||
// One root session's memory on disk: main/ and tree/ day logs (append-only), view.json, and a lock.
|
||||
import {
|
||||
appendFile,
|
||||
mkdir,
|
||||
open,
|
||||
readFile,
|
||||
readdir,
|
||||
rename,
|
||||
rm,
|
||||
writeFile,
|
||||
} from "node:fs/promises";
|
||||
import { join } from "node:path";
|
||||
import { z } from "zod";
|
||||
import { bytes, key, pieces, type Kind, type Line, type View } from "./optchat-core.js";
|
||||
|
||||
const kind = z.enum(["user", "agent", "tool", "echo"]);
|
||||
const record = z.object({
|
||||
i: z.number().int().nonnegative(),
|
||||
kind,
|
||||
text: z.string(),
|
||||
size: z.number(),
|
||||
date: z.string(),
|
||||
/** The history part this message came from; it is never logged twice. */
|
||||
src: z.string(),
|
||||
});
|
||||
const node = z.object({
|
||||
l: z.number().int().nonnegative(),
|
||||
i: z.number().int().nonnegative(),
|
||||
text: z.string(),
|
||||
size: z.number(),
|
||||
});
|
||||
const line = z.tuple([z.number().int().nonnegative(), z.number().int().nonnegative()]);
|
||||
const view = z.object({ lines: z.array(line), batching: z.boolean() });
|
||||
export type Message = z.infer<typeof record>;
|
||||
|
||||
async function lines(directory: string) {
|
||||
const names = await readdir(directory).catch(() => []);
|
||||
const out: string[] = [];
|
||||
for (const name of names.filter((name) => name.endsWith(".jsonl")).sort())
|
||||
out.push(...(await readFile(join(directory, name), "utf8")).split("\n"));
|
||||
// A torn last line from a crash is skipped; its message or node is redone.
|
||||
return out
|
||||
.filter((text) => text.trim())
|
||||
.flatMap((text) => {
|
||||
try {
|
||||
return [JSON.parse(text)];
|
||||
} catch {
|
||||
return [];
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
async function lock(file: string) {
|
||||
for (let attempt = 0; attempt < 2; attempt++) {
|
||||
try {
|
||||
const handle = await open(file, "wx");
|
||||
await handle.writeFile(String(process.pid));
|
||||
await handle.close();
|
||||
return;
|
||||
} catch (error) {
|
||||
if (!(error instanceof Error && "code" in error && error.code === "EEXIST"))
|
||||
throw error;
|
||||
}
|
||||
const pid = Number(await readFile(file, "utf8").catch(() => ""));
|
||||
if (Number.isInteger(pid) && pid > 0 && pid !== process.pid) {
|
||||
try {
|
||||
process.kill(pid, 0);
|
||||
throw new Error(`optchat memory is in use by process ${pid}`);
|
||||
} catch (error) {
|
||||
if (!(error instanceof Error && "code" in error && error.code === "ESRCH"))
|
||||
throw error;
|
||||
}
|
||||
} else if (pid === process.pid) throw new Error("optchat memory is already open");
|
||||
await rm(file, { force: true });
|
||||
}
|
||||
throw new Error("optchat memory lock is contended");
|
||||
}
|
||||
|
||||
export class Memory {
|
||||
readonly messages: Message[] = [];
|
||||
readonly #tree = new Map<string, string>();
|
||||
readonly #sources = new Set<string>();
|
||||
#writes: Promise<void> = Promise.resolve();
|
||||
view: View = { lines: [], batching: false };
|
||||
|
||||
private constructor(readonly directory: string) {}
|
||||
|
||||
/** Takes the lock and loads the log, tree and saved view. */
|
||||
static async open(directory: string) {
|
||||
await mkdir(join(directory, "main"), { recursive: true });
|
||||
await mkdir(join(directory, "tree"), { recursive: true });
|
||||
await lock(join(directory, "lock"));
|
||||
const memory = new Memory(directory);
|
||||
try {
|
||||
for (const value of await lines(join(directory, "main"))) {
|
||||
const parsed = record.safeParse(value);
|
||||
if (!parsed.success || parsed.data.i !== memory.messages.length) continue;
|
||||
memory.messages.push(parsed.data);
|
||||
memory.#sources.add(parsed.data.src);
|
||||
}
|
||||
for (const value of await lines(join(directory, "tree"))) {
|
||||
const parsed = node.safeParse(value);
|
||||
if (parsed.success && (parsed.data.l > 0 || parsed.data.i < memory.messages.length))
|
||||
memory.#tree.set(key([parsed.data.l, parsed.data.i]), parsed.data.text);
|
||||
}
|
||||
const saved = await readFile(join(directory, "view.json"), "utf8").catch(
|
||||
() => undefined,
|
||||
);
|
||||
if (saved !== undefined) memory.view = view.parse(JSON.parse(saved));
|
||||
// Messages logged after the last saved view (a crash in between) are appended in order.
|
||||
const covered = memory.view.lines.reduce((sum, [l]) => sum + 2 ** l, 0);
|
||||
for (let i = covered; i < memory.messages.length; i++) memory.view.lines.push([0, i]);
|
||||
} catch (error) {
|
||||
await memory.close();
|
||||
throw error;
|
||||
}
|
||||
return memory;
|
||||
}
|
||||
|
||||
has(source: string) {
|
||||
return this.#sources.has(source);
|
||||
}
|
||||
|
||||
text(line: Line) {
|
||||
return this.#tree.get(key(line));
|
||||
}
|
||||
|
||||
nodes() {
|
||||
return [...this.#tree.keys()].map((name): Line => {
|
||||
const [l, i] = name.split(":").map(Number);
|
||||
return [l!, i!];
|
||||
});
|
||||
}
|
||||
|
||||
/** Logs one history part as one or more messages, in order; returns their ids. */
|
||||
log(kind: Kind, text: string, source: string) {
|
||||
if (this.#sources.has(source)) return [];
|
||||
this.#sources.add(source);
|
||||
const date = new Date().toISOString();
|
||||
return pieces(kind, text).map((text, k) => {
|
||||
const message: Message = {
|
||||
i: this.messages.length,
|
||||
kind,
|
||||
text,
|
||||
size: bytes(text),
|
||||
date,
|
||||
src: k === 0 ? source : `${source}#${k}`,
|
||||
};
|
||||
this.messages.push(message);
|
||||
this.#append("main", date, message);
|
||||
return message.i;
|
||||
});
|
||||
}
|
||||
|
||||
node([l, i]: Line, text: string) {
|
||||
this.#tree.set(key([l, i]), text);
|
||||
this.#append("tree", new Date().toISOString(), { l, i, text, size: bytes(text) });
|
||||
}
|
||||
|
||||
saveView() {
|
||||
const snapshot = JSON.stringify(this.view);
|
||||
this.#write(async () => {
|
||||
const file = join(this.directory, "view.json");
|
||||
await writeFile(`${file}.tmp`, snapshot);
|
||||
await rename(`${file}.tmp`, file);
|
||||
});
|
||||
}
|
||||
|
||||
/** Resolves once every write so far reached the file. */
|
||||
flushed() {
|
||||
return this.#writes;
|
||||
}
|
||||
|
||||
async close() {
|
||||
await this.#writes;
|
||||
await rm(join(this.directory, "lock"), { force: true });
|
||||
}
|
||||
|
||||
#append(directory: "main" | "tree", date: string, value: Message | z.infer<typeof node>) {
|
||||
const file = join(this.directory, directory, `${date.slice(0, 10)}.jsonl`);
|
||||
const text = JSON.stringify(value) + "\n";
|
||||
this.#write(() => appendFile(file, text));
|
||||
}
|
||||
|
||||
#write(write: () => Promise<void>) {
|
||||
// One writer, in order; a failed write is reported and later writes still run.
|
||||
this.#writes = this.#writes
|
||||
.then(write)
|
||||
.catch((error) => console.error("optchat write failed", error));
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user