smarta2a 0.2.1__py3-none-any.whl → 0.2.3__py3-none-any.whl

smarta2a/server/state_manager.py

@@ -0,0 +1,34 @@
+# Library imports
+from typing import Optional, Dict, Any
+from uuid import uuid4
+
+# Local imports
+from smarta2a.state_stores.base_state_store import BaseStateStore
+from smarta2a.history_update_strategies.history_update_strategy import HistoryUpdateStrategy
+from smarta2a.utils.types import Message, StateData
+
+class StateManager:
+    def __init__(self, store: Optional[BaseStateStore], history_strategy: HistoryUpdateStrategy):
+        self.store = store
+        self.strategy = history_strategy
+
+    def init_or_get(self, session_id: Optional[str], message: Message, metadata: Dict[str, Any]) -> StateData:
+        sid = session_id or str(uuid4())
+        if not self.store:
+            return StateData(sessionId=sid, history=[message], metadata=metadata or {})
+        existing = self.store.get_state(sid) or StateData(sessionId=sid, history=[], metadata={})
+        existing.history.append(message)
+        existing.metadata = {**(existing.metadata or {}), **(metadata or {})}
+        self.store.update_state(sid, existing)
+        return existing
+
+    def update(self, state: StateData):
+        if self.store:
+            self.store.update_state(state.sessionId, state)
+    
+    def get_store(self) -> Optional[BaseStateStore]:
+        return self.store
+    
+    def get_strategy(self) -> HistoryUpdateStrategy:
+        return self.strategy
+    

smarta2a/server/subscription_service.py

@@ -0,0 +1,109 @@
+# Library imports
+from typing import Optional, List, Dict, Any, AsyncGenerator, Union
+from datetime import datetime
+from collections import defaultdict
+from uuid import uuid4
+from fastapi.responses import StreamingResponse
+from sse_starlette.sse import EventSourceResponse
+
+# Local imports
+from smarta2a.server.handler_registry import HandlerRegistry
+from smarta2a.server.state_manager import StateManager
+from smarta2a.utils.types import (
+    Message, StateData, SendTaskStreamingRequest, SendTaskStreamingResponse,
+    TaskSendParams, A2AStatus, A2AStreamResponse, TaskStatusUpdateEvent,
+    TaskStatus, TaskState, TaskArtifactUpdateEvent, Artifact, TextPart,
+    FilePart, DataPart, FileContent, MethodNotFoundError, TaskNotFoundError,
+    InternalError
+)
+
+class SubscriptionService:
+    def __init__(self, registry: HandlerRegistry, state_mgr: StateManager):
+        self.registry = registry
+        self.state_mgr = state_mgr
+
+    async def subscribe(self, request: SendTaskStreamingRequest, state: Optional[StateData]) -> StreamingResponse:
+        handler = self.registry.get_subscription("tasks/sendSubscribe")
+        if not handler:
+            err = SendTaskStreamingResponse(jsonrpc="2.0", id=request.id, error=MethodNotFoundError()).model_dump_json()
+            return EventSourceResponse(err)
+
+        session_id = state.sessionId if state else request.params.sessionId or str(uuid4())
+        history = state.history.copy() if state else [request.params.message]
+        metadata = state.metadata.copy() if state else (request.params.metadata or {})
+
+        async def event_stream():
+            try:
+                events = handler(request, state) if state else handler(request)
+                async for ev in self._normalize(request.params, events, history.copy(), metadata.copy(), session_id):
+                    yield f"data: {ev}\n\n"
+            except Exception as e:
+                err = TaskNotFoundError() if 'not found' in str(e).lower() else InternalError(data=str(e))
+                msg = SendTaskStreamingResponse(jsonrpc="2.0", id=request.id, error=err).model_dump_json()
+                yield f"data: {msg}\n\n"
+
+        return StreamingResponse(event_stream(), media_type="text/event-stream; charset=utf-8")
+
+    async def _normalize(
+        self,
+        params: TaskSendParams,
+        events: AsyncGenerator,
+        history: List[Message],
+        metadata: Dict[str, Any],
+        session_id: str
+    ) -> AsyncGenerator[str, None]:
+        artifact_state = defaultdict(lambda: {"index": 0, "last_chunk": False})
+        async for item in events:
+            if isinstance(item, SendTaskStreamingResponse):
+                yield item.model_dump_json()
+                continue
+
+            if isinstance(item, A2AStatus):
+                te = TaskStatusUpdateEvent(
+                    id=params.id,
+                    status=TaskStatus(state=TaskState(item.status), timestamp=datetime.now()),
+                    final=item.final or (item.status.lower() == TaskState.COMPLETED),
+                    metadata=item.metadata
+                )
+                yield SendTaskStreamingResponse(jsonrpc="2.0", id=params.id, result=te).model_dump_json()
+                continue
+
+            content_item = item
+            if not isinstance(item, A2AStreamResponse):
+                content_item = A2AStreamResponse(content=item)
+
+            parts: List[Union[TextPart, FilePart, DataPart]] = []
+            cont = content_item.content
+            if isinstance(cont, str): parts.append(TextPart(text=cont))
+            elif isinstance(cont, bytes): parts.append(FilePart(file=FileContent(bytes=cont)))
+            elif isinstance(cont, (TextPart, FilePart, DataPart)): parts.append(cont)
+            elif isinstance(cont, Artifact): parts.extend(cont.parts)
+            elif isinstance(cont, list):
+                for elem in cont:
+                    if isinstance(elem, str): parts.append(TextPart(text=elem))
+                    elif isinstance(elem, (TextPart, FilePart, DataPart)): parts.append(elem)
+                    elif isinstance(elem, Artifact): parts.extend(elem.parts)
+
+            idx = content_item.index
+            state = artifact_state[idx]
+            evt = TaskArtifactUpdateEvent(
+                id=params.id,
+                artifact=Artifact(
+                    parts=parts,
+                    index=idx,
+                    append=content_item.append or (state["index"] == idx),
+                    lastChunk=content_item.final or state["last_chunk"],
+                    metadata=content_item.metadata
+                )
+            )
+            if content_item.final:
+                state["last_chunk"] = True
+            state["index"] += 1
+
+            agent_msg = Message(role="agent", parts=evt.artifact.parts, metadata=evt.artifact.metadata)
+            new_hist = self.state_mgr.strategy.update_history(history, [agent_msg])
+            metadata = {**metadata, **(evt.artifact.metadata or {})}
+            self.state_mgr.update(StateData(session_id, new_hist, metadata))
+            history = new_hist
+
+            yield SendTaskStreamingResponse(jsonrpc="2.0", id=params.id, result=evt).model_dump_json()

