solace-agent-mesh 1.0.8__py3-none-any.whl → 1.1.0__py3-none-any.whl

solace_agent_mesh/gateway/http_sse/routers/tasks.py

@@ -8,30 +8,28 @@ from fastapi import (
     HTTPException,
     Request as FastAPIRequest,
     status,
-    Form,
-    File,
-    UploadFile,
 )
-from pydantic import BaseModel, Field
-from typing import List
+from typing import Union
 
 from solace_ai_connector.common.log import log
 
 from ....gateway.http_sse.session_manager import SessionManager
 from ....gateway.http_sse.services.task_service import TaskService
 
-from ....common.types import (
-    JSONRPCResponse,
-    InternalError,
-    InvalidRequestError,
+from a2a.types import (
+    CancelTaskRequest,
+    SendMessageRequest,
+    SendStreamingMessageRequest,
+    SendMessageSuccessResponse,
+    SendStreamingMessageSuccessResponse,
 )
+from ....common import a2a
 
 from ....gateway.http_sse.dependencies import (
     get_session_manager,
     get_sac_component,
     get_task_service,
 )
-from ....gateway.http_sse.routers.users import get_current_user
 
 from typing import TYPE_CHECKING
 
@@ -41,29 +39,26 @@ if TYPE_CHECKING:
 router = APIRouter()
 
 
-class CancelTaskApiPayload(BaseModel):
-    """Request body for the task cancellation endpoint."""
+async def _submit_task(
+    request: FastAPIRequest,
+    payload: Union[SendMessageRequest, SendStreamingMessageRequest],
+    session_manager: SessionManager,
+    component: "WebUIBackendComponent",
+    is_streaming: bool,
+):
+    """Helper to submit a task, handling both streaming and non-streaming cases."""
+    log_prefix = f"[POST /api/v1/message:{'stream' if is_streaming else 'send'}] "
 
-    agent_name: str = Field(
-        ..., description="The name of the agent currently handling the task."
-    )
-    task_id: str = Field(..., description="The ID of the task to cancel.")
+    agent_name = None
+    if payload.params and payload.params.message and payload.params.message.metadata:
+        agent_name = payload.params.message.metadata.get("agent_name")
 
+    if not agent_name:
+        raise HTTPException(
+            status_code=status.HTTP_400_BAD_REQUEST,
+            detail="Missing 'agent_name' in request payload message metadata.",
+        )
 
-@router.post("/send", response_model=JSONRPCResponse)
-async def send_task_to_agent(
-    request: FastAPIRequest,
-    agent_name: str = Form(...),
-    message: str = Form(...),
-    files: List[UploadFile] = File([]),
-    session_manager: SessionManager = Depends(get_session_manager),
-    component: "WebUIBackendComponent" = Depends(get_sac_component),
-):
-    """
-    Submits a non-streaming task request to the specified agent.
-    Accepts multipart/form-data.
-    """
-    log_prefix = "[POST /api/v1/tasks/send] "
     log.info("%sReceived request for agent: %s", log_prefix, agent_name)
 
     try:
@@ -87,207 +82,159 @@ async def send_task_to_agent(
             "%sUsing ClientID: %s, SessionID: %s", log_prefix, client_id, session_id
         )
 
-        external_event_data = {
-            "agent_name": agent_name,
-            "message": message,
-            "files": files,
-            "client_id": client_id,
+        # Use the helper to get the unwrapped parts from the incoming message.
+        a2a_parts = a2a.get_parts_from_message(payload.params.message)
+
+        external_req_ctx = {
+            "app_name_for_artifacts": component.gateway_id,
+            "user_id_for_artifacts": client_id,
             "a2a_session_id": session_id,
+            "user_id_for_a2a": client_id,
+            "target_agent_name": agent_name,
         }
 
-        target_agent, a2a_parts, external_req_ctx = (
-            await component._translate_external_input(external_event_data)
-        )
-
         task_id = await component.submit_a2a_task(
-            target_agent_name=target_agent,
+            target_agent_name=agent_name,
             a2a_parts=a2a_parts,
             external_request_context=external_req_ctx,
             user_identity=user_identity,
-            is_streaming=False,
+            is_streaming=is_streaming,
         )
 
         log.info("%sTask submitted successfully. TaskID: %s", log_prefix, task_id)
 
-        return JSONRPCResponse(result={"taskId": task_id})
-
-    except InvalidRequestError as e:
-        log.warning("%sInvalid request: %s", log_prefix, e.message, exc_info=True)
-        raise HTTPException(
-            status_code=status.HTTP_400_BAD_REQUEST,
-            detail=e.model_dump(exclude_none=True),
+        task_object = a2a.create_initial_task(
+            task_id=task_id,
+            context_id=session_id,
+            agent_name=agent_name,
         )
+
+        if is_streaming:
+            return a2a.create_send_streaming_message_success_response(
+                result=task_object, request_id=payload.id
+            )
+        else:
+            return a2a.create_send_message_success_response(
+                result=task_object, request_id=payload.id
+            )
+
     except PermissionError as pe:
         log.warning("%sPermission denied: %s", log_prefix, str(pe))
         raise HTTPException(
             status_code=status.HTTP_403_FORBIDDEN,
             detail=str(pe),
         )
-    except InternalError as e:
-        log.error(
-            "%sInternal error submitting task: %s", log_prefix, e.message, exc_info=True
-        )
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=e.model_dump(exclude_none=True),
-        )
     except Exception as e:
         log.exception("%sUnexpected error submitting task: %s", log_prefix, e)
-        error_resp = InternalError(message="Unexpected server error: %s" % e)
+        error_resp = a2a.create_internal_error(
+            message="Unexpected server error: %s" % e
+        )
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail=error_resp.model_dump(exclude_none=True),
         )
 
 
-@router.post("/subscribe", response_model=JSONRPCResponse)
-async def subscribe_task_from_agent(
+@router.post("/message:send", response_model=SendMessageSuccessResponse)
+async def send_task_to_agent(
     request: FastAPIRequest,
-    agent_name: str = Form(...),
-    message: str = Form(...),
-    files: List[UploadFile] = File([]),
+    payload: SendMessageRequest,
     session_manager: SessionManager = Depends(get_session_manager),
     component: "WebUIBackendComponent" = Depends(get_sac_component),
-    user: dict = Depends(get_current_user),
 ):
     """
-    Submits a streaming task request (`tasks/sendSubscribe`) to the specified agent.
-    Accepts multipart/form-data.
-    The client should subsequently connect to the SSE endpoint using the returned taskId.
+    Submits a non-streaming task request to the specified agent.
+    Accepts application/json.
     """
-    log_prefix = "[POST /api/v1/tasks/subscribe] "
-    log.info("%sReceived streaming request for agent: %s", log_prefix, agent_name)
-
-    try:
-        user_identity = await component.authenticate_and_enrich_user(request)
-        if user_identity is None:
-            log.warning("%sUser authentication failed. Denying request.", log_prefix)
-            raise HTTPException(
-                status_code=status.HTTP_401_UNAUTHORIZED,
-                detail="User authentication failed or identity not found.",
-            )
-        log.info(
-            "%sAuthenticated user identity: %s",
-            log_prefix,
-            user_identity.get("id", "unknown"),
-        )
-
-        client_id = session_manager.get_a2a_client_id(request)
-        session_id = session_manager.ensure_a2a_session(request)
-
-        log.info(
-            "%sUsing ClientID: %s, SessionID: %s", log_prefix, client_id, session_id
-        )
-
-        external_event_data = {
-            "agent_name": agent_name,
-            "message": message,
-            "files": files,
-            "client_id": client_id,
-            "a2a_session_id": session_id,
-        }
+    return await _submit_task(
+        request=request,
+        payload=payload,
+        session_manager=session_manager,
+        component=component,
+        is_streaming=False,
+    )
 
-        target_agent, a2a_parts, external_req_ctx = (
-            await component._translate_external_input(external_event_data)
-        )
 
-        task_id = await component.submit_a2a_task(
-            target_agent_name=target_agent,
-            a2a_parts=a2a_parts,
-            external_request_context=external_req_ctx,
-            user_identity=user_identity,
-            is_streaming=True,
-        )
+@router.post("/message:stream", response_model=SendStreamingMessageSuccessResponse)
+async def subscribe_task_from_agent(
+    request: FastAPIRequest,
+    payload: SendStreamingMessageRequest,
+    session_manager: SessionManager = Depends(get_session_manager),
+    component: "WebUIBackendComponent" = Depends(get_sac_component),
+):
+    """
+    Submits a streaming task request to the specified agent.
+    Accepts application/json.
+    The client should subsequently connect to the SSE endpoint using the returned taskId.
+    """
+    return await _submit_task(
+        request=request,
+        payload=payload,
+        session_manager=session_manager,
+        component=component,
+        is_streaming=True,
+    )
 
-        log.info(
-            "%sStreaming task submitted successfully. TaskID: %s", log_prefix, task_id
-        )
 
-        return JSONRPCResponse(result={"taskId": task_id})
+@router.post("/tasks/{taskId}:cancel", status_code=status.HTTP_202_ACCEPTED)
+async def cancel_agent_task(
+    request: FastAPIRequest,
+    taskId: str,
+    payload: CancelTaskRequest,
+    session_manager: SessionManager = Depends(get_session_manager),
+    task_service: TaskService = Depends(get_task_service),
+    component: "WebUIBackendComponent" = Depends(get_sac_component),
+):
+    """
+    Sends a cancellation request for a specific task to the specified agent.
+    Returns 202 Accepted, as cancellation is asynchronous.
+    """
+    log_prefix = f"[POST /api/v1/tasks/{taskId}:cancel] "
+    log.info("%sReceived cancellation request.", log_prefix)
 
-    except InvalidRequestError as e:
-        log.warning("%sInvalid request: %s", log_prefix, e.message, exc_info=True)
+    if taskId != payload.params.id:
         raise HTTPException(
             status_code=status.HTTP_400_BAD_REQUEST,
-            detail=e.model_dump(exclude_none=True),
+            detail="Task ID in URL path does not match task ID in payload.",
         )
-    except PermissionError as pe:
-        log.warning("%sPermission denied: %s", log_prefix, str(pe))
+
+    context = component.task_context_manager.get_context(taskId)
+    if not context:
         raise HTTPException(
-            status_code=status.HTTP_403_FORBIDDEN,
-            detail=str(pe),
+            status_code=status.HTTP_404_NOT_FOUND,
+            detail=f"No active task context found for task ID: {taskId}",
         )
-    except InternalError as e:
+
+    agent_name = context.get("target_agent_name")
+    if not agent_name:
         log.error(
-            "%sInternal error submitting streaming task: %s",
+            "%sCould not determine target agent for task %s. Context is missing 'target_agent_name'.",
             log_prefix,
-            e.message,
-            exc_info=True,
-        )
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=e.model_dump(exclude_none=True),
+            taskId,
         )
-    except Exception as e:
-        log.exception("%sUnexpected error submitting streaming task: %s", log_prefix, e)
-        error_resp = InternalError(message="Unexpected server error: %s" % e)
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=error_resp.model_dump(exclude_none=True),
+            detail="Could not determine target agent for the task.",
         )
 
