feat(pi): add session resume shorthand (#113)

This commit is contained in:
banteg
2026-01-13 05:42:17 +04:00
committed by GitHub
parent c1205cd5a8
commit d4aad8e068
6 changed files with 198 additions and 18 deletions
+2 -2
View File
@@ -45,7 +45,7 @@ The core handler module containing:
**Key patterns:** **Key patterns:**
- Progress edits are best-effort and only run when new events arrive (Telegram outbox handles rate limiting/coalescing) - Progress edits are best-effort and only run when new events arrive (Telegram outbox handles rate limiting/coalescing)
- Resume tokens are runner-formatted command lines (e.g., `` `codex resume <token>` ``, `` `claude --resume <token>` ``, `` `pi --session <path>` ``) - Resume tokens are runner-formatted command lines (e.g., `` `codex resume <token>` ``, `` `claude --resume <token>` ``, `` `pi --session <id>` ``)
- Resume lines are stripped from the prompt before invoking the runner - Resume lines are stripped from the prompt before invoking the runner
- Errors/cancellation render final status while preserving resume tokens when known - Errors/cancellation render final status while preserving resume tokens when known
@@ -356,7 +356,7 @@ transport.send()/edit() final message, delete progress if needed
### Resume Flow ### Resume Flow
Same as above; auto-router polls all runners to extract resume tokens: Same as above; auto-router polls all runners to extract resume tokens:
- Router returns first matching token (e.g. `` `claude --resume <id>` `` routes to Claude, `` `pi --session <path>` `` routes to Pi) - Router returns first matching token (e.g. `` `claude --resume <id>` `` routes to Claude, `` `pi --session <id>` `` routes to Pi)
- Selected runner spawns with resume (e.g. `codex exec --json resume <token> -`, `pi --print --mode json --session <path> <prompt>`) - Selected runner spawns with resume (e.g. `codex exec --json resume <token> -`, `pi --print --mode json --session <path> <prompt>`)
- Per-token lock serializes concurrent resumes on the same thread - Per-token lock serializes concurrent resumes on the same thread
+4 -2
View File
@@ -31,13 +31,15 @@ Provide the **`pi`** engine backend so Takopi can:
Takopi appends a **single backticked** resume line at the end of the message, like: Takopi appends a **single backticked** resume line at the end of the message, like:
```text ```text
`pi --session /home/user/.pi/agent/sessions/--repo--/2026-01-02T12-34-56-789Z_abcd.jsonl` `pi --session ccd569e0`
``` ```
Notes: Notes:
* `pi --resume/-r` opens an interactive session picker, so Takopi uses `--session <path>` instead. * `pi --resume/-r` opens an interactive session picker, so Takopi uses `--session <path>` instead.
* The resume token is the **session file path** (JSONL), treated as an opaque string. * The resume token is the **session id** (short prefix), derived from the first JSON
object in the session file. If the id cannot be read, Takopi falls back to the
session file path.
* If the path contains spaces, the runner will quote it. * If the path contains spaces, the runner will quote it.
### Non-interactive runs ### Non-interactive runs
+3 -2
View File
@@ -30,10 +30,11 @@ Notes:
- Canonical resume line (embedded in chat): - Canonical resume line (embedded in chat):
``` ```
`pi --session <path>` `pi --session <id>`
``` ```
The token is the **session JSONL file path**. The token is the **short session id**, derived from the first JSON object in the
session file. If the id cannot be read, Takopi falls back to the session file path.
Why not `--resume`? Why not `--resume`?
- `--resume/-r` opens an interactive session picker; it does not accept a - `--resume/-r` opens an interactive session picker; it does not accept a
+1 -1
View File
@@ -41,7 +41,7 @@ The canonical ResumeLine embedded in chat MUST be the engines CLI resume comm
- `codex resume <id>` - `codex resume <id>`
- `claude --resume <id>` - `claude --resume <id>`
- `pi --session <path>` - `pi --session <token>`
ResumeLine MUST resume the interactive session when the engine offers both interactive and headless modes. It MUST NOT point to a headless/batch command that requires a new prompt (e.g., a `run` subcommand that errors without a message). ResumeLine MUST resume the interactive session when the engine offers both interactive and headless modes. It MUST NOT point to a headless/batch command that requires a new prompt (e.g., a `run` subcommand that errors without a message).
+98 -3
View File
@@ -1,7 +1,9 @@
from __future__ import annotations from __future__ import annotations
import json
import os import os
import re import re
from collections.abc import AsyncIterator
from dataclasses import dataclass, field from dataclasses import dataclass, field
from datetime import datetime, UTC from datetime import datetime, UTC
from pathlib import Path, PurePath from pathlib import Path, PurePath
@@ -36,10 +38,14 @@ ENGINE: EngineId = "pi"
_RESUME_RE = re.compile(r"(?im)^\s*`?pi\s+--session\s+(?P<token>.+?)`?\s*$") _RESUME_RE = re.compile(r"(?im)^\s*`?pi\s+--session\s+(?P<token>.+?)`?\s*$")
_SESSION_ID_PREFIX_LEN = 8
@dataclass(slots=True) @dataclass(slots=True)
class PiStreamState: class PiStreamState:
resume: ResumeToken resume: ResumeToken
session_path: str | None = None
allow_id_promotion: bool = False
pending_actions: dict[str, Action] = field(default_factory=dict) pending_actions: dict[str, Action] = field(default_factory=dict)
last_assistant_text: str | None = None last_assistant_text: str | None = None
last_assistant_error: str | None = None last_assistant_error: str | None = None
@@ -48,6 +54,69 @@ class PiStreamState:
note_seq: int = 0 note_seq: int = 0
def _looks_like_session_path(token: str) -> bool:
if not token:
return False
if token.endswith(".jsonl"):
return True
if "/" in token or "\\" in token:
return True
return token.startswith("~")
def _short_session_id(session_id: str) -> str:
if not session_id:
return session_id
if "-" in session_id:
return session_id.split("-", 1)[0]
if len(session_id) > _SESSION_ID_PREFIX_LEN:
return session_id[:_SESSION_ID_PREFIX_LEN]
return session_id
def _session_id_from_line(line: str) -> str | None:
try:
data = json.loads(line)
except json.JSONDecodeError:
return None
if not isinstance(data, dict):
return None
event_type = data.get("type")
if event_type is not None and event_type != "session":
return None
session_id = data.get("id")
if isinstance(session_id, str) and session_id:
return _short_session_id(session_id)
return None
def _session_id_from_path(path: Path) -> str | None:
path = path.expanduser()
try:
with path.open("r", encoding="utf-8") as handle:
for raw_line in handle:
line = raw_line.strip()
if not line:
continue
return _session_id_from_line(line)
except OSError:
return None
return None
def _maybe_promote_session_id(state: PiStreamState) -> None:
if not state.allow_id_promotion:
return
session_path = state.session_path
if not session_path:
return
if state.resume.value != session_path:
return
session_id = _session_id_from_path(Path(session_path))
if session_id:
state.resume = ResumeToken(engine=ENGINE, value=session_id)
def _action_event( def _action_event(
*, *,
phase: ActionPhase, phase: ActionPhase,
@@ -117,6 +186,7 @@ def translate_pi_event(
state: PiStreamState, state: PiStreamState,
) -> list[TakopiEvent]: ) -> list[TakopiEvent]:
out: list[TakopiEvent] = [] out: list[TakopiEvent] = []
_maybe_promote_session_id(state)
if not state.started: if not state.started:
out.append( out.append(
StartedEvent( StartedEvent(
@@ -240,6 +310,11 @@ class PiRunner(ResumeTokenMixin, JsonlSubprocessRunner):
raise RuntimeError(f"resume token is for engine {token.engine!r}") raise RuntimeError(f"resume token is for engine {token.engine!r}")
return f"`pi --session {self._quote_token(token.value)}`" return f"`pi --session {self._quote_token(token.value)}`"
def run(
self, prompt: str, resume: ResumeToken | None
) -> AsyncIterator[TakopiEvent]:
return super().run(prompt, self._normalize_resume_token(resume))
def extract_resume(self, text: str | None) -> ResumeToken | None: def extract_resume(self, text: str | None) -> ResumeToken | None:
if not text: if not text:
return None return None
@@ -254,8 +329,24 @@ class PiRunner(ResumeTokenMixin, JsonlSubprocessRunner):
found = token found = token
if not found: if not found:
return None return None
if _looks_like_session_path(found):
session_id = _session_id_from_path(Path(found))
if session_id:
found = session_id
return ResumeToken(engine=self.engine, value=found) return ResumeToken(engine=self.engine, value=found)
def _normalize_resume_token(self, resume: ResumeToken | None) -> ResumeToken | None:
if resume is None:
return None
if resume.engine != ENGINE:
return resume
if not _looks_like_session_path(resume.value):
return resume
session_id = _session_id_from_path(Path(resume.value))
if session_id:
return ResumeToken(engine=ENGINE, value=session_id)
return resume
def command(self) -> str: def command(self) -> str:
return "pi" return "pi"
@@ -294,9 +385,13 @@ class PiRunner(ResumeTokenMixin, JsonlSubprocessRunner):
if resume is None: if resume is None:
session_path = self._new_session_path() session_path = self._new_session_path()
token = ResumeToken(engine=ENGINE, value=session_path) token = ResumeToken(engine=ENGINE, value=session_path)
else: return PiStreamState(
token = resume resume=token,
return PiStreamState(resume=token) session_path=session_path,
allow_id_promotion=True,
)
session_path = resume.value if _looks_like_session_path(resume.value) else None
return PiStreamState(resume=resume, session_path=session_path)
def translate( def translate(
self, self,
+90 -8
View File
@@ -29,22 +29,24 @@ def _load_fixture(name: str) -> list[pi_schema.PiEvent]:
return events return events
def test_pi_resume_format_and_extract() -> None: def test_pi_resume_format_and_extract(tmp_path: Path) -> None:
runner = PiRunner( runner = PiRunner(
extra_args=[], extra_args=[],
model=None, model=None,
provider=None, provider=None,
) )
token = ResumeToken(engine=ENGINE, value="/tmp/pi/session.jsonl") session_path = tmp_path / "session.jsonl"
token = ResumeToken(engine=ENGINE, value=str(session_path))
assert runner.format_resume(token) == "`pi --session /tmp/pi/session.jsonl`" assert runner.format_resume(token) == f"`pi --session {session_path}`"
assert runner.extract_resume("`pi --session /tmp/pi/session.jsonl`") == token assert runner.extract_resume(f"`pi --session {session_path}`") == token
assert runner.extract_resume('pi --session "/tmp/pi/session.jsonl"') == token assert runner.extract_resume(f'pi --session "{session_path}"') == token
assert runner.extract_resume("`codex resume sid`") is None assert runner.extract_resume("`codex resume sid`") is None
spaced = ResumeToken(engine=ENGINE, value="/tmp/pi session.jsonl") spaced_path = tmp_path / "pi session.jsonl"
assert runner.format_resume(spaced) == '`pi --session "/tmp/pi session.jsonl"`' spaced = ResumeToken(engine=ENGINE, value=str(spaced_path))
assert runner.extract_resume('`pi --session "/tmp/pi session.jsonl"`') == spaced assert runner.format_resume(spaced) == f'`pi --session "{spaced_path}"`'
assert runner.extract_resume(f'`pi --session "{spaced_path}"`') == spaced
def test_translate_success_fixture() -> None: def test_translate_success_fixture() -> None:
@@ -97,6 +99,86 @@ def test_translate_error_fixture() -> None:
assert completed.answer == "Request failed." assert completed.answer == "Request failed."
def test_session_id_promotion_from_file(tmp_path: Path) -> None:
session_path = tmp_path / "session.jsonl"
session_path.write_text(
'{"type":"session","version":3,'
'"id":"ccd569e0-4e1b-4c7d-a981-637ed4107310",'
'"timestamp":"2026-01-13T00:33:34.702Z",'
'"cwd":"/tmp"}\n',
encoding="utf-8",
)
runner = PiRunner(
extra_args=[],
model=None,
provider=None,
)
with patch(
"takopi.runners.pi.PiRunner._new_session_path",
return_value=str(session_path),
):
state = runner.new_state("prompt", None)
events = translate_pi_event(
pi_schema.AgentStart(), title="pi", meta=None, state=state
)
started = next(evt for evt in events if isinstance(evt, StartedEvent))
assert started.resume.value == "ccd569e0"
def test_extract_resume_prefers_session_id(tmp_path: Path) -> None:
session_path = tmp_path / "session.jsonl"
session_path.write_text(
'{"type":"session","version":3,'
'"id":"ccd569e0-4e1b-4c7d-a981-637ed4107310",'
'"timestamp":"2026-01-13T00:33:34.702Z",'
'"cwd":"/tmp"}\n',
encoding="utf-8",
)
runner = PiRunner(
extra_args=[],
model=None,
provider=None,
)
token = runner.extract_resume(f"pi --session {session_path}")
assert token is not None
assert token.value == "ccd569e0"
@pytest.mark.anyio
async def test_run_normalizes_resume_path(tmp_path: Path) -> None:
session_path = tmp_path / "session.jsonl"
session_path.write_text(
'{"type":"session","version":3,'
'"id":"ccd569e0-4e1b-4c7d-a981-637ed4107310",'
'"timestamp":"2026-01-13T00:33:34.702Z",'
'"cwd":"/tmp"}\n',
encoding="utf-8",
)
runner = PiRunner(
extra_args=[],
model=None,
provider=None,
)
seen_resume: ResumeToken | None = None
async def run_stub(_prompt: str, resume: ResumeToken | None):
nonlocal seen_resume
seen_resume = resume
yield CompletedEvent(
engine=ENGINE,
resume=resume,
ok=True,
answer="ok",
)
runner.run_impl = run_stub # type: ignore[assignment]
resume = ResumeToken(engine=ENGINE, value=str(session_path))
async for _event in runner.run("test", resume):
pass
assert seen_resume is not None
assert seen_resume.value == "ccd569e0"
@pytest.mark.anyio @pytest.mark.anyio
async def test_run_serializes_same_session() -> None: async def test_run_serializes_same_session() -> None:
runner = PiRunner( runner = PiRunner(