Fix chat analysis-model selection, AG-UI example state, and chat docs

Drive analysis-only chat with analysis.model (falling back to qa.model) so
the running model matches the one the analysis capability configures,
including its vision flag; RAG-bearing chats still run on qa.model.

Give the AG-UI example state-bearing deps and a final STATE_SNAPSHOT so
registered citations reach the client, mirroring the app backend.

Update chat docs: haiku-rag chat uses --capability/-c, and drop the removed
"View state" command-palette entry.

docs/chat.md

@@ -49,7 +49,6 @@ haiku-rag visualize <chunk_id>
 | Filter documents | Restrict searches to selected documents |
 | Show visual grounding | Visual grounding for a citation |
 | Database info | Document and chunk counts, storage stats |
-| View state | Current session state, citations, and intermediate tool results |
 
 ## Capabilities
 
@@ -57,10 +56,10 @@ The default capability is `rag`. Enable `analysis` when the question needs compu
 
 ```bash
 # both capabilities (the agent routes between them)
-haiku-rag chat -s rag -s analysis
+haiku-rag chat -c rag -c analysis
 
 # analysis only
-haiku-rag chat -s analysis
+haiku-rag chat -c analysis
 ```
 
 The `analysis` capability mounts every document as a virtual filesystem at `/documents/{id}/` (with `metadata.json`, `content.txt`, `items.jsonl`, and `toc.json`) and runs Python in a sandboxed interpreter with `search` and `list_documents` as awaitable functions. It's the right choice for questions like:

docs/cli.md

@@ -196,7 +196,7 @@ haiku-rag chat
 haiku-rag chat --db /path/to/database.lancedb
 
 # Enable the analysis capability (code execution)
-haiku-rag chat -s rag -s analysis
+haiku-rag chat -c rag -c analysis
 ```
 
 !!! note

examples/custom_agent_agui.py

@@ -13,8 +13,11 @@ Usage:
 
 import os
 import sys
+from dataclasses import dataclass, field
 from pathlib import Path
+from typing import Any
 
+from ag_ui.core import EventType, StateSnapshotEvent
 from pydantic_ai import Agent
 from pydantic_ai.ui import SSE_CONTENT_TYPE
 from pydantic_ai.ui.ag_ui import AGUIAdapter
@@ -23,7 +26,7 @@ from starlette.requests import Request
 from starlette.responses import JSONResponse, Response, StreamingResponse
 from starlette.routing import Route
 
-from haiku.rag.capabilities.rag import create_capability
+from haiku.rag.capabilities.rag import RAGState, create_capability
 
 db_path = os.environ.get("DB_PATH")
 if not db_path:
@@ -34,9 +37,16 @@ if not db_path:
 
 capability = create_capability(db_path=Path(db_path), defer_loading=False)
 
+
+@dataclass
+class AppDeps:
+    state: dict[str, Any] = field(default_factory=dict)
+
+
 agent = Agent(
     "anthropic:claude-haiku-4-5-20251001",
     capabilities=[capability],
+    deps_type=AppDeps,
 )
 
 
@@ -47,8 +57,21 @@ async def stream_chat(request: Request) -> Response:
 
     adapter = AGUIAdapter(agent=agent, run_input=run_input, accept=accept)
 
+    incoming_state = run_input.state if isinstance(run_input.state, dict) else {}
+    incoming_state.setdefault("rag", RAGState().model_dump(mode="json"))
+    deps = AppDeps(state=incoming_state)
+
     async def event_stream():
-        async for chunk in adapter.encode_stream(adapter.run_stream()):
+        async def with_final_state():
+            async for event in adapter.run_stream(deps=deps):
+                if getattr(event, "type", None) == EventType.RUN_FINISHED:
+                    yield StateSnapshotEvent(
+                        type=EventType.STATE_SNAPSHOT,
+                        snapshot=deps.state,
+                    )
+                yield event
+
+        async for chunk in adapter.encode_stream(with_final_state()):
             yield chunk
 
     return StreamingResponse(

haiku_rag_slim/haiku/rag/chat/__init__.py

@@ -56,10 +56,18 @@ def run_chat(
             )
         )
 
+    # Drive with the analysis model when analysis is the only capability, so the
+    # running model matches the one the analysis capability configures (including
+    # its vision flag). RAG runs on the QA model.
+    if "rag" not in enabled and "analysis" in enabled:
+        driving_model = config.analysis.model or config.qa.model
+    else:
+        driving_model = config.qa.model
+
     app = ChatApp(
         db_path,
         capabilities=capability_list,
         read_only=read_only,
-        model=model or get_model(config.qa.model, config),
+        model=model or get_model(driving_model, config),
     )
     app.run()

tests/chat/test_chat_app.py

@@ -46,6 +46,41 @@ def test_run_chat_defers_multiple_capabilities(temp_db_path: Path):
     assert all(capability.defer_loading for capability in attached)
 
 
+@pytest.mark.parametrize(
+    ("enabled", "expected_model"),
+    [
+        (["analysis"], "analysis-model"),
+        (["rag"], "qa-model"),
+        (["rag", "analysis"], "qa-model"),
+    ],
+)
+def test_run_chat_drives_analysis_only_with_analysis_model(
+    temp_db_path: Path, enabled, expected_model
+):
+    """Analysis-only chat runs on analysis.model; otherwise on qa.model."""
+    from haiku.rag.config.models import AppConfig, ModelConfig
+
+    config = AppConfig()
+    config.qa.model = ModelConfig(provider="openai", name="qa-model")
+    config.analysis.model = ModelConfig(provider="openai", name="analysis-model")
+    captured: dict[str, str] = {}
+
+    def fake_get_model(model_config, _config):
+        captured["name"] = model_config.name
+        return "resolved-model"
+
+    with (
+        patch("haiku.rag.chat.app.ChatApp"),
+        patch("haiku.rag.config.get_config", return_value=config),
+        patch("haiku.rag.utils.get_model", side_effect=fake_get_model),
+    ):
+        from haiku.rag.chat import run_chat
+
+        run_chat(db_path=temp_db_path, capabilities=enabled)
+
+    assert captured["name"] == expected_model
+
+
 def _make_mock_client():
     """Create a mock HaikuRAG client."""
     mock_client = AsyncMock()