-
-@router.post("/cancel", status_code=status.HTTP_202_ACCEPTED)
-async def cancel_agent_task(
-    request: FastAPIRequest,
-    payload: CancelTaskApiPayload,
-    session_manager: SessionManager = Depends(get_session_manager),
-    task_service: TaskService = Depends(get_task_service),
-):
-    """
-    Sends a cancellation request for a specific task to the specified agent.
-    Returns 202 Accepted, as cancellation is asynchronous.
-    """
-    log_prefix = "[POST /api/v1/tasks/cancel][Task:%s] " % payload.task_id
-    log.info(
-        "%sReceived cancellation request for agent: %s", log_prefix, payload.agent_name
-    )
+    log.info("%sTarget agent for cancellation is '%s'", log_prefix, agent_name)
 
     try:
         client_id = session_manager.get_a2a_client_id(request)
 
         log.info("%sUsing ClientID: %s", log_prefix, client_id)
 
-        await task_service.cancel_task(
-            payload.agent_name, payload.task_id, client_id, client_id
-        )
+        await task_service.cancel_task(agent_name, taskId, client_id, client_id)
 
         log.info("%sCancellation request published successfully.", log_prefix)
 
         return {"message": "Cancellation request sent"}
 
-    except InvalidRequestError as e:
-        log.warning(
-            "%sInvalid cancellation request: %s", log_prefix, e.message, exc_info=True
-        )
-        raise HTTPException(
-            status_code=status.HTTP_400_BAD_REQUEST,
-            detail=e.model_dump(exclude_none=True),
-        )
-    except InternalError as e:
-        log.error(
-            "%sInternal error sending cancellation: %s",
-            log_prefix,
-            e.message,
-            exc_info=True,
-        )
-        raise HTTPException(
-            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
-            detail=e.model_dump(exclude_none=True),
-        )
     except Exception as e:
         log.exception("%sUnexpected error sending cancellation: %s", log_prefix, e)
