langgraph-runtime-inmem 0.6.4__py3-none-any.whl → 0.18.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
@@ -27,17 +26,24 @@ from starlette.exceptions import HTTPException
 
 from langgraph_runtime_inmem.checkpoint import Checkpointer
 from langgraph_runtime_inmem.database import InMemConnectionProto, connect
-from langgraph_runtime_inmem.inmem_stream import Message, get_stream_manager
+from langgraph_runtime_inmem.inmem_stream import (
+    THREADLESS_KEY,
+    ContextQueue,
+    Message,
+    get_stream_manager,
+)
 
 if typing.TYPE_CHECKING:
     from langgraph_api.asyncio import ValueEvent
     from langgraph_api.config import ThreadTTLConfig
     from langgraph_api.schema import (
         Assistant,
+        AssistantSelectField,
         Checkpoint,
         Config,
         Context,
         Cron,
+        CronSelectField,
         DeprecatedInterrupt,
         IfNotExists,
         MetadataInput,
@@ -46,15 +52,19 @@ if typing.TYPE_CHECKING:
         OnConflictBehavior,
         QueueStats,
         Run,
+        RunSelectField,
         RunStatus,
         StreamMode,
         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__)
 
@@ -136,6 +146,7 @@ class Assistants(Authenticated):
         offset: int,
         sort_by: str | None = None,
         sort_order: str | None = None,
+        select: list[AssistantSelectField] | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> tuple[AsyncIterator[Assistant], int]:
         metadata = metadata if metadata is not None else {}
@@ -157,9 +168,6 @@ class Assistants(Authenticated):
             and (not filters or _check_filter_match(assistant["metadata"], filters))
         ]
 