smarta2a/server/task_service.py

@@ -0,0 +1,155 @@
+# Library imports
+from typing import Optional, Union, Any
+from uuid import uuid4
+from fastapi import HTTPException
+from pydantic import ValidationError
+
+# Local imports
+from smarta2a.server.handler_registry import HandlerRegistry
+from smarta2a.server.state_manager import StateManager
+from smarta2a.utils.task_builder import TaskBuilder
+from smarta2a.utils.types import (
+    Message, StateData, SendTaskRequest, SendTaskResponse,
+    GetTaskRequest, GetTaskResponse, CancelTaskRequest, CancelTaskResponse,
+    SetTaskPushNotificationRequest, GetTaskPushNotificationRequest,
+    SetTaskPushNotificationResponse, GetTaskPushNotificationResponse,
+    TaskPushNotificationConfig, TaskState, A2AStatus,
+    JSONRPCError, MethodNotFoundError, InternalError, InvalidParamsError,
+    TaskNotCancelableError, UnsupportedOperationError
+)
+
+class TaskService:
+    def __init__(self, registry: HandlerRegistry, state_mgr: StateManager):
+        self.registry = registry
+        self.state_mgr = state_mgr
+        self.builder = TaskBuilder(default_status=TaskState.COMPLETED)
+
+    def send(self, request: SendTaskRequest, state: Optional[StateData]) -> SendTaskResponse:
+        handler = self.registry.get_handler("tasks/send")
+        if not handler:
+            return SendTaskResponse(id=request.id, error=MethodNotFoundError())
+
+        session_id = state.sessionId if state else request.params.sessionId or str(uuid4())
+        history = state.history.copy() if state else [request.params.message]
+        metadata = state.metadata.copy() if state else (request.params.metadata or {})
+
+        try:
+            raw = handler(request, state) if state else handler(request)
+            if isinstance(raw, SendTaskResponse):
+                return raw
+
+            task = self.builder.build(
+                content=raw,
+                task_id=request.params.id,
+                session_id=session_id,
+                metadata=metadata,
+                history=history
+            )
+
+            if task.artifacts:
+                parts = [p for a in task.artifacts for p in a.parts]
+                agent_msg = Message(role="agent", parts=parts, metadata=task.metadata)
+                new_hist = self.state_mgr.strategy.update_history(history, [agent_msg])
+                task.history = new_hist
+                self.state_mgr.update(StateData(sessionId=session_id, history=new_hist, metadata=metadata))
+
+            return SendTaskResponse(id=request.id, result=task)
+        except JSONRPCError as e:
+            return SendTaskResponse(id=request.id, error=e)
+        except Exception as e:
+            return SendTaskResponse(id=request.id, error=InternalError(data=str(e)))
+
+    def get(self, request: GetTaskRequest) -> GetTaskResponse:
+        handler = self.registry.get_handler("tasks/get")
+        if not handler:
+            return GetTaskResponse(id=request.id, error=MethodNotFoundError())
+        try:
+            raw = handler(request)
+            if isinstance(raw, GetTaskResponse):
+                return self._validate(raw, request)
+
+            task = self.builder.build(
+                content=raw,
+                task_id=request.params.id,
+                metadata=request.params.metadata or {}
+            )
+            return self._finalize(request, task)
+        except JSONRPCError as e:
+            return GetTaskResponse(id=request.id, error=e)
+        except Exception as e:
+            return GetTaskResponse(id=request.id, error=InternalError(data=str(e)))
+
+    def cancel(self, request: CancelTaskRequest) -> CancelTaskResponse:
+        handler = self.registry.get_handler("tasks/cancel")
+        if not handler:
+            return CancelTaskResponse(id=request.id, error=MethodNotFoundError())
+        try:
+            raw = handler(request)
+            if isinstance(raw, CancelTaskResponse):
+                return self._validate(raw, request)
+
+            if isinstance(raw, A2AStatus):
+                task = self.builder.normalize_from_status(status=raw.status, task_id=request.params.id, metadata=raw.metadata or {})
+            else:
+                task = self.builder.build(content=raw, task_id=request.params.id, metadata=raw.metadata or {})
+
+            if task.id != request.params.id:
+                raise InvalidParamsError(data=f"Task ID mismatch: {task.id} vs {request.params.id}")
+            if task.status.state not in [TaskState.CANCELED, TaskState.COMPLETED]:
+                raise TaskNotCancelableError()
+
+            return CancelTaskResponse(id=request.id, result=task)
+        except JSONRPCError as e:
+            return CancelTaskResponse(id=request.id, error=e)
+        except (InvalidParamsError, TaskNotCancelableError) as e:
+            return CancelTaskResponse(id=request.id, error=e)
+        except HTTPException as e:
+            if e.status_code == 405:
+                return CancelTaskResponse(id=request.id, error=UnsupportedOperationError())
+            return CancelTaskResponse(id=request.id, error=InternalError(data=str(e)))
+        except Exception as e:
+            return CancelTaskResponse(id=request.id, error=InternalError(data=str(e)))
+
+    def set_notification(self, request: SetTaskPushNotificationRequest) -> SetTaskPushNotificationResponse:
+        handler = self.registry.get_handler("tasks/pushNotification/set")
+        if not handler:
+            return SetTaskPushNotificationResponse(id=request.id, error=MethodNotFoundError())
+        try:
+            raw = handler(request)
+            if raw is None:
+                return SetTaskPushNotificationResponse(id=request.id, result=request.params)
+            if isinstance(raw, SetTaskPushNotificationResponse):
+                return raw
+        except JSONRPCError as e:
+            return SetTaskPushNotificationResponse(id=request.id, error=e)
+        except Exception as e:
+            return SetTaskPushNotificationResponse(id=request.id, error=InternalError(data=str(e)))
+
+    def get_notification(self, request: GetTaskPushNotificationRequest) -> GetTaskPushNotificationResponse:
+        handler = self.registry.get_handler("tasks/pushNotification/get")
+        if not handler:
+            return GetTaskPushNotificationResponse(id=request.id, error=MethodNotFoundError())
+        try:
+            raw = handler(request)
+            if isinstance(raw, GetTaskPushNotificationResponse):
+                return raw
+            cfg = TaskPushNotificationConfig.model_validate(raw)
+            return GetTaskPushNotificationResponse(id=request.id, result=cfg)
+        except ValidationError as e:
+            return GetTaskPushNotificationResponse(id=request.id, error=InvalidParamsError(data=e.errors()))
+        except JSONRPCError as e:
+            return GetTaskPushNotificationResponse(id=request.id, error=e)
+        except Exception as e:
+            return GetTaskPushNotificationResponse(id=request.id, error=InternalError(data=str(e)))
+
+    def _validate(self, resp: Union[SendTaskResponse, GetTaskResponse, CancelTaskResponse], req) -> Any:
+        if resp.result and resp.result.id != req.params.id:
+            return type(resp)(id=req.id, error=InvalidParamsError(data=f"Task ID mismatch: {resp.result.id} vs {req.params.id}"))
+        return resp
+
+    def _finalize(self, request: GetTaskRequest, task) -> GetTaskResponse:
+        if task.id != request.params.id:
+            return GetTaskResponse(id=request.id, error=InvalidParamsError(data=f"Task ID mismatch: {task.id} vs {request.params.id}"))
+        if request.params.historyLength and task.history:
+            task.history = task.history[-request.params.historyLength:]
+        return GetTaskResponse(id=request.id, result=task)