-        error_resp = InternalError(message="Unexpected server error: %s" % e)
+        error_resp = a2a.create_internal_error(
+            message="Unexpected server error: %s" % e
+        )
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail=error_resp.model_dump(exclude_none=True),

solace_agent_mesh/gateway/http_sse/services/agent_service.py

@@ -8,7 +8,7 @@ from typing import List, Optional
 from solace_ai_connector.common.log import log
 
 from ....common.agent_registry import AgentRegistry
-from ....common.types import AgentCard
+from a2a.types import AgentCard
 
 
 class AgentService:

solace_agent_mesh/gateway/http_sse/services/services_llm.txt

@@ -1,179 +1,133 @@
-Here is the developer guide for the `services` directory.
+# DEVELOPER GUIDE: services
 
 ## Quick Summary
-The `services` directory contains the business logic layer for the HTTP SSE Gateway. It encapsulates core functionalities by providing high-level services for agent management, user identity searches, and A2A (Agent-to-Agent) task operations. These services abstract the complexities of interacting with the agent registry, external identity providers, and the underlying A2A messaging protocol.
+The `services` directory contains the business logic layer for the HTTP SSE Gateway. It provides high-level services for agent management (discovering and retrieving A2A agents), user identity operations (searching for users), and task management (cancelling A2A tasks). These services abstract the complexities of interacting with agent registries, identity providers, and A2A messaging protocols.
 
 ## Files Overview
