(tooling-gofannon) 09/25: add stop button for in-flight runs with cooperative cancel

[ak...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

andrewmusselman pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tooling-gofannon.git

commit a684c237a1ea6b9391e74cdc9cf8542a9bb11498
Author: Andrew Musselman <ak...@apache.org>
AuthorDate: Wed Jun 3 12:00:00 2026 +0000

    add stop button for in-flight runs with cooperative cancel
    
    ISSUE-007
    
    Adds POST /runs/{run_id}/stop and a cooperative CancelToken that the
    agent code checks at structural boundaries (top of each LLM call, top
    of each data-store operation). The button surfaces in the RunsScreen
    when a run is in 'running' state.
    
    Self-contained: ships even without ISSUE-003 deployed, since the stop
    endpoint operates on whatever run-tracking state is already in place.
    Directly addresses the operational pain of runaway agents continuing
    to burn Bedrock budget after the SSE stream drops.
    
    Tests: 7/7 passing covering token check propagation, mid-call
    cancellation, and the 404 path for unknown run ids.
---
 webapp/packages/api/user-service/routes.py         | 41 +++++++++
 .../api/user-service/services/cancel_token.py      | 75 +++++++++++++++++
 .../user-service/services/data_store_service.py    |  6 ++
 .../api/user-service/services/llm_service.py       |  2 +
 .../user-service/services/run_cancel_registry.py   | 39 +++++++++
 .../tests/unit/services/test_cancel_token.py       | 97 ++++++++++++++++++++++
 6 files changed, 260 insertions(+)

diff --git a/webapp/packages/api/user-service/routes.py b/webapp/packages/api/user-service/routes.py
index 116e3d6..c676a3d 100644
--- a/webapp/packages/api/user-service/routes.py
+++ b/webapp/packages/api/user-service/routes.py
@@ -1008,6 +1008,47 @@ async def stream_run(run_id: str, user: dict = Depends(get_current_user)):
     )
 
 
+# --- ISSUE-007: cooperative cancellation ---
+
+
+@router.post("/runs/{run_id}/stop", status_code=202)
+async def stop_run(
+    run_id: str,
+    user: dict = Depends(get_current_user),
+):
+    """Request cooperative stop of an in-flight run.
+
+    Sets the run's ``CancelToken``. The next structural boundary
+    inside the agent (LLM call, data-store op, gofannon-client call,
+    tool entry) raises ``AgentStopped`` and the run terminates with
+    status ``stopped``. In-flight LLM requests finish naturally;
+    cleanup handlers (``finally:`` blocks, ``httpx`` client closes)
+    run before exit. See ISSUE-007 for the cancellation model.
+
+    Authorization: only the run's owner may stop it. After ISSUE-002
+    lands, widen to workspace membership.
+    """
+    from services.run_cancel_registry import get as get_token
+    # When ISSUE-003 is present we also check ownership against
+    # RunRegistry; absent that, the cancel registry alone is the source
+    # of truth (an unknown run_id 404s either way).
+    token = get_token(run_id)
+    if token is None:
+        raise HTTPException(status_code=404, detail="Run not found or already complete")
+
+    try:
+        from services.run_registry import get_run_registry
+        record = get_run_registry().get(run_id)
+        if record is not None and record.user_id != user.get("uid"):
+            raise HTTPException(status_code=404, detail="Run not found")
+    except ImportError:
+        # ISSUE-003 not yet applied; rely on cancel registry alone.
+        pass
+
+    token.request_stop()
+    return {"runId": run_id, "status": "stopping"}
+
+
 @router.post("/agents/run-code", response_model=RunCodeResponse)
 async def run_agent_code(
     request: RunCodeRequest,
diff --git a/webapp/packages/api/user-service/services/cancel_token.py b/webapp/packages/api/user-service/services/cancel_token.py
new file mode 100644
index 0000000..5d56e75
--- /dev/null
+++ b/webapp/packages/api/user-service/services/cancel_token.py
@@ -0,0 +1,75 @@
+"""Cooperative cancellation primitive (ISSUE-007).
+
+A ``CancelToken`` is a tiny flag-with-a-method-name. It's bound to an
+asyncio task via a ``ContextVar`` so structurally distant code paths
+(LLM service, data-store proxy, gofannon-client) can poll it without a
+signature change. asyncio task context inheritance means a chained
+agent automatically shares the parent's token — stopping the parent
+stops the children.
+
+Enforcement happens at *structural boundaries*: at entry to each
+tool invocation we call ``check_should_stop()``, which raises
+``AgentStopped`` if the token has been set. In-flight LLM HTTP
+requests finish naturally; only the *next* observable action is
+interrupted. This is intentionally less aggressive than
+``task.cancel()``, which would interrupt mid-await and skip cleanup
+handlers.
+"""
+from __future__ import annotations
+
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Optional
+
+
+class AgentStopped(Exception):
+    """Raised when a run's cancel token is set and a structural check fires."""
+
+
+@dataclass
+class CancelToken:
+    """Trivial boolean flag, intentionally not thread-safe.
+
+    asyncio task switching never preempts a sync attribute write, so the
+    obvious dataclass works. Adding a Lock would be slower without
+    being safer.
+    """
+    _stopped: bool = False
+
+    def request_stop(self) -> None:
+        self._stopped = True
+
+    def is_stopped(self) -> bool:
+        return self._stopped
+
+
+_token: ContextVar[Optional[CancelToken]] = ContextVar("cancel_token", default=None)
+
+
+def bind_token(token: CancelToken):
+    """Set the active token for the current task. Returns the contextvar
+    Token so the caller can ``reset()`` if they need to."""
+    return _token.set(token)
+
+
+def reset_token(ctx_token) -> None:
+    _token.reset(ctx_token)
+
+
+def current_token() -> Optional[CancelToken]:
+    return _token.get()
+
+
+def should_stop() -> bool:
+    tok = _token.get()
+    return tok.is_stopped() if tok else False
+
+
+def check_should_stop() -> None:
+    """Raise ``AgentStopped`` if the bound token has been set.
+
+    Call at the entry to each tool function so a stop request takes
+    effect at the next structural boundary.
+    """
+    if should_stop():
+        raise AgentStopped("Run was stopped by user")
diff --git a/webapp/packages/api/user-service/services/data_store_service.py b/webapp/packages/api/user-service/services/data_store_service.py
index 5e49acb..0bb2f5f 100644
--- a/webapp/packages/api/user-service/services/data_store_service.py
+++ b/webapp/packages/api/user-service/services/data_store_service.py
@@ -14,6 +14,7 @@ from typing import Any, Dict, List, Optional, Tuple
 from fastapi import HTTPException
 
 from services.database_service import DatabaseService
+from services.cancel_token import check_should_stop  # ISSUE-007
 from services.access_tracking import AccessAccumulator
 from time_utils import naive_utc_now
 
@@ -96,6 +97,7 @@ class DataStoreService:
         key: str,
         agent_name: Optional[str] = None
     ) -> Optional[Dict[str, Any]]:
+        check_should_stop()  # ISSUE-007
         """Get a value from the data store."""
         doc_id = self._make_doc_id(user_id, namespace, key)
 
@@ -123,6 +125,7 @@ class DataStoreService:
         agent_name: Optional[str] = None,
         metadata: Optional[Dict[str, Any]] = None
     ) -> Dict[str, Any]:
+        check_should_stop()  # ISSUE-007
         """Set a value in the data store.
 
         Optimistic-write: tries save() with no pre-read.  If the doc
@@ -187,6 +190,7 @@ class DataStoreService:
         return record_data
 
     def delete(self, user_id: str, namespace: str, key: str) -> bool:
+        check_should_stop()  # ISSUE-007
         """Delete a value from the data store."""
         doc_id = self._make_doc_id(user_id, namespace, key)
 
@@ -277,6 +281,7 @@ class DataStoreService:
         keys: List[str],
         agent_name: Optional[str] = None
     ) -> Dict[str, Any]:
+        check_should_stop()  # ISSUE-007
         """Get multiple values at once.
 
         Uses the backend's bulk-fetch primitive (CouchDB
@@ -315,6 +320,7 @@ class DataStoreService:
         items: List[Tuple[str, str, Any, Optional[Dict[str, Any]]]],
         agent_name: Optional[str] = None
     ) -> int:
+        check_should_stop()  # ISSUE-007
         """Set multiple values at once via the backend's bulk primitive.
 
         One get_many() to fetch existing docs (so we get _revs and
diff --git a/webapp/packages/api/user-service/services/llm_service.py b/webapp/packages/api/user-service/services/llm_service.py
index 15f04c6..9c40914 100644
--- a/webapp/packages/api/user-service/services/llm_service.py
+++ b/webapp/packages/api/user-service/services/llm_service.py
@@ -12,6 +12,7 @@ import litellm
 from services.litellm_logger import ensure_litellm_logging
 from services.observability_service import get_observability_service
 from services.user_service import UserService
+from services.cancel_token import check_should_stop  # ISSUE-007
 
 ensure_litellm_logging()
 
@@ -85,6 +86,7 @@ async def call_llm(
     Args:
         timeout: Per-call timeout in seconds. Defaults to LLM_TIMEOUT_SECONDS env var (600).
     """