smarta2a/state_stores/__init__.py

@@ -0,0 +1,8 @@
+"""
+State store implementations for managing conversation state.
+"""
+
+from .base_state_store import BaseStateStore
+from .inmemory_state_store import InMemoryStateStore
+
+__all__ = ['BaseStateStore', 'InMemoryStateStore'] 

smarta2a/state_stores/base_state_store.py

@@ -0,0 +1,20 @@
+# Library imports
+from abc import ABC, abstractmethod
+from typing import Optional, List, Dict, Any
+
+# Local imports
+from smarta2a.utils.types import StateData, Message
+
+class BaseStateStore(ABC):
+    
+    @abstractmethod
+    async def get_state(self, session_id: str) -> Optional[StateData]:
+        pass
+    
+    @abstractmethod
+    async def update_state(self, session_id: str, state_data: StateData) -> None:
+        pass
+    
+    @abstractmethod
+    async def delete_state(self, session_id: str) -> None:
+        pass

smarta2a/state_stores/inmemory_state_store.py

@@ -0,0 +1,21 @@
+# Library imports
+from typing import Dict, Any, Optional, List
+import uuid
+
+# Local imports
+from smarta2a.state_stores.base_state_store import BaseStateStore
+from smarta2a.utils.types import StateData, Message
+
+class InMemoryStateStore(BaseStateStore):
+    def __init__(self):
+        self.states: Dict[str, StateData] = {}
+    
+    def get_state(self, session_id: str) -> Optional[StateData]:
+        return self.states.get(session_id)
+    
+    def update_state(self, session_id: str, state_data: StateData):
+        self.states[session_id] = state_data
+    
+    def delete_state(self, session_id: str):
+        if session_id in self.states:
+            del self.states[session_id]

