refactor(chat): implement turn cancellation and status management in new chat routes for improved user experience and error handling

surfsense_backend/tests/unit/agents/new_chat/test_busy_mutex.py

@@ -7,7 +7,9 @@ import pytest
 from app.agents.new_chat.errors import BusyError
 from app.agents.new_chat.middleware.busy_mutex import (
     BusyMutexMiddleware,
+    end_turn,
     get_cancel_event,
+    is_cancel_requested,
     manager,
     request_cancel,
     reset_cancel,
@@ -88,3 +90,31 @@ async def test_no_thread_id_skipped_when_not_required() -> None:
 def test_reset_cancel_idempotent() -> None:
     # Should not raise even if event was never created
     reset_cancel("never-seen")
+
+
+def test_request_cancel_creates_event_for_unseen_thread() -> None:
+    thread_id = "never-seen-cancel"
+    reset_cancel(thread_id)
+
+    assert request_cancel(thread_id) is True
+    assert get_cancel_event(thread_id).is_set()
+    assert is_cancel_requested(thread_id) is True
+
+
+@pytest.mark.asyncio
+async def test_end_turn_force_clears_lock_and_cancel_state() -> None:
+    thread_id = "forced-end-turn"
+    mw = BusyMutexMiddleware()
+    runtime = _Runtime(thread_id)
+
+    await mw.abefore_agent({}, runtime)
+    assert manager.lock_for(thread_id).locked()
+
+    request_cancel(thread_id)
+    assert is_cancel_requested(thread_id) is True
+
+    end_turn(thread_id)
+
+    assert not manager.lock_for(thread_id).locked()
+    assert not get_cancel_event(thread_id).is_set()
+    assert is_cancel_requested(thread_id) is False

surfsense_backend/tests/unit/test_stream_new_chat_contract.py

@@ -8,6 +8,7 @@ import pytest
 
 import app.tasks.chat.stream_new_chat as stream_new_chat_module
 from app.agents.new_chat.errors import BusyError
+from app.agents.new_chat.middleware.busy_mutex import request_cancel, reset_cancel
 from app.tasks.chat.stream_new_chat import (
     StreamResult,
     _classify_stream_exception,
@@ -147,7 +148,7 @@ def test_stream_exception_classifies_rate_limited():
     exc = Exception(
         '{"error":{"type":"rate_limit_error","message":"Rate limited. Please try again later."}}'
     )
-    kind, code, severity, is_expected, user_message = _classify_stream_exception(
+    kind, code, severity, is_expected, user_message, extra = _classify_stream_exception(
         exc, flow_label="chat"
     )
     assert kind == "rate_limited"
@@ -155,11 +156,12 @@ def test_stream_exception_classifies_rate_limited():
     assert severity == "warn"
     assert is_expected is True
     assert "temporarily rate-limited" in user_message
+    assert extra is None
 
 
 def test_stream_exception_classifies_thread_busy():
     exc = BusyError(request_id="thread-123")
-    kind, code, severity, is_expected, user_message = _classify_stream_exception(
+    kind, code, severity, is_expected, user_message, extra = _classify_stream_exception(
         exc, flow_label="chat"
     )
     assert kind == "thread_busy"
@@ -167,11 +169,12 @@ def test_stream_exception_classifies_thread_busy():
     assert severity == "warn"
     assert is_expected is True
     assert "still finishing for this thread" in user_message
+    assert extra is None
 
 
 def test_stream_exception_classifies_thread_busy_from_message():
     exc = Exception("Thread is busy with another request")
-    kind, code, severity, is_expected, user_message = _classify_stream_exception(
+    kind, code, severity, is_expected, user_message, extra = _classify_stream_exception(
         exc, flow_label="chat"
     )
     assert kind == "thread_busy"
@@ -179,6 +182,24 @@ def test_stream_exception_classifies_thread_busy_from_message():
     assert severity == "warn"
     assert is_expected is True
     assert "still finishing for this thread" in user_message
+    assert extra is None
+
+
+def test_stream_exception_classifies_turn_cancelling_when_cancel_requested():
+    thread_id = "thread-cancelling-1"
+    reset_cancel(thread_id)
+    request_cancel(thread_id)
+    exc = BusyError(request_id=thread_id)
+    kind, code, severity, is_expected, user_message, extra = _classify_stream_exception(
+        exc, flow_label="chat"
+    )
+    assert kind == "thread_busy"
+    assert code == "TURN_CANCELLING"
+    assert severity == "info"
+    assert is_expected is True
+    assert "stopping" in user_message
+    assert isinstance(extra, dict)
+    assert "retry_after_ms" in extra
 
 
 def test_premium_classification_is_error_code_driven():
@@ -219,33 +240,33 @@ def test_toast_only_pre_accept_policy_has_no_inline_failed_marker():
 def test_network_send_failures_use_unified_retry_toast_message():
     classifier_path = Path(__file__).resolve().parents[3] / "surfsense_web/lib/chat/chat-error-classifier.ts"
     classifier_source = classifier_path.read_text(encoding="utf-8")
-    page_path = (
-        Path(__file__).resolve().parents[3]
-        / "surfsense_web/app/dashboard/[search_space_id]/new-chat/[[...chat_id]]/page.tsx"
+    request_errors_path = (
+        Path(__file__).resolve().parents[3] / "surfsense_web/lib/chat/chat-request-errors.ts"
     )
-    page_source = page_path.read_text(encoding="utf-8")
+    request_errors_source = request_errors_path.read_text(encoding="utf-8")
 
     assert '"send_failed_pre_accept"' in classifier_source
     assert 'errorCode === "SEND_FAILED_PRE_ACCEPT"' in classifier_source
+    assert 'errorCode === "TURN_CANCELLING"' in classifier_source
     assert "if (withCode.code) return withCode.code;" in classifier_source
     assert 'userMessage: "Message not sent. Please retry."' in classifier_source
     assert 'userMessage: "Connection issue. Please try again."' in classifier_source
-    assert "tagPreAcceptSendFailure(error)" in page_source
-    assert "const passthroughCodes = new Set([" in page_source
-    assert '"PREMIUM_QUOTA_EXHAUSTED"' in page_source
-    assert '"THREAD_BUSY"' in page_source
-    assert '"AUTH_EXPIRED"' in page_source
-    assert '"UNAUTHORIZED"' in page_source
-    assert '"RATE_LIMITED"' in page_source
-    assert '"NETWORK_ERROR"' in page_source
-    assert '"STREAM_PARSE_ERROR"' in page_source
-    assert '"TOOL_EXECUTION_ERROR"' in page_source
-    assert '"PERSIST_MESSAGE_FAILED"' in page_source
-    assert '"SERVER_ERROR"' in page_source
-    assert "passthroughCodes.has(existingCode)" in page_source
-    assert 'errorCode: "SEND_FAILED_PRE_ACCEPT"' in page_source
-    assert 'errorCode: "NETWORK_ERROR"' not in page_source
-    assert "Failed to start chat. Please try again." not in page_source
+    assert "const passthroughCodes = new Set([" in request_errors_source
+    assert '"PREMIUM_QUOTA_EXHAUSTED"' in request_errors_source
+    assert '"THREAD_BUSY"' in request_errors_source
+    assert '"TURN_CANCELLING"' in request_errors_source
+    assert '"AUTH_EXPIRED"' in request_errors_source
+    assert '"UNAUTHORIZED"' in request_errors_source
+    assert '"RATE_LIMITED"' in request_errors_source
+    assert '"NETWORK_ERROR"' in request_errors_source
+    assert '"STREAM_PARSE_ERROR"' in request_errors_source
+    assert '"TOOL_EXECUTION_ERROR"' in request_errors_source
+    assert '"PERSIST_MESSAGE_FAILED"' in request_errors_source
+    assert '"SERVER_ERROR"' in request_errors_source
+    assert "passthroughCodes.has(existingCode)" in request_errors_source
+    assert 'errorCode: "SEND_FAILED_PRE_ACCEPT"' in request_errors_source
+    assert 'errorCode: "NETWORK_ERROR"' not in request_errors_source
+    assert "Failed to start chat. Please try again." not in classifier_source
 
 
 def test_pre_post_accept_abort_contract_exists_for_new_resume_regenerate_flows():
@@ -269,3 +290,75 @@ def test_pre_post_accept_abort_contract_exists_for_new_resume_regenerate_flows()
 
     # New flow persists only when accepted and not already persisted.
     assert "if (newAccepted && !userPersisted) {" in source
+    assert "const fetchWithTurnCancellingRetry = useCallback(" in source
+    assert "computeFallbackTurnCancellingRetryDelay" in source
+    assert 'withMeta.errorCode === "TURN_CANCELLING"' in source
+    assert 'withMeta.errorCode === "THREAD_BUSY"' in source
+    assert "await fetchWithTurnCancellingRetry(() =>" in source
+
+
+def test_cancel_active_turn_route_contract_exists():
+    routes_path = (
+        Path(__file__).resolve().parents[3]
+        / "surfsense_backend/app/routes/new_chat_routes.py"
+    )
+    source = routes_path.read_text(encoding="utf-8")
+
+    assert '@router.post(\n    "/threads/{thread_id}/cancel-active-turn",' in source
+    assert "response_model=CancelActiveTurnResponse" in source
+    assert 'status="cancelling",' in source
+    assert 'error_code="TURN_CANCELLING",' in source
+    assert "retry_after_ms=retry_after_ms if retry_after_ms > 0 else None," in source
+    assert "retry_after_at=" in source
+    assert 'status="idle",' in source
+    assert 'error_code="NO_ACTIVE_TURN",' in source
+
+
+def test_turn_status_route_contract_exists():
+    routes_path = (
+        Path(__file__).resolve().parents[3]
+        / "surfsense_backend/app/routes/new_chat_routes.py"
+    )
+    source = routes_path.read_text(encoding="utf-8")
+
+    assert '@router.get(\n    "/threads/{thread_id}/turn-status",' in source
+    assert "response_model=TurnStatusResponse" in source
+    assert "_build_turn_status_payload(thread_id)" in source
+    assert "Permission.CHATS_READ.value" in source
+    assert "_raise_if_thread_busy_for_start(" in source
+
+
+def test_turn_cancelling_retry_policy_contract_exists():
+    routes_path = (
+        Path(__file__).resolve().parents[3]
+        / "surfsense_backend/app/routes/new_chat_routes.py"
+    )
+    source = routes_path.read_text(encoding="utf-8")
+
+    assert "TURN_CANCELLING_INITIAL_DELAY_MS = 200" in source
+    assert "TURN_CANCELLING_BACKOFF_FACTOR = 2" in source
+    assert "TURN_CANCELLING_MAX_DELAY_MS = 1500" in source
+    assert "def _compute_turn_cancelling_retry_delay(" in source
+    assert "retry-after-ms" in source
+    assert '"Retry-After"' in source
+    assert '"errorCode": "TURN_CANCELLING"' in source
+
+
+def test_turn_status_sse_contract_exists():
+    stream_source = (
+        Path(__file__).resolve().parents[3]
+        / "surfsense_backend/app/tasks/chat/stream_new_chat.py"
+    ).read_text(encoding="utf-8")
+    state_source = (
+        Path(__file__).resolve().parents[3] / "surfsense_web/lib/chat/streaming-state.ts"
+    ).read_text(encoding="utf-8")
+    pipeline_source = (
+        Path(__file__).resolve().parents[3] / "surfsense_web/lib/chat/stream-pipeline.ts"
+    ).read_text(encoding="utf-8")
+
+    assert '"turn-status"' in stream_source
+    assert '"status": "busy"' in stream_source
+    assert '"status": "idle"' in stream_source
+    assert "type: \"data-turn-status\"" in state_source
+    assert "case \"data-turn-status\":" in pipeline_source
+    assert "end_turn(str(chat_id))" in stream_source