-- `__init__.py` - Marks the directory as a Python package.
-- `agent_service.py` - Service for retrieving information about discovered A2A agents.
-- `people_service.py` - Service for searching for users via a configured identity service.
-- `task_service.py` - Service for handling the cancellation of tasks with A2A agents.
+- `__init__.py` - Package initialization file marking the directory as a Python package
+- `agent_service.py` - Service for retrieving information about discovered A2A agents from the registry
+- `people_service.py` - Service for searching users via configured identity services
+- `task_service.py` - Service for handling A2A task cancellation operations
 
 ## Developer API Reference
 
 ### __init__.py
-**Purpose:** This file marks the `services` directory as a Python package, allowing its modules to be imported.
-**Import:** N/A - No public interfaces.
-
----
+**Purpose:** Marks the services directory as a Python package
+**Import:** N/A - No public interfaces
 
 ### agent_service.py
-**Purpose:** Provides high-level methods for accessing information about discovered A2A agents from the shared `AgentRegistry`.
-**Import:** `from src.solace_agent_mesh.gateway.http_sse.services.agent_service import AgentService`
+**Purpose:** Provides methods for accessing information about discovered A2A agents from the shared AgentRegistry
+**Import:** `from solace_agent_mesh.gateway.http_sse.services.agent_service import AgentService`
 
 **Classes:**
