Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 15 additions & 3 deletions src/agentex/lib/adk/_modules/acp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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,
),
Expand All @@ -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,
)
39 changes: 27 additions & 12 deletions src/agentex/lib/core/services/adk/acp/acp.py
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down Expand Up @@ -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,
)
39 changes: 39 additions & 0 deletions tests/test_task_cancel.py
Original file line number Diff line number Diff line change
@@ -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
Loading