smarta2a/utils/__init__.py

@@ -0,0 +1,32 @@
+from .types import *
+
+__all__ = [
+    "TaskSendParams",
+    "SendTaskRequest",
+    "GetTaskRequest",
+    "CancelTaskRequest",
+    "CancelTaskResponse",
+    "Task",
+    "TaskStatus",
+    "TaskState",
+    "Artifact",
+    "TextPart",
+    "FilePart",
+    "FileContent",
+    "A2AResponse",
+    "A2ARequest",
+    "TaskQueryParams",
+    "TaskStatusUpdateEvent",
+    "TaskArtifactUpdateEvent",
+    "A2AStatus",
+    "A2AStreamResponse",
+    "SendTaskResponse",
+    "Message",
+    "InternalError",
+    "TaskNotFoundError",
+    "SetTaskPushNotificationRequest",
+    "GetTaskPushNotificationRequest",
+    "SetTaskPushNotificationResponse",
+    "GetTaskPushNotificationResponse",
+    "TaskPushNotificationConfig"
+]

smarta2a/utils/prompt_helpers.py

@@ -0,0 +1,38 @@
+# Library imports
+from typing import Optional, List
+
+# Local imports
+from smarta2a.client.tools_manager import ToolsManager
+from smarta2a.utils.types import AgentCard
+
+def build_system_prompt(
+    base_prompt: Optional[str],
+    tools_manager: ToolsManager,
+    mcp_server_urls_or_paths: Optional[List[str]] = None,
+    agent_cards: Optional[List[AgentCard]] = None
+) -> str:
+    """
+    Compose the final system prompt by combining the base prompt
+    with a clear listing of available tools.
+    """
+    header = base_prompt or "You are a helpful assistant with access to the following tools:"
+    
+    if mcp_server_urls_or_paths:
+        mcp_tools_desc = tools_manager.describe_tools("mcp")
+        header += f"\n\nAvailable tools:\n{mcp_tools_desc}"
+    
+    if agent_cards:
+        a2a_tools_desc = tools_manager.describe_tools("a2a")
+        header += f"\n\nIf needed, you can delegate parts of your task to other agents. The Agents you can use are:\n{_print_agent_list(agent_cards)}\n\nUse the following tools to send tasks to an agent:\n{a2a_tools_desc}"
+    
+    return header 
+
+
+def _print_agent_list(agents: List[AgentCard]) -> None:
+    """Prints multiple agents with separators"""
+    separator = "---"
+    agent_strings = [agent.pretty_print(include_separators=False) for agent in agents]
+    full_output = [separator]
+    full_output.extend(agent_strings)
+    full_output.append(separator)
+    print("\n".join(full_output))

