feat: 支持 agent 运行取消功能

2026-03-25 18:33:25 +08:00
parent 599c597e69
commit 96fc4a1e77
21 changed files with 778 additions and 85 deletions
@@ -1,6 +1,7 @@
 from __future__ import annotations

-from typing import Any, Protocol
+import asyncio
+from typing import Any, Awaitable, Callable, Protocol

 from ag_ui.core.types import RunAgentInput
 from agentscope.message import Msg
@@ -28,6 +29,7 @@ class RunnerLike(Protocol):
        runtime_config: RuntimeConfig,
        user_memory: UserMemoryContent | None,
        work_memory: WorkProfileContent | None,
+        cancel_checker: Callable[[], Awaitable[bool]] | None = None,
    ) -> dict[str, Any]: ...


@@ -53,6 +55,7 @@ class AgentScopeRuntimeOrchestrator:
        runtime_config: RuntimeConfig,
        user_memory: UserMemoryContent | None = None,
        work_memory: WorkProfileContent | None = None,
+        cancel_checker: Callable[[], Awaitable[bool]] | None = None,
    ) -> dict[str, Any]:
        thread_id = run_input.thread_id
        run_id = run_input.run_id
@@ -74,6 +77,7 @@ class AgentScopeRuntimeOrchestrator:
                runtime_config=runtime_config,
                user_memory=user_memory,
                work_memory=work_memory,
+                cancel_checker=cancel_checker,
            )

            await self._pipeline.emit(
@@ -85,6 +89,23 @@ class AgentScopeRuntimeOrchestrator:
                },
            )
            return result if isinstance(result, dict) else {}
