langgraph-api 0.4.1__py3-none-any.whl → 0.7.3__py3-none-any.whl

langgraph_api/utils/stream_codec.py

@@ -0,0 +1,315 @@
+from __future__ import annotations
+
+import base64
+from dataclasses import dataclass
+
+import orjson
+import structlog
+
+PROTOCOL_VERSION = 1
+"""
+---
+Version 1:
+Byte Offsets
+0        1                  3                5                    5+N                5+N+M
++--------+------------------+----------------+------------------+------------------+--------------------+
+| version| stream_id_len    | event_len      | stream_id        | event            | message            |
++--------+------------------+----------------+------------------+------------------+--------------------+
+   1 B         2 B                2 B              N B                 M B               variable
+
+---- Old (to be dropped soon / multiple formats)
+Version 0 (old):
+1) b"$:" + <stream_id> + b"$:" + <event> + b"$:" + <raw_json>
+2) b"$:" + <stream_id> + b"$:" + <raw_json>
+"""
+
+BYTE_MASK = 0xFF
+HEADER_LEN = 5
+logger = structlog.stdlib.get_logger(__name__)
+
+
+class StreamFormatError(ValueError):
+    """Raised when a stream frame fails validation."""
+
+
+@dataclass(slots=True)
+class StreamPacket:
+    version: int
+    event: memoryview | bytes
+    message: memoryview | bytes
+    stream_id: memoryview | bytes | None
+
+    @property
+    def event_bytes(self) -> bytes:
+        return (
+            self.event.tobytes() if isinstance(self.event, memoryview) else self.event
+        )
+
+    @property
+    def message_bytes(self) -> bytes:
+        return (
+            self.message.tobytes()
+            if isinstance(self.message, memoryview)
+            else self.message
+        )
+
+    @property
+    def resumable(self) -> bool:
+        return self.stream_id is not None
+
+    @property
+    def stream_id_bytes(self) -> bytes | None:
+        if self.stream_id is None:
+            return None
+        if isinstance(self.stream_id, bytes):
+            return self.stream_id
+        return self.stream_id.tobytes()
+
+
+class StreamCodec:
+    """Codec for encoding and decoding stream packets."""
+
+    __slots__ = ("_version",)
+
+    def __init__(self, *, protocol_version: int = PROTOCOL_VERSION) -> None:
+        self._version = protocol_version & BYTE_MASK
+
+    def encode(
+        self,
+        event: str,
+        message: bytes,
+        *,
+        stream_id: str | None = None,
+    ) -> bytes:
+        if not event:
+            raise StreamFormatError("event cannot be empty")
+        event_bytes = event.encode("utf-8")
+        if len(event_bytes) > 0xFFFF:
+            raise StreamFormatError("event exceeds 65535 bytes; cannot encode")
+        if not event_bytes:
+            raise StreamFormatError("event cannot be empty")
+
+        if stream_id:
+            # It's a resumable stream
+            stream_id_bytes = stream_id.encode("utf-8")
+            if len(stream_id_bytes) > 0xFFFF:
+                raise StreamFormatError("stream_id exceeds 65535 bytes; cannot encode")
+        else:
+            stream_id_bytes = None
+        stream_id_len = len(stream_id_bytes) if stream_id_bytes else 0
+        event_len = len(event_bytes)
+        frame = bytearray(HEADER_LEN + stream_id_len + event_len + len(message))
+        frame[0] = self._version
+        frame[1:3] = stream_id_len.to_bytes(2, "big")
+        frame[3:5] = event_len.to_bytes(2, "big")
+
+        cursor = HEADER_LEN
+        if stream_id_bytes is not None:
+            frame[cursor : cursor + stream_id_len] = stream_id_bytes
+            cursor += stream_id_len
+
+        frame[cursor : cursor + event_len] = event_bytes
+        cursor += event_len
+        frame[cursor:] = message
+        return bytes(frame)
+
+    def decode(self, data: bytes | bytearray | memoryview) -> StreamPacket:
+        view = data if isinstance(data, memoryview) else memoryview(data)
+        if len(view) < HEADER_LEN:
+            raise StreamFormatError("frame too short")
+
+        version = view[0]
+        if version != self._version:
+            raise StreamFormatError(f"unsupported protocol version: {version}")
+
+        stream_id_len = int.from_bytes(view[1:3], "big")
+        event_len = int.from_bytes(view[3:5], "big")
+        if event_len == 0:
+            raise StreamFormatError("event cannot be empty")
+        offset = HEADER_LEN
+        if stream_id_len > 0:
+            stream_id_view = view[offset : offset + stream_id_len]
+            offset += stream_id_len
+        else:
+            # Not resumable
+            stream_id_view = None
+        if len(view) < offset + event_len:
+            raise StreamFormatError("truncated event payload")
+        event_view = view[offset : offset + event_len]
+        offset += event_len
+        message_view = view[offset:]
+        return StreamPacket(
+            version=version,
+            event=event_view,
+            message=message_view,
+            stream_id=stream_id_view,
+        )
+
+    def decode_safe(self, data: bytes | bytearray | memoryview) -> StreamPacket | None:
+        try:
+            return self.decode(data)
+        except StreamFormatError as e:
+            logger.warning(f"Failed to decode as version {self._version}", error=e)
+            return None
+
+
+STREAM_CODEC = StreamCodec()
+
+
+def decode_stream_message(
+    data: bytes | bytearray | memoryview,
+    *,
+    channel: bytes | str | None = None,
+) -> StreamPacket:
+    if isinstance(data, memoryview):
+        view = data
+    elif isinstance(data, (bytes, bytearray)):
+        view = memoryview(data)
+    else:
+        logger.warning("Unknown type for stream message", type=type(data))
+        view = memoryview(bytes(data))
+
+    # Current protocol version
+    if packet := STREAM_CODEC.decode_safe(view):
+        return packet
+    logger.debug("Attempting to decode a v0 formatted stream message")
+    # Legacy codecs. Yuck. Won't be hit unless you have stale pods running (or for a brief period during upgrade).
+    # Schedule for removal in next major release.
+    if packet := _decode_v0_resumable_format(view, channel):
+        return packet
+
+    # Non-resumable format.
+    if packet := _decode_v0_live_format(view, channel):
+        return packet
+    raise StreamFormatError("failed to decode stream message")
+
+
+_STREAMING_DELIMITER = b"$:"
+_STREAMING_DELIMITER_LEN = len(_STREAMING_DELIMITER)
+
+
+def _decode_v0_resumable_format(
+    view: memoryview,
+    channel: bytes | str | None = None,
+) -> StreamPacket | None:
+    """
+    Legacy v0 resumable format:
+      1) b"$:" + <stream_id> + b"$:" + <event> + b"$:" + <raw_json>
+      2) b"$:" + <stream_id> + b"$:" + <raw_json>
+    """
+
+    # must start with "$:"
+    if (
+        len(view) < _STREAMING_DELIMITER_LEN
+        or view[:_STREAMING_DELIMITER_LEN] != _STREAMING_DELIMITER
+    ):
+        return None
+
+    # "$:<stream_id>$:"
+    first = _find_delim(view, _STREAMING_DELIMITER_LEN, _STREAMING_DELIMITER)
+    if first == -1:
+        return None
+    stream_view = view[_STREAMING_DELIMITER_LEN:first]
+
+    # try "$:<event>$:"
+    second = _find_delim(view, first + _STREAMING_DELIMITER_LEN, _STREAMING_DELIMITER)
+    if second != -1:
+        event_view = view[first + _STREAMING_DELIMITER_LEN : second]
+        msg_view = view[second + _STREAMING_DELIMITER_LEN :]
+        return StreamPacket(
+            version=0,
+            event=event_view,
+            message=msg_view,
+            stream_id=stream_view,
+        )
+
+    chan_bytes = channel.encode("utf-8") if isinstance(channel, str) else channel
+
+    if chan_bytes:
+        marker = b":stream:"
+        idx = chan_bytes.rfind(marker)
+        event_bytes = chan_bytes[idx + len(marker) :] if idx != -1 else chan_bytes
+    else:
+        event_bytes = b""
+
+    msg_view = view[first + _STREAMING_DELIMITER_LEN :]
+    return StreamPacket(
+        version=0,
+        event=memoryview(event_bytes),
+        message=msg_view,
+        stream_id=stream_view,
+    )
+
+
+def _decode_v0_live_format(
+    view: memoryview, channel: bytes | str | None = None
+) -> StreamPacket | None:
+    try:
+        package = orjson.loads(view)
+    except orjson.JSONDecodeError:
+        return _decode_v0_flat_format(view, channel)
+    if (
+        not isinstance(package, dict)
+        or "event" not in package
+        or "message" not in package
+    ):
+        return _decode_v0_flat_format(view, channel)
+    event_obj = package.get("event")
+    message_obj = package.get("message")
+    if event_obj is None:
+        event_bytes = b""
+    elif isinstance(event_obj, str):
+        event_bytes = event_obj.encode()
+    elif isinstance(event_obj, (bytes, bytearray, memoryview)):
+        event_bytes = bytes(event_obj)
+    else:
+        event_bytes = orjson.dumps(event_obj)
+
+    if isinstance(message_obj, (bytes, bytearray, memoryview)):
+        message_view = memoryview(bytes(message_obj))
+    elif isinstance(message_obj, str):
+        try:
+            message_view = memoryview(base64.b64decode(message_obj))
+        except Exception:
+            message_view = memoryview(message_obj.encode())
+    elif message_obj is None:
+        message_view = memoryview(b"")
+    else:
+        message_view = memoryview(orjson.dumps(message_obj))
+
+    return StreamPacket(
+        event=event_bytes,
+        message=message_view,
+        stream_id=None,
+        version=0,
+    )
+
+
+def _decode_v0_flat_format(
+    view: memoryview, channel: bytes | str | None = None
+) -> StreamPacket | None:
+    packet = bytes(view)
+    stream_id = None
+    if channel is None:
+        return
+    if packet.startswith(b"$:"):
+        _, stream_id, packet = packet.split(b":", 2)
+    channel = channel.encode("utf-8") if isinstance(channel, str) else channel
+    channel = channel.split(b":")[-1]
+    return StreamPacket(
+        version=0,
+        event=memoryview(channel),
+        message=memoryview(packet),
+        stream_id=stream_id,
+    )
+
+
+def _find_delim(view: memoryview, start: int, delimiter: bytes) -> int:
+    delim_len = len(delimiter)
+    end = len(view) - delim_len
+    i = start
+    while i <= end:
+        if view[i : i + delim_len] == delimiter:
+            return i
+        i += 1
+    return -1

