# Login Element Retry Implementation Plan

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

**Goal:** Make the nodriver login engine resilient to page transitions (redirects, SPAs, lazy-rendered forms) by retrying element selection through `ProtocolException`, and add an optional `wait_for` field to `LoginConfig` for explicit pre-interaction waits.

**Architecture:** Add a `_select_with_retry` async helper in `login_engine.py` that wraps `tab.select()` with a deadline-based retry loop catching nodriver's `ProtocolException`. Replace the three pre-submit `tab.select()` calls with this helper. Add `wait_for: str = ""` to `LoginConfig` and use it before field filling when set.

**Tech Stack:** Python, nodriver, asyncio, pytest

---

### Task 1: Add `wait_for` field to `LoginConfig`

**Files:**
- Modify: `src/graftpunk/plugins/cli_plugin.py:182-203`
- Test: `tests/unit/test_cli_plugin.py`

**Step 1: Write the failing tests**

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

```python
def test_wait_for_default_empty(self) -> None:
    """LoginConfig.wait_for defaults to empty string."""
    cfg = LoginConfig(url="/login", fields={"u": "#u"}, submit="#b")
    assert cfg.wait_for == ""

def test_wait_for_stores_value(self) -> None:
    """LoginConfig stores wait_for selector."""
    cfg = LoginConfig(
        url="/login",
        fields={"u": "#u"},
        submit="#b",
        wait_for="input#signInName",
    )
    assert cfg.wait_for == "input#signInName"
```

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

Run: `uv run pytest tests/unit/test_cli_plugin.py::TestLoginConfig::test_wait_for_default_empty tests/unit/test_cli_plugin.py::TestLoginConfig::test_wait_for_stores_value -v`
Expected: FAIL — `LoginConfig.__init__() got an unexpected keyword argument 'wait_for'`

**Step 3: Implement the change**

In `src/graftpunk/plugins/cli_plugin.py`, add `wait_for` to `LoginConfig`:

```python
@dataclass(frozen=True)
class LoginConfig:
    """Declarative browser-automated login configuration.

    All three required fields (url, fields, submit) must be non-empty.

    Attributes:
        url: Login page path (appended to base_url).
        fields: Maps credential names to CSS selectors for form inputs.
        submit: CSS selector for the submit button.
        failure: Text on the page indicating login failure.
        success: CSS selector for an element indicating login success.
        wait_for: CSS selector to wait for before interacting with the page.
            Useful when the login URL triggers a redirect or the form renders
            asynchronously. Empty string (default) means no explicit wait.
    """

    url: str
    fields: dict[str, str]
    submit: str
    failure: str = ""
    success: str = ""
    wait_for: str = ""

    def __post_init__(self) -> None:
        if not self.url:
            raise ValueError("LoginConfig.url must be non-empty")
        if not self.fields:
            raise ValueError("LoginConfig.fields must be non-empty")
        if not self.submit:
            raise ValueError("LoginConfig.submit must be non-empty")
        # Defensive copy: prevent external mutation of fields dict
        object.__setattr__(self, "fields", dict(self.fields))
```

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

Run: `uv run pytest tests/unit/test_cli_plugin.py::TestLoginConfig -v`
Expected: PASS (all LoginConfig tests)

**Step 5: Run full test suite**

Run: `uv run pytest tests/ -v`
Expected: All tests pass. The new field defaults to `""` so no existing code breaks.

**Step 6: Commit**