+        except asyncio.CancelledError:
+            logger.info(
+                "agentscope runtime execution canceled",
+                thread_id=thread_id,
+                run_id=run_id,
+            )
+            await self._pipeline.emit(
+                session_id=thread_id,
+                event={
+                    "type": "RUN_ERROR",
+                    "threadId": thread_id,
+                    "runId": run_id,
+                    "message": "run canceled by user",
+                    "code": "RUN_CANCELED",
+                },
+            )
+            raise
        except Exception:
            logger.exception(
                "agentscope runtime execution failed",
@@ -1,8 +1,10 @@
 from __future__ import annotations

+import asyncio
+import contextlib
 from dataclasses import dataclass
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Awaitable, Callable
 from uuid import UUID

 from ag_ui.core.types import RunAgentInput
@@ -64,6 +66,8 @@ class AgentScopeRunner:
    def __init__(self, *, litellm_service: LiteLLMService | None = None) -> None:
        patch_agentscope_json_repair_compat()
        self._litellm_service: LiteLLMService = litellm_service or LiteLLMService()
+        self._active_agent: JsonReActAgent | None = None
+        self._active_agent_lock = asyncio.Lock()

    async def execute(
        self,
@@ -75,51 +79,99 @@ class AgentScopeRunner:
        runtime_config: RuntimeConfig,
        user_memory: UserMemoryContent | None = None,
        work_memory: WorkProfileContent | None = None,
+        cancel_checker: Callable[[], Awaitable[bool]] | None = None,
    ) -> dict[str, Any]:
        owner_id = UUID(user_context.id)
        runtime_client_time = self._resolve_runtime_client_time(run_input=run_input)
        runtime_mode = self._resolve_runtime_mode(run_input=run_input)
+        stop_cancel_watch = asyncio.Event()
+        cancel_watch_task: asyncio.Task[None] | None = None
+        run_task = asyncio.current_task()

-        async with AsyncSessionLocal() as session:
-            router_config = await self._load_stage_config(
-                session=session,
-                agent_type=AgentType.ROUTER,
-            )
-            worker_config = await self._load_stage_config(
-                session=session,
-                agent_type=AgentType.WORKER,
-            )
-            worker_toolkit = self._build_toolkit(
-                session=session,
-                owner_id=owner_id,
-                enabled_tools=runtime_config.enabled_tools,
+        if cancel_checker is not None and run_task is not None:
+            cancel_watch_task = asyncio.create_task(
+                self._watch_cancel_signal(
+                    cancel_checker=cancel_checker,
+                    stop_signal=stop_cancel_watch,
+                    run_task=run_task,
+                )
            )

-            router_output = await self._execute_router_step(
-                pipeline=pipeline,
-                run_input=run_input,
-                user_context=user_context,
-                context_messages=context_messages,
-                stage_config=router_config,
-                runtime_client_time=runtime_client_time,
-                runtime_mode=runtime_mode,
-                user_memory=user_memory,
-            )
-            worker_output = await self._execute_worker_step(
-                pipeline=pipeline,
-                run_input=run_input,
-                user_context=user_context,
-                router_output=router_output,
-                toolkit=worker_toolkit,
-                stage_config=worker_config,
-                runtime_client_time=runtime_client_time,
-                runtime_mode=runtime_mode,
-                work_memory=work_memory,
-            )
-            return {
-                "router": router_output.model_dump(mode="json", exclude_none=True),
-                "worker": worker_output.model_dump(mode="json", exclude_none=True),
-            }
+        try:
+            async with AsyncSessionLocal() as session:
+                router_config = await self._load_stage_config(
+                    session=session,
+                    agent_type=AgentType.ROUTER,
+                )
+                worker_config = await self._load_stage_config(
+                    session=session,
+                    agent_type=AgentType.WORKER,
+                )
+                worker_toolkit = self._build_toolkit(
+                    session=session,
+                    owner_id=owner_id,
+                    enabled_tools=runtime_config.enabled_tools,
+                )
+
+                router_output = await self._execute_router_step(
+                    pipeline=pipeline,
+                    run_input=run_input,
+                    user_context=user_context,
+                    context_messages=context_messages,
+                    stage_config=router_config,
+                    runtime_client_time=runtime_client_time,
+                    runtime_mode=runtime_mode,
+                    user_memory=user_memory,
+                )
+                if cancel_checker is not None and await cancel_checker():
+                    raise asyncio.CancelledError("run canceled by user")
+                worker_output = await self._execute_worker_step(
+                    pipeline=pipeline,
+                    run_input=run_input,
+                    user_context=user_context,
+                    router_output=router_output,
+                    toolkit=worker_toolkit,
+                    stage_config=worker_config,
+                    runtime_client_time=runtime_client_time,
+                    runtime_mode=runtime_mode,
+                    work_memory=work_memory,
+                )
+                return {
+                    "router": router_output.model_dump(mode="json", exclude_none=True),
+                    "worker": worker_output.model_dump(mode="json", exclude_none=True),
+                }
+        finally:
+            stop_cancel_watch.set()
+            if cancel_watch_task is not None:
+                cancel_watch_task.cancel()
+                with contextlib.suppress(asyncio.CancelledError):
+                    await cancel_watch_task
+
+    async def _watch_cancel_signal(
+        self,
+        *,
+        cancel_checker: Callable[[], Awaitable[bool]],
+        stop_signal: asyncio.Event,
+        run_task: asyncio.Task[object],
+    ) -> None:
+        while not stop_signal.is_set():
+            should_cancel = False
+            try:
+                should_cancel = await cancel_checker()
+            except Exception:
+                should_cancel = False
+
+            if should_cancel:
+                async with self._active_agent_lock:
+                    active_agent = self._active_agent
+                if active_agent is not None:
+                    with contextlib.suppress(Exception):
+                        await active_agent.interrupt()
+                if not run_task.done():
+                    run_task.cancel("run canceled by user")
+                return
+
+            await asyncio.sleep(0.2)

    def _build_toolkit(
        self,
@@ -373,9 +425,16 @@ class AgentScopeRunner:
            model=tracking_model,
            emitter=emitter,
        )
-        response_msg = await agent.reply_json(
-            input_messages, output_model=worker_output_model
-        )
+        async with self._active_agent_lock:
+            self._active_agent = agent
+        try:
+            response_msg = await agent.reply_json(
+                input_messages, output_model=worker_output_model
+            )
+        finally:
+            async with self._active_agent_lock:
+                if self._active_agent is agent:
+                    self._active_agent = None
        worker_payload = worker_output_model.model_validate(response_msg.metadata or {})
        response_metadata = self._litellm_service.build_usage_metadata(
            model=stage_config.model_code,
@@ -257,6 +257,7 @@ async def run_agentscope_task(command: dict[str, Any]) -> dict[str, object]:
    thread_id = run_input.thread_id
    run_id = run_input.run_id
    owner_id = UUID(raw_owner_id)
+    cancel_key = f"agent:cancel:{thread_id}:{run_id}"

    if command_type != "run":
        raise ValueError("invalid command type")
@@ -278,6 +279,15 @@ async def run_agentscope_task(command: dict[str, Any]) -> dict[str, object]:
        work_memory: WorkProfileContent | None = memories_result.get("work_memory")

        redis_client = await get_or_init_redis_client()
+
+        async def _cancel_checker() -> bool:
+            exists_fn = getattr(redis_client, "exists", None)
+            if not callable(exists_fn):
+                return False
+            exists_call = cast(Any, exists_fn)(cancel_key)
+            result = await exists_call
+            return bool(result)
+
        bus = RedisStreamBus(
            client=redis_client,
            stream_prefix=config.agent_runtime.redis_stream_prefix,
@@ -302,14 +312,21 @@ async def run_agentscope_task(command: dict[str, Any]) -> dict[str, object]:
            context_config=runtime_config.context,
        )

-        await runtime.run(
-            run_input=run_input,
-            context_messages=context_messages,
-            user_context=user_context,
-            runtime_config=runtime_config,
-            user_memory=user_memory,
-            work_memory=work_memory,
-        )
+        try:
+            await runtime.run(
+                run_input=run_input,
+                context_messages=context_messages,
+                user_context=user_context,
+                runtime_config=runtime_config,
+                user_memory=user_memory,
+                work_memory=work_memory,
+                cancel_checker=_cancel_checker,
+            )
+        finally:
+            delete_fn = getattr(redis_client, "delete", None)
+            if callable(delete_fn):
+                delete_call = cast(Any, delete_fn)(cancel_key)
+                await delete_call
    logger.info(
        "agentscope runtime task completed",
        command_type=command_type,