Skip to content

Feat/pi restore support#2791

Open
nikhilachale wants to merge 10 commits into
AgentWrapper:mainfrom
nikhilachale:feat/pi-restore-support
Open

Feat/pi restore support#2791
nikhilachale wants to merge 10 commits into
AgentWrapper:mainfrom
nikhilachale:feat/pi-restore-support

Conversation

@nikhilachale

Copy link
Copy Markdown
Collaborator

Summary

Implements true PI restore support by installing an AO-managed project-local PI extension that captures PI’s native session UUID and reports it through the existing ao hooks activity endpoint.

No daemon API or storage schema changes are required. PI restore continues to use the existing MetadataKeyAgentSessionID path and resumes with:

pi --append-system-prompt ... --session <native-pi-session-id>

Changes

  • Adds an embedded PI TypeScript extension at .pi/extensions/ao-activity.ts.

  • Installs the extension from backend/internal/adapters/agent/pi.GetAgentHooks.

  • Captures PI’s native session UUID via ctx.sessionManager.getSessionId() on session_start.

  • Reports PI hook events through ao hooks pi ...:

    • session-start captures metadata only.
    • before_agent_start maps to user-prompt-submit.
    • agent_settled maps to stop.
    • session_shutdown maps to session-end only for terminal quit shutdowns.
  • Adds sentinel-based overwrite protection so AO refuses to replace a non-AO file at .pi/extensions/ao-activity.ts.

  • Adds .pi/.gitignore management for the AO-installed extension.

  • Registers PI in activitydispatch.Derivers.

  • Adds a PI-specific activity deriver:

    • session-start: metadata-only, no activity state.
    • user-prompt-submit: active.
    • stop: idle.
    • session-end: exited.
  • Extends restore manager coverage so PI sessions with captured native ids use native resume.

Why

PI already supported native resume with pi --session , and AO’s PI adapter already read agentSessionId metadata in GetRestoreCommand. The missing piece was capturing PI’s native session UUID during normal AO-managed launches.

This change fills that metadata-capture path through PI’s official project-local extension API, avoiding any new daemon routes or schema changes.

Behavior

With this change:

  • New AO-managed PI sessions install .pi/extensions/ao-activity.ts.
  • When PI starts a native session, AO receives ao hooks pi session-start with:

{"session_id":""}

  • AO persists that UUID as agentSessionId.
  • Later, ao session restore and boot-time restore can resume PI with pi --session .
  • If no PI native UUID has been captured yet, PI keeps the same fallback behavior as other resumable agents: restore falls back to a fresh launch with the saved prompt when valid.

Safety Notes

  • The extension is AO-owned and guarded by a sentinel comment.
  • Re-running hook installation is idempotent.
  • A user-created file at the managed path is not overwritten.
  • Activity reporting is best-effort and must not break the PI session.
  • session_start is metadata-only so process startup does not incorrectly mark a turn active.
  • stop is emitted from agent_settled, not agent_end, to avoid reporting idle while PI may still auto-retry, compact, or process queued follow-ups.
  • session-end is only reported for session_shutdown.reason === "quit" to avoid marking AO exited during PI /new, /resume, /fork, or reload flows.

Tests

Added/updated coverage for:

  • PI hook installation writes .pi/extensions/ao-activity.ts.

  • Installed extension contains the AO sentinel.

  • Hook installation is idempotent.

  • Non-AO files at the managed path are not overwritten.

  • .pi/.gitignore includes the AO-managed extension.

  • Embedded extension includes all managed hook commands:

    • ao hooks pi session-start
    • ao hooks pi user-prompt-submit
    • ao hooks pi stop
    • ao hooks pi session-end
  • Embedded extension uses agent_settled for stop instead of agent_end.

  • Embedded extension gates session-end to terminal quit shutdowns.

  • PI activity derivation maps events correctly.

  • activitydispatch recognizes PI as hook-capable.

  • ao hooks pi session-start persists agentSessionId without an activity state.

  • Restore manager passes PI’s captured agentSessionId into native resume config.

@whoisasx whoisasx left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The overall implementation follows AO's adapter and hook boundaries, and the selected Pi lifecycle events, project-local extension location, native session-ID capture, and restore command are consistent with Pi's documented extension API. The focused Pi, activity-dispatch, CLI-hook, and restore tests pass, and the current GitHub checks are green.

The remaining functional blocker is Windows execution: the embedded extension depends on a POSIX shell even though the Pi adapter and AO support native Windows. GitHub also currently reports this PR as conflicting with main; the concrete conflict is in backend/internal/adapters/agent/activitydispatch/dispatch_test.go. Please make hook dispatch platform-aware, add an executed Windows hook test, resolve the conflict, and rerun checks.

const command = HOOK_COMMANDS[hookName];
if (!command) return;
try {
spawnSync("sh", ["-c", `if ! command -v ao >/dev/null 2>&1; then exit 0; fi; exec ${command}`], {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P1] This makes Pi restore/activity silently unavailable on supported native Windows installations. sh.exe is not guaranteed to exist on Windows, so every callback can fail before ao hooks runs; session-start then never persists agentSessionId, restore falls back to a fresh launch, and activity transitions are lost. Please use a platform-aware runner (ComSpec/cmd.exe on Windows and sh on Unix), preserve stdin/timeout behavior, inspect the spawn result, and add a Windows-executed test with a stub ao receiver.

…port

# Conflicts:
#	backend/internal/adapters/agent/activitydispatch/dispatch_test.go
@nikhilachale

Copy link
Copy Markdown
Collaborator Author
  • Added platform-aware Pi hook execution: %ComSpec%/cmd.exe on Windows, sh on Unix.
  • Preserved synchronous stdin and timeout behavior; spawn results are inspected.
  • Added a native Windows execution test using a stub ao.cmd receiver.

@nikhilachale
nikhilachale requested a review from whoisasx July 19, 2026 16:40
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants