langgraph-runtime-inmem 0.9.0__py3-none-any.whl → 0.20.1__py3-none-any.whl

langgraph_runtime_inmem/ops.py

@@ -11,7 +11,6 @@ import uuid
 from collections import defaultdict
 from collections.abc import AsyncIterator, Sequence
 from contextlib import asynccontextmanager
-from copy import deepcopy
 from datetime import UTC, datetime, timedelta
 from typing import Any, Literal, cast
 from uuid import UUID, uuid4
@@ -29,6 +28,7 @@ from langgraph_runtime_inmem.checkpoint import Checkpointer
 from langgraph_runtime_inmem.database import InMemConnectionProto, connect
 from langgraph_runtime_inmem.inmem_stream import (
     THREADLESS_KEY,
+    ContextQueue,
     Message,
     get_stream_manager,
 )
@@ -58,12 +58,13 @@ if typing.TYPE_CHECKING:
         Thread,
         ThreadSelectField,
         ThreadStatus,
+        ThreadStreamMode,
         ThreadUpdateResponse,
     )
     from langgraph_api.schema import Interrupt as InterruptSchema
-    from langgraph_api.serde import Fragment
     from langgraph_api.utils import AsyncConnectionProto
 
+StreamHandler = ContextQueue
 
 logger = structlog.stdlib.get_logger(__name__)
 
@@ -140,6 +141,7 @@ class Assistants(Authenticated):
         conn: InMemConnectionProto,
         *,
         graph_id: str | None,
+        name: str | None,
         metadata: MetadataInput,
         limit: int,
         offset: int,
@@ -163,6 +165,7 @@ class Assistants(Authenticated):
             assistant
             for assistant in assistants
             if (not graph_id or assistant["graph_id"] == graph_id)
+            and (not name or name.lower() in assistant["name"].lower())
             and (not metadata or is_jsonb_contained(assistant["metadata"], metadata))
             and (not filters or _check_filter_match(assistant["metadata"], filters))
         ]
@@ -228,7 +231,7 @@ class Assistants(Authenticated):
                 if assistant["assistant_id"] == assistant_id and (
                     not filters or _check_filter_match(assistant["metadata"], filters)
                 ):
-                    yield assistant
+                    yield copy.deepcopy(assistant)
 
         return _yield_result()
 
@@ -247,6 +250,8 @@ class Assistants(Authenticated):
         description: str | None = None,
     ) -> AsyncIterator[Assistant]:
         """Insert an assistant."""
+        from langgraph_api.graph import GRAPHS
+
         assistant_id = _ensure_uuid(assistant_id)
         metadata = metadata if metadata is not None else {}
         filters = await Assistants.handle_event(
@@ -268,6 +273,9 @@ class Assistants(Authenticated):
                 detail="Cannot specify both configurable and context. Prefer setting context alone. Context was introduced in LangGraph 0.6.0 and is the long term planned replacement for configurable.",
             )
 
+        if graph_id not in GRAPHS:
+            raise HTTPException(status_code=404, detail=f"Graph {graph_id} not found")
+
         # Keep config and context up to date with one another
         if config.get("configurable"):
             context = config["configurable"]
@@ -555,6 +563,8 @@ class Assistants(Authenticated):
                 "metadata": version_data["metadata"],
                 "version": version_data["version"],
                 "updated_at": datetime.now(UTC),
+                "name": version_data["name"],
+                "description": version_data["description"],
             }
         )
 
@@ -618,6 +628,7 @@ class Assistants(Authenticated):
         conn: InMemConnectionProto,
         *,
         graph_id: str | None = None,
+        name: str | None = None,
         metadata: MetadataInput = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> int:
@@ -635,6 +646,7 @@ class Assistants(Authenticated):
         for assistant in conn.store["assistants"]:
             if (
                 (not graph_id or assistant["graph_id"] == graph_id)
+                and (not name or name.lower() in assistant["name"].lower())
                 and (
                     not metadata or is_jsonb_contained(assistant["metadata"], metadata)
                 )
@@ -738,6 +750,7 @@ class Threads(Authenticated):
     async def search(
         conn: InMemConnectionProto,
         *,
+        ids: list[str] | list[UUID] | None = None,
         metadata: MetadataInput,
         values: MetadataInput,
         status: ThreadStatus | None,
@@ -765,7 +778,19 @@ class Threads(Authenticated):
         )
 
         # Apply filters
+        id_set: set[UUID] | None = None
+        if ids:
+            id_set = set()
+            for i in ids:
+                try:
+                    id_set.add(_ensure_uuid(i))
+                except Exception:
+                    raise HTTPException(
+                        status_code=400, detail="Invalid thread ID " + str(i)
+                    ) from None
         for thread in threads:
+            if id_set is not None and thread.get("thread_id") not in id_set:
+                continue
             if filters and not _check_filter_match(thread["metadata"], filters):
                 continue
 
@@ -943,6 +968,7 @@ class Threads(Authenticated):
         thread_id: UUID,
         *,
         metadata: MetadataValue,
+        ttl: ThreadTTLConfig | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Thread]:
         """Update a thread."""