-- `AgentService(agent_registry: AgentRegistry)` - A service that provides methods for accessing information about discovered A2A agents.
-  - `get_all_agents() -> List[AgentCard]` - Retrieves all currently discovered and registered agent cards from the registry.
-  - `get_agent_by_name(agent_name: str) -> Optional[AgentCard]` - Retrieves a specific agent card by its unique name. Returns `None` if the agent is not found.
+- `AgentService(agent_registry: AgentRegistry)` - Service for accessing discovered A2A agent information
+  - `get_all_agents() -> List[AgentCard]` - Retrieves all currently discovered and registered agent cards
+  - `get_agent_by_name(agent_name: str) -> Optional[AgentCard]` - Retrieves a specific agent card by name, returns None if not found
 
 **Usage Examples:**
 ```python
-from typing import List, Optional
-from src.solace_agent_mesh.gateway.http_sse.services.agent_service import AgentService
-from src.solace_agent_mesh.common.agent_registry import AgentRegistry
-from src.solace_agent_mesh.common.types import AgentCard
-
-# In a real application, AgentRegistry would be a shared instance.
-# For this example, we'll create a new one and populate it.
-agent_registry = AgentRegistry()
-my_agent_card = AgentCard(name="data-analyzer", description="Analyzes data files.")
-agent_registry.register_agent(my_agent_card)
-
-# 1. Initialize the service with the agent registry
+from solace_agent_mesh.gateway.http_sse.services.agent_service import AgentService
+from solace_agent_mesh.common.agent_registry import AgentRegistry
+
+# Initialize with shared agent registry
+agent_registry = AgentRegistry()  # Usually injected as shared instance
 agent_service = AgentService(agent_registry=agent_registry)
 
-# 2. Get all available agents
+# Get all available agents
 all_agents = agent_service.get_all_agents()
-print(f"Found {len(all_agents)} agent(s).")
-for agent in all_agents:
-    print(f"- Agent: {agent.name}, Description: {agent.description}")
-
-# 3. Get a specific agent by name
-found_agent = agent_service.get_agent_by_name("data-analyzer")
-if found_agent:
-    print(f"\nSuccessfully retrieved agent: {found_agent.name}")
-else:
-    print("\nCould not find agent 'data-analyzer'.")
+print(f"Found {len(all_agents)} agents")
 
-# 4. Try to get a non-existent agent
-missing_agent = agent_service.get_agent_by_name("non-existent-agent")
-print(f"Result for 'non-existent-agent': {missing_agent}")
+# Get specific agent by name
+agent = agent_service.get_agent_by_name("data-processor")
+if agent:
+    print(f"Found agent: {agent.name}")
+else:
+    print("Agent not found")
 ```
 
----
-
 ### people_service.py
-**Purpose:** Acts as a layer on top of a configured `IdentityService` to provide user search functionality. If no identity service is configured, it gracefully returns empty results.
-**Import:** `from src.solace_agent_mesh.gateway.http_sse.services.people_service import PeopleService`
+**Purpose:** Provides user search functionality via configured identity services
+**Import:** `from solace_agent_mesh.gateway.http_sse.services.people_service import PeopleService`
 
 **Classes:**
