Skip to content

Commit a89f921

Browse files
ualtinokalfonso-magic-context
ualtinok
and
alfonso-magic-context
committed
mason: harden dreamer manifests and leases
Co-authored-by: Alfonso [Magic Context] <288211368+alfonso-magic-context@users.noreply.github.com>
1 parent 131de2c commit a89f921

28 files changed

Lines changed: 872 additions & 203 deletions

‎packages/pi-plugin/src/dreamer/retrospective-raw-provider-pi.test.ts‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,12 @@ describe("PiRetrospectiveRawProvider", () => {
1414
path: "/sessions/s1.jsonl",
1515
modified: 30,
1616
},
17+
{
18+
id: "s0",
19+
cwd: "/repo/project",
20+
path: "/sessions/s0.jsonl",
21+
modified: 20,
22+
},
1723
{
1824
id: "s2",
1925
cwd: "/repo/other",
@@ -25,6 +31,7 @@ describe("PiRetrospectiveRawProvider", () => {
2531
});
2632

2733
expect(await provider.listProjectSessions("identity")).toEqual([
34+
{ sessionId: "s0", path: "/sessions/s0.jsonl", updatedAt: 20 },
2835
{ sessionId: "s1", path: "/sessions/s1.jsonl", updatedAt: 30 },
2936
]);
3037
});

‎packages/pi-plugin/src/dreamer/retrospective-raw-provider-pi.ts‎

Lines changed: 15 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -73,7 +73,7 @@ export class PiRetrospectiveRawProvider implements RetrospectiveRawProvider {
7373
});
7474
}
7575

76-
return result.sort((a, b) => (b.updatedAt ?? 0) - (a.updatedAt ?? 0));
76+
return result.sort((a, b) => (a.updatedAt ?? 0) - (b.updatedAt ?? 0));
7777
}
7878

7979
async readUserMessagesSince(
@@ -98,6 +98,20 @@ export class PiRetrospectiveRawProvider implements RetrospectiveRawProvider {
9898
};
9999
}
100100

101+
async readOldestMessageTimesSince(
102+
sessionIds: readonly string[],
103+
sinceMs: number,
104+
): Promise<Map<string, number>> {
105+
const out = new Map<string, number>();
106+
for (const sessionId of sessionIds) {
107+
const oldest = (await this.loadUserEntries(sessionId))
108+
.filter((message) => message.ts > sinceMs)
109+
.sort((a, b) => a.ts - b.ts || a.ordinal - b.ordinal)[0];
110+
if (oldest) out.set(sessionId, oldest.ts);
111+
}
112+
return out;
113+
}
114+
101115
async readUserMessagesBefore(
102116
sessionId: string,
103117
beforeMs: number,

‎packages/plugin/src/features/magic-context/dreamer/classify-prompt.test.ts‎

Lines changed: 21 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -18,19 +18,29 @@ describe("parseClassifyManifest", () => {
1818
]);
1919
});
2020

