gbrain As Memory

// src/gbrain-memory.ts
import { createMCPClient } from "@ai-sdk/mcp";
import type { HarnessMemoryBackend } from "@little-stack/harness";

/**
 * A memory kept in gbrain. Each memory is one gbrain source, reached with a token scoped to it.
 * On load, each active fact becomes one file, at the path the agent first gave it (kept in the
 * fact's provenance). On store, a new file is a `remember`, a deleted file a `forget`, and a
 * changed file both, so gbrain keeps the old version in its audit trail.
 */
export function gbrainMemory<TExtraBody>(options: {
  /** The gbrain server's MCP endpoint, e.g. http://127.0.0.1:3131/mcp. */
  url: string;
  /** This turn's tokens: one that reads the source, and one that writes it when the turn may. */
  tokens: (ctx: { extraBody?: TExtraBody | undefined }) => { read: string; write?: string };
  /** The gbrain fact kind for what this memory saves. */
  kind: "fact" | "preference";
}): HarnessMemoryBackend<TExtraBody> {
  return {
    async load({ extraBody }) {
      const { read } = options.tokens({ extraBody });
      const recalled = await call(options.url, read, "recall", { limit: 100 });
      const facts = (recalled.facts ?? []) as GbrainFact[];
      const files: Record<string, string> = {};
      const ids: Record<string, string> = {};
      for (const fact of facts) {
        const wanted = /(?:^|; )path: ([^;]+)/.exec(fact.provenance ?? "")?.[1]?.trim();
        const path =
          wanted && !(wanted in files) && /^[\w./-]+\.md$/.test(wanted) && !wanted.includes("..")
            ? wanted
            : `gbrain-fact-${fact.fact_id}.md`;
        const { description, body } = splitFact(fact.fact);
        files[path] = `---\ndescription: ${JSON.stringify(description)}\n---\n\n${body}\n`;
        ids[path] = fact.fact_id;
      }
      // The cursor comes back to store, so it knows which fact each file was.
      return { files, cursor: { ids } };
    },
    async store({ changes, extraBody, session, cursor }) {
      const ids = (cursor as { ids?: Record<string, string> } | undefined)?.ids ?? {};
      const written = Object.entries({ ...changes.created, ...changes.updated });
      const gone = [...changes.deleted, ...Object.keys(changes.updated)].flatMap((path) => ids[path] ?? []);
      if (written.length === 0 && gone.length === 0) return;
      const { write } = options.tokens({ extraBody });
      if (write === undefined) throw new Error("gbrainMemory: this turn has no write token for the memory it changed");
      for (const id of gone) {
        await call(options.url, write, "forget", { id, reason: `changed or deleted in Little Harness session ${session.id}` });
      }
      for (const [path, bytes] of written) {
        const fact = factText(new TextDecoder().decode(bytes));
        if (fact.length === 0) continue;
        await call(options.url, write, "remember", {
          fact,
          provenance: `path: ${path}; Little Harness session ${session.id}`.slice(0, 500),
          kind: options.kind,
          infer_entity: false,
        });
      }
    },
  };
}

type GbrainFact = { fact_id: string; fact: string; provenance?: string; valid_from?: string; created_at?: string };

/** A memory file as one gbrain fact: its description, then its body, without the frontmatter. */
function factText(text: string): string {
  const match = /^---\r?\n([\s\S]*?)\r?\n---[ \t]*(?:\r?\n|$)([\s\S]*)$/.exec(text);
  const body = (match ? match[2]! : text).trim();
  const raw = match ? /^description:[ \t]*(.*)$/m.exec(match[1]!)?.[1]?.trim() : undefined;
  const description = raw?.startsWith('"') ? (JSON.parse(raw) as string) : raw;
  if (!description) return body;
  return body.length === 0 || body.startsWith(description) ? body || description : `${description}\n\n${body}`;
}

/** The first paragraph of a fact as the file's description; the whole fact as its body. */
function splitFact(fact: string): { description: string; body: string } {
  const first = fact.trim().split(/\n\s*\n/)[0] ?? fact;
  const description = first.replace(/\s+/g, " ").trim();
  return { description: description.length <= 600 ? description : `${description.slice(0, 600)}…`, body: fact.trim() };
}

async function call(url: string, token: string, name: string, args: Record<string, unknown>): Promise<Record<string, any>> {
  const client = await createMCPClient({ transport: { type: "http", url, headers: { Authorization: `Bearer ${token}` } } });
  try {
    const result = await client.callTool({ name, arguments: args });
    const text = (result.content as { type: string; text?: string }[]).find((part) => part.type === "text")?.text ?? "{}";
    const parsed = JSON.parse(text) as Record<string, any>;
    if (result.isError || parsed.error) throw new Error(`gbrain ${name}: ${parsed.code ?? parsed.error}: ${parsed.message ?? text}`);
    return parsed;
  } finally {
    await client.close();
  }
}

import { createHarness, localHost, memory } from "@little-stack/harness";
import { gbrainMemory } from "./gbrain-memory";

type DeskUser = { userId: string; role: "officer" | "underwriter" };

const url = process.env.GBRAIN_URL ?? "http://127.0.0.1:3131/mcp";
const tokens = await loadGbrainTokens(); // from your secrets store

const harness = createHarness<typeof tools, DeskUser>({
  host: localHost({ dataDir: ".little-harness" }),
  model,
  tools,
  memory: [
    memory<DeskUser>({
      harnessDir: "/persistent/company",
      description: "The bank's shared lending knowledge. Only underwriters may change it.",
      canWrite: ({ extraBody }) => extraBody?.role === "underwriter",
      backend: gbrainMemory<DeskUser>({
        url,
        kind: "fact",
        tokens: ({ extraBody }) => ({
          read: tokens["company-reader"]!,
          ...(extraBody?.role === "underwriter" ? { write: tokens["company-writer"]! } : {}),
        }),
      }),
    }),
    memory<DeskUser>({
      harnessDir: "/persistent/me",
      description: "This person's own preferences, such as how they like answers. Never lending rules.",
      backend: gbrainMemory<DeskUser>({
        url,
        kind: "preference",
        tokens: ({ extraBody }) => {
          const token = extraBody && tokens[`user-${extraBody.userId}`];
          if (!token) throw new Error("This turn needs a userId with a gbrain token.");
          return { read: token, write: token };
        },
      }),
    }),
  ],
});