feat(pi): add full lifecycle hooks to pi integration#604
Open
paulodearaujo wants to merge 1 commit into
Open
Conversation
|
@paulodearaujo is attempting to deploy a commit to the rohitg00's projects Team on Vercel. A member of the Team first needs to authorize it. |
Contributor
📝 WalkthroughWalkthroughThe PI extension import is updated; helpers for base URL normalization, SHA‑256 hashing, image-scrubbing, and client-side deduplication (5‑minute window) are added. callAgentMemory gains timeout support. Session and tool lifecycle hooks are reorganized to post deduplicated observations, trigger background context/summarize, and sequence summarize then session/end on shutdown. ChangesPI Extension Lifecycle and Deduplication
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Warning There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure. 🔧 ESLint
ESLint skipped: no ESLint configuration detected in root package.json. To enable, add Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@integrations/pi/index.ts`:
- Around line 332-348: The background observe posts started from the
pi.on("tool_result") (and the similar agent_end handler) can still be in-flight
when session_shutdown summarizes/closes; track and await these before
finalizing. Create a shared pendingObserves Set/array used by
callAgentMemory("observe") callers: when you call callAgentMemory(...) from the
tool_result and agent_end handlers, push the returned Promise into
pendingObserves and remove it on settle; in session_shutdown (and any
summary/close function) await Promise.all([...pendingObserves]) (or drain the
set) before proceeding to summarize/close to ensure no observe requests are
dropped. Ensure the same pattern is applied to the other observe emitters
mentioned.
- Around line 399-420: The session_shutdown handler awaits
callAgentMemory("summarize", ...) and callAgentMemory("session/end", ...)
without a timeout or AbortSignal, so PI shutdown can hang; modify the handler in
integrations/pi/index.ts (the pi.on("session_shutdown", ...) block) to create an
AbortController with a bounded timeout (e.g., setTimeout to controller.abort
after N ms), pass controller.signal into callAgentMemory (and update
callAgentMemory to accept and forward an AbortSignal to fetch if it doesn't
already), and ensure you abort/clear the timer after the requests complete (use
Promise.race or Promise.allSettled to await both with the signal) so the
shutdown path always finishes promptly even if the server never responds.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
paulodearaujo
force-pushed
the
feat/pi-full-lifecycle-hooks
branch
from
225716c to
1e847e3
Compare
Wire tool_result, session_shutdown, session_before_compact, and session/start hooks that pi's extension API already supports but were not used. Fix deprecated @mariozechner import. Before: 3 hooks (session_start, before_agent_start, agent_end). After: 6 hooks matching Codex parity. Changes: - tool_result: POST /observe per tool call with dedup - session_shutdown: await POST /summarize + /session/end - session_before_compact: POST /summarize before context loss - session_start: POST /session/start for project profile loading - Fix import to @earendil-works/pi-coding-agent - Extract observePayload() helper to reduce repetition - Add client-side SHA-256 dedup (5min window, mirrors server)
paulodearaujo
force-pushed
the
feat/pi-full-lifecycle-hooks
branch
from
1e847e3 to
f4a4a47
Compare
Contributor
There was a problem hiding this comment.
♻️ Duplicate comments (1)
integrations/pi/index.ts (1)
357-373:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDrain in-flight
observeposts before finalsummarize/session/end.
tool_resultandagent_endenqueue backgroundobservecalls, butsession_shutdowncan finalize before they settle. That can drop the last observations or summarize without them.Suggested minimal patch
+ const pendingObserves = new Set<Promise<unknown>>(); + + function trackObserve<T>(promise: Promise<T>): Promise<T> { + pendingObserves.add(promise); + void promise.finally(() => pendingObserves.delete(promise)); + return promise; + } + pi.on("tool_result", (event) => { @@ - void callAgentMemory("observe", { + void trackObserve(callAgentMemory("observe", { body: observePayload({ tool_name: event.toolName, tool_input: input.slice(0, 8000), tool_output: output.slice(0, 8000), is_error: event.isError ?? false, }), timeoutMs: 3000, - }); + })); }); @@ - void callAgentMemory("observe", { + void trackObserve(callAgentMemory("observe", { body: observePayload({ tool_name: "conversation", tool_input: lastPrompt.slice(0, 8000), tool_output: assistantText.slice(0, 8000), }), timeoutMs: 3000, - }); + })); }); @@ pi.on("session_shutdown", async () => { if (!lastHealthOk) return; + await Promise.allSettled([...pendingObserves]); await callAgentMemory("summarize", { body: { sessionId }, timeoutMs: 120_000, });Also applies to: 400-415, 424-440
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@integrations/pi/index.ts` around lines 357 - 373, The background observe calls started in the pi.on("tool_result") handler (and similar handlers around agent_end/session_shutdown) can still be in-flight when session finalization runs, so add a simple inflight tracker: create a Set or array (e.g., inflightObserves) and when you call callAgentMemory("observe", ...) from the observed places (tool_result, agent_end), push the resulting Promise into inflightObserves and attach .finally() to remove it; then update the session_shutdown / summarize / session_end flow to await Promise.race([Promise.all(inflightObserves), timeoutPromise(ms)]) or otherwise drain the Set with a short timeout to ensure the last observes complete before finalizing; reference the event handler pi.on("tool_result"), callAgentMemory("observe"), observePayload, and the session_shutdown/agent_end summarization entry points when making these changes.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Duplicate comments:
In `@integrations/pi/index.ts`:
- Around line 357-373: The background observe calls started in the
pi.on("tool_result") handler (and similar handlers around
agent_end/session_shutdown) can still be in-flight when session finalization
runs, so add a simple inflight tracker: create a Set or array (e.g.,
inflightObserves) and when you call callAgentMemory("observe", ...) from the
observed places (tool_result, agent_end), push the resulting Promise into
inflightObserves and attach .finally() to remove it; then update the
session_shutdown / summarize / session_end flow to await
Promise.race([Promise.all(inflightObserves), timeoutPromise(ms)]) or otherwise
drain the Set with a short timeout to ensure the last observes complete before
finalizing; reference the event handler pi.on("tool_result"),
callAgentMemory("observe"), observePayload, and the session_shutdown/agent_end
summarization entry points when making these changes.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
integrations/pi/index.ts (1)
102-126:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMake helper timeouts opt-out, not opt-in.
timeoutMsonly protects callers that remember to pass it. In this file, interactive paths likememory_searchandmemory_savestill omit it, so a stalled agentmemory request can hang PI indefinitely. A bounded default here is safer, and the long-running hooks can still override it explicitly.Suggested fix
async function callAgentMemory<T>( pathname: string, options?: { method?: "GET" | "POST"; body?: unknown; baseUrl?: string; timeoutMs?: number; }, ): Promise<T | null> { const baseUrl = options?.baseUrl?.replace(/\/+$/, "") || getBaseUrl(); const method = options?.method || "POST"; + const timeoutMs = options?.timeoutMs ?? 5000; const url = `${baseUrl}/agentmemory/${pathname.replace(/^\/+/, "")}`; const headers: Record<string, string> = {}; const secret = process.env.AGENTMEMORY_SECRET; guardPlaintextBearerAuth(baseUrl, secret); if (options?.body !== undefined) headers["Content-Type"] = "application/json"; if (secret) headers.Authorization = `Bearer ${secret}`; @@ const response = await fetch(url, { method, headers, body: options?.body !== undefined ? JSON.stringify(options.body) : undefined, - signal: options?.timeoutMs ? AbortSignal.timeout(options.timeoutMs) : undefined, + signal: AbortSignal.timeout(timeoutMs), });🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@integrations/pi/index.ts` around lines 102 - 126, callAgentMemory currently leaves requests unbounded unless callers pass timeoutMs, which lets interactive paths hang; change it to apply a sensible default timeout when options?.timeoutMs is undefined (e.g., set const timeout = options?.timeoutMs ?? DEFAULT_TIMEOUT_MS) and pass AbortSignal.timeout(timeout) into the fetch call so callers (memory_search, memory_save, etc.) can still override by providing timeoutMs; keep existing logic for headers, baseUrl and guardPlaintextBearerAuth, and choose a constant DEFAULT_TIMEOUT_MS near other service timeouts in the codebase.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@integrations/pi/index.ts`:
- Around line 311-319: The background call to callAgentMemory("session/start")
(with sessionId/currentProject) can race with before_agent_start/smart-search
and cause first-turn recall to miss newly-loaded profiles; change this to await
the call (or otherwise wait for its completion) instead of firing it void with
an 800ms timeout so session/context initialization finishes before
before_agent_start runs — update the callAgentMemory invocation near the
lastHealthOk branch to await the promise (and remove or extend the timeout) or
gate before_agent_start on the callAgentMemory promise so the server has
finished loading profiles/context before first-turn recall executes.
- Around line 332-354: The health check (refreshStatus(ctx)) must run before
blocking agentmemory calls so an outage doesn't cause repeated turn latency; in
before_agent_start call refreshStatus(ctx) prior to invoking
trackPost/callAgentMemory("observe") and before awaiting
callAgentMemory("smart-search"), and if refreshStatus indicates the backend is
unhealthy, skip or short-circuit the smart-search/recall logic (i.e., avoid
awaiting callAgentMemory and set recallBlock = ""), keeping the rest of the flow
(lastPrompt, results/recallBlock usage) unchanged.
---
Outside diff comments:
In `@integrations/pi/index.ts`:
- Around line 102-126: callAgentMemory currently leaves requests unbounded
unless callers pass timeoutMs, which lets interactive paths hang; change it to
apply a sensible default timeout when options?.timeoutMs is undefined (e.g., set
const timeout = options?.timeoutMs ?? DEFAULT_TIMEOUT_MS) and pass
AbortSignal.timeout(timeout) into the fetch call so callers (memory_search,
memory_save, etc.) can still override by providing timeoutMs; keep existing
logic for headers, baseUrl and guardPlaintextBearerAuth, and choose a constant
DEFAULT_TIMEOUT_MS near other service timeouts in the codebase.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
Comment on lines
+311
to
+319
| if (lastHealthOk) { | ||
| void callAgentMemory("session/start", { | ||
| body: { | ||
| sessionId, | ||
| project: currentProject, | ||
| cwd: currentProject, | ||
| }, | ||
| timeoutMs: 800, | ||
| }); |
Contributor
There was a problem hiding this comment.
Await session/start before relying on first-turn recall.
This hook is supposed to initialize session context, but it is fired in the background with an 800 ms timeout. before_agent_start can race ahead and query smart-search before the server has finished loading profiles/context, so the first recall block can miss the very data this PR is adding.
Suggested fix
if (lastHealthOk) {
- void callAgentMemory("session/start", {
+ await callAgentMemory("session/start", {
body: {
sessionId,
project: currentProject,
cwd: currentProject,
},
- timeoutMs: 800,
+ timeoutMs: 3000,
});
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| if (lastHealthOk) { | |
| void callAgentMemory("session/start", { | |
| body: { | |
| sessionId, | |
| project: currentProject, | |
| cwd: currentProject, | |
| }, | |
| timeoutMs: 800, | |
| }); | |
| if (lastHealthOk) { | |
| await callAgentMemory("session/start", { | |
| body: { | |
| sessionId, | |
| project: currentProject, | |
| cwd: currentProject, | |
| }, | |
| timeoutMs: 3000, | |
| }); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@integrations/pi/index.ts` around lines 311 - 319, The background call to
callAgentMemory("session/start") (with sessionId/currentProject) can race with
before_agent_start/smart-search and cause first-turn recall to miss newly-loaded
profiles; change this to await the call (or otherwise wait for its completion)
instead of firing it void with an 800ms timeout so session/context
initialization finishes before before_agent_start runs — update the
callAgentMemory invocation near the lastHealthOk branch to await the promise
(and remove or extend the timeout) or gate before_agent_start on the
callAgentMemory promise so the server has finished loading profiles/context
before first-turn recall executes.
Comment on lines
+332
to
354
| // Observe the prompt so the server knows what was asked. | ||
| trackPost(callAgentMemory("observe", { | ||
| body: { | ||
| hookType: "prompt_submit", | ||
| sessionId, | ||
| project: currentProject, | ||
| cwd: currentProject, | ||
| timestamp: new Date().toISOString(), | ||
| data: { prompt: lastPrompt }, | ||
| }, | ||
| timeoutMs: 3000, | ||
| })); | ||
|
|
||
| const result = await callAgentMemory<{ results?: SmartSearchResult[] }>("smart-search", { | ||
| body: { query: lastPrompt, limit: 5 }, | ||
| timeoutMs: 3000, | ||
| }); | ||
| const results = result?.results || []; | ||
| const recallBlock = results.length | ||
| ? [ | ||
| "Relevant long-term memory from agentmemory:", | ||
| formatSearchResults(results), | ||
| ].join("\n") | ||
| ? ["Relevant long-term memory from agentmemory:", formatSearchResults(results)].join("\n") | ||
| : ""; | ||
|
|
||
| await refreshStatus(ctx); |
Contributor
There was a problem hiding this comment.
Refresh health before the blocking recall path.
before_agent_start starts observe and then awaits smart-search before it checks health. When agentmemory is down, every turn pays the search timeout first and only marks the backend unhealthy afterward, which turns an outage into repeated prompt-start latency.
Suggested fix
pi.on("before_agent_start", async (event, ctx) => {
currentProject = event.systemPromptOptions.cwd || process.cwd();
lastPrompt = event.prompt?.trim() || "";
if (!lastPrompt) return;
+
+ await refreshStatus(ctx);
+ if (!lastHealthOk) {
+ return {
+ systemPrompt: [event.systemPrompt, TOOL_GUIDANCE].filter(Boolean).join("\n\n"),
+ };
+ }
// Observe the prompt so the server knows what was asked.
trackPost(callAgentMemory("observe", {
@@
const results = result?.results || [];
const recallBlock = results.length
? ["Relevant long-term memory from agentmemory:", formatSearchResults(results)].join("\n")
: "";
- await refreshStatus(ctx);
return {
systemPrompt: [event.systemPrompt, TOOL_GUIDANCE, recallBlock].filter(Boolean).join("\n\n"),
};
});🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@integrations/pi/index.ts` around lines 332 - 354, The health check
(refreshStatus(ctx)) must run before blocking agentmemory calls so an outage
doesn't cause repeated turn latency; in before_agent_start call
refreshStatus(ctx) prior to invoking trackPost/callAgentMemory("observe") and
before awaiting callAgentMemory("smart-search"), and if refreshStatus indicates
the backend is unhealthy, skip or short-circuit the smart-search/recall logic
(i.e., avoid awaiting callAgentMemory and set recallBlock = ""), keeping the
rest of the flow (lastPrompt, results/recallBlock usage) unchanged.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Add the missing lifecycle hooks to the pi integration, bringing it from
3 hooks to 6 — matching Codex CLI parity.
Problem
The pi integration only used
session_start,before_agent_start, andagent_end. Pi's extension API exposes 29 lifecycle events, but mostwere unused. As a result:
turn instead of per-tool-call granularity.
context.
Changes
session_startPOST /session/startbefore_agent_startPOST /smart-searchtool_resultPOST /observesession_before_compactPOST /summarizeagent_endPOST /observesession_shutdownPOST /summarize+POST /session/endAdditional improvements:
@mariozechner/pi-coding-agent→@earendil-works/pi-coding-agentdedup, avoids unnecessary HTTP round-trips during rapid tool bursts
observePayload()helper: reduces repetition acrossobserve calls
session_shutdownawaits both HTTP callssequentially instead of fire-and-forget, preventing data loss when the
process exits immediately after the hook returns
How to verify
agentmemoryintegrations/pi/to~/.pi/agent/extensions/agentmemory/http://localhost:3113— observations shouldappear for each tool call, not just one per agent turn
Note on system prompt injection
The
before_agent_starthook still mutates the system prompt per turn(same as before this PR). A follow-up could explore one-shot injection
via
pi.sendMessageatsession_startto preserve prefix cache, butthat is a separate concern and out of scope here.
Summary by CodeRabbit
New Features
Improvements