feat(core): add createBookContextTransform for truth file injection

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
fanghanjun
2026-04-17 03:01:18 -07:00
parent 95ad429646
commit 4fbf4ee66f
2 changed files with 155 additions and 0 deletions
@@ -0,0 +1,82 @@
import { describe, it, expect, beforeEach, afterEach } from "vitest";
import { mkdtemp, mkdir, writeFile, rm } from "node:fs/promises";
import { join } from "node:path";
import { tmpdir } from "node:os";
import { createBookContextTransform } from "../agent/context-transform.js";
describe("createBookContextTransform", () => {
let projectRoot: string;
const bookId = "test-book";
beforeEach(async () => {
projectRoot = await mkdtemp(join(tmpdir(), "ctx-test-"));
const storyDir = join(projectRoot, "books", bookId, "story");
await mkdir(storyDir, { recursive: true });
await writeFile(join(storyDir, "story_bible.md"), "# Story Bible\nA hero's journey.");
await writeFile(join(storyDir, "current_focus.md"), "Focus on chapter 3.");
});
afterEach(async () => {
await rm(projectRoot, { recursive: true, force: true });
});
it("returns messages unchanged when bookId is null", async () => {
const transform = createBookContextTransform(null, projectRoot);
const messages = [
{ role: "user" as const, content: "hello", timestamp: Date.now() },
];
const result = await transform(messages);
expect(result).toBe(messages);
});
it("prepends a user message with truth file contents", async () => {
const transform = createBookContextTransform(bookId, projectRoot);
const original = [
{ role: "user" as const, content: "写下一章", timestamp: Date.now() },
];
const result = await transform(original);
expect(original).toHaveLength(1);
expect(result).toHaveLength(2);
const injected = result[0] as { role: string; content: string };
expect(injected.role).toBe("user");
expect(injected.content).toContain("story_bible.md");
expect(injected.content).toContain("A hero's journey.");
expect(injected.content).toContain("current_focus.md");
expect(injected.content).toContain("Focus on chapter 3.");
expect(result[1]).toBe(original[0]);
});
it("sorts truth files in priority order", async () => {
const storyDir = join(projectRoot, "books", bookId, "story");
await writeFile(join(storyDir, "volume_outline.md"), "# Volume Outline");
await writeFile(join(storyDir, "book_rules.md"), "# Book Rules");
await writeFile(join(storyDir, "extra_notes.md"), "# Extra");
const transform = createBookContextTransform(bookId, projectRoot);
const result = await transform([
{ role: "user" as const, content: "test", timestamp: Date.now() },
]);
const content = (result[0] as { content: string }).content;
const bibleIdx = content.indexOf("story_bible.md");
const outlineIdx = content.indexOf("volume_outline.md");
const rulesIdx = content.indexOf("book_rules.md");
const focusIdx = content.indexOf("current_focus.md");
const extraIdx = content.indexOf("extra_notes.md");
expect(bibleIdx).toBeLessThan(outlineIdx);
expect(outlineIdx).toBeLessThan(rulesIdx);
expect(rulesIdx).toBeLessThan(focusIdx);
expect(focusIdx).toBeLessThan(extraIdx);
});
it("returns original messages when story/ directory does not exist", async () => {
const transform = createBookContextTransform("nonexistent-book", projectRoot);
const original = [
{ role: "user" as const, content: "test", timestamp: Date.now() },
];
const result = await transform(original);
expect(result).toBe(original);
});
});
@@ -0,0 +1,73 @@
import type { AgentMessage } from "@mariozechner/pi-agent-core";
import type { UserMessage } from "@mariozechner/pi-ai";
import { readdir, readFile } from "node:fs/promises";
import { join } from "node:path";
/** Files read in this order; anything else in story/ comes after, sorted alphabetically. */
const PRIORITY_FILES = [
"story_bible.md",
"volume_outline.md",
"book_rules.md",
"current_focus.md",
];
export function createBookContextTransform(
bookId: string | null,
projectRoot: string,
): (messages: AgentMessage[], signal?: AbortSignal) => Promise<AgentMessage[]> {
if (bookId === null) {
return async (messages) => messages;
}
const storyDir = join(projectRoot, "books", bookId, "story");
return async (messages) => {
const sections = await readTruthFiles(storyDir);
if (sections.length === 0) return messages;
const body =
"[以下是当前书籍的真相文件,每次对话时自动从磁盘读取注入。请基于这些内容进行创作和判断。]\n\n" +
sections.map((s) => `=== ${s.name} ===\n${s.content}`).join("\n\n");
const injected: UserMessage = {
role: "user",
content: body,
timestamp: Date.now(),
};
return [injected, ...messages];
};
}
interface TruthFileSection {
name: string;
content: string;
}
async function readTruthFiles(storyDir: string): Promise<TruthFileSection[]> {
let entries: string[];
try {
entries = await readdir(storyDir);
} catch {
return [];
}
const mdFiles = entries.filter((f) => f.endsWith(".md"));
if (mdFiles.length === 0) return [];
const prioritySet = new Set(PRIORITY_FILES);
const prioritized = PRIORITY_FILES.filter((f) => mdFiles.includes(f));
const rest = mdFiles.filter((f) => !prioritySet.has(f)).sort();
const ordered = [...prioritized, ...rest];
const sections: TruthFileSection[] = [];
for (const fileName of ordered) {
try {
const content = await readFile(join(storyDir, fileName), "utf-8");
sections.push({ name: fileName, content });
} catch {
// skip unreadable files
}
}
return sections;
}