-- `PeopleService(identity_service: Optional[BaseIdentityService])` - A service for searching and retrieving user information.
-  - `search_for_users(query: str, limit: int = 10) -> List[Dict[str, Any]]` - Asynchronously searches for users via the identity service. Returns an empty list if the query is too short, no identity service is configured, or an error occurs.
+- `PeopleService(identity_service: Optional[BaseIdentityService])` - Service for searching and retrieving user information
+  - `search_for_users(query: str, limit: int = 10) -> List[Dict[str, Any]]` - Asynchronously searches for users, returns empty list if no identity service configured or query too short
 
 **Usage Examples:**
 ```python
 import asyncio
-from typing import Any, Dict, List, Optional
-from src.solace_agent_mesh.gateway.http_sse.services.people_service import PeopleService
-from src.solace_agent_mesh.common.services.identity_service import BaseIdentityService
-
-# Define a mock identity service for the example
-class MockIdentityService(BaseIdentityService):
-    async def search_users(self, query: str, limit: int) -> List[Dict[str, Any]]:
-        print(f"MockIdentityService: Searching for '{query}' with limit {limit}")
-        all_users = [
-            {"id": "jdoe", "name": "John Doe", "email": "j.doe@example.com"},
-            {"id": "jsmith", "name": "Jane Smith", "email": "j.smith@example.com"},
-        ]
-        return [user for user in all_users if query.lower() in user["name"].lower()][:limit]
-
-async def main():
-    # 1. Initialize with a configured identity service
-    identity_service = MockIdentityService()
-    people_service = PeopleService(identity_service=identity_service)
-
-    # 2. Search for users
+from solace_agent_mesh.gateway.http_sse.services.people_service import PeopleService
+from solace_agent_mesh.common.services.identity_service import BaseIdentityService
+
+# Initialize with identity service
+identity_service = SomeIdentityService()  # Your identity service implementation
+people_service = PeopleService(identity_service=identity_service)
+
+async def search_users():
+    # Search for users
     users = await people_service.search_for_users("john", limit=5)
-    print(f"Found {len(users)} user(s): {users}")
+    for user in users:
+        print(f"User: {user.get('name')} - {user.get('email')}")
 
-    # 3. Initialize without an identity service
-    people_service_no_id = PeopleService(identity_service=None)
-    
-    # 4. Search will return an empty list
-    empty_results = await people_service_no_id.search_for_users("jane")
-    print(f"Results with no identity service: {empty_results}")
+# Initialize without identity service (graceful degradation)
+people_service_no_id = PeopleService(identity_service=None)
+# search_for_users will return empty list
 
-if __name__ == "__main__":
-    asyncio.run(main())
+asyncio.run(search_users())
 ```
 
----
-
 ### task_service.py
-**Purpose:** Handles the business logic for cancelling A2A tasks. It uses `CoreA2AService` to construct the cancellation message and a provided publisher function to send it over the messaging fabric.
-**Import:** `from src.solace_agent_mesh.gateway.http_sse.services.task_service import TaskService, PublishFunc`
-
-**Classes:**
-- `TaskService(core_a2a_service: CoreA2AService, publish_func: PublishFunc, namespace: str, gateway_id: str, sse_manager: SSEManager, task_context_map: Dict[str, Dict], task_context_lock: threading.Lock, app_name: str)` - A service for managing A2A task operations.
-  - `cancel_task(agent_name: str, task_id: str, client_id: str, user_id: str = "web_user") -> None` - Asynchronously constructs and publishes an A2A `CancelTaskRequest` message for a specific task.
+**Purpose:** Handles A2A task operations, specifically task cancellation using CoreA2AService and message publishing
+**Import:** `from solace_agent_mesh.gateway.http_sse.services.task_service import TaskService, PublishFunc`
 
 **Type Aliases:**
