# TUI First Repaint Before SDK Submit Implementation Plan

> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking.

**Goal:** Ensure that after the user presses Enter in `comate_cli`, the local TUI has completed at least one visible loading/staging repaint before any SDK first-turn work that can block or wait for provider/MCP startup begins.

**Architecture:** Add a small prompt_toolkit repaint acknowledgement mechanism in `TerminalAgentTUI`, backed by `Application.after_render`. Staged idle submits will wait for the specific render generation requested by `_enter_staging()` before calling `_submit_user_message()`, with a short timeout so submit cannot hang if prompt_toolkit fails to render. This keeps SDK queue/query semantics unchanged and makes “local feedback before SDK work” a TUI contract.

**Tech Stack:** Python 3.11, asyncio, prompt_toolkit `Application.after_render`, existing `FirstTurnProfiler`, pytest/unittest, ruff.

---

## Context And Constraints

- User-visible problem: on Windows cmd/PowerShell, first Enter appears stuck even when local submit profiling shows `_submit_from_input()` and `session.send()` are fast.
- Current weak mitigation: `_submit_staged_user_message_after_first_repaint()` only does `await asyncio.sleep(0)`. That yields to asyncio once but does **not** prove prompt_toolkit rendered anything.
- prompt_toolkit documentation confirms:
  - `Application.run_async()` is where the asyncio event loop runs.
  - `Application.invalidate()` schedules redraw; it does not synchronously draw.
  - `Application` exposes callbacks including `before_render` and `after_render`.
  - rendering can be postponed when the event loop is busy, controlled by render scheduling and `max_render_postpone_time`.
- Project rule: this is prompt_toolkit/TUI work, so keep the main TUI path non-blocking and bounded.
- Do not change SDK queue, `ChatSession`, MCP, provider, or context semantics in this plan.

## Success Criteria

- Pressing Enter for an idle normal user message immediately creates local staging/loading state.
- The background submit coroutine does not enter `_submit_user_message()` until either:
  - prompt_toolkit `after_render` confirms the render generation containing the staging/loading state completed, or
  - a short bounded timeout expires.
- On timeout, submit still proceeds, with first-turn profile evidence showing timeout.
- Busy queue submit, slash command execution, question UI, and consumed queue display behavior are unchanged.
- Existing queue suppression behavior remains intact: a direct submit’s own SDK queue item must not show as “queued message / Press ↑ to edit”.

## Files

- Modify: `packages/comate_cli/comate_cli/terminal_agent/tui.py`
  - Owns `TerminalAgentTUI`, prompt_toolkit `Application` construction, staging commit, background submit, and first-turn profile marks.
- Modify: `packages/comate_cli/tests/test_tui_send_undo_window.py`
  - Existing lightweight fake TUI tests for staging submit behavior.
- Optional Modify: `packages/comate_cli/docs/startup-performance.md`
  - Only if the final implementation changes the first-turn profile phase names or Windows troubleshooting instructions.

## Design Details

### Render Generation Model

Add generation counters to `TerminalAgentTUI`:

```python
self._render_request_generation = 0
self._render_completed_generation = 0
self._render_waiters: list[tuple[int, asyncio.Future[None]]] = []
self._submit_repaint_timeout_s = read_env_float(
    "COMATE_TUI_SUBMIT_REPAINT_TIMEOUT_S",
    0.15,
)
```

`_invalidate()` increments `_render_request_generation` before calling `app.invalidate()`. `Application.after_render` calls `_on_after_render()`, which advances `_render_completed_generation` and resolves waiters whose target generation has completed.

### Submit Waiting Model

`_commit_staging()` captures the current render generation after `_enter_staging()` has invalidated the UI:

```python
target_generation = int(getattr(self, "_render_request_generation", 0))
self._schedule_background(
    self._submit_staged_user_message_after_first_repaint(
        payload.submit_text,
        display_text=payload.display_text,
        defer_user_seed=True,
        render_generation=target_generation,
    )
)
```

