smarta2a 0.2.0__tar.gz → 0.2.1__tar.gz

{smarta2a-0.2.0 → smarta2a-0.2.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: smarta2a
-Version: 0.2.0
+Version: 0.2.1
 Summary: A Python package for creating servers and clients following Google's Agent2Agent protocol
 Project-URL: Homepage, https://github.com/siddharthsma/smarta2a
 Project-URL: Bug Tracker, https://github.com/siddharthsma/smarta2a/issues

{smarta2a-0.2.0 → smarta2a-0.2.1}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "smarta2a"
-version = "0.2.0"
+version = "0.2.1"
 authors = [
     { name = "Siddharth Ambegaonkar", email = "siddharthsma@gmail.com" },
 ]

{smarta2a-0.2.0 → smarta2a-0.2.1}/smarta2a/server.py

@@ -48,6 +48,11 @@ from .types import (
     A2AStatus,
     A2AStreamResponse,
     TaskSendParams,
+    SetTaskPushNotificationRequest,
+    GetTaskPushNotificationRequest,
+    SetTaskPushNotificationResponse,
+    GetTaskPushNotificationResponse,
+    TaskPushNotificationConfig,
 )
 
 class SmartA2A:
@@ -129,6 +134,18 @@ class SmartA2A:
             self._register_handler("tasks/cancel", func, "task_cancel", "handler")
             return func
         return decorator
+    
+    def set_notification(self):
+        def decorator(func: Callable[[SetTaskPushNotificationRequest], None]) -> Callable:
+            self._register_handler("tasks/pushNotification/set", func, "set_notification", "handler")
+            return func
+        return decorator
+    
+    def get_notification(self):
+        def decorator(func: Callable[[GetTaskPushNotificationRequest], Union[TaskPushNotificationConfig, GetTaskPushNotificationResponse]]):
+            self._register_handler("tasks/pushNotification/get", func, "get_notification", "handler")
+            return func
+        return decorator
 
     async def process_request(self, request_data: dict) -> JSONRPCResponse:
         try:
@@ -141,6 +158,10 @@ class SmartA2A:
                 return self._handle_get_task(request_data)
             elif method == "tasks/cancel":
                 return self._handle_cancel_task(request_data)
+            elif method == "tasks/pushNotification/set":
+                return self._handle_set_notification(request_data)
+            elif method == "tasks/pushNotification/get":
+                return self._handle_get_notification(request_data)
             else:
                 return self._error_response(
                     request_data.get("id"),
@@ -422,6 +443,102 @@ class SmartA2A:
                 error=InternalError(data=str(e))
             )
 
+    def _handle_set_notification(self, request_data: dict) -> SetTaskPushNotificationResponse:
+        try:
+            request = SetTaskPushNotificationRequest.model_validate(request_data)
+            handler = self.handlers.get("tasks/pushNotification/set")
+            
+            if not handler:
+                return SetTaskPushNotificationResponse(
+                    id=request.id,
+                    error=MethodNotFoundError()
+                )
+
+            try:
+                # Execute handler (may or may not return something)
+                raw_result = handler(request)
+                
+                # If handler returns nothing - build success response from request params
+                if raw_result is None:
+                    return SetTaskPushNotificationResponse(
+                        id=request.id,
+                        result=request.params
+                    )
+                
+                # If handler returns a full response object
+                if isinstance(raw_result, SetTaskPushNotificationResponse):
+                    return raw_result
+                    
+
+            except Exception as e:
+                if isinstance(e, JSONRPCError):
+                    return SetTaskPushNotificationResponse(
+                        id=request.id,
+                        error=e
+                    )
+                return SetTaskPushNotificationResponse(
+                    id=request.id,
+                    error=InternalError(data=str(e))
+                )
+
+        except ValidationError as e:
+            return SetTaskPushNotificationResponse(
+                id=request_data.get("id"),
+                error=InvalidRequestError(data=e.errors())
+            )
+                      
+
+    def _handle_get_notification(self, request_data: dict) -> GetTaskPushNotificationResponse:
+        try:
+            request = GetTaskPushNotificationRequest.model_validate(request_data)
+            handler = self.handlers.get("tasks/pushNotification/get")
+            
+            if not handler:
+                return GetTaskPushNotificationResponse(
+                    id=request.id,
+                    error=MethodNotFoundError()
+                )
+            
+            try:
+                raw_result = handler(request)
+                
+                if isinstance(raw_result, GetTaskPushNotificationResponse):
+                    return raw_result
+                else:
+                    # Validate raw_result as TaskPushNotificationConfig
+                    config = TaskPushNotificationConfig.model_validate(raw_result)
+                    return GetTaskPushNotificationResponse(
+                        id=request.id,
+                        result=config
+                    )
+            except ValidationError as e:
+                return GetTaskPushNotificationResponse(
+                    id=request.id,
+                    error=InvalidParamsError(data=e.errors())
+                )
+            except Exception as e:
+                if isinstance(e, JSONRPCError):
+                    return GetTaskPushNotificationResponse(
+                        id=request.id,
+                        error=e
+                    )
+                return GetTaskPushNotificationResponse(
+                    id=request.id,
+                    error=InternalError(data=str(e))
+                )
+
+        except ValidationError as e:
+            return GetTaskPushNotificationResponse(
+                id=request_data.get("id"),
+                error=InvalidRequestError(data=e.errors())
+            )
+        except json.JSONDecodeError as e:
+            return GetTaskPushNotificationResponse(
+                id=request_data.get("id"),
+                error=JSONParseError(data=str(e))
+            )
+
+
     def _normalize_artifacts(self, content: Any) -> List[Artifact]:
         """Handle both A2AResponse content and regular returns"""
         if isinstance(content, Artifact):

{smarta2a-0.2.0 → smarta2a-0.2.1}/tests/test_server.py

@@ -24,7 +24,14 @@ from smarta2a.types import (
     A2AStatus,
     A2AStreamResponse,
     SendTaskResponse,  
-    Message
+    Message,
+    InternalError,
+    TaskNotFoundError,
+    SetTaskPushNotificationRequest,
+    GetTaskPushNotificationRequest,
+    SetTaskPushNotificationResponse,
+    GetTaskPushNotificationResponse,
+    TaskPushNotificationConfig
 )
 
 @pytest.fixture
@@ -500,4 +507,145 @@ def test_send_task_content_access():
     assert request.content == request.params.message.parts
 
 
+def test_set_notification_success(a2a_server, client):
+    # Test basic success case with no return value
+    @a2a_server.set_notification()
+    def handle_set(req: SetTaskPushNotificationRequest):
+        # No return needed - just validate request
+        assert req.params.id == "test123"
+    
+    request_data = {
+        "jsonrpc": "2.0",
+        "id": 1,
+        "method": "tasks/pushNotification/set",
+        "params": {
+            "id": "test123",
+            "pushNotificationConfig": {
+                "url": "https://example.com/callback",
+                "authentication": {
+                    "schemes": ["jwt"]
+                }
+            }
+        }
+    }
+    
+    response = client.post("/", json=request_data).json()
+    
+    assert response["result"]["id"] == "test123"
+    assert response["result"]["pushNotificationConfig"]["url"] == request_data["params"]["pushNotificationConfig"]["url"]
+    assert response["result"]["pushNotificationConfig"]["authentication"]["schemes"] == ["jwt"]
+
+def test_set_notification_custom_response(a2a_server, client):
+    # Test handler returning custom response
+    @a2a_server.set_notification()
+    def handle_set(req):
+        return SetTaskPushNotificationResponse(
+            id=req.id,
+            result=TaskPushNotificationConfig(
+                id="test123",
+                pushNotificationConfig={
+                    "url": "custom-url",
+                    "token": "secret"
+                }
+            )
+        )
+    
+    response = client.post("/", json={
+        "jsonrpc": "2.0",
+        "id": 2,
+        "method": "tasks/pushNotification/set",
+        "params": {
+            "id": "test123",
+            "pushNotificationConfig": {"url": "https://example.com"}
+        }
+    }).json()
+    
+    assert response["result"]["pushNotificationConfig"]["url"] == "custom-url"
+    assert "secret" in response["result"]["pushNotificationConfig"]["token"]
+
+
+# --- Get Notification Tests ---
+
+def test_get_notification_success(a2a_server, client):
+    # Test successful config retrieval
+    @a2a_server.get_notification()
+    def handle_get(req: GetTaskPushNotificationRequest):
+        return TaskPushNotificationConfig(
+            id=req.params.id,
+            pushNotificationConfig={
+                "url": "https://test.com",
+                "token": "abc123"
+            }
+        )
+    
+    request_data = {
+        "jsonrpc": "2.0",
+        "id": 4,
+        "method": "tasks/pushNotification/get",
+        "params": {"id": "test456"}
+    }
+    
+    response = client.post("/", json=request_data).json()
+    
+    assert response["result"]["id"] == "test456"
+    assert response["result"]["pushNotificationConfig"]["url"] == "https://test.com"
+
+def test_get_notification_direct_response(a2a_server, client):
+    # Test handler returning full response object
+    @a2a_server.get_notification()
+    def handle_get(req):
+        return GetTaskPushNotificationResponse(
+            id=req.id,
+            result=TaskPushNotificationConfig(
+                id=req.params.id,
+                pushNotificationConfig={
+                    "url": "direct-response.example",
+                    "authentication": {"schemes": ["basic"]}
+                }
+            )
+        )
+    
+    response = client.post("/", json={
+        "jsonrpc": "2.0",
+        "id": 5,
+        "method": "tasks/pushNotification/get",
+        "params": {"id": "test789"}
+    }).json()
+    
+    assert "direct-response" in response["result"]["pushNotificationConfig"]["url"]
+    assert "basic" in response["result"]["pushNotificationConfig"]["authentication"]["schemes"]
+
+def test_get_notification_validation_error(a2a_server, client):
+    # Test invalid response from handler
+    @a2a_server.get_notification()
+    def handle_get(req):
+        return {"invalid": "config"}
+    
+    response = client.post("/", json={
+        "jsonrpc": "2.0",
+        "id": 6,
+        "method": "tasks/pushNotification/get",
+        "params": {"id": "test999"}
+    }).json()
+    
+    assert response["error"]["code"] == -32602  # Invalid params
+    
+
+def test_get_notification_error_propagation(a2a_server, client):
+    # Test exception handling
+    @a2a_server.get_notification()
+    def handle_get(req):
+        raise InternalError(message="Storage failure")
+    
+    response = client.post("/", json={
+        "jsonrpc": "2.0",
+        "id": 7,
+        "method": "tasks/pushNotification/get",
+        "params": {"id": "test-error"}
+    }).json()
+    
+    assert response["error"]["code"] == -32603  # Internal error code
+    
+
+