Make Chat abort actually stop the agent turn.

Wire should_stop from aborted runs, snapshot abort before clearing flags so finals are not emitted after stop, and keep fail-path turn_uuid aligned for bubble splitting.

Co-authored-by: Cursor <cursoragent@cursor.com>

tests/test_turn_runner_abort.py

@@ -0,0 +1,164 @@
+"""WS turn abort must stop the agent and not emit a full final reply."""
+
+from __future__ import annotations
+
+import asyncio
+from dataclasses import dataclass
+from typing import Any
+
+import pytest
+
+
+@dataclass
+class DummyResult:
+    run_id: str = "run-abort-1"
+    reply_text: str = "FULL_REPLY_SHOULD_NOT_EMIT"
+    elapsed_ms: int = 10
+    turn_uuid: str = "tu-1"
+    mode: str = "sync_direct"
+    trace_id: str = ""
+
+
+class DummyLock:
+    def __enter__(self) -> None:
+        return None
+
+    def __exit__(self, exc_type, exc, tb) -> bool:
+        return False
+
+
+class DummyStore:
+    def __init__(self) -> None:
+        self.messages: list[dict[str, Any]] = []
+
+    def get_messages(self, session_id: str, limit: int = 200) -> list[Any]:
+        return []
+
+    def add_message(self, **kwargs: Any) -> None:
+        self.messages.append(dict(kwargs))
+
+    def oclaw_run_get(self, **kwargs: Any) -> None:
+        return None
+
+    def get_setting(self, key: str) -> str:
+        return ""
+
+
+class DummyConn:
+    def __init__(self) -> None:
+        self.auth_ctx = {"tenant_id": "t", "user_id": "u", "username": "n", "lang": "zh"}
+        self._abort_lock = DummyLock()
+        self._active_run_session: dict[str, str] = {}
+        self._aborted_run_ids: set[str] = set()
+        self._is_webchat_client = True
+        self._subscribed_sessions_changed = False
+        self.chat_calls: list[dict[str, Any]] = []
+        self.events: list[tuple[str, Any]] = []
+
+    async def emit_agent_event(self, **kwargs: Any) -> None:
+        return None
+
+    async def emit_chat_event(self, **kwargs: Any) -> None:
+        self.chat_calls.append(dict(kwargs))
+
+    async def send_event(self, event: str, payload: Any) -> None:
+        self.events.append((str(event), payload))
+
+    async def send_res(self, *args: Any, **kwargs: Any) -> None:
+        return None
+
+
+def _run_turn(monkeypatch: pytest.MonkeyPatch, *, gateway: Any, conn: DummyConn, run_id: str) -> DummyStore:
+    from interfaces.ws import turn_runner
+
+    store = DummyStore()
+    monkeypatch.setattr(turn_runner, "SqliteStore", lambda _p: store)
+    monkeypatch.setattr(turn_runner, "db_path", lambda: "dummy.sqlite")
+    monkeypatch.setattr(turn_runner, "OclawGateway", lambda store: gateway)
+    monkeypatch.setattr(turn_runner, "build_gateway_executor", lambda *a, **k: object())
+    monkeypatch.setattr(turn_runner, "get_assistant_store", lambda: store)
+    monkeypatch.setattr(turn_runner, "persist_assistant_text_if_turn_missing", lambda **k: False)
+
+    asyncio.run(
+        turn_runner.run_agent_turn_via_bridge(
+            conn=conn,
+            req_id="req-1",
+            p={"message": "hello", "runId": run_id, "lang": "zh"},
+            session_id="s1",
+            send_response=False,
+            normalize_ws_attachments=lambda _a: [],
+            validate_relay_share_envelope=lambda _e: (True, "", {}),
+            now_ms=lambda: 1,
+            error_shape=lambda c, m: {"code": c, "message": m},
+        )
+    )
+    return store
+
+
+def test_abort_before_finish_emits_aborted_not_final(monkeypatch: pytest.MonkeyPatch) -> None:
+    run_id = "run-abort-1"
+    seen_should_stop: list[bool] = []
+
+    class Gateway:
+        def handle_turn(self, **kwargs: Any) -> DummyResult:
+            assert callable(kwargs.get("should_stop"))
+            seen_should_stop.append(bool(kwargs["should_stop"]()))
+            return DummyResult(run_id=run_id, reply_text="FULL_REPLY_SHOULD_NOT_EMIT")
+
+    conn = DummyConn()
+    conn._aborted_run_ids.add(run_id)
+    store = _run_turn(monkeypatch, gateway=Gateway(), conn=conn, run_id=run_id)
+
+    states = [c.get("state") for c in conn.chat_calls]
+    assert "aborted" in states
+    assert "final" not in states
+    assert seen_should_stop == [True]
+    assert any("已中断" in str(m.get("content") or "") for m in store.messages)
+    assert not any("FULL_REPLY" in str(m.get("content") or "") for m in store.messages)
+
+
+def test_on_token_raises_when_aborted_mid_stream(monkeypatch: pytest.MonkeyPatch) -> None:
+    run_id = "run-abort-2"
+    conn = DummyConn()
+
+    class Gateway:
+        def handle_turn(self, **kwargs: Any) -> DummyResult:
+            on_token = kwargs["on_token"]
+            on_token("partial")
+            conn._aborted_run_ids.add(run_id)
+            with pytest.raises(RuntimeError, match="interrupted"):
+                on_token("more")
+            raise RuntimeError("generation interrupted by user")
+
+    store = _run_turn(monkeypatch, gateway=Gateway(), conn=conn, run_id=run_id)
+    states = [c.get("state") for c in conn.chat_calls]
+    assert "aborted" in states
+    assert "final" not in states
+    assert any("已中断" in str(m.get("content") or "") for m in store.messages)
+
+
+def test_abort_chat_session_marks_active_runs(monkeypatch: pytest.MonkeyPatch) -> None:
+    from interfaces.ws import server_methods_bridge
+
+    monkeypatch.setattr(server_methods_bridge, "get_assistant_store", lambda: DummyStore())
+    lock = DummyLock()
+    active: dict[str, str] = {"r1": "s1", "r2": "s2"}
+    aborted: set[str] = set()
+
+    ctx = server_methods_bridge.build_gateway_context(
+        conn_id="c1",
+        subscribed_sessions_changed=False,
+        subscribed_message_keys=set(),
+        abort_lock=lock,
+        active_run_session=active,
+        aborted_run_ids=aborted,
+        run_agent_turn=lambda *a, **k: None,
+        normalize_ws_attachments=lambda _a: [],
+        validate_relay_share_envelope=lambda _e: (True, "", {}),
+        now_ms=lambda: 1,
+    )
+    assert ctx["abort_chat_session"]("s1") is True
+    assert "r1" in aborted
+    assert "r2" not in aborted
+    assert ctx["abort_chat_run"]("r2") is True
+    assert "r2" in aborted