smarta2a/utils/task_builder.py

@@ -0,0 +1,153 @@
+# Library imports
+from uuid import uuid4
+from datetime import datetime
+from typing import Any, List, Optional, Dict, Union
+from pydantic import ValidationError
+
+# Local imports
+from smarta2a.utils.types import (
+    Task,
+    TaskStatus,
+    TaskState,
+    Artifact,
+    Part,
+    TextPart,
+    FilePart,
+    DataPart,
+    Message,
+    A2AResponse,
+)
+
+class TaskBuilder:
+    def __init__(
+        self,
+        default_status: TaskState = TaskState.COMPLETED,
+    ):
+        self.default_status = default_status
+
+    def build(
+        self,
+        content: Any,
+        task_id: str,
+        session_id: Optional[str] = None,
+        metadata: Optional[Dict[str,Any]] = None,
+        history: Optional[List[Message]] = None,
+    ) -> Task:
+        """Universal task construction from various return types."""
+        history = history or []
+        metadata = metadata or {}
+
+        # 1) If the handler already gave us a full Task, just fix IDs & history:
+        if isinstance(content, Task):
+            content.sessionId = content.sessionId or session_id
+            content.history = history + (content.history or [])
+            content.metadata = content.metadata or metadata
+            return content
+
+        # 2) If they returned an A2AResponse, extract status/content:
+        if isinstance(content, A2AResponse):
+            # prefer the sessionId inside the A2AResponse
+            sid = content.sessionId or session_id
+            # merge metadata from builder-call and from A2AResponse
+            md = {**(metadata or {}), **(content.metadata or {})}
+            status = (
+                content.status
+                if isinstance(content.status, TaskStatus)
+                else TaskStatus(state=content.status)
+            )
+            artifacts = self._normalize_content(content.content)
+            return Task(
+                id=task_id,
+                sessionId=sid,
+                status=status,
+                artifacts=artifacts,
+                metadata=md,
+                history=history,
+            )
+
+        # 3) If they returned a plain dict describing a Task:
+        if isinstance(content, dict):
+            try:
+                return Task(
+                    **content,
+                    sessionId=session_id or content.get("sessionId"),
+                    metadata=metadata or content.get("metadata", {}),
+                    history=history,
+                )
+            except ValidationError:
+                pass
+
+        # 4) Fallback: treat whatever they returned as “artifact content”:
+        artifacts = self._normalize_content(content)
+        return Task(
+            id=task_id,
+            sessionId=session_id,
+            status=TaskStatus(state=self.default_status),
+            artifacts=artifacts,
+            metadata=metadata,
+            history=history,
+        )
+
+    def normalize_from_status(
+        self, status: TaskState, task_id: str, metadata: Dict[str,Any]
+    ) -> Task:
+        """Build a Task when only a cancellation or status‐only event occurs."""
+        return Task(
+            id=task_id,
+            sessionId="",
+            status=TaskStatus(state=status, timestamp=datetime.now()),
+            artifacts=[],
+            metadata=metadata,
+            history=[],
+        )
+
+    def _normalize_content(self, content: Any) -> List[Artifact]:
+        """Turn any handler return value into a list of Artifact."""
+        if isinstance(content, Artifact):
+            return [content]
+
+        if isinstance(content, list) and all(isinstance(a, Artifact) for a in content):
+            return content
+
+        if isinstance(content, list):
+            return [Artifact(parts=self._parts_from_mixed(content))]
+
+        if isinstance(content, str):
+            return [Artifact(parts=[TextPart(text=content)])]
+
+        if isinstance(content, dict):
+            # raw artifact dict
+            return [Artifact.model_validate(content)]
+
+        # explicit `Part` subclasses
+        if isinstance(content, (TextPart, FilePart, DataPart)):
+            return [Artifact(parts=[content])]
+
+        # “unknown” object: try Pydantic → dict → fallback to text
+        try:
+            return [Artifact.model_validate(content)]
+        except ValidationError:
+            return [Artifact(parts=[TextPart(text=str(content))])]
+
+    def _parts_from_mixed(self, items: List[Any]) -> List[Part]:
+        parts: List[Part] = []
+        for item in items:
+            if isinstance(item, Artifact):
+                parts.extend(item.parts)
+            else:
+                parts.append(self._create_part(item))
+        return parts
+
+    def _create_part(self, item: Any) -> Part:
+        from smarta2a.utils.types import Part as UnionPart
+        # guard against Union alias
+        if isinstance(item, (TextPart, FilePart, DataPart)):
+            return item
+        if isinstance(item, str):
+            return TextPart(text=item)
+        if isinstance(item, dict):
+            try:
+                return UnionPart.model_validate(item)
+            except ValidationError:
+                return TextPart(text=str(item))
+        return TextPart(text=str(item))

