openbox-langgraph-sdk-python 1.0.0__py3-none-any.whl → 1.2.0__py3-none-any.whl

openbox_langgraph/__init__.py

@@ -17,7 +17,7 @@ Example:
     ... )
 """
 
-__version__ = "1.0.0"
+__version__ = "1.2.0"
 
 from openbox_langgraph.client import GovernanceClient, build_auth_headers
 from openbox_langgraph.config import (
@@ -38,6 +38,7 @@ from openbox_langgraph.errors import (
     OpenBoxError,
     OpenBoxInsecureURLError,
     OpenBoxNetworkError,
+    OpenBoxSigningError,
 )
 from openbox_langgraph.hitl import poll_until_decision
 from openbox_langgraph.identity import (
@@ -112,6 +113,7 @@ __all__ = [
     "OpenBoxLangGraphHandler",
     "OpenBoxLangGraphHandlerOptions",
     "OpenBoxNetworkError",
+    "OpenBoxSigningError",
     "Verdict",
     "VerdictContext",
     "WorkflowEventType",

openbox_langgraph/activity_approval.py

@@ -0,0 +1,156 @@
+"""Wait at the governed operation, preserving its stack and activity identity."""
+
+from __future__ import annotations
+
+import asyncio
+import threading
+from dataclasses import dataclass, field
+
+from openbox_core.client import EvaluationClient
+from openbox_core.contracts.context import ActivityContext
+from openbox_core.contracts.results import ApprovalResult, EvaluationResult, Verdict
+
+from openbox_langgraph.errors import (
+    ApprovalExpiredError,
+    ApprovalRejectedError,
+    OpenBoxConfigError,
+    _raise_core_error,
+)
+from openbox_langgraph.types import HITLConfig
+
+
+@dataclass
+class _ApprovalTurn:
+    cancelled: threading.Event = field(default_factory=threading.Event)
+    decisions: dict[
+        str, tuple[EvaluationResult, ApprovalExpiredError | ApprovalRejectedError | None]
+    ] = field(default_factory=dict)
+
+
+class ActivityApprovalWaiter:
+    """Share approval decisions between the sync and async tool callbacks.
+
+    The callbacks can enforce the same stashed evaluation twice. Only that
+    exact evaluation is reusable: a later hook verdict on the same activity
+    must get its own decision. All state is discarded when the turn ends.
+    """
+
+    def __init__(self, client: EvaluationClient, config: HITLConfig) -> None:
+        self._client = client
+        self._config = config
+        self._lock = threading.Lock()
+        self._turns: dict[tuple[str, str], _ApprovalTurn] = {}
+
+    def begin_turn(self, workflow_id: str, run_id: str) -> None:
+        with self._lock:
+            self._turns[(workflow_id, run_id)] = _ApprovalTurn()
+
+    def end_turn(self, workflow_id: str) -> None:
+        with self._lock:
+            for key in list(self._turns):
+                if key[0] == workflow_id:
+                    self._turns.pop(key).cancelled.set()
+
+    def _turn(self, ctx: ActivityContext) -> _ApprovalTurn:
+        workflow_id, run_id, _ = self._identity(ctx)
+        with self._lock:
+            turn = self._turns.get((workflow_id, run_id))
+        if turn is None:
+            raise OpenBoxConfigError("Approval requires an active governed turn")
+        return turn
+
+    @staticmethod
+    def _identity(ctx: ActivityContext) -> tuple[str, str, str]:
+        if not ctx.workflow_id or not ctx.run_id or not ctx.activity_id:
+            raise OpenBoxConfigError("Approval requires workflow, run, and activity IDs")
+        return ctx.workflow_id, ctx.run_id, ctx.activity_id
+
+    @staticmethod
+    def _check_cancelled(turn: _ApprovalTurn) -> None:
+        # Cancelling an asyncio task does not stop its executor thread. Wake
+        # the sync waiter too, so a late approval cannot run an abandoned tool.
+        if turn.cancelled.is_set():
+            raise asyncio.CancelledError
+
+    def _reuse_decision(
+        self, turn: _ApprovalTurn, ctx: ActivityContext, result: EvaluationResult
+    ) -> bool:
+        self._check_cancelled(turn)
+        with self._lock:
+            decision = turn.decisions.get(self._identity(ctx)[2])
+        if decision is None or decision[0] is not result:
+            return False
+        if decision[1] is not None:
+            raise decision[1]
+        return True
+
+    def _record_response(
+        self,
+        turn: _ApprovalTurn,
+        ctx: ActivityContext,
+        result: EvaluationResult,
+        response: ApprovalResult | None,
+    ) -> bool:
+        self._check_cancelled(turn)
+        try:
+            approved = self._approved(response, ctx)
+        except (ApprovalExpiredError, ApprovalRejectedError) as exc:
+            with self._lock:
+                turn.decisions[self._identity(ctx)[2]] = (result, exc)
+            raise
+        if not approved:
+            return False
+        with self._lock:
+            turn.decisions[self._identity(ctx)[2]] = (result, None)
+        return True
+
+    @staticmethod
+    def _approved(response: ApprovalResult | None, ctx: ActivityContext) -> bool:
+        if response is None:
+            return False
+        if response.expired:
+            raise ApprovalExpiredError(f"Approval expired for {ctx.activity_type}")
+        if response.verdict in (Verdict.BLOCK, Verdict.HALT):
+            raise ApprovalRejectedError(
+                response.reason or f"Approval rejected for {ctx.activity_type}"
+            )
+        return response.verdict is Verdict.ALLOW
+
+    async def wait(self, result: EvaluationResult, ctx: ActivityContext) -> None:
+        turn = self._turn(ctx)
+        if self._reuse_decision(turn, ctx, result):
+            return
+        while True:
+            self._check_cancelled(turn)
+            try:
+                response = await self._client.apoll_approval(*self._identity(ctx))
+            except Exception as exc:
+                _raise_core_error(exc)
+            self._check_cancelled(turn)
+            if self._record_response(turn, ctx, result, response):
+                return
+            await asyncio.sleep(self._config.poll_interval_ms / 1000.0)
+
+    def wait_sync(self, result: EvaluationResult, ctx: ActivityContext) -> None:
+        turn = self._turn(ctx)
+        if self._reuse_decision(turn, ctx, result):
+            return
+        try:
+            asyncio.get_running_loop()
+        except RuntimeError:
+            pass  # A sync tool runs in LangChain's executor thread.
+        else:
+            raise OpenBoxConfigError(
+                "A synchronous operation requiring approval cannot wait on the event-loop "
+                "thread. Use its async API or run it with asyncio.to_thread()."
+            )
+        while True:
+            self._check_cancelled(turn)
+            try:
+                response = self._client.poll_approval(*self._identity(ctx))
+            except Exception as exc:
+                _raise_core_error(exc)
+            self._check_cancelled(turn)
+            if self._record_response(turn, ctx, result, response):
+                return
+            turn.cancelled.wait(self._config.poll_interval_ms / 1000.0)

openbox_langgraph/client.py

@@ -13,14 +13,20 @@ from openbox_core.contracts.results import EvaluationResult
 from openbox_core.contracts.results import Verdict as _CoreVerdict
 from openbox_core.errors import ContractError as _CoreContractError
 from openbox_core.errors import GovernanceAPIError as _CoreGovernanceAPIError
-from openbox_core.errors import OpenBoxNetworkError as _CoreOpenBoxNetworkError
+from openbox_core.errors import OpenBoxConfigError as _CoreOpenBoxConfigError
 
 from openbox_langgraph.core_events import to_envelope
-from openbox_langgraph.errors import OpenBoxConfigError, OpenBoxNetworkError
+from openbox_langgraph.errors import (
+    OpenBoxConfigError,
+    OpenBoxError,
+    OpenBoxNetworkError,
+    _raise_core_error,
+)
 from openbox_langgraph.identity import (
     AgentIdentityConfig,
     create_agent_identity_headers,
     parse_optional_agent_identity_config,
+    parse_optional_workload_private_key,
 )
 from openbox_langgraph.types import (
     ApprovalResponse,
@@ -32,9 +38,10 @@ from openbox_langgraph.types import (
 )
 
 if TYPE_CHECKING:
+    from openbox_core.client import EvaluationClient
     from openbox_core.gate import GovernanceGate
 
-_SDK_PACKAGE_VERSION = "1.0.0"
+_SDK_PACKAGE_VERSION = "1.2.0"
 _SDK_IDENTIFIER = f"openbox-langgraph-python-v{_SDK_PACKAGE_VERSION}"
 
 
@@ -114,9 +121,11 @@ async def _gate_evaluate(
       independent of `on_api_error`. Enforcing fail_closed here would let an
       SDK-side mapping defect block a user's graph for a reason their OWN policy
       never produced — strictly worse than dropping one governance event.
-    - `GovernanceAPIError` / `OpenBoxNetworkError` (network-shaped failure) ->
-      this SDK's `OpenBoxNetworkError`, same public exception the legacy path
-      raises under fail_closed.
+    - Base authentication, signing, and configuration errors always propagate
+      through this SDK's public errors, including under fail_open.
+    - `GovernanceAPIError` / `OpenBoxNetworkError` -> this SDK's
+      `OpenBoxNetworkError`. Bootstrap/token-exchange failures must propagate;
+      the base client applies fail_open only to ordinary governance transport errors.
     - Any OTHER exception (e.g. a malformed Core 200 body the base parser
       cannot decode) is a transport-shaped fault, NOT a governance verdict:
       routed through `_network_fallback_result` so fail_open returns `None`
@@ -127,8 +136,8 @@ async def _gate_evaluate(
         result = await gate.aevaluate(to_envelope(event))
     except _CoreContractError:
         return None
-    except (_CoreGovernanceAPIError, _CoreOpenBoxNetworkError) as e:
-        raise OpenBoxNetworkError(str(e)) from e
+    except (_CoreGovernanceAPIError, _CoreOpenBoxConfigError) as e:
+        _raise_core_error(e)
     except Exception as e:
         return _network_fallback_result(on_api_error, f"Governance gate error: {e}")
     return _collapse_client_synthesized_fallback(result)
@@ -144,7 +153,8 @@ def build_auth_headers(
 ) -> dict[str, str]:
     """Build standard auth headers for governance API calls.
 
