@mainahq/core 0.2.0

package/src/context/__tests__/treesitter.test.ts

@@ -0,0 +1,229 @@
+import { describe, expect, test } from "bun:test";
+import {
+	extractEntities,
+	extractExports,
+	extractImports,
+	parseFile,
+} from "../treesitter";
+
+describe("extractImports", () => {
+	test("parses named imports correctly", () => {
+		const content = `import { Foo, Bar, Baz } from "some-module";`;
+		const imports = extractImports(content);
+		expect(imports).toHaveLength(1);
+		expect(imports[0]?.source).toBe("some-module");
+		expect(imports[0]?.specifiers).toEqual(["Foo", "Bar", "Baz"]);
+		expect(imports[0]?.isDefault).toBe(false);
+	});
+
+	test("parses default imports", () => {
+		const content = `import MyDefault from "./my-file";`;
+		const imports = extractImports(content);
+		expect(imports).toHaveLength(1);
+		expect(imports[0]?.source).toBe("./my-file");
+		expect(imports[0]?.specifiers).toEqual(["MyDefault"]);
+		expect(imports[0]?.isDefault).toBe(true);
+	});
+
+	test("parses namespace imports (import * as X)", () => {
+		const content = `import * as Utils from "../utils";`;
+		const imports = extractImports(content);
+		expect(imports).toHaveLength(1);
+		expect(imports[0]?.source).toBe("../utils");
+		expect(imports[0]?.specifiers).toEqual(["Utils"]);
+		expect(imports[0]?.isDefault).toBe(false);
+	});
+
+	test("parses type imports", () => {
+		const content = `import type { SomeType, AnotherType } from "types-module";`;
+		const imports = extractImports(content);
+		expect(imports).toHaveLength(1);
+		expect(imports[0]?.source).toBe("types-module");
+		expect(imports[0]?.specifiers).toEqual(["SomeType", "AnotherType"]);
+		expect(imports[0]?.isDefault).toBe(false);
+	});
+
+	test("parses multiple imports", () => {
+		const content = [
+			`import { A } from "mod-a";`,
+			`import B from "mod-b";`,
+			`import type { C } from "mod-c";`,
+		].join("\n");
+		const imports = extractImports(content);
+		expect(imports).toHaveLength(3);
+		expect(imports[0]?.source).toBe("mod-a");
+		expect(imports[1]?.source).toBe("mod-b");
+		expect(imports[2]?.source).toBe("mod-c");
+	});
+
+	test("returns empty array when no imports", () => {
+		const content = `const x = 1;`;
+		const imports = extractImports(content);
+		expect(imports).toHaveLength(0);
+	});
+});
+
+describe("extractExports", () => {
+	test("finds exported functions", () => {
+		const content = [
+			`export function doSomething() {}`,
+			`export async function fetchData() {}`,
+		].join("\n");
+		const exports = extractExports(content);
+		const names = exports.map((e) => e.name);
+		expect(names).toContain("doSomething");
+		expect(names).toContain("fetchData");
+	});
+
+	test("finds exported classes and interfaces", () => {
+		const content = [
+			`export class MyService {}`,
+			`export interface MyInterface {}`,
+		].join("\n");
+		const exports = extractExports(content);
+		const names = exports.map((e) => e.name);
+		expect(names).toContain("MyService");
+		expect(names).toContain("MyInterface");
+
+		const classExport = exports.find((e) => e.name === "MyService");
+		expect(classExport?.kind).toBe("class");
+
+		const interfaceExport = exports.find((e) => e.name === "MyInterface");
+		expect(interfaceExport?.kind).toBe("interface");
+	});
+
+	test("finds exported type aliases and const", () => {
+		const content = [
+			`export type MyAlias = string | number;`,
+			`export const MY_CONST = 42;`,
+		].join("\n");
+		const exports = extractExports(content);
+		const names = exports.map((e) => e.name);
+		expect(names).toContain("MyAlias");
+		expect(names).toContain("MY_CONST");
+	});
+
+	test("finds export default", () => {
+		const content = `export default class DefaultClass {}`;
+		const exports = extractExports(content);
+		const defaultExport = exports.find((e) => e.name === "default");
+		expect(defaultExport).toBeDefined();
+	});
+
+	test("finds named re-exports: export { X, Y }", () => {
+		const content = `export { Foo, Bar } from "./somewhere";`;
+		const exports = extractExports(content);
+		const names = exports.map((e) => e.name);
+		expect(names).toContain("Foo");
+		expect(names).toContain("Bar");
+	});
+
+	test("returns empty array when no exports", () => {
+		const content = `const x = 1;`;
+		const exports = extractExports(content);
+		expect(exports).toHaveLength(0);
+	});
+});
+
+describe("extractEntities", () => {
+	test("finds function declarations with line numbers", () => {
+		const content = [
+			``,
+			`function greet(name: string): string {`,
+			`  return "hello " + name;`,
+			`}`,
+		].join("\n");
+		const entities = extractEntities(content);
+		const fn = entities.find((e) => e.name === "greet");
+		expect(fn).toBeDefined();
+		expect(fn?.kind).toBe("function");
+		expect(fn?.startLine).toBe(2);
+	});
+
+	test("finds class declarations", () => {
+		const content = [`class MyClass {`, `  constructor() {}`, `}`].join("\n");
+		const entities = extractEntities(content);
+		const cls = entities.find((e) => e.name === "MyClass");
+		expect(cls).toBeDefined();
+		expect(cls?.kind).toBe("class");
+		expect(cls?.startLine).toBe(1);
+	});
+
+	test("finds interface declarations", () => {
+		const content = `interface Shape {\n  area(): number;\n}`;
+		const entities = extractEntities(content);
+		const iface = entities.find((e) => e.name === "Shape");
+		expect(iface).toBeDefined();
+		expect(iface?.kind).toBe("interface");
+	});
+
+	test("finds type alias declarations", () => {
+		const content = `type ID = string | number;`;
+		const entities = extractEntities(content);
+		const typeAlias = entities.find((e) => e.name === "ID");
+		expect(typeAlias).toBeDefined();
+		expect(typeAlias?.kind).toBe("type");
+	});
+
+	test("finds top-level const declarations", () => {
+		const content = `const MAX_SIZE = 100;`;
+		const entities = extractEntities(content);
+		const variable = entities.find((e) => e.name === "MAX_SIZE");
+		expect(variable).toBeDefined();
+		expect(variable?.kind).toBe("variable");
+	});
+
+	test("endLine is set (at least equal to startLine)", () => {
+		const content = `function hello() {}`;
+		const entities = extractEntities(content);
+		const fn = entities.find((e) => e.name === "hello");
+		expect(fn).toBeDefined();
+		expect(fn?.endLine ?? 0).toBeGreaterThanOrEqual(fn?.startLine ?? 0);
+	});
+
+	test("returns empty array when no entities", () => {
+		const content = `// just a comment`;
+		const entities = extractEntities(content);
+		expect(entities).toHaveLength(0);
+	});
+});
+
+describe("parseFile", () => {
+	test("reads an actual TS file and returns imports/exports/entities", async () => {
+		const filePath = `${process.cwd()}/packages/core/src/git/index.ts`;
+		const result = await parseFile(filePath);
+
+		expect(result).toHaveProperty("imports");
+		expect(result).toHaveProperty("exports");
+		expect(result).toHaveProperty("entities");
+
+		expect(Array.isArray(result.imports)).toBe(true);
+		expect(Array.isArray(result.exports)).toBe(true);
+		expect(Array.isArray(result.entities)).toBe(true);
+
+		// git/index.ts exports several functions
+		const exportNames = result.exports.map((e) => e.name);
+		expect(exportNames).toContain("getCurrentBranch");
+		expect(exportNames).toContain("getRepoRoot");
+		expect(exportNames).toContain("getRecentCommits");
+
+		// git/index.ts has function/interface entities
+		const entityNames = result.entities.map((e) => e.name);
+		expect(entityNames).toContain("Commit");
+		expect(
+			entityNames.some((n) =>
+				["exec", "getCurrentBranch", "getRepoRoot"].includes(n),
+			),
+		).toBe(true);
+	});
+
+	test("returns empty arrays for a file with no imports/exports/entities", async () => {
+		// Write a minimal temp file and parse it
+		const tmpPath = "/tmp/maina-test-empty.ts";
+		await Bun.write(tmpPath, "// empty file\n");
+		const result = await parseFile(tmpPath);
+		expect(result.imports).toHaveLength(0);
+		expect(result.exports).toHaveLength(0);
+		expect(result.entities).toHaveLength(0);
+	});
+});

