finish a2a
This commit is contained in:
zachary62
@@ -0,0 +1,4 @@
|
||||
from .client import A2AClient
|
||||
from .card_resolver import A2ACardResolver
|
||||
|
||||
__all__ = ["A2AClient", "A2ACardResolver"]
|
||||
@@ -0,0 +1,21 @@
|
||||
import httpx
|
||||
from common.types import (
|
||||
AgentCard,
|
||||
A2AClientJSONError,
|
||||
)
|
||||
import json
|
||||
|
||||
|
||||
class A2ACardResolver:
|
||||
def __init__(self, base_url, agent_card_path="/.well-known/agent.json"):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.agent_card_path = agent_card_path.lstrip("/")
|
||||
|
||||
def get_agent_card(self) -> AgentCard:
|
||||
with httpx.Client() as client:
|
||||
response = client.get(self.base_url + "/" + self.agent_card_path)
|
||||
response.raise_for_status()
|
||||
try:
|
||||
return AgentCard(**response.json())
|
||||
except json.JSONDecodeError as e:
|
||||
raise A2AClientJSONError(str(e)) from e
|
||||
@@ -0,0 +1,182 @@
|
||||
import httpx
|
||||
from httpx_sse import connect_sse
|
||||
from typing import Any, AsyncIterable
|
||||
from common.types import (
|
||||
AgentCard,
|
||||
GetTaskRequest,
|
||||
SendTaskRequest,
|
||||
SendTaskResponse,
|
||||
JSONRPCRequest,
|
||||
JSONRPCResponse,
|
||||
JSONRPCError,
|
||||
GetTaskResponse,
|
||||
CancelTaskResponse,
|
||||
CancelTaskRequest,
|
||||
SetTaskPushNotificationRequest,
|
||||
SetTaskPushNotificationResponse,
|
||||
GetTaskPushNotificationRequest,
|
||||
GetTaskPushNotificationResponse,
|
||||
A2AClientHTTPError,
|
||||
A2AClientJSONError,
|
||||
SendTaskStreamingRequest,
|
||||
SendTaskStreamingResponse,
|
||||
Task,
|
||||
TaskPushNotificationConfig,
|
||||
TaskStatusUpdateEvent,
|
||||
TaskArtifactUpdateEvent,
|
||||
)
|
||||
import json
|
||||
import logging
|
||||
|
||||
# Configure a logger specific to the client
|
||||
logger = logging.getLogger("A2AClient")
|
||||
|
||||
class A2AClientError(Exception):
|
||||
"""Base class for A2A client errors"""
|
||||
def __init__(self, message):
|
||||
super().__init__(message)
|
||||
|
||||
class RpcError(Exception):
|
||||
code: int
|
||||
data: Any = None
|
||||
def __init__(self, code: int, message: str, data: Any = None):
|
||||
super().__init__(message)
|
||||
self.name = "RpcError"
|
||||
self.code = code
|
||||
self.data = data
|
||||
|
||||
class A2AClient:
|
||||
def __init__(self, agent_card: AgentCard = None, url: str = None):
|
||||
if agent_card:
|
||||
self.url = agent_card.url.rstrip("/")
|
||||
elif url:
|
||||
self.url = url.rstrip("/")
|
||||
else:
|
||||
raise ValueError("Must provide either agent_card or url")
|
||||
self.fetchImpl = httpx.AsyncClient(timeout=None)
|
||||
|
||||
def _generateRequestId(self):
|
||||
import time
|
||||
return int(time.time() * 1000)
|
||||
|
||||
async def _send_request(self, request: JSONRPCRequest) -> dict[str, Any]:
|
||||
req_id = request.id
|
||||
req_method = request.method
|
||||
req_dump = request.model_dump(exclude_none=True)
|
||||
|
||||
logger.info(f"-> Sending Request (ID: {req_id}, Method: {req_method}):\n{json.dumps(req_dump, indent=2)}")
|
||||
|
||||
try:
|
||||
response = await self.fetchImpl.post(
|
||||
self.url, json=req_dump, timeout=60.0
|
||||
)
|
||||
logger.info(f"<- Received HTTP Status {response.status_code} for Request (ID: {req_id})")
|
||||
response_text = await response.aread()
|
||||
logger.debug(f"Raw Response Body (ID: {req_id}):\n{response_text.decode('utf-8', errors='replace')}")
|
||||
|
||||
response.raise_for_status()
|
||||
|
||||
try:
|
||||
json_response = json.loads(response_text)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Failed to decode JSON response (ID: {req_id}): {e}")
|
||||
raise A2AClientJSONError(f"Failed to decode JSON: {e}") from e
|
||||
|
||||
if "error" in json_response and json_response["error"] is not None:
|
||||
rpc_error = json_response["error"]
|
||||
logger.warning(f"<- Received JSON-RPC Error (ID: {req_id}): Code={rpc_error.get('code')}, Msg='{rpc_error.get('message')}'")
|
||||
raise RpcError(rpc_error.get("code", -32000), rpc_error.get("message", "Unknown RPC Error"), rpc_error.get("data"))
|
||||
|
||||
logger.info(f"<- Received Success Response (ID: {req_id}):\n{json.dumps(json_response, indent=2)}")
|
||||
return json_response
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.error(f"HTTP Error for Request (ID: {req_id}): {e.response.status_code} - {e.request.url}")
|
||||
error_body = await e.response.aread()
|
||||
raise A2AClientHTTPError(e.response.status_code, f"{e}. Body: {error_body.decode('utf-8', errors='replace')}") from e
|
||||
except httpx.RequestError as e:
|
||||
logger.error(f"Request Error for (ID: {req_id}): {e}")
|
||||
raise A2AClientError(f"Network or request error: {e}") from e
|
||||
except RpcError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error during request (ID: {req_id}): {e}", exc_info=True)
|
||||
raise A2AClientError(f"Unexpected error: {e}") from e
|
||||
|
||||
async def send_task(self, payload: dict[str, Any]) -> SendTaskResponse:
|
||||
request = SendTaskRequest(params=payload)
|
||||
response_dict = await self._send_request(request)
|
||||
return SendTaskResponse(**response_dict)
|
||||
|
||||
async def send_task_streaming(
|
||||
self, payload: dict[str, Any]
|
||||
) -> AsyncIterable[SendTaskStreamingResponse]:
|
||||
request = SendTaskStreamingRequest(params=payload)
|
||||
req_id = request.id
|
||||
req_dump = request.model_dump(exclude_none=True)
|
||||
|
||||
logger.info(f"-> Sending Streaming Request (ID: {req_id}, Method: {request.method}):\n{json.dumps(req_dump, indent=2)}")
|
||||
|
||||
try:
|
||||
async with self.fetchImpl.stream("POST", self.url, json=req_dump, timeout=None) as response:
|
||||
logger.info(f"<- Received HTTP Status {response.status_code} for Streaming Request (ID: {req_id})")
|
||||
response.raise_for_status()
|
||||
|
||||
buffer = ""
|
||||
async for line in response.aiter_lines():
|
||||
if not line:
|
||||
if buffer.startswith("data:"):
|
||||
data_str = buffer[len("data:"):].strip()
|
||||
logger.debug(f"Received SSE Data Line (ID: {req_id}): {data_str}")
|
||||
try:
|
||||
sse_data_dict = json.loads(data_str)
|
||||
yield SendTaskStreamingResponse(**sse_data_dict)
|
||||
except json.JSONDecodeError as e:
|
||||
logger.error(f"Failed to decode SSE JSON (ID: {req_id}): {e}. Data: '{data_str}'")
|
||||
except Exception as e:
|
||||
logger.error(f"Error processing SSE data (ID: {req_id}): {e}. Data: '{data_str}'", exc_info=True)
|
||||
elif buffer:
|
||||
logger.debug(f"Received non-data SSE line (ID: {req_id}): {buffer}")
|
||||
buffer = ""
|
||||
else:
|
||||
buffer += line + "\n"
|
||||
|
||||
if buffer:
|
||||
logger.warning(f"SSE stream ended with partial data in buffer (ID: {req_id}): {buffer}")
|
||||
|
||||
logger.info(f"SSE Stream ended for request ID: {req_id}")
|
||||
|
||||
except httpx.HTTPStatusError as e:
|
||||
logger.error(f"HTTP Error during streaming connection (ID: {req_id}): {e.response.status_code} - {e.request.url}")
|
||||
error_body = await e.response.aread()
|
||||
raise A2AClientHTTPError(e.response.status_code, f"{e}. Body: {error_body.decode('utf-8', errors='replace')}") from e
|
||||
except httpx.RequestError as e:
|
||||
logger.error(f"Request Error during streaming (ID: {req_id}): {e}")
|
||||
raise A2AClientError(f"Network or request error during streaming: {e}") from e
|
||||
except Exception as e:
|
||||
logger.error(f"Unexpected error during streaming (ID: {req_id}): {e}", exc_info=True)
|
||||
raise A2AClientError(f"Unexpected streaming error: {e}") from e
|
||||
|
||||
async def get_task(self, payload: dict[str, Any]) -> GetTaskResponse:
|
||||
request = GetTaskRequest(params=payload)
|
||||
response_dict = await self._send_request(request)
|
||||
return GetTaskResponse(**response_dict)
|
||||
|
||||
async def cancel_task(self, payload: dict[str, Any]) -> CancelTaskResponse:
|
||||
request = CancelTaskRequest(params=payload)
|
||||
response_dict = await self._send_request(request)
|
||||
return CancelTaskResponse(**response_dict)
|
||||
|
||||
async def set_task_callback(
|
||||
self, payload: dict[str, Any]
|
||||
) -> SetTaskPushNotificationResponse:
|
||||
request = SetTaskPushNotificationRequest(params=payload)
|
||||
response_dict = await self._send_request(request)
|
||||
return SetTaskPushNotificationResponse(**response_dict)
|
||||
|
||||
async def get_task_callback(
|
||||
self, payload: dict[str, Any]
|
||||
) -> GetTaskPushNotificationResponse:
|
||||
request = GetTaskPushNotificationRequest(params=payload)
|
||||
response_dict = await self._send_request(request)
|
||||
return GetTaskPushNotificationResponse(**response_dict)
|
||||
@@ -0,0 +1,4 @@
|
||||
from .server import A2AServer
|
||||
from .task_manager import TaskManager, InMemoryTaskManager
|
||||
|
||||
__all__ = ["A2AServer", "TaskManager", "InMemoryTaskManager"]
|
||||
@@ -0,0 +1,168 @@
|
||||
from starlette.applications import Starlette
|
||||
from starlette.responses import JSONResponse
|
||||
from sse_starlette.sse import EventSourceResponse
|
||||
from starlette.requests import Request
|
||||
from common.types import (
|
||||
A2ARequest,
|
||||
JSONRPCResponse,
|
||||
InvalidRequestError,
|
||||
JSONParseError,
|
||||
GetTaskRequest,
|
||||
CancelTaskRequest,
|
||||
SendTaskRequest,
|
||||
SetTaskPushNotificationRequest,
|
||||
GetTaskPushNotificationRequest,
|
||||
InternalError,
|
||||
AgentCard,
|
||||
TaskResubscriptionRequest,
|
||||
SendTaskStreamingRequest,
|
||||
Message,
|
||||
)
|
||||
from pydantic import ValidationError
|
||||
import json
|
||||
from typing import AsyncIterable, Any
|
||||
from common.server.task_manager import TaskManager
|
||||
|
||||
import logging
|
||||
|
||||
# Configure a logger specific to the server
|
||||
logger = logging.getLogger("A2AServer")
|
||||
|
||||
|
||||
class A2AServer:
|
||||
def __init__(
|
||||
self,
|
||||
host="0.0.0.0",
|
||||
port=5000,
|
||||
endpoint="/",
|
||||
agent_card: AgentCard = None,
|
||||
task_manager: TaskManager = None,
|
||||
):
|
||||
self.host = host
|
||||
self.port = port
|
||||
self.endpoint = endpoint
|
||||
self.task_manager = task_manager
|
||||
self.agent_card = agent_card
|
||||
self.app = Starlette()
|
||||
self.app.add_route(self.endpoint, self._process_request, methods=["POST"])
|
||||
self.app.add_route(
|
||||
"/.well-known/agent.json", self._get_agent_card, methods=["GET"]
|
||||
)
|
||||
|
||||
def start(self):
|
||||
if self.agent_card is None:
|
||||
raise ValueError("agent_card is not defined")
|
||||
|
||||
if self.task_manager is None:
|
||||
raise ValueError("request_handler is not defined")
|
||||
|
||||
import uvicorn
|
||||
|
||||
# Basic logging config moved to __main__.py for application-level control
|
||||
uvicorn.run(self.app, host=self.host, port=self.port)
|
||||
|
||||
def _get_agent_card(self, request: Request) -> JSONResponse:
|
||||
logger.info("Serving Agent Card request")
|
||||
return JSONResponse(self.agent_card.model_dump(exclude_none=True))
|
||||
|
||||
async def _process_request(self, request: Request):
|
||||
request_id_for_log = "N/A" # Default if parsing fails early
|
||||
raw_body = b""
|
||||
try:
|
||||
# Log raw body first
|
||||
raw_body = await request.body()
|
||||
body = json.loads(raw_body) # Attempt parsing
|
||||
request_id_for_log = body.get("id", "N/A") # Get ID if possible
|
||||
logger.info(f"<- Received Request (ID: {request_id_for_log}):\n{json.dumps(body, indent=2)}")
|
||||
|
||||
json_rpc_request = A2ARequest.validate_python(body)
|
||||
|
||||
# Route based on method (same as before)
|
||||
if isinstance(json_rpc_request, GetTaskRequest):
|
||||
result = await self.task_manager.on_get_task(json_rpc_request)
|
||||
elif isinstance(json_rpc_request, SendTaskRequest):
|
||||
result = await self.task_manager.on_send_task(json_rpc_request)
|
||||
elif isinstance(json_rpc_request, SendTaskStreamingRequest):
|
||||
result = await self.task_manager.on_send_task_subscribe(
|
||||
json_rpc_request
|
||||
)
|
||||
elif isinstance(json_rpc_request, CancelTaskRequest):
|
||||
result = await self.task_manager.on_cancel_task(json_rpc_request)
|
||||
elif isinstance(json_rpc_request, SetTaskPushNotificationRequest):
|
||||
result = await self.task_manager.on_set_task_push_notification(json_rpc_request)
|
||||
elif isinstance(json_rpc_request, GetTaskPushNotificationRequest):
|
||||
result = await self.task_manager.on_get_task_push_notification(json_rpc_request)
|
||||
elif isinstance(json_rpc_request, TaskResubscriptionRequest):
|
||||
result = await self.task_manager.on_resubscribe_to_task(
|
||||
json_rpc_request
|
||||
)
|
||||
else:
|
||||
logger.warning(f"Unexpected request type: {type(json_rpc_request)}")
|
||||
raise ValueError(f"Unexpected request type: {type(request)}")
|
||||
|
||||
return self._create_response(result) # Pass result to response creation
|
||||
|
||||
except json.decoder.JSONDecodeError as e:
|
||||
logger.error(f"JSON Parse Error for Request body: <<<{raw_body.decode('utf-8', errors='replace')}>>>\nError: {e}")
|
||||
return self._handle_exception(e, request_id_for_log) # Pass ID if known
|
||||
except ValidationError as e:
|
||||
logger.error(f"Request Validation Error (ID: {request_id_for_log}): {e.json()}")
|
||||
return self._handle_exception(e, request_id_for_log)
|
||||
except Exception as e:
|
||||
logger.error(f"Unhandled Exception processing request (ID: {request_id_for_log}): {e}", exc_info=True)
|
||||
return self._handle_exception(e, request_id_for_log) # Pass ID if known
|
||||
|
||||
def _handle_exception(self, e: Exception, req_id=None) -> JSONResponse: # Accept req_id
|
||||
if isinstance(e, json.decoder.JSONDecodeError):
|
||||
json_rpc_error = JSONParseError()
|
||||
elif isinstance(e, ValidationError):
|
||||
json_rpc_error = InvalidRequestError(data=json.loads(e.json()))
|
||||
else:
|
||||
# Log the full exception details
|
||||
logger.error(f"Internal Server Error (ReqID: {req_id}): {e}", exc_info=True)
|
||||
json_rpc_error = InternalError(message=f"Internal Server Error: {type(e).__name__}")
|
||||
|
||||
response = JSONRPCResponse(id=req_id, error=json_rpc_error)
|
||||
response_dump = response.model_dump(exclude_none=True)
|
||||
logger.info(f"-> Sending Error Response (ReqID: {req_id}):\n{json.dumps(response_dump, indent=2)}")
|
||||
# A2A errors are still sent with HTTP 200
|
||||
return JSONResponse(response_dump, status_code=200)
|
||||
|
||||
def _create_response(self, result: Any) -> JSONResponse | EventSourceResponse:
|
||||
if isinstance(result, AsyncIterable):
|
||||
# Streaming response
|
||||
async def event_generator(result_stream) -> AsyncIterable[dict[str, str]]:
|
||||
stream_request_id = None # Capture ID from the first event if possible
|
||||
try:
|
||||
async for item in result_stream:
|
||||
# Log each streamed item
|
||||
response_json = item.model_dump_json(exclude_none=True)
|
||||
stream_request_id = item.id # Update ID
|
||||
logger.info(f"-> Sending SSE Event (ID: {stream_request_id}):\n{json.dumps(json.loads(response_json), indent=2)}")
|
||||
yield {"data": response_json}
|
||||
logger.info(f"SSE Stream ended for request ID: {stream_request_id}")
|
||||
except Exception as e:
|
||||
logger.error(f"Error during SSE generation (ReqID: {stream_request_id}): {e}", exc_info=True)
|
||||
# Optionally yield an error event if the protocol allows/requires it
|
||||
# error_payload = JSONRPCResponse(id=stream_request_id, error=InternalError(message=f"SSE Error: {e}"))
|
||||
# yield {"data": error_payload.model_dump_json(exclude_none=True)}
|
||||
|
||||
logger.info("Starting SSE stream...") # Log stream start
|
||||
return EventSourceResponse(event_generator(result))
|
||||
elif isinstance(result, JSONRPCResponse):
|
||||
# Standard JSON response
|
||||
response_dump = result.model_dump(exclude_none=True)
|
||||
log_id = result.id if result.id is not None else "N/A (Notification?)"
|
||||
log_prefix = "->"
|
||||
log_type = "Response"
|
||||
if result.error:
|
||||
log_prefix = "-> Sending Error"
|
||||
log_type = "Error Response"
|
||||
|
||||
logger.info(f"{log_prefix} {log_type} (ID: {log_id}):\n{json.dumps(response_dump, indent=2)}")
|
||||
return JSONResponse(response_dump)
|
||||
else:
|
||||
# This should ideally not happen if task manager returns correctly
|
||||
logger.error(f"Task manager returned unexpected type: {type(result)}")
|
||||
err_resp = JSONRPCResponse(id=None, error=InternalError(message="Invalid internal response type"))
|
||||
return JSONResponse(err_resp.model_dump(exclude_none=True), status_code=500)
|
||||
@@ -0,0 +1,277 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Union, AsyncIterable, List
|
||||
from common.types import Task
|
||||
from common.types import (
|
||||
JSONRPCResponse,
|
||||
TaskIdParams,
|
||||
TaskQueryParams,
|
||||
GetTaskRequest,
|
||||
TaskNotFoundError,
|
||||
SendTaskRequest,
|
||||
CancelTaskRequest,
|
||||
TaskNotCancelableError,
|
||||
SetTaskPushNotificationRequest,
|
||||
GetTaskPushNotificationRequest,
|
||||
GetTaskResponse,
|
||||
CancelTaskResponse,
|
||||
SendTaskResponse,
|
||||
SetTaskPushNotificationResponse,
|
||||
GetTaskPushNotificationResponse,
|
||||
PushNotificationNotSupportedError,
|
||||
TaskSendParams,
|
||||
TaskStatus,
|
||||
TaskState,
|
||||
TaskResubscriptionRequest,
|
||||
SendTaskStreamingRequest,
|
||||
SendTaskStreamingResponse,
|
||||
Artifact,
|
||||
PushNotificationConfig,
|
||||
TaskStatusUpdateEvent,
|
||||
JSONRPCError,
|
||||
TaskPushNotificationConfig,
|
||||
InternalError,
|
||||
)
|
||||
from common.server.utils import new_not_implemented_error
|
||||
import asyncio
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
class TaskManager(ABC):
|
||||
@abstractmethod
|
||||
async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def on_send_task_subscribe(
|
||||
self, request: SendTaskStreamingRequest
|
||||
) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def on_set_task_push_notification(
|
||||
self, request: SetTaskPushNotificationRequest
|
||||
) -> SetTaskPushNotificationResponse:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def on_get_task_push_notification(
|
||||
self, request: GetTaskPushNotificationRequest
|
||||
) -> GetTaskPushNotificationResponse:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def on_resubscribe_to_task(
|
||||
self, request: TaskResubscriptionRequest
|
||||
) -> Union[AsyncIterable[SendTaskResponse], JSONRPCResponse]:
|
||||
pass
|
||||
|
||||
|
||||
class InMemoryTaskManager(TaskManager):
|
||||
def __init__(self):
|
||||
self.tasks: dict[str, Task] = {}
|
||||
self.push_notification_infos: dict[str, PushNotificationConfig] = {}
|
||||
self.lock = asyncio.Lock()
|
||||
self.task_sse_subscribers: dict[str, List[asyncio.Queue]] = {}
|
||||
self.subscriber_lock = asyncio.Lock()
|
||||
|
||||
async def on_get_task(self, request: GetTaskRequest) -> GetTaskResponse:
|
||||
logger.info(f"Getting task {request.params.id}")
|
||||
task_query_params: TaskQueryParams = request.params
|
||||
|
||||
async with self.lock:
|
||||
task = self.tasks.get(task_query_params.id)
|
||||
if task is None:
|
||||
return GetTaskResponse(id=request.id, error=TaskNotFoundError())
|
||||
|
||||
task_result = self.append_task_history(
|
||||
task, task_query_params.historyLength
|
||||
)
|
||||
|
||||
return GetTaskResponse(id=request.id, result=task_result)
|
||||
|
||||
async def on_cancel_task(self, request: CancelTaskRequest) -> CancelTaskResponse:
|
||||
logger.info(f"Cancelling task {request.params.id}")
|
||||
task_id_params: TaskIdParams = request.params
|
||||
|
||||
async with self.lock:
|
||||
task = self.tasks.get(task_id_params.id)
|
||||
if task is None:
|
||||
return CancelTaskResponse(id=request.id, error=TaskNotFoundError())
|
||||
|
||||
return CancelTaskResponse(id=request.id, error=TaskNotCancelableError())
|
||||
|
||||
@abstractmethod
|
||||
async def on_send_task(self, request: SendTaskRequest) -> SendTaskResponse:
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
async def on_send_task_subscribe(
|
||||
self, request: SendTaskStreamingRequest
|
||||
) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
|
||||
pass
|
||||
|
||||
async def set_push_notification_info(self, task_id: str, notification_config: PushNotificationConfig):
|
||||
async with self.lock:
|
||||
task = self.tasks.get(task_id)
|
||||
if task is None:
|
||||
raise ValueError(f"Task not found for {task_id}")
|
||||
|
||||
self.push_notification_infos[task_id] = notification_config
|
||||
|
||||
return
|
||||
|
||||
async def get_push_notification_info(self, task_id: str) -> PushNotificationConfig:
|
||||
async with self.lock:
|
||||
task = self.tasks.get(task_id)
|
||||
if task is None:
|
||||
raise ValueError(f"Task not found for {task_id}")
|
||||
|
||||
return self.push_notification_infos[task_id]
|
||||
|
||||
return
|
||||
|
||||
async def has_push_notification_info(self, task_id: str) -> bool:
|
||||
async with self.lock:
|
||||
return task_id in self.push_notification_infos
|
||||
|
||||
|
||||
async def on_set_task_push_notification(
|
||||
self, request: SetTaskPushNotificationRequest
|
||||
) -> SetTaskPushNotificationResponse:
|
||||
logger.info(f"Setting task push notification {request.params.id}")
|
||||
task_notification_params: TaskPushNotificationConfig = request.params
|
||||
|
||||
try:
|
||||
await self.set_push_notification_info(task_notification_params.id, task_notification_params.pushNotificationConfig)
|
||||
except Exception as e:
|
||||
logger.error(f"Error while setting push notification info: {e}")
|
||||
return JSONRPCResponse(
|
||||
id=request.id,
|
||||
error=InternalError(
|
||||
message="An error occurred while setting push notification info"
|
||||
),
|
||||
)
|
||||
|
||||
return SetTaskPushNotificationResponse(id=request.id, result=task_notification_params)
|
||||
|
||||
async def on_get_task_push_notification(
|
||||
self, request: GetTaskPushNotificationRequest
|
||||
) -> GetTaskPushNotificationResponse:
|
||||
logger.info(f"Getting task push notification {request.params.id}")
|
||||
task_params: TaskIdParams = request.params
|
||||
|
||||
try:
|
||||
notification_info = await self.get_push_notification_info(task_params.id)
|
||||
except Exception as e:
|
||||
logger.error(f"Error while getting push notification info: {e}")
|
||||
return GetTaskPushNotificationResponse(
|
||||
id=request.id,
|
||||
error=InternalError(
|
||||
message="An error occurred while getting push notification info"
|
||||
),
|
||||
)
|
||||
|
||||
return GetTaskPushNotificationResponse(id=request.id, result=TaskPushNotificationConfig(id=task_params.id, pushNotificationConfig=notification_info))
|
||||
|
||||
async def upsert_task(self, task_send_params: TaskSendParams) -> Task:
|
||||
logger.info(f"Upserting task {task_send_params.id}")
|
||||
async with self.lock:
|
||||
task = self.tasks.get(task_send_params.id)
|
||||
if task is None:
|
||||
task = Task(
|
||||
id=task_send_params.id,
|
||||
sessionId = task_send_params.sessionId,
|
||||
messages=[task_send_params.message],
|
||||
status=TaskStatus(state=TaskState.SUBMITTED),
|
||||
history=[task_send_params.message],
|
||||
)
|
||||
self.tasks[task_send_params.id] = task
|
||||
else:
|
||||
task.history.append(task_send_params.message)
|
||||
|
||||
return task
|
||||
|
||||
async def on_resubscribe_to_task(
|
||||
self, request: TaskResubscriptionRequest
|
||||
) -> Union[AsyncIterable[SendTaskStreamingResponse], JSONRPCResponse]:
|
||||
return new_not_implemented_error(request.id)
|
||||
|
||||
async def update_store(
|
||||
self, task_id: str, status: TaskStatus, artifacts: list[Artifact]
|
||||
) -> Task:
|
||||
async with self.lock:
|
||||
try:
|
||||
task = self.tasks[task_id]
|
||||
except KeyError:
|
||||
logger.error(f"Task {task_id} not found for updating the task")
|
||||
raise ValueError(f"Task {task_id} not found")
|
||||
|
||||
task.status = status
|
||||
|
||||
if status.message is not None:
|
||||
task.history.append(status.message)
|
||||
|
||||
if artifacts is not None:
|
||||
if task.artifacts is None:
|
||||
task.artifacts = []
|
||||
task.artifacts.extend(artifacts)
|
||||
|
||||
return task
|
||||
|
||||
def append_task_history(self, task: Task, historyLength: int | None):
|
||||
new_task = task.model_copy()
|
||||
if historyLength is not None and historyLength > 0:
|
||||
new_task.history = new_task.history[-historyLength:]
|
||||
else:
|
||||
new_task.history = []
|
||||
|
||||
return new_task
|
||||
|
||||
async def setup_sse_consumer(self, task_id: str, is_resubscribe: bool = False):
|
||||
async with self.subscriber_lock:
|
||||
if task_id not in self.task_sse_subscribers:
|
||||
if is_resubscribe:
|
||||
raise ValueError("Task not found for resubscription")
|
||||
else:
|
||||
self.task_sse_subscribers[task_id] = []
|
||||
|
||||
sse_event_queue = asyncio.Queue(maxsize=0) # <=0 is unlimited
|
||||
self.task_sse_subscribers[task_id].append(sse_event_queue)
|
||||
return sse_event_queue
|
||||
|
||||
async def enqueue_events_for_sse(self, task_id, task_update_event):
|
||||
async with self.subscriber_lock:
|
||||
if task_id not in self.task_sse_subscribers:
|
||||
return
|
||||
|
||||
current_subscribers = self.task_sse_subscribers[task_id]
|
||||
for subscriber in current_subscribers:
|
||||
await subscriber.put(task_update_event)
|
||||
|
||||
async def dequeue_events_for_sse(
|
||||
self, request_id, task_id, sse_event_queue: asyncio.Queue
|
||||
) -> AsyncIterable[SendTaskStreamingResponse] | JSONRPCResponse:
|
||||
try:
|
||||
while True:
|
||||
event = await sse_event_queue.get()
|
||||
if isinstance(event, JSONRPCError):
|
||||
yield SendTaskStreamingResponse(id=request_id, error=event)
|
||||
break
|
||||
|
||||
yield SendTaskStreamingResponse(id=request_id, result=event)
|
||||
if isinstance(event, TaskStatusUpdateEvent) and event.final:
|
||||
break
|
||||
finally:
|
||||
async with self.subscriber_lock:
|
||||
if task_id in self.task_sse_subscribers:
|
||||
self.task_sse_subscribers[task_id].remove(sse_event_queue)
|
||||
|
||||
@@ -0,0 +1,28 @@
|
||||
from common.types import (
|
||||
JSONRPCResponse,
|
||||
ContentTypeNotSupportedError,
|
||||
UnsupportedOperationError,
|
||||
)
|
||||
from typing import List
|
||||
|
||||
|
||||
def are_modalities_compatible(
|
||||
server_output_modes: List[str], client_output_modes: List[str]
|
||||
):
|
||||
"""Modalities are compatible if they are both non-empty
|
||||
and there is at least one common element."""
|
||||
if client_output_modes is None or len(client_output_modes) == 0:
|
||||
return True
|
||||
|
||||
if server_output_modes is None or len(server_output_modes) == 0:
|
||||
return True
|
||||
|
||||
return any(x in server_output_modes for x in client_output_modes)
|
||||
|
||||
|
||||
def new_incompatible_types_error(request_id):
|
||||
return JSONRPCResponse(id=request_id, error=ContentTypeNotSupportedError())
|
||||
|
||||
|
||||
def new_not_implemented_error(request_id):
|
||||
return JSONRPCResponse(id=request_id, error=UnsupportedOperationError())
|
||||
@@ -0,0 +1,365 @@
|
||||
from typing import Union, Any
|
||||
from pydantic import BaseModel, Field, TypeAdapter
|
||||
from typing import Literal, List, Annotated, Optional
|
||||
from datetime import datetime
|
||||
from pydantic import model_validator, ConfigDict, field_serializer
|
||||
from uuid import uuid4
|
||||
from enum import Enum
|
||||
from typing_extensions import Self
|
||||
|
||||
|
||||
class TaskState(str, Enum):
|
||||
SUBMITTED = "submitted"
|
||||
WORKING = "working"
|
||||
INPUT_REQUIRED = "input-required"
|
||||
COMPLETED = "completed"
|
||||
CANCELED = "canceled"
|
||||
FAILED = "failed"
|
||||
UNKNOWN = "unknown"
|
||||
|
||||
|
||||
class TextPart(BaseModel):
|
||||
type: Literal["text"] = "text"
|
||||
text: str
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class FileContent(BaseModel):
|
||||
name: str | None = None
|
||||
mimeType: str | None = None
|
||||
bytes: str | None = None
|
||||
uri: str | None = None
|
||||
|
||||
@model_validator(mode="after")
|
||||
def check_content(self) -> Self:
|
||||
if not (self.bytes or self.uri):
|
||||
raise ValueError("Either 'bytes' or 'uri' must be present in the file data")
|
||||
if self.bytes and self.uri:
|
||||
raise ValueError(
|
||||
"Only one of 'bytes' or 'uri' can be present in the file data"
|
||||
)
|
||||
return self
|
||||
|
||||
|
||||
class FilePart(BaseModel):
|
||||
type: Literal["file"] = "file"
|
||||
file: FileContent
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class DataPart(BaseModel):
|
||||
type: Literal["data"] = "data"
|
||||
data: dict[str, Any]
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
Part = Annotated[Union[TextPart, FilePart, DataPart], Field(discriminator="type")]
|
||||
|
||||
|
||||
class Message(BaseModel):
|
||||
role: Literal["user", "agent"]
|
||||
parts: List[Part]
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class TaskStatus(BaseModel):
|
||||
state: TaskState
|
||||
message: Message | None = None
|
||||
timestamp: datetime = Field(default_factory=datetime.now)
|
||||
|
||||
@field_serializer("timestamp")
|
||||
def serialize_dt(self, dt: datetime, _info):
|
||||
return dt.isoformat()
|
||||
|
||||
|
||||
class Artifact(BaseModel):
|
||||
name: str | None = None
|
||||
description: str | None = None
|
||||
parts: List[Part]
|
||||
metadata: dict[str, Any] | None = None
|
||||
index: int = 0
|
||||
append: bool | None = None
|
||||
lastChunk: bool | None = None
|
||||
|
||||
|
||||
class Task(BaseModel):
|
||||
id: str
|
||||
sessionId: str | None = None
|
||||
status: TaskStatus
|
||||
artifacts: List[Artifact] | None = None
|
||||
history: List[Message] | None = None
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class TaskStatusUpdateEvent(BaseModel):
|
||||
id: str
|
||||
status: TaskStatus
|
||||
final: bool = False
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class TaskArtifactUpdateEvent(BaseModel):
|
||||
id: str
|
||||
artifact: Artifact
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class AuthenticationInfo(BaseModel):
|
||||
model_config = ConfigDict(extra="allow")
|
||||
|
||||
schemes: List[str]
|
||||
credentials: str | None = None
|
||||
|
||||
|
||||
class PushNotificationConfig(BaseModel):
|
||||
url: str
|
||||
token: str | None = None
|
||||
authentication: AuthenticationInfo | None = None
|
||||
|
||||
|
||||
class TaskIdParams(BaseModel):
|
||||
id: str
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class TaskQueryParams(TaskIdParams):
|
||||
historyLength: int | None = None
|
||||
|
||||
|
||||
class TaskSendParams(BaseModel):
|
||||
id: str
|
||||
sessionId: str = Field(default_factory=lambda: uuid4().hex)
|
||||
message: Message
|
||||
acceptedOutputModes: Optional[List[str]] = None
|
||||
pushNotification: PushNotificationConfig | None = None
|
||||
historyLength: int | None = None
|
||||
metadata: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class TaskPushNotificationConfig(BaseModel):
|
||||
id: str
|
||||
pushNotificationConfig: PushNotificationConfig
|
||||
|
||||
|
||||
## RPC Messages
|
||||
|
||||
|
||||
class JSONRPCMessage(BaseModel):
|
||||
jsonrpc: Literal["2.0"] = "2.0"
|
||||
id: int | str | None = Field(default_factory=lambda: uuid4().hex)
|
||||
|
||||
|
||||
class JSONRPCRequest(JSONRPCMessage):
|
||||
method: str
|
||||
params: dict[str, Any] | None = None
|
||||
|
||||
|
||||
class JSONRPCError(BaseModel):
|
||||
code: int
|
||||
message: str
|
||||
data: Any | None = None
|
||||
|
||||
|
||||
class JSONRPCResponse(JSONRPCMessage):
|
||||
result: Any | None = None
|
||||
error: JSONRPCError | None = None
|
||||
|
||||
|
||||
class SendTaskRequest(JSONRPCRequest):
|
||||
method: Literal["tasks/send"] = "tasks/send"
|
||||
params: TaskSendParams
|
||||
|
||||
|
||||
class SendTaskResponse(JSONRPCResponse):
|
||||
result: Task | None = None
|
||||
|
||||
|
||||
class SendTaskStreamingRequest(JSONRPCRequest):
|
||||
method: Literal["tasks/sendSubscribe"] = "tasks/sendSubscribe"
|
||||
params: TaskSendParams
|
||||
|
||||
|
||||
class SendTaskStreamingResponse(JSONRPCResponse):
|
||||
result: TaskStatusUpdateEvent | TaskArtifactUpdateEvent | None = None
|
||||
|
||||
|
||||
class GetTaskRequest(JSONRPCRequest):
|
||||
method: Literal["tasks/get"] = "tasks/get"
|
||||
params: TaskQueryParams
|
||||
|
||||
|
||||
class GetTaskResponse(JSONRPCResponse):
|
||||
result: Task | None = None
|
||||
|
||||
|
||||
class CancelTaskRequest(JSONRPCRequest):
|
||||
method: Literal["tasks/cancel",] = "tasks/cancel"
|
||||
params: TaskIdParams
|
||||
|
||||
|
||||
class CancelTaskResponse(JSONRPCResponse):
|
||||
result: Task | None = None
|
||||
|
||||
|
||||
class SetTaskPushNotificationRequest(JSONRPCRequest):
|
||||
method: Literal["tasks/pushNotification/set",] = "tasks/pushNotification/set"
|
||||
params: TaskPushNotificationConfig
|
||||
|
||||
|
||||
class SetTaskPushNotificationResponse(JSONRPCResponse):
|
||||
result: TaskPushNotificationConfig | None = None
|
||||
|
||||
|
||||
class GetTaskPushNotificationRequest(JSONRPCRequest):
|
||||
method: Literal["tasks/pushNotification/get",] = "tasks/pushNotification/get"
|
||||
params: TaskIdParams
|
||||
|
||||
|
||||
class GetTaskPushNotificationResponse(JSONRPCResponse):
|
||||
result: TaskPushNotificationConfig | None = None
|
||||
|
||||
|
||||
class TaskResubscriptionRequest(JSONRPCRequest):
|
||||
method: Literal["tasks/resubscribe",] = "tasks/resubscribe"
|
||||
params: TaskIdParams
|
||||
|
||||
|
||||
A2ARequest = TypeAdapter(
|
||||
Annotated[
|
||||
Union[
|
||||
SendTaskRequest,
|
||||
GetTaskRequest,
|
||||
CancelTaskRequest,
|
||||
SetTaskPushNotificationRequest,
|
||||
GetTaskPushNotificationRequest,
|
||||
TaskResubscriptionRequest,
|
||||
SendTaskStreamingRequest,
|
||||
],
|
||||
Field(discriminator="method"),
|
||||
]
|
||||
)
|
||||
|
||||
## Error types
|
||||
|
||||
|
||||
class JSONParseError(JSONRPCError):
|
||||
code: int = -32700
|
||||
message: str = "Invalid JSON payload"
|
||||
data: Any | None = None
|
||||
|
||||
|
||||
class InvalidRequestError(JSONRPCError):
|
||||
code: int = -32600
|
||||
message: str = "Request payload validation error"
|
||||
data: Any | None = None
|
||||
|
||||
|
||||
class MethodNotFoundError(JSONRPCError):
|
||||
code: int = -32601
|
||||
message: str = "Method not found"
|
||||
data: None = None
|
||||
|
||||
|
||||
class InvalidParamsError(JSONRPCError):
|
||||
code: int = -32602
|
||||
message: str = "Invalid parameters"
|
||||
data: Any | None = None
|
||||
|
||||
|
||||
class InternalError(JSONRPCError):
|
||||
code: int = -32603
|
||||
message: str = "Internal error"
|
||||
data: Any | None = None
|
||||
|
||||
|
||||
class TaskNotFoundError(JSONRPCError):
|
||||
code: int = -32001
|
||||
message: str = "Task not found"
|
||||
data: None = None
|
||||
|
||||
|
||||
class TaskNotCancelableError(JSONRPCError):
|
||||
code: int = -32002
|
||||
message: str = "Task cannot be canceled"
|
||||
data: None = None
|
||||
|
||||
|
||||
class PushNotificationNotSupportedError(JSONRPCError):
|
||||
code: int = -32003
|
||||
message: str = "Push Notification is not supported"
|
||||
data: None = None
|
||||
|
||||
|
||||
class UnsupportedOperationError(JSONRPCError):
|
||||
code: int = -32004
|
||||
message: str = "This operation is not supported"
|
||||
data: None = None
|
||||
|
||||
|
||||
class ContentTypeNotSupportedError(JSONRPCError):
|
||||
code: int = -32005
|
||||
message: str = "Incompatible content types"
|
||||
data: None = None
|
||||
|
||||
|
||||
class AgentProvider(BaseModel):
|
||||
organization: str
|
||||
url: str | None = None
|
||||
|
||||
|
||||
class AgentCapabilities(BaseModel):
|
||||
streaming: bool = False
|
||||
pushNotifications: bool = False
|
||||
stateTransitionHistory: bool = False
|
||||
|
||||
|
||||
class AgentAuthentication(BaseModel):
|
||||
schemes: List[str]
|
||||
credentials: str | None = None
|
||||
|
||||
|
||||
class AgentSkill(BaseModel):
|
||||
id: str
|
||||
name: str
|
||||
description: str | None = None
|
||||
tags: List[str] | None = None
|
||||
examples: List[str] | None = None
|
||||
inputModes: List[str] | None = None
|
||||
outputModes: List[str] | None = None
|
||||
|
||||
|
||||
class AgentCard(BaseModel):
|
||||
name: str
|
||||
description: str | None = None
|
||||
url: str
|
||||
provider: AgentProvider | None = None
|
||||
version: str
|
||||
documentationUrl: str | None = None
|
||||
capabilities: AgentCapabilities
|
||||
authentication: AgentAuthentication | None = None
|
||||
defaultInputModes: List[str] = ["text"]
|
||||
defaultOutputModes: List[str] = ["text"]
|
||||
skills: List[AgentSkill]
|
||||
|
||||
|
||||
class A2AClientError(Exception):
|
||||
pass
|
||||
|
||||
|
||||
class A2AClientHTTPError(A2AClientError):
|
||||
def __init__(self, status_code: int, message: str):
|
||||
self.status_code = status_code
|
||||
self.message = message
|
||||
super().__init__(f"HTTP Error {status_code}: {message}")
|
||||
|
||||
|
||||
class A2AClientJSONError(A2AClientError):
|
||||
def __init__(self, message: str):
|
||||
self.message = message
|
||||
super().__init__(f"JSON Error: {message}")
|
||||
|
||||
|
||||
class MissingAPIKeyError(Exception):
|
||||
"""Exception for missing API key."""
|
||||
|
||||
pass
|
||||
@@ -0,0 +1,109 @@
|
||||
"""In Memory Cache utility."""
|
||||
|
||||
import threading
|
||||
import time
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
|
||||
class InMemoryCache:
|
||||
"""A thread-safe Singleton class to manage cache data.
|
||||
|
||||
Ensures only one instance of the cache exists across the application.
|
||||
"""
|
||||
|
||||
_instance: Optional["InMemoryCache"] = None
|
||||
_lock: threading.Lock = threading.Lock()
|
||||
_initialized: bool = False
|
||||
|
||||
def __new__(cls):
|
||||
"""Override __new__ to control instance creation (Singleton pattern).
|
||||
|
||||
Uses a lock to ensure thread safety during the first instantiation.
|
||||
|
||||
Returns:
|
||||
The singleton instance of InMemoryCache.
|
||||
"""
|
||||
if cls._instance is None:
|
||||
with cls._lock:
|
||||
if cls._instance is None:
|
||||
cls._instance = super().__new__(cls)
|
||||
return cls._instance
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the cache storage.
|
||||
|
||||
Uses a flag (_initialized) to ensure this logic runs only on the very first
|
||||
creation of the singleton instance.
|
||||
"""
|
||||
if not self._initialized:
|
||||
with self._lock:
|
||||
if not self._initialized:
|
||||
# print("Initializing SessionCache storage")
|
||||
self._cache_data: Dict[str, Dict[str, Any]] = {}
|
||||
self._ttl: Dict[str, float] = {}
|
||||
self._data_lock: threading.Lock = threading.Lock()
|
||||
self._initialized = True
|
||||
|
||||
def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
|
||||
"""Set a key-value pair.
|
||||
|
||||
Args:
|
||||
key: The key for the data.
|
||||
value: The data to store.
|
||||
ttl: Time to live in seconds. If None, data will not expire.
|
||||
"""
|
||||
with self._data_lock:
|
||||
self._cache_data[key] = value
|
||||
|
||||
if ttl is not None:
|
||||
self._ttl[key] = time.time() + ttl
|
||||
else:
|
||||
if key in self._ttl:
|
||||
del self._ttl[key]
|
||||
|
||||
def get(self, key: str, default: Any = None) -> Any:
|
||||
"""Get the value associated with a key.
|
||||
|
||||
Args:
|
||||
key: The key for the data within the session.
|
||||
default: The value to return if the session or key is not found.
|
||||
|
||||
Returns:
|
||||
The cached value, or the default value if not found.
|
||||
"""
|
||||
with self._data_lock:
|
||||
if key in self._ttl and time.time() > self._ttl[key]:
|
||||
del self._cache_data[key]
|
||||
del self._ttl[key]
|
||||
return default
|
||||
return self._cache_data.get(key, default)
|
||||
|
||||
def delete(self, key: str) -> None:
|
||||
"""Delete a specific key-value pair from a cache.
|
||||
|
||||
Args:
|
||||
key: The key to delete.
|
||||
|
||||
Returns:
|
||||
True if the key was found and deleted, False otherwise.
|
||||
"""
|
||||
|
||||
with self._data_lock:
|
||||
if key in self._cache_data:
|
||||
del self._cache_data[key]
|
||||
if key in self._ttl:
|
||||
del self._ttl[key]
|
||||
return True
|
||||
return False
|
||||
|
||||
def clear(self) -> bool:
|
||||
"""Remove all data.
|
||||
|
||||
Returns:
|
||||
True if the data was cleared, False otherwise.
|
||||
"""
|
||||
with self._data_lock:
|
||||
self._cache_data.clear()
|
||||
self._ttl.clear()
|
||||
return True
|
||||
return False
|
||||
@@ -0,0 +1,135 @@
|
||||
from jwcrypto import jwk
|
||||
import uuid
|
||||
from starlette.responses import JSONResponse
|
||||
from starlette.requests import Request
|
||||
from typing import Any
|
||||
|
||||
import jwt
|
||||
import time
|
||||
import json
|
||||
import hashlib
|
||||
import httpx
|
||||
import logging
|
||||
|
||||
from jwt import PyJWK, PyJWKClient
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
AUTH_HEADER_PREFIX = 'Bearer '
|
||||
|
||||
class PushNotificationAuth:
|
||||
def _calculate_request_body_sha256(self, data: dict[str, Any]):
|
||||
"""Calculates the SHA256 hash of a request body.
|
||||
|
||||
This logic needs to be same for both the agent who signs the payload and the client verifier.
|
||||
"""
|
||||
body_str = json.dumps(
|
||||
data,
|
||||
ensure_ascii=False,
|
||||
allow_nan=False,
|
||||
indent=None,
|
||||
separators=(",", ":"),
|
||||
)
|
||||
return hashlib.sha256(body_str.encode()).hexdigest()
|
||||
|
||||
class PushNotificationSenderAuth(PushNotificationAuth):
|
||||
def __init__(self):
|
||||
self.public_keys = []
|
||||
self.private_key_jwk: PyJWK = None
|
||||
|
||||
@staticmethod
|
||||
async def verify_push_notification_url(url: str) -> bool:
|
||||
async with httpx.AsyncClient(timeout=10) as client:
|
||||
try:
|
||||
validation_token = str(uuid.uuid4())
|
||||
response = await client.get(
|
||||
url,
|
||||
params={"validationToken": validation_token}
|
||||
)
|
||||
response.raise_for_status()
|
||||
is_verified = response.text == validation_token
|
||||
|
||||
logger.info(f"Verified push-notification URL: {url} => {is_verified}")
|
||||
return is_verified
|
||||
except Exception as e:
|
||||
logger.warning(f"Error during sending push-notification for URL {url}: {e}")
|
||||
|
||||
return False
|
||||
|
||||
def generate_jwk(self):
|
||||
key = jwk.JWK.generate(kty='RSA', size=2048, kid=str(uuid.uuid4()), use="sig")
|
||||
self.public_keys.append(key.export_public(as_dict=True))
|
||||
self.private_key_jwk = PyJWK.from_json(key.export_private())
|
||||
|
||||
def handle_jwks_endpoint(self, _request: Request):
|
||||
"""Allow clients to fetch public keys.
|
||||
"""
|
||||
return JSONResponse({
|
||||
"keys": self.public_keys
|
||||
})
|
||||
|
||||
def _generate_jwt(self, data: dict[str, Any]):
|
||||
"""JWT is generated by signing both the request payload SHA digest and time of token generation.
|
||||
|
||||
Payload is signed with private key and it ensures the integrity of payload for client.
|
||||
Including iat prevents from replay attack.
|
||||
"""
|
||||
|
||||
iat = int(time.time())
|
||||
|
||||
return jwt.encode(
|
||||
{"iat": iat, "request_body_sha256": self._calculate_request_body_sha256(data)},
|
||||
key=self.private_key_jwk,
|
||||
headers={"kid": self.private_key_jwk.key_id},
|
||||
algorithm="RS256"
|
||||
)
|
||||
|
||||
async def send_push_notification(self, url: str, data: dict[str, Any]):
|
||||
jwt_token = self._generate_jwt(data)
|
||||
headers = {'Authorization': f"Bearer {jwt_token}"}
|
||||
async with httpx.AsyncClient(timeout=10) as client:
|
||||
try:
|
||||
response = await client.post(
|
||||
url,
|
||||
json=data,
|
||||
headers=headers
|
||||
)
|
||||
response.raise_for_status()
|
||||
logger.info(f"Push-notification sent for URL: {url}")
|
||||
except Exception as e:
|
||||
logger.warning(f"Error during sending push-notification for URL {url}: {e}")
|
||||
|
||||
class PushNotificationReceiverAuth(PushNotificationAuth):
|
||||
def __init__(self):
|
||||
self.public_keys_jwks = []
|
||||
self.jwks_client = None
|
||||
|
||||
async def load_jwks(self, jwks_url: str):
|
||||
self.jwks_client = PyJWKClient(jwks_url)
|
||||
|
||||
async def verify_push_notification(self, request: Request) -> bool:
|
||||
auth_header = request.headers.get("Authorization")
|
||||
if not auth_header or not auth_header.startswith(AUTH_HEADER_PREFIX):
|
||||
print("Invalid authorization header")
|
||||
return False
|
||||
|
||||
token = auth_header[len(AUTH_HEADER_PREFIX):]
|
||||
signing_key = self.jwks_client.get_signing_key_from_jwt(token)
|
||||
|
||||
decode_token = jwt.decode(
|
||||
token,
|
||||
signing_key,
|
||||
options={"require": ["iat", "request_body_sha256"]},
|
||||
algorithms=["RS256"],
|
||||
)
|
||||
|
||||
actual_body_sha256 = self._calculate_request_body_sha256(await request.json())
|
||||
if actual_body_sha256 != decode_token["request_body_sha256"]:
|
||||
# Payload signature does not match the digest in signed token.
|
||||
raise ValueError("Invalid request body")
|
||||
|
||||
if time.time() - decode_token["iat"] > 60 * 5:
|
||||
# Do not allow push-notifications older than 5 minutes.
|
||||
# This is to prevent replay attack.
|
||||
raise ValueError("Token is expired")
|
||||
|
||||
return True
|
||||
Reference in New Issue
Block a user