```bash
git add src/graftpunk/plugins/cli_plugin.py tests/unit/test_cli_plugin.py
git commit -m "feat: add wait_for field to LoginConfig

Optional CSS selector that the login engine waits for before filling
form fields. Useful when the login URL triggers a redirect or the
form renders asynchronously (Azure AD B2C, Okta, SPAs, etc.).

Closes part of #63

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

---

### Task 2: Implement `_select_with_retry` helper

**Files:**
- Modify: `src/graftpunk/plugins/login_engine.py:1-25` (imports and constants)
- Modify: `src/graftpunk/plugins/login_engine.py` (new function after constants)
- Test: `tests/unit/test_login_engine.py`

**Context:** nodriver's `tab.select()` has a built-in retry loop for when an element doesn't exist yet (returns None). But during page transitions, `query_selector` raises `ProtocolException` instead of returning None — the document node itself is invalid. This exception escapes `select()` entirely. The helper wraps `tab.select()` with a deadline-based retry loop that catches `ProtocolException`.

**Step 1: Write the failing tests**

Add a new test class to `tests/unit/test_login_engine.py`:

```python
class TestSelectWithRetry:
    """Tests for _select_with_retry helper."""

    @pytest.mark.asyncio
    async def test_returns_element_on_first_try(self) -> None:
        """Returns element immediately when select succeeds."""
        from graftpunk.plugins.login_engine import _select_with_retry

        mock_tab = AsyncMock()
        mock_element = AsyncMock()
        mock_tab.select = AsyncMock(return_value=mock_element)

        result = await _select_with_retry(mock_tab, "input#name")
        assert result is mock_element
        mock_tab.select.assert_called_once()

    @pytest.mark.asyncio
    async def test_retries_on_protocol_exception(self) -> None:
        """Retries when ProtocolException is raised, succeeds on later attempt."""
        from nodriver.core.connection import ProtocolException

        from graftpunk.plugins.login_engine import _select_with_retry

        mock_tab = AsyncMock()
        mock_element = AsyncMock()

        # Fail twice with ProtocolException, succeed on third
        exc = ProtocolException({"code": -32000, "message": "Could not find node"})
        mock_tab.select = AsyncMock(side_effect=[exc, exc, mock_element])

        result = await _select_with_retry(mock_tab, "input#name", timeout=10, interval=0.01)
        assert result is mock_element
        assert mock_tab.select.call_count == 3

    @pytest.mark.asyncio
    async def test_raises_after_timeout(self) -> None:
        """Raises last ProtocolException when timeout expires."""
        from nodriver.core.connection import ProtocolException

        from graftpunk.plugins.login_engine import _select_with_retry

        mock_tab = AsyncMock()
        exc = ProtocolException({"code": -32000, "message": "Could not find node"})
        mock_tab.select = AsyncMock(side_effect=exc)

        with pytest.raises(ProtocolException):
            await _select_with_retry(mock_tab, "input#name", timeout=0.1, interval=0.01)

    @pytest.mark.asyncio
    async def test_returns_none_when_select_returns_none(self) -> None:
        """Returns None when select returns None and timeout expires (no exception)."""
        from graftpunk.plugins.login_engine import _select_with_retry

        mock_tab = AsyncMock()
        mock_tab.select = AsyncMock(return_value=None)

        result = await _select_with_retry(mock_tab, "input#gone", timeout=0.1, interval=0.01)
        assert result is None

    @pytest.mark.asyncio
    async def test_non_protocol_exception_propagates(self) -> None:
        """Non-ProtocolException errors propagate immediately (no retry)."""
        from graftpunk.plugins.login_engine import _select_with_retry

        mock_tab = AsyncMock()
        mock_tab.select = AsyncMock(side_effect=RuntimeError("unexpected"))

        with pytest.raises(RuntimeError, match="unexpected"):
            await _select_with_retry(mock_tab, "input#name")

        # Only called once — no retry for non-protocol errors
        mock_tab.select.assert_called_once()
```

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

Run: `uv run pytest tests/unit/test_login_engine.py::TestSelectWithRetry -v`
Expected: FAIL — `cannot import name '_select_with_retry'`

**Step 3: Implement `_select_with_retry`**

In `src/graftpunk/plugins/login_engine.py`, add constants after `_POST_SUBMIT_DELAY` (line 24):

```python
_ELEMENT_WAIT_TIMEOUT = 30    # seconds to wait for element during page transitions
_ELEMENT_RETRY_INTERVAL = 1.0  # seconds between retry attempts
```

Add the helper function after the constants (before `_warn_no_login_validation`):

```python
async def _select_with_retry(
    tab: Any,
    selector: str,
    *,
    timeout: float = _ELEMENT_WAIT_TIMEOUT,
    interval: float = _ELEMENT_RETRY_INTERVAL,
) -> Any:
    """Wait for a CSS selector, retrying through page transitions.

    nodriver's tab.select() handles the case where an element doesn't exist
    yet (returns None, retries internally). But during cross-origin redirects
    or page transitions, the document node itself becomes invalid, causing a
    ProtocolException that bypasses select()'s retry loop.

    This wrapper catches ProtocolException and retries the entire select()
    call, giving the browser time to complete redirects and render the form.

    Args:
        tab: nodriver tab instance.
        selector: CSS selector string.
        timeout: Total seconds to wait before giving up.
        interval: Seconds between retry attempts.

    Returns:
        The matched element, or None if not found within timeout.

    Raises:
        ProtocolException: If timeout expires and last failure was a protocol error.
    """
    from nodriver.core.connection import ProtocolException

    loop = asyncio.get_running_loop()
    deadline = loop.time() + timeout
    last_exc: ProtocolException | None = None

    while loop.time() < deadline:
        remaining = deadline - loop.time()
        try:
            per_attempt = min(5.0, remaining)
            element = await tab.select(selector, timeout=per_attempt)
            if element is not None:
                return element
        except ProtocolException as exc:
            last_exc = exc
            LOG.debug(
                "login_element_retry",
                selector=selector,
                error=str(exc),
                remaining=f"{remaining:.1f}s",
            )
        if loop.time() >= deadline:
            break
        await asyncio.sleep(interval)

    if last_exc is not None:
        raise last_exc
    return None