package/src/context/__tests__/working.test.ts

@@ -0,0 +1,236 @@
+import {
+	afterAll,
+	afterEach,
+	beforeEach,
+	describe,
+	expect,
+	it,
+	mock,
+} from "bun:test";
+import { mkdirSync, rmSync } from "node:fs";
+import { tmpdir } from "node:os";
+import { join, resolve } from "node:path";
+
+// We mock getCurrentBranch before importing the module under test.
+// Use the absolute path so it matches how working.ts resolves its import.
+const GIT_MODULE = resolve(import.meta.dir, "../../git/index");
+
+// Import real module first so we can spread its exports
+const realGit = await import("../../git/index");
+
+const mockGetCurrentBranch = mock(async (_cwd?: string) => "main");
+
+mock.module(GIT_MODULE, () => ({
+	...realGit,
+	getCurrentBranch: mockGetCurrentBranch,
+}));
+
+afterAll(() => {
+	mock.restore();
+});
+
+// Import after mocking
+const {
+	loadWorkingContext,
+	saveWorkingContext,
+	trackFile,
+	setVerificationResult,
+	resetWorkingContext,
+	assembleWorkingText,
+} = await import("../working");
+
+// ── helpers ──────────────────────────────────────────────────────────────────
+
+function makeTempDir(): string {
+	const dir = join(
+		tmpdir(),
+		`maina-working-test-${Date.now()}-${Math.random().toString(36).slice(2)}`,
+	);
+	mkdirSync(dir, { recursive: true });
+	return dir;
+}
+
+// ── tests ─────────────────────────────────────────────────────────────────────
+
+describe("loadWorkingContext", () => {
+	let mainaDir: string;
+	let repoRoot: string;
+
+	beforeEach(() => {
+		mainaDir = makeTempDir();
+		repoRoot = makeTempDir();
+		mockGetCurrentBranch.mockImplementation(async () => "main");
+	});
+
+	afterEach(() => {
+		rmSync(mainaDir, { recursive: true, force: true });
+		rmSync(repoRoot, { recursive: true, force: true });
+	});
+
+	it("returns empty context when no file exists", async () => {
+		const ctx = await loadWorkingContext(mainaDir, repoRoot);
+		expect(ctx.branch).toBe("main");
+		expect(ctx.planContent).toBeNull();
+		expect(ctx.touchedFiles).toEqual([]);
+		expect(ctx.lastVerification).toBeNull();
+		expect(typeof ctx.updatedAt).toBe("string");
+	});
+
+	it("round-trips save and load", async () => {
+		const ctx = await loadWorkingContext(mainaDir, repoRoot);
+		ctx.touchedFiles = ["src/foo.ts", "src/bar.ts"];
+		ctx.branch = "main";
+		await saveWorkingContext(mainaDir, ctx);
+
+		const loaded = await loadWorkingContext(mainaDir, repoRoot);
+		expect(loaded.touchedFiles).toEqual(["src/foo.ts", "src/bar.ts"]);
+		expect(loaded.branch).toBe("main");
+	});
+
+	it("resets context when branch has changed since last save", async () => {
+		// Save a context recorded on branch "feature"
+		const ctx = resetWorkingContext(mainaDir);
+		ctx.branch = "feature";
+		ctx.touchedFiles = ["old-file.ts"];
+		await saveWorkingContext(mainaDir, ctx);
+
+		// Now git says we are on "main"
+		mockGetCurrentBranch.mockImplementation(async () => "main");
+
+		const loaded = await loadWorkingContext(mainaDir, repoRoot);
+		expect(loaded.branch).toBe("main");
+		expect(loaded.touchedFiles).toEqual([]);
+	});
+
+	it("loads PLAN.md content from repoRoot if it exists", async () => {
+		await Bun.write(join(repoRoot, "PLAN.md"), "# My Plan\nDo the thing.");
+		const ctx = await loadWorkingContext(mainaDir, repoRoot);
+		expect(ctx.planContent).toBe("# My Plan\nDo the thing.");
+	});
+
+	it("keeps planContent null when PLAN.md is absent", async () => {
+		const ctx = await loadWorkingContext(mainaDir, repoRoot);
+		expect(ctx.planContent).toBeNull();
+	});
+});
+
+describe("trackFile", () => {
+	let mainaDir: string;
+	let repoRoot: string;
+
+	beforeEach(() => {
+		mainaDir = makeTempDir();
+		repoRoot = makeTempDir();
+		mockGetCurrentBranch.mockImplementation(async () => "main");
+	});
+
+	afterEach(() => {
+		rmSync(mainaDir, { recursive: true, force: true });
+		rmSync(repoRoot, { recursive: true, force: true });
+	});
+
+	it("adds a file to touchedFiles and persists", async () => {
+		const ctx = await trackFile(mainaDir, repoRoot, "src/hello.ts");
+		expect(ctx.touchedFiles).toContain("src/hello.ts");
+
+		const reloaded = await loadWorkingContext(mainaDir, repoRoot);
+		expect(reloaded.touchedFiles).toContain("src/hello.ts");
+	});
+
+	it("does not duplicate files already tracked", async () => {
+		await trackFile(mainaDir, repoRoot, "src/hello.ts");
+		const ctx = await trackFile(mainaDir, repoRoot, "src/hello.ts");
+		const occurrences = ctx.touchedFiles.filter((f) => f === "src/hello.ts");
+		expect(occurrences.length).toBe(1);
+	});
+});
+
+describe("setVerificationResult", () => {
+	let mainaDir: string;
+	let repoRoot: string;
+
+	beforeEach(() => {
+		mainaDir = makeTempDir();
+		repoRoot = makeTempDir();
+		mockGetCurrentBranch.mockImplementation(async () => "main");
+	});
+
+	afterEach(() => {
+		rmSync(mainaDir, { recursive: true, force: true });
+		rmSync(repoRoot, { recursive: true, force: true });
+	});
+
+	it("stores the last verification result and persists it", async () => {
+		const result = {
+			passed: true,
+			checks: [{ name: "lint", passed: true, output: "ok" }],
+			timestamp: new Date().toISOString(),
+		};
+
+		const ctx = await setVerificationResult(mainaDir, repoRoot, result);
+		expect(ctx.lastVerification).toEqual(result);
+
+		const reloaded = await loadWorkingContext(mainaDir, repoRoot);
+		expect(reloaded.lastVerification).toEqual(result);
+	});
+});
+
+describe("assembleWorkingText", () => {
+	it("includes branch name and touched files count", () => {
+		const ctx = {
+			branch: "feature/foo",
+			planContent: null,
+			workflowContext: null,
+			touchedFiles: ["a.ts", "b.ts"],
+			lastVerification: null,
+			updatedAt: new Date().toISOString(),
+		};
+		const text = assembleWorkingText(ctx);
+		expect(text).toContain("feature/foo");
+		expect(text).toContain("2");
+	});
+
+	it("includes individual touched file paths", () => {
+		const ctx = {
+			branch: "main",
+			planContent: null,
+			workflowContext: null,
+			touchedFiles: ["src/foo.ts", "src/bar.ts"],
+			lastVerification: null,
+			updatedAt: new Date().toISOString(),
+		};
+		const text = assembleWorkingText(ctx);
+		expect(text).toContain("src/foo.ts");
+		expect(text).toContain("src/bar.ts");
+	});
+
+	it("includes PLAN.md content when present", () => {
+		const ctx = {
+			branch: "main",
+			planContent: "# Plan\nDo stuff.",
+			workflowContext: null,
+			touchedFiles: [],
+			lastVerification: null,
+			updatedAt: new Date().toISOString(),
+		};
+		const text = assembleWorkingText(ctx);
+		expect(text).toContain("# Plan\nDo stuff.");
+	});
+
+	it("includes verification status when present", () => {
+		const ctx = {
+			branch: "main",
+			planContent: null,
+			workflowContext: null,
+			touchedFiles: [],
+			lastVerification: {
+				passed: false,
+				checks: [{ name: "typecheck", passed: false, output: "3 errors" }],
+				timestamp: new Date().toISOString(),
+			},
+			updatedAt: new Date().toISOString(),
+		};
+		const text = assembleWorkingText(ctx);
+		expect(text.toLowerCase()).toContain("fail");
+	});
+});