langgraph_api/utils/uuids.py

@@ -1,32 +1,35 @@
-import os
-import time
-from uuid import UUID, SafeUUID
+"""UUID utility functions.
 
-_last_timestamp_v7 = None
-_last_counter_v7 = 0  # 42-bit counter
-_RFC_4122_VERSION_7_FLAGS = (7 << 76) | (0x8000 << 48)
+This module exports a uuid7 function to generate monotonic, time-ordered UUIDs
+for tracing and similar operations.
+"""
 
+from __future__ import annotations
 
-def _uuid7_get_counter_and_tail():
-    rand = int.from_bytes(os.urandom(10))
-    # 42-bit counter with MSB set to 0
-    counter = (rand >> 32) & 0x1FF_FFFF_FFFF
-    # 32-bit random data
-    tail = rand & 0xFFFF_FFFF
-    return counter, tail
+import typing
+from uuid import UUID
 
+from uuid_utils.compat import uuid7 as _uuid_utils_uuid7
 
-def _from_int(value: int) -> UUID:
-    uid = object.__new__(UUID)
-    object.__setattr__(uid, "int", value)
-    object.__setattr__(uid, "is_safe", SafeUUID.unknown)
-    return uid
+if typing.TYPE_CHECKING:
+    from uuid import UUID
 