21-
it("omits invalid scope and keeps the rest", () => {
22-
const out = parseClassifyManifest(
23-
`<memory id="5" importance="50" scope="bogus" shareable="true"/>`,
24-
);
25-
expect(out).toEqual([{ id: 5, importance: 50, shareable: true }]);
26-
});
27-
28-
it("skips an entry that carries no classification fields", () => {
29-
expect(parseClassifyManifest(`<memory id="9"/>`)).toEqual([]);
21+
it("rejects invalid scope", () => {
22+
expect(() =>
23+
parseClassifyManifest(
24+
`<classify><memory id="5" importance="50" scope="bogus" shareable="true"/></classify>`,
25+
),
26+
).toThrow(/invalid scope/);
3027
});
3128

32-
it("skips a non-numeric id", () => {
33-
expect(parseClassifyManifest(`<memory id="x" importance="50"/>`)).toEqual([]);
29+
it("rejects truncated, duplicate, and invalid entries", () => {
30+
expect(() => parseClassifyManifest(`<classify><memory id="5" importance="50"/>`)).toThrow(
31+
/closing root/,
32+
);
33+
expect(() =>
34+
parseClassifyManifest(
35+
`<classify><memory id="5" importance="50"/><memory id="5" shareable="true"/></classify>`,
36+
),
37+
).toThrow(/duplicate id/);
38+
expect(() =>
39+
parseClassifyManifest(`<classify><memory id="x" importance="50"/></classify>`),
40+
).toThrow(/numeric id/);
41+
expect(() => parseClassifyManifest(`<classify><memory id="9"/></classify>`)).toThrow(
42+
/classification fields/,
43+
);
3444
});
3545
});
3646

‎packages/plugin/src/features/magic-context/dreamer/classify-prompt.ts‎

Lines changed: 19 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@
1111
* 229 importances assigned, correct discrimination, 4/4 private controls held).
1212
*/
1313

14+
import { assertNoDuplicateManifestIds, extractCompleteManifestBody } from "./manifest-parser";
15+
1416
export interface ClassifyPromptMemory {
1517
id: number;
1618
category: string;
@@ -116,16 +118,17 @@ export interface ParsedClassification {
116118

117119
const SCOPES = new Set(["project", "ecosystem", "universe"]);
118120

119-
/** Parse the agent's `<classify>` manifest. Tolerant of attribute order; a
120-
* memory missing a valid attribute simply omits that field (host skips it). */
121+
/** Parse the agent's complete `<classify>` manifest. A missing root close tag is
122+
* treated as truncation and rejects the whole batch. */
121123
export function parseClassifyManifest(text: string): ParsedClassification[] {
122124
const out: ParsedClassification[] = [];
123-
for (const m of text.matchAll(/<memory\b([^>]*)\/?>/g)) {
125+
const body = extractCompleteManifestBody(text, "classify");
126+
for (const m of body.matchAll(/<memory\b([^>]*)\/?>/g)) {
124127
const attrs = m[1];
125128
const idMatch = attrs.match(/\bid\s*=\s*"(\d+)"/);
126-
if (!idMatch) continue;
129+
if (!idMatch) throw new Error("classify manifest entry missing numeric id");
127130
const id = Number.parseInt(idMatch[1], 10);
128-
if (!Number.isInteger(id)) continue;
131+
if (!Number.isInteger(id)) throw new Error("classify manifest entry missing numeric id");
129132

130133
const entry: ParsedClassification = { id };
131134
const impMatch = attrs.match(/\bimportance\s*=\s*"(\d+)"/);
@@ -134,18 +137,24 @@ export function parseClassifyManifest(text: string): ParsedClassification[] {
134137
if (Number.isInteger(imp)) entry.importance = Math.max(1, Math.min(100, imp));
135138
}
136139
const scopeMatch = attrs.match(/\bscope\s*=\s*"([a-z]+)"/i);
137-
if (scopeMatch && SCOPES.has(scopeMatch[1].toLowerCase())) {
138-
entry.scope = scopeMatch[1].toLowerCase() as ParsedClassification["scope"];
140+
if (scopeMatch) {
141+
const scope = scopeMatch[1].toLowerCase();
142+
if (!SCOPES.has(scope)) throw new Error(`classify manifest invalid scope ${scope}`);
143+
entry.scope = scope as ParsedClassification["scope"];
139144
}
140145
const shareMatch = attrs.match(/\bshareable\s*=\s*"(true|false|1|0)"/i);
141146
if (shareMatch) {
142147
const v = shareMatch[1].toLowerCase();
143148
entry.shareable = v === "true" || v === "1";
144149
}
145-
// Only keep an entry that carries at least one classification field.
146-
if (entry.importance !== undefined || entry.scope || entry.shareable !== undefined) {
147-
out.push(entry);
150+
if (entry.importance === undefined && !entry.scope && entry.shareable === undefined) {
151+
throw new Error(`classify manifest entry ${id} missing classification fields`);
148152
}
153+
out.push(entry);
149154
}
155+
assertNoDuplicateManifestIds(
156+
out.map((entry) => entry.id),
157+
"classify",
158+
);
150159
return out;
151160
}

‎packages/plugin/src/features/magic-context/dreamer/classify.test.ts‎

Lines changed: 99 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,99 @@
1+
/// <reference types="bun-types" />
2+
3+
import { describe, expect, test } from "bun:test";
4+
5+
import { Database } from "../../../shared/sqlite";
6+
import { closeQuietly } from "../../../shared/sqlite-helpers";
7+
import { getMemoryById, insertMemory } from "../memory";
8+
import { runMigrations } from "../migrations";
9+
import { initializeDatabase } from "../storage-db";
10+
import { applyClassifications, type ClassifyArgs } from "./classify";
11+
import { acquireLease } from "./lease";
12+
13+
function freshDb(): Database {
14+
const db = new Database(":memory:");
15+
initializeDatabase(db);
16+
runMigrations(db);
17+
return db;
18+
}
19+
20+
function classifyArgs(db: Database, projectIdentity: string): ClassifyArgs {
21+
const holderId = "classify-holder";
22+
const leaseKey = `classify-${Math.random()}`;
23+
expect(acquireLease(db, holderId, leaseKey)).toBe(true);
24+
return {
25+
db,
26+
client: {} as never,
27+
projectIdentity,
28+
parentSessionId: undefined,
29+
sessionDirectory: process.cwd(),
30+
holderId,
31+
leaseKey,
32+
deadline: Date.now() + 60_000,
33+
};
34+
}
35+
36+
describe("applyClassifications", () => {
37+
test("complete manifest applies classification fields", () => {
38+
const db = freshDb();
39+
try {
40+
const projectIdentity = "git:test";
41+
const memory = insertMemory(db, {
42+
projectPath: projectIdentity,
43+
category: "ARCHITECTURE",
44+
content: "Important project fact.",
45+
sourceSessionId: "ses",
46+
});
47+
48+
const result = applyClassifications(
49+
classifyArgs(db, projectIdentity),
50+
[memory],
51+
`<classify><memory id="${memory.id}" importance="85" scope="project" shareable="true"/></classify>`,
52+
);
53+
54+
expect(result.classified).toBe(1);
55+
const after = getMemoryById(db, memory.id);
56+
expect(after?.importance).toBe(85);
57+
expect(after?.scope).toBe("project");
58+
expect(after?.shareable).toBe(1);
59+
} finally {
60+
closeQuietly(db);
61+
}
62+
});
63+
64+
test("truncated manifest rejects before stamping classified_at", () => {
65+
const db = freshDb();
66+
try {
67+
const projectIdentity = "git:test";
68+
const memory = insertMemory(db, {
69+
projectPath: projectIdentity,
70+
category: "ARCHITECTURE",
71+
content: "Important project fact.",
72+
sourceSessionId: "ses",
73+
});
74+
const before = getMemoryById(db, memory.id);
75+
const beforeRow = db
76+
.prepare("SELECT classified_at FROM memories WHERE id = ?")
77+
.get(memory.id) as { classified_at?: number | null } | undefined;
78+
79+
expect(() =>
80+
applyClassifications(
81+
classifyArgs(db, projectIdentity),
82+
[memory],
83+
`<classify><memory id="${memory.id}" importance="85"`,
84+
),
85+
).toThrow(/closing root/);
86+
87+
const after = getMemoryById(db, memory.id);
88+
expect(after?.importance).toBe(before?.importance);
89+
expect(after?.scope).toBe(before?.scope);
90+
expect(after?.shareable).toBe(before?.shareable);
91+
const afterRow = db
92+
.prepare("SELECT classified_at FROM memories WHERE id = ?")
93+
.get(memory.id) as { classified_at?: number | null } | undefined;
94+
expect(afterRow?.classified_at).toBe(beforeRow?.classified_at);
95+
} finally {
96+
closeQuietly(db);
97+
}
98+
});
99+
});

‎packages/plugin/src/features/magic-context/dreamer/classify.ts‎

Lines changed: 21 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,10 @@
11
import { DREAMER_CLASSIFIER_AGENT } from "../../../agents/dreamer";
22
import type { PluginContext } from "../../../plugin/types";
33
import * as shared from "../../../shared";
4-
import { extractLatestAssistantText } from "../../../shared/assistant-message-extractor";
4+
import {
5+
extractLatestAssistantText,
6+
hasLengthCappedOutput,
7+
} from "../../../shared/assistant-message-extractor";
58
import { describeError, getErrorMessage } from "../../../shared/error-message";
69
import { shouldKeepSubagents } from "../../../shared/keep-subagents";
710
import { log } from "../../../shared/logger";
@@ -22,7 +25,8 @@ import {
2225
type ClassifyPromptMemory,
2326
parseClassifyManifest,
2427
} from "./classify-prompt";
25-
import { peekLeaseHolderAndExpiry, startLeaseHeartbeat } from "./lease";
28+
import { runLeaseGuardedWrite, startLeaseHeartbeat } from "./lease";
29+
import { assertManifestCoversExactly } from "./manifest-parser";
2630

2731
/**
2832
* classify-memories: a NON-agentic single-shot transform. Scores each project
@@ -238,8 +242,12 @@ async function classifyOneChunk(
238242
});
239243
},
240244
validateOutput: (messages) => {
245+
if (hasLengthCappedOutput(messages)) {
246+
throw new Error("classify returned length-capped output");
247+
}
241248
const text = extractLatestAssistantText(messages);
242249
if (!text) throw new Error("classify returned no output");
250+
parseClassifyManifest(text);
243251
return text;
244252
},
245253
},
@@ -273,26 +281,26 @@ async function classifyOneChunk(
273281
}
274282
}
275283

276-
/** Apply the manifest host-side: only ids that were IN this chunk; shareable
277-
* fails closed against sensitive text. setMemoryClassification stamps
284+
/** Apply the manifest host-side: the manifest must cover exactly this chunk;
285+
* shareable fails closed against sensitive text. setMemoryClassification stamps
278286
* classified_at (the run-gate) and is cache-neutral. */
279-
function applyClassifications(
287+
export function applyClassifications(
280288
args: ClassifyArgs,
281289
chunk: Memory[],
282290
manifestText: string,
283291
): { classified: number; changed: number } {
284292
const byId = new Map(chunk.map((m) => [m.id, m]));
285-
const parsed = parseClassifyManifest(manifestText).filter((p) => byId.has(p.id));
293+
const parsed = parseClassifyManifest(manifestText);
294+
assertManifestCoversExactly(
295+
parsed.map((entry) => entry.id),
296+
new Set(byId.keys()),
297+
"classify",
298+
);
286299
if (parsed.length === 0) return { classified: 0, changed: 0 };
287300

288301
let classified = 0;
289302
let changed = 0;
290-
let leaseLost = false;
291-
args.db.transaction(() => {
292-
if (!peekLeaseHolderAndExpiry(args.db, args.holderId, args.leaseKey)) {
293-
leaseLost = true;
294-
return;
295-
}
303+
runLeaseGuardedWrite(args.db, args.holderId, args.leaseKey, () => {
296304
for (const p of parsed) {
297305
const memory = byId.get(p.id);
298306
if (!memory) continue;
@@ -310,8 +318,7 @@ function applyClassifications(
310318
classified += 1; // stamped classified_at (run-gate satisfied)
311319
if (didChange) changed += 1; // an actual column value moved
312320
}
313-
})();
314-
if (leaseLost) throw new Error("Dream lease lost during classify commit");
321+
});
315322
return { classified, changed };
316323
}
317324

0 commit comments

Comments
 (0)