-- `PublishFunc: Callable[[str, Dict, Optional[Dict]], None]` - A callable that sends a message. It takes a topic, a payload dictionary, and optional user properties.
+- `PublishFunc: Callable[[str, Dict, Optional[Dict]], None]` - Function type for publishing messages (topic, payload, user_properties)
+
+**Classes:**
+- `TaskService(core_a2a_service: CoreA2AService, publish_func: PublishFunc, namespace: str, gateway_id: str, sse_manager: SSEManager, task_context_map: Dict[str, Dict], task_context_lock: threading.Lock, app_name: str)` - Service for managing A2A task operations
+  - `cancel_task(agent_name: str, task_id: str, client_id: str, user_id: str = "web_user") -> None` - Asynchronously cancels a task by publishing A2A CancelTaskRequest message
 
 **Usage Examples:**
 ```python
 import asyncio
 import threading
-from typing import Callable, Dict, Optional
-
-from src.solace_agent_mesh.gateway.http_sse.services.task_service import TaskService, PublishFunc
-from src.solace_agent_mesh.core_a2a.service import CoreA2AService
-from src.solace_agent_mesh.gateway.http_sse.sse_manager import SSEManager
-
-# Mock dependencies for the example
-class MockCoreA2AService(CoreA2AService):
-    def cancel_task(self, agent_name, task_id, client_id, user_id):
-        topic = f"a2a/request/{agent_name}/cancelTask"
-        payload = {"taskId": task_id}
-        user_props = {"clientId": client_id, "userId": user_id}
-        return topic, payload, user_props
-
-def my_publish_func(topic: str, payload: Dict, user_properties: Optional[Dict]):
-    print("\n--- Publishing Message ---")
-    print(f"Topic: {topic}")
-    print(f"Payload: {payload}")
-    print(f"User Properties: {user_properties}")
-    print("------------------------")
-
-async def main():
-    # 1. Set up dependencies
-    core_a2a_service = MockCoreA2AService()
-    sse_manager = SSEManager()
-    task_context_map = {}
-    task_context_lock = threading.Lock()
-
-    # 2. Initialize TaskService
-    task_service = TaskService(
-        core_a2a_service=core_a2a_service,
-        publish_func=my_publish_func,
-        namespace="my-namespace",
-        gateway_id="gateway-01",
-        sse_manager=sse_manager,
-        task_context_map=task_context_map,
-        task_context_lock=task_context_lock,
-        app_name="my-app"
-    )
-
-    # 3. Call the cancel_task method
-    print("Requesting task cancellation...")
+from solace_agent_mesh.gateway.http_sse.services.task_service import TaskService, PublishFunc
+from solace_agent_mesh.core_a2a.service import CoreA2AService
+from solace_agent_mesh.gateway.http_sse.sse_manager import SSEManager
+
+# Define publish function
+def my_publish_func(topic: str, payload: dict, user_properties: dict = None):
+    print(f"Publishing to {topic}: {payload}")
+    # Your actual message publishing logic here
+
+# Initialize dependencies
+core_a2a_service = CoreA2AService()  # Your core A2A service
+sse_manager = SSEManager()
+task_context_map = {}
+task_context_lock = threading.Lock()
+
+# Create task service
+task_service = TaskService(
+    core_a2a_service=core_a2a_service,
+    publish_func=my_publish_func,
+    namespace="my-namespace",
+    gateway_id="gateway-01",
+    sse_manager=sse_manager,
+    task_context_map=task_context_map,
+    task_context_lock=task_context_lock,
+    app_name="my-app"
+)
+
+async def cancel_task_example():
+    # Cancel a task
     await task_service.cancel_task(
-        agent_name="report-generator",
-        task_id="task-12345",
-        client_id="client-abcde",
-        user_id="test.user"
+        agent_name="data-processor",
+        task_id="task-123",
+        client_id="client-456",
+        user_id="user@example.com"
     )
-    print("Task cancellation request sent.")
 
-if __name__ == "__main__":
-    asyncio.run(main())
-```
+asyncio.run(cancel_task_example())
+```
+
+# content_hash: 83ddf6b403dc50598ed550e4b3a5445f832b3956dad75f7a3fbbb7e6e5c6c115