-    Single source of truth for the SDK's outbound governance requests.
+    Compatibility helper for API-key/DID requests. Workload headers are composed
+    by the base SDK after bootstrap and token exchange.
     """
     headers = {
         "Authorization": f"Bearer {api_key}",
@@ -194,6 +204,8 @@ class GovernanceClient:
         agent_did: str | None = None,
         agent_private_key: str | None = None,
         gate: GovernanceGate | None = None,
+        workload_private_key: str | None = None,
+        core_client: EvaluationClient | None = None,
     ) -> None:
         self._api_url = api_url.rstrip("/")
         self._api_key = api_key
@@ -205,13 +217,37 @@ class GovernanceClient:
             did=agent_did,
             private_key=agent_private_key,
         )
+        workload_private_key = parse_optional_workload_private_key(workload_private_key)
+        self._core_client = core_client
+        self._owns_core_client = core_client is None and workload_private_key is not None
+        if self._owns_core_client:
+            from openbox_core.client import EvaluationClient
+            from openbox_core.identity import AgentIdentity
+
+            self._core_client = EvaluationClient(
+                self._api_url,
+                self._api_key,
+                timeout_seconds=timeout,
+                on_api_error=on_api_error,
+                identity=(
+                    AgentIdentity.from_private_key(
+                        self._agent_identity.did, self._agent_identity.private_key
+                    )
+                    if self._agent_identity
+                    else None
+                ),
+                workload_private_key=workload_private_key,
+                sdk_version=_SDK_PACKAGE_VERSION,
+                sdk_engine="langgraph",
+            )
         # Optional base-SDK gate. When wired (by the handler, from a core
         # runtime built off the SAME api_url/api_key/timeout/on_api_error),
         # `evaluate_event`'s ASYNC path routes lifecycle events through it
         # instead of this client's own httpx transport — see `evaluate_event`.
         # `None` (the default) preserves the exact legacy transport/serialization
         # for every existing caller that constructs a bare `GovernanceClient()`.
-        # `evaluate_event_sync` (sync middleware hooks) is unaffected either way.
+        # The handler also lends its runtime client for sync/raw calls and
+        # approval polling, so every path shares the same workload token cache.
         self._gate = gate
         # Deduplication: prevent sending the same (activity_id, event_type) twice
         # within the same workflow run. Keyed by (workflow_id, run_id) so it resets
@@ -234,7 +270,9 @@ class GovernanceClient:
         return self._sync_client
 
     async def close(self) -> None:
-        """Close the underlying HTTP clients."""
+        """Close owned HTTP clients; a borrowed runtime client belongs to its runtime."""
+        if self._owns_core_client and self._core_client is not None:
+            await self._core_client.aclose()
         if self._client and not self._client.is_closed:
             await self._client.aclose()
         self._client = None
@@ -255,6 +293,13 @@ class GovernanceClient:
         """
         from openbox_langgraph.errors import OpenBoxAuthError
 
