fix: implement extension-driven session switching (/pr, /work resume) - #582
fix: implement extension-driven session switching (/pr, /work resume)#582tongdu99 wants to merge 1 commit into
Conversation
Extension commands that replace the running session (`ctx.switchSession()`,
`ctx.newSession()` — used by /pr, /work, etc.) always failed in Pi Web: both
command-context actions were stubbed to return `{ cancelled: true }`, so
answering "Resume it?" immediately reported "Resume cancelled" and the fresh
session path reported "New session cancelled".
Pi Web has no AgentSessionRuntime, so rpc-manager now reimplements that flow:
- emit `session_before_switch` so only a real extension guard (e.g. the
dirty-worktree prompt) can cancel
- emit a `session_switch` event to connected browsers BEFORE teardown, since
the wrapper's listeners disappear with it
- retire any wrapper already holding the target, abort + `session_shutdown`
the current one, then start the replacement with `sessionStartEvent.reason`
`new`/`resume` so `session_start` handlers (branch checkout) run like in the TUI
- run the extension's `withSession()` against the replacement's
`createReplacedSessionContext()` after `waitUntilReady()`
Sessions created by `ctx.newSession()` are flushed to disk immediately, because
pi delays the first flush until an assistant message exists while the browser
navigates to the new session right away.
Client side, `useAgentSession` handles `session_switch` and follows it through
`onSessionSwitched` -> AppShell's existing select/hydrate/URL-replace path.
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1120d1d619
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| const openTarget = getRegistry().get(sessionId); | ||
| if (openTarget && openTarget !== this && openTarget.isAlive()) await openTarget.shutdown(); |
There was a problem hiding this comment.
When /work resumes a session whose wrapper is currently running in another tab, this unconditionally calls shutdown(), which destroys that wrapper without waiting for its active prompt to finish. The in-flight target run and its event stream can therefore be cut off merely by switching into it; check openTarget.isRunning() and cancel/reject the switch, or explicitly coordinate with the active run, rather than retiring a busy target.
Useful? React with 👍 / 👎.
| if (typeof targetSessionId === "string" && targetSessionId !== sessionIdRef.current) { | ||
| onSessionSwitched?.(targetSessionId); |
There was a problem hiding this comment.
When an extension calls ctx.switchSession() with the active session's own file, the server still destroys and recreates its wrapper, but this guard suppresses the navigation/remount because the ID is unchanged. The existing SSE remains attached to the destroyed wrapper, so a withSession run will not stream live events and may leave the UI stale until reconciliation; follow the switch even when the target ID equals the current ID so the client reconnects.
AGENTS.md reference: AGENTS.md:L129-L129
Useful? React with 👍 / 👎.
2 participants
Problem
Extension commands that replace the running session —
ctx.switchSession()andctx.newSession(), used by extensions such as/prand/work— always failed in Pi Web.Both command-context actions in
lib/rpc-manager.tswere stubbed to return{ cancelled: true }, so:Resume cancelledNew session cancelledFix
Pi Web does not use the SDK's
AgentSessionRuntime, sorpc-manager.tsnow reimplements that replacement flow:session_before_switch, so only a real extension guard (e.g. a dirty-worktree prompt) can cancelsession_switchevent to connected browsers before teardown, since the wrapper's event listeners disappear with itabort()+ emitsession_shutdownfor the current one, then start the replacement withsessionStartEvent.reasonnew/resumesosession_starthandlers (branch checkout, etc.) run exactly like they do in the TUIwithSession()against the replacement'screateReplacedSessionContext()afterwaitUntilReady()startRpcSession()gained two options for this:sessionManager(adopt an already-open manager) andsessionStartEvent(report the switch reason). Sessions created viactx.newSession()are flushed to disk immediately (persistSessionFileIfMissing) because pi otherwise delays the first flush until an assistant message exists, while the browser navigates to the new session right away.Client side,
useAgentSessionhandles the newsession_switchevent and follows it viaonSessionSwitched->AppShell's existing select / hydrate / URL-replace path (the same navigation used after a fork).Files
lib/rpc-manager.ts— replacement flow:emitBeforeSwitch,replaceWithNewSession,replaceWithExistingSession,finishSessionReplacement,teardownForReplacement,persistSessionFileIfMissing, new start optionslib/pi-types.ts— widenedExtensionRunnerLike.emit, addedhasHandlersandcreateReplacedSessionContexthooks/useAgentSession.ts,components/ChatWindow.tsx,components/AppShell.tsx—session_switchhandling /onSessionSwitchedplumbinglib/rpc-manager-session-switch.test.mjs— new source-assertion testsAGENTS.md— documents the flowVerification
tsc --noEmit,npm run lint,npm test(593 tests) all pass/pr <number>-> "Resume" now switches the browser to the existing PR session and checks out its branch