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

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

@@ -1,11 +1,13 @@
 Metadata-Version: 2.4
 Name: harness-sdk-python
-Version: 0.1.0
+Version: 0.2.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
+Provides-Extra: postgres
+Requires-Dist: langgraph-checkpoint-postgres>=2.0.0; extra == 'postgres'
 Description-Content-Type: text/markdown
 
 # harness-sdk-python

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

@@ -1,12 +1,15 @@
 [project]
 name = "harness-sdk-python"
-version = "0.1.0"
+version = "0.2.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"]
 
+[project.optional-dependencies]
+postgres = ["langgraph-checkpoint-postgres>=2.0.0"]
+
 [project.urls]
 Repository = "https://github.com/assistant-ui/harness-sdk"
 

harness_sdk_python-0.2.0/src/harness_sdk/fenced_postgres.py

@@ -0,0 +1,220 @@
+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.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-Za-z_][A-Za-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 an unquoted 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: errors.Error) -> bool:
+    return exc.sqlstate == FENCE_LOST_SQLSTATE
+
+
+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,
+            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 not _is_fence_lost(exc):
+                raise
+            raise self._fence_lost() from exc
+
+    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)

harness_sdk_python-0.2.0/tests/test_fenced_postgres.py

@@ -0,0 +1,255 @@
+import os
+import shutil
+import socket
+import subprocess
+import time
+import uuid
+
+import psycopg
+import pytest
+from langgraph.checkpoint.base import empty_checkpoint
+from psycopg.rows import dict_row
+
+from harness_sdk.fenced_postgres import AsyncFencedPostgresSaver, FenceLost
+
+THREAD = "thread-1"
+
+
+def _free_port():
+    with socket.socket() as s:
+        s.bind(("127.0.0.1", 0))
+        return s.getsockname()[1]
+
+
+def _wait_ready(dsn, timeout=120.0):
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        try:
+            psycopg.connect(dsn, connect_timeout=2).close()
+            return
+        except psycopg.OperationalError:
+            time.sleep(0.5)
+    raise RuntimeError(f"postgres at {dsn} did not become ready within {timeout}s")
+
+
+@pytest.fixture(scope="session")
+def postgres_dsn():
+    if dsn := os.environ.get("HARNESS_TEST_POSTGRES_DSN"):
+        yield dsn
+        return
+    docker = shutil.which("docker")
+    if docker is None or subprocess.run(
+        [docker, "info"], capture_output=True
+    ).returncode:
+        pytest.skip(
+            "postgres unavailable: set HARNESS_TEST_POSTGRES_DSN or start docker"
+        )
+    port = _free_port()
+    container = subprocess.run(
+        [
+            docker, "run", "--rm", "-d",
+            "-e", "POSTGRES_PASSWORD=harness",
+            "-p", f"127.0.0.1:{port}:5432",
+            "postgres:16",
+        ],
+        capture_output=True, text=True, check=True,
+    ).stdout.strip()
+    dsn = f"postgresql://postgres:harness@127.0.0.1:{port}/postgres"
+    try:
+        _wait_ready(dsn)
+        yield dsn
+    finally:
+        subprocess.run([docker, "rm", "-f", container], capture_output=True)
+
+
+@pytest.fixture
+def dsn(postgres_dsn):
+    dbname = f"fence_test_{uuid.uuid4().hex[:12]}"
+    with psycopg.connect(postgres_dsn, autocommit=True) as admin:
+        admin.execute(f'CREATE DATABASE "{dbname}"')
+    yield postgres_dsn.rsplit("/", 1)[0] + f"/{dbname}"
+    with psycopg.connect(postgres_dsn, autocommit=True) as admin:
+        admin.execute(f'DROP DATABASE "{dbname}" WITH (FORCE)')
+
+
+@pytest.fixture
+async def aconnect(dsn):
+    conns = []
+
+    async def _connect():
+        conn = await psycopg.AsyncConnection.connect(
+            dsn, autocommit=True, prepare_threshold=0, row_factory=dict_row
+        )
+        conns.append(conn)
+        return conn
+
+    yield _connect
+    for conn in conns:
+        await conn.close()
+
+
+def make_config(thread_id=THREAD, checkpoint_id=None):
+    return {
+        "configurable": {
+            "thread_id": thread_id,
+            "checkpoint_ns": "",
+            "checkpoint_id": checkpoint_id or str(uuid.uuid4()),
+        }
+    }
+
+
+def thread_config(thread_id=THREAD):
+    return {"configurable": {"thread_id": thread_id}}
+
+
+async def aput_checkpoint(saver, thread_id=THREAD):
+    checkpoint = empty_checkpoint()
+    checkpoint["id"] = str(uuid.uuid4())
+    return await saver.aput(
+        make_config(thread_id), checkpoint, {"source": "loop", "step": 1}, {}
+    )
+
+
+async def aacquire(aconnect, thread_id=THREAD, **kwargs):
+    saver = await AsyncFencedPostgresSaver.acquire(
+        await aconnect(), thread_id, create_fence_table=True, **kwargs
+    )
+    await saver.setup()
+    return saver
+
+
+class TestFencing:
+    async def test_acquire_bumps_epoch_monotonically(self, aconnect):
+        first = await aacquire(aconnect)
+        second = await aacquire(aconnect)
+        assert first.epoch == 1
+        assert second.epoch == 2
+
+    async def test_same_instance_writes_keep_working(self, aconnect):
+        saver = await aacquire(aconnect)
+        for _ in range(10):
+            config = await aput_checkpoint(saver)
+            await saver.aput_writes(config, [("channel", "value")], "task-1")
+        assert len([c async for c in saver.alist(thread_config())]) == 10
+        await saver.adelete_thread(THREAD)
+        assert [c async for c in saver.alist(thread_config())] == []
+
+    async def test_new_instance_fences_old_writes(self, aconnect):
+        old = await aacquire(aconnect)
+        config = await aput_checkpoint(old)
+        new = await aacquire(aconnect)
+        with pytest.raises(FenceLost):
+            await aput_checkpoint(old)
+        with pytest.raises(FenceLost):
+            await old.aput_writes(config, [("channel", "value")], "task-1")
+        with pytest.raises(FenceLost):
+            await old.adelete_thread(THREAD)
+        await aput_checkpoint(new)
+
+    async def test_fenced_write_leaves_no_trace(self, aconnect):
+        old = await aacquire(aconnect)
+        await aput_checkpoint(old)
+        await aacquire(aconnect)
+        with pytest.raises(FenceLost):
+            await aput_checkpoint(old)
+        assert len([c async for c in old.alist(thread_config())]) == 1
+
+    async def test_reacquire_kills_old_handle_of_same_worker(self, aconnect):
+        conn = await aconnect()
+        old = await AsyncFencedPostgresSaver.acquire(
+            conn, THREAD, create_fence_table=True
+        )
+        await old.setup()
+        reacquired = await AsyncFencedPostgresSaver.acquire(
+            conn, THREAD, create_fence_table=True
+        )
+        assert reacquired.epoch == old.epoch + 1
+        with pytest.raises(FenceLost):
+            await aput_checkpoint(old)
+        await aput_checkpoint(reacquired)
+
+    async def test_fences_are_per_thread(self, aconnect):
+        other = await aacquire(aconnect, "thread-2")
+        mine = await aacquire(aconnect)
+        await aacquire(aconnect, "thread-2")
+        await aput_checkpoint(mine)
+        with pytest.raises(FenceLost):
+            await aput_checkpoint(other, "thread-2")
+
+    async def test_rejects_writes_for_other_threads(self, aconnect):
+        saver = await aacquire(aconnect)
+        with pytest.raises(ValueError):
+            await aput_checkpoint(saver, "thread-2")
+        with pytest.raises(ValueError):
+            await saver.aput_writes(make_config("thread-2"), [("c", "v")], "task-1")
+        with pytest.raises(ValueError):
+            await saver.adelete_thread("thread-2")
+
+    async def test_acquire_rejects_invalid_thread_id(self, aconnect):
+        with pytest.raises(ValueError):
+            await AsyncFencedPostgresSaver.acquire(
+                await aconnect(), "", create_fence_table=True
+            )
+        with pytest.raises(ValueError):
+            await AsyncFencedPostgresSaver.acquire(
+                await aconnect(), None, create_fence_table=True
+            )
+
+
+class TestFenceOptions:
+    async def test_custom_fence_table(self, aconnect):
+        old = await aacquire(aconnect, fence_table="my_fence")
+        new = await aacquire(aconnect, fence_table="my_fence")
+        assert new.epoch == old.epoch + 1
+        with pytest.raises(FenceLost):
+            await aput_checkpoint(old)
+        await aput_checkpoint(new)
+
+    async def test_acquire_without_create_requires_existing_table(self, aconnect):
+        with pytest.raises(RuntimeError, match="does not exist"):
+            await AsyncFencedPostgresSaver.acquire(
+                await aconnect(), THREAD, create_fence_table=False
+            )
+
+    async def test_acquire_without_create_requires_existing_function(self, aconnect):
+        await aacquire(aconnect)
+        conn = await aconnect()
+        await conn.execute("DROP FUNCTION checkpoint_fence_assert(text, bigint)")
+        with pytest.raises(RuntimeError, match="fence function"):
+            await AsyncFencedPostgresSaver.acquire(
+                conn, THREAD, create_fence_table=False
+            )
+
+    async def test_acquire_without_create_uses_existing_table(self, aconnect):
+        await aacquire(aconnect)
+        saver = await AsyncFencedPostgresSaver.acquire(
+            await aconnect(), THREAD, create_fence_table=False
+        )
+        assert saver.epoch == 2
+
+    async def test_acquire_rejects_invalid_fence_table(self, aconnect):
+        with pytest.raises(ValueError):
+            await AsyncFencedPostgresSaver.acquire(
+                await aconnect(), THREAD, create_fence_table=True, fence_table="bad table;"
+            )
+
+
+class TestOnFenceLost:
+    async def test_callback_fires_once_across_write_attempts(self, aconnect):
+        calls = []
+        old = await aacquire(aconnect, on_fence_lost=lambda: calls.append("lost"))
+        await aacquire(aconnect)
+        for _ in range(3):
+            with pytest.raises(FenceLost):
+                await aput_checkpoint(old)
+        with pytest.raises(FenceLost):
+            await old.adelete_thread(THREAD)
+        assert calls == ["lost"]
+
+    async def test_no_callback_still_raises(self, aconnect):
+        old = await aacquire(aconnect)
+        await aacquire(aconnect)
+        with pytest.raises(FenceLost) as excinfo:
+            await aput_checkpoint(old)
+        assert excinfo.value.thread_id == THREAD
+        assert excinfo.value.epoch == 1