From 6a1001426229348fa394e6c27b9350cf88211880 Mon Sep 17 00:00:00 2001 From: Shaishav Pidadi Date: Sun, 25 Oct 2026 06:07:54 -0400 Subject: [PATCH] fix(sdk): confirmation client speaks the platform API; policy enrichment opt-in MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ConfirmationClient used /confirmation/... (no /api/v1 prefix), expected a flat response with an 'approved' flag, polled for 'rejected' and cancelled with DELETE — none of which the platform does. Now: /api/v1/confirmation/ create|:id|cancel, nested 'confirmation' records, denied/cancelled/consumed - consume_confirmation(): single-use claim of an approved confirmation - ConfirmationStatus carries request_type/request_payload and is_final - send_policy_config (default False): precheck resolves the org policy - bump to 1.1.0 (not published) --- governs_ai/__init__.py | 2 +- governs_ai/client.py | 10 +++- governs_ai/clients/confirmation.py | 40 +++++++++++-- governs_ai/clients/precheck.py | 16 ++++- governs_ai/models/confirmation.py | 52 ++++++++++++---- pyproject.toml | 2 +- tests/test_confirmation_client.py | 96 ++++++++++++++++++++++++++++++ tests/test_precheck_client.py | 31 +++++++++- 8 files changed, 223 insertions(+), 26 deletions(-) create mode 100644 tests/test_confirmation_client.py diff --git a/governs_ai/__init__.py b/governs_ai/__init__.py index 7189ec6..27d879c 100644 --- a/governs_ai/__init__.py +++ b/governs_ai/__init__.py @@ -54,7 +54,7 @@ AnalyticsError, ) -__version__ = "1.0.0" +__version__ = "1.1.0" __author__ = "GovernsAI" __email__ = "support@governs.ai" diff --git a/governs_ai/client.py b/governs_ai/client.py index d00ba71..1c8569b 100644 --- a/governs_ai/client.py +++ b/governs_ai/client.py @@ -41,6 +41,10 @@ class GovernsAIConfig: retry_delay: int = 1000 logger: Optional[Any] = None http_client: Optional[HTTPClient] = None + # Fetch the org's policy and send it as policy_config. Off by default: + # precheck resolves the org's dashboard policy from the API key and + # ignores a payload policy whenever the org has one. + send_policy_config: bool = False def __post_init__(self): """Validate configuration after initialization.""" @@ -113,7 +117,11 @@ def __init__( ) # Initialize feature clients - self.precheck = PrecheckClient(self.http_client, self.logger) + self.precheck = PrecheckClient( + self.http_client, + self.logger, + send_policy_config=self.config.send_policy_config, + ) self.confirmation = ConfirmationClient(self.http_client, self.logger) self.budget = BudgetClient(self.http_client, self.logger) self.tools = ToolClient(self.http_client, self.logger) diff --git a/governs_ai/clients/confirmation.py b/governs_ai/clients/confirmation.py index b2991c7..a43c760 100644 --- a/governs_ai/clients/confirmation.py +++ b/governs_ai/clients/confirmation.py @@ -42,7 +42,7 @@ async def create_confirmation(self, request: ConfirmationRequest) -> Confirmatio try: self.logger.debug(f"Creating confirmation: {request.request_type}") response = await self.http_client.post( - "/confirmation", + "/api/v1/confirmation/create", data=request.to_dict(), ) return ConfirmationResponse.from_dict(response.data) @@ -63,7 +63,7 @@ async def get_confirmation_status(self, correlation_id: str) -> ConfirmationStat try: self.logger.debug(f"Getting confirmation status: {correlation_id}") response = await self.http_client.get( - f"/confirmation/{correlation_id}", + f"/api/v1/confirmation/{correlation_id}", ) return ConfirmationStatus.from_dict(response.data) except Exception as e: @@ -100,11 +100,11 @@ async def poll_confirmation( try: status = await self.get_confirmation_status(correlation_id) - if status.approved: + if status.status == "approved": self.logger.info(f"Confirmation approved: {correlation_id}") return status - - if status.status == "rejected": + + if status.status in ("denied", "cancelled", "consumed", "rejected"): self.logger.warning(f"Confirmation rejected: {correlation_id}") raise ConfirmationRejectedError( f"Confirmation was rejected: {correlation_id}", @@ -137,6 +137,31 @@ async def poll_confirmation( self.logger.error(f"Poll confirmation failed: {str(e)}") raise ConfirmationError(f"Poll confirmation failed: {str(e)}") + async def consume_confirmation(self, correlation_id: str) -> ConfirmationStatus: + """ + Claim an approved confirmation for exactly one execution. + + Call this right before running the approved action: it moves the + confirmation from approved to consumed atomically, and a second call + fails, so one approval cannot authorise two runs. + + Args: + correlation_id: Correlation ID of the confirmation + + Returns: + ConfirmationStatus with status "consumed" + """ + try: + self.logger.debug(f"Consuming confirmation: {correlation_id}") + response = await self.http_client.post( + f"/api/v1/confirmation/{correlation_id}/consume", + data={}, + ) + return ConfirmationStatus.from_dict(response.data) + except Exception as e: + self.logger.error(f"Consume confirmation failed: {str(e)}") + raise ConfirmationError(f"Consume confirmation failed: {str(e)}") + async def cancel_confirmation(self, correlation_id: str) -> bool: """ Cancel a confirmation request. @@ -149,7 +174,10 @@ async def cancel_confirmation(self, correlation_id: str) -> bool: """ try: self.logger.debug(f"Cancelling confirmation: {correlation_id}") - response = await self.http_client.delete(f"/confirmation/{correlation_id}") + response = await self.http_client.post( + "/api/v1/confirmation/cancel", + data={"correlationId": correlation_id}, + ) return response.is_success except Exception as e: self.logger.error(f"Cancel confirmation failed: {str(e)}") diff --git a/governs_ai/clients/precheck.py b/governs_ai/clients/precheck.py index b0a6d14..ad4fde4 100644 --- a/governs_ai/clients/precheck.py +++ b/governs_ai/clients/precheck.py @@ -15,10 +15,20 @@ class PrecheckClient: """Client for precheck operations.""" - def __init__(self, http_client: HTTPClient, logger: GovernsAILogger): - """Initialize precheck client.""" + def __init__( + self, + http_client: HTTPClient, + logger: GovernsAILogger, + send_policy_config: bool = False, + ): + """Initialize precheck client. + + send_policy_config: also fetch and send the org policy. Precheck + already resolves it from the API key, so this is off by default. + """ self.http_client = http_client self.logger = logger + self.send_policy_config = send_policy_config async def check_request(self, request: PrecheckRequest) -> PrecheckResponse: """ @@ -44,7 +54,7 @@ async def check_request(self, request: PrecheckRequest) -> PrecheckResponse: async def _enrich_request(self, request: PrecheckRequest) -> PrecheckRequest: """Enrich request with platform policy/tool/budget context when missing.""" - needs_policy = request.policy_config is None + needs_policy = self.send_policy_config and request.policy_config is None needs_tool = request.tool_config is None and bool(request.tool) needs_budget = request.budget_context is None diff --git a/governs_ai/models/confirmation.py b/governs_ai/models/confirmation.py index 44de720..58dc5c9 100644 --- a/governs_ai/models/confirmation.py +++ b/governs_ai/models/confirmation.py @@ -42,6 +42,16 @@ def to_dict(self) -> Dict[str, Any]: } +def _unwrap(data: Dict[str, Any]) -> Dict[str, Any]: + """The platform nests the record under `confirmation`; accept both forms.""" + inner = data.get("confirmation") + return inner if isinstance(inner, dict) else data + + +# Statuses after which nothing changes. `consumed` = approved and used once. +FINAL_STATUSES = ("approved", "denied", "expired", "cancelled", "consumed") + + @dataclass class ConfirmationResponse: """Response model for confirmation operations.""" @@ -52,12 +62,13 @@ class ConfirmationResponse: @classmethod def from_dict(cls, data: Dict[str, Any]) -> "ConfirmationResponse": - """Create from dictionary.""" + """Create from the platform's `/confirmation/create` response.""" + record = _unwrap(data) return cls( - correlation_id=data["correlationId"], - confirmation_url=data["confirmationUrl"], - status=data["status"], - expires_at=data.get("expiresAt"), + correlation_id=record["correlationId"], + confirmation_url=data.get("confirmationUrl") or record.get("confirmationUrl", ""), + status=record["status"], + expires_at=record.get("expiresAt"), ) def to_dict(self) -> Dict[str, Any]: @@ -72,24 +83,37 @@ def to_dict(self) -> Dict[str, Any]: @dataclass class ConfirmationStatus: - """Status model for confirmation tracking.""" + """Status model for confirmation tracking. + + status is one of pending, approved, denied, expired, cancelled, consumed. + """ correlation_id: str status: str approved: bool approved_at: Optional[str] = None rejected_at: Optional[str] = None expires_at: Optional[str] = None + request_type: Optional[str] = None + request_payload: Optional[Dict[str, Any]] = None + + @property + def is_final(self) -> bool: + return self.status in FINAL_STATUSES @classmethod def from_dict(cls, data: Dict[str, Any]) -> "ConfirmationStatus": - """Create from dictionary.""" + """Create from `GET /api/v1/confirmation/:id` (nested or flat).""" + record = _unwrap(data) + status = record["status"] return cls( - correlation_id=data["correlationId"], - status=data["status"], - approved=data["approved"], - approved_at=data.get("approvedAt"), - rejected_at=data.get("rejectedAt"), - expires_at=data.get("expiresAt"), + correlation_id=record["correlationId"], + status=status, + approved=bool(record.get("approved", status in ("approved", "consumed"))), + approved_at=record.get("approvedAt"), + rejected_at=record.get("rejectedAt"), + expires_at=record.get("expiresAt"), + request_type=record.get("requestType"), + request_payload=record.get("requestPayload"), ) def to_dict(self) -> Dict[str, Any]: @@ -101,4 +125,6 @@ def to_dict(self) -> Dict[str, Any]: "approvedAt": self.approved_at, "rejectedAt": self.rejected_at, "expiresAt": self.expires_at, + "requestType": self.request_type, + "requestPayload": self.request_payload, } diff --git a/pyproject.toml b/pyproject.toml index 79083fe..90f6b96 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "governs-ai-sdk" -version = "1.0.0" +version = "1.1.0" description = "Python SDK for GovernsAI - AI governance and compliance platform" readme = "README.md" requires-python = ">=3.8" diff --git a/tests/test_confirmation_client.py b/tests/test_confirmation_client.py new file mode 100644 index 0000000..a9d054d --- /dev/null +++ b/tests/test_confirmation_client.py @@ -0,0 +1,96 @@ +# SPDX-License-Identifier: MIT +# Copyright (c) 2024 GovernsAI. All rights reserved. +"""ConfirmationClient against the platform's real routes and response shapes.""" + +from unittest.mock import AsyncMock + +import pytest + +from governs_ai.clients.confirmation import ConfirmationClient +from governs_ai.exceptions.confirmation import ConfirmationRejectedError +from governs_ai.models.confirmation import ConfirmationRequest, ConfirmationStatus +from governs_ai.utils.http import HTTPResponse +from governs_ai.utils.logging import GovernsAILogger + + +def _resp(data, status=200): + return HTTPResponse(status_code=status, data=data, headers={}, url="") + + +def _record(status): + return { + "confirmation": { + "correlationId": "corr-1", + "status": status, + "requestType": "tool_call", + "requestPayload": {"tool": "payment_process", "args": {"amount": 5}}, + "expiresAt": "2026-10-25T10:00:00Z", + } + } + + +@pytest.fixture +def http(): + return AsyncMock() + + +@pytest.mark.asyncio +async def test_create_posts_to_platform_route(http): + http.post.return_value = _resp( + { + "success": True, + "confirmation": {"correlationId": "corr-1", "status": "pending"}, + "confirmationUrl": "http://localhost:3002/confirm/corr-1", + } + ) + client = ConfirmationClient(http, GovernsAILogger()) + created = await client.create_confirmation( + ConfirmationRequest( + request_type="tool_call", + request_desc="Run payment_process", + request_payload={"tool": "payment_process"}, + reasons=["tool-specific.confirm"], + user_id="u1", + correlation_id="corr-1", + ) + ) + assert http.post.call_args.args[0] == "/api/v1/confirmation/create" + assert created.correlation_id == "corr-1" + assert created.confirmation_url.endswith("/confirm/corr-1") + + +@pytest.mark.asyncio +async def test_status_unwraps_nested_record(http): + http.get.return_value = _resp(_record("approved")) + client = ConfirmationClient(http, GovernsAILogger()) + status = await client.get_confirmation_status("corr-1") + assert http.get.call_args.args[0] == "/api/v1/confirmation/corr-1" + assert status.approved is True + assert status.is_final + assert status.request_payload["tool"] == "payment_process" + + +@pytest.mark.asyncio +async def test_poll_raises_on_denied(http): + http.get.return_value = _resp(_record("denied")) + client = ConfirmationClient(http, GovernsAILogger()) + with pytest.raises(ConfirmationRejectedError): + await client.poll_confirmation("corr-1", max_duration=1000, poll_interval=10) + + +@pytest.mark.asyncio +async def test_consume_and_cancel_routes(http): + http.post.return_value = _resp(_record("consumed")) + client = ConfirmationClient(http, GovernsAILogger()) + consumed = await client.consume_confirmation("corr-1") + assert http.post.call_args.args[0] == "/api/v1/confirmation/corr-1/consume" + assert consumed.status == "consumed" + + http.post.return_value = _resp({"success": True}) + assert await client.cancel_confirmation("corr-1") is True + assert http.post.call_args.args[0] == "/api/v1/confirmation/cancel" + assert http.post.call_args.kwargs["data"] == {"correlationId": "corr-1"} + + +def test_pending_is_not_final(): + assert not ConfirmationStatus.from_dict(_record("pending")).is_final diff --git a/tests/test_precheck_client.py b/tests/test_precheck_client.py index bebfab6..60ef06b 100644 --- a/tests/test_precheck_client.py +++ b/tests/test_precheck_client.py @@ -67,7 +67,7 @@ async def get_side_effect(endpoint, params=None, headers=None): } ) - client = PrecheckClient(mock_http_client, GovernsAILogger()) + client = PrecheckClient(mock_http_client, GovernsAILogger(), send_policy_config=True) request = PrecheckRequest( tool="model.chat", scope="net.external", @@ -89,3 +89,32 @@ async def get_side_effect(endpoint, params=None, headers=None): assert "policy_config" in post_payload assert "tool_config" in post_payload assert "budget_context" in post_payload + + +@pytest.mark.asyncio +async def test_policy_config_not_fetched_by_default(): + """Precheck resolves the org policy itself; the SDK only adds tool/budget context.""" + mock_http_client = AsyncMock() + requested = [] + + async def get_side_effect(endpoint, params=None, headers=None): + requested.append(endpoint) + if endpoint == "/api/v1/tools/model.chat/metadata": + return _mock_response({"metadata": {"tool_name": "model.chat"}}) + if endpoint == "/api/v1/budget/context": + return _mock_response({"monthly_limit": 10.0}) + raise AssertionError(f"Unexpected GET endpoint: {endpoint}") + + mock_http_client.get.side_effect = get_side_effect + mock_http_client.post.return_value = _mock_response({"decision": "allow", "reasons": []}) + + client = PrecheckClient(mock_http_client, GovernsAILogger()) + await client.check_request( + PrecheckRequest(tool="model.chat", scope="net.external", raw_text="hi", payload={}, tags=[]) + ) + + assert "/api/v1/policies" not in requested + post_payload = mock_http_client.post.call_args.kwargs.get("data") or mock_http_client.post.call_args.args[1] + assert post_payload.get("policy_config") is None + assert post_payload["tool_config"] == {"tool_name": "model.chat"} +