harness-sdk-python 0.1.0__tar.gz → 0.3.0__tar.gz

{harness_sdk_python-0.1.0 → harness_sdk_python-0.3.0}/PKG-INFO

@@ -1,11 +1,13 @@
 Metadata-Version: 2.4
 Name: harness-sdk-python
-Version: 0.1.0
+Version: 0.3.0
 Summary: RunManager: the harness-sdk runs subsystem for Python Statewire hosts
 Project-URL: Repository, https://github.com/assistant-ui/harness-sdk
 License-Expression: MIT
 Requires-Python: <4.0,>=3.11
-Requires-Dist: statewire<0.3,>=0.2.0
+Requires-Dist: statewire<0.4,>=0.3.0
+Provides-Extra: postgres
+Requires-Dist: langgraph-checkpoint-postgres>=2.0.0; extra == 'postgres'
 Description-Content-Type: text/markdown
 
 # harness-sdk-python
@@ -19,7 +21,7 @@ class MyHost(Statewire):
     async def lifespan(self):
         self.state = initial_state()
         self.runs = RunManager(
-            state=self.state["runs"],
+            state=self.state,
             start=self._start,
             get_message_meta=self._get_message_meta,
             create_task=self.create_task,

{harness_sdk_python-0.1.0 → harness_sdk_python-0.3.0}/README.md

@@ -9,7 +9,7 @@ class MyHost(Statewire):
     async def lifespan(self):
         self.state = initial_state()
         self.runs = RunManager(
-            state=self.state["runs"],
+            state=self.state,
             start=self._start,
             get_message_meta=self._get_message_meta,
             create_task=self.create_task,

{harness_sdk_python-0.1.0 → harness_sdk_python-0.3.0}/pyproject.toml

@@ -1,11 +1,14 @@
 [project]
 name = "harness-sdk-python"
-version = "0.1.0"
+version = "0.3.0"
 description = "RunManager: the harness-sdk runs subsystem for Python Statewire hosts"
 readme = "README.md"
 license = "MIT"
 requires-python = ">=3.11,<4.0"
-dependencies = ["statewire>=0.2.0,<0.3"]
+dependencies = ["statewire>=0.3.0,<0.4"]
+
+[project.optional-dependencies]
+postgres = ["langgraph-checkpoint-postgres>=2.0.0"]
 
 [project.urls]
 Repository = "https://github.com/assistant-ui/harness-sdk"

harness_sdk_python-0.3.0/src/harness_sdk/fenced_postgres.py

@@ -0,0 +1,247 @@
+import re
+from collections.abc import AsyncIterator, Callable, Sequence
+from contextlib import asynccontextmanager
+from typing import Any
+
+try:
+    from langchain_core.runnables import RunnableConfig
+    from langgraph.checkpoint.base import SerializerProtocol
+    from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver, _ainternal
+    from psycopg import AsyncPipeline, errors
+    from psycopg.pq import TransactionStatus
+    from psycopg.cursor_async import AsyncCursor
+    from psycopg.rows import DictRow, dict_row
+    from psycopg.sql import SQL, Identifier, Literal
+except ImportError as exc:
+    raise ImportError(
+        "harness_sdk.fenced_postgres requires the 'postgres' extra: "
+        "pip install 'harness-sdk-python[postgres]'"
+    ) from exc
+
+DEFAULT_FENCE_TABLE = "checkpoint_fence"
+FENCE_LOST_SQLSTATE = "FL001"
+
+_IDENTIFIER_RE = re.compile(r"^[a-z_][a-z0-9_]*$")
+
+
+class FenceLost(RuntimeError):
+    def __init__(self, thread_id: str, epoch: int) -> None:
+        super().__init__(
+            f"fence lost for thread {thread_id!r}: epoch {epoch} is no longer current"
+        )
+        self.thread_id = thread_id
+        self.epoch = epoch
+
+
+def _require_thread_id(thread_id: str) -> str:
+    if not isinstance(thread_id, str) or not thread_id:
+        raise ValueError(f"thread_id must be a non-empty string, got {thread_id!r}")
+    return thread_id
+
+
+def _require_fence_table(fence_table: str) -> str:
+    if not isinstance(fence_table, str) or not _IDENTIFIER_RE.match(fence_table):
+        raise ValueError(
+            f"fence_table must be a lowercase SQL identifier, got {fence_table!r}"
+        )
+    return fence_table
+
+
+def _assert_fn(fence_table: str) -> str:
+    return f"{fence_table}_assert"
+
+
+def _create_sql(fence_table: str) -> SQL:
+    return SQL(
+        "CREATE TABLE IF NOT EXISTS {} (thread_id TEXT PRIMARY KEY, epoch BIGINT NOT NULL)"
+    ).format(Identifier(fence_table))
+
+
+# FOR UPDATE holds the fence row against concurrent epoch bumps until the
+# surrounding write transaction ends; the RAISE aborts the pipeline/transaction
+def _create_fn_sql(fence_table: str) -> SQL:
+    return SQL(
+        "CREATE OR REPLACE FUNCTION {fn}(p_thread_id TEXT, p_epoch BIGINT) "
+        "RETURNS void LANGUAGE plpgsql AS $$\n"
+        "BEGIN\n"
+        "  PERFORM 1 FROM {t} WHERE thread_id = p_thread_id AND epoch = p_epoch FOR UPDATE;\n"
+        "  IF NOT FOUND THEN\n"
+        "    RAISE EXCEPTION 'fence lost for thread %: epoch % is not current', "
+        "p_thread_id, p_epoch USING ERRCODE = {code};\n"
+        "  END IF;\n"
+        "END;\n"
+        "$$"
+    ).format(
+        fn=Identifier(_assert_fn(fence_table)),
+        t=Identifier(fence_table),
+        code=Literal(FENCE_LOST_SQLSTATE),
+    )
+
+
+def _bump_sql(fence_table: str) -> SQL:
+    return SQL(
+        "INSERT INTO {t} (thread_id, epoch) VALUES (%s, 1) "
+        "ON CONFLICT (thread_id) DO UPDATE SET epoch = {t}.epoch + 1 "
+        "RETURNING epoch"
+    ).format(t=Identifier(fence_table))
+
+
+def _assert_sql(fence_table: str) -> SQL:
+    return SQL("SELECT {}(%s, %s)").format(Identifier(_assert_fn(fence_table)))
+
+
+def _is_fence_lost(exc: BaseException | None) -> bool:
+    seen: set[int] = set()
+    while exc is not None and id(exc) not in seen:
+        seen.add(id(exc))
+        if isinstance(exc, errors.Error) and exc.sqlstate == FENCE_LOST_SQLSTATE:
+            return True
+        exc = exc.__cause__ or exc.__context__
+    return False
+
+
+def _missing_fence_error(missing: str) -> RuntimeError:
+    return RuntimeError(
+        f"{missing} does not exist; "
+        "pass create_fence_table=True or create it before acquiring"
+    )
+
+
+class AsyncFencedPostgresSaver(AsyncPostgresSaver):
+    def __init__(
+        self,
+        conn: _ainternal.Conn,
+        thread_id: str,
+        epoch: int,
+        *,
+        pipe: AsyncPipeline | None = None,
+        serde: SerializerProtocol | None = None,
+        fence_table: str = DEFAULT_FENCE_TABLE,
+        on_fence_lost: Callable[[], None] | None = None,
+    ) -> None:
+        super().__init__(conn, pipe=pipe, serde=serde)
+        self.thread_id = _require_thread_id(thread_id)
+        self.epoch = epoch
+        self.fence_table = _require_fence_table(fence_table)
+        self.on_fence_lost = on_fence_lost
+        self._fence_lost_notified = False
+
+    @classmethod
+    async def acquire(
+        cls,
+        conn: _ainternal.Conn,
+        thread_id: str,
+        *,
+        create_fence_table: bool,
+        pipe: AsyncPipeline | None = None,
+        serde: SerializerProtocol | None = None,
+        fence_table: str = DEFAULT_FENCE_TABLE,
+        on_fence_lost: Callable[[], None] | None = None,
+    ) -> "AsyncFencedPostgresSaver":
+        _require_thread_id(thread_id)
+        _require_fence_table(fence_table)
+        async with _ainternal.get_connection(conn) as c:
+            if c.info.transaction_status != TransactionStatus.IDLE:
+                raise RuntimeError(
+                    "acquire requires a connection with no transaction in progress"
+                )
+            async with (
+                c.transaction(),
+                c.cursor(binary=True, row_factory=dict_row) as cur,
+            ):
+                if create_fence_table:
+                    await cur.execute(_create_sql(fence_table))
+                    await cur.execute(_create_fn_sql(fence_table))
+                else:
+                    await cur.execute(
+                        "SELECT to_regclass(%s) AS tbl, to_regprocedure(%s) AS fn",
+                        (fence_table, f"{_assert_fn(fence_table)}(text, bigint)"),
+                    )
+                    row = await cur.fetchone()
+                    if row["tbl"] is None:
+                        raise _missing_fence_error(f"fence table {fence_table!r}")
+                    if row["fn"] is None:
+                        raise _missing_fence_error(
+                            f"fence function {_assert_fn(fence_table)!r}"
+                        )
+                await cur.execute(_bump_sql(fence_table), (thread_id,))
+                epoch = (await cur.fetchone())["epoch"]
+        if pipe is not None:
+            await pipe.sync()
+        return cls(
+            conn,
+            thread_id,
+            epoch,
+            pipe=pipe,
+            serde=serde,
+            fence_table=fence_table,
+            on_fence_lost=on_fence_lost,
+        )
+
+    def _fence_lost(self) -> FenceLost:
+        if self.on_fence_lost is not None and not self._fence_lost_notified:
+            self._fence_lost_notified = True
+            self.on_fence_lost()
+        return FenceLost(self.thread_id, self.epoch)
+
+    def _bound_to(self, thread_id: Any) -> None:
+        if str(thread_id) != self.thread_id:
+            raise ValueError(
+                f"saver is fenced to thread {self.thread_id!r}, got {thread_id!r}"
+            )
+
+    @asynccontextmanager
+    async def _cursor(
+        self, *, pipeline: bool = False
+    ) -> AsyncIterator[AsyncCursor[DictRow]]:
+        if not pipeline:
+            async with super()._cursor(pipeline=False) as cur:
+                yield cur
+            return
+        try:
+            async with super()._cursor(pipeline=True) as cur:
+                await cur.execute(_assert_sql(self.fence_table), (self.thread_id, self.epoch))
+                yield cur
+        except errors.Error as exc:
+            if _is_fence_lost(exc):
+                raise self._fence_lost() from exc
+            # An aborted pipeline can surface before the assert's FL001 does;
+            # the fence row decides which it was.
+            if isinstance(exc, errors.PipelineAborted) and await self._fence_stale():
+                raise self._fence_lost() from exc
+            raise
+
+    async def _fence_stale(self) -> bool:
+        async with super()._cursor(pipeline=False) as cur:
+            await cur.execute(
+                SQL("SELECT epoch FROM {} WHERE thread_id = %s").format(
+                    Identifier(self.fence_table)
+                ),
+                (self.thread_id,),
+            )
+            row = await cur.fetchone()
+        return row is None or row["epoch"] != self.epoch
+
+    async def aput(
+        self,
+        config: RunnableConfig,
+        checkpoint: Any,
+        metadata: Any,
+        new_versions: Any,
+    ) -> RunnableConfig:
+        self._bound_to(config["configurable"]["thread_id"])
+        return await super().aput(config, checkpoint, metadata, new_versions)
+
+    async def aput_writes(
+        self,
+        config: RunnableConfig,
+        writes: Sequence[tuple[str, Any]],
+        task_id: str,
+        task_path: str = "",
+    ) -> None:
+        self._bound_to(config["configurable"]["thread_id"])
+        await super().aput_writes(config, writes, task_id, task_path)
+
+    async def adelete_thread(self, thread_id: str) -> None:
+        self._bound_to(thread_id)
+        await super().adelete_thread(thread_id)