smarta2a/utils/task_request_builder.py

@@ -0,0 +1,114 @@
+# Library imports
+from typing import Any, Literal
+from uuid import uuid4
+
+# Local imports
+from smarta2a.utils.types import (
+    TaskPushNotificationConfig,
+    PushNotificationConfig,
+    TaskSendParams,
+    TextPart,
+    DataPart,
+    FilePart,
+    FileContent,
+    Message,
+    Part,
+    TaskQueryParams,
+    TaskIdParams,
+    GetTaskRequest,
+    CancelTaskRequest,
+    SetTaskPushNotificationRequest,
+    GetTaskPushNotificationRequest,
+    AuthenticationInfo,
+)
+
+class TaskRequestBuilder:
+    @staticmethod
+    def build_send_task_request(
+        *,
+        id: str,
+        role: Literal["user", "agent"] = "user",
+        text: str | None = None,
+        data: dict[str, Any] | None = None,
+        file_uri: str | None = None,
+        session_id: str | None = None,
+        accepted_output_modes: list[str] | None = None,
+        push_notification: PushNotificationConfig | None = None,
+        history_length: int | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> TaskSendParams:
+        parts: list[Part] = []
+
+        if text is not None:
+            parts.append(TextPart(text=text))
+
+        if data is not None:
+            parts.append(DataPart(data=data))
+
+        if file_uri is not None:
+            file_content = FileContent(uri=file_uri)
+            parts.append(FilePart(file=file_content))
+
+        message = Message(role=role, parts=parts)
+
+        return TaskSendParams(
+            id=id,
+            sessionId=session_id or uuid4().hex,
+            message=message,
+            acceptedOutputModes=accepted_output_modes,
+            pushNotification=push_notification,
+            historyLength=history_length,
+            metadata=metadata,
+        )
+    
+    @staticmethod
+    def get_task(
+        id: str,
+        history_length: int | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> GetTaskRequest:
+        params = TaskQueryParams(
+            id=id,
+            historyLength=history_length,
+            metadata=metadata,
+        )
+        return GetTaskRequest(params=params)
+
+    @staticmethod
+    def cancel_task(
+        id: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> CancelTaskRequest:
+        params = TaskIdParams(id=id, metadata=metadata)
+        return CancelTaskRequest(params=params)
+
+    @staticmethod
+    def set_push_notification(
+        id: str,
+        url: str,
+        token: str | None = None,
+        authentication: AuthenticationInfo | dict[str, Any] | None = None,
+    ) -> SetTaskPushNotificationRequest:
+        # allow passing AuthenticationInfo _or_ raw dict
+        auth = (
+            authentication
+            if isinstance(authentication, AuthenticationInfo)
+            else (AuthenticationInfo(**authentication) if authentication else None)
+        )
+        push_cfg = TaskPushNotificationConfig(
+            id=id,
+            pushNotificationConfig=PushNotificationConfig(
+                url=url,
+                token=token,
+                authentication=auth,
+            )
+        )
+        return SetTaskPushNotificationRequest(params=push_cfg)
+
+    @staticmethod
+    def get_push_notification(
+        id: str,
+        metadata: dict[str, Any] | None = None,
+    ) -> GetTaskPushNotificationRequest:
+        params = TaskIdParams(id=id, metadata=metadata)
+        return GetTaskPushNotificationRequest(params=params)