-        # Get total count before sorting and pagination
-        total_count = len(filtered_assistants)
-
         # Sort based on sort_by and sort_order
         sort_by = sort_by.lower() if sort_by else None
         if sort_by and sort_by in (
@@ -181,22 +189,31 @@ class Assistants(Authenticated):
             else:
                 filtered_assistants.sort(key=lambda x: x.get(sort_by), reverse=reverse)
         else:
+            sort_by = "created_at"
             # Default sorting by created_at in descending order
             filtered_assistants.sort(key=lambda x: x["created_at"], reverse=True)
 
         # Apply pagination
         paginated_assistants = filtered_assistants[offset : offset + limit]
+        cur = offset + limit if len(filtered_assistants) > offset + limit else None
 
         async def assistant_iterator() -> AsyncIterator[Assistant]:
             for assistant in paginated_assistants:
-                yield assistant
+                if select:
+                    # Filter to only selected fields
+                    filtered_assistant = {
+                        k: v for k, v in assistant.items() if k in select
+                    }
+                    yield filtered_assistant
+                else:
+                    yield assistant
 
-        return assistant_iterator(), total_count
+        return assistant_iterator(), cur
 
     @staticmethod
     async def get(
         conn: InMemConnectionProto,
-        assistant_id: UUID,
+        assistant_id: UUID | str,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Assistant]:
         """Get an assistant by ID."""
@@ -212,14 +229,14 @@ 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()
 
     @staticmethod
     async def put(
         conn: InMemConnectionProto,
-        assistant_id: UUID,
+        assistant_id: UUID | str,
         *,
         graph_id: str,
         config: Config,
@@ -231,6 +248,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(
@@ -245,6 +264,22 @@ class Assistants(Authenticated):
                 name=name,
             ),
         )
+
+        if config.get("configurable") and context:
+            raise HTTPException(
+                status_code=400,
+                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"]
+        elif context:
+            config["configurable"] = context
+
         existing_assistant = next(
             (a for a in conn.store["assistants"] if a["assistant_id"] == assistant_id),
             None,
@@ -330,6 +365,7 @@ class Assistants(Authenticated):
         """
         assistant_id = _ensure_uuid(assistant_id)
         metadata = metadata if metadata is not None else {}
+        config = config if config is not None else {}
         filters = await Assistants.handle_event(
             ctx,
             "update",
@@ -342,6 +378,19 @@ class Assistants(Authenticated):
                 name=name,
             ),
         )
+
+        if config.get("configurable") and context:
+            raise HTTPException(
+                status_code=400,
+                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.",
+            )
+
+        # Keep config and context up to date with one another
+        if config.get("configurable"):
+            context = config["configurable"]
+        elif context:
+            config["configurable"] = context
+
         assistant = next(
             (a for a in conn.store["assistants"] if a["assistant_id"] == assistant_id),
             None,
@@ -367,19 +416,17 @@ class Assistants(Authenticated):
             else 1
         )
 
-        # Update assistant_versions table
-        if metadata:
-            metadata = {
-                **assistant["metadata"],
-                **metadata,
-            }
         new_version_entry = {
             "assistant_id": assistant_id,
             "version": new_version,
             "graph_id": graph_id if graph_id is not None else assistant["graph_id"],
-            "config": config if config is not None else assistant["config"],
-            "context": context if context is not None else assistant["context"],
-            "metadata": metadata if metadata is not None else assistant["metadata"],
+            "config": config if config else assistant["config"],
+            "context": context if context is not None else assistant.get("context", {}),
+            "metadata": (
+                {**assistant["metadata"], **metadata}
+                if metadata is not None
+                else assistant["metadata"]
+            ),
             "created_at": now,
             "name": name if name is not None else assistant["name"],
             "description": (
@@ -514,6 +561,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"],
             }
         )
 
@@ -572,6 +621,37 @@ class Assistants(Authenticated):
 
         return _yield_versions()
 
+    @staticmethod
+    async def count(
+        conn: InMemConnectionProto,
+        *,
+        graph_id: str | None = None,
+        metadata: MetadataInput = None,
+        ctx: Auth.types.BaseAuthContext | None = None,
+    ) -> int:
+        """Get count of assistants."""
+        metadata = metadata if metadata is not None else {}
+        filters = await Assistants.handle_event(
+            ctx,
+            "search",
+            Auth.types.AssistantsSearch(
+                graph_id=graph_id, metadata=metadata, limit=0, offset=0
+            ),
+        )
+
+        count = 0
+        for assistant in conn.store["assistants"]:
+            if (
+                (not graph_id or assistant["graph_id"] == graph_id)
+                and (
+                    not metadata or is_jsonb_contained(assistant["metadata"], metadata)
+                )
+                and (not filters or _check_filter_match(assistant["metadata"], filters))
+            ):
+                count += 1
+
+        return count
+
 
 def is_jsonb_contained(superset: dict[str, Any], subset: dict[str, Any]) -> bool:
     """
@@ -649,13 +729,13 @@ def _patch_interrupt(
             interrupt = Interrupt(**interrupt)
 
         return {
-            "id": interrupt.interrupt_id
-            if hasattr(interrupt, "interrupt_id")
-            else None,
+            "id": (
+                interrupt.interrupt_id if hasattr(interrupt, "interrupt_id") else None
+            ),
             "value": interrupt.value,
             "resumable": interrupt.resumable,
             "ns": interrupt.ns,
-            "when": interrupt.when,
+            "when": interrupt.when,  # type: ignore[unresolved-attribute]
         }
 
 
@@ -666,6 +746,7 @@ class Threads(Authenticated):
     async def search(
         conn: InMemConnectionProto,
         *,
+        ids: list[str] | list[UUID] | None = None,
         metadata: MetadataInput,
         values: MetadataInput,
         status: ThreadStatus | None,
@@ -673,6 +754,7 @@ class Threads(Authenticated):
         offset: int,
         sort_by: str | None = None,
         sort_order: str | None = None,
+        select: list[ThreadSelectField] | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> tuple[AsyncIterator[Thread], int]:
         threads = conn.store["threads"]
@@ -692,7 +774,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
 
@@ -710,8 +804,6 @@ class Threads(Authenticated):
                 continue
 
             filtered_threads.append(thread)
-        # Get total count before pagination
-        total_count = len(filtered_threads)
 
         if sort_by and sort_by in [
             "thread_id",
@@ -724,20 +816,26 @@ class Threads(Authenticated):
                 filtered_threads, key=lambda x: x.get(sort_by), reverse=reverse
             )
         else:
+            sort_by = "created_at"
             # Default sorting by created_at in descending order
             sorted_threads = sorted(
-                filtered_threads, key=lambda x: x["created_at"], reverse=True
+                filtered_threads, key=lambda x: x["updated_at"], reverse=True
             )
 
         # Apply limit and offset
         paginated_threads = sorted_threads[offset : offset + limit]
+        cursor = offset + limit if len(sorted_threads) > offset + limit else None
 
         async def thread_iterator() -> AsyncIterator[Thread]:
             for thread in paginated_threads:
-                yield thread
+                if select:
+                    # Filter to only selected fields
+                    filtered_thread = {k: v for k, v in thread.items() if k in select}
+                    yield filtered_thread
+                else:
+                    yield thread
 
-        # Return both the iterator and the total count
-        return thread_iterator(), total_count
+        return thread_iterator(), cursor
 
     @staticmethod
     async def _get_with_filters(
@@ -799,7 +897,7 @@ class Threads(Authenticated):
     @staticmethod
     async def put(
         conn: InMemConnectionProto,
-        thread_id: UUID,
+        thread_id: UUID | str,
         *,
         metadata: MetadataInput,
         if_exists: OnConflictBehavior,
@@ -866,6 +964,7 @@ class Threads(Authenticated):
         thread_id: UUID,
         *,
         metadata: MetadataValue,
+        ttl: ThreadTTLConfig | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Thread]:
         """Update a thread."""
@@ -978,6 +1077,7 @@ class Threads(Authenticated):
         thread_id: UUID,
         run_id: UUID,
         run_status: RunStatus | Literal["rollback"],
+        graph_id: str,
         checkpoint: CheckpointPayload | None = None,
         exception: BaseException | None = None,
     ) -> None:
@@ -1057,6 +1157,7 @@ class Threads(Authenticated):
                 final_thread_status = "busy"
             else:
                 final_thread_status = base_thread_status
+        thread["metadata"]["graph_id"] = graph_id
         thread.update(
             {
                 "updated_at": now,
@@ -1136,13 +1237,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(
@@ -1151,8 +1262,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()
 
@@ -1161,7 +1272,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": {},
             }
@@ -1248,9 +1359,24 @@ 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", {}),
+                },
+            }
 
-            if graph_id := metadata.get("graph_id"):
+            # Fallback to graph_id from run if not in thread metadata
+            graph_id = metadata.get("graph_id")
+            if not graph_id:
+                for run in conn.store["runs"]:
+                    if run["thread_id"] == thread_id:
+                        graph_id = run["kwargs"]["config"]["configurable"]["graph_id"]
+                        break
+
+            if graph_id:
                 # format latest checkpoint for response
                 checkpointer.latest_iter = checkpoint
                 async with get_graph(
@@ -1290,6 +1416,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
 
@@ -1327,8 +1454,23 @@ 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", {}),
+                },
+            }
 
-            if graph_id := metadata.get("graph_id"):
+            # Fallback to graph_id from run if not in thread metadata
+            graph_id = metadata.get("graph_id")
+            if not graph_id:
+                for run in conn.store["runs"]:
+                    if run["thread_id"] == thread_id:
+                        graph_id = run["kwargs"]["config"]["configurable"]["graph_id"]
+                        break
+
+            if graph_id:
                 config["configurable"].setdefault("graph_id", graph_id)
 
                 checkpointer.latest_iter = checkpoint
@@ -1359,6 +1501,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
@@ -1366,7 +1521,11 @@ class Threads(Authenticated):
                         checkpoint_id=next_config["configurable"]["checkpoint_id"],
                     )
             else:
-                raise HTTPException(status_code=400, detail="Thread has no graph ID.")
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Thread '{thread['thread_id']}' has no assigned graph ID. This usually occurs when no runs have been made on this particular thread."
+                    " This operation requires a graph ID. Please ensure a run has been made for the thread or manually update the thread metadata (by setting the 'graph_id' field) before running this operation.",
+                )
 
         @staticmethod
         async def bulk(
@@ -1397,7 +1556,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:
@@ -1444,18 +1610,35 @@ 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"],
                     )
             else:
-                raise HTTPException(status_code=400, detail="Thread has no graph ID")
+                raise HTTPException(
+                    status_code=400,
+                    detail=f"Thread '{thread['thread_id']}' has no assigned graph ID. This usually occurs when no runs have been made on this particular thread."
+                    " This operation requires a graph ID. Please ensure a run has been made for the thread or manually update the thread metadata (by setting the 'graph_id' field) before running this operation.",
+                )
 
         @staticmethod
         async def list(
             conn: InMemConnectionProto,
             *,
             config: Config,
-            limit: int = 10,
+            limit: int = 1,
             before: str | Checkpoint | None = None,
             metadata: MetadataInput = None,
             ctx: Auth.types.BaseAuthContext | None = None,
@@ -1481,7 +1664,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(
@@ -1510,6 +1700,265 @@ class Threads(Authenticated):
 
             return []
 
+    class Stream(Authenticated):
+        resource = "threads"
+
+        @staticmethod
+        async def subscribe(
+            conn: InMemConnectionProto | AsyncConnectionProto,
+            thread_id: UUID,
+            seen_runs: set[UUID],
+        ) -> list[tuple[UUID, asyncio.Queue]]:
+            """Subscribe to the thread stream, creating queues for unseen runs."""
+            stream_manager = get_stream_manager()
+            queues = []
+
+            # 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"]
+                    if run_id not in seen_runs:
+                        queue = await stream_manager.add_queue(run_id, thread_id)
+                        queues.append((run_id, queue))
+                        seen_runs.add(run_id)
+
+            return queues
+
+        @staticmethod
+        async def join(
+            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."""
+            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()
+            created_queues: list[tuple[UUID, asyncio.Queue]] = []
+
+            try:
+                async with connect() as conn:
+                    await logger.ainfo(
+                        "Joined thread stream",
+                        thread_id=str(thread_id),
+                    )
+
+                    # Restore messages if resuming from a specific event
+                    if last_event_id is not None:
+                        # Collect all events from all message stores for this thread
+                        all_events = []
+                        for run_id in stream_manager.message_stores.get(
+                            str(thread_id), []
+                        ):
+                            for message in stream_manager.restore_messages(
+                                run_id, thread_id, last_event_id
+                            ):
+                                all_events.append((message, run_id))
+
+                        # Sort by message ID (which is ms-seq format)
+                        all_events.sort(key=lambda x: x[0].id.decode())
+
+                        # Yield sorted events
+                        for message, run_id in all_events:
+                            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}
+                                    )
+                            if not should_filter_event(
+                                event_bytes.decode("utf-8"), message_bytes
+                            ):
+                                yield (
+                                    event_bytes,
+                                    message_bytes,
+                                    message.id,
+                                )
+
+                    # Listen for live messages from all queues
+                    while True:
+                        # Refresh queues to pick up any new runs that joined this thread
+                        new_queue_tuples = await Threads.Stream.subscribe(
+                            conn, thread_id, seen_runs
+                        )
+                        # Track new queues for cleanup
+                        for run_id, queue in new_queue_tuples:
+                            created_queues.append((run_id, queue))
+
+                        for run_id, queue in created_queues:
+                            try:
+                                message = await asyncio.wait_for(
+                                    queue.get(), timeout=0.2
+                                )
+                                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
+                            except (ValueError, KeyError):
+                                continue
+
+                        # Yield execution to other tasks to prevent event loop starvation
+                        await asyncio.sleep(0)
+
+            except WrappedHTTPException as e:
+                raise e.http_exception from None
+            except asyncio.CancelledError:
+                await logger.awarning(
+                    "Thread stream client disconnected",
+                    thread_id=str(thread_id),
+                )
+                raise
+            except:
+                raise
+            finally:
+                # Clean up all created queues
+                for run_id, queue in created_queues:
+                    try:
+                        await stream_manager.remove_queue(run_id, thread_id, queue)
+                    except Exception:
+                        # 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,
+        *,
+        metadata: MetadataInput = None,
+        values: MetadataInput = None,
+        status: ThreadStatus | None = None,
+        ctx: Auth.types.BaseAuthContext | None = None,
+    ) -> int:
+        """Get count of threads."""
+        threads = conn.store["threads"]
+        metadata = metadata if metadata is not None else {}
+        values = values if values is not None else {}
+        filters = await Threads.handle_event(
+            ctx,
+            "search",
+            Auth.types.ThreadsSearch(
+                metadata=metadata,
+                values=values,
+                status=status,
+                limit=0,
+                offset=0,
+            ),
+        )
+
+        count = 0
+        for thread in threads:
+            if filters and not _check_filter_match(thread["metadata"], filters):
+                continue
+
+            if metadata and not is_jsonb_contained(thread["metadata"], metadata):
+                continue
+
+            if (
+                values
+                and "values" in thread
+                and not is_jsonb_contained(thread["values"], values)
+            ):
+                continue
+
+            if status and thread.get("status") != status:
+                continue
+
+            count += 1
+
+        return count
+
 
 RUN_LOCK = asyncio.Lock()
 
@@ -1526,38 +1975,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
@@ -1621,38 +2069,51 @@ class Runs(Authenticated):
     @asynccontextmanager
     @staticmethod
     async def enter(
-        run_id: UUID, 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 queue for this run
-        queue = await Runs.Stream.subscribe(run_id)
+        # Get control queue for this run (normal queue is created during run creation)
+        control_queue = await stream_manager.add_control_queue(run_id, thread_id)
 
         async with SimpleTaskGroup(cancel=True, taskgroup_name="Runs.enter") as tg:
             done = ValueEvent()
-            tg.create_task(listen_for_cancellation(queue, run_id, done))
+            tg.create_task(
+                listen_for_cancellation(control_queue, run_id, thread_id, done)
+            )
 
             # Give done event to caller
             yield done
-            # Signal done to all subscribers
+            # Store the control message for late subscribers
             control_message = Message(
                 topic=f"run:{run_id}:control".encode(), data=b"done"
             )
+            await stream_manager.put(run_id, thread_id, control_message)
 
-            # Store the control message for late subscribers
-            await stream_manager.put(run_id, control_message)
-            stream_manager.control_queues[run_id].append(control_message)
-            # Clean up this queue
-            await stream_manager.remove_queue(run_id, queue)
+            # Signal done to all subscribers using stream codec
+            stream_message = Message(
+                topic=f"run:{run_id}:stream".encode(),
+                data=STREAM_CODEC.encode("control", b"done"),
+            )
+            await stream_manager.put(
+                run_id, thread_id, stream_message, resumable=resumable
+            )
+
+            # Remove the control_queue (normal queue is cleaned up during run deletion)
+            await stream_manager.remove_control_queue(run_id, thread_id, control_queue)
 
     @staticmethod
-    async def sweep(conn: InMemConnectionProto) -> list[UUID]:
+    async def sweep() -> None:
         """Sweep runs that are no longer running"""
-        return []
+        pass
 
     @staticmethod
     def _merge_jsonb(*objects: dict) -> dict:
@@ -1670,7 +2131,7 @@ class Runs(Authenticated):
 
     @staticmethod
     async def put(
-        conn: InMemConnectionProto,
+        conn: InMemConnectionProto | AsyncConnectionProto,
         assistant_id: UUID,
         kwargs: dict,
         *,
@@ -1701,6 +2162,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(
@@ -1710,7 +2172,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,
@@ -1731,6 +2193,7 @@ class Runs(Authenticated):
             # Create new thread
             if thread_id is None:
                 thread_id = uuid4()
+
             thread = Thread(
                 thread_id=thread_id,
                 status="busy",
@@ -1746,7 +2209,6 @@ class Runs(Authenticated):
                     {
                         "configurable": Runs._merge_jsonb(
                             Runs._get_configurable(assistant["config"]),
-                            Runs._get_configurable(config),
                         )
                     },
                 ),
@@ -1754,6 +2216,7 @@ class Runs(Authenticated):
                 updated_at=datetime.now(UTC),
                 values=b"",
             )
+
             await logger.ainfo("Creating thread", thread_id=thread_id)
             conn.store["threads"].append(thread)
         elif existing_thread:
@@ -1775,7 +2238,6 @@ class Runs(Authenticated):
                         "configurable": Runs._merge_jsonb(
                             Runs._get_configurable(assistant["config"]),
                             Runs._get_configurable(existing_thread["config"]),
-                            Runs._get_configurable(config),
                         )
                     },
                 )
@@ -1920,6 +2382,7 @@ class Runs(Authenticated):
             if not thread:
                 return _empty_generator()
         _delete_checkpoints_for_thread(thread_id, conn, run_id=run_id)
+
         found = False
         for i, run in enumerate(conn.store["runs"]):
             if run["run_id"] == run_id and run["thread_id"] == thread_id:
@@ -1935,54 +2398,10 @@ 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
-        ):
-            if mode == b"values":
-                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"])})
-                    )
-                return thread["values"]
-
     @staticmethod
     async def cancel(
-        conn: InMemConnectionProto,
-        run_ids: Sequence[UUID] | None = None,
+        conn: InMemConnectionProto | AsyncConnectionProto,
+        run_ids: Sequence[UUID | str] | None = None,
         *,
         action: Literal["interrupt", "rollback"] = "interrupt",
         thread_id: UUID | None = None,
@@ -2086,9 +2505,9 @@ class Runs(Authenticated):
                 topic=f"run:{run_id}:control".encode(),
                 data=action.encode(),
             )
-            coros.append(stream_manager.put(run_id, control_message))
+            coros.append(stream_manager.put(run_id, thread_id, control_message))
 
-            queues = stream_manager.get_queues(run_id)
+            queues = stream_manager.get_queues(run_id, thread_id)
 
             if run["status"] in ("pending", "running"):
                 cancelable_runs.append(run)
@@ -2146,6 +2565,7 @@ class Runs(Authenticated):
         limit: int = 10,
         offset: int = 0,
         status: RunStatus | None = None,
+        select: list[RunSelectField] | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Run]:
         """List all runs by thread."""
@@ -2173,7 +2593,12 @@ class Runs(Authenticated):
 
         async def _return():
             for run in sliced_runs:
-                yield run
+                if select:
+                    # Filter to only selected fields
+                    filtered_run = {k: v for k, v in run.items() if k in select}
+                    yield filtered_run
+                else:
+                    yield run
 
         return _return()
 
@@ -2197,73 +2622,81 @@ class Runs(Authenticated):
         @staticmethod
         async def subscribe(
             run_id: UUID,
-            *,
-            stream_mode: StreamMode | None = None,
-        ) -> asyncio.Queue:
+            thread_id: UUID | None = None,
+        ) -> ContextQueue:
             """Subscribe to the run stream, returning a queue."""
             stream_manager = get_stream_manager()
-            queue = await stream_manager.add_queue(_ensure_uuid(run_id))
+            queue = await stream_manager.add_queue(_ensure_uuid(run_id), thread_id)
 
             # If there's a control message already stored, send it to the new subscriber
-            if control_messages := stream_manager.control_queues.get(run_id):
-                for control_msg in control_messages:
-                    await queue.put(control_msg)
+            if thread_id is None:
+                thread_id = THREADLESS_KEY
+            if control_queues := stream_manager.control_queues.get(thread_id, {}).get(
+                run_id
+            ):
+                for control_queue in control_queues:
+                    try:
+                        while True:
+                            control_msg = control_queue.get()
+                            await queue.put(control_msg)
+                    except asyncio.QueueEmpty:
+                        pass
             return queue
 
         @staticmethod
         async def join(
             run_id: UUID,
             *,
+            stream_channel: asyncio.Queue,
             thread_id: UUID,
             ignore_404: bool = False,
             cancel_on_disconnect: bool = False,
-            stream_mode: StreamMode | 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_dumpb
+            from langgraph_api.utils.stream_codec import decode_stream_message
 
-            queue = (
-                stream_mode
-                if isinstance(stream_mode, asyncio.Queue)
-                else await Runs.Stream.subscribe(run_id, stream_mode=stream_mode)
-            )
-
+            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"
-                                )
-                            )
-                    channel_prefix = f"run:{run_id}:stream:"
-                    len_prefix = len(channel_prefix.encode())
+                    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, last_event_id
+                        run_id, thread_id, last_event_id
                     ):
-                        topic, data, id = message.topic, message.data, message.id
-                        if topic.decode() == f"run:{run_id}:control":
-                            if data == b"done":
+                        data, id = message.data, message.id
+                        decoded = decode_stream_message(data, channel=message.topic)
+                        mode = decoded.event_bytes.decode("utf-8")
+                        payload = decoded.message_bytes
+
+                        if mode == "control":
+                            if payload == b"done":
                                 return
-                        else:
-                            yield topic[len_prefix:], data, id
+                        elif (
+                            not stream_mode
+                            or mode in stream_mode
+                            or (
+                                (
+                                    "messages" in stream_mode
+                                    or "messages-tuple" in stream_mode
+                                )
+                                and mode.startswith("messages")
+                            )
+                        ):
+                            yield mode.encode(), payload, id
                             logger.debug(
                                 "Replayed run event",
                                 run_id=str(run_id),
                                 message_id=id,
-                                stream_mode=topic[len_prefix:],
+                                stream_mode=mode,
                                 data=data,
                             )
 
@@ -2271,20 +2704,32 @@ class Runs(Authenticated):
                         try:
                             # Wait for messages with a timeout
                             message = await asyncio.wait_for(queue.get(), timeout=0.5)
-                            topic, data, id = message.topic, message.data, message.id
+                            data, id = message.data, message.id
+                            decoded = decode_stream_message(data, channel=message.topic)
+                            mode = decoded.event_bytes.decode("utf-8")
+                            payload = decoded.message_bytes
 
-                            if topic.decode() == f"run:{run_id}:control":
-                                if data == b"done":
+                            if mode == "control":
+                                if payload == b"done":
                                     break
-                            else:
-                                # Extract mode from topic
-                                yield topic[len_prefix:], data, id
+                            elif (
+                                not stream_mode
+                                or mode in stream_mode
+                                or (
+                                    (
+                                        "messages" in stream_mode
+                                        or "messages-tuple" in stream_mode
+                                    )
+                                    and mode.startswith("messages")
+                                )
+                            ):
+                                yield mode.encode(), payload, id
                                 logger.debug(
                                     "Streamed run event",
                                     run_id=str(run_id),
-                                    stream_mode=topic[len_prefix:],
+                                    stream_mode=mode,
                                     message_id=id,
-                                    data=data,
+                                    data=payload,
                                 )
                         except TimeoutError:
                             # Check if the run is still pending
@@ -2298,8 +2743,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,
                                 )
@@ -2314,45 +2761,68 @@ class Runs(Authenticated):
                 raise
             finally:
                 stream_manager = get_stream_manager()
-                await stream_manager.remove_queue(run_id, queue)
+                await stream_manager.remove_queue(run_id, thread_id, queue)
 
         @staticmethod
-        async def publish(
+        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,
             event: str,
             message: bytes,
             *,
+            thread_id: UUID | str | None = None,
             resumable: bool = False,
         ) -> None:
             """Publish a message to all subscribers of the run stream."""
-            topic = f"run:{run_id}:stream:{event}".encode()
+            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
+            # Send to all queues subscribed to this run_id using protocol frame
+            payload = STREAM_CODEC.encode(event, message)
             await stream_manager.put(
-                run_id, Message(topic=topic, data=message), resumable
+                run_id, thread_id, Message(topic=topic, data=payload), resumable
             )
 
 
-async def listen_for_cancellation(queue: asyncio.Queue, run_id: UUID, done: ValueEvent):
+async def listen_for_cancellation(
+    queue: asyncio.Queue, run_id: UUID, thread_id: UUID | None, done: ValueEvent
+):
     """Listen for cancellation messages and set the done event accordingly."""
     from langgraph_api.errors import UserInterrupt, UserRollback
 
     stream_manager = get_stream_manager()
-    control_key = f"run:{run_id}:control"
 
-    if existing_queue := stream_manager.control_queues.get(run_id):
-        for message in existing_queue:
-            payload = message.data
-            if payload == b"rollback":
-                done.set(UserRollback())
-            elif payload == b"interrupt":
-                done.set(UserInterrupt())
+    if control_key := stream_manager.get_control_key(run_id, thread_id):
+        payload = control_key.data
+        if payload == b"rollback":
+            done.set(UserRollback())
+        elif payload == b"interrupt":
+            done.set(UserInterrupt())
 
     while not done.is_set():
         try:
             # This task gets cancelled when Runs.enter exits anyway,
-            # so we can have a pretty length timeout here
+            # so we can have a pretty lengthy timeout here
             message = await asyncio.wait_for(queue.get(), timeout=240)
             payload = message.data
             if payload == b"rollback":
@@ -2362,10 +2832,6 @@ async def listen_for_cancellation(queue: asyncio.Queue, run_id: UUID, done: Valu
             elif payload == b"done":
                 done.set()
                 break
-
-            # Store control messages for late subscribers
-            if message.topic.decode() == control_key:
-                stream_manager.control_queues[run_id].append(message)
         except TimeoutError:
             break
 
@@ -2379,6 +2845,7 @@ 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,
         ctx: Auth.types.BaseAuthContext | None = None,
     ) -> AsyncIterator[Cron]:
@@ -2417,10 +2884,24 @@ class Crons:
         thread_id: UUID | None,
         limit: int,
         offset: int,
+        select: list[CronSelectField] | None = None,
         ctx: Auth.types.BaseAuthContext | None = None,
-    ) -> AsyncIterator[Cron]:
+        sort_by: str | None = None,
+        sort_order: Literal["asc", "desc"] | None = None,
+    ) -> tuple[AsyncIterator[Cron], int]:
         raise NotImplementedError
 
+    @staticmethod
+    async def count(
+        conn: InMemConnectionProto,
+        *,
+        assistant_id: UUID | None = None,
+        thread_id: UUID | None = None,
+        ctx: Auth.types.BaseAuthContext | None = None,
+    ) -> int:
+        """Get count of crons."""
+        raise NotImplementedError("The in-mem server does not implement Crons.")
+
 
 async def cancel_run(
     thread_id: UUID, run_id: UUID, ctx: Auth.types.BaseAuthContext | None = None
@@ -2478,11 +2959,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
@@ -2498,6 +2986,7 @@ async def _empty_generator():
 
 
 __all__ = [
+    "StreamHandler",
     "Assistants",
     "Crons",
     "Runs",