onceonly-sdk 2.0.1__py3-none-any.whl → 3.0.0__py3-none-any.whl

onceonly/__init__.py

@@ -1,6 +1,14 @@
 from .version import __version__
 from .client import OnceOnly, create_client
-from .models import CheckLockResult
+from .models import (
+    CheckLockResult,
+    Policy,
+    AgentStatus,
+    AgentLogItem,
+    AgentMetrics,
+)
+from .decorators import idempotent, idempotent_ai
+from .ai_models import AiToolResult
 from .exceptions import (
     OnceOnlyError,
     UnauthorizedError,
@@ -13,12 +21,29 @@ from .exceptions import (
 __all__ = [
     "OnceOnly",
     "create_client",
+
+    # Core idempotency
     "CheckLockResult",
+
+    # Agent Governance
+    "Policy",
+    "AgentStatus",
+    "AgentLogItem",
+    "AgentMetrics",
+
+    # Decorators
+    "idempotent",
+    "idempotent_ai",
+
+    "AiToolResult",
+
+    # Errors
     "OnceOnlyError",
     "UnauthorizedError",
     "OverLimitError",
     "RateLimitError",
     "ValidationError",
     "ApiError",
+
     "__version__",
 ]

onceonly/_http.py

@@ -78,10 +78,21 @@ def request_with_retries_sync(
     base_backoff: float,
     max_backoff: float,
 ) -> httpx.Response:
+    RETRYABLE_STATUS = {429, 500, 502, 503, 504}
+
     attempt = 0
     while True:
-        resp = fn()
-        if resp.status_code != 429 or attempt >= max_retries:
+        try:
+            resp = fn()
+        except httpx.RequestError:
+            if attempt >= max_retries:
+                raise
+            sleep_s = min(max_backoff, base_backoff * (2**attempt))
+            time.sleep(max(0.0, float(sleep_s)))
+            attempt += 1
+            continue
+
+        if resp.status_code not in RETRYABLE_STATUS or attempt >= max_retries:
             return resp
 
         ra = _parse_retry_after(resp)
@@ -97,10 +108,21 @@ async def request_with_retries_async(
     base_backoff: float,
     max_backoff: float,
 ) -> httpx.Response:
+    RETRYABLE_STATUS = {429, 500, 502, 503, 504}
+
     attempt = 0
     while True:
-        resp = await fn()
-        if resp.status_code != 429 or attempt >= max_retries:
+        try:
+            resp = await fn()
+        except httpx.RequestError:
+            if attempt >= max_retries:
+                raise
+            sleep_s = min(max_backoff, base_backoff * (2**attempt))
+            await asyncio.sleep(max(0.0, float(sleep_s)))
+            attempt += 1
+            continue
+
+        if resp.status_code not in RETRYABLE_STATUS or attempt >= max_retries:
             return resp
 
         ra = _parse_retry_after(resp)

onceonly/_util.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+import json
+
 from dataclasses import is_dataclass, asdict
 from typing import Any, Dict, Mapping, Optional, Union
 
@@ -39,7 +41,7 @@ def to_metadata_dict(metadata: Optional[MetadataLike]) -> Optional[Dict[str, Any
     # mapping
     if isinstance(metadata, Mapping):
         try:
-            return dict(metadata)
+            return json.loads(json.dumps(dict(metadata), ensure_ascii=False, default=str))
         except Exception:
             return {"value": str(metadata)}
 

onceonly/ai.py

@@ -2,10 +2,11 @@ from __future__ import annotations
 
 import asyncio
 import time
+import inspect
 import logging
-from typing import Any, Dict, Optional, Awaitable, Callable
-
 import httpx
+import threading
+from typing import Any, Dict, Optional, Awaitable, Callable, Union
 
 from ._http import (
     parse_json_or_raise,
@@ -13,7 +14,7 @@ from ._http import (
     request_with_retries_async,
 )
 from ._util import to_metadata_dict, MetadataLike
-from .ai_models import AiRun, AiStatus, AiResult
+from .ai_models import AiRun, AiStatus, AiResult, AiToolResult
 
 logger = logging.getLogger("onceonly")
 
@@ -51,17 +52,111 @@ class AiClient:
         self._retry_backoff = float(retry_backoff)
         self._retry_max_backoff = float(retry_max_backoff)
 
+    @staticmethod
+    def _result_to_dict(value: Any) -> Optional[Dict[str, Any]]:
+        if value is None:
+            return None
+        if isinstance(value, dict):
+            return value
+
+        md = getattr(value, "model_dump", None)  # pydantic v2
+        if callable(md):
+            try:
+                out = md()
+                return out if isinstance(out, dict) else {"data": out}
+            except Exception:
+                return {"value": str(value)}
+
+        dct = getattr(value, "dict", None)  # pydantic v1
+        if callable(dct):
+            try:
+                out = dct()
+                return out if isinstance(out, dict) else {"data": out}
+            except Exception:
+                return {"value": str(value)}
+
+        if hasattr(value, "__dataclass_fields__"):
+            try:
+                import dataclasses
+                out = dataclasses.asdict(value)
+                return out if isinstance(out, dict) else {"data": out}
+            except Exception:
+                return {"value": str(value)}
+
+        return {"value": str(value)}
+
+    def _start_heartbeat_thread(
+            self,
+            *,
+            key: str,
+            lease_id: str,
+            ttl: Optional[int],
+            extend_every: float,
+    ) -> "threading.Event":
+        stop = threading.Event()
+
+        def _loop() -> None:
+            while not stop.is_set():
+                try:
+                    self.extend(key=key, lease_id=lease_id, ttl=ttl)
+                except Exception:
+                    pass
+                stop.wait(max(1.0, float(extend_every)))
+
+        t = threading.Thread(target=_loop, name="onceonly-ai-heartbeat", daemon=True)
+        t.start()
+        return stop
+
+    async def _start_heartbeat_task(
+            self,
+            *,
+            key: str,
+            lease_id: str,
+            ttl: Optional[int],
+            extend_every: float,
+    ) -> "asyncio.Task[None]":
+        async def _loop() -> None:
+            while True:
+                try:
+                    await self.extend_async(key=key, lease_id=lease_id, ttl=ttl)
+                except Exception:
+                    pass
+                await asyncio.sleep(max(1.0, float(extend_every)))
+
+        return asyncio.create_task(_loop())
+
     # ------------------------------------------------------------------
     # High-level: /ai/run + /ai/status + /ai/result
     # ------------------------------------------------------------------
 
-    def run(self, key: str, ttl: Optional[int] = None, metadata: Optional[MetadataLike] = None) -> AiRun:
-        payload: Dict[str, Any] = {"key": key}
-        if ttl is not None:
-            payload["ttl"] = int(ttl)
-        md = to_metadata_dict(metadata)
-        if md is not None:
-            payload["metadata"] = md
+    def run(
+        self,
+        key: Optional[str] = None,
+        ttl: Optional[int] = None,
+        metadata: Optional[MetadataLike] = None,
+        *,
+        agent_id: Optional[str] = None,
+        tool: Optional[str] = None,
+        args: Optional[Dict[str, Any]] = None,
+        spend_usd: Optional[float] = None,
+    ) -> Union[AiRun, AiToolResult]:
+        if key is None:
+            if not agent_id or not tool:
+                raise ValueError("ai.run requires key=... OR agent_id=... and tool=...")
+            payload: Dict[str, Any] = {"agent_id": str(agent_id), "tool": str(tool)}
+            if args is not None:
+                payload["args"] = dict(args)
+            if spend_usd is not None:
+                payload["spend_usd"] = float(spend_usd)
+        else:
+            if agent_id or tool or args or spend_usd is not None:
+                raise ValueError("ai.run: provide either key=... OR agent_id/tool, not both")
+            payload = {"key": key}
+            if ttl is not None:
+                payload["ttl"] = int(ttl)
+            md = to_metadata_dict(metadata)
+            if md is not None:
+                payload["metadata"] = md
 
         resp = request_with_retries_sync(
             lambda: self._c.post("/ai/run", json=payload),
@@ -71,6 +166,9 @@ class AiClient:
         )
         data = parse_json_or_raise(resp)
 
+        if "allowed" in data or "decision" in data:
+            return AiToolResult.from_dict(data)
+
         logger.debug("ai.run key=%s status=%s version=%s", key, data.get("status"), data.get("version"))
         return AiRun.from_dict(data)
 
@@ -105,15 +203,14 @@ class AiClient:
         timeout: float = 60.0,
         poll_min: float = 0.5,
         poll_max: float = 5.0,
+        auto_extend: bool = True,
+        extend_every: float = 30.0,
+        lease_id: Optional[str] = None,
+        ttl: Optional[int] = None,
     ) -> AiResult:
-        """
-        Polls /ai/status until completed/failed or timeout.
-
-        NOTE: We do NOT return status="timeout" because backend model statuses
-        are usually limited to: not_found|in_progress|completed|failed.
-        Instead we return status="failed" with error_code="timeout".
-        """
         t0 = time.time()
+        last_ext = 0.0
+
         while True:
             st = self.status(key)
             if st.status in ("completed", "failed"):
@@ -122,32 +219,218 @@ class AiClient:
             if time.time() - t0 >= timeout:
                 return AiResult(ok=False, status="failed", key=key, error_code="timeout")
 
+            # heartbeat (best-effort)
+            if auto_extend and lease_id:
+                now = time.time()
+                if now - last_ext >= float(extend_every):
+                    try:
+                        self.extend(key=key, lease_id=lease_id, ttl=ttl)
+                    except Exception:
+                        pass
+                    last_ext = now
+
             sleep_s = st.retry_after_sec if isinstance(st.retry_after_sec, int) else poll_min
             sleep_s = max(poll_min, min(poll_max, float(sleep_s)))
             time.sleep(sleep_s)
 
     def run_and_wait(
         self,
-        key: str,
+        key: Optional[str] = None,
         *,
         ttl: Optional[int] = None,
         metadata: Optional[MetadataLike] = None,
+        agent_id: Optional[str] = None,
+        tool: Optional[str] = None,
+        args: Optional[Dict[str, Any]] = None,
+        spend_usd: Optional[float] = None,
+        timeout: float = 60.0,
+        poll_min: float = 0.5,
+        poll_max: float = 5.0,
+        auto_extend: bool = True,
+        extend_every: float = 30.0,
+    ) -> Union[AiResult, AiToolResult]:
+        if key is None and (agent_id or tool):
+            return self.run(
+                key=None,
+                agent_id=agent_id,
+                tool=tool,
+                args=args,
+                spend_usd=spend_usd,
+            )
+
+        if key is None:
+            raise ValueError("ai.run_and_wait requires key=... OR agent_id/tool for tool execution")
+
+        run = self.run(key=key, ttl=ttl, metadata=metadata)
+        return self.wait(
+            key=key,
+            timeout=timeout,
+            poll_min=poll_min,
+            poll_max=poll_max,
+            auto_extend=auto_extend,
+            extend_every=extend_every,
+            lease_id=run.lease_id,
+            ttl=ttl,
+        )
+
+    def run_tool(
+        self,
+        *,
+        agent_id: str,
+        tool: str,
+        args: Optional[Dict[str, Any]] = None,
+        spend_usd: Optional[float] = None,
+    ) -> AiToolResult:
+        res = self.run(
+            key=None,
+            agent_id=agent_id,
+            tool=tool,
+            args=args,
+            spend_usd=spend_usd,
+        )
+        assert isinstance(res, AiToolResult)
+        return res
+
+    def run_fn(
+        self,
+        key: str,
+        fn: Callable[[], Any],
+        *,
+        ttl: int = 300,
+        metadata: Optional[MetadataLike] = None,
+        extend_every: float = 30.0,
+        wait_on_conflict: bool = True,
+        timeout: float = 60.0,
+        poll_min: float = 0.5,
+        poll_max: float = 5.0,
+        error_code: str = "fn_error",
+    ) -> AiResult:
+        """
+        Local execution, exactly-once:
+          - POST /ai/lease (charged only if acquired)
+          - Heartbeat: /ai/extend while fn runs (best effort)
+          - POST /ai/complete or /ai/fail
+          - Returns /ai/result (typed)
+        """
+        lease = self.lease(key=key, ttl=int(ttl), metadata=metadata)
+        status = str(lease.get("status") or "").lower()
+
+        if status == "acquired":
+            lease_id = lease.get("lease_id")
+            if not lease_id:
+                return AiResult(ok=False, status="failed", key=key, error_code="missing_lease_id")
+
+            stop = self._start_heartbeat_thread(key=key, lease_id=str(lease_id), ttl=int(ttl),
+                                                extend_every=extend_every)
+            try:
+                out = fn()
+                res_dict = self._result_to_dict(out)
+                self.complete(key=key, lease_id=str(lease_id), result=res_dict)
+            except Exception:
+                try:
+                    self.fail(key=key, lease_id=str(lease_id), error_code=error_code)
+                except Exception:
+                    pass
+                raise
+            finally:
+                stop.set()
+
+            return self.result(key)
+
+        if status in ("completed", "failed"):
+            return self.result(key)
+
+        # in_progress / locked / etc.
+        if wait_on_conflict:
+            return self.wait(key=key, timeout=timeout, poll_min=poll_min, poll_max=poll_max)
+
+        return AiResult(ok=False, status=status or "in_progress", key=key, error_code="not_acquired")
+
+    async def run_fn_async(
+        self,
+        key: str,
+        fn: Callable[[], Any],
+        *,
+        ttl: int = 300,
+        metadata: Optional[MetadataLike] = None,
+        extend_every: float = 30.0,
+        wait_on_conflict: bool = True,
         timeout: float = 60.0,
         poll_min: float = 0.5,
         poll_max: float = 5.0,
+        error_code: str = "fn_error",
     ) -> AiResult:
-        self.run(key=key, ttl=ttl, metadata=metadata)
-        return self.wait(key=key, timeout=timeout, poll_min=poll_min, poll_max=poll_max)
+        lease = await self.lease_async(key=key, ttl=int(ttl), metadata=metadata)
+        status = str(lease.get("status") or "").lower()
+
+        if status == "acquired":
+            lease_id = lease.get("lease_id")
+            if not lease_id:
+                return AiResult(ok=False, status="failed", key=key, error_code="missing_lease_id")
+
+            hb = await self._start_heartbeat_task(
+                key=key, lease_id=str(lease_id), ttl=int(ttl), extend_every=extend_every
+            )
+            try:
+                out = fn()
+                if inspect.isawaitable(out):
+                    out = await out
+
+                res_dict = self._result_to_dict(out)
+                await self.complete_async(key=key, lease_id=str(lease_id), result=res_dict)
+            except Exception:
+                try:
+                    await self.fail_async(key=key, lease_id=str(lease_id), error_code=error_code)
+                except Exception:
+                    pass
+                raise
+            finally:
+                hb.cancel()
+                try:
+                    await hb
+                except Exception:
+                    pass
+
+            return await self.result_async(key)
+
+        if status in ("completed", "failed"):
+            return await self.result_async(key)
+
+        if wait_on_conflict:
+            return await self.wait_async(key=key, timeout=timeout, poll_min=poll_min, poll_max=poll_max)
+
+        return AiResult(ok=False, status=status or "in_progress", key=key, error_code="not_acquired")
 
     # -------------------- async high-level --------------------
 
-    async def run_async(self, key: str, ttl: Optional[int] = None, metadata: Optional[MetadataLike] = None) -> AiRun:
-        payload: Dict[str, Any] = {"key": key}
-        if ttl is not None:
-            payload["ttl"] = int(ttl)
-        md = to_metadata_dict(metadata)
-        if md is not None:
-            payload["metadata"] = md
+    async def run_async(
+        self,
+        key: Optional[str] = None,
+        ttl: Optional[int] = None,
+        metadata: Optional[MetadataLike] = None,
+        *,
+        agent_id: Optional[str] = None,
+        tool: Optional[str] = None,
+        args: Optional[Dict[str, Any]] = None,
+        spend_usd: Optional[float] = None,
+    ) -> Union[AiRun, AiToolResult]:
+        if key is None:
+            if not agent_id or not tool:
+                raise ValueError("ai.run_async requires key=... OR agent_id=... and tool=...")
+            payload: Dict[str, Any] = {"agent_id": str(agent_id), "tool": str(tool)}
+            if args is not None:
+                payload["args"] = dict(args)
+            if spend_usd is not None:
+                payload["spend_usd"] = float(spend_usd)
+        else:
+            if agent_id or tool or args or spend_usd is not None:
+                raise ValueError("ai.run_async: provide either key=... OR agent_id/tool, not both")
+            payload = {"key": key}
+            if ttl is not None:
+                payload["ttl"] = int(ttl)
+            md = to_metadata_dict(metadata)
+            if md is not None:
+                payload["metadata"] = md
 
         c = await self._get_ac()
         resp = await request_with_retries_async(
@@ -158,6 +441,9 @@ class AiClient:
         )
         data = parse_json_or_raise(resp)
 
+        if "allowed" in data or "decision" in data:
+            return AiToolResult.from_dict(data)
+
         logger.debug("ai.run_async key=%s status=%s version=%s", key, data.get("status"), data.get("version"))
         return AiRun.from_dict(data)
 
@@ -194,8 +480,14 @@ class AiClient:
         timeout: float = 60.0,
         poll_min: float = 0.5,
         poll_max: float = 5.0,
+        auto_extend: bool = True,
+        extend_every: float = 30.0,
+        lease_id: Optional[str] = None,
+        ttl: Optional[int] = None,
     ) -> AiResult:
         t0 = time.time()
+        last_ext = 0.0
+
         while True:
             st = await self.status_async(key)
             if st.status in ("completed", "failed"):
@@ -204,22 +496,77 @@ class AiClient:
             if time.time() - t0 >= timeout:
                 return AiResult(ok=False, status="failed", key=key, error_code="timeout")
 
+            # heartbeat (best-effort)
+            if auto_extend and lease_id:
+                now = time.time()
+                if now - last_ext >= float(extend_every):
+                    try:
+                        await self.extend_async(key=key, lease_id=lease_id, ttl=ttl)
+                    except Exception:
+                        pass
+                    last_ext = now
+
             sleep_s = st.retry_after_sec if isinstance(st.retry_after_sec, int) else poll_min
             sleep_s = max(poll_min, min(poll_max, float(sleep_s)))
             await asyncio.sleep(sleep_s)
 
     async def run_and_wait_async(
         self,
-        key: str,
+        key: Optional[str] = None,
         *,
         ttl: Optional[int] = None,
         metadata: Optional[MetadataLike] = None,
+        agent_id: Optional[str] = None,
+        tool: Optional[str] = None,
+        args: Optional[Dict[str, Any]] = None,
+        spend_usd: Optional[float] = None,
         timeout: float = 60.0,
         poll_min: float = 0.5,
         poll_max: float = 5.0,
-    ) -> AiResult:
-        await self.run_async(key=key, ttl=ttl, metadata=metadata)
-        return await self.wait_async(key=key, timeout=timeout, poll_min=poll_min, poll_max=poll_max)
+        auto_extend: bool = True,
+        extend_every: float = 30.0,
+    ) -> Union[AiResult, AiToolResult]:
+        if key is None and (agent_id or tool):
+            return await self.run_async(
+                key=None,
+                agent_id=agent_id,
+                tool=tool,
+                args=args,
+                spend_usd=spend_usd,
+            )
+
+        if key is None:
+            raise ValueError("ai.run_and_wait_async requires key=... OR agent_id/tool for tool execution")
+
+        run = await self.run_async(key=key, ttl=ttl, metadata=metadata)
+        return await self.wait_async(
+            key=key,
+            timeout=timeout,
+            poll_min=poll_min,
+            poll_max=poll_max,
+            auto_extend=auto_extend,
+            extend_every=extend_every,
+            lease_id=run.lease_id,
+            ttl=ttl,
+        )
+
+    async def run_tool_async(
+        self,
+        *,
+        agent_id: str,
+        tool: str,
+        args: Optional[Dict[str, Any]] = None,
+        spend_usd: Optional[float] = None,
+    ) -> AiToolResult:
+        res = await self.run_async(
+            key=None,
+            agent_id=agent_id,
+            tool=tool,
+            args=args,
+            spend_usd=spend_usd,
+        )
+        assert isinstance(res, AiToolResult)
+        return res
 
     # ------------------------------------------------------------------
     # Low-level lease API (sync) - returns raw dicts (backend models)

onceonly/ai_models.py

@@ -15,6 +15,9 @@ class AiRun:
     key: str
     lease_id: Optional[str] = None
     version: int = 0
+    ttl: Optional[int] = None
+    ttl_left: Optional[int] = None
+    first_seen_at: Optional[str] = None
     charged: Optional[int] = None
     usage: Optional[int] = None
     limit: Optional[int] = None
@@ -32,6 +35,9 @@ class AiRun:
             key=str(d.get("key") or ""),
             lease_id=d.get("lease_id"),
             version=int(d.get("version") or 0),
+            ttl=d.get("ttl"),
+            ttl_left=d.get("ttl_left"),
+            first_seen_at=d.get("first_seen_at"),
             charged=d.get("charged"),
             usage=d.get("usage"),
             limit=d.get("limit"),
@@ -95,3 +101,24 @@ class AiResult:
             error_code=d.get("error_code"),
             done_at=d.get("done_at"),
         )
+
+
+@dataclass(frozen=True)
+class AiToolResult:
+    ok: bool
+    allowed: bool
+    decision: str
+    policy_reason: Optional[str] = None
+    risk_level: Optional[str] = None
+    result: Optional[Dict[str, Any]] = None
+
+    @staticmethod
+    def from_dict(d: Dict[str, Any]) -> "AiToolResult":
+        return AiToolResult(
+            ok=bool(d.get("ok", False)),
+            allowed=bool(d.get("allowed", False)),
+            decision=str(d.get("decision") or ""),
+            policy_reason=d.get("policy_reason"),
+            risk_level=d.get("risk_level"),
+            result=_dict_or_none(d.get("result")),
+        )