```

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

Run: `uv run pytest tests/unit/test_login_engine.py::TestSelectWithRetry -v`
Expected: PASS (all 5 tests)

**Step 5: Run full test suite**

Run: `uv run pytest tests/ -v`
Expected: All tests pass. The helper is defined but not yet used by `_generate_nodriver_login`.

**Step 6: Commit**

```bash
git add src/graftpunk/plugins/login_engine.py tests/unit/test_login_engine.py
git commit -m "feat: add _select_with_retry helper for resilient element selection

Wraps tab.select() with a deadline-based retry loop that catches
nodriver's ProtocolException. During page transitions (redirects,
SPAs), the DOM document node becomes invalid and select() throws
instead of returning None. This helper retries through those
transient failures.

Ref #63

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

---

### Task 3: Wire `_select_with_retry` and `wait_for` into the login flow

**Files:**
- Modify: `src/graftpunk/plugins/login_engine.py:202-261`
- Test: `tests/unit/test_login_engine.py`

**Context:** Replace three `tab.select()` calls in `_generate_nodriver_login` with `_select_with_retry()` (field filling, submit button). Add the `wait_for` check before field filling. The success selector check after submit stays as bare `tab.select()` — by that point the page has settled.

**Step 1: Write the failing tests**

Add a new test class to `tests/unit/test_login_engine.py`:

```python
class TestLoginRetryIntegration:
    """Tests for _select_with_retry integration in the login flow."""

    @pytest.mark.asyncio
    async def test_login_retries_through_protocol_exception(self) -> None:
        """Login succeeds when select fails transiently then recovers."""
        from nodriver.core.connection import ProtocolException

        from graftpunk.plugins.login_engine import generate_login_method

        plugin = DeclarativeHN()
        login_method = generate_login_method(plugin)

        mock_tab = AsyncMock()
        mock_element = AsyncMock()

        # select() fails with ProtocolException on first call, succeeds after
        exc = ProtocolException({"code": -32000, "message": "Could not find node"})
        mock_tab.select = AsyncMock(side_effect=[exc, mock_element, mock_element, mock_element])
        mock_tab.get_content = AsyncMock(return_value="<html>Welcome</html>")

        mock_bs, instance = _make_nodriver_mock_bs()
        instance.driver = MagicMock()
        instance.driver.get = AsyncMock(return_value=mock_tab)
        instance.transfer_nodriver_cookies_to_session = AsyncMock()

        with (
            patch("graftpunk.plugins.login_engine.BrowserSession", mock_bs),
            patch("graftpunk.plugins.login_engine.cache_session"),
            patch("graftpunk.plugins.login_engine.asyncio.sleep", new_callable=AsyncMock),
        ):
            result = await login_method({"username": "user", "password": "test"})  # noqa: S106

        assert result is True


class TestLoginWaitFor:
    """Tests for LoginConfig.wait_for integration."""

    @pytest.mark.asyncio
    async def test_wait_for_called_before_fields(self) -> None:
        """When wait_for is set, the selector is awaited before filling fields."""
        from graftpunk.plugins.login_engine import generate_login_method

        class WaitForPlugin(SitePlugin):
            site_name = "waitfor"
            session_name = "waitfor"
            help_text = "WF"
            base_url = "https://example.com"
            backend = "nodriver"
            login_config = LoginConfig(
                url="/login",
                fields={"username": "#user"},
                submit="#btn",
                wait_for="#login-form",
            )

        plugin = WaitForPlugin()
        login_method = generate_login_method(plugin)

        mock_tab = AsyncMock()
        mock_element = AsyncMock()
        select_calls: list[str] = []

        async def tracking_select(selector, timeout=10):
            select_calls.append(selector)
            return mock_element

        mock_tab.select = tracking_select
        mock_tab.get_content = AsyncMock(return_value="<html>OK</html>")

        mock_bs, instance = _make_nodriver_mock_bs()
        instance.driver = MagicMock()
        instance.driver.get = AsyncMock(return_value=mock_tab)
        instance.transfer_nodriver_cookies_to_session = AsyncMock()

        with (
            patch("graftpunk.plugins.login_engine.BrowserSession", mock_bs),
            patch("graftpunk.plugins.login_engine.cache_session"),
            patch("graftpunk.plugins.login_engine.asyncio.sleep", new_callable=AsyncMock),
        ):
            result = await login_method({"username": "user"})

        assert result is True
        # wait_for selector should be first
        assert select_calls[0] == "#login-form"
        # Then the field selector, then submit
        assert "#user" in select_calls
        assert "#btn" in select_calls

    @pytest.mark.asyncio
    async def test_wait_for_timeout_raises_plugin_error(self) -> None:
        """When wait_for times out, raises PluginError."""
        from graftpunk.plugins.login_engine import generate_login_method

        class WaitForPlugin(SitePlugin):
            site_name = "waitfor"
            session_name = "waitfor"
            help_text = "WF"
            base_url = "https://example.com"
            backend = "nodriver"
            login_config = LoginConfig(
                url="/login",
                fields={"username": "#user"},
                submit="#btn",
                wait_for="#login-form",
            )

        plugin = WaitForPlugin()
        login_method = generate_login_method(plugin)

        mock_tab = AsyncMock()
        # select always returns None (element never appears)
        mock_tab.select = AsyncMock(return_value=None)

        mock_bs, instance = _make_nodriver_mock_bs()
        instance.driver = MagicMock()
        instance.driver.get = AsyncMock(return_value=mock_tab)

        with (
            patch("graftpunk.plugins.login_engine.BrowserSession", mock_bs),
            patch("graftpunk.plugins.login_engine.asyncio.sleep", new_callable=AsyncMock),
            pytest.raises(PluginError, match="Timed out waiting for"),
        ):
            await login_method({"username": "user"})

    @pytest.mark.asyncio
    async def test_no_wait_for_skips_wait(self) -> None:
        """When wait_for is empty, no extra select call is made."""
        from graftpunk.plugins.login_engine import generate_login_method

        plugin = DeclarativeHN()  # no wait_for
        login_method = generate_login_method(plugin)

        mock_tab = AsyncMock()
        mock_element = AsyncMock()
        mock_tab.select = AsyncMock(return_value=mock_element)
        mock_tab.get_content = AsyncMock(return_value="<html>Welcome</html>")

        mock_bs, instance = _make_nodriver_mock_bs()
        instance.driver = MagicMock()
        instance.driver.get = AsyncMock(return_value=mock_tab)
        instance.transfer_nodriver_cookies_to_session = AsyncMock()

        with (
            patch("graftpunk.plugins.login_engine.BrowserSession", mock_bs),
            patch("graftpunk.plugins.login_engine.cache_session"),
            patch("graftpunk.plugins.login_engine.asyncio.sleep", new_callable=AsyncMock),
        ):
            result = await login_method({"username": "user", "password": "test"})  # noqa: S106

        assert result is True
        # select should NOT have been called with wait_for selector
        for call in mock_tab.select.call_args_list:
            assert call[0][0] != "#login-form"
```

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

Run: `uv run pytest tests/unit/test_login_engine.py::TestLoginRetryIntegration tests/unit/test_login_engine.py::TestLoginWaitFor -v`
Expected: FAIL — login still uses bare `tab.select()` which doesn't retry, and `wait_for` is not wired in.

**Step 3: Implement the changes**