@@ -1215,13 +1241,23 @@ class Threads(Authenticated):
         """Create a copy of an existing thread."""
         thread_id = _ensure_uuid(thread_id)
         new_thread_id = uuid4()
-        filters = await Threads.handle_event(
+        read_filters = await Threads.handle_event(
             ctx,
             "read",
             Auth.types.ThreadsRead(
+                thread_id=thread_id,
+            ),
+        )
+        # Assert that the user has permissions to create a new thread.
+        # (We don't actually need the filters.)
+        await Threads.handle_event(
+            ctx,
+            "create",
+            Auth.types.ThreadsCreate(
                 thread_id=new_thread_id,
             ),
         )
+
         async with conn.pipeline():
             # Find the original thread in our store
             original_thread = next(
@@ -1230,8 +1266,8 @@ class Threads(Authenticated):
 
             if not original_thread:
                 return _empty_generator()
-            if filters and not _check_filter_match(
-                original_thread["metadata"], filters
+            if read_filters and not _check_filter_match(
+                original_thread["metadata"], read_filters
             ):
                 return _empty_generator()
 
@@ -1240,7 +1276,7 @@ class Threads(Authenticated):
                 "thread_id": new_thread_id,
                 "created_at": datetime.now(tz=UTC),
                 "updated_at": datetime.now(tz=UTC),
-                "metadata": deepcopy(original_thread["metadata"]),
+                "metadata": copy.deepcopy(original_thread["metadata"]),
                 "status": "idle",
                 "config": {},
             }
@@ -1327,7 +1363,14 @@ class Threads(Authenticated):
                 )
 
             metadata = thread.get("metadata", {})
-            thread_config = thread.get("config", {})
+            thread_config = cast(dict[str, Any], thread.get("config", {}))
+            thread_config = {
+                **thread_config,
+                "configurable": {
+                    **thread_config.get("configurable", {}),
+                    **config.get("configurable", {}),
+                },
+            }
 
             # Fallback to graph_id from run if not in thread metadata
             graph_id = metadata.get("graph_id")
@@ -1377,6 +1420,7 @@ class Threads(Authenticated):
             """Add state to a thread."""
             from langgraph_api.graph import get_graph
             from langgraph_api.schema import ThreadUpdateResponse
+            from langgraph_api.state import state_snapshot_to_thread_state
             from langgraph_api.store import get_store
             from langgraph_api.utils import fetchone
 
@@ -1414,6 +1458,13 @@ class Threads(Authenticated):
                     status_code=409,
                     detail=f"Thread {thread_id} has in-flight runs: {pending_runs}",
                 )
+            thread_config = {
+                **thread_config,
+                "configurable": {
+                    **thread_config.get("configurable", {}),
+                    **config.get("configurable", {}),
+                },
+            }
 
             # Fallback to graph_id from run if not in thread metadata
             graph_id = metadata.get("graph_id")
@@ -1454,6 +1505,19 @@ class Threads(Authenticated):
                             thread["values"] = state.values
                             break
 
+                    # Publish state update event
+                    from langgraph_api.serde import json_dumpb
+
+                    event_data = {
+                        "state": state_snapshot_to_thread_state(state),
+                        "thread_id": str(thread_id),
+                    }
+                    await Threads.Stream.publish(
+                        thread_id,
+                        "state_update",
+                        json_dumpb(event_data),
+                    )
+
                     return ThreadUpdateResponse(
                         checkpoint=next_config["configurable"],
                         # Including deprecated fields
@@ -1496,7 +1560,14 @@ class Threads(Authenticated):
                 thread_iter, not_found_detail=f"Thread {thread_id} not found."
             )
 
-            thread_config = thread["config"]
+            thread_config = cast(dict[str, Any], thread["config"])
+            thread_config = {
+                **thread_config,
+                "configurable": {
+                    **thread_config.get("configurable", {}),
+                    **config.get("configurable", {}),
+                },
+            }
             metadata = thread["metadata"]
 
             if not thread:
@@ -1543,6 +1614,19 @@ class Threads(Authenticated):
                             thread["values"] = state.values
                             break
 
+                    # Publish state update event
+                    from langgraph_api.serde import json_dumpb
+
+                    event_data = {
+                        "state": state,
+                        "thread_id": str(thread_id),
+                    }
+                    await Threads.Stream.publish(
+                        thread_id,
+                        "state_update",
+                        json_dumpb(event_data),
+                    )
+
                     return ThreadUpdateResponse(
                         checkpoint=next_config["configurable"],
                     )
@@ -1584,7 +1668,14 @@ class Threads(Authenticated):
             if not _check_filter_match(thread_metadata, filters):
                 return []
 
