# Session Persistence After Plugin Commands — Implementation Plan

> **For Claude:** REQUIRED SUB-SKILL: Use superpowers:executing-plans to implement this plan task-by-task.

**Goal:** Allow plugin commands to persist session cookie/header updates back to the encrypted cache, via `@command(saves_session=True)` and/or `ctx.save_session()`.

**Architecture:** Add `saves_session` flag to `CommandMetadata`/`CommandSpec`/`@command()`. Make `CommandContext` non-frozen with a `save_session()` method that sets a dirty flag. After successful command execution, the `_create_plugin_command` callback checks the flag and persists cookies back via a new `update_session_cookies()` function in `cache.py` that does load-modify-save to preserve browser metadata.

**Tech Stack:** Python 3.13, dataclasses, dill (pickle), pytest, typer/click

---

### Task 1: Add `saves_session` to `CommandMetadata` and `CommandSpec`

**Files:**
- Modify: `src/graftpunk/plugins/cli_plugin.py:91-98` (CommandMetadata)
- Modify: `src/graftpunk/plugins/cli_plugin.py:142-164` (CommandSpec)
- Test: `tests/unit/test_cli_plugin.py`

**Step 1: Write failing tests**

Add to `tests/unit/test_cli_plugin.py`:

```python
class TestSavesSessionField:
    """Tests for saves_session field on CommandMetadata and CommandSpec."""

    def test_command_metadata_default_false(self) -> None:
        meta = CommandMetadata(name="test", help_text="help")
        assert meta.saves_session is False

    def test_command_metadata_explicit_true(self) -> None:
        meta = CommandMetadata(name="test", help_text="help", saves_session=True)
        assert meta.saves_session is True

    def test_command_spec_default_false(self) -> None:
        spec = CommandSpec(name="test", handler=lambda ctx: None)
        assert spec.saves_session is False

    def test_command_spec_explicit_true(self) -> None:
        spec = CommandSpec(name="test", handler=lambda ctx: None, saves_session=True)
        assert spec.saves_session is True
```

**Step 2: Run tests to verify they fail**

Run: `uv run pytest tests/unit/test_cli_plugin.py::TestSavesSessionField -v`
Expected: FAIL with `TypeError: __init__() got an unexpected keyword argument 'saves_session'`

**Step 3: Implement**

In `cli_plugin.py`, add `saves_session: bool = False` to both dataclasses:

`CommandMetadata` (after line 98):
```python
@dataclass(frozen=True)
class CommandMetadata:
    """Metadata stored on @command-decorated methods."""

    name: str
    help_text: str
    params: tuple[PluginParamSpec, ...] = ()
    parent: type | None = None
    requires_session: bool | None = None
    saves_session: bool = False
```

`CommandSpec` (after line 153):
```python
@dataclass(frozen=True)
class CommandSpec:
    """Specification for a single CLI command."""

    name: str
    handler: Callable[..., Any]
    help_text: str = ""
    params: tuple[PluginParamSpec, ...] = ()
    timeout: float | None = None
    max_retries: int = 0
    rate_limit: float | None = None
    requires_session: bool | None = None
    group: str | None = None
    saves_session: bool = False
```

**Step 4: Run tests to verify they pass**

Run: `uv run pytest tests/unit/test_cli_plugin.py::TestSavesSessionField -v`
Expected: PASS

**Step 5: Commit**

```bash
git add src/graftpunk/plugins/cli_plugin.py tests/unit/test_cli_plugin.py
git commit -m "feat: add saves_session field to CommandMetadata and CommandSpec

Co-authored-by: stavxyz <hi@stav.xyz>"
```

---

### Task 2: Add `saves_session` parameter to `@command()` decorator

**Files:**
- Modify: `src/graftpunk/plugins/cli_plugin.py:357-422` (`command()` function)
- Modify: `src/graftpunk/plugins/cli_plugin.py:530-544` (`_build_command_spec()`)
- Test: `tests/unit/test_cli_plugin.py`

**Step 1: Write failing tests**

Add to `tests/unit/test_cli_plugin.py`:

