mirror of
https://github.com/turnstonelabs/turnstone.git
synced 2026-08-12 23:12:23 -06:00
fix: cancel button race condition with stream abort and force cancel (#202)
The cancel endpoint emitted a 'cancelled' SSE event before the worker thread terminated. The frontend transitioned to "send" mode prematurely, so the next send got rejected with "Already processing a request." Backend: - Providers expose SDK stream handle via cancel_ref parameter so cancel() can close the HTTP connection and unblock iteration - Generation counter prevents orphaned threads from mutating messages or clearing cancel state after force cancel - _check_cancelled() added between retry attempts in _try_stream - Server polls (async, non-blocking) for cancelled worker to exit - Force cancel (force:true) abandons stuck worker, keeps cancel event set so subprocesses are killed, guards against spurious SSE events Frontend: - 'cancelled' shows "Cancelling..." then escalates to "Force Stop" after 2s for a harder cancel that abandons the worker immediately - 10s safety timeout auto-recovers if stream_end never arrives - busy_error re-enables stop button instead of showing send - Timeout cleanup in disconnectSSE, stream_end, and force .then() - Layout shift prevention (min-width, white-space: nowrap) - aria-label updates for accessibility Tests: - 7 new tests: stream close, error suppression, cancel_ref population, transport error conversion, non-cancel exception propagation, retry cancellation check
This commit is contained in:
Patrick Buckley
+21
-6
@@ -452,9 +452,12 @@ after `/clear` or `/new` commands).
|
||||
{"type": "clear_ui"}
|
||||
```
|
||||
|
||||
**`cancelled`** -- the generation was cancelled by the user (via the Stop
|
||||
button or `POST /v1/api/cancel`). The client should finalize any in-progress
|
||||
assistant message with whatever partial content was streamed.
|
||||
**`cancelled`** -- a cancel request was acknowledged (via the Stop button or
|
||||
`POST /v1/api/cancel`). This signals that cancellation is in progress, not
|
||||
that it is complete. The worker thread may still be finishing — wait for
|
||||
`stream_end` before transitioning to a ready state. The client should clear
|
||||
any in-progress assistant rendering but not re-enable the send button until
|
||||
`stream_end` arrives.
|
||||
|
||||
```json
|
||||
{"type": "cancelled"}
|
||||
@@ -793,23 +796,35 @@ containing the resumed session's messages.
|
||||
|
||||
Cancels the active generation in a workstream. Sets a cooperative cancellation
|
||||
flag that is checked at multiple points in the generation loop (per streaming
|
||||
chunk, before tool execution, inside bash commands). The session transitions to
|
||||
`idle` state and preserves any partial content already streamed.
|
||||
chunk, before tool execution, inside bash commands). Also closes the underlying
|
||||
HTTP stream to the LLM provider, unblocking any pending read immediately.
|
||||
The session transitions to `idle` state and preserves any partial content
|
||||
already streamed.
|
||||
|
||||
If the workstream is waiting for tool approval or plan review, the pending
|
||||
prompt is automatically denied/rejected to unblock the worker thread.
|
||||
|
||||
Calling this endpoint when the workstream is already idle is a harmless no-op.
|
||||
|
||||
**Force cancel:** When `force` is `true`, the server abandons the stuck worker
|
||||
thread immediately and transitions the workstream to `idle`. The abandoned
|
||||
thread continues to wind down in the background (killing any running
|
||||
subprocesses and exiting at the next cancellation checkpoint). During this
|
||||
wind-down it may emit a final `stream_end` event which the server suppresses
|
||||
for the orphaned thread. Use force cancel when cooperative cancel has not
|
||||
resolved within a few seconds — the web UI offers this as a "Force Stop"
|
||||
button automatically.
|
||||
|
||||
**Request body:**
|
||||
|
||||
```json
|
||||
{"ws_id": "abc123"}
|
||||
{"ws_id": "abc123", "force": false}
|
||||
```
|
||||
|
||||
| Field | Type | Required | Description |
|
||||
|--------|--------|----------|----------------------|
|
||||
| `ws_id`| string | yes | Target workstream ID |
|
||||
| `force`| bool | no | Abandon stuck worker immediately (default: `false`) |
|
||||
|
||||
**Response:**
|
||||
|
||||
|
||||
@@ -40,12 +40,23 @@ running --> error : Exception during\ntool execution
|
||||
|
||||
error --> thinking : New send() call\n_emit_state("thinking")
|
||||
|
||||
thinking --> idle : cancel() called\n_emit_state("idle")
|
||||
thinking --> idle : cancel() called\nstream aborted\n_emit_state("idle")
|
||||
|
||||
running --> idle : cancel() called\n_emit_state("idle")
|
||||
|
||||
attention --> idle : cancel() unblocks\napproval/plan wait\n_emit_state("idle")
|
||||
|
||||
note left of idle
|
||||
**Cancel escalation:**
|
||||
1. **Cooperative**: cancel() sets event + closes
|
||||
SDK stream → worker exits at next checkpoint
|
||||
2. **Force**: force=true abandons the worker
|
||||
thread, emits stream_end immediately.
|
||||
Orphaned thread still kills subprocesses
|
||||
but skips message mutations (generation
|
||||
counter prevents stale writes).
|
||||
end note
|
||||
|
||||
note right of thinking
|
||||
**Emitted via:**
|
||||
session._emit_state(state)
|
||||
|
||||
@@ -1,3 +1,3 @@
|
||||
version https://git-lfs.github.com/spec/v1
|
||||
oid sha256:7896c6e041b6dbb89d034468fa980c8fe645df5eb969d45ef966ccc6399edac2
|
||||
size 200083
|
||||
oid sha256:7bf27afa267d5b8d6da38e83213ed1b8e87639d5105a0a1ccc2e5a4bf4d3b67e
|
||||
size 185282
|
||||
|
||||
+1
-1
@@ -75,7 +75,7 @@ Both `TurnstoneServer` (sync) and `AsyncTurnstoneServer` (async) expose:
|
||||
| | `approve(*, ws_id, approved, feedback, always)` | `StatusResponse` |
|
||||
| | `plan_feedback(*, ws_id, feedback)` | `StatusResponse` |
|
||||
| | `command(*, ws_id, command)` | `StatusResponse` |
|
||||
| | `cancel(ws_id)` | `StatusResponse` |
|
||||
| | `cancel(ws_id, *, force=False)` | `StatusResponse` |
|
||||
| **Streaming** | `stream_events(ws_id)` | `Iterator[ServerEvent]` |
|
||||
| | `stream_global_events()` | `Iterator[ServerEvent]` |
|
||||
| **High-level** | `send_and_wait(message, ws_id, *, timeout, on_event)` | `TurnResult` |
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"openapi": "3.1.0",
|
||||
"info": {
|
||||
"title": "turnstone Console API",
|
||||
"version": "0.8.4",
|
||||
"version": "0.9.0",
|
||||
"description": "Cluster-wide visibility and control across all turnstone nodes."
|
||||
},
|
||||
"paths": {
|
||||
@@ -3002,7 +3002,7 @@
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/StatusResponse"
|
||||
"$ref": "#/components/schemas/DeleteSettingResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -3453,6 +3453,126 @@
|
||||
}
|
||||
}
|
||||
},
|
||||
"/v1/api/admin/tls/ca": {
|
||||
"get": {
|
||||
"summary": "CA status: initialization state, CN, cert count, cert inventory",
|
||||
"operationId": "v1_api_admin_tls_ca_get",
|
||||
"tags": [
|
||||
"Admin"
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Success"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/v1/api/admin/tls/ca.pem": {
|
||||
"get": {
|
||||
"summary": "Download CA root certificate (PEM format)",
|
||||
"operationId": "v1_api_admin_tls_ca.pem_get",
|
||||
"tags": [
|
||||
"Admin"
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Success"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/v1/api/admin/tls/certs": {
|
||||
"get": {
|
||||
"summary": "List all issued TLS certificates",
|
||||
"operationId": "v1_api_admin_tls_certs_get",
|
||||
"tags": [
|
||||
"Admin"
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Success"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/v1/api/admin/tls/certs/{domain}/renew": {
|
||||
"post": {
|
||||
"summary": "Force-renew a certificate by domain",
|
||||
"operationId": "v1_api_admin_tls_certs_{domain}_renew_post",
|
||||
"tags": [
|
||||
"Admin"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "domain",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Success"
|
||||
},
|
||||
"404": {
|
||||
"description": "Error 404",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"500": {
|
||||
"description": "Error 500",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/v1/api/admin/tls/certs/{domain}": {
|
||||
"delete": {
|
||||
"summary": "Delete a certificate by domain",
|
||||
"operationId": "v1_api_admin_tls_certs_{domain}_delete",
|
||||
"tags": [
|
||||
"Admin"
|
||||
],
|
||||
"parameters": [
|
||||
{
|
||||
"name": "domain",
|
||||
"in": "path",
|
||||
"required": true,
|
||||
"schema": {
|
||||
"type": "string"
|
||||
}
|
||||
}
|
||||
],
|
||||
"responses": {
|
||||
"200": {
|
||||
"description": "Success"
|
||||
},
|
||||
"404": {
|
||||
"description": "Error 404",
|
||||
"content": {
|
||||
"application/json": {
|
||||
"schema": {
|
||||
"$ref": "#/components/schemas/ErrorResponse"
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
},
|
||||
"/health": {
|
||||
"get": {
|
||||
"summary": "Console health check",
|
||||
@@ -3507,6 +3627,34 @@
|
||||
"title": "StatusResponse",
|
||||
"type": "object"
|
||||
},
|
||||
"DeleteSettingResponse": {
|
||||
"description": "DELETE /v1/api/admin/settings/{key} response.",
|
||||
"properties": {
|
||||
"status": {
|
||||
"default": "ok",
|
||||
"examples": [
|
||||
"ok"
|
||||
],
|
||||
"title": "Status",
|
||||
"type": "string"
|
||||
},
|
||||
"key": {
|
||||
"description": "Dotted setting key that was reset",
|
||||
"title": "Key",
|
||||
"type": "string"
|
||||
},
|
||||
"default": {
|
||||
"description": "Registry default value the setting reverted to",
|
||||
"title": "Default"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
"key",
|
||||
"default"
|
||||
],
|
||||
"title": "DeleteSettingResponse",
|
||||
"type": "object"
|
||||
},
|
||||
"AuthLoginRequest": {
|
||||
"description": "POST /v1/api/auth/login request body.\n\nEither username+password or token must be provided.",
|
||||
"properties": {
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
"openapi": "3.1.0",
|
||||
"info": {
|
||||
"title": "turnstone Server API",
|
||||
"version": "0.8.4",
|
||||
"version": "0.9.0",
|
||||
"description": "Single-node workstream management, chat interaction, and real-time streaming."
|
||||
},
|
||||
"paths": {
|
||||
@@ -1223,6 +1223,12 @@
|
||||
"description": "Target workstream ID",
|
||||
"title": "Ws Id",
|
||||
"type": "string"
|
||||
},
|
||||
"force": {
|
||||
"default": false,
|
||||
"description": "Force cancel: abandon the stuck worker thread immediately. Use when cooperative cancel has not resolved within a few seconds.",
|
||||
"title": "Force",
|
||||
"type": "boolean"
|
||||
}
|
||||
},
|
||||
"required": [
|
||||
|
||||
@@ -93,10 +93,13 @@ export class TurnstoneServer extends BaseClient {
|
||||
});
|
||||
}
|
||||
|
||||
async cancel(wsId: string): Promise<StatusResponse> {
|
||||
return this.request("POST", "/v1/api/cancel", {
|
||||
json: { ws_id: wsId },
|
||||
});
|
||||
async cancel(
|
||||
wsId: string,
|
||||
opts?: { force?: boolean },
|
||||
): Promise<StatusResponse> {
|
||||
const body: Record<string, unknown> = { ws_id: wsId };
|
||||
if (opts?.force) body.force = true;
|
||||
return this.request("POST", "/v1/api/cancel", { json: body });
|
||||
}
|
||||
|
||||
// -- Streaming ------------------------------------------------------------
|
||||
|
||||
+292
-1
@@ -7,7 +7,7 @@ from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from turnstone.core.session import ChatSession, GenerationCancelled
|
||||
from turnstone.core.session import ChatSession, GenerationCancelled, _CancelRef
|
||||
|
||||
|
||||
class NullUI:
|
||||
@@ -407,3 +407,294 @@ class TestStreamFlushBeforeToolCalls:
|
||||
stream_end_idx = next(i for i, e in enumerate(events) if e[0] == "stream_end")
|
||||
late_content = [e for e in events[stream_end_idx + 1 :] if e[0] == "content"]
|
||||
assert late_content == [], f"Content after stream_end: {late_content}"
|
||||
|
||||
|
||||
class TestStreamAbort:
|
||||
"""Tests for cancel() closing the underlying SDK stream."""
|
||||
|
||||
def test_cancel_closes_cancel_stream(self, tmp_db):
|
||||
"""cancel() calls .close() on the stored SDK stream handle."""
|
||||
session = _make_session()
|
||||
mock_stream = MagicMock()
|
||||
session._cancel_stream = mock_stream
|
||||
session.cancel()
|
||||
mock_stream.close.assert_called_once()
|
||||
assert session._cancel_event.is_set()
|
||||
|
||||
def test_cancel_without_stream_is_safe(self, tmp_db):
|
||||
"""cancel() with no active stream just sets the event."""
|
||||
session = _make_session()
|
||||
assert session._cancel_stream is None
|
||||
session.cancel() # Should not raise
|
||||
assert session._cancel_event.is_set()
|
||||
|
||||
def test_cancel_stream_close_error_suppressed(self, tmp_db):
|
||||
"""Errors from stream.close() are suppressed."""
|
||||
session = _make_session()
|
||||
mock_stream = MagicMock()
|
||||
mock_stream.close.side_effect = RuntimeError("already closed")
|
||||
session._cancel_stream = mock_stream
|
||||
session.cancel() # Should not raise
|
||||
assert session._cancel_event.is_set()
|
||||
|
||||
def test_cancel_ref_populated_after_first_chunk(self, tmp_db):
|
||||
"""_cancel_ref is populated by the provider after the first chunk
|
||||
arrives (lazy generator evaluation)."""
|
||||
ui = NullUI()
|
||||
session = _make_session(ui=ui)
|
||||
|
||||
@dataclass
|
||||
class FakeChunk:
|
||||
content_delta: str = ""
|
||||
reasoning_delta: str = ""
|
||||
tool_call_deltas: list = field(default_factory=list)
|
||||
usage: None = None
|
||||
finish_reason: str = "stop"
|
||||
info_delta: str = ""
|
||||
provider_blocks: list = field(default_factory=list)
|
||||
|
||||
sdk_stream = MagicMock()
|
||||
|
||||
def fake_provider_stream():
|
||||
# Simulate provider appending to cancel_ref before first yield
|
||||
session._cancel_ref.append(sdk_stream)
|
||||
yield FakeChunk(content_delta="hi", finish_reason="stop")
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
session,
|
||||
"_create_stream_with_retry",
|
||||
return_value=fake_provider_stream(),
|
||||
),
|
||||
patch.object(session, "_full_messages", return_value=[]),
|
||||
):
|
||||
session.send("test")
|
||||
|
||||
# After stream completes, cancel_stream should be cleared
|
||||
assert session._cancel_stream is None
|
||||
assert len(session._cancel_ref) == 0
|
||||
|
||||
def test_transport_error_during_cancel_becomes_generation_cancelled(self, tmp_db):
|
||||
"""When cancel() closes the stream, the resulting transport error
|
||||
is converted to GenerationCancelled."""
|
||||
ui = NullUI()
|
||||
session = _make_session(ui=ui)
|
||||
|
||||
@dataclass
|
||||
class FakeChunk:
|
||||
content_delta: str = ""
|
||||
reasoning_delta: str = ""
|
||||
tool_call_deltas: list = field(default_factory=list)
|
||||
usage: None = None
|
||||
finish_reason: str = ""
|
||||
info_delta: str = ""
|
||||
provider_blocks: list = field(default_factory=list)
|
||||
|
||||
def stream_that_errors():
|
||||
yield FakeChunk(content_delta="Hello")
|
||||
session._cancel_event.set()
|
||||
raise ConnectionError("stream closed")
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
session,
|
||||
"_create_stream_with_retry",
|
||||
return_value=stream_that_errors(),
|
||||
),
|
||||
patch.object(session, "_full_messages", return_value=[]),
|
||||
):
|
||||
session.send("test")
|
||||
|
||||
# Should complete as cancelled, not error
|
||||
assert "idle" in ui.states
|
||||
assert any("cancelled" in i.lower() for i in ui.infos)
|
||||
# Partial content preserved
|
||||
assistant_msgs = [m for m in session.messages if m["role"] == "assistant"]
|
||||
assert len(assistant_msgs) == 1
|
||||
assert assistant_msgs[0]["content"] == "Hello"
|
||||
|
||||
def test_non_cancel_exception_not_swallowed(self, tmp_db):
|
||||
"""Exceptions during streaming that aren't caused by cancel
|
||||
should propagate normally."""
|
||||
ui = NullUI()
|
||||
session = _make_session(ui=ui)
|
||||
|
||||
@dataclass
|
||||
class FakeChunk:
|
||||
content_delta: str = ""
|
||||
reasoning_delta: str = ""
|
||||
tool_call_deltas: list = field(default_factory=list)
|
||||
usage: None = None
|
||||
finish_reason: str = ""
|
||||
info_delta: str = ""
|
||||
provider_blocks: list = field(default_factory=list)
|
||||
|
||||
def stream_that_errors():
|
||||
yield FakeChunk(content_delta="Hello")
|
||||
raise ValueError("unexpected error")
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
session,
|
||||
"_create_stream_with_retry",
|
||||
return_value=stream_that_errors(),
|
||||
),
|
||||
patch.object(session, "_full_messages", return_value=[]),
|
||||
pytest.raises(ValueError, match="unexpected error"),
|
||||
):
|
||||
session.send("test")
|
||||
|
||||
def test_check_cancelled_between_retries(self, tmp_db):
|
||||
"""_try_stream checks for cancellation between retry attempts."""
|
||||
session = _make_session()
|
||||
session.cancel()
|
||||
|
||||
with pytest.raises(GenerationCancelled):
|
||||
session._try_stream(
|
||||
client=MagicMock(),
|
||||
model="test",
|
||||
msgs=[],
|
||||
)
|
||||
|
||||
|
||||
class TestCancelRef:
|
||||
"""Tests for the _CancelRef list proxy."""
|
||||
|
||||
def test_append_sets_cancel_stream(self, tmp_db):
|
||||
"""Appending a stream handle to _CancelRef sets _cancel_stream eagerly."""
|
||||
session = _make_session()
|
||||
mock_stream = MagicMock()
|
||||
assert session._cancel_stream is None
|
||||
|
||||
session._cancel_ref.append(mock_stream)
|
||||
|
||||
assert session._cancel_stream is mock_stream
|
||||
|
||||
def test_append_closes_stream_when_already_cancelled(self, tmp_db):
|
||||
"""If cancel is already set when a stream is appended, it is closed immediately."""
|
||||
session = _make_session()
|
||||
session.cancel() # Set cancel event before stream is created
|
||||
|
||||
mock_stream = MagicMock()
|
||||
session._cancel_ref.append(mock_stream)
|
||||
|
||||
mock_stream.close.assert_called_once()
|
||||
|
||||
def test_append_does_not_close_stream_when_not_cancelled(self, tmp_db):
|
||||
"""Stream is not closed if cancel hasn't been requested."""
|
||||
session = _make_session()
|
||||
mock_stream = MagicMock()
|
||||
|
||||
session._cancel_ref.append(mock_stream)
|
||||
|
||||
mock_stream.close.assert_not_called()
|
||||
assert session._cancel_stream is mock_stream
|
||||
|
||||
def test_append_close_error_suppressed(self, tmp_db):
|
||||
"""Errors from stream.close() during eager close are suppressed."""
|
||||
session = _make_session()
|
||||
session.cancel()
|
||||
|
||||
mock_stream = MagicMock()
|
||||
mock_stream.close.side_effect = RuntimeError("already closed")
|
||||
|
||||
session._cancel_ref.append(mock_stream) # Should not raise
|
||||
|
||||
def test_cancel_ref_is_cancel_ref_instance(self, tmp_db):
|
||||
"""ChatSession._cancel_ref is a _CancelRef instance."""
|
||||
session = _make_session()
|
||||
assert isinstance(session._cancel_ref, _CancelRef)
|
||||
|
||||
def test_cancel_ref_cleared_after_stream_ends(self, tmp_db):
|
||||
"""_cancel_ref is cleared in the send() finally block after streaming."""
|
||||
ui = NullUI()
|
||||
session = _make_session(ui=ui)
|
||||
mock_stream = MagicMock()
|
||||
session._cancel_ref.append(mock_stream)
|
||||
assert len(session._cancel_ref) == 1
|
||||
|
||||
@dataclass
|
||||
class FakeChunk:
|
||||
content_delta: str = ""
|
||||
reasoning_delta: str = ""
|
||||
tool_call_deltas: list = field(default_factory=list)
|
||||
usage: None = None
|
||||
finish_reason: str = "stop"
|
||||
info_delta: str = ""
|
||||
provider_blocks: list = field(default_factory=list)
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
session,
|
||||
"_create_stream_with_retry",
|
||||
return_value=iter([FakeChunk(content_delta="hi", finish_reason="stop")]),
|
||||
),
|
||||
patch.object(session, "_full_messages", return_value=[]),
|
||||
):
|
||||
session.send("test")
|
||||
|
||||
# After send() completes, _cancel_ref is cleared in the finally block
|
||||
assert len(session._cancel_ref) == 0
|
||||
|
||||
|
||||
class TestForceCancelGeneration:
|
||||
"""Tests for per-generation tracking that prevents orphaned-thread side-effects."""
|
||||
|
||||
def test_check_cancelled_raises_for_orphaned_generation(self, tmp_db):
|
||||
"""_check_cancelled raises GenerationCancelled when my_generation is stale."""
|
||||
session = _make_session()
|
||||
session._generation = 2 # Simulate two generations having run
|
||||
|
||||
with pytest.raises(GenerationCancelled):
|
||||
session._check_cancelled(my_generation=1) # Generation 1 is orphaned
|
||||
|
||||
def test_check_cancelled_ok_for_current_generation(self, tmp_db):
|
||||
"""_check_cancelled does not raise when my_generation matches current."""
|
||||
session = _make_session()
|
||||
session._generation = 3
|
||||
session._check_cancelled(my_generation=3) # Should not raise
|
||||
|
||||
def test_force_cancel_orphaned_thread_does_not_mutate_messages(self, tmp_db):
|
||||
"""An abandoned generation (force-cancel) cannot append to session.messages."""
|
||||
ui = NullUI()
|
||||
session = _make_session(ui=ui)
|
||||
|
||||
# We can't trivially test the full threading scenario in a unit test,
|
||||
# so directly verify that _check_cancelled raises when my_generation
|
||||
# is stale, which is what guards _stream_response against orphaned
|
||||
# (force-cancelled) threads continuing to mutate messages.
|
||||
session._generation = 5
|
||||
with pytest.raises(GenerationCancelled):
|
||||
session._check_cancelled(my_generation=4) # orphaned generation
|
||||
|
||||
def test_new_cancel_event_per_generation_in_send(self, tmp_db):
|
||||
"""send() replaces _cancel_event with a fresh Event each generation."""
|
||||
ui = NullUI()
|
||||
session = _make_session(ui=ui)
|
||||
|
||||
@dataclass
|
||||
class FakeChunk:
|
||||
content_delta: str = ""
|
||||
reasoning_delta: str = ""
|
||||
tool_call_deltas: list = field(default_factory=list)
|
||||
usage: None = None
|
||||
finish_reason: str = "stop"
|
||||
info_delta: str = ""
|
||||
provider_blocks: list = field(default_factory=list)
|
||||
|
||||
original_event = session._cancel_event
|
||||
|
||||
with (
|
||||
patch.object(
|
||||
session,
|
||||
"_create_stream_with_retry",
|
||||
return_value=iter([FakeChunk(content_delta="hi", finish_reason="stop")]),
|
||||
),
|
||||
patch.object(session, "_full_messages", return_value=[]),
|
||||
):
|
||||
session.send("test")
|
||||
|
||||
# After send() completes, _cancel_event should be a NEW Event
|
||||
# (not the same object as before the call).
|
||||
assert session._cancel_event is not original_event
|
||||
assert not session._cancel_event.is_set()
|
||||
|
||||
@@ -41,6 +41,11 @@ class CommandRequest(BaseModel):
|
||||
|
||||
class CancelRequest(BaseModel):
|
||||
ws_id: str = Field(description="Target workstream ID")
|
||||
force: bool = Field(
|
||||
default=False,
|
||||
description="Force cancel: abandon the stuck worker thread immediately. "
|
||||
"Use when cooperative cancel has not resolved within a few seconds.",
|
||||
)
|
||||
|
||||
|
||||
class CreateWorkstreamRequest(BaseModel):
|
||||
|
||||
@@ -459,6 +459,7 @@ class AnthropicProvider:
|
||||
reasoning_effort: str = "medium",
|
||||
extra_params: dict[str, Any] | None = None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
cancel_ref: list[Any] | None = None,
|
||||
) -> Iterator[StreamChunk]:
|
||||
_ensure_anthropic()
|
||||
caps = self.get_capabilities(model)
|
||||
@@ -477,6 +478,8 @@ class AnthropicProvider:
|
||||
)
|
||||
|
||||
with client.messages.stream(**kwargs) as stream:
|
||||
if cancel_ref is not None:
|
||||
cancel_ref.append(stream)
|
||||
yield from self._iter_anthropic_stream(stream)
|
||||
|
||||
def _iter_anthropic_stream(self, stream: Any) -> Iterator[StreamChunk]:
|
||||
|
||||
@@ -310,6 +310,7 @@ class OpenAIProvider:
|
||||
reasoning_effort: str = "medium",
|
||||
extra_params: dict[str, Any] | None = None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
cancel_ref: list[Any] | None = None,
|
||||
) -> Iterator[StreamChunk]:
|
||||
caps = self.get_capabilities(model)
|
||||
messages = self._sanitize_messages(messages)
|
||||
@@ -330,6 +331,8 @@ class OpenAIProvider:
|
||||
kwargs["extra_body"] = extra_params
|
||||
|
||||
stream = client.chat.completions.create(**kwargs)
|
||||
if cancel_ref is not None:
|
||||
cancel_ref.append(stream)
|
||||
yield from self._iter_stream(stream)
|
||||
|
||||
def _iter_stream(self, stream: Any) -> Iterator[StreamChunk]:
|
||||
|
||||
@@ -125,8 +125,15 @@ class LLMProvider(Protocol):
|
||||
reasoning_effort: str = "medium",
|
||||
extra_params: dict[str, Any] | None = None,
|
||||
deferred_names: frozenset[str] | None = None,
|
||||
cancel_ref: list[Any] | None = None,
|
||||
) -> Iterator[StreamChunk]:
|
||||
"""Create a streaming request, yielding normalized StreamChunks."""
|
||||
"""Create a streaming request, yielding normalized StreamChunks.
|
||||
|
||||
If *cancel_ref* is provided the provider appends the underlying SDK
|
||||
stream object (which has a ``.close()`` method) before yielding the
|
||||
first chunk. The caller can then close it from another thread to
|
||||
abort a blocked HTTP read immediately.
|
||||
"""
|
||||
...
|
||||
|
||||
def create_completion(
|
||||
|
||||
+146
-9
@@ -117,6 +117,35 @@ class GenerationCancelled(BaseException):
|
||||
"""
|
||||
|
||||
|
||||
class _CancelRef(list[Any]):
|
||||
"""List proxy used for ``ChatSession._cancel_ref``.
|
||||
|
||||
Providers call ``cancel_ref.append(stream_handle)`` inside a generator
|
||||
body — the generator body doesn't execute until the first ``next()`` call,
|
||||
i.e. just before the first chunk is yielded. By overriding ``append`` we
|
||||
update ``ChatSession._cancel_stream`` eagerly at that moment. If
|
||||
cancellation was already requested before the first chunk arrived (e.g.
|
||||
the model was slow to start responding), the stream is closed immediately
|
||||
so the blocked ``for chunk in stream`` iteration is unblocked.
|
||||
"""
|
||||
|
||||
__slots__ = ("_session",)
|
||||
|
||||
def __init__(self, session: ChatSession) -> None:
|
||||
super().__init__()
|
||||
self._session = session
|
||||
|
||||
def append(self, stream: Any) -> None:
|
||||
super().append(stream)
|
||||
self._session._cancel_stream = stream
|
||||
# If cancel was requested before the first chunk arrived (the worker
|
||||
# thread is blocked inside the provider generator waiting for the HTTP
|
||||
# response), close the stream immediately to unblock it.
|
||||
if self._session._cancel_event.is_set():
|
||||
with contextlib.suppress(Exception):
|
||||
stream.close()
|
||||
|
||||
|
||||
# Image extensions handled as vision content (SVG excluded — it's XML text)
|
||||
_IMAGE_EXTENSIONS: frozenset[str] = frozenset(
|
||||
{".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp", ".tiff", ".tif", ".ico"}
|
||||
@@ -318,6 +347,11 @@ class ChatSession:
|
||||
self._recent_tool_sigs: set[str] = set()
|
||||
# Cooperative cancellation: set from outside to stop generation
|
||||
self._cancel_event = threading.Event()
|
||||
self._cancel_ref: _CancelRef = _CancelRef(self) # provider appends SDK stream here
|
||||
self._cancel_stream: Any = None # closeable SDK stream handle
|
||||
self._generation: int = 0 # monotonic counter; orphaned threads skip cleanup
|
||||
self._active_procs: set[subprocess.Popen[str]] = set() # for force-kill
|
||||
self._procs_lock = threading.Lock()
|
||||
self._cancelled_partial_msg: dict[str, Any] | None = None
|
||||
# Intent validation judge (lazy-initialized)
|
||||
self._judge_config: JudgeConfig | None = judge_config
|
||||
@@ -1173,6 +1207,8 @@ class ChatSession:
|
||||
prov = provider or self._provider
|
||||
last_err: Exception | None = None
|
||||
for attempt in range(self._MAX_RETRIES + 1):
|
||||
self._check_cancelled()
|
||||
self._cancel_ref.clear() # discard stale handle from prior attempt
|
||||
try:
|
||||
return prov.create_streaming(
|
||||
client=client,
|
||||
@@ -1184,6 +1220,7 @@ class ChatSession:
|
||||
reasoning_effort=self.reasoning_effort,
|
||||
extra_params=self._provider_extra_params(provider=prov),
|
||||
deferred_names=self._get_deferred_names(),
|
||||
cancel_ref=self._cancel_ref,
|
||||
)
|
||||
except Exception as e:
|
||||
ename = type(e).__name__
|
||||
@@ -1205,11 +1242,36 @@ class ChatSession:
|
||||
while the worker thread is inside ``send()``.
|
||||
"""
|
||||
self._cancel_event.set()
|
||||
# Close the underlying SDK stream to unblock the iteration
|
||||
# immediately. Without this the worker thread stays blocked in
|
||||
# ``for chunk in stream`` until the next SSE chunk arrives from
|
||||
# the LLM provider (can be seconds during extended thinking).
|
||||
s = self._cancel_stream
|
||||
if s is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
s.close()
|
||||
# Kill all tracked subprocesses (bash tool). This is the
|
||||
# last line of defense — ensures destructive commands are
|
||||
# stopped even if the worker thread is stuck.
|
||||
with self._procs_lock:
|
||||
procs = list(self._active_procs)
|
||||
for proc in procs:
|
||||
if proc.poll() is not None:
|
||||
continue # already exited
|
||||
try:
|
||||
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
|
||||
except (OSError, ProcessLookupError):
|
||||
with contextlib.suppress(OSError, ProcessLookupError):
|
||||
proc.kill()
|
||||
|
||||
def _check_cancelled(self) -> None:
|
||||
"""Raise ``GenerationCancelled`` if cancellation has been requested."""
|
||||
def _check_cancelled(self, my_generation: int = 0) -> None:
|
||||
"""Raise ``GenerationCancelled`` if cancellation has been requested
|
||||
or if this thread belongs to an orphaned generation (force cancel).
|
||||
"""
|
||||
if self._cancel_event.is_set():
|
||||
raise GenerationCancelled()
|
||||
if my_generation and my_generation != self._generation:
|
||||
raise GenerationCancelled()
|
||||
|
||||
# -- Main generation loop ------------------------------------------------
|
||||
|
||||
@@ -1234,7 +1296,12 @@ class ChatSession:
|
||||
self._budget_exhausted = False
|
||||
self._budget_warned = False
|
||||
self._notify_count = 0
|
||||
self._cancel_event.clear()
|
||||
self._generation += 1
|
||||
my_generation = self._generation
|
||||
# Fresh cancel event per generation. The old event object stays
|
||||
# set for any abandoned thread — _exec_bash captures a local
|
||||
# reference so subprocesses from old generations are still killed.
|
||||
self._cancel_event = threading.Event()
|
||||
self._cancelled_partial_msg = None
|
||||
self.messages.append({"role": "user", "content": user_input})
|
||||
self._msg_tokens.append(max(1, int(len(user_input) / self._chars_per_token)))
|
||||
@@ -1248,7 +1315,7 @@ class ChatSession:
|
||||
|
||||
try:
|
||||
while True:
|
||||
self._check_cancelled()
|
||||
self._check_cancelled(my_generation)
|
||||
msgs = self._full_messages()
|
||||
|
||||
if self.debug:
|
||||
@@ -1258,10 +1325,19 @@ class ChatSession:
|
||||
self.ui.on_thinking_start()
|
||||
try:
|
||||
stream = self._create_stream_with_retry(msgs)
|
||||
assistant_msg = self._stream_response(stream)
|
||||
assistant_msg = self._stream_response(stream, my_generation)
|
||||
finally:
|
||||
# Only clear if this generation is still active —
|
||||
# an orphaned thread must not clobber a newer stream.
|
||||
if self._generation == my_generation:
|
||||
self._cancel_stream = None
|
||||
self._cancel_ref.clear()
|
||||
self.ui.on_thinking_stop()
|
||||
|
||||
# Bail if this generation was superseded (force cancel).
|
||||
if self._generation != my_generation:
|
||||
return
|
||||
|
||||
self._update_token_table(assistant_msg)
|
||||
self.messages.append(assistant_msg)
|
||||
self._msg_tokens.append(
|
||||
@@ -1314,6 +1390,8 @@ class ChatSession:
|
||||
f"\n[Auto-compacting: prompt exceeds {pct_display}% of context window]"
|
||||
)
|
||||
self._compact_messages(auto=True)
|
||||
# Update status bar with post-compaction token counts
|
||||
self._print_status_line()
|
||||
# Auto-title session after first exchange
|
||||
if not self._title_generated:
|
||||
self._title_generated = True
|
||||
@@ -1328,6 +1406,10 @@ class ChatSession:
|
||||
self._emit_state("running")
|
||||
results, user_feedback = self._execute_tools(tool_calls)
|
||||
|
||||
# Bail if generation was superseded during tool execution.
|
||||
if self._generation != my_generation:
|
||||
return
|
||||
|
||||
# Repeat detection: warn when a tool is called with identical args.
|
||||
# Skip error outputs — retrying a failed tool is valid.
|
||||
# Skip JSON outputs (MCP structured results) — appending
|
||||
@@ -1457,6 +1539,10 @@ class ChatSession:
|
||||
self.messages.append({"role": "user", "content": user_feedback})
|
||||
self._msg_tokens.append(max(1, int(len(user_feedback) / self._chars_per_token)))
|
||||
except GenerationCancelled:
|
||||
# If a newer send() has started (force cancel), this thread is
|
||||
# orphaned — skip all message mutations and state changes.
|
||||
if self._generation != my_generation:
|
||||
return
|
||||
# Cooperative cancellation — preserve partial content if available.
|
||||
if self._cancelled_partial_msg:
|
||||
# _stream_response was interrupted — save partial assistant msg
|
||||
@@ -1485,7 +1571,8 @@ class ChatSession:
|
||||
self.messages.pop()
|
||||
if self._msg_tokens:
|
||||
self._msg_tokens.pop()
|
||||
self._cancel_event.clear()
|
||||
# No need to clear _cancel_event — it's replaced per-generation
|
||||
# in send(), so this generation's event is simply discarded.
|
||||
self.ui.on_info("[Generation cancelled]")
|
||||
self._emit_state("idle")
|
||||
# Do NOT re-raise — return normally so server worker thread
|
||||
@@ -1530,7 +1617,9 @@ class ChatSession:
|
||||
_THINK_CLOSE_TAGS = ("</think>", "</reasoning>")
|
||||
_MAX_TAG_LEN = max(len(t) for t in _THINK_OPEN_TAGS + _THINK_CLOSE_TAGS)
|
||||
|
||||
def _stream_response(self, stream: Iterator[StreamChunk]) -> dict[str, Any]:
|
||||
def _stream_response(
|
||||
self, stream: Iterator[StreamChunk], my_generation: int = 0
|
||||
) -> dict[str, Any]:
|
||||
"""Stream response, dispatching tokens to the UI as they arrive.
|
||||
|
||||
Handles two reasoning delivery mechanisms:
|
||||
@@ -1621,7 +1710,13 @@ class ChatSession:
|
||||
finish_reason = None
|
||||
try:
|
||||
for chunk in stream:
|
||||
self._check_cancelled()
|
||||
# _cancel_stream is set eagerly by _CancelRef.append() when the
|
||||
# provider creates the SDK stream handle (before the first chunk
|
||||
# is returned). This fallback handles providers that use a
|
||||
# plain list for cancel_ref (e.g. some test fakes).
|
||||
if self._cancel_ref and self._cancel_stream is None:
|
||||
self._cancel_stream = self._cancel_ref[0]
|
||||
self._check_cancelled(my_generation)
|
||||
# Track finish_reason (e.g. "stop", "length", "tool_calls")
|
||||
if chunk.finish_reason:
|
||||
finish_reason = chunk.finish_reason
|
||||
@@ -1735,6 +1830,22 @@ class ChatSession:
|
||||
partial["_provider_content"] = provider_blocks
|
||||
self._cancelled_partial_msg = partial
|
||||
raise
|
||||
except Exception:
|
||||
# cancel() closed the underlying SDK stream, aborting the HTTP
|
||||
# connection. The blocked next() call on the iterator raises a
|
||||
# transport-level error (httpx, httpcore, etc.). Convert to
|
||||
# GenerationCancelled if a cancel was requested.
|
||||
if self._cancel_event.is_set():
|
||||
if pending:
|
||||
_flush_text(pending, in_think)
|
||||
self.ui.on_stream_end()
|
||||
partial = {"role": "assistant"}
|
||||
partial["content"] = "".join(content_parts) or ""
|
||||
if provider_blocks:
|
||||
partial["_provider_content"] = provider_blocks
|
||||
self._cancelled_partial_msg = partial
|
||||
raise GenerationCancelled() from None
|
||||
raise
|
||||
|
||||
# Flush any remaining buffered text
|
||||
if pending:
|
||||
@@ -2089,6 +2200,14 @@ class ChatSession:
|
||||
self._msg_tokens = [su_tok, sa_tok]
|
||||
after_tokens = self._system_tokens + sum(self._msg_tokens)
|
||||
|
||||
# Update usage estimate so the status bar reflects post-compaction state
|
||||
if self._last_usage:
|
||||
self._last_usage = {
|
||||
**self._last_usage,
|
||||
"prompt_tokens": after_tokens,
|
||||
"total_tokens": after_tokens,
|
||||
}
|
||||
|
||||
self.ui.on_info(f"[compacted: ~{before_tokens:,} -> ~{after_tokens:,} tokens]")
|
||||
separator = "\u2500" * 60
|
||||
lines = [separator]
|
||||
@@ -2291,6 +2410,7 @@ class ChatSession:
|
||||
def run_one(
|
||||
item: dict[str, Any],
|
||||
) -> tuple[str, str | list[dict[str, Any]]]:
|
||||
self._check_cancelled()
|
||||
if item.get("error"):
|
||||
self.ui.on_tool_result(
|
||||
item["call_id"], item.get("func_name", "unknown"), item["error"]
|
||||
@@ -3510,6 +3630,7 @@ class ChatSession:
|
||||
|
||||
def _exec_mcp_tool(self, item: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Execute an MCP tool call via the MCPClientManager."""
|
||||
self._check_cancelled()
|
||||
call_id: str = item["call_id"]
|
||||
func_name: str = item["mcp_func_name"]
|
||||
args: dict[str, Any] = item["mcp_args"]
|
||||
@@ -3583,6 +3704,7 @@ class ChatSession:
|
||||
|
||||
def _exec_read_resource(self, item: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Read an MCP resource by URI."""
|
||||
self._check_cancelled()
|
||||
call_id: str = item["call_id"]
|
||||
uri: str = item["resource_uri"]
|
||||
|
||||
@@ -3660,6 +3782,7 @@ class ChatSession:
|
||||
|
||||
def _exec_use_prompt(self, item: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Invoke an MCP prompt and return expanded messages."""
|
||||
self._check_cancelled()
|
||||
call_id: str = item["call_id"]
|
||||
name: str = item["prompt_name"]
|
||||
arguments: dict[str, str] = item["prompt_arguments"]
|
||||
@@ -3686,6 +3809,10 @@ class ChatSession:
|
||||
|
||||
def _exec_bash(self, item: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Execute a bash command via temp script, streaming stdout."""
|
||||
self._check_cancelled()
|
||||
# Capture cancel event locally so force-cancel (which replaces
|
||||
# _cancel_event with a fresh instance) doesn't disarm this check.
|
||||
cancel = self._cancel_event
|
||||
call_id, command = item["call_id"], item["command"]
|
||||
try:
|
||||
with tempfile.NamedTemporaryFile(mode="w", suffix=".sh", delete=False) as f:
|
||||
@@ -3702,6 +3829,8 @@ class ChatSession:
|
||||
start_new_session=True,
|
||||
env=scrubbed_env(),
|
||||
)
|
||||
with self._procs_lock:
|
||||
self._active_procs.add(proc)
|
||||
# Drain stderr in background thread to avoid pipe deadlock
|
||||
stderr_lines: list[str] = []
|
||||
|
||||
@@ -3739,7 +3868,7 @@ class ChatSession:
|
||||
except Exception:
|
||||
log.debug("UI callback error during tool output", exc_info=True)
|
||||
# Check cancellation during long-running commands
|
||||
if self._cancel_event.is_set():
|
||||
if cancel.is_set():
|
||||
with contextlib.suppress(OSError, ProcessLookupError):
|
||||
try:
|
||||
os.killpg(os.getpgid(proc.pid), signal.SIGKILL)
|
||||
@@ -3756,6 +3885,8 @@ class ChatSession:
|
||||
log.warning("Process did not exit after SIGKILL, pid=%d", proc.pid)
|
||||
stderr_thread.join(timeout=5)
|
||||
finally:
|
||||
with self._procs_lock:
|
||||
self._active_procs.discard(proc)
|
||||
os.unlink(script_path)
|
||||
|
||||
if timed_out.is_set():
|
||||
@@ -4615,6 +4746,7 @@ class ChatSession:
|
||||
|
||||
def _exec_notify(self, item: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Send a notification directly to the channel gateway via HTTP."""
|
||||
self._check_cancelled()
|
||||
call_id = item["call_id"]
|
||||
|
||||
if self._notify_count >= 5:
|
||||
@@ -4999,6 +5131,7 @@ class ChatSession:
|
||||
|
||||
def _exec_write_file(self, item: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Write content to a file, creating parent directories as needed."""
|
||||
self._check_cancelled()
|
||||
call_id = item["call_id"]
|
||||
path, content, resolved = item["path"], item["content"], item["resolved"]
|
||||
try:
|
||||
@@ -5020,6 +5153,7 @@ class ChatSession:
|
||||
When near_line is set, picks the occurrence nearest that line
|
||||
instead of requiring uniqueness.
|
||||
"""
|
||||
self._check_cancelled()
|
||||
call_id = item["call_id"]
|
||||
path, old_string, new_string = (
|
||||
item["path"],
|
||||
@@ -5071,6 +5205,7 @@ class ChatSession:
|
||||
|
||||
def _exec_man(self, item: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Look up a man or info page."""
|
||||
self._check_cancelled()
|
||||
call_id = item["call_id"]
|
||||
page = item["page"]
|
||||
section = item.get("section", "")
|
||||
@@ -5128,6 +5263,7 @@ class ChatSession:
|
||||
|
||||
def _exec_web_fetch(self, item: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Fetch a URL, then summarize/extract using an API call."""
|
||||
self._check_cancelled()
|
||||
call_id, url = item["call_id"], item["url"]
|
||||
question = item.get("question", "Summarize the key content of this page.")
|
||||
|
||||
@@ -5215,6 +5351,7 @@ class ChatSession:
|
||||
|
||||
def _exec_web_search(self, item: dict[str, Any]) -> tuple[str, str]:
|
||||
"""Search the web via the configured backend (Tavily, DDG, or MCP)."""
|
||||
self._check_cancelled()
|
||||
call_id = item["call_id"]
|
||||
query = item["query"]
|
||||
max_results = item.get("max_results", 5)
|
||||
|
||||
@@ -162,11 +162,14 @@ class AsyncTurnstoneServer(_BaseClient):
|
||||
response_model=StatusResponse,
|
||||
)
|
||||
|
||||
async def cancel(self, ws_id: str) -> StatusResponse:
|
||||
async def cancel(self, ws_id: str, *, force: bool = False) -> StatusResponse:
|
||||
body: dict[str, object] = {"ws_id": ws_id}
|
||||
if force:
|
||||
body["force"] = True
|
||||
return await self._request(
|
||||
"POST",
|
||||
"/v1/api/cancel",
|
||||
json_body={"ws_id": ws_id},
|
||||
json_body=body,
|
||||
response_model=StatusResponse,
|
||||
)
|
||||
|
||||
@@ -473,8 +476,8 @@ class TurnstoneServer:
|
||||
def command(self, *, ws_id: str, command: str) -> StatusResponse:
|
||||
return self._runner.run(self._async.command(ws_id=ws_id, command=command))
|
||||
|
||||
def cancel(self, ws_id: str) -> StatusResponse:
|
||||
return self._runner.run(self._async.cancel(ws_id))
|
||||
def cancel(self, ws_id: str, *, force: bool = False) -> StatusResponse:
|
||||
return self._runner.run(self._async.cancel(ws_id, force=force))
|
||||
|
||||
# -- streaming -----------------------------------------------------------
|
||||
|
||||
|
||||
+33
-7
@@ -1142,6 +1142,15 @@ async def send_message(request: Request) -> JSONResponse:
|
||||
return JSONResponse({"error": "Unknown workstream"}, status_code=404)
|
||||
# Atomically check-and-start to prevent two concurrent workers on the
|
||||
# same session (ChatSession.send() is not thread-safe).
|
||||
# If cancel was requested, poll briefly for the worker to exit before
|
||||
# rejecting. Snapshot the thread ref since force-cancel can set it to
|
||||
# None concurrently. Uses async sleep to avoid blocking the event loop.
|
||||
worker = ws.worker_thread
|
||||
if worker and worker.is_alive() and ws.session and ws.session._cancel_event.is_set():
|
||||
for _ in range(30): # up to 3s in 100ms steps
|
||||
await asyncio.sleep(0.1)
|
||||
if not worker.is_alive():
|
||||
break
|
||||
with ws._lock:
|
||||
if ws.worker_thread and ws.worker_thread.is_alive():
|
||||
ui._enqueue(
|
||||
@@ -1156,16 +1165,21 @@ async def send_message(request: Request) -> JSONResponse:
|
||||
|
||||
def run() -> None:
|
||||
assert ui is not None
|
||||
me = threading.current_thread()
|
||||
try:
|
||||
session.send(message)
|
||||
except GenerationCancelled:
|
||||
# Safety net — send() normally handles this internally.
|
||||
ui._enqueue({"type": "stream_end"})
|
||||
ui.on_state_change("idle")
|
||||
# If this thread was force-abandoned, ws.worker_thread will
|
||||
# have been set to None — don't emit spurious events.
|
||||
if ws.worker_thread is me:
|
||||
ui._enqueue({"type": "stream_end"})
|
||||
ui.on_state_change("idle")
|
||||
except Exception as e:
|
||||
ui.on_error(f"Error: {e}")
|
||||
ui._enqueue({"type": "stream_end"})
|
||||
ui.on_state_change("error")
|
||||
if ws.worker_thread is me:
|
||||
ui.on_error(f"Error: {e}")
|
||||
ui._enqueue({"type": "stream_end"})
|
||||
ui.on_state_change("error")
|
||||
|
||||
t = threading.Thread(target=run, daemon=True)
|
||||
ws.worker_thread = t
|
||||
@@ -1237,6 +1251,7 @@ async def cancel_generation(request: Request) -> JSONResponse:
|
||||
session = ws.session
|
||||
if session is None:
|
||||
return JSONResponse({"error": "No session"}, status_code=400)
|
||||
force = body.get("force", False) is True
|
||||
# Only act if generation is actually in progress
|
||||
if ws.worker_thread and ws.worker_thread.is_alive():
|
||||
# Set the cooperative cancel flag (worker thread checks at checkpoints)
|
||||
@@ -1244,8 +1259,19 @@ async def cancel_generation(request: Request) -> JSONResponse:
|
||||
# Unblock any pending approval/plan review waits
|
||||
ui.resolve_approval(False, "Cancelled by user")
|
||||
ui.resolve_plan("reject")
|
||||
# Emit cancelled SSE event so SDK consumers get a typed signal
|
||||
ui._enqueue({"type": "cancelled"})
|
||||
if force:
|
||||
# Force cancel: abandon the stuck worker thread (daemon, will
|
||||
# die on process exit or stream timeout) and emit stream_end
|
||||
# so the UI and session recover immediately. The per-generation
|
||||
# cancel event stays set so the abandoned thread still kills
|
||||
# subprocesses at its next checkpoint.
|
||||
with ws._lock:
|
||||
ws.worker_thread = None
|
||||
ui._enqueue({"type": "stream_end"})
|
||||
ui.on_state_change("idle")
|
||||
else:
|
||||
# Emit cancelled SSE event so SDK consumers get a typed signal
|
||||
ui._enqueue({"type": "cancelled"})
|
||||
return JSONResponse({"status": "ok"})
|
||||
|
||||
|
||||
|
||||
+85
-10
@@ -30,6 +30,8 @@ function Pane(wsId) {
|
||||
this.model = "";
|
||||
this.modelAlias = "";
|
||||
this.statusText = "";
|
||||
this._cancelTimeout = null;
|
||||
this._forceTimeout = null;
|
||||
this._createDOM();
|
||||
}
|
||||
|
||||
@@ -193,6 +195,14 @@ Pane.prototype.updateWsName = function () {
|
||||
};
|
||||
|
||||
Pane.prototype.disconnectSSE = function () {
|
||||
if (this._cancelTimeout) {
|
||||
clearTimeout(this._cancelTimeout);
|
||||
this._cancelTimeout = null;
|
||||
}
|
||||
if (this._forceTimeout) {
|
||||
clearTimeout(this._forceTimeout);
|
||||
this._forceTimeout = null;
|
||||
}
|
||||
if (this.evtSource) {
|
||||
this.evtSource.close();
|
||||
this.evtSource = null;
|
||||
@@ -205,6 +215,9 @@ Pane.prototype.setBusy = function (b) {
|
||||
this.sendBtn.style.display = b ? "none" : "";
|
||||
this.stopBtn.style.display = b ? "" : "none";
|
||||
this.stopBtn.disabled = !b;
|
||||
this.stopBtn.textContent = "\u25a0 Stop";
|
||||
this.stopBtn.setAttribute("aria-label", "Stop generation");
|
||||
delete this.stopBtn.dataset.forceCancel;
|
||||
};
|
||||
|
||||
Pane.prototype.showEmptyState = function () {
|
||||
@@ -357,8 +370,18 @@ Pane.prototype.handleEvent = function (evt) {
|
||||
break;
|
||||
|
||||
case "stream_end":
|
||||
if (this._cancelTimeout) {
|
||||
clearTimeout(this._cancelTimeout);
|
||||
this._cancelTimeout = null;
|
||||
}
|
||||
if (this._forceTimeout) {
|
||||
clearTimeout(this._forceTimeout);
|
||||
this._forceTimeout = null;
|
||||
}
|
||||
// Render final markdown for the assistant message (existing code).
|
||||
// Note: renderMarkdown is the project's sanitizing markdown renderer.
|
||||
if (this.currentAssistantEl && this.contentBuffer) {
|
||||
this.currentAssistantEl.innerHTML = renderMarkdown(this.contentBuffer);
|
||||
this.currentAssistantEl.innerHTML = renderMarkdown(this.contentBuffer); // sanitized by renderMarkdown
|
||||
postRenderMarkdown(this.currentAssistantEl);
|
||||
}
|
||||
this.currentAssistantEl = null;
|
||||
@@ -420,17 +443,51 @@ Pane.prototype.handleEvent = function (evt) {
|
||||
break;
|
||||
|
||||
case "busy_error":
|
||||
// Server is still busy — don't transition to send mode.
|
||||
// Re-enable the stop button so the user can try cancelling.
|
||||
this.addErrorMessage(evt.message);
|
||||
this.setBusy(false);
|
||||
this.stopBtn.textContent = "\u25a0 Stop";
|
||||
this.stopBtn.setAttribute("aria-label", "Stop generation");
|
||||
delete this.stopBtn.dataset.forceCancel;
|
||||
this.stopBtn.disabled = false;
|
||||
break;
|
||||
|
||||
case "cancelled":
|
||||
// Cancel requested but worker thread may still be finishing.
|
||||
// Show "Cancelling..." state; stream_end will transition to ready.
|
||||
// If stream_end already arrived (busy is false), the cancel is
|
||||
// already handled — don't re-enter the cancelling state.
|
||||
if (!this.busy) break;
|
||||
// Clear any prior timeouts first (duplicate cancelled events).
|
||||
clearTimeout(this._cancelTimeout);
|
||||
clearTimeout(this._forceTimeout);
|
||||
this.currentAssistantEl = null;
|
||||
this.currentReasoningEl = null;
|
||||
this.contentBuffer = "";
|
||||
this.setBusy(false);
|
||||
this.inputEl.focus();
|
||||
this.stopBtn.disabled = true;
|
||||
this.stopBtn.textContent = "Cancelling\u2026";
|
||||
this.stopBtn.setAttribute("aria-label", "Cancelling generation");
|
||||
this.scrollToBottom(true);
|
||||
// After 2s, offer "Force Stop" for a harder cancel that abandons
|
||||
// the stuck worker thread. Safety timeout at 10s auto-recovers
|
||||
// if stream_end never arrives (connection drop).
|
||||
var self = this;
|
||||
this._cancelTimeout = setTimeout(function () {
|
||||
if (self.busy) {
|
||||
self.stopBtn.disabled = false;
|
||||
self.stopBtn.textContent = "\u26a0 Force Stop";
|
||||
self.stopBtn.setAttribute("aria-label", "Force stop generation");
|
||||
self.stopBtn.dataset.forceCancel = "true";
|
||||
}
|
||||
}, 2000);
|
||||
this._forceTimeout = setTimeout(function () {
|
||||
if (self.busy) {
|
||||
self.addInfoMessage(
|
||||
"Cancel didn\u2019t complete in time. You may need to resend your last message.",
|
||||
);
|
||||
self.setBusy(false);
|
||||
}
|
||||
}, 10000);
|
||||
break;
|
||||
|
||||
case "connected":
|
||||
@@ -1080,17 +1137,35 @@ Pane.prototype.sendMessage = function () {
|
||||
};
|
||||
|
||||
Pane.prototype.cancelGeneration = function () {
|
||||
if (!this.busy || !this.wsId) return;
|
||||
if (!this.busy || !this.wsId || this.stopBtn.disabled) return;
|
||||
var self = this;
|
||||
var isForce = this.stopBtn.dataset.forceCancel === "true";
|
||||
this.stopBtn.disabled = true;
|
||||
authFetch("/v1/api/cancel", {
|
||||
method: "POST",
|
||||
headers: { "Content-Type": "application/json" },
|
||||
body: JSON.stringify({ ws_id: this.wsId }),
|
||||
}).catch(function (err) {
|
||||
self.addErrorMessage("Cancel error: " + err.message);
|
||||
self.stopBtn.disabled = false;
|
||||
});
|
||||
body: JSON.stringify({ ws_id: this.wsId, force: isForce }),
|
||||
})
|
||||
.then(function () {
|
||||
if (isForce) {
|
||||
// Force cancel abandons the worker — transition immediately.
|
||||
// Clear timeouts to prevent stale timers firing on next send.
|
||||
if (self._cancelTimeout) {
|
||||
clearTimeout(self._cancelTimeout);
|
||||
self._cancelTimeout = null;
|
||||
}
|
||||
if (self._forceTimeout) {
|
||||
clearTimeout(self._forceTimeout);
|
||||
self._forceTimeout = null;
|
||||
}
|
||||
self.addInfoMessage("Force stopped. Previous generation abandoned.");
|
||||
self.setBusy(false);
|
||||
}
|
||||
})
|
||||
.catch(function (err) {
|
||||
self.addErrorMessage("Cancel error: " + err.message);
|
||||
self.stopBtn.disabled = false;
|
||||
});
|
||||
};
|
||||
|
||||
Pane.prototype._autoResize = function () {
|
||||
|
||||
@@ -667,7 +667,7 @@ body { position: static; }
|
||||
}
|
||||
.pane-input-area button:hover { filter: brightness(1.1); }
|
||||
.pane-input-area button:disabled { opacity: 0.35; cursor: not-allowed; filter: none; }
|
||||
.pane-stop { background: var(--red, #c94040); }
|
||||
.pane-stop { background: var(--red, #c94040); min-width: 120px; text-align: center; white-space: nowrap; }
|
||||
.pane-stop:focus-visible { outline: 2px solid var(--fg-bright, #e8ecf4); outline-offset: 2px; }
|
||||
[data-theme="light"] .pane-stop { color: #fff; }
|
||||
|
||||
|
||||
Reference in New Issue
Block a user