polyrob 0.5.0__py3-none-any.whl

api/a2a/endpoints.py

@@ -0,0 +1,375 @@
+"""A2A Protocol Endpoints - JSON-RPC over HTTP.
+
+Implements the A2A JSON-RPC 2.0 interface for task management.
+All methods follow the A2A specification naming convention.
+
+Supported Methods:
+- message/send: Create new task or continue existing
+- tasks/get: Get task status
+- tasks/list: List tasks with pagination
+- tasks/cancel: Cancel a task
+- tasks/resubscribe: Resubscribe to streaming events
+- tasks/pushNotificationConfig/*: Push notification management
+
+Reference: https://a2a-protocol.org/latest/specification/
+"""
+
+import logging
+from typing import Optional
+from fastapi import APIRouter, Request, HTTPException, Depends
+from fastapi.responses import JSONResponse
+
+from api.a2a.models import (
+    JSONRPCRequest, JSONRPCResponse, JSONRPCError, A2AErrorCode,
+    SendMessageRequest, GetTaskRequest, ListTasksRequest,
+    CancelTaskRequest, TaskResubscriptionRequest,
+    SetPushNotificationConfigRequest, PushNotificationConfig,
+    A2AMessage, A2ATask, ListTasksResponse
+)
+from api.a2a.task_handler import A2ATaskHandler
+from api.payment_verification import verify_payment_for_request, payment_required_response
+
+logger = logging.getLogger(__name__)
+router = APIRouter(prefix="/a2a", tags=["a2a"])
+
+# A2A-specific error code for payment required
+A2A_ERROR_PAYMENT_REQUIRED = -32402
+
+
+def get_task_handler(request: Request) -> A2ATaskHandler:
+    """Get A2A task handler from container."""
+    from core.container import DependencyContainer
+    container = DependencyContainer.get_instance()
+    return A2ATaskHandler(container)
+
+
+# Permissive auth policy (accepts x402, JWT, and API-key auth).
+# Delegates to the canonical implementation in api.dependencies.
+from api.dependencies import get_user_permissive as get_authenticated_user
+
+
+@router.post("/rpc")
+async def a2a_rpc_endpoint(
+    rpc_request: JSONRPCRequest,
+    request: Request,
+    handler: A2ATaskHandler = Depends(get_task_handler)
+):
+    """A2A JSON-RPC 2.0 endpoint.
+
+    Handles all A2A protocol methods via JSON-RPC.
+    """
+    request_id = rpc_request.id
+
+    try:
+        user_id = await get_authenticated_user(request)
+        method = rpc_request.method
+        params = rpc_request.params or {}
+
+        logger.info(f"A2A RPC: {method} from user {user_id}")
+
+        # Route to appropriate handler (pass request for payment verification)
+        result = await _handle_rpc_method(
+            method=method,
+            params=params,
+            user_id=user_id,
+            handler=handler,
+            request=request
+        )
+
+        return JSONRPCResponse(
+            id=request_id,
+            result=result.dict() if hasattr(result, 'dict') else result
+        )
+
+    except HTTPException as e:
+        # Handle 402 Payment Required specially for JSON-RPC
+        if e.status_code == 402:
+            return JSONRPCResponse(
+                id=request_id,
+                error=JSONRPCError(
+                    code=A2A_ERROR_PAYMENT_REQUIRED,
+                    message="Payment required",
+                    data=e.detail if isinstance(e.detail, dict) else {"message": e.detail}
+                )
+            )
+        return JSONRPCResponse(
+            id=request_id,
+            error=JSONRPCError(
+                code=A2AErrorCode.AUTHENTICATION_REQUIRED if e.status_code == 401 else A2AErrorCode.INTERNAL_ERROR,
+                message=e.detail
+            )
+        )
+    except ValueError as e:
+        return JSONRPCResponse(
+            id=request_id,
+            error=JSONRPCError(
+                code=A2AErrorCode.INVALID_PARAMS,
+                message=str(e)
+            )
+        )
+    except Exception as e:
+        logger.error(f"A2A RPC error: {e}", exc_info=True)
+        return JSONRPCResponse(
+            id=request_id,
+            error=JSONRPCError(
+                code=A2AErrorCode.INTERNAL_ERROR,
+                message="Internal error",
+                data=str(e)
+            )
+        )
+
+
+async def _handle_rpc_method(
+    method: str,
+    params: dict,
+    user_id: str,
+    handler: A2ATaskHandler,
+    request: Request
+):
+    """Handle individual RPC methods with payment verification.
+
+    Args:
+        method: RPC method name
+        params: Method parameters
+        user_id: Authenticated user
+        handler: A2A task handler
+        request: FastAPI request for payment verification
+
+    Returns:
+        Method result
+
+    Raises:
+        ValueError: Invalid params
+        HTTPException: Method not found or payment required
+    """
+    # Message operations
+    if method == "message/send":
+        message = A2AMessage(**params.get("message", {}))
+        config = params.get("configuration")
+        msg_request = SendMessageRequest(message=message, configuration=config)
+
+        # Check if this is a new task or continuation
+        if message.taskId:
+            # Continuation of existing task - already paid
+            return await handler.send_message(message.taskId, message, user_id)
+        else:
+            # NEW TASK: Verify payment before creation
+            try:
+                payment_method, details = await verify_payment_for_request(
+                    request=request,
+                    cost_credits=1  # Session creation = 1 credit
+                )
+                logger.info(f"A2A RPC task payment verified via {payment_method}: {details}")
+            except HTTPException as e:
+                if e.status_code == 402:
+                    # Return 402 with payment options
+                    raise HTTPException(
+                        status_code=402,
+                        detail=payment_required_response(request, cost_credits=1)
+                    )
+                raise
+
+            return await handler.create_task(msg_request, user_id, message.contextId)
+
+    # Task operations
+    elif method == "tasks/get":
+        task_id = params.get("id")
+        if not task_id:
+            raise ValueError("Missing task id")
+        history_length = params.get("historyLength")
+        return await handler.get_task(task_id, history_length, user_id=user_id)
+
+    elif method == "tasks/list":
+        context_id = params.get("contextId")
+        page_token = params.get("pageToken")
+        page_size = params.get("pageSize", 20)
+        tasks, next_token = await handler.list_tasks(
+            user_id, context_id, page_token, page_size
+        )
+        return ListTasksResponse(tasks=tasks, nextPageToken=next_token)
+
+    elif method == "tasks/cancel":
+        task_id = params.get("id")
+        if not task_id:
+            raise ValueError("Missing task id")
+        return await handler.cancel_task(task_id, user_id)
+
+    # Push notification operations
+    elif method == "tasks/pushNotificationConfig/set":
+        task_id = params.get("taskId")
+        config_data = params.get("config", {})
+        config = PushNotificationConfig(**config_data)
+        await handler.set_push_notification_config(task_id, config, user_id=user_id)
+        return {"success": True}
+
+    elif method == "tasks/pushNotificationConfig/get":
+        task_id = params.get("taskId")
+        config = await handler.get_push_notification_config(task_id, user_id=user_id)
+        return config.dict() if config else None
+
+    elif method == "tasks/pushNotificationConfig/delete":
+        task_id = params.get("taskId")
+        success = await handler.delete_push_notification_config(task_id, user_id=user_id)
+        return {"success": success}
+
+    else:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Method '{method}' not found"
+        )
+
+
+# =============================================================================
+@router.post("/tasks")
+async def create_task(
+    request_body: SendMessageRequest,
+    request: Request,
+    handler: A2ATaskHandler = Depends(get_task_handler)
+):
+    """Create a new A2A task with payment verification.
+
+    RESTful convenience endpoint for task creation.
+    Equivalent to message/send RPC method.
+    """
+    user_id = await get_authenticated_user(request)
+
+    # Verify payment for new task creation
+    try:
+        payment_method, details = await verify_payment_for_request(
+            request=request,
+            cost_credits=1
+        )
+        logger.info(f"A2A REST task payment verified via {payment_method}")
+    except HTTPException as e:
+        if e.status_code == 402:
+            return JSONResponse(
+                status_code=402,
+                content=payment_required_response(request, cost_credits=1)
+            )
+        raise
+
+    return await handler.create_task(
+        request_body,
+        user_id,
+        request_body.message.contextId
+    )
+
+
+@router.get("/tasks/{task_id}")
+async def get_task(
+    task_id: str,
+    historyLength: Optional[int] = None,
+    request: Request = None,
+    handler: A2ATaskHandler = Depends(get_task_handler)
+):
+    """Get A2A task status.
+
+    RESTful convenience endpoint for getting task status.
+    Equivalent to tasks/get RPC method.
+    """
+    # Authenticate and enforce ownership (was previously unauthenticated — IDOR).
+    user_id = await get_authenticated_user(request)
+    try:
+        return await handler.get_task(task_id, historyLength, user_id=user_id)
+    except ValueError as e:
+        raise HTTPException(status_code=404, detail=str(e))
+
+
+@router.post("/tasks/{task_id}/send")
+async def send_to_task(
+    task_id: str,
+    message: A2AMessage,
+    request: Request,
+    handler: A2ATaskHandler = Depends(get_task_handler)
+):
+    """Send a message to an existing A2A task.
+
+    RESTful convenience endpoint for continuing tasks.
+    Equivalent to message/send RPC method with taskId.
+    """
+    user_id = await get_authenticated_user(request)
+    message.taskId = task_id  # Ensure ID matches path
+    try:
+        return await handler.send_message(task_id, message, user_id)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+@router.post("/tasks/{task_id}/cancel")
+async def cancel_task(
+    task_id: str,
+    request: Request,
+    handler: A2ATaskHandler = Depends(get_task_handler)
+):
+    """Cancel an A2A task.
+
+    RESTful convenience endpoint for task cancellation.
+    Equivalent to tasks/cancel RPC method.
+    """
+    user_id = await get_authenticated_user(request)
+    try:
+        return await handler.cancel_task(task_id, user_id)
+    except ValueError as e:
+        raise HTTPException(status_code=400, detail=str(e))
+
+
+@router.get("/tasks")
+async def list_tasks(
+    contextId: Optional[str] = None,
+    pageToken: Optional[str] = None,
+    pageSize: int = 20,
+    request: Request = None,
+    handler: A2ATaskHandler = Depends(get_task_handler)
+):
+    """List A2A tasks with pagination.
+
+    RESTful convenience endpoint for listing tasks.
+    Equivalent to tasks/list RPC method.
+    """
+    user_id = await get_authenticated_user(request)
+    tasks, next_token = await handler.list_tasks(
+        user_id, contextId, pageToken, pageSize
+    )
+    return ListTasksResponse(tasks=tasks, nextPageToken=next_token)
+
+
+# =============================================================================
+@router.post("/tasks/{task_id}/push-config")
+async def set_push_config(
+    task_id: str,
+    config: PushNotificationConfig,
+    request: Request,
+    handler: A2ATaskHandler = Depends(get_task_handler)
+):
+    """Set push notification config for a task."""
+    user_id = await get_authenticated_user(request)
+    await handler.set_push_notification_config(task_id, config, user_id=user_id)
+    return {"success": True, "taskId": task_id}
+
+
+@router.get("/tasks/{task_id}/push-config")
+async def get_push_config(
+    task_id: str,
+    request: Request,
+    handler: A2ATaskHandler = Depends(get_task_handler)
+):
+    """Get push notification config for a task."""
+    user_id = await get_authenticated_user(request)
+    config = await handler.get_push_notification_config(task_id, user_id=user_id)
+    if not config:
+        raise HTTPException(status_code=404, detail="No push config found")
+    return config
+
+
+@router.delete("/tasks/{task_id}/push-config")
+async def delete_push_config(
+    task_id: str,
+    request: Request,
+    handler: A2ATaskHandler = Depends(get_task_handler)
+):
+    """Delete push notification config for a task."""
+    user_id = await get_authenticated_user(request)
+    success = await handler.delete_push_notification_config(task_id, user_id=user_id)
+    if not success:
+        raise HTTPException(status_code=404, detail="No push config found")
+    return {"success": True, "taskId": task_id}