+        if self._core_client is not None:
+            try:
+                await self._core_client.avalidate_api_key()
+            except _CoreOpenBoxConfigError as exc:
+                _raise_core_error(exc)
+            return
+
         try:
             client = self._get_client()
             response = await client.get(
@@ -341,13 +386,25 @@ class GovernanceClient:
             )
 
         if self._gate is not None:
-            return await _gate_evaluate(self._gate, event, self._on_api_error)
+            try:
+                return await _gate_evaluate(self._gate, event, self._on_api_error)
+            except OpenBoxError:
+                self._forget_failed_event(event, server_event_type)
+                raise
 
         payload = event.to_dict()
         payload["event_type"] = server_event_type
         payload["task_queue"] = event.task_queue or "langgraph"
         payload["source"] = "workflow-telemetry"
 
+        if self._core_client is not None:
+            try:
+                result = await self._core_client.aevaluate(payload)
+            except (_CoreOpenBoxConfigError, _CoreGovernanceAPIError) as exc:
+                self._forget_failed_event(event, server_event_type)
+                _raise_core_error(exc)
+            return _collapse_client_synthesized_fallback(result)
+
         try:
             client = self._get_client()
             body = _json_body(payload)
@@ -372,9 +429,7 @@ class GovernanceClient:
         except OpenBoxNetworkError:
             raise
         except Exception as e:
-            return _network_fallback_result(
-                self._on_api_error, f"Governance API unreachable: {e}"
-            )
+            return _network_fallback_result(self._on_api_error, f"Governance API unreachable: {e}")
 
     def evaluate_event_sync(
         self, event: LangChainGovernanceEvent
@@ -395,6 +450,14 @@ class GovernanceClient:
         payload["task_queue"] = event.task_queue or "langgraph"
         payload["source"] = "workflow-telemetry"
 
+        if self._core_client is not None:
+            try:
+                result = self._core_client.evaluate(payload)
+            except (_CoreOpenBoxConfigError, _CoreGovernanceAPIError) as exc:
+                self._forget_failed_event(event, server_event_type)
+                _raise_core_error(exc)
+            return _collapse_client_synthesized_fallback(result)
+
         if os.environ.get("OPENBOX_DEBUG") == "1":
             import json
 
@@ -427,18 +490,26 @@ class GovernanceClient:
         except OpenBoxNetworkError:
             raise
         except Exception as e:
-            return _network_fallback_result(
-                self._on_api_error, f"Governance API unreachable: {e}"
-            )
+            return _network_fallback_result(self._on_api_error, f"Governance API unreachable: {e}")
 
     async def poll_approval(self, params: ApprovalPollParams) -> ApprovalResponse | None:
         """Poll for HITL approval status.
 
-        Returns `None` on network failure so the caller can retry.
+        Returns `None` on ordinary polling transport failure so the caller can retry.
+        Workload authentication and bootstrap failures always raise.
 
         Args:
             params: Identifiers for the pending approval.
         """
+        if self._core_client is not None:
+            try:
+                result = await self._core_client.apoll_approval(
+                    params.workflow_id, params.run_id, params.activity_id
+                )
+            except _CoreOpenBoxConfigError as exc:
+                _raise_core_error(exc)
+            return ApprovalResponse.from_result(result) if result is not None else None
+
         try:
             client = self._get_client()
             body = _json_body(
@@ -482,8 +553,17 @@ class GovernanceClient:
         by the caller (no event_type translation needed).
 
         Args:
-            payload: The raw dict to POST to `/api/v1/governance/evaluate`.
+            payload: The raw dict to POST to the identity-appropriate evaluate route.
         """
+        if self._core_client is not None:
+            try:
+                result = await self._core_client.aevaluate(payload)
+            except (_CoreOpenBoxConfigError, _CoreGovernanceAPIError) as exc:
+                _raise_core_error(exc)
+            if _collapse_client_synthesized_fallback(result) is None:
+                return None
+            return dict(result.raw)
+
         if os.environ.get("OPENBOX_DEBUG") == "1":
             import json
 
@@ -535,6 +615,11 @@ class GovernanceClient:
     # Private helpers
     # ─────────────────────────────────────────────────────────────
 
+    def _forget_failed_event(self, event: LangChainGovernanceEvent, event_type: str) -> None:
+        """A caller retrying a rejected event must authenticate again, never dedup to ALLOW."""
+        if event.activity_id and self._dedup_run == (event.workflow_id, event.run_id):
+            self._dedup_sent.discard((event.activity_id, event_type))
+
     def _headers(self, *, method: str, pathname: str, body: bytes | str | None) -> dict[str, str]:
         return build_auth_headers(
             self._api_key,

openbox_langgraph/config.py

@@ -14,10 +14,12 @@ from openbox_langgraph.errors import (
     OpenBoxConfigError,
     OpenBoxInsecureURLError,
     OpenBoxNetworkError,
+    _raise_core_error,
 )
 from openbox_langgraph.identity import (
     AgentIdentityConfig,
     parse_optional_agent_identity_config,
+    parse_optional_workload_private_key,
 )
 from openbox_langgraph.types import DEFAULT_HITL_CONFIG, HITLConfig
 
@@ -193,6 +195,7 @@ class _GlobalConfigState:
     governance_timeout: float = 30.0  # seconds
     agent_did: str | None = None
     agent_private_key: str | None = None
+    workload_private_key: str | None = field(default=None, repr=False)
 
     def configure(
         self,
@@ -200,12 +203,14 @@ class _GlobalConfigState:
         api_key: str,
         governance_timeout: float = 30.0,
         agent_identity: AgentIdentityConfig | None = None,
+        workload_private_key: str | None = None,
     ) -> None:
         self.api_url = api_url.rstrip("/")
         self.api_key = api_key
         self.governance_timeout = governance_timeout
         self.agent_did = agent_identity.did if agent_identity else None
         self.agent_private_key = agent_identity.private_key if agent_identity else None
+        self.workload_private_key = workload_private_key
 
     def __repr__(self) -> str:
         if self.api_key and len(self.api_key) > 8:
@@ -293,6 +298,41 @@ def _validate_api_key_with_server(
 # ═══════════════════════════════════════════════════════════════════
 
 
+def _validate_workload_identity_with_server(
+    api_url: str,
+    api_key: str,
+    timeout: float,
+    workload_private_key: str,
+    agent_identity: AgentIdentityConfig | None,
+) -> None:
+    """Let the base SDK bootstrap, exchange a token, and validate the active identity."""
+    from openbox_core.client import EvaluationClient
+    from openbox_core.errors import OpenBoxConfigError as CoreOpenBoxConfigError
+    from openbox_core.identity import AgentIdentity
+
+    from openbox_langgraph.client import _SDK_PACKAGE_VERSION
+
+    client = EvaluationClient(
+        api_url,
+        api_key,
+        timeout_seconds=timeout,
+        workload_private_key=workload_private_key,
+        identity=(
+            AgentIdentity.from_private_key(agent_identity.did, agent_identity.private_key)
+            if agent_identity
+            else None
+        ),
+        sdk_version=_SDK_PACKAGE_VERSION,
+        sdk_engine="langgraph",
+    )
+    try:
+        client.validate_api_key()
+    except CoreOpenBoxConfigError as exc:
+        _raise_core_error(exc)
+    finally:
+        client.close()
+
+
 def initialize(
     api_url: str,
     api_key: str,
@@ -300,6 +340,7 @@ def initialize(
     validate: bool = True,
     agent_did: str | None = None,
     agent_private_key: str | None = None,
+    workload_private_key: str | None = None,
 ) -> None:
     """Initialize the OpenBox LangGraph SDK with credentials.
 
@@ -314,6 +355,8 @@ def initialize(
         agent_did: Optional OpenBox agent DID. Falls back to `OPENBOX_AGENT_DID`.
         agent_private_key: Optional raw Ed25519 private key seed. Falls back to
             `OPENBOX_AGENT_PRIVATE_KEY`.
+        workload_private_key: PKCS8 PEM RSA key for IAM v3. Falls back to
+            `OPENBOX_LANGGRAPH_WORKLOAD_PRIVATE_KEY`, then `OPENBOX_WORKLOAD_PRIVATE_KEY`.
     """
     import os
 
@@ -337,7 +380,17 @@ def initialize(
     except OpenBoxConfigError:
         raise
 
-    if validate:
+    if workload_private_key is None:
+        workload_private_key = os.environ.get(
+            "OPENBOX_LANGGRAPH_WORKLOAD_PRIVATE_KEY", os.environ.get("OPENBOX_WORKLOAD_PRIVATE_KEY")
+        )
+    workload_private_key = parse_optional_workload_private_key(workload_private_key)
+
+    if validate and workload_private_key is not None:
+        _validate_workload_identity_with_server(
+            api_url.rstrip("/"), api_key, governance_timeout, workload_private_key, agent_identity
+        )
+    elif validate:
         _validate_api_key_with_server(
             api_url.rstrip("/"),
             api_key,
@@ -350,6 +403,7 @@ def initialize(
         api_key,
         governance_timeout,
         agent_identity,
+        workload_private_key,
     )
 
     _get_logger().info(f"OpenBox LangGraph SDK initialized with API URL: {api_url}")