In `src/graftpunk/plugins/login_engine.py`, modify `_generate_nodriver_login`. After the header capture setup (line 226) and before the field-filling loop (line 228), add the `wait_for` block:

```python
            # Wait for a specific element before interacting (e.g., redirect targets)
            wait_for_selector = plugin.login_config.wait_for
            if wait_for_selector:
                wait_el = await _select_with_retry(tab, wait_for_selector)
                if wait_el is None:
                    raise PluginError(
                        f"Timed out waiting for '{wait_for_selector}' to appear. "
                        f"The page may not have loaded or redirected as expected."
                    )
```

In the field-filling loop, replace `tab.select(selector)` with `_select_with_retry(tab, selector)` (line 232):

```python
                    element = await _select_with_retry(tab, selector)
```

In the submit button section, replace `tab.select(submit_selector)` with `_select_with_retry(tab, submit_selector)` (line 249):

```python
                submit = await _select_with_retry(tab, submit_selector)
```

The success selector check (line 271) stays as `await tab.select(success_selector)` — no retry needed post-submit.

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

Run: `uv run pytest tests/unit/test_login_engine.py::TestLoginRetryIntegration tests/unit/test_login_engine.py::TestLoginWaitFor -v`
Expected: PASS

**Step 5: Run full test suite**

Run: `uv run pytest tests/ -v`
Expected: All tests pass. Existing login tests should still pass since `_select_with_retry` returns immediately when `tab.select()` succeeds on first attempt.

**Step 6: Commit**

```bash
git add src/graftpunk/plugins/login_engine.py tests/unit/test_login_engine.py
git commit -m "feat: wire _select_with_retry and wait_for into nodriver login flow

Replace bare tab.select() with _select_with_retry() for field filling
and submit button selection. Add wait_for support that waits for a
selector before interacting with the page.

The success selector check post-submit intentionally stays as bare
tab.select() since the page has already settled by that point.

Closes #63

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

---

### Task 4: Lint, format, type check, and create PR

**Files:**
- All modified files from tasks 1-3

**Step 1: Run quality checks**

```bash
uvx ruff check --fix . && uvx ruff format . && uvx ty check src/
```

Fix any issues.

**Step 2: Run full test suite one final time**

```bash
uv run pytest tests/ -v
```

Expected: All tests pass.

**Step 3: Push and create PR**

```bash
git push -u origin fix/63-login-element-retry
gh pr create --title "feat: resilient element selection in login engine" --body "$(cat <<'EOF'
## Summary

The nodriver login engine assumed the login form was present in the DOM immediately after `tab.get()` returned. When the login URL triggers a redirect (Azure AD B2C, Okta, SPAs), the form doesn't exist yet and `tab.select()` throws a `ProtocolException` that bypasses nodriver's built-in retry loop. The login fails instantly rather than waiting for the redirect to complete.

### Changes

- **`_select_with_retry` helper** (`login_engine.py`): Wraps `tab.select()` with a 30-second deadline-based retry loop catching `ProtocolException`. During page transitions, the DOM document node becomes invalid and `select()` throws instead of returning None. This helper retries through those transient failures with 1-second intervals.

- **`LoginConfig.wait_for`** (`cli_plugin.py`): Optional CSS selector field. When set, the login engine waits for this selector to appear before filling any form fields. Useful when plugin authors know the login page needs time to redirect/render.

- **Integration** (`login_engine.py`): Field filling and submit button selection now use `_select_with_retry`. The success selector check post-submit intentionally stays as bare `tab.select()` since the page has already settled.

Closes #63

## Test plan

- [ ] `_select_with_retry`: returns element on first try
- [ ] `_select_with_retry`: retries through ProtocolException, succeeds later
- [ ] `_select_with_retry`: raises after timeout expires
- [ ] `_select_with_retry`: returns None when select returns None
- [ ] `_select_with_retry`: non-ProtocolException propagates immediately
- [ ] Integration: login succeeds after transient ProtocolException
- [ ] `wait_for`: selector waited for before field filling
- [ ] `wait_for`: timeout raises PluginError
- [ ] `wait_for`: empty string skips wait
- [ ] `LoginConfig.wait_for`: defaults to empty string
- [ ] `LoginConfig.wait_for`: stores value
- [ ] All existing login engine tests still pass
EOF
)"
```
