Skip to content

Commit 230e8c6

Browse files
phernandezphernandez
committed
feat(context-engine): move recall and capture into native lifecycle
1 parent bda2fd6 commit 230e8c6

28 files changed

+563
-55
lines changed

‎context-engine/basic-memory-context-engine.test.ts‎

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
import { beforeEach, describe, expect, it, jest } from "bun:test"
2+
import type { AgentMessage } from "@mariozechner/pi-agent-core"
3+
import type { BmClient } from "../bm-client.ts"
4+
import type { BasicMemoryConfig } from "../config.ts"
5+
import { BasicMemoryContextEngine } from "./basic-memory-context-engine.ts"
6+
7+
function makeConfig(
8+
overrides?: Partial<BasicMemoryConfig>,
9+
): BasicMemoryConfig {
10+
return {
11+
project: "test-project",
12+
bmPath: "bm",
13+
memoryDir: "memory/",
14+
memoryFile: "MEMORY.md",
15+
projectPath: "/tmp/test-project",
16+
autoCapture: true,
17+
captureMinChars: 10,
18+
autoRecall: true,
19+
recallPrompt:
20+
"Check for active tasks and recent activity. Summarize anything relevant to the current session.",
21+
debug: false,
22+
...overrides,
23+
}
24+
}
25+
26+
function makeMessages(messages: Array<Record<string, unknown>>): AgentMessage[] {
27+
return messages as AgentMessage[]
28+
}
29+
30+
describe("BasicMemoryContextEngine", () => {
31+
let mockClient: {
32+
search: jest.Mock
33+
recentActivity: jest.Mock
34+
indexConversation: jest.Mock
35+
}
36+
37+
beforeEach(() => {
38+
mockClient = {
39+
search: jest.fn().mockResolvedValue([
40+
{
41+
title: "Fix auth rollout",
42+
permalink: "fix-auth-rollout",
43+
content: "Continue staging verification for auth rollout.",
44+
file_path: "memory/tasks/fix-auth-rollout.md",
45+
},
46+
]),
47+
recentActivity: jest.fn().mockResolvedValue([
48+
{
49+
title: "API review",
50+
permalink: "api-review",
51+
file_path: "memory/api-review.md",
52+
created_at: "2026-03-09T12:00:00Z",
53+
},
54+
]),
55+
indexConversation: jest.fn().mockResolvedValue(undefined),
56+
}
57+
})
58+
59+
it("bootstraps recall state from active tasks and recent activity", async () => {
60+
const engine = new BasicMemoryContextEngine(
61+
mockClient as unknown as BmClient,
62+
makeConfig(),
63+
)
64+
65+
await expect(
66+
engine.bootstrap({
67+
sessionId: "session-1",
68+
sessionFile: "/tmp/session-1.jsonl",
69+
}),
70+
).resolves.toEqual({ bootstrapped: true })
71+
expect(mockClient.search).toHaveBeenCalledWith(undefined, 5, undefined, {
72+
note_types: ["Task"],
73+
status: "active",
74+
})
75+
expect(mockClient.recentActivity).toHaveBeenCalledWith("1d")
76+
})
77+
78+
it("skips bootstrap when recall is disabled", async () => {
79+
const engine = new BasicMemoryContextEngine(
80+
mockClient as unknown as BmClient,
81+
makeConfig({ autoRecall: false }),
82+
)
83+
84+
await expect(
85+
engine.bootstrap({
86+
sessionId: "session-2",
87+
sessionFile: "/tmp/session-2.jsonl",
88+
}),
89+
).resolves.toEqual({
90+
bootstrapped: false,
91+
reason: "autoRecall disabled",
92+
})
93+
94+
const result = await engine.assemble({
95+
sessionId: "session-2",
96+
messages: makeMessages([{ role: "user", content: "hello" }]),
97+
})
98+
99+
expect(result).toEqual({
100+
messages: makeMessages([{ role: "user", content: "hello" }]),
101+
estimatedTokens: 0,
102+
})
103+
})
104+
105+
it("returns a no-op bootstrap result when there is no recall context", async () => {
106+
mockClient.search.mockResolvedValue([])
107+
mockClient.recentActivity.mockResolvedValue([])
108+
109+
const engine = new BasicMemoryContextEngine(
110+
mockClient as unknown as BmClient,
111+
makeConfig(),
112+
)
113+
114+
await expect(
115+
engine.bootstrap({
116+
sessionId: "session-3",
117+
sessionFile: "/tmp/session-3.jsonl",
118+
}),
119+
).resolves.toEqual({
120+
bootstrapped: false,
121+
reason: "no recall context found",
122+
})
123+
})
124+
125+
it("captures only the current turn after prePromptMessageCount", async () => {
126+
const engine = new BasicMemoryContextEngine(
127+
mockClient as unknown as BmClient,
128+
makeConfig(),
129+
)
130+
131+
await engine.afterTurn({
132+
sessionId: "session-4",
133+
sessionFile: "/tmp/session-4.jsonl",
134+
prePromptMessageCount: 2,
135+
messages: makeMessages([
136+
{ role: "user", content: "Old question" },
137+
{ role: "assistant", content: "Old answer" },
138+
{ role: "user", content: "Current question with enough detail" },
139+
{ role: "assistant", content: "Current answer with enough detail" },
140+
]),
141+
})
142+
143+
expect(mockClient.indexConversation).toHaveBeenCalledWith(
144+
"Current question with enough detail",
145+
"Current answer with enough detail",
146+
)
147+
})
148+
149+
it("respects captureMinChars for afterTurn capture", async () => {
150+
const engine = new BasicMemoryContextEngine(
151+
mockClient as unknown as BmClient,
152+
makeConfig({ captureMinChars: 50 }),
153+
)
154+
155+
await engine.afterTurn({
156+
sessionId: "session-5",
157+
sessionFile: "/tmp/session-5.jsonl",
158+
prePromptMessageCount: 0,
159+
messages: makeMessages([
160+
{ role: "user", content: "short" },
161+
{ role: "assistant", content: "tiny" },
162+
]),
163+
})
164+
165+
expect(mockClient.indexConversation).not.toHaveBeenCalled()
166+
})
167+
168+
it("swallows capture failures in afterTurn", async () => {
169+
mockClient.indexConversation.mockRejectedValue(new Error("BM down"))
170+
const engine = new BasicMemoryContextEngine(
171+
mockClient as unknown as BmClient,
172+
makeConfig(),
173+
)
174+
175+
await expect(
176+
engine.afterTurn({
177+
sessionId: "session-6",
178+
sessionFile: "/tmp/session-6.jsonl",
179+
prePromptMessageCount: 0,
180+
messages: makeMessages([
181+
{ role: "user", content: "Current question with enough detail" },
182+
{ role: "assistant", content: "Current answer with enough detail" },
183+
]),
184+
}),
185+
).resolves.toBeUndefined()
186+
})
187+
})