The submit wrapper waits for that generation:

```python
await self._wait_for_render_generation(
    render_generation,
    timeout_s=self._submit_repaint_timeout_s,
)
await self._submit_user_message(...)
```

If there is no `Application` yet, or the app is closing, the wait is a no-op plus `await asyncio.sleep(0)`.

---

### Task 1: Replace Weak Yield Test With Render-Ack Regression Test

**Files:**
- Modify: `packages/comate_cli/tests/test_tui_send_undo_window.py`

- [ ] **Step 1: Replace the current weak-yield test with a render-ack test**

Find the existing test named:

```python
def test_commit_staging_yields_once_before_submitting_for_first_repaint(self) -> None:
```

Replace it with:

```python
def test_commit_staging_waits_for_prompt_toolkit_render_before_submitting(
    self,
) -> None:
    tui, _session = _build_tui()
    scheduled = _capture_scheduled(tui)
    observed: list[tuple[str, str | None, bool]] = []

    class _FakeApp:
        def __init__(self) -> None:
            self.invalidate_calls = 0

        def invalidate(self) -> None:
            self.invalidate_calls += 1

    async def _submit_user_message(
        text: str,
        *,
        display_text: str | None = None,
        defer_user_seed: bool = False,
    ) -> None:
        observed.append((text, display_text, defer_user_seed))

    async def _drive_scheduled() -> None:
        task = asyncio.create_task(scheduled[0])
        await asyncio.sleep(0)
        self.assertEqual(observed, [])
        tui._on_after_render(tui._app)
        await task

    tui._app = _FakeApp()
    tui._render_request_generation = 0
    tui._render_completed_generation = 0
    tui._render_waiters = []
    tui._submit_repaint_timeout_s = 1.0
    tui._submit_user_message = _submit_user_message
    tui._enter_staging(raw_text="hello", submit_text="hello", display_text="hello")

    self.assertTrue(tui._commit_staging())
    self.assertEqual(len(scheduled), 1)
    asyncio.run(_drive_scheduled())
    self.assertEqual(observed, [("hello", "hello", True)])
```

- [ ] **Step 2: Run the regression test and verify it fails**

Run:

```bash
uv run pytest packages/comate_cli/tests/test_tui_send_undo_window.py::TestTUISendUndoWindow::test_commit_staging_waits_for_prompt_toolkit_render_before_submitting
```

Expected:

```text
FAILED ... observed ... [('hello', 'hello', True)] != []
```

The failure proves current `await asyncio.sleep(0)` is insufficient because it allows submit before render acknowledgement.

- [ ] **Step 3: Commit only the failing test if working in commits**

Use a Chinese commit message:

```bash
git add packages/comate_cli/tests/test_tui_send_undo_window.py
git commit -m "test: 覆盖首轮提交必须等待 TUI 完成首帧渲染"
```

---

### Task 2: Add Render Generation State And after_render Hook

**Files:**
- Modify: `packages/comate_cli/comate_cli/terminal_agent/tui.py`

- [ ] **Step 1: Initialize render ack fields in `TerminalAgentTUI.__init__`**

Near the existing state initialization around `_queued_display_by_message_id`, add:

```python
self._render_request_generation = 0
self._render_completed_generation = 0
self._render_waiters: list[tuple[int, asyncio.Future[None]]] = []
self._submit_repaint_timeout_s = read_env_float(
    "COMATE_TUI_SUBMIT_REPAINT_TIMEOUT_S",
    0.15,
)
```

- [ ] **Step 2: Wire prompt_toolkit `after_render` callback**

Find the `Application(` construction in `TerminalAgentTUI.__init__`.

Add the callback argument:

```python
after_render=self._on_after_render,
```

The callback method must accept the prompt_toolkit `Application` object.

- [ ] **Step 3: Add `_on_after_render()` and waiter resolution helpers**

Add these methods near `_invalidate()` or near `_first_turn_mark()`:

```python
def _on_after_render(self, _app: Application[None]) -> None:
    completed = int(getattr(self, "_render_request_generation", 0))
    self._render_completed_generation = max(
        int(getattr(self, "_render_completed_generation", 0)),
        completed,
    )
    self._first_turn_mark(f"render.after_render.generation.{completed}")
    self._resolve_render_waiters()

def _resolve_render_waiters(self) -> None:
    completed = int(getattr(self, "_render_completed_generation", 0))
    waiters = getattr(self, "_render_waiters", [])
    if not isinstance(waiters, list):
        self._render_waiters = []
        return

    pending: list[tuple[int, asyncio.Future[None]]] = []
    for target_generation, future in waiters:
        if future.done():
            continue
        if target_generation <= completed:
            future.set_result(None)
        else:
            pending.append((target_generation, future))
    self._render_waiters = pending
```

- [ ] **Step 4: Make `_invalidate()` advance the requested generation**

Replace `_invalidate()` with:

```python
def _invalidate(self) -> None:
    self._render_request_generation = (
        int(getattr(self, "_render_request_generation", 0)) + 1
    )
    app = self._app
    if app is None:
        return
    app.invalidate()
```

- [ ] **Step 5: Run the existing regression test and confirm it still fails**

Run:

```bash
uv run pytest packages/comate_cli/tests/test_tui_send_undo_window.py::TestTUISendUndoWindow::test_commit_staging_waits_for_prompt_toolkit_render_before_submitting
```

Expected:

```text
FAILED ... observed ... [('hello', 'hello', True)] != []
```

This task only adds render accounting; submit does not wait yet.

- [ ] **Step 6: Commit render accounting if working in commits**

Use a Chinese commit message:

```bash
git add packages/comate_cli/comate_cli/terminal_agent/tui.py
git commit -m "feat: 为 TUI 增加 prompt_toolkit 首帧渲染确认机制"
```

---

### Task 3: Wait For Staging Repaint Before SDK Submit

**Files:**
- Modify: `packages/comate_cli/comate_cli/terminal_agent/tui.py`

- [ ] **Step 1: Add bounded render wait helper**

Add this method near `_on_after_render()`:

```python
async def _wait_for_render_generation(
    self,
    target_generation: int,
    *,
    timeout_s: float,
) -> None:
    if target_generation <= int(getattr(self, "_render_completed_generation", 0)):
        self._first_turn_mark("submit_staged.wait_for_repaint.already_done")
        return

    app = getattr(self, "_app", None)
    if app is None or bool(getattr(self, "_closing", False)):
        self._first_turn_mark("submit_staged.wait_for_repaint.no_app")
        await asyncio.sleep(0)
        return

    loop = asyncio.get_running_loop()
    future: asyncio.Future[None] = loop.create_future()
    waiters = getattr(self, "_render_waiters", None)
    if not isinstance(waiters, list):
        waiters = []
        self._render_waiters = waiters
    waiters.append((int(target_generation), future))

    self._first_turn_mark(
        f"submit_staged.wait_for_repaint.start.generation.{target_generation}"
    )
    self._invalidate()
    try:
        await asyncio.wait_for(future, timeout=max(0.001, float(timeout_s)))
    except asyncio.TimeoutError:
        if not future.done():
            future.cancel()
        self._first_turn_mark(
            f"submit_staged.wait_for_repaint.timeout.generation.{target_generation}"
        )
    else:
        self._first_turn_mark(
            f"submit_staged.wait_for_repaint.done.generation.{target_generation}"
        )
    finally:
        self._render_waiters = [
            (generation, waiter)
            for generation, waiter in getattr(self, "_render_waiters", [])
            if waiter is not future and not waiter.done()
        ]
```

- [ ] **Step 2: Pass render generation from `_commit_staging()`**

In the non-busy branch of `_commit_staging()`, replace the scheduled call with:

```python
target_generation = int(getattr(self, "_render_request_generation", 0))
self._schedule_background(
    self._submit_staged_user_message_after_first_repaint(
        payload.submit_text,
        display_text=payload.display_text,
        defer_user_seed=True,
        render_generation=target_generation,
    )
)
```

- [ ] **Step 3: Replace weak `sleep(0)` in `_submit_staged_user_message_after_first_repaint()`**

Change the method signature and body to:

```python
async def _submit_staged_user_message_after_first_repaint(
    self,
    text: str,
    *,
    display_text: str | None,
    defer_user_seed: bool,
    render_generation: int,
) -> None:
    await self._wait_for_render_generation(
        render_generation,
        timeout_s=float(getattr(self, "_submit_repaint_timeout_s", 0.15)),
    )
    await self._submit_user_message(
        text,
        display_text=display_text,
        defer_user_seed=defer_user_seed,
    )
```

- [ ] **Step 4: Run the render-ack regression test and verify it passes**

Run:

```bash
uv run pytest packages/comate_cli/tests/test_tui_send_undo_window.py::TestTUISendUndoWindow::test_commit_staging_waits_for_prompt_toolkit_render_before_submitting
```

Expected:

```text
1 passed
```

- [ ] **Step 5: Commit the wait implementation if working in commits**

Use a Chinese commit message:

```bash
git add packages/comate_cli/comate_cli/terminal_agent/tui.py packages/comate_cli/tests/test_tui_send_undo_window.py
git commit -m "fix: 首轮提交前等待 TUI 完成本地 loading 首帧渲染"
```

---

### Task 4: Cover Timeout And No-App Paths

**Files:**
- Modify: `packages/comate_cli/tests/test_tui_send_undo_window.py`

- [ ] **Step 1: Add timeout-path test**

Add this test next to the render-ack regression test:

```python
def test_commit_staging_proceeds_after_repaint_timeout(self) -> None:
    tui, _session = _build_tui()
    scheduled = _capture_scheduled(tui)
    observed: list[tuple[str, str | None, bool]] = []

    class _FakeApp:
        def invalidate(self) -> None:
            return None

    async def _submit_user_message(
        text: str,
        *,
        display_text: str | None = None,
        defer_user_seed: bool = False,
    ) -> None:
        observed.append((text, display_text, defer_user_seed))

    tui._app = _FakeApp()
    tui._render_request_generation = 0
    tui._render_completed_generation = 0
    tui._render_waiters = []
    tui._submit_repaint_timeout_s = 0.001
    tui._submit_user_message = _submit_user_message
    tui._enter_staging(raw_text="hello", submit_text="hello", display_text="hello")

    self.assertTrue(tui._commit_staging())
    self.assertEqual(len(scheduled), 1)
    asyncio.run(scheduled[0])
    self.assertEqual(observed, [("hello", "hello", True)])
```

- [ ] **Step 2: Add no-app path test**

Add:

```python
def test_commit_staging_without_app_does_not_block_submit(self) -> None:
    tui, _session = _build_tui()
    scheduled = _capture_scheduled(tui)
    observed: list[tuple[str, str | None, bool]] = []

    async def _submit_user_message(
        text: str,
        *,
        display_text: str | None = None,
        defer_user_seed: bool = False,
    ) -> None:
        observed.append((text, display_text, defer_user_seed))

    tui._app = None
    tui._render_request_generation = 0
    tui._render_completed_generation = 0
    tui._render_waiters = []
    tui._submit_repaint_timeout_s = 1.0
    tui._submit_user_message = _submit_user_message
    tui._enter_staging(raw_text="hello", submit_text="hello", display_text="hello")

    self.assertTrue(tui._commit_staging())
    self.assertEqual(len(scheduled), 1)
    asyncio.run(scheduled[0])
    self.assertEqual(observed, [("hello", "hello", True)])
```

- [ ] **Step 3: Run the focused tests**

Run:

```bash
uv run pytest \
  packages/comate_cli/tests/test_tui_send_undo_window.py::TestTUISendUndoWindow::test_commit_staging_waits_for_prompt_toolkit_render_before_submitting \
  packages/comate_cli/tests/test_tui_send_undo_window.py::TestTUISendUndoWindow::test_commit_staging_proceeds_after_repaint_timeout \
  packages/comate_cli/tests/test_tui_send_undo_window.py::TestTUISendUndoWindow::test_commit_staging_without_app_does_not_block_submit
```

Expected:

```text
3 passed
```

- [ ] **Step 4: Commit timeout/no-app coverage if working in commits**

Use a Chinese commit message:

```bash
git add packages/comate_cli/tests/test_tui_send_undo_window.py
git commit -m "test: 覆盖首帧渲染等待的超时和无应用降级路径"
```

---

### Task 5: Preserve Existing Queue Preview Suppression Behavior

**Files:**
- Modify: `packages/comate_cli/tests/test_tui_send_undo_window.py`
- Read-only check: `packages/comate_cli/comate_cli/terminal_agent/tui_parts/render_panels.py`
- Read-only check: `packages/comate_cli/comate_cli/terminal_agent/tui.py`

- [ ] **Step 1: Confirm the existing direct-submit queue suppression test remains present**

The test must still exist:

```python
def test_direct_submit_suppresses_own_queue_preview_while_sdk_queue_catches_up(
    self,
) -> None:
    tui, session = _build_tui()
    session.reflect_sent_in_queue = True

    asyncio.run(tui._submit_user_message("hello"))

    self.assertEqual([message.message_id for message in session.peek_queue()], ["msg-1"])
    self.assertFalse(tui._should_show_queue_panel())
    self.assertIsNone(tui._input_placeholder_hint())
```

- [ ] **Step 2: Run the queue suppression and busy queue tests**

Run:

```bash
uv run pytest \
  packages/comate_cli/tests/test_tui_send_undo_window.py::TestTUISendUndoWindow::test_direct_submit_suppresses_own_queue_preview_while_sdk_queue_catches_up \
  packages/comate_cli/tests/test_tui_queue_sdk_source.py
```

Expected:

```text
15 passed
```

If the exact count changes because tests were added, accept only zero failures.

- [ ] **Step 3: Commit only if this task required test changes**

Use a Chinese commit message:

```bash
git add packages/comate_cli/tests/test_tui_send_undo_window.py
git commit -m "test: 保持直接提交不显示自身队列预览"
```

If no files changed in this task, do not create a commit.

---

### Task 6: Add Optional Troubleshooting Documentation

**Files:**
- Modify: `packages/comate_cli/docs/startup-performance.md`

- [ ] **Step 1: Update first-turn troubleshooting notes only if phase names changed**

If the implementation introduces the phase names from this plan, add these bullets under “排查首条消息 Enter 卡顿”:

```markdown
- `tui.submit_staged.wait_for_repaint.*`：TUI 是否在进入 SDK submit 前等到了首帧 render。
- `tui.render.after_render.generation.*`：prompt_toolkit 是否真的完成过本地 repaint。
```

- [ ] **Step 2: Run a docs diff sanity check**

Run:

```bash
git diff -- packages/comate_cli/docs/startup-performance.md
```

Expected:

```text
Only the first-turn troubleshooting section changed.
```

- [ ] **Step 3: Commit docs if changed**

Use a Chinese commit message:

```bash
git add packages/comate_cli/docs/startup-performance.md
git commit -m "docs: 补充首条消息首帧渲染排障指标"
```

---

### Task 7: Full Verification

**Files:**
- Verify only; no source changes.

- [ ] **Step 1: Run lint**

Run:

```bash
uv run ruff check \
  packages/comate_cli/comate_cli/terminal_agent/tui.py \
  packages/comate_cli/comate_cli/terminal_agent/tui_parts/render_panels.py \
  packages/comate_cli/tests/test_tui_send_undo_window.py \
  packages/comate_cli/tests/test_tui_queue_sdk_source.py
```

Expected:

```text
All checks passed!
```

- [ ] **Step 2: Run focused TUI tests**

Run:

```bash
uv run pytest \
  packages/comate_cli/tests/test_tui_send_undo_window.py \
  packages/comate_cli/tests/test_tui_queue_sdk_source.py \
  packages/comate_cli/tests/test_tui_paste_placeholder.py \
  packages/comate_cli/tests/test_tui_startup_latency.py \
  packages/comate_cli/tests/test_startup_profile.py::TestStartupProfiler
```

Expected:

```text
All selected tests pass with 0 failures.
```

- [ ] **Step 3: Run SDK guard tests touched by first-turn profiling**

Run:

```bash
uv run pytest \
  comate_agent_sdk/utils/tests/test_first_turn_profile.py \
  comate_agent_sdk/context/tests/test_token_counter_tiktoken_cache.py \
  comate_agent_sdk/context/tests/test_budget_prewarm.py \
  comate_agent_sdk/agent/tests/test_chat_session_persistence.py::TestChatSessionPersistence::test_resume_requires_header_snapshot \
  comate_agent_sdk/agent/tests/test_chat_session_persistence.py::TestChatSessionPersistence::test_resume_restores_header_snapshot_over_runtime_initialization
```

Expected:

```text
All selected tests pass with 0 failures.
```

- [ ] **Step 4: Manual Windows verification**

On Windows cmd or PowerShell:

```powershell
$env:COMATE_FIRST_ENTER_PROFILE="1"
comate
```

Submit one short message immediately after TUI appears.

Expected user-visible result:

```text
The local staged message/loading indicator appears before any multi-second SDK/provider/MCP wait.
No duplicate queued-message hint appears for the same direct submit.
```

Expected log shape:

```text
tui.input.submit_from_input.done
tui.render.after_render.generation.<N>
tui.submit_staged.wait_for_repaint.done.generation.<N>
tui.submit_user_message.start
...
chat_session.run_turn_stream.before_query_stream
```

If `tui.submit_staged.wait_for_repaint.timeout.generation.<N>` appears, submit should still proceed; treat that as a Windows render scheduling signal to investigate, not a functional failure.

- [ ] **Step 5: Final commit if working in commits**

Use a Chinese commit message:

```bash
git status --short
git add packages/comate_cli/comate_cli/terminal_agent/tui.py \
  packages/comate_cli/tests/test_tui_send_undo_window.py \
  packages/comate_cli/docs/startup-performance.md
git commit -m "fix: 首条消息提交前确保 TUI loading 已完成首帧渲染

问题背景：
- Windows cmd/PowerShell 首条消息 Enter 后用户仍感到 TUI 卡顿
- 现有 sleep(0) 只能让出事件循环，不能证明 prompt_toolkit 已实际重绘

实现：
- 通过 prompt_toolkit after_render 建立 render generation ack
- idle staged submit 进入 SDK 前等待对应 generation 完成渲染
- 增加短 timeout 防止极端渲染失败阻塞提交
- 保持 busy queue、direct submit queue suppression 和 SDK 语义不变

验证：
- ruff check 通过
- TUI staging/queue/startup 相关测试通过
- SDK first-turn profile/token/header snapshot 护栏测试通过"
```

---

## Self-Review

- Spec coverage:
  - “先 loading 再允许 SDK 阻塞” covered by Tasks 1-4.
  - prompt_toolkit `after_render` based acknowledgement covered by Tasks 2-3.
  - No SDK semantic changes covered by file scope and Task 5 regression.
  - Windows troubleshooting covered by Task 6 and Task 7 manual verification.
- Placeholder scan:
  - No `TBD`, `TODO`, “implement later”, or unbound “add appropriate” steps.
- Type consistency:
  - `render_generation`, `_render_request_generation`, `_render_completed_generation`, and `_render_waiters` are consistently named across tests and implementation.
  - `_on_after_render()` signature accepts the prompt_toolkit `Application` object as documented.