package/src/context/budget.ts

@@ -0,0 +1,130 @@
+export type BudgetMode = "focused" | "default" | "explore";
+
+export interface BudgetAllocation {
+	working: number;
+	episodic: number;
+	semantic: number;
+	retrieval: number;
+	total: number;
+	headroom: number; // reserved for AI reasoning
+}
+
+export interface LayerContent {
+	name: string;
+	text: string;
+	tokens: number;
+	priority: number; // lower = higher priority (working=0, semantic=1, episodic=2, retrieval=3)
+}
+
+const DEFAULT_MODEL_CONTEXT_WINDOW = 200_000;
+
+/**
+ * Approximate token count: ~1 token per 3.5 characters.
+ */
+export function calculateTokens(text: string): number {
+	if (text.length === 0) return 0;
+	return Math.ceil(text.length / 3.5);
+}
+
+/**
+ * Returns the fraction of the model context window available for layer content.
+ * The remainder is reserved as headroom for AI reasoning.
+ */
+export function getBudgetRatio(mode: BudgetMode): number {
+	switch (mode) {
+		case "focused":
+			return 0.4;
+		case "default":
+			return 0.6;
+		case "explore":
+			return 0.8;
+	}
+}
+
+/**
+ * Calculates per-layer token allocations for a given mode and model context window.
+ *
+ * Layer proportions within the usable budget:
+ *   working:   ~25%
+ *   semantic:  ~33%
+ *   episodic:  ~25%
+ *   retrieval: ~17%  (remainder after the above three)
+ */
+export function assembleBudget(
+	mode: BudgetMode,
+	modelContextWindow: number = DEFAULT_MODEL_CONTEXT_WINDOW,
+): BudgetAllocation {
+	const ratio = getBudgetRatio(mode);
+	const budget = Math.floor(modelContextWindow * ratio);
+	const headroom = modelContextWindow - budget;
+
+	const working = Math.floor(budget * 0.25);
+	const semantic = Math.floor(budget * 0.33);
+	const episodic = Math.floor(budget * 0.25);
+	// retrieval gets the exact remainder so all layer tokens sum to budget
+	const retrieval = budget - working - semantic - episodic;
+
+	return {
+		working,
+		semantic,
+		episodic,
+		retrieval,
+		headroom,
+		total: modelContextWindow,
+	};
+}
+
+/**
+ * Trims layers so the total token count fits within the budget.
+ * Drops lowest-priority layers first (retrieval → episodic → semantic).
+ * The working layer (priority 0) is never removed.
+ *
+ * Returns the surviving layers in ascending priority order.
+ */
+export function truncateToFit(
+	layers: LayerContent[],
+	budget: BudgetAllocation,
+): LayerContent[] {
+	// Sort by priority ascending so highest-priority layers come first
+	const sorted = [...layers].sort((a, b) => a.priority - b.priority);
+
+	const totalTokens = sorted.reduce((sum, l) => sum + l.tokens, 0);
+	if (totalTokens <= budget.total) {
+		return sorted;
+	}
+
+	// Drop layers from lowest priority (highest numeric value) until we fit,
+	// but never drop priority-0 (working) layers.
+	const mutable = [...sorted];
+	let current = totalTokens;
+
+	while (current > budget.total) {
+		// Find the lowest-priority non-working layer
+		let dropIndex = -1;
+		let maxPriority = -1;
+		for (let i = 0; i < mutable.length; i++) {
+			const layer = mutable[i];
+			if (
+				layer !== undefined &&
+				layer.priority > 0 &&
+				layer.priority > maxPriority
+			) {
+				maxPriority = layer.priority;
+				dropIndex = i;
+			}
+		}
+
+		if (dropIndex === -1) {
+			// Only working layers remain — stop
+			break;
+		}
+
+		const dropped = mutable[dropIndex];
+		if (dropped !== undefined) {
+			current -= dropped.tokens;
+		}
+		mutable.splice(dropIndex, 1);
+	}
+
+	return mutable;
+}