-            thread_config = thread["config"]
+            thread_config = cast(dict[str, Any], thread["config"])
+            thread_config = {
+                **thread_config,
+                "configurable": {
+                    **thread_config.get("configurable", {}),
+                    **config.get("configurable", {}),
+                },
+            }
             # If graph_id exists, get state history
             if graph_id := thread_metadata.get("graph_id"):
                 async with get_graph(
@@ -1613,7 +1704,9 @@ class Threads(Authenticated):
 
             return []
 
-    class Stream:
+    class Stream(Authenticated):
+        resource = "threads"
+
         @staticmethod
         async def subscribe(
             conn: InMemConnectionProto | AsyncConnectionProto,
@@ -1626,6 +1719,13 @@ class Threads(Authenticated):
 
             # Create new queues only for runs not yet seen
             thread_id = _ensure_uuid(thread_id)
+
+            # Add thread stream queue
+            if thread_id not in seen_runs:
+                queue = await stream_manager.add_thread_stream(thread_id)
+                queues.append((thread_id, queue))
+                seen_runs.add(thread_id)
+
             for run in conn.store["runs"]:
                 if run["thread_id"] == thread_id:
                     run_id = run["run_id"]
@@ -1641,9 +1741,32 @@ class Threads(Authenticated):
             thread_id: UUID,
             *,
             last_event_id: str | None = None,
+            stream_modes: list[ThreadStreamMode],
+            ctx: Auth.types.BaseAuthContext | None = None,
         ) -> AsyncIterator[tuple[bytes, bytes, bytes | None]]:
             """Stream the thread output."""
-            from langgraph_api.serde import json_loads
+            await Threads.Stream.check_thread_stream_auth(thread_id, ctx)
+
+            from langgraph_api.utils.stream_codec import (
+                decode_stream_message,
+            )
+
+            def should_filter_event(event_name: str, message_bytes: bytes) -> bool:
+                """Check if an event should be filtered out based on stream_modes."""
+                if "run_modes" in stream_modes and event_name != "state_update":
+                    return False
+                if "state_update" in stream_modes and event_name == "state_update":
+                    return False
+                if "lifecycle" in stream_modes and event_name == "metadata":
+                    try:
+                        message_data = orjson.loads(message_bytes)
+                        if message_data.get("status") == "run_done":
+                            return False
+                        if "attempt" in message_data and "run_id" in message_data:
+                            return False
+                    except (orjson.JSONDecodeError, TypeError):
+                        pass
+                return True
 
             stream_manager = get_stream_manager()
             seen_runs: set[UUID] = set()
@@ -1673,23 +1796,24 @@ class Threads(Authenticated):
 
                         # Yield sorted events
                         for message, run_id in all_events:
-                            data = json_loads(message.data)
-                            event_name = data["event"]
-                            message_content = data["message"]
-
-                            if event_name == "control":
-                                if message_content == b"done":
-                                    yield (
-                                        b"metadata",
-                                        orjson.dumps(
-                                            {"status": "run_done", "run_id": run_id}
-                                        ),
-                                        message.id,
+                            decoded = decode_stream_message(
+                                message.data, channel=message.topic
+                            )
+                            event_bytes = decoded.event_bytes
+                            message_bytes = decoded.message_bytes
+
+                            if event_bytes == b"control":
+                                if message_bytes == b"done":
+                                    event_bytes = b"metadata"
+                                    message_bytes = orjson.dumps(
+                                        {"status": "run_done", "run_id": run_id}
                                     )
-                            else:
+                            if not should_filter_event(
+                                event_bytes.decode("utf-8"), message_bytes
+                            ):
                                 yield (
-                                    event_name.encode(),
-                                    base64.b64decode(message_content),
+                                    event_bytes,
+                                    message_bytes,
                                     message.id,
                                 )
 
@@ -1708,28 +1832,27 @@ class Threads(Authenticated):
                                 message = await asyncio.wait_for(
                                     queue.get(), timeout=0.2
                                 )
-                                data = json_loads(message.data)
-                                event_name = data["event"]
-                                message_content = data["message"]
-
-                                if event_name == "control":
-                                    if message_content == b"done":
-                                        # Extract run_id from topic
-                                        topic = message.topic.decode()
-                                        run_id = topic.split("run:")[1].split(":")[0]
-                                        yield (
-                                            b"metadata",
-                                            orjson.dumps(
-                                                {"status": "run_done", "run_id": run_id}
-                                            ),
-                                            message.id,
-                                        )
-                                else:
-                                    yield (
-                                        event_name.encode(),
-                                        base64.b64decode(message_content),
-                                        message.id,
+                                decoded = decode_stream_message(
+                                    message.data, channel=message.topic
+                                )
+                                event = decoded.event_bytes
+                                event_name = event.decode("utf-8")
+                                payload = decoded.message_bytes
+
+                                if event == b"control" and payload == b"done":
+                                    topic = message.topic.decode()
+                                    run_id = topic.split("run:")[1].split(":")[0]
+                                    meta_event = b"metadata"
+                                    meta_payload = orjson.dumps(
+                                        {"status": "run_done", "run_id": run_id}
                                     )
+                                    if not should_filter_event(
+                                        "metadata", meta_payload
+                                    ):
+                                        yield (meta_event, meta_payload, message.id)
+                                else:
+                                    if not should_filter_event(event_name, payload):
+                                        yield (event, payload, message.id)
 
                             except TimeoutError:
                                 continue
@@ -1758,6 +1881,41 @@ class Threads(Authenticated):
                         # Ignore cleanup errors
                         pass
 
+        @staticmethod
+        async def publish(
+            thread_id: UUID | str,
+            event: str,
+            message: bytes,
+        ) -> None:
+            """Publish a thread-level event to the thread stream."""
+            from langgraph_api.utils.stream_codec import STREAM_CODEC
+
+            topic = f"thread:{thread_id}:stream".encode()
+
+            stream_manager = get_stream_manager()
+            payload = STREAM_CODEC.encode(event, message)
+            await stream_manager.put_thread(
+                str(thread_id), Message(topic=topic, data=payload)
+            )
+
+        @staticmethod
+        async def check_thread_stream_auth(
+            thread_id: UUID,
+            ctx: Auth.types.BaseAuthContext | None = None,
+        ) -> None:
+            async with connect() as conn:
+                filters = await Threads.Stream.handle_event(
+                    ctx,
+                    "read",
+                    Auth.types.ThreadsRead(thread_id=thread_id),
+                )
+                if filters:
+                    thread = await Threads._get_with_filters(
+                        cast(InMemConnectionProto, conn), thread_id, filters
+                    )
+                    if not thread:
+                        raise HTTPException(status_code=404, detail="Thread not found")
+
     @staticmethod
     async def count(
         conn: InMemConnectionProto,
@@ -1821,38 +1979,37 @@ class Runs(Authenticated):
         if not pending_runs and not running_runs:
             return {
                 "n_pending": 0,
-                "max_age_secs": None,
-                "med_age_secs": None,
+                "pending_runs_wait_time_max_secs": None,
+                "pending_runs_wait_time_med_secs": None,
                 "n_running": 0,
             }
 
-        # Get all creation timestamps
-        created_times = [run.get("created_at") for run in (pending_runs + running_runs)]
-        created_times = [
-            t for t in created_times if t is not None
-        ]  # Filter out None values
-
-        if not created_times:
-            return {
-                "n_pending": len(pending_runs),
-                "n_running": len(running_runs),
-                "max_age_secs": None,
-                "med_age_secs": None,
-            }
-
-        # Find oldest (max age)
-        oldest_time = min(created_times)  # Earliest timestamp = oldest run
-
-        # Find median age
-        sorted_times = sorted(created_times)
-        median_idx = len(sorted_times) // 2
-        median_time = sorted_times[median_idx]
+        now = datetime.now(UTC)
+        pending_waits: list[float] = []
+        for run in pending_runs:
+            created_at = run.get("created_at")
+            if not isinstance(created_at, datetime):
+                continue
+            if created_at.tzinfo is None:
+                created_at = created_at.replace(tzinfo=UTC)
+            pending_waits.append((now - created_at).total_seconds())
+
+        max_pending_wait = max(pending_waits) if pending_waits else None
+        if pending_waits:
+            sorted_waits = sorted(pending_waits)
+            half = len(sorted_waits) // 2
+            if len(sorted_waits) % 2 == 1:
+                med_pending_wait = sorted_waits[half]
+            else:
+                med_pending_wait = (sorted_waits[half - 1] + sorted_waits[half]) / 2
+        else:
+            med_pending_wait = None
 
         return {
             "n_pending": len(pending_runs),
             "n_running": len(running_runs),
-            "max_age_secs": oldest_time,
-            "med_age_secs": median_time,
+            "pending_runs_wait_time_max_secs": max_pending_wait,
+            "pending_runs_wait_time_med_secs": med_pending_wait,
         }
 
     @staticmethod
@@ -1916,12 +2073,16 @@ class Runs(Authenticated):
     @asynccontextmanager
     @staticmethod
     async def enter(
-        run_id: UUID, thread_id: UUID | None, loop: asyncio.AbstractEventLoop
+        run_id: UUID,
+        thread_id: UUID | None,
+        loop: asyncio.AbstractEventLoop,
+        resumable: bool,
     ) -> AsyncIterator[ValueEvent]:
         """Enter a run, listen for cancellation while running, signal when done."
         This method should be called as a context manager by a worker executing a run.
         """
         from langgraph_api.asyncio import SimpleTaskGroup, ValueEvent
+        from langgraph_api.utils.stream_codec import STREAM_CODEC
 
         stream_manager = get_stream_manager()
         # Get control queue for this run (normal queue is created during run creation)
@@ -1941,12 +2102,14 @@ class Runs(Authenticated):
             )
             await stream_manager.put(run_id, thread_id, control_message)
 
-            # Signal done to all subscribers
+            # Signal done to all subscribers using stream codec
             stream_message = Message(
                 topic=f"run:{run_id}:stream".encode(),
-                data={"event": "control", "message": b"done"},
+                data=STREAM_CODEC.encode("control", b"done"),
+            )
+            await stream_manager.put(
+                run_id, thread_id, stream_message, resumable=resumable
             )
-            await stream_manager.put(run_id, thread_id, stream_message)
 
             # Remove the control_queue (normal queue is cleaned up during run deletion)
             await stream_manager.remove_control_queue(run_id, thread_id, control_queue)
@@ -1988,7 +2151,6 @@ class Runs(Authenticated):
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Run]:
         """Create a run."""
-        from langgraph_api.config import FF_RICH_THREADS
         from langgraph_api.schema import Run, Thread
 
         assistant_id = _ensure_uuid(assistant_id)
@@ -2004,6 +2166,7 @@ class Runs(Authenticated):
         run_id = _ensure_uuid(run_id) if run_id else None
         metadata = metadata if metadata is not None else {}
         config = kwargs.get("config", {})
+        temporary = kwargs.get("temporary", False)
 
         # Handle thread creation/update
         existing_thread = next(
@@ -2013,7 +2176,7 @@ class Runs(Authenticated):
             ctx,
             "create_run",
             Auth.types.RunsCreate(
-                thread_id=thread_id,
+                thread_id=None if temporary else thread_id,
                 assistant_id=assistant_id,
                 run_id=run_id,
                 status=status,
@@ -2034,49 +2197,35 @@ class Runs(Authenticated):
             # Create new thread
             if thread_id is None:
                 thread_id = uuid4()
-            if FF_RICH_THREADS:
-                thread = Thread(
-                    thread_id=thread_id,
-                    status="busy",
-                    metadata={
-                        "graph_id": assistant["graph_id"],
-                        "assistant_id": str(assistant_id),
-                        **(config.get("metadata") or {}),
-                        **metadata,
-                    },
-                    config=Runs._merge_jsonb(
-                        assistant["config"],
-                        config,
-                        {
-                            "configurable": Runs._merge_jsonb(
-                                Runs._get_configurable(assistant["config"]),
-                            )
-                        },
-                    ),
-                    created_at=datetime.now(UTC),
-                    updated_at=datetime.now(UTC),
-                    values=b"",
-                )
-            else:
-                thread = Thread(
-                    thread_id=thread_id,
-                    status="idle",
-                    metadata={
-                        "graph_id": assistant["graph_id"],
-                        "assistant_id": str(assistant_id),
-                        **(config.get("metadata") or {}),
-                        **metadata,
+
+            thread = Thread(
+                thread_id=thread_id,
+                status="busy",
+                metadata={
+                    "graph_id": assistant["graph_id"],
+                    "assistant_id": str(assistant_id),
+                    **(config.get("metadata") or {}),
+                    **metadata,
+                },
+                config=Runs._merge_jsonb(
+                    assistant["config"],
+                    config,
+                    {
+                        "configurable": Runs._merge_jsonb(
+                            Runs._get_configurable(assistant["config"]),
+                        )
                     },
-                    config={},
-                    created_at=datetime.now(UTC),
-                    updated_at=datetime.now(UTC),
-                    values=b"",
-                )
+                ),
+                created_at=datetime.now(UTC),
+                updated_at=datetime.now(UTC),
+                values=b"",
+            )
+
             await logger.ainfo("Creating thread", thread_id=thread_id)
             conn.store["threads"].append(thread)
         elif existing_thread:
             # Update existing thread
-            if FF_RICH_THREADS and existing_thread["status"] != "busy":
+            if existing_thread["status"] != "busy":
                 existing_thread["status"] = "busy"
                 existing_thread["metadata"] = Runs._merge_jsonb(
                     existing_thread["metadata"],
@@ -2253,66 +2402,6 @@ class Runs(Authenticated):
 
         return _yield_deleted()
 
-    @staticmethod
-    async def join(
-        run_id: UUID,
-        *,
-        thread_id: UUID,
-        ctx: Auth.types.BaseAuthContext | None = None,
-    ) -> Fragment:
-        """Wait for a run to complete. If already done, return immediately.
-
-        Returns:
-            the final state of the run.
-        """
-        from langgraph_api.serde import Fragment
-        from langgraph_api.utils import fetchone
-
-        async with connect() as conn:
-            # Validate ownership
-            thread_iter = await Threads.get(conn, thread_id, ctx=ctx)
-            await fetchone(thread_iter)
-        last_chunk: bytes | None = None
-        # wait for the run to complete
-        # Rely on this join's auth
-        async for mode, chunk, _ in Runs.Stream.join(
-            run_id,
-            thread_id=thread_id,
-            ctx=ctx,
-            ignore_404=True,
-            stream_mode=["values", "updates", "error"],
-        ):
-            if mode == b"values":
-                last_chunk = chunk
-            elif mode == b"updates" and b"__interrupt__" in chunk:
-                last_chunk = chunk
-            elif mode == b"error":
-                last_chunk = orjson.dumps({"__error__": orjson.Fragment(chunk)})
-        # if we received a final chunk, return it
-        if last_chunk is not None:
-            # ie. if the run completed while we were waiting for it
-            return Fragment(last_chunk)
-        else:
-            # otherwise, the run had already finished, so fetch the state from thread
-            async with connect() as conn:
-                thread_iter = await Threads.get(conn, thread_id, ctx=ctx)
-                thread = await fetchone(thread_iter)
-                if thread["status"] == "error":
-                    return Fragment(
-                        orjson.dumps({"__error__": orjson.Fragment(thread["error"])})
-                    )
-                if thread["status"] == "interrupted":
-                    # Get an interrupt for the thread. There is the case where there are multiple interrupts for the same run and we may not show the same
-                    # interrupt, but we'll always show one. Long term we should show all of them.
-                    try:
-                        interrupt_map = thread["interrupts"]
-                        interrupt = [next(iter(interrupt_map.values()))[0]]
-                        return Fragment(orjson.dumps({"__interrupt__": interrupt}))
-                    except Exception:
-                        # No interrupt, but status is interrupted from a before/after block. Default back to values.
-                        pass
-                return thread["values"]
-
     @staticmethod
     async def cancel(
         conn: InMemConnectionProto | AsyncConnectionProto,
@@ -2538,7 +2627,7 @@ class Runs(Authenticated):
         async def subscribe(
             run_id: UUID,
             thread_id: UUID | None = None,
-        ) -> asyncio.Queue:
+        ) -> ContextQueue:
             """Subscribe to the run stream, returning a queue."""
             stream_manager = get_stream_manager()
             queue = await stream_manager.add_queue(_ensure_uuid(run_id), thread_id)
@@ -2562,54 +2651,38 @@ class Runs(Authenticated):
         async def join(
             run_id: UUID,
             *,
+            stream_channel: asyncio.Queue,
             thread_id: UUID,
             ignore_404: bool = False,
             cancel_on_disconnect: bool = False,
-            stream_channel: asyncio.Queue | None = None,
             stream_mode: list[StreamMode] | StreamMode | None = None,
             last_event_id: str | None = None,
             ctx: Auth.types.BaseAuthContext | None = None,
         ) -> AsyncIterator[tuple[bytes, bytes, bytes | None]]:
             """Stream the run output."""
             from langgraph_api.asyncio import create_task
-            from langgraph_api.serde import json_loads
-
-            queue = (
-                stream_channel
-                if stream_channel
-                else await Runs.Stream.subscribe(run_id, thread_id)
-            )
+            from langgraph_api.serde import json_dumpb
+            from langgraph_api.utils.stream_codec import decode_stream_message
 
+            queue = stream_channel
             try:
                 async with connect() as conn:
-                    filters = await Runs.handle_event(
-                        ctx,
-                        "read",
-                        Auth.types.ThreadsRead(thread_id=thread_id),
-                    )
-                    if filters:
-                        thread = await Threads._get_with_filters(
-                            cast(InMemConnectionProto, conn), thread_id, filters
-                        )
-                        if not thread:
-                            raise WrappedHTTPException(
-                                HTTPException(
-                                    status_code=404, detail="Thread not found"
-                                )
-                            )
+                    try:
+                        await Runs.Stream.check_run_stream_auth(run_id, thread_id, ctx)
+                    except HTTPException as e:
+                        raise WrappedHTTPException(e) from None
                     run = await Runs.get(conn, run_id, thread_id=thread_id, ctx=ctx)
 
                     for message in get_stream_manager().restore_messages(
                         run_id, thread_id, last_event_id
                     ):
                         data, id = message.data, message.id
-
-                        data = json_loads(data)
-                        mode = data["event"]
-                        message = data["message"]
+                        decoded = decode_stream_message(data, channel=message.topic)
+                        mode = decoded.event_bytes.decode("utf-8")
+                        payload = decoded.message_bytes
 
                         if mode == "control":
-                            if message == b"done":
+                            if payload == b"done":
                                 return
                         elif (
                             not stream_mode
@@ -2622,7 +2695,7 @@ class Runs(Authenticated):
                                 and mode.startswith("messages")
                             )
                         ):
-                            yield mode.encode(), base64.b64decode(message), id
+                            yield mode.encode(), payload, id
                             logger.debug(
                                 "Replayed run event",
                                 run_id=str(run_id),
@@ -2636,13 +2709,12 @@ class Runs(Authenticated):
                             # Wait for messages with a timeout
                             message = await asyncio.wait_for(queue.get(), timeout=0.5)
                             data, id = message.data, message.id
-
-                            data = json_loads(data)
-                            mode = data["event"]
-                            message = data["message"]
+                            decoded = decode_stream_message(data, channel=message.topic)
+                            mode = decoded.event_bytes.decode("utf-8")
+                            payload = decoded.message_bytes
 
                             if mode == "control":
-                                if message == b"done":
+                                if payload == b"done":
                                     break
                             elif (
                                 not stream_mode
@@ -2655,13 +2727,13 @@ class Runs(Authenticated):
                                     and mode.startswith("messages")
                                 )
                             ):
-                                yield mode.encode(), base64.b64decode(message), id
+                                yield mode.encode(), payload, id
                                 logger.debug(
                                     "Streamed run event",
                                     run_id=str(run_id),
                                     stream_mode=mode,
                                     message_id=id,
-                                    data=message,
+                                    data=payload,
                                 )
                         except TimeoutError:
                             # Check if the run is still pending
@@ -2675,8 +2747,10 @@ class Runs(Authenticated):
                             elif run is None:
                                 yield (
                                     b"error",
-                                    HTTPException(
-                                        status_code=404, detail="Run not found"
+                                    json_dumpb(
+                                        HTTPException(
+                                            status_code=404, detail="Run not found"
+                                        )
                                     ),
                                     None,
                                 )
@@ -2693,6 +2767,25 @@ class Runs(Authenticated):
                 stream_manager = get_stream_manager()
                 await stream_manager.remove_queue(run_id, thread_id, queue)
 
+        @staticmethod
+        async def check_run_stream_auth(
+            run_id: UUID,
+            thread_id: UUID,
+            ctx: Auth.types.BaseAuthContext | None = None,
+        ) -> None:
+            async with connect() as conn:
+                filters = await Runs.handle_event(
+                    ctx,
+                    "read",
+                    Auth.types.ThreadsRead(thread_id=thread_id),
+                )
+                if filters:
+                    thread = await Threads._get_with_filters(
+                        cast(InMemConnectionProto, conn), thread_id, filters
+                    )
+                    if not thread:
+                        raise HTTPException(status_code=404, detail="Thread not found")
+
         @staticmethod
         async def publish(
             run_id: UUID | str,
@@ -2703,18 +2796,13 @@ class Runs(Authenticated):
             resumable: bool = False,
         ) -> None:
             """Publish a message to all subscribers of the run stream."""
-            from langgraph_api.serde import json_dumpb
+            from langgraph_api.utils.stream_codec import STREAM_CODEC
 
             topic = f"run:{run_id}:stream".encode()
 
             stream_manager = get_stream_manager()
-            # Send to all queues subscribed to this run_id
-            payload = json_dumpb(
-                {
-                    "event": event,
-                    "message": message,
-                }
-            )
+            # Send to all queues subscribed to this run_id using protocol frame
+            payload = STREAM_CODEC.encode(event, message)
             await stream_manager.put(
                 run_id, thread_id, Message(topic=topic, data=payload), resumable
             )
@@ -2761,7 +2849,9 @@ class Crons:
         schedule: str,
         cron_id: UUID | None = None,
         thread_id: UUID | None = None,
+        on_run_completed: Literal["delete", "keep"] | None = None,
         end_time: datetime | None = None,
+        metadata: dict | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Cron]:
         raise NotImplementedError
@@ -2852,19 +2942,154 @@ def _delete_checkpoints_for_thread(
         )
 
 
-def _check_filter_match(metadata: dict, filters: Auth.types.FilterType | None) -> bool:
+def _validate_filter_structure(
+    filters: Auth.types.FilterType | None,
+    nesting_level: int = 0,
+) -> None:
+    """Validate the structure of filter conditions without checking matches.
+
+    Args:
+        filters: The filter conditions to validate
+        nesting_level: Current depth of nested operators (max 2)
+
+    Raises:
+        HTTPException: If the filter structure is invalid
+    """
+    if nesting_level > 2:
+        raise HTTPException(
+            status_code=500,
+            detail="Your auth handler returned a filter with too many nested operators. The maximum depth for nested operators is 2. Please simplify your filter.",
+        )
+
+    if not filters:
+        return
+
+    # Handle $or operator
+    if "$or" in filters:
+        or_groups = filters["$or"]
+        if not isinstance(or_groups, list) or not len(or_groups) >= 2:
+            raise HTTPException(
+                status_code=500,
+                detail="Your auth handler returned a filter with an invalid $or operator. The $or operator must be a list of at least 2 filter objects. Check the filter returned by your auth handler.",
+            )
+
+        # Recursively validate all groups
+        for group in or_groups:
+            _validate_filter_structure(group, nesting_level=nesting_level + 1)
+
+        # Validate remaining filters (implicit AND with the $or)
+        remaining_filters = {k: v for k, v in filters.items() if k != "$or"}
+        if remaining_filters:
+            _validate_filter_structure(
+                remaining_filters, nesting_level=nesting_level + 1
+            )
+
+    # Handle $and operator
+    if "$and" in filters:
+        and_groups = filters["$and"]
+        if not isinstance(and_groups, list) or not len(and_groups) >= 2:
+            raise HTTPException(
+                status_code=500,
+                detail="Your auth handler returned a filter with an invalid $and operator. The $and operator must be a list of at least 2 filter objects. Check the filter returned by your auth handler.",
+            )
+
+        # Recursively validate all groups
+        for group in and_groups:
+            _validate_filter_structure(group, nesting_level=nesting_level + 1)
+
+        # Validate remaining filters (implicit AND with the $and)
+        remaining_filters = {k: v for k, v in filters.items() if k != "$and"}
+        if remaining_filters:
+            _validate_filter_structure(
+                remaining_filters, nesting_level=nesting_level + 1
+            )
+
+
+def _check_filter_match(
+    metadata: dict,
+    filters: Auth.types.FilterType | None,
+    nesting_level: int = 0,
+) -> bool:
     """Check if metadata matches the filter conditions.
 
     Args:
         metadata: The metadata to check
         filters: The filter conditions to apply
+        nesting_level: Current depth of nested operators (max 2)
 
     Returns:
         True if the metadata matches all filter conditions, False otherwise
     """
+    if nesting_level > 2:
+        raise HTTPException(
+            status_code=500,
+            detail="Your auth handler returned a filter with too many nested operators. The maximum depth for nested operators is 2. Please simplify your filter.",
+        )
+
     if not filters:
         return True
 
+    # Handle $or operator
+    if "$or" in filters:
+        or_groups = filters["$or"]
+        if not isinstance(or_groups, list) or not len(or_groups) >= 2:
+            raise HTTPException(
+                status_code=500,
+                detail="Your auth handler returned a filter with an invalid $or operator. The $or operator must be a list of at least 2 filter objects. Check the filter returned by your auth handler.",
+            )
+
+        # Validate all groups first to ensure nesting limits are respected
+        # (even if we short-circuit during matching)
+        for group in or_groups:
+            _validate_filter_structure(group, nesting_level=nesting_level + 1)
+
+        # At least one group must match
+        or_match = False
+        for group in or_groups:
+            if _check_filter_match(metadata, group, nesting_level=nesting_level + 1):
+                or_match = True
+                break
+
+        if not or_match:
+            return False
+
+        # Check remaining filters (implicit AND with the $or)
+        remaining_filters = {k: v for k, v in filters.items() if k != "$or"}
+        if remaining_filters:
+            return _check_filter_match(
+                metadata, remaining_filters, nesting_level=nesting_level + 1
+            )
+        return True
+
+    # Handle $and operator
+    if "$and" in filters:
+        and_groups = filters["$and"]
+        if not isinstance(and_groups, list) or not len(and_groups) >= 2:
+            raise HTTPException(
+                status_code=500,
+                detail="Your auth handler returned a filter with an invalid $and operator. The $and operator must be a list of at least 2 filter objects. Check the filter returned by your auth handler.",
+            )
+
+        # Validate all groups first to ensure nesting limits are respected
+        for group in and_groups:
+            _validate_filter_structure(group, nesting_level=nesting_level + 1)
+
+        # All groups must match
+        for group in and_groups:
+            if not _check_filter_match(
+                metadata, group, nesting_level=nesting_level + 1
+            ):
+                return False
+
+        # Check remaining filters (implicit AND with the $and)
+        remaining_filters = {k: v for k, v in filters.items() if k != "$and"}
+        if remaining_filters:
+            return _check_filter_match(
+                metadata, remaining_filters, nesting_level=nesting_level + 1
+            )
+        return True
+
+    # Regular filter logic (implicit AND)
     for key, value in filters.items():
         if isinstance(value, dict):
             op = next(iter(value))
@@ -2874,11 +3099,18 @@ def _check_filter_match(metadata: dict, filters: Auth.types.FilterType | None) -
                 if key not in metadata or metadata[key] != filter_value:
                     return False
             elif op == "$contains":
-                if (
-                    key not in metadata
-                    or not isinstance(metadata[key], list)
-                    or filter_value not in metadata[key]
-                ):
+                if key not in metadata or not isinstance(metadata[key], list):
+                    return False
+
+                if isinstance(filter_value, list):
+                    # Mimick Postgres containment operator behavior.
+                    # It would be more efficient to use set operations here,
+                    # but we can't assume that elements are hashable.
+                    # The Postgres algorithm is also O(n^2).
+                    for filter_element in filter_value:
+                        if filter_element not in metadata[key]:
+                            return False
+                elif filter_value not in metadata[key]:
                     return False
         else:
             # Direct equality
@@ -2894,6 +3126,7 @@ async def _empty_generator():
 
 
 __all__ = [
+    "StreamHandler",
     "Assistants",
     "Crons",
     "Runs",