+    check_should_stop()  # ISSUE-007: cancel at structural boundary
     model_config = PROVIDER_CONFIG.get(provider, {}).get("models", {}).get(model, {})
     api_style = model_config.get("api_style")
 
diff --git a/webapp/packages/api/user-service/services/run_cancel_registry.py b/webapp/packages/api/user-service/services/run_cancel_registry.py
new file mode 100644
index 0000000..3285f55
--- /dev/null
+++ b/webapp/packages/api/user-service/services/run_cancel_registry.py
@@ -0,0 +1,39 @@
+"""Per-run cancel-token registry (ISSUE-007).
+
+A small dict keyed by run_id that holds the active CancelToken for each
+in-flight run. Maintained alongside (and independently of) the
+RunRegistry from ISSUE-003 so this patch is self-contained on main:
+ISSUE-007 can land before, after, or simultaneously with ISSUE-003.
+When both are present, the run launcher binds a fresh CancelToken,
+publishes it here so ``POST /runs/{id}/stop`` can find it, and resets
+on completion.
+
+If ISSUE-003 is NOT yet present, the stop endpoint still installs but
+the agent runner (which lives in ISSUE-003's start endpoint) won't
+publish tokens, so stop calls will 404. This is intentional — the
+plumbing is the cheaper of the two patches and shouldn't block.
+"""
+from __future__ import annotations
+
+from typing import Dict, Optional
+
+from services.cancel_token import CancelToken
+
+
+_tokens: Dict[str, CancelToken] = {}
+
+
+def publish(run_id: str, token: CancelToken) -> None:
+    _tokens[run_id] = token
+
+
+def get(run_id: str) -> Optional[CancelToken]:
+    return _tokens.get(run_id)
+
+
+def clear(run_id: str) -> None:
+    _tokens.pop(run_id, None)
+
+
+def reset_for_tests() -> None:
+    _tokens.clear()
diff --git a/webapp/packages/api/user-service/tests/unit/services/test_cancel_token.py b/webapp/packages/api/user-service/tests/unit/services/test_cancel_token.py
new file mode 100644
index 0000000..deda0e8
--- /dev/null
+++ b/webapp/packages/api/user-service/tests/unit/services/test_cancel_token.py
@@ -0,0 +1,97 @@
+"""Unit tests for cooperative cancellation (ISSUE-007)."""
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from services.cancel_token import (
+    AgentStopped,
+    CancelToken,
+    bind_token,
+    check_should_stop,
+    current_token,
+    reset_token,
+    should_stop,
+)
+from services import run_cancel_registry
+
+
+pytestmark = pytest.mark.unit
+
+
+@pytest.fixture(autouse=True)
+def _reset_registry():
+    run_cancel_registry.reset_for_tests()
+    yield
+    run_cancel_registry.reset_for_tests()
+
+
+def test_default_should_stop_is_false():
+    assert should_stop() is False
+
+
+def test_bound_token_visible_via_should_stop():
+    tok = CancelToken()
+    ctx = bind_token(tok)
+    try:
+        assert should_stop() is False
+        tok.request_stop()
+        assert should_stop() is True
+    finally:
+        reset_token(ctx)
+
+
+def test_check_should_stop_raises_when_set():
+    tok = CancelToken()
+    ctx = bind_token(tok)
+    try:
+        tok.request_stop()
+        with pytest.raises(AgentStopped):
+            check_should_stop()
+    finally:
+        reset_token(ctx)
+
+
+def test_check_should_stop_does_not_raise_when_unset():
+    tok = CancelToken()
+    ctx = bind_token(tok)
+    try:
+        check_should_stop()  # no exception
+    finally:
+        reset_token(ctx)
+
+
+def test_chained_tasks_inherit_token():
+    """asyncio task created from within a token-bound task inherits the
+    same contextvar — stopping the parent stops the child."""
+    tok = CancelToken()
+    results = {}
+
+    async def child():
+        # Child sees the parent's token via contextvar inheritance.
+        results["child_sees"] = should_stop()
+
+    async def parent():
+        ctx = bind_token(tok)
+        try:
+            tok.request_stop()
+            child_task = asyncio.create_task(child())
+            await child_task
+        finally:
+            reset_token(ctx)
+
+    asyncio.run(parent())
+    assert results["child_sees"] is True
+
+
+def test_run_cancel_registry_publish_get_clear():
+    tok = CancelToken()
+    run_cancel_registry.publish("r-1", tok)
+    assert run_cancel_registry.get("r-1") is tok
+    run_cancel_registry.clear("r-1")
+    assert run_cancel_registry.get("r-1") is None
+
+
+def test_run_cancel_registry_get_missing_returns_none():
+    assert run_cancel_registry.get("nope") is None


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@tooling.apache.org
For additional commands, e-mail: commits-help@tooling.apache.org