```python
class TestCommandDecoratorSavesSession:
    """Tests for saves_session parameter on @command decorator."""

    def test_default_false(self) -> None:
        @command(help="test")
        def my_cmd(self, ctx: Any) -> None:
            pass

        assert my_cmd._command_meta.saves_session is False

    def test_explicit_true(self) -> None:
        @command(help="test", saves_session=True)
        def my_cmd(self, ctx: Any) -> None:
            pass

        assert my_cmd._command_meta.saves_session is True

    def test_saves_session_flows_to_command_spec(self) -> None:
        """saves_session propagates from CommandMetadata through _build_command_spec."""

        class TestPlugin(SitePlugin):
            site_name = "test"
            session_name = "test"
            help_text = "test"

            @command(help="test", saves_session=True)
            def my_cmd(self, ctx: Any) -> dict:
                return {}

        plugin = TestPlugin()
        specs = plugin.get_commands()
        my_spec = next(s for s in specs if s.name == "my_cmd")
        assert my_spec.saves_session is True
```

**Step 2: Run tests to verify they fail**

Run: `uv run pytest tests/unit/test_cli_plugin.py::TestCommandDecoratorSavesSession -v`
Expected: FAIL with `TypeError: command() got an unexpected keyword argument 'saves_session'`

**Step 3: Implement**

Add `saves_session` parameter to the `command()` function signature (line 357):

```python
def command(
    help: str = "",  # noqa: A002 - shadows builtin but matches typer convention
    params: list[PluginParamSpec] | None = None,
    parent: type | None = None,
    requires_session: bool | None = None,
    saves_session: bool = False,
) -> Callable[..., Any]:
```

In the function decorator path (line 413), pass it through:

```python
            target._command_meta = CommandMetadata(
                name=target.__name__,
                help_text=help,
                params=tuple(params) if params else (),
                parent=parent,
                requires_session=requires_session,
                saves_session=saves_session,
            )
```

In the class decorator's auto-discovery path (line 403), keep `saves_session=False` (the default — auto-discovered group methods don't save sessions).

In `_build_command_spec()` (line 537), pass it through:

```python
        return CommandSpec(
            name=meta.name,
            handler=handler,
            help_text=meta.help_text,
            params=self._resolve_params(meta, handler),
            requires_session=meta.requires_session,
            group=group,
            saves_session=meta.saves_session,
        )
```

**Step 4: Run tests to verify they pass**

Run: `uv run pytest tests/unit/test_cli_plugin.py::TestCommandDecoratorSavesSession -v`
Expected: PASS

**Step 5: Commit**

```bash
git add src/graftpunk/plugins/cli_plugin.py tests/unit/test_cli_plugin.py
git commit -m "feat: add saves_session parameter to @command() decorator

Co-authored-by: stavxyz <hi@stav.xyz>"
```

---

### Task 3: Make `CommandContext` non-frozen with `save_session()` method

**Files:**
- Modify: `src/graftpunk/plugins/cli_plugin.py:105-127` (CommandContext)
- Test: `tests/unit/test_cli_plugin.py`

**Step 1: Write failing tests**

Add to `tests/unit/test_cli_plugin.py`:

```python
class TestCommandContextSaveSession:
    """Tests for CommandContext.save_session() method."""

    def test_session_not_dirty_by_default(self) -> None:
        ctx = CommandContext(
            session=MagicMock(),
            plugin_name="test",
            command_name="cmd",
            api_version=1,
            _session_name="testsession",
        )
        assert ctx._session_dirty is False

    def test_save_session_sets_dirty_flag(self) -> None:
        ctx = CommandContext(
            session=MagicMock(),
            plugin_name="test",
            command_name="cmd",
            api_version=1,
            _session_name="testsession",
        )
        ctx.save_session()
        assert ctx._session_dirty is True

    def test_save_session_requires_session_name(self) -> None:
        ctx = CommandContext(
            session=MagicMock(),
            plugin_name="test",
            command_name="cmd",
            api_version=1,
        )
        with pytest.raises(ValueError, match="No session name"):
            ctx.save_session()

    def test_session_name_default_empty(self) -> None:
        ctx = CommandContext(
            session=MagicMock(),
            plugin_name="test",
            command_name="cmd",
            api_version=1,
        )
        assert ctx._session_name == ""
```