tests/test_turn_runner_turn_uuid.py

@@ -0,0 +1,53 @@
+"""turn_uuid recovery / final payload alignment for WS turn runner."""
+
+from __future__ import annotations
+
+import uuid
+
+from interfaces.ws.turn_runner import _latest_turn_uuid
+from svc.persistence.sqlite_store import SqliteStore
+
+
+def _seed_session(store: SqliteStore) -> str:
+    t = store.create_tenant("T")
+    tid = str(t["id"])
+    store.create_user_account(
+        tenant_id=tid,
+        username="u",
+        display_name="U",
+        role="owner",
+        password_hash="x",
+        is_active=True,
+    )
+    uid = str(store.get_user_by_username(tenant_id=tid, username="u")["id"])
+    s = store.create_session_for_user(title="s", tenant_id=tid, user_id=uid)
+    return str(s.id)
+
+
+def test_latest_turn_uuid_prefers_latest_user(tmp_path) -> None:
+    store = SqliteStore(str(tmp_path / "t.sqlite"))
+    sid = _seed_session(store)
+    older = uuid.uuid4().hex
+    newer = uuid.uuid4().hex
+    store.add_message(session_id=sid, role="user", content="a", turn_uuid=older, event_type="user_text")
+    store.add_message(
+        session_id=sid, role="assistant", content="ok", turn_uuid=older, event_type="assistant_text"
+    )
+    store.add_message(session_id=sid, role="user", content="b", turn_uuid=newer, event_type="user_text")
+    assert _latest_turn_uuid(store, sid) == newer
+
+
+def test_latest_turn_uuid_falls_back_to_assistant(tmp_path) -> None:
+    store = SqliteStore(str(tmp_path / "u.sqlite"))
+    sid = _seed_session(store)
+    tu = uuid.uuid4().hex
+    store.add_message(
+        session_id=sid, role="assistant", content="partial", turn_uuid=tu, event_type="assistant_text"
+    )
+    assert _latest_turn_uuid(store, sid) == tu
+
+
+def test_latest_turn_uuid_empty_session(tmp_path) -> None:
+    store = SqliteStore(str(tmp_path / "e.sqlite"))
+    sid = _seed_session(store)
+    assert _latest_turn_uuid(store, sid) == ""