Merge pull request #102 from scaleapi/fix/task-cancellation

Fix: task cancellation bug - Separate agent and task identification

Author: prassanna-ravishankar

src/agentex/lib/adk/_modules/acp.py

@@ -205,18 +205,22 @@ async def cancel_task(
         self,
         task_id: str | None = None,
         task_name: str | None = None,
+        agent_id: str | None = None,
+        agent_name: str | None = None,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
         start_to_close_timeout: timedelta = timedelta(seconds=5),
         heartbeat_timeout: timedelta = timedelta(seconds=5),
         retry_policy: RetryPolicy = DEFAULT_RETRY_POLICY,
     ) -> Task:
         """
-        Cancel a task.
+        Cancel a task by sending cancel request to the agent that owns the task.
 
         Args:
-            task_id: The ID of the task to cancel.
-            task_name: The name of the task to cancel.
+            task_id: ID of the task to cancel.
+            task_name: Name of the task to cancel.
+            agent_id: ID of the agent that owns the task.
+            agent_name: Name of the agent that owns the task.
             trace_id: The trace ID for the task.
             parent_span_id: The parent span ID for the task.
             start_to_close_timeout: The start to close timeout for the task.
@@ -225,13 +229,19 @@ async def cancel_task(
 
         Returns:
             The task entry.
+            
+        Raises:
+            ValueError: If neither agent_name nor agent_id is provided,
+                       or if neither task_name nor task_id is provided
         """
         if in_temporal_workflow():
             return await ActivityHelpers.execute_activity(
                 activity_name=ACPActivityName.TASK_CANCEL,
                 request=TaskCancelParams(
                     task_id=task_id,
                     task_name=task_name,
+                    agent_id=agent_id,
+                    agent_name=agent_name,
                     trace_id=trace_id,
                     parent_span_id=parent_span_id,
                 ),
@@ -244,6 +254,8 @@ async def cancel_task(
             return await self._acp_service.task_cancel(
                 task_id=task_id,
                 task_name=task_name,
+                agent_id=agent_id,
+                agent_name=agent_name,
                 trace_id=trace_id,
                 parent_span_id=parent_span_id,
             )

src/agentex/lib/core/services/adk/acp/acp.py

@@ -180,37 +180,52 @@ async def task_cancel(
         self,
         task_id: str | None = None,
         task_name: str | None = None,
+        agent_id: str | None = None,
+        agent_name: str | None = None,
         trace_id: str | None = None,
         parent_span_id: str | None = None,
-    ) -> Task:
+    ) -> Task:        
+        # Require agent identification
+        if not agent_name and not agent_id:
+            raise ValueError("Either agent_name or agent_id must be provided to identify the agent that owns the task")
+            
+        # Require task identification
+        if not task_name and not task_id:
+            raise ValueError("Either task_name or task_id must be provided to identify the task to cancel")
+        
         trace = self._tracer.trace(trace_id=trace_id)
         async with trace.span(
             parent_id=parent_span_id,
             name="task_cancel",
             input={
                 "task_id": task_id,
                 "task_name": task_name,
+                "agent_id": agent_id,
+                "agent_name": agent_name,
             },
         ) as span:
             heartbeat_if_in_workflow("task cancel")
+            
+            # Build params for the agent (task identification)
+            params = {}
+            if task_id:
+                params["task_id"] = task_id
             if task_name:
+                params["task_name"] = task_name
+            
+            # Send cancel request to the correct agent
+            if agent_name:
                 json_rpc_response = await self._agentex_client.agents.rpc_by_name(
-                    agent_name=task_name,
+                    agent_name=agent_name,
                     method="task/cancel",
-                    params={
-                        "task_name": task_name,
-                    },
+                    params=params,
                 )
-            elif task_id:
+            else:  # agent_id is provided (validated above)
                 json_rpc_response = await self._agentex_client.agents.rpc(
-                    agent_id=task_id,
+                    agent_id=agent_id,
                     method="task/cancel",
-                    params={
-                        "task_id": task_id,
-                    },
+                    params=params,
                 )
-            else:
-                raise ValueError("Either task_name or task_id must be provided")
 
             task_entry = Task.model_validate(json_rpc_response.result)
             if span:

src/agentex/lib/core/temporal/activities/adk/acp/acp_activities.py

@@ -45,6 +45,8 @@ class EventSendParams(BaseModelWithTraceParams):
 class TaskCancelParams(BaseModelWithTraceParams):
     task_id: str | None = None
     task_name: str | None = None
+    agent_id: str | None = None
+    agent_name: str | None = None
 
 
 class ACPActivities:
@@ -83,4 +85,8 @@ async def task_cancel(self, params: TaskCancelParams) -> Task:
         return await self._acp_service.task_cancel(
             task_id=params.task_id,
             task_name=params.task_name,
+            agent_id=params.agent_id,
+            agent_name=params.agent_name,
+            trace_id=params.trace_id,
+            parent_span_id=params.parent_span_id,
         )

tests/test_task_cancel.py

@@ -0,0 +1,39 @@
+"""Tests for task cancellation bug fix."""
+
+import os
+import pytest
+
+from agentex import AsyncAgentex
+from tests.utils import assert_matches_type
+from agentex.types import Task
+
+base_url = os.environ.get("TEST_API_BASE_URL", "http://127.0.0.1:4010")
+
+
+class TestTaskCancelBugFix:
+    """Test that task cancellation bug is fixed - agent identification is required."""
+    parametrize = pytest.mark.parametrize("client", [False, True], indirect=True, ids=["loose", "strict"])
+
+    @pytest.mark.skip(reason="Integration test - demonstrates the fix for task cancel bug")
+    @parametrize  
+    async def test_task_cancel_requires_agent_and_task_identification(self, client: AsyncAgentex) -> None:
+        """
+        Test that demonstrates the task cancellation bug fix.
+        
+        Previously: task_cancel(task_name="my-task") incorrectly treated task_name as agent_name
+        Fixed: task_cancel(task_name="my-task", agent_name="my-agent") correctly identifies both
+        """
+        # This test documents the correct usage pattern
+        # In practice, you would need a real agent and task for this to work
+        try:
+            task = await client.agents.cancel_task(
+                agent_name="test-agent",  # REQUIRED: Agent that owns the task
+                params={
+                    "task_id": "test-task-123"  # REQUIRED: Task to cancel
+                }
+            )
+            assert_matches_type(Task, task, path=["response"])
+        except Exception:
+            # Expected to fail in test environment without real agents/tasks
+            # The important thing is that the API now requires both parameters
+            pass