diff --git a/docs/api-reference.md b/docs/api-reference.md index 500065bc..0bd47559 100644 --- a/docs/api-reference.md +++ b/docs/api-reference.md @@ -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:** diff --git a/docs/diagrams/09-workstream-states.puml b/docs/diagrams/09-workstream-states.puml index 6d3d2774..06399bf9 100644 --- a/docs/diagrams/09-workstream-states.puml +++ b/docs/diagrams/09-workstream-states.puml @@ -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) diff --git a/docs/diagrams/png/09-workstream-states.png b/docs/diagrams/png/09-workstream-states.png index 3e580534..8d70e702 100644 --- a/docs/diagrams/png/09-workstream-states.png +++ b/docs/diagrams/png/09-workstream-states.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:7896c6e041b6dbb89d034468fa980c8fe645df5eb969d45ef966ccc6399edac2 -size 200083 +oid sha256:7bf27afa267d5b8d6da38e83213ed1b8e87639d5105a0a1ccc2e5a4bf4d3b67e +size 185282 diff --git a/docs/sdk.md b/docs/sdk.md index f622e09a..4e0e8d81 100644 --- a/docs/sdk.md +++ b/docs/sdk.md @@ -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` | diff --git a/sdk/typescript/openapi-console.json b/sdk/typescript/openapi-console.json index f3408ec2..7ead9946 100644 --- a/sdk/typescript/openapi-console.json +++ b/sdk/typescript/openapi-console.json @@ -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": { diff --git a/sdk/typescript/openapi-server.json b/sdk/typescript/openapi-server.json index 77c7eca9..a430bd28 100644 --- a/sdk/typescript/openapi-server.json +++ b/sdk/typescript/openapi-server.json @@ -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": [ diff --git a/sdk/typescript/src/server.ts b/sdk/typescript/src/server.ts index 1a9e6ae3..b8ae41b0 100644 --- a/sdk/typescript/src/server.ts +++ b/sdk/typescript/src/server.ts @@ -93,10 +93,13 @@ export class TurnstoneServer extends BaseClient { }); } - async cancel(wsId: string): Promise { - return this.request("POST", "/v1/api/cancel", { - json: { ws_id: wsId }, - }); + async cancel( + wsId: string, + opts?: { force?: boolean }, + ): Promise { + const body: Record = { ws_id: wsId }; + if (opts?.force) body.force = true; + return this.request("POST", "/v1/api/cancel", { json: body }); } // -- Streaming ------------------------------------------------------------ diff --git a/tests/test_cancel.py b/tests/test_cancel.py index b484a757..54ff4013 100644 --- a/tests/test_cancel.py +++ b/tests/test_cancel.py @@ -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() diff --git a/turnstone/api/server_schemas.py b/turnstone/api/server_schemas.py index 2c200698..c3afbf3a 100644 --- a/turnstone/api/server_schemas.py +++ b/turnstone/api/server_schemas.py @@ -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): diff --git a/turnstone/core/providers/_anthropic.py b/turnstone/core/providers/_anthropic.py index fe8d2983..84cea5a4 100644 --- a/turnstone/core/providers/_anthropic.py +++ b/turnstone/core/providers/_anthropic.py @@ -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]: diff --git a/turnstone/core/providers/_openai.py b/turnstone/core/providers/_openai.py index f75f3266..aae797ef 100644 --- a/turnstone/core/providers/_openai.py +++ b/turnstone/core/providers/_openai.py @@ -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]: diff --git a/turnstone/core/providers/_protocol.py b/turnstone/core/providers/_protocol.py index c01ab4c5..e9129162 100644 --- a/turnstone/core/providers/_protocol.py +++ b/turnstone/core/providers/_protocol.py @@ -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( diff --git a/turnstone/core/session.py b/turnstone/core/session.py index d517141e..430b3693 100644 --- a/turnstone/core/session.py +++ b/turnstone/core/session.py @@ -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 = ("", "") _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) diff --git a/turnstone/sdk/server.py b/turnstone/sdk/server.py index 9d793177..053c21cc 100644 --- a/turnstone/sdk/server.py +++ b/turnstone/sdk/server.py @@ -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 ----------------------------------------------------------- diff --git a/turnstone/server.py b/turnstone/server.py index 839aee03..bf65cd63 100644 --- a/turnstone/server.py +++ b/turnstone/server.py @@ -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"}) diff --git a/turnstone/ui/static/app.js b/turnstone/ui/static/app.js index af9209fa..84607e35 100644 --- a/turnstone/ui/static/app.js +++ b/turnstone/ui/static/app.js @@ -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 () { diff --git a/turnstone/ui/static/style.css b/turnstone/ui/static/style.css index ee98a867..c3aa2215 100644 --- a/turnstone/ui/static/style.css +++ b/turnstone/ui/static/style.css @@ -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; }