api/a2a/models.py

@@ -0,0 +1,268 @@
+"""A2A Protocol Models - Pydantic models for A2A communication.
+
+Based on Google's A2A Protocol Specification:
+https://a2a-protocol.org/latest/specification/
+
+All models follow the official A2A schema definitions.
+"""
+
+from typing import Optional, List, Dict, Any, Union, Literal
+from pydantic import BaseModel, Field
+from enum import Enum
+from datetime import datetime
+import uuid
+
+
+class A2ATaskState(str, Enum):
+    """A2A Task lifecycle states.
+
+    Non-terminal (active) states:
+    - submitted: Initial state, typically transitions immediately to working
+    - working: Agent actively processing the task
+
+    Interrupted states (can resume):
+    - input-required: Awaiting client information
+    - auth-required: Awaiting external authentication
+
+    Terminal states (immutable):
+    - completed: Task finished successfully
+    - failed: Task processing failed
+    - canceled: User-initiated cancellation
+    - rejected: Agent declined execution
+    - unknown: Unknown/undefined terminal state
+    """
+    SUBMITTED = "submitted"
+    WORKING = "working"
+    INPUT_REQUIRED = "input-required"
+    AUTH_REQUIRED = "auth-required"
+    COMPLETED = "completed"
+    FAILED = "failed"
+    CANCELED = "canceled"
+    REJECTED = "rejected"
+    UNKNOWN = "unknown"
+
+    @classmethod
+    def is_terminal(cls, state: 'A2ATaskState') -> bool:
+        """Check if state is terminal (no further transitions allowed)."""
+        return state in [
+            cls.COMPLETED, cls.FAILED, cls.CANCELED,
+            cls.REJECTED, cls.UNKNOWN
+        ]
+
+    @classmethod
+    def is_interrupted(cls, state: 'A2ATaskState') -> bool:
+        """Check if state is interrupted (awaiting input)."""
+        return state in [cls.INPUT_REQUIRED, cls.AUTH_REQUIRED]
+
+
+class TextPart(BaseModel):
+    """Text content part."""
+    kind: Literal['text'] = 'text'
+    text: str
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class FileWithBytes(BaseModel):
+    """File content with inline bytes."""
+    bytes: str  # base64-encoded content
+    name: Optional[str] = None
+    mimeType: Optional[str] = None
+
+
+class FileWithUri(BaseModel):
+    """File content with URI reference."""
+    uri: str
+    name: Optional[str] = None
+    mimeType: Optional[str] = None
+
+
+class FilePart(BaseModel):
+    """File content part."""
+    kind: Literal['file'] = 'file'
+    file: Union[FileWithBytes, FileWithUri]
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class DataPart(BaseModel):
+    """Structured JSON data part."""
+    kind: Literal['data'] = 'data'
+    data: Dict[str, Any]
+    metadata: Optional[Dict[str, Any]] = None
+
+
+# Part union type
+Part = Union[TextPart, FilePart, DataPart]
+
+
+class A2AMessage(BaseModel):
+    """A2A Message - Communication payload.
+
+    Messages represent a single turn of communication between
+    user/client and agent.
+    """
+    messageId: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    role: Literal["user", "agent"]
+    parts: List[Dict[str, Any]]  # Mixed parts (text, file, data)
+    kind: Literal['message'] = 'message'
+    taskId: Optional[str] = None
+    contextId: Optional[str] = None
+    referenceTaskIds: Optional[List[str]] = None
+    extensions: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class A2AArtifact(BaseModel):
+    """A2A Artifact - Output produced by agent.
+
+    Artifacts represent outputs/results of task execution.
+    Can contain mixed content types via parts array.
+    """
+    artifactId: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    parts: List[Dict[str, Any]]
+    name: Optional[str] = None
+    description: Optional[str] = None
+    extensions: Optional[List[str]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+    # Streaming support
+    index: Optional[int] = None  # Index within owning task
+    append: Optional[bool] = None  # For streaming: append to existing?
+    lastChunk: Optional[bool] = None  # For streaming: final chunk?
+
+
+class A2ATaskStatus(BaseModel):
+    """Task status and related messages."""
+    state: A2ATaskState
+    message: Optional[A2AMessage] = None
+    timestamp: Optional[str] = Field(default_factory=lambda: datetime.now().isoformat())
+
+
+class A2ATask(BaseModel):
+    """A2A Task - Fundamental unit of work.
+
+    Tasks encapsulate the entire lifecycle of an agent interaction,
+    from initial request to final result.
+    """
+    id: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    contextId: str = Field(default_factory=lambda: str(uuid.uuid4()))
+    status: A2ATaskStatus
+    kind: Literal['task'] = 'task'
+    history: Optional[List[A2AMessage]] = None
+    artifacts: Optional[List[A2AArtifact]] = None
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class SendMessageRequest(BaseModel):
+    """Request to send a message (create or continue task)."""
+    message: A2AMessage
+    configuration: Optional[Dict[str, Any]] = None  # pushNotificationConfig, etc.
+
+
+class SendMessageResponse(BaseModel):
+    """Response from sending a message."""
+    task: Optional[A2ATask] = None
+    message: Optional[A2AMessage] = None
+
+
+class GetTaskRequest(BaseModel):
+    """Request to get task status."""
+    id: str
+    historyLength: Optional[int] = None  # None = all, 0 = none, >0 = last N
+    metadata: Optional[Dict[str, Any]] = None
+
+
+class ListTasksRequest(BaseModel):
+    """Request to list tasks."""
+    contextId: Optional[str] = None
+    pageToken: Optional[str] = None
+    pageSize: Optional[int] = 20
+    includeArtifacts: Optional[bool] = False
+
+
+class ListTasksResponse(BaseModel):
+    """Response for listing tasks."""
+    tasks: List[A2ATask]
+    nextPageToken: Optional[str] = None
+
+
+class CancelTaskRequest(BaseModel):
+    """Request to cancel a task."""
+    id: str
+
+
+class TaskResubscriptionRequest(BaseModel):
+    """Request to resubscribe to task events."""
+    id: str
+    historyLength: Optional[int] = None
+
+
+class AuthenticationInfo(BaseModel):
+    """Authentication info for webhook callbacks."""
+    schemes: List[str]  # e.g., ["Bearer"]
+    credentials: Optional[str] = None
+
+
+class PushNotificationConfig(BaseModel):
+    """Configuration for webhook push notifications."""
+    url: str  # HTTPS webhook endpoint
+    token: Optional[str] = None  # Client-generated validation token
+    authentication: Optional[AuthenticationInfo] = None
+
+
+class SetPushNotificationConfigRequest(BaseModel):
+    """Request to set push notification config."""
+    taskId: str
+    config: PushNotificationConfig
+
+
+class TaskStatusUpdateEvent(BaseModel):
+    """Task status update event in streaming response."""
+    task: A2ATask
+    final: bool = False  # True if this is the final update
+
+
+class TaskArtifactUpdateEvent(BaseModel):
+    """Artifact update event in streaming response."""
+    taskId: str
+    artifact: A2AArtifact
+
+
+class JSONRPCRequest(BaseModel):
+    """JSON-RPC 2.0 Request."""
+    jsonrpc: str = "2.0"
+    method: str
+    params: Optional[Dict[str, Any]] = None
+    id: Optional[Union[str, int]] = None
+
+
+class JSONRPCError(BaseModel):
+    """JSON-RPC 2.0 Error."""
+    code: int
+    message: str
+    data: Optional[Any] = None
+
+
+class JSONRPCResponse(BaseModel):
+    """JSON-RPC 2.0 Response."""
+    jsonrpc: str = "2.0"
+    result: Optional[Any] = None
+    error: Optional[JSONRPCError] = None
+    id: Optional[Union[str, int]] = None
+
+
+class A2AErrorCode:
+    """A2A-specific error codes (JSON-RPC server error range)."""
+    # Standard JSON-RPC errors
+    PARSE_ERROR = -32700
+    INVALID_REQUEST = -32600
+    METHOD_NOT_FOUND = -32601
+    INVALID_PARAMS = -32602
+    INTERNAL_ERROR = -32603
+
+    # A2A-specific errors (-32000 to -32099)
+    TASK_NOT_FOUND = -32000
+    TASK_NOT_CANCELABLE = -32001
+    PUSH_NOTIFICATION_NOT_SUPPORTED = -32002
+    UNSUPPORTED_OPERATION = -32003
+    AUTHENTICATION_REQUIRED = -32004
+    PAYMENT_REQUIRED = -32005