**Step 2: Run tests to verify they fail**

Run: `uv run pytest tests/unit/test_cli_plugin.py::TestCommandContextSaveSession -v`
Expected: FAIL (no `_session_name` field, no `save_session` method)

**Step 3: Implement**

Change `CommandContext` from `frozen=True` to plain `@dataclass`. Add `_session_name` and `_session_dirty` private fields with `repr=False`. Add `save_session()` method:

```python
@dataclass
class CommandContext:
    """Execution context passed to command handlers."""

    session: requests.Session
    plugin_name: str
    command_name: str
    api_version: int
    base_url: str = ""
    config: PluginConfig | None = None
    observe: ObservabilityContext = field(default_factory=NoOpObservabilityContext)
    _session_name: str = field(default="", repr=False)
    _session_dirty: bool = field(default=False, repr=False, init=False)

    def __post_init__(self) -> None:
        if self.api_version not in SUPPORTED_API_VERSIONS:
            raise ValueError(
                f"api_version {self.api_version} not supported. "
                f"Supported: {sorted(SUPPORTED_API_VERSIONS)}"
            )
        if not self.plugin_name:
            raise ValueError("plugin_name must be non-empty")
        if not self.command_name:
            raise ValueError("command_name must be non-empty")

    def save_session(self) -> None:
        """Mark the session for persistence after the command completes.

        The actual save happens in the command wrapper after successful
        completion, not here. This ensures failed commands don't corrupt
        the cached session.

        Raises:
            ValueError: If no session name is configured.
        """
        if not self._session_name:
            raise ValueError(
                "No session name configured on this context. "
                "Cannot save session without a name."
            )
        self._session_dirty = True
```

**Important:** `_session_dirty` uses `init=False` so it's always initialized to `False` and cannot be set by the constructor. `_session_name` uses `init=True` (the default) so it can be set during construction.

**Step 4: Run tests to verify they pass**

Run: `uv run pytest tests/unit/test_cli_plugin.py::TestCommandContextSaveSession -v`
Expected: PASS