‎context-engine/basic-memory-context-engine.ts‎

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
1+
import { createRequire } from "node:module"
2+
import { dirname, resolve } from "node:path"
3+
import { pathToFileURL } from "node:url"
4+
import type { AgentMessage } from "@mariozechner/pi-agent-core"
5+
import type {
6+
AssembleResult,
7+
BootstrapResult,
8+
CompactResult,
9+
ContextEngine,
10+
} from "openclaw/plugin-sdk"
11+
import type { BmClient } from "../bm-client.ts"
12+
import type { BasicMemoryConfig } from "../config.ts"
13+
import { selectCaptureTurn } from "../hooks/capture.ts"
14+
import { loadRecallState } from "../hooks/recall.ts"
15+
import { log } from "../logger.ts"
16+
17+
const require = createRequire(import.meta.url)
18+
19+
interface SessionMemoryState {
20+
recallContext: string
21+
}
22+
23+
type LegacyContextEngineModule = {
24+
LegacyContextEngine: new () => {
25+
compact(params: {
26+
sessionId: string
27+
sessionFile: string
28+
tokenBudget?: number
29+
force?: boolean
30+
currentTokenCount?: number
31+
compactionTarget?: "budget" | "threshold"
32+
customInstructions?: string
33+
runtimeContext?: Record<string, unknown>
34+
}): Promise<CompactResult>
35+
}
36+
}
37+
38+
async function loadLegacyContextEngine(): Promise<
39+
LegacyContextEngineModule["LegacyContextEngine"]
40+
> {
41+
const pluginSdkPath = require.resolve("openclaw/plugin-sdk")
42+
const legacyPath = resolve(
43+
dirname(pluginSdkPath),
44+
"context-engine",
45+
"legacy.js",
46+
)
47+
const module = (await import(
48+
pathToFileURL(legacyPath).href
49+
)) as LegacyContextEngineModule
50+
return module.LegacyContextEngine
51+
}
52+
53+
export class BasicMemoryContextEngine implements ContextEngine {
54+
readonly info = {
55+
id: "openclaw-basic-memory",
56+
name: "Basic Memory Context Engine",
57+
version: "0.1.5",
58+
ownsCompaction: false,
59+
} as const
60+
61+
private readonly sessionState = new Map<string, SessionMemoryState>()
62+
private legacyContextEnginePromise:
63+
| Promise<InstanceType<LegacyContextEngineModule["LegacyContextEngine"]>>
64+
| null = null
65+
66+
constructor(
67+
private readonly client: BmClient,
68+
private readonly cfg: BasicMemoryConfig,
69+
) {}
70+
71+
async bootstrap(params: {
72+
sessionId: string
73+
sessionFile: string
74+
}): Promise<BootstrapResult> {
75+
if (!this.cfg.autoRecall) {
76+
this.sessionState.delete(params.sessionId)
77+
return { bootstrapped: false, reason: "autoRecall disabled" }
78+
}
79+
80+
try {
81+
const recall = await loadRecallState(this.client, this.cfg)
82+
if (!recall) {
83+
this.sessionState.delete(params.sessionId)
84+
return { bootstrapped: false, reason: "no recall context found" }
85+
}
86+
87+
this.sessionState.set(params.sessionId, {
88+
recallContext: recall.context,
89+
})
90+
91+
log.debug(
92+
`context-engine bootstrap: session=${params.sessionId} tasks=${recall.tasks.length} recent=${recall.recent.length}`,
93+
)
94+
95+
return { bootstrapped: true }
96+
} catch (err) {
97+
this.sessionState.delete(params.sessionId)
98+
log.error("context-engine bootstrap failed", err)
99+
return { bootstrapped: false, reason: "recall failed" }
100+
}
101+
}
102+
103+
async ingest(): Promise<{ ingested: boolean }> {
104+
return { ingested: false }
105+
}
106+
107+
async assemble(params: {
108+
sessionId: string
109+
messages: AgentMessage[]
110+
tokenBudget?: number
111+
}): Promise<AssembleResult> {
112+
return {
113+
messages: params.messages,
114+
estimatedTokens: 0,
115+
}
116+
}
117+
118+
async afterTurn(params: {
119+
sessionId: string
120+
sessionFile: string
121+
messages: AgentMessage[]
122+
prePromptMessageCount: number
123+
autoCompactionSummary?: string
124+
isHeartbeat?: boolean
125+
tokenBudget?: number
126+
runtimeContext?: Record<string, unknown>
127+
}): Promise<void> {
128+
if (!this.cfg.autoCapture) return
129+
130+
const newMessages = params.messages.slice(params.prePromptMessageCount)
131+
const turn =
132+
selectCaptureTurn(newMessages, this.cfg.captureMinChars) ??
133+
selectCaptureTurn(params.messages, this.cfg.captureMinChars)
134+
135+
if (!turn) return
136+
137+
log.debug(
138+
`context-engine afterTurn: session=${params.sessionId} user=${turn.userText.length} assistant=${turn.assistantText.length}`,
139+
)
140+
141+
try {
142+
await this.client.indexConversation(turn.userText, turn.assistantText)
143+
} catch (err) {
144+
log.error("context-engine capture failed", err)
145+
}
146+
}
147+
148+
async compact(params: {
149+
sessionId: string
150+
sessionFile: string
151+
tokenBudget?: number
152+
force?: boolean
153+
currentTokenCount?: number
154+
compactionTarget?: "budget" | "threshold"
155+
customInstructions?: string
156+
runtimeContext?: Record<string, unknown>
157+
}): Promise<CompactResult> {
158+
const legacy = await this.getLegacyContextEngine()
159+
return legacy.compact(params)
160+
}
161+
162+
async dispose(): Promise<void> {
163+
this.sessionState.clear()
164+
}
165+
166+
private async getLegacyContextEngine(): Promise<
167+
InstanceType<LegacyContextEngineModule["LegacyContextEngine"]>
168+
> {
169+
if (!this.legacyContextEnginePromise) {
170+
this.legacyContextEnginePromise = loadLegacyContextEngine().then(
171+
(LegacyContextEngine) => new LegacyContextEngine(),
172+
)
173+
}
174+
175+
return this.legacyContextEnginePromise
176+
}
177+
}

0 commit comments

Comments
 (0)