+_NANOS_PER_SECOND: typing.Final = 1_000_000_000
 
-def uuid7():
-    """Generate a UUID from a Unix timestamp in milliseconds and random bits.
+
+def _to_timestamp_and_nanos(nanoseconds: int) -> tuple[int, int]:
+    """Split a nanosecond timestamp into seconds and remaining nanoseconds."""
+    seconds, nanos = divmod(nanoseconds, _NANOS_PER_SECOND)
+    return seconds, nanos
+
+
+def uuid7(nanoseconds: int | None = None) -> UUID:
+    """Generate a UUID from a Unix timestamp in nanoseconds and random bits.
 
     UUIDv7 objects feature monotonicity within a millisecond.
+
+    Args:
+        nanoseconds: Optional ns timestamp. If not provided, uses current time.
     """
     # --- 48 ---   -- 4 --   --- 12 ---   -- 2 --   --- 30 ---   - 32 -
     # unix_ts_ms | version | counter_hi | variant | counter_lo | random
@@ -41,47 +44,11 @@ def uuid7():
     # advanced and the counter is reset to a random 42-bit integer with MSB
     # set to 0.
 
-    global _last_timestamp_v7
-    global _last_counter_v7
-
-    nanoseconds = time.time_ns()
-    timestamp_ms = nanoseconds // 1_000_000
-
-    if _last_timestamp_v7 is None or timestamp_ms > _last_timestamp_v7:
-        counter, tail = _uuid7_get_counter_and_tail()
-    else:
-        if timestamp_ms < _last_timestamp_v7:
-            timestamp_ms = _last_timestamp_v7 + 1
-        # advance the 42-bit counter
-        counter = _last_counter_v7 + 1
-        if counter > 0x3FF_FFFF_FFFF:
-            # advance the 48-bit timestamp
-            timestamp_ms += 1
-            counter, tail = _uuid7_get_counter_and_tail()
-        else:
-            # 32-bit random data
-            tail = int.from_bytes(os.urandom(4))
-
-    unix_ts_ms = timestamp_ms & 0xFFFF_FFFF_FFFF
-    counter_msbs = counter >> 30
-    # keep 12 counter's MSBs and clear variant bits
-    counter_hi = counter_msbs & 0x0FFF
-    # keep 30 counter's LSBs and clear version bits
-    counter_lo = counter & 0x3FFF_FFFF
-    # ensure that the tail is always a 32-bit integer (by construction,
-    # it is already the case, but future interfaces may allow the user
-    # to specify the random tail)
-    tail &= 0xFFFF_FFFF
+    # For now, just delegate to the uuid_utils implementation
+    if nanoseconds is None:
+        return _uuid_utils_uuid7()
+    seconds, nanos = _to_timestamp_and_nanos(nanoseconds)
+    return _uuid_utils_uuid7(timestamp=seconds, nanos=nanos)
 
-    int_uuid_7 = unix_ts_ms << 80
-    int_uuid_7 |= counter_hi << 64
-    int_uuid_7 |= counter_lo << 32
-    int_uuid_7 |= tail
-    # by construction, the variant and version bits are already cleared
-    int_uuid_7 |= _RFC_4122_VERSION_7_FLAGS
-    res = _from_int(int_uuid_7)
 
-    # defer global update until all computations are done
-    _last_timestamp_v7 = timestamp_ms
-    _last_counter_v7 = counter
-    return res
+__all__ = ["uuid7"]

langgraph_api/validation.py

@@ -8,6 +8,9 @@ with open(pathlib.Path(__file__).parent.parent / "openapi.json") as f:
 
 openapi = orjson.loads(openapi_str)
 
+ConfigValidator = jsonschema_rs.validator_for(
+    openapi["components"]["schemas"]["Config"]
+)
 AssistantVersionsSearchRequest = jsonschema_rs.validator_for(
     openapi["components"]["schemas"]["AssistantVersionsSearchRequest"]
 )
@@ -49,6 +52,9 @@ ThreadCreate = jsonschema_rs.validator_for(
 ThreadPatch = jsonschema_rs.validator_for(
     openapi["components"]["schemas"]["ThreadPatch"]
 )
+ThreadPruneRequest = jsonschema_rs.validator_for(
+    openapi["components"]["schemas"]["ThreadPruneRequest"]
+)
 ThreadStateUpdate = jsonschema_rs.validator_for(
     {
         **openapi["components"]["schemas"]["ThreadStateUpdate"],
@@ -121,6 +127,9 @@ RunCreateStateful = jsonschema_rs.validator_for(
 )
 RunsCancel = jsonschema_rs.validator_for(openapi["components"]["schemas"]["RunsCancel"])
 CronCreate = jsonschema_rs.validator_for(openapi["components"]["schemas"]["CronCreate"])
+ThreadCronCreate = jsonschema_rs.validator_for(
+    openapi["components"]["schemas"]["ThreadCronCreate"]
+)
 CronSearch = jsonschema_rs.validator_for(openapi["components"]["schemas"]["CronSearch"])
 CronCountRequest = jsonschema_rs.validator_for(
     openapi["components"]["schemas"]["CronCountRequest"]

langgraph_api/webhook.py

@@ -1,14 +1,117 @@
+import asyncio
+import ipaddress
+import socket
 from datetime import UTC, datetime
+from typing import TYPE_CHECKING
+from urllib.parse import urlparse
 
 import structlog
+from starlette.exceptions import HTTPException
 
-from langgraph_api.config import HTTP_CONFIG
-from langgraph_api.http import get_http_client, get_loopback_client, http_request
-from langgraph_api.worker import WorkerResult
+from langgraph_api.config import HTTP_CONFIG, WEBHOOKS_CONFIG
+from langgraph_api.config.schemas import WebhookUrlPolicy
+from langgraph_api.http import ensure_http_client, get_loopback_client, http_request
+
+if TYPE_CHECKING:
+    from langgraph_api.worker import WorkerResult
 
 logger = structlog.stdlib.get_logger(__name__)
 
 
+async def validate_webhook_url_or_raise(url: str) -> None:
+    """Validate a user-provided webhook URL against configured policy.
+
+    No-ops when WEBHOOKS_CONFIG is not set (preserves legacy behavior).
+    """
+    cfg = WEBHOOKS_CONFIG
+    if not cfg:
+        return
+
+    policy = WebhookUrlPolicy(cfg.get("url") or {})
+    allowed_domains = policy.get("allowed_domains") or []
+    allowed_ports = policy.get("allowed_ports")
+    max_url_length = int(policy.get("max_url_length", 4096))
+    # TODO: We should flip this in the next minor release
+    require_https = bool(policy.get("require_https", False))
+    disable_loopback = bool(policy.get("disable_loopback", False))
+
+    if len(url) > max_url_length:
+        raise HTTPException(status_code=422, detail="Webhook URL too long")
+
+    # Relative loopback URL (internal route)
+    if url.startswith("/"):
+        if disable_loopback:
+            raise HTTPException(
+                status_code=422, detail="Loopback webhooks are disabled"
+            )
+        # The other checks would fail here, so we can just return
+        return
+
+    parsed = urlparse(url)
+    if require_https and parsed.scheme.lower() != "https":
+        raise HTTPException(status_code=422, detail="Webhook must use https")
+
+    # Port policy: only enforce if configured; omit default enforcement otherwise
+    if allowed_ports:
+        if parsed.port is not None:
+            port = parsed.port
+        else:
+            port = 443 if parsed.scheme == "https" else 80
+        if port not in allowed_ports:
+            raise HTTPException(
+                status_code=422, detail=f"Webhook port {port} not allowed"
+            )
+
+    host = parsed.hostname or ""
+    if not host:
+        raise HTTPException(
+            status_code=422, detail=f"Invalid webhook hostname '{host}'"
+        )
+
+    # Domain allowlist
+    if allowed_domains:
+        host_allowed = False
+        for pattern in allowed_domains:
+            pattern = pattern.strip().lower()
+            if pattern.startswith("*."):
+                base = pattern[2:]
+                if host.lower().endswith("." + base):
+                    host_allowed = True
+                    break
+            else:
+                if host.lower() == pattern:
+                    host_allowed = True
+                    break
+        if not host_allowed:
+            raise HTTPException(status_code=422, detail="Webhook domain not allowed")
+
+    # Note we don't do default SSRF protections mainly because it would require a minor bump since it could break valid use cases.
+    try:
+        infos = await asyncio.to_thread(socket.getaddrinfo, host, None)
+    except Exception as e:
+        raise HTTPException(
+            status_code=422, detail="Failed to resolve webhook host"
+        ) from e
+
+    for info in infos:
+        ip_str = info[4][0]
+        try:
+            ip = ipaddress.ip_address(ip_str)
+        except ValueError:
+            # Skip non-IP entries just in case
+            continue
+        if (
+            ip.is_private
+            or ip.is_loopback
+            or ip.is_link_local
+            or ip.is_multicast
+            or ip.is_reserved
+        ):
+            raise HTTPException(
+                status_code=422, detail="Webhook host resolves to a disallowed IP"
+            )
+
+
 async def call_webhook(result: "WorkerResult") -> None:
     if HTTP_CONFIG and HTTP_CONFIG.get("disable_webhooks"):
         logger.info(
@@ -30,16 +133,27 @@ async def call_webhook(result: "WorkerResult") -> None:
     webhook = result.get("webhook")
     if webhook:
         try:
+            # We've already validated on ingestion, but you could technically have an issue if you re-deployed with a different environment
+            await validate_webhook_url_or_raise(webhook)
+            # Note: header templates should have already been evaluated against the env at load time.
+            headers = WEBHOOKS_CONFIG.get("headers") if WEBHOOKS_CONFIG else None
+
             if webhook.startswith("/"):
                 # Call into this own app
                 webhook_client = get_loopback_client()
             else:
-                webhook_client = get_http_client()
-            await http_request("POST", webhook, json=payload, client=webhook_client)
+                webhook_client = await ensure_http_client()
+            await http_request(
+                "POST",
+                webhook,
+                json=payload,
+                headers=headers,
+                client=webhook_client,
+            )
             await logger.ainfo(
                 "Background worker called webhook",
                 webhook=result["webhook"],
-                run_id=result["run"]["run_id"],
+                run_id=str(result["run"]["run_id"]),
             )
         except Exception as exc:
             logger.exception(