Also run existing CommandContext tests to ensure nothing broke:
Run: `uv run pytest tests/unit/test_cli_plugin.py::TestCommandContext -v`
Run: `uv run pytest tests/unit/test_cli_plugin.py::TestCommandContextValidation -v`
Run: `uv run pytest tests/unit/test_cli_plugin.py::TestCommandContextExpansion -v`
Expected: PASS (existing tests don't depend on frozen behavior)

**Step 5: Commit**

```bash
git add src/graftpunk/plugins/cli_plugin.py tests/unit/test_cli_plugin.py
git commit -m "feat: add save_session() method to CommandContext

Makes CommandContext non-frozen, adds _session_name and _session_dirty
fields, and save_session() method that marks the session for persistence
after successful command completion.

Co-authored-by: stavxyz <hi@stav.xyz>"
```

---

### Task 4: Add `update_session_cookies()` to `cache.py`

**Files:**
- Modify: `src/graftpunk/cache.py` (add new function after `load_session_for_api`)
- Test: `tests/unit/test_cache.py`

**Step 1: Write failing tests**

Add to `tests/unit/test_cache.py`:

```python
class TestUpdateSessionCookies:
    """Tests for update_session_cookies() — persisting API session changes."""

    def test_updates_cookies_on_cached_session(self, tmp_path: Path) -> None:
        """Cookies from the API session are merged into the cached session."""
        from unittest.mock import MagicMock, patch

        import requests

        from graftpunk.cache import update_session_cookies

        # Create a mock cached session with original cookies
        cached_session = MagicMock()
        cached_session.cookies = requests.cookies.RequestsCookieJar()
        cached_session.cookies.set("original", "value")
        cached_session.headers = {}

        # API session has new cookies
        api_session = requests.Session()
        api_session.cookies.set("original", "updated")
        api_session.cookies.set("new_cookie", "new_value")

        with (
            patch("graftpunk.cache.load_session") as mock_load,
            patch("graftpunk.cache.cache_session") as mock_cache,
        ):
            mock_load.return_value = cached_session
            update_session_cookies(api_session, "testsession")

            mock_load.assert_called_once_with("testsession")
            mock_cache.assert_called_once_with(cached_session, "testsession")
            # Verify cookies were updated on the original session object
            cached_session.cookies.update.assert_called_once_with(api_session.cookies)

    def test_load_failure_logs_warning_and_returns(self) -> None:
        """If loading the cached session fails, log a warning and return silently."""
        from unittest.mock import patch

        import requests

        from graftpunk.cache import update_session_cookies

        api_session = requests.Session()

        with patch("graftpunk.cache.load_session", side_effect=Exception("corrupt")):
            # Should not raise
            update_session_cookies(api_session, "badsession")

    def test_cache_failure_logs_warning_and_returns(self) -> None:
        """If re-caching fails, log a warning and return silently."""
        from unittest.mock import MagicMock, patch

        import requests

        from graftpunk.cache import update_session_cookies

        cached_session = MagicMock()
        cached_session.cookies = requests.cookies.RequestsCookieJar()
        api_session = requests.Session()

        with (
            patch("graftpunk.cache.load_session", return_value=cached_session),
            patch("graftpunk.cache.cache_session", side_effect=OSError("disk full")),
        ):
            # Should not raise
            update_session_cookies(api_session, "testsession")
```

**Step 2: Run tests to verify they fail**

Run: `uv run pytest tests/unit/test_cache.py::TestUpdateSessionCookies -v`
Expected: FAIL with `ImportError: cannot import name 'update_session_cookies'`

**Step 3: Implement**

Add to `cache.py` after `load_session_for_api()` (after line 386):

```python
def update_session_cookies(api_session: requests.Session, session_name: str) -> None:
    """Persist an API session's cookies back to the session cache.

    Loads the original cached session (preserving browser metadata),
    updates its cookies from the API session, and saves it back.
    This is best-effort — failures are logged but do not raise.

    Args:
        api_session: The API session with potentially updated cookies.
        session_name: Name of the cached session to update.
    """
    try:
        original = load_session(session_name)
    except Exception as exc:  # noqa: BLE001 — best-effort save
        LOG.warning(
            "session_save_skipped_load_failed",
            session_name=session_name,
            error=str(exc),
        )
        return

    try:
        original.cookies.update(api_session.cookies)
        cache_session(original, session_name)
        LOG.info("session_cookies_updated", session_name=session_name)
    except Exception as exc:  # noqa: BLE001 — best-effort save
        LOG.warning(
            "session_save_failed",
            session_name=session_name,
            error=str(exc),
        )
```

**Step 4: Run tests to verify they pass**

Run: `uv run pytest tests/unit/test_cache.py::TestUpdateSessionCookies -v`
Expected: PASS

**Step 5: Commit**

```bash
git add src/graftpunk/cache.py tests/unit/test_cache.py
git commit -m "feat: add update_session_cookies() for persisting API session changes

Loads the original cached BrowserSession, merges cookies from the API
session, and saves it back. Preserves all browser metadata. Best-effort
— failures are logged but don't raise.

Co-authored-by: stavxyz <hi@stav.xyz>"
```

---

### Task 5: Wire up session persistence in `_create_plugin_command`

**Files:**
- Modify: `src/graftpunk/cli/plugin_commands.py:365-393` (callback function)
- Test: `tests/unit/test_plugin_commands.py`

**Step 1: Write failing tests**

Add to `tests/unit/test_plugin_commands.py`:

```python
class TestSessionPersistence:
    """Tests for saves_session decorator and ctx.save_session() persistence."""

    def test_saves_session_decorator_persists_after_success(
        self, isolated_config: Path
    ) -> None:
        """@command(saves_session=True) triggers update_session_cookies after handler."""

        class PersistPlugin(SitePlugin):
            site_name = "persisttest"
            session_name = "persistsession"
            help_text = "test"

            @command(help="Save test", saves_session=True)
            def save_cmd(self, ctx: Any) -> dict:
                return {"ok": True}

        plugin = PersistPlugin()
        app = GraftpunkApp()

        with (
            patch("graftpunk.cli.plugin_commands.discover_site_plugins", return_value={}),
            patch(
                "graftpunk.cli.plugin_commands.discover_python_plugins",
                return_value=PythonDiscoveryResult(plugins={"persisttest": plugin}),
            ),
            patch.object(plugin, "get_session", return_value=requests.Session()),
            patch(
                "graftpunk.cli.plugin_commands.update_session_cookies"
            ) as mock_update,
        ):
            app.register_plugins()
            result = runner.invoke(app.app, ["persisttest", "save-cmd"])
            assert result.exit_code == 0
            mock_update.assert_called_once()
            call_args = mock_update.call_args
            assert call_args[1]["session_name"] == "persistsession" or call_args[0][1] == "persistsession"

    def test_ctx_save_session_persists_after_success(
        self, isolated_config: Path
    ) -> None:
        """ctx.save_session() triggers update_session_cookies after handler."""
        from graftpunk.plugins import CommandContext

        class ExplicitPlugin(SitePlugin):
            site_name = "explicittest"
            session_name = "explicitsession"
            help_text = "test"

            @command(help="Explicit save")
            def explicit_cmd(self, ctx: CommandContext) -> dict:
                ctx.save_session()
                return {"saved": True}

        plugin = ExplicitPlugin()
        app = GraftpunkApp()

        with (
            patch("graftpunk.cli.plugin_commands.discover_site_plugins", return_value={}),
            patch(
                "graftpunk.cli.plugin_commands.discover_python_plugins",
                return_value=PythonDiscoveryResult(plugins={"explicittest": plugin}),
            ),
            patch.object(plugin, "get_session", return_value=requests.Session()),
            patch(
                "graftpunk.cli.plugin_commands.update_session_cookies"
            ) as mock_update,
        ):
            app.register_plugins()
            result = runner.invoke(app.app, ["explicittest", "explicit-cmd"])
            assert result.exit_code == 0
            mock_update.assert_called_once()

    def test_no_save_when_not_requested(self, isolated_config: Path) -> None:
        """Without saves_session or ctx.save_session(), no persistence happens."""

        class NoSavePlugin(SitePlugin):
            site_name = "nosavetest"
            session_name = "nosavesession"
            help_text = "test"

            @command(help="No save")
            def no_save_cmd(self, ctx: Any) -> dict:
                return {"ok": True}

        plugin = NoSavePlugin()
        app = GraftpunkApp()

        with (
            patch("graftpunk.cli.plugin_commands.discover_site_plugins", return_value={}),
            patch(
                "graftpunk.cli.plugin_commands.discover_python_plugins",
                return_value=PythonDiscoveryResult(plugins={"nosavetest": plugin}),
            ),
            patch.object(plugin, "get_session", return_value=requests.Session()),
            patch(
                "graftpunk.cli.plugin_commands.update_session_cookies"
            ) as mock_update,
        ):
            app.register_plugins()
            result = runner.invoke(app.app, ["nosavetest", "no-save-cmd"])
            assert result.exit_code == 0
            mock_update.assert_not_called()

    def test_no_save_on_command_failure(self, isolated_config: Path) -> None:
        """Session is NOT saved when the command raises an exception."""

        class FailPlugin(SitePlugin):
            site_name = "failtest"
            session_name = "failsession"
            help_text = "test"

            @command(help="Failing cmd", saves_session=True)
            def fail_cmd(self, ctx: Any) -> dict:
                raise RuntimeError("command failed")

        plugin = FailPlugin()
        app = GraftpunkApp()

        with (
            patch("graftpunk.cli.plugin_commands.discover_site_plugins", return_value={}),
            patch(
                "graftpunk.cli.plugin_commands.discover_python_plugins",
                return_value=PythonDiscoveryResult(plugins={"failtest": plugin}),
            ),
            patch.object(plugin, "get_session", return_value=requests.Session()),
            patch(
                "graftpunk.cli.plugin_commands.update_session_cookies"
            ) as mock_update,
        ):
            app.register_plugins()
            result = runner.invoke(app.app, ["failtest", "fail-cmd"])
            assert result.exit_code != 0
            mock_update.assert_not_called()
```

**Step 2: Run tests to verify they fail**

Run: `uv run pytest tests/unit/test_plugin_commands.py::TestSessionPersistence -v`
Expected: FAIL (no `update_session_cookies` import, no wiring)

**Step 3: Implement**

In `plugin_commands.py`, add the import at the top:

```python
from graftpunk.cache import update_session_cookies
```

In `_create_plugin_command`, modify the `CommandContext` construction (line 366) to pass `_session_name`:

```python
            ctx = CommandContext(
                session=session,
                plugin_name=plugin.site_name,
                command_name=cmd_spec.name,
                api_version=plugin.api_version,
                base_url=getattr(plugin, "base_url", ""),
                config=getattr(plugin, "_plugin_config", None),
                observe=observe_ctx,
                _session_name=plugin.session_name if needs_session else "",
            )
```

After the handler returns successfully (after line 393, but before the except blocks), add session persistence. The key is that `format_output` runs after `result` is obtained, and we save after the handler succeeds but the save must happen before or after formatting. Place it right after `result` is obtained (after line 389 in the 403 retry path and line 377 in the normal path):

Replace lines 376-393 with:

```python
            try:
                result = _execute_with_limits(cmd_spec.handler, ctx, cmd_spec, **kwargs)
            except requests.exceptions.HTTPError as exc:
                if (
                    exc.response is not None
                    and exc.response.status_code == 403
                    and token_config is not None
                ):
                    from graftpunk.tokens import clear_cached_tokens
                    from graftpunk.tokens import prepare_session as _prep

                    clear_cached_tokens(session)
                    _prep(session, token_config, getattr(plugin, "base_url", ""))
                    result = _execute_with_limits(cmd_spec.handler, ctx, cmd_spec, **kwargs)
                else:
                    raise

            # Persist session if requested (decorator or explicit ctx.save_session())
            if (cmd_spec.saves_session or ctx._session_dirty) and needs_session:
                update_session_cookies(session, plugin.session_name)

            format_output(result, output_format, _format_console)
```

**Step 4: Run tests to verify they pass**

Run: `uv run pytest tests/unit/test_plugin_commands.py::TestSessionPersistence -v`
Expected: PASS

Run full suite: `uv run pytest tests/ -v`
Expected: All pass

**Step 5: Commit**

```bash
git add src/graftpunk/cli/plugin_commands.py tests/unit/test_plugin_commands.py
git commit -m "feat: wire up session persistence in plugin command execution

After successful command execution, checks saves_session decorator flag
and ctx._session_dirty flag. If either is set, calls
update_session_cookies() to persist cookie changes back to cache.
Session is NOT saved on command failure.

Co-authored-by: stavxyz <hi@stav.xyz>"
```

---

### Task 6: Final verification and lint cleanup

**Files:**
- All modified files from previous tasks

**Step 1: Run full test suite**

Run: `uv run pytest tests/ -v`
Expected: All pass (1285+ tests)

**Step 2: Run lint**

Run: `uvx ruff check .`
Expected: Clean

**Step 3: Run format check**

Run: `uvx ruff format --check .`
Expected: Clean (run `uvx ruff format .` if needed)

**Step 4: Run type check**

Run: `uvx ty check src/`
Expected: No new diagnostics beyond the pre-existing 8

**Step 5: Push**

```bash
git push
```

---

## Summary of changes

| File | Change |
|------|--------|
| `src/graftpunk/plugins/cli_plugin.py` | Add `saves_session` to `CommandMetadata`, `CommandSpec`, `@command()`. Make `CommandContext` non-frozen, add `_session_name`, `_session_dirty`, `save_session()`. Pass `saves_session` through `_build_command_spec()`. |
| `src/graftpunk/cache.py` | Add `update_session_cookies()` — load-modify-save that preserves browser metadata. |
| `src/graftpunk/cli/plugin_commands.py` | Import `update_session_cookies`. Pass `_session_name` to `CommandContext`. After successful handler execution, check flags and call `update_session_cookies()`. |
| `tests/unit/test_cli_plugin.py` | Tests for `saves_session` field, decorator param, `CommandContext.save_session()` method. |
| `tests/unit/test_cache.py` | Tests for `update_session_cookies()` — success, load failure, cache failure. |
| `tests/unit/test_plugin_commands.py` | Integration tests for decorator save, explicit save, no-save, and failure-no-save paths. |
