pugmark 0.1.0__py3-none-any.whl

pugmark/__init__.py

@@ -0,0 +1,1676 @@
+"""
+Pugmark Python SDK for subprocess communication.
+
+This module provides inter-process communication functionality for the Pugmark session log
+storage system.
+It enables communication between the main Pugmark process and Python subprocesses,
+allowing for data transformation and processing workflows.
+
+The SDK implements a JSON-based protocol for passing object data between processes via stdin/stdout.
+
+Architecture:
+- The main process serves objects via HTTP and sends input records via stdin
+- Python subprocesses receive input records, fetch object data via HTTP, process it,
+  and send output records via stdout
+- The main process collects output records and converts them back to Pugmark objects
+
+Example usage:
+
+    import pugmark
+
+    for obj in pugmark.read():
+        # Access data through the unified Object interface
+        data = obj.body()
+        if obj.content_type().startswith("text/"):
+            text_data = data.decode("utf-8")
+            processed = text_data.upper()
+        elif obj.content_type() == "application/json":
+            import json
+            json_data = json.loads(data.decode("utf-8"))
+            processed = {"processed": json_data}
+        else:
+            processed = {"binary_length": len(data)}
+
+        # Send output (automatically encodes based on type)
+        pugmark.write(processed)
+
+        # Optionally pause execution for interactive workflows
+        if needs_approval:
+            pugmark.pause("Processing complete, awaiting user approval")
+
+Key Functions:
+- pugmark.read(): Read input objects from the parent process (returns RemoteObject instances)
+- pugmark.write(): Send output objects to the parent process (returns LocalObject instances)
+- pugmark.pause(): Pause execution with an optional reason
+- pugmark.load(): Load objects by name from session state
+- pugmark.store(): Store objects by name in session state
+- pugmark.open(): Context manager for reading/writing named objects
+
+Object Types:
+- Object: Abstract base class for all objects
+- RemoteObject: Objects backed by Input data (from HTTP URLs)
+- LocalObject: Objects backed by Output data (local data)
+
+The API now matches the Go ipc package structure with consistent Input/Output types.
+"""
+
+import base64
+import builtins
+import gzip
+import json
+import os
+import sys
+import tempfile
+import time
+import urllib.error
+import urllib.parse
+import urllib.request
+from contextlib import contextmanager
+from contextvars import Token
+from dataclasses import dataclass, field
+from typing import IO, Any, BinaryIO, Dict, Iterator, List, Optional, Protocol, TextIO, Union
+
+import zstandard as zstd
+from zstandard import ZstdDecompressionReader
+
+__version__ = "0.1.0"
+
+# Optional OpenTelemetry imports for trace context propagation
+try:
+    from opentelemetry import context as otel_context
+    from opentelemetry import propagate as otel_propagate
+    from opentelemetry.instrumentation.urllib import URLLibInstrumentor
+
+    _HAS_OTEL = True
+except ImportError:
+    _HAS_OTEL = False
+
+# Optional OpenTelemetry trace imports for span creation
+try:
+    from opentelemetry import trace as otel_trace
+
+    _HAS_OTEL_TRACE = True
+except ImportError:
+    _HAS_OTEL_TRACE = False
+
+
+def _get_tracer() -> Optional["otel_trace.Tracer"]:
+    """Get the OpenTelemetry tracer for pugmark."""
+    if not _HAS_OTEL_TRACE:
+        return None
+    return otel_trace.get_tracer("github.com/firetiger-oss/pugmark")
+
+
+# Context token for cleanup (used when trace context is attached)
+_otel_context_token: Optional[Token[object]] = None
+
+
+def _init_trace_context_from_env() -> None:
+    """Extract and attach trace context from environment variables at startup.
+
+    This function reads W3C trace context (TRACEPARENT, TRACESTATE) from
+    environment variables and attaches it to the current OpenTelemetry context.
+    This enables trace continuity from Go parent processes to Python subprocesses.
+    """
+    global _otel_context_token
+    if not _HAS_OTEL:
+        return
+
+    traceparent = os.environ.get("TRACEPARENT", "")
+    tracestate = os.environ.get("TRACESTATE", "")
+
+    if not traceparent:
+        return
+
+    carrier = {
+        "traceparent": traceparent,
+        "tracestate": tracestate,
+    }
+    ctx = otel_propagate.extract(carrier)
+    _otel_context_token = otel_context.attach(ctx)
+
+
+def _init_urllib_instrumentation() -> None:
+    """Instrument urllib for automatic trace propagation in HTTP requests.
+
+    When OpenTelemetry is available, this instruments the urllib library so that
+    outgoing HTTP requests (such as fetching objects from the Go object server)
+    automatically include trace context headers.
+    """
+    if not _HAS_OTEL:
+        return
+    try:
+        URLLibInstrumentor().instrument()
+    except Exception:
+        pass  # Already instrumented or unavailable
+
+
+# Initialize trace context when module is imported (for subprocess mode)
+_init_trace_context_from_env()
+_init_urllib_instrumentation()
+
+_print = builtins.print
+
+
+def print(*args: Any, **kwargs: Any) -> Any:
+    kwargs.setdefault("file", sys.stderr)
+    return _print(*args, **kwargs)
+
+
+def _parse_media_type(content_type: str) -> str:
+    """Parse media type from content type string, ignoring parameters."""
+    if not content_type:
+        return ""
+    return content_type.split(";")[0].strip()
+
+
+# Override built-in print to redirect to stderr because stdout is used to
+# output pugmark records, this helps avoid confusion when the program uses
+# print statements for debugging.
+builtins.print = print
+
+
+@dataclass
+class Input:
+    """Represents an input record sent from the main Pugmark process."""
+
+    url: str = ""
+    size: int = 0
+    event: str = ""
+    content_type: str = ""
+    content_encoding: str = ""
+    schema_url: str = ""
+    name: str = ""
+    metadata: Optional[Dict[str, str]] = None
+    data: Optional[bytes] = None
+
+    def __post_init__(self) -> None:
+        if self.metadata is None:
+            self.metadata = {}
+
+
+@dataclass
+class Output:
+    """Represents an output record sent to the main Pugmark process.
+
+    Either ``data`` or ``uri`` must be set, but not both.  Use ``uri`` to
+    reference an object already stored at a remote location (e.g.
+    ``s3://bucket/path/to/object``); pugmark will fetch it, compute its SHA-256,
+    and store it in the session without the data travelling through this process.
+    """
+
+    data: bytes = b""
+    uri: str = ""
+    name: str = ""
+    content_type: str = ""
+    content_encoding: str = ""
+    schema_url: str = ""
+    metadata: Optional[Dict[str, str]] = field(default_factory=dict)
+    fork: bool = False  # Marks this as a fork message that creates a child session
+    session: str = ""  # Target session ID for cross-session writes
+
+
+class Object(Protocol):
+    """
+    Protocol defining the interface for Pugmark objects.
+
+    Objects can be either local (backed by Output data) or remote (backed by Input data).
+    This unified interface matches the Go ipc package structure.
+    """
+
+    def body(self) -> bytes:
+        """Get the object's data as bytes."""
+        ...
+
+    def content_type(self) -> str:
+        """Get the content type of the object."""
+        ...
+
+    def schema_url(self) -> str:
+        """Get the schema URL of the object."""
+        ...
+
+    def name(self) -> str:
+        """Get the name of the object."""
+        ...
+
+    def event(self) -> str:
+        """Get the event type associated with this object."""
+        ...
+
+    def metadata(self) -> Dict[str, str]:
+        """Get the metadata dictionary of the object."""
+        ...
+
+    def fork(self) -> bool:
+        """Return True if this object is a fork message that creates a child session."""
+        ...
+
+    def decode_data(self) -> bytes:
+        """Get the object's data as raw bytes (alias for body())."""
+        ...
+
+    def decode_text(self) -> str:
+        """Decode the object's data as UTF-8 text."""
+        ...
+
+    def decode_json(self) -> Any:
+        """Decode the object's data as JSON."""
+        ...
+
+
+# Event classes for the events API
+class Event:
+    """
+    Base class for pugmark session events.
+
+    All events provide a string representation through the __str__ method,
+    matching the fmt.Stringer interface from the Go implementation.
+    """
+
+    def __str__(self) -> str:
+        """Return the string representation of the event."""
+        raise NotImplementedError
+
+
+class StartEvent(Event):
+    """
+    Indicates that a new root session has been created.
+
+    This event is emitted when a session is started without a parent session,
+    marking the beginning of a new conversation or workflow.
+    """
+
+    def __str__(self) -> str:
+        """Return the string representation of a StartEvent."""
+        return "start"
+
+
+class ForkEvent(Event):
+    """
+    Indicates that a new child session has been created from a parent session.
+
+    This event is emitted when a session is created with a parent reference,
+    allowing for branching conversations and parallel execution paths.
+    """
+
+    def __str__(self) -> str:
+        """Return the string representation of a ForkEvent."""
+        return "fork"
+
+
+class WakeEvent(Event):
+    """
+    Indicates that a session has been resumed or awakened.
+
+    This event is typically emitted when a session becomes active after
+    being paused or sleeping.
+    """
+
+    def __str__(self) -> str:
+        """Return the string representation of a WakeEvent."""
+        return "wake"
+
+
+class PushEvent(Event):
+    """
+    Indicates that an object has been added to the session.
+
+    This event contains the Object that was pushed, allowing access to
+    its content, metadata, and other properties.
+    """
+
+    def __init__(self, obj: Object):
+        """
+        Initialize a PushEvent with an Object.
+
+        Args:
+            obj: The Object that was pushed to the session
+        """
+        self.object = obj
+
+    def __str__(self) -> str:
+        """Return the string representation of a PushEvent."""
+        return "push"
+
+
+class RemoteObject:
+    """
+    Represents a remote object accessible via HTTP.
+
+    RemoteObjects are created from Input records and fetch data from remote URLs.
+    This matches the concept of remote objects in the Go ipc package.
+    """
+
+    def __init__(self, input_record: Input):
+        """
+        Initialize a RemoteObject from an Input record.
+
+        Args:
+            input_record: Input record containing URL and metadata
+            name: Optional name for the object
+        """
+        self._input = input_record
+        self._body: Optional[bytes] = None
+
+    def body(self) -> bytes:
+        """
+        Get the object's data as bytes.
+
+        Returns:
+            Raw bytes of the object data
+        """
+        if self._body is not None:
+            return self._body
+
+        # Return empty bytes if URL is not provided (for events without data)
+        if not self._input.url:
+            self._body = b""
+            return self._body
+
+        with self._read() as response:
+            self._body = response.read()
+            return self._body
+
+    def content_type(self) -> str:
+        """Get the content type of the object."""
+        return self._input.content_type
+
+    def schema_url(self) -> str:
+        """Get the schema URL of the object."""
+        return self._input.schema_url
+
+    def name(self) -> str:
+        """Get the name of the object."""
+        return self._input.name
+
+    def event(self) -> str:
+        """Get the event type associated with this object."""
+        return self._input.event
+
+    def metadata(self) -> Dict[str, str]:
+        """Get the metadata dictionary of the object."""
+        return self._input.metadata or {}
+
+    def fork(self) -> bool:
+        """Return False - remote objects from input are never fork markers."""
+        return False
+
+    def decode_data(self) -> bytes:
+        """Get the object's data as raw bytes (alias for body())."""
+        return self.body()
+
+    def decode_text(self) -> str:
+        """Decode the object's data as UTF-8 text."""
+        return self.body().decode("utf-8")
+
+    def decode_json(self) -> Any:
+        """Decode the object's data as JSON."""
+        return json.loads(self.body())
+
+    def __str__(self) -> str:
+        """Return a string representation of the object."""
+        return self.name() or self._input.url or "<remote object>"
+
+    def __repr__(self) -> str:
+        return f"RemoteObject(body={self.body()}, content_type={self.content_type()})"
+
+    def _read(self) -> Union[gzip.GzipFile, ZstdDecompressionReader, BinaryIO]:
+        """
+        Return a file-like object for streaming the object's data.
+
+        The returned reader handles content encoding decompression automatically based
+        on the Input record's content_encoding field. Retries on ConnectionRefusedError
+        to handle the race condition where the Go-side HTTP server is shutting down
+        while Python still has pending reads.
+
+        Returns:
+            File-like object for reading object data
+        """
+        last_error: Optional[Exception] = None
+        for attempt in range(3):
+            try:
+                response = urllib.request.urlopen(self._input.url)
+            except urllib.error.URLError as exc:
+                if isinstance(exc.reason, ConnectionRefusedError) and attempt < 2:
+                    time.sleep(0.1 * (2**attempt))
+                    last_error = exc
+                    continue
+                raise
+            content_encoding = self._input.content_encoding
+            if content_encoding == "gzip":
+                return gzip.GzipFile(fileobj=response)
+            elif content_encoding == "zstd":
+                decompressor = zstd.ZstdDecompressor()
+                return decompressor.stream_reader(response)
+            else:
+                return response  # type: ignore[no-any-return]
+        raise last_error  # type: ignore[misc]
+
+
+class LocalObject:
+    """
+    Represents a local object backed by Output data.
+
+    LocalObjects are created from Output records and contain data locally.
+    This matches the concept of local objects in the Go ipc package.
+    """
+
+    def __init__(self, output_record: Output, event: str = "push"):
+        """
+        Initialize a LocalObject from an Output record.
+
+        Args:
+            output_record: Output record containing data and metadata
+            event: Event type for this object
+        """
+        self._output = output_record
+        self._event = event
+        self._body: Optional[bytes] = None
+
+    def body(self) -> bytes:
+        """Get the object's data as bytes, automatically decompressed if needed."""
+        if self._body is not None:
+            return self._body
+
+        data = self._output.data
+        content_encoding = self._output.content_encoding
+
+        if content_encoding == "gzip":
+            import gzip
+
+            self._body = gzip.decompress(data)
+        elif content_encoding == "zstd":
+            decompressor = zstd.ZstdDecompressor()
+            self._body = decompressor.decompress(data)
+        else:
+            self._body = data
+
+        return self._body
+
+    def content_type(self) -> str:
+        """Get the content type of the object."""
+        return self._output.content_type
+
+    def schema_url(self) -> str:
+        """Get the schema URL of the object."""
+        return self._output.schema_url
+
+    def name(self) -> str:
+        """Get the name of the object."""
+        return self._output.name
+
+    def event(self) -> str:
+        """Get the event type associated with this object."""
+        return self._event
+
+    def metadata(self) -> Dict[str, str]:
+        """Get the metadata dictionary of the object."""
+        return self._output.metadata or {}
+
+    def fork(self) -> bool:
+        """Return True if this object is a fork message that creates a child session."""
+        return self._output.fork
+
+    def decode_data(self) -> bytes:
+        """Get the object's data as raw bytes (alias for body())."""
+        return self.body()
+
+    def decode_text(self) -> str:
+        """Decode the object's data as UTF-8 text."""
+        return self.body().decode("utf-8")
+
+    def decode_json(self) -> Any:
+        """Decode the object's data as JSON."""
+        return json.loads(self.body())
+
+    def __str__(self) -> str:
+        """Return a string representation of the object."""
+        return self.name() or "<local object>"
+
+
+def _is_object(value: Any) -> bool:
+    """Duck-typed check for whether a value already implements the Object protocol."""
+    return all(
+        callable(getattr(value, attr, None))
+        for attr in ("body", "name", "content_type", "metadata")
+    )
+
+
+def _local_object(
+    value: Any,
+    *,
+    name: str = "",
+    content_type: Optional[str] = None,
+    content_encoding: str = "",
+    schema_url: str = "",
+    metadata: Optional[Dict[str, str]] = None,
+) -> "LocalObject":
+    """
+    Build a LocalObject from a raw Python value. Content type is auto-detected
+    when not provided:
+
+      - str    → "text/plain"
+      - bytes  → "application/octet-stream"
+      - dict, list, int, float, bool, None → "application/json"
+
+    Raises TypeError for unsupported types.
+    """
+    if isinstance(value, str):
+        data = value.encode("utf-8")
+        ct = content_type or "text/plain"
+    elif isinstance(value, (bytes, bytearray)):
+        data = bytes(value)
+        ct = content_type or "application/octet-stream"
+    elif isinstance(value, (dict, list, int, float, bool)) or value is None:
+        data = json.dumps(value).encode("utf-8")
+        ct = content_type or "application/json"
+    else:
+        raise TypeError(
+            f"cannot write value of type {type(value).__name__}; "
+            "expected str, bytes, dict, list, or pugmark.Object"
+        )
+    return LocalObject(
+        Output(
+            data=data,
+            name=name,
+            content_type=ct,
+            content_encoding=content_encoding,
+            schema_url=schema_url,
+            metadata=metadata or {},
+        )
+    )
+
+
+class Reader:
+    """
+    Provides an iterator interface for reading Input records from a stream
+    and converting them to Object instances for processing.
+    """
+
+    def __init__(self, stream: TextIO):
+        """
+        Initialize a Reader instance.
+
+        Args:
+            stream: Source stream for Input records
+        """
+        self.stream = stream
+
+    def read(self) -> Iterator[Object]:
+        """
+        Return an iterator that yields Object instances from Input records.
+
+        Each Input record is read from the stream, decoded as JSON, and converted to a RemoteObject.
+        The iterator stops on the first decoding error or when the stream is exhausted.
+
+        Yields:
+            Object instances created from Input record data
+
+        Raises:
+            json.JSONDecodeError: If input JSON is malformed
+        """
+        for line in self.stream:
+            line = line.strip()
+            if not line:
+                continue
+
+            try:
+                data = json.loads(line)
+                record = Input(
+                    url=data.get("url", ""),
+                    size=data.get("size", 0),
+                    event=data.get("event", ""),
+                    content_type=data.get("content-type", ""),
+                    content_encoding=data.get("content-encoding", ""),
+                    schema_url=data.get("schema-url", ""),
+                    name=data.get("name", ""),
+                    metadata=data.get("metadata", {}),
+                )
+
+                if record.data is None:
+                    obj = RemoteObject(input_record=record)
+                else:
+                    obj = LocalObject(
+                        output_record=Output(
+                            data=record.data,
+                            name=record.name,
+                            content_type=record.content_type,
+                            content_encoding=record.content_encoding,
+                            schema_url=record.schema_url,
+                            metadata=record.metadata or {},
+                        ),
+                        event=record.event,
+                    )
+
+                yield obj
+            except json.JSONDecodeError as e:
+                raise RuntimeError(f"Failed to decode input record: {e}") from e
+
+
+class Writer:
+    """
+    Provides methods for sending Output records to a stream.
+
+    Handles JSON encoding, buffering, and automatic compression.
+    """
+
+    def __init__(self, stream: TextIO):
+        """
+        Initialize a Writer instance.
+
+        Args:
+            stream: Destination stream for Output records
+        """
+        self.stream = stream
+        self._initialized = False
+
+    def write(self, obj: Object) -> None:
+        """
+        Write an Object to the stream.
+
+        Args:
+            obj: The Object to write
+        """
+        # Get data and metadata from the object
+        data = obj.body()
+        name = obj.name()
+        content_type = obj.content_type()
+        schema_url = obj.schema_url()
+        metadata = obj.metadata()
+        is_fork = obj.fork()
+
+        # Create Output record from Object data
+        output = Output(
+            data=data,
+            name=name,
+            content_type=content_type,
+            content_encoding="",  # Content encoding handled internally by Object
+            schema_url=schema_url,
+            metadata=metadata,
+            fork=is_fork,
+        )
+
+        # Convert to JSON format expected by main process
+        json_data: Dict[str, Any] = {
+            "data": base64.b64encode(output.data).decode("ascii"),  # Base64 encode binary data
+        }
+
+        if output.name:
+            json_data["name"] = output.name
+        if output.content_type:
+            json_data["content-type"] = output.content_type
+        if output.content_encoding:
+            json_data["content-encoding"] = output.content_encoding
+        if output.schema_url:
+            json_data["schema-url"] = output.schema_url
+        if output.metadata:
+            json_data["metadata"] = output.metadata
+        if output.fork:
+            json_data["fork"] = True
+        if output.session:
+            json_data["session"] = output.session
+
+        # Write JSON line to stream
+        json.dump(json_data, self.stream, separators=(",", ":"))
+        self.stream.write("\n")
+        self.stream.flush()
+
+    def write_remote(
+        self,
+        uri: str,
+        *,
+        name: str = "",
+        content_type: str = "",
+        content_encoding: str = "",
+        schema_url: str = "",
+        metadata: Optional[Dict[str, str]] = None,
+        session: str = "",
+    ) -> None:
+        """
+        Write a remote object reference to the stream.
+
+        Instead of sending the object's bytes inline, pugmark will fetch the
+        object at ``uri`` (e.g. ``s3://bucket/path/to/object``), compute its
+        SHA-256, and store it in the session.
+
+        Args:
+            uri: Remote object URI (e.g. ``s3://bucket/path/to/object``)
+            name: Optional name for the object
+            content_type: Optional MIME type override (defaults to remote value)
+            content_encoding: Optional encoding override (defaults to remote value)
+            schema_url: Optional schema URL
+            metadata: Optional metadata key-value pairs
+            session: Optional target session ID for cross-session writes
+        """
+        json_data: Dict[str, Any] = {"uri": uri}
+        if name:
+            json_data["name"] = name
+        if content_type:
+            json_data["content-type"] = content_type
+        if content_encoding:
+            json_data["content-encoding"] = content_encoding
+        if schema_url:
+            json_data["schema-url"] = schema_url
+        if metadata:
+            json_data["metadata"] = metadata
+        if session:
+            json_data["session"] = session
+
+        json.dump(json_data, self.stream, separators=(",", ":"))
+        self.stream.write("\n")
+        self.stream.flush()
+
+
+@dataclass
+class State:
+    """
+    Represents the current state of all objects read from stdin.
+    This class maintains a dictionary of named objects and a list of all objects
+    to allow easy access and iteration over the objects processed during the session.
+    """
+
+    _dict: Dict[str, Object]
+    _list: List[Object]
+    _memory_namespaces: Dict[str, str]  # namespace -> URL mapping
+
+    def __init__(self) -> None:
+        self._dict = {}
+        self._list = []
+        self._memory_namespaces = {}
+
+    def __iter__(self) -> Iterator[Object]:
+        """
+        Iterate over all objects in the state.
+
+        Yields:
+            Object instances stored in the state
+        """
+        return iter(self._list)
+
+    def __len__(self) -> int:
+        """
+        Get the number of objects in the state.
+
+        Returns:
+            The count of Object instances stored in the state
+        """
+        return len(self._dict)
+
+    def __getitem__(self, key: str) -> Object:
+        """
+        Get an object by name from the state.
+
+        Args:
+            key: The name of the object to retrieve
+
+        Returns:
+            The Object instance associated with the given name
+        """
+        return self._dict[key]
+
+    def __setitem__(self, key: str, value: Object) -> None:
+        """
+        Set an object in the state by name.
+
+        Args:
+            key: The name of the object to set
+            value: The Object instance to associate with the name
+        """
+        self._dict[key] = value
+        self._list.append(value)
+
+    def __contains__(self, key: str) -> bool:
+        """
+        Check if an object with the given name exists in the state.
+
+        Args:
+            key: The name of the object to check
+
+        Returns:
+            True if the object exists, False otherwise
+        """
+        return key in self._dict
+
+    def append(self, value: Object) -> None:
+        if value.name():
+            self._dict[value.name()] = value
+        self._list.append(value)
+
+    def add_memory_namespace(self, namespace: str, url: str) -> None:
+        """
+        Add a memory namespace URL mapping.
+
+        Args:
+            namespace: The namespace name
+            url: The base URL for this namespace
+        """
+        self._memory_namespaces[namespace] = url
+
+    def get_memory_url(self, name: str) -> Optional[str]:
+        """
+        Get the memory URL for an object name if it's in a memory namespace.
+
+        Args:
+            name: The object name to check (e.g., "agents/gpt4/config.txt")
+
+        Returns:
+            The full URL if the name matches a memory namespace, None otherwise
+        """
+        if "/" not in name:
+            return None
+
+        # Sort keys and search for the first match in reverse order
+        # This finds the longest matching namespace
+        keys = sorted(self._memory_namespaces.keys(), reverse=True)
+        for key in keys:
+            if name.startswith(key):
+                object_name = name[len(key) :]
+                url = self._memory_namespaces[key] + object_name
+                return url
+
+        return None
+
+
+class Session:
+    """
+    Encapsulates the pugmark communication session with configurable I/O streams.
+
+    This allows for dependency injection of reader/writer streams for testing
+    while maintaining the same API as the global functions.
+    """
+
+    def __init__(
+        self,
+        reader_stream: Optional[TextIO] = None,
+        writer_stream: Optional[TextIO] = None,
+        agent: Optional[str] = None,
+        session_id: Optional[str] = None,
+        snapshot: Optional[int] = None,
+        metadata: Optional[Dict[str, str]] = None,
+    ):
+        """
+        Initialize a Session with optional custom streams.
+
+        Args:
+            reader_stream: Input stream for reading records (defaults to sys.stdin)
+            writer_stream: Output stream for writing records (defaults to sys.stdout)
+            agent: Agent identifier (defaults to PUGMARK_AGENT_ID env var)
+            session_id: Session identifier (defaults to PUGMARK_SESSION_ID env var)
+            snapshot: Snapshot version number
+            metadata: Session metadata from manifest (defaults to PUGMARK_SESSION_METADATA env var)
+        """
+        self.reader = Reader(reader_stream or sys.stdin)
+        self.writer = Writer(writer_stream or sys.stdout)
+        self._state = State()
+        # Session ID, agent, and snapshot can be passed directly or via environment variables
+        self._session_id: str = (
+            session_id if session_id is not None else os.environ.get("PUGMARK_SESSION_ID", "")
+        )
+        self._agent: str = agent if agent is not None else os.environ.get("PUGMARK_AGENT_ID", "")
+        self._snapshot: Optional[int] = snapshot
+        # Metadata can be passed directly or via environment variable (JSON-encoded)
+        if metadata is not None:
+            self._metadata: Dict[str, str] = metadata
+        else:
+            metadata_json = os.environ.get("PUGMARK_SESSION_METADATA", "")
+            if metadata_json:
+                try:
+                    self._metadata = json.loads(metadata_json)
+                except json.JSONDecodeError:
+                    self._metadata = {}
+            else:
+                self._metadata = {}
+
+    def events(self) -> Iterator[Event]:
+        """
+        Return an iterator over Event instances parsed from the session's input stream.
+
+        This method provides a higher-level abstraction over the raw read() method by
+        interpreting Objects as typed events based on their event() field.
+
+        The iterator yields events in the order they are received from the underlying
+        input stream. Events are parsed according to the following mapping:
+        - "start" → StartEvent: Session creation without a parent
+        - "fork" → ForkEvent: Session creation with a parent
+        - "wake" → WakeEvent: Session resumption or awakening
+        - "push" → PushEvent: Object addition (contains the actual Object)
+        - Other event types are silently ignored
+
+        Example usage:
+            for event in session.events():
+                match event:
+                    case StartEvent():
+                        print("Session started")
+                    case ForkEvent():
+                        print("Session forked")
+                    case WakeEvent():
+                        print("Session awakened")
+                    case PushEvent() as push:
+                        print(f"Object pushed: {push.object.name()}")
+
+        Returns:
+            Iterator yielding Event instances from the input stream
+        """
+        for obj in self.read():
+            event_type = obj.event()
+            if event_type == "start":
+                yield StartEvent()
+            elif event_type == "fork":
+                yield ForkEvent()
+            elif event_type == "wake":
+                yield WakeEvent()
+            elif event_type == "push":
+                yield PushEvent(obj)
+
+    def read(self) -> Iterator[Object]:
+        """
+        Read objects from the input stream.
+
+        Returns:
+            Iterator yielding Object instances from the input stream
+        """
+        for obj in self.reader.read():
+            # Handle memory namespace events
+            if obj.event() == "memory":
+                # Extract namespace from URL path
+                url = obj._input.url  # Access the URL directly
+                if url.endswith("/"):
+                    # Extract namespace from URL like "http://127.0.0.1:port/namespace/"
+                    namespace = url.split("/")[-2] + "/"  # Keep the trailing slash
+                    self._state.add_memory_namespace(namespace, url)
+                # Don't yield memory events to the user
+                continue
+
+            self._state.append(obj)
+            yield obj
+
+    def write(
+        self,
+        value: Any,
+        *,
+        name: str = "",
+        content_type: Optional[str] = None,
+        content_encoding: str = "",
+        schema_url: str = "",
+        metadata: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """
+        Write a value to the output stream.
+
+        ``value`` may be a :class:`pugmark.Object` (written as-is) or a raw
+        Python value (str, bytes, dict, list, or any JSON-serializable scalar),
+        in which case it is wrapped in a :class:`LocalObject`. Content type is
+        auto-detected from the value's type unless explicitly provided.
+        """
+        if _is_object(value):
+            obj = value
+        else:
+            obj = _local_object(
+                value,
+                name=name,
+                content_type=content_type,
+                content_encoding=content_encoding,
+                schema_url=schema_url,
+                metadata=metadata,
+            )
+        self.writer.write(obj)
+        self._state.append(obj)
+
+    def write_text(
+        self,
+        text: str,
+        *,
+        name: str = "",
+        content_type: str = "text/plain",
+        metadata: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """Write a string as a text/plain object (or other text/* content type)."""
+        self.write(text, name=name, content_type=content_type, metadata=metadata)
+
+    def write_json(
+        self,
+        data: Any,
+        *,
+        name: str = "",
+        content_type: str = "application/json",
+        metadata: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """Write a JSON-serializable value as an application/json object."""
+        self.write(data, name=name, content_type=content_type, metadata=metadata)
+
+    def write_bytes(
+        self,
+        data: Union[bytes, bytearray],
+        *,
+        name: str = "",
+        content_type: str = "application/octet-stream",
+        metadata: Optional[Dict[str, str]] = None,
+    ) -> None:
+        """Write raw bytes with a configurable content type."""
+        self.write(bytes(data), name=name, content_type=content_type, metadata=metadata)
+
+    def write_remote(
+        self,
+        uri: str,
+        *,
+        name: str = "",
+        content_type: str = "",
+        content_encoding: str = "",
+        schema_url: str = "",
+        metadata: Optional[Dict[str, str]] = None,
+        session: str = "",
+    ) -> None:
+        """
+        Write a remote object reference to the output stream.
+
+        Instead of sending the object's bytes inline, pugmark will fetch the
+        object at ``uri`` (e.g. ``s3://bucket/path/to/object``), compute its
+        SHA-256, and store it in the session.
+
+        Args:
+            uri: Remote object URI (e.g. ``s3://bucket/path/to/object``)
+            name: Optional name for the object
+            content_type: Optional MIME type override (defaults to remote value)
+            content_encoding: Optional encoding override (defaults to remote value)
+            schema_url: Optional schema URL
+            metadata: Optional metadata key-value pairs
+            session: Optional target session ID for cross-session writes
+        """
+        self.writer.write_remote(
+            uri,
+            name=name,
+            content_type=content_type,
+            content_encoding=content_encoding,
+            schema_url=schema_url,
+            metadata=metadata,
+            session=session,
+        )
+
+    def state(self) -> "State":
+        """
+        Get the current state of objects read from the input stream.
+
+        Returns:
+            The State instance containing all objects seen so far
+        """
+        for _ in self.read():
+            pass
+        return self._state
+
+    def load(self, name: str) -> Optional[Object]:
+        """
+        Load an object by name from the state.
+
+        Args:
+            name: The name of the object to retrieve
+
+        Returns:
+            The Object instance, or None if not found
+        """
+        objects = self.state()
+
+        # Check if the name corresponds to a memory namespace object
+        memory_url = objects.get_memory_url(name)
+        if memory_url:
+            # Create a RemoteObject that points to the memory URL
+            memory_input = Input(url=memory_url, name=name)
+            return RemoteObject(memory_input)
+
+        if name in objects:
+            return objects[name]
+
+        return None
+
+    def load_text(self, name: str) -> Optional[str]:
+        """Load a named object and decode it as UTF-8 text."""
+        obj = self.load(name)
+        return obj.decode_text() if obj is not None else None
+
+    def load_json(self, name: str) -> Any:
+        """Load a named object and decode it as JSON. Returns None if absent."""
+        obj = self.load(name)
+        return obj.decode_json() if obj is not None else None
+
+    def load_bytes(self, name: str) -> Optional[bytes]:
+        """Load a named object and return its raw bytes."""
+        obj = self.load(name)
+        return obj.decode_data() if obj is not None else None
+
+    def store(self, obj: Object) -> None:
+        """
+        Store an object in the state.
+
+        Args:
+            obj: The Object instance to store
+        """
+        self.write(obj)
+
+    def pause(self, reason: str = "") -> None:
+        """
+        Pause the execution of the current Pugmark subprocess.
+
+        Args:
+            reason: An optional string explaining why execution is being paused
+        """
+        output = Output(
+            data=reason.encode("utf-8"),
+            name="",
+            content_type="application/vnd.pugmark.pause+text",
+            content_encoding="",
+            schema_url="",
+            metadata={},
+        )
+        obj = LocalObject(output)
+        self.write(obj)
+
+    def sleep(self, seconds: float, reason: str = "") -> None:
+        """
+        Sleep for a specified number of seconds, pausing execution.
+
+        Args:
+            seconds: Number of seconds to sleep
+            reason: Optional reason for the pause
+        """
+        output = Output(
+            data=reason.encode("utf-8"),
+            name="",
+            content_type="application/vnd.pugmark.pause+text",
+            content_encoding="",
+            schema_url="",
+            metadata={
+                "duration": f"{seconds:.2f}s",
+            },
+        )
+        obj = LocalObject(output)
+        self.write(obj)
+
+    def session_id(self) -> Optional[str]:
+        """
+        Get the pugmark session ID.
+
+        The session ID is read from the PUGMARK_SESSION_ID environment variable,
+        which is automatically set by the pugmark runner when spawning subprocesses.
+
+        Returns:
+            The session ID string if available, None otherwise
+        """
+        return self._session_id if self._session_id else None
+
+    def agent(self) -> Optional[str]:
+        """
+        Get the agent identity.
+
+        The agent ID is read from the PUGMARK_AGENT_ID environment variable,
+        which is automatically set by the pugmark runner when spawning subprocesses.
+        This enables multi-agent routing where handlers can route
+        requests to different implementations based on the agent.
+
+        Returns:
+            The agent identity string if available, None otherwise
+        """
+        return self._agent if self._agent else None
+
+    def snapshot(self) -> Optional[int]:
+        """
+        Get the snapshot version number.
+
+        The snapshot is the version of the session at the time of invocation.
+        This is passed via query parameters when running in RPC mode.
+
+        Returns:
+            The snapshot version number if available, None otherwise
+        """
+        return self._snapshot
+
+    def metadata(self) -> Dict[str, str]:
+        """
+        Get the session metadata from the manifest.
+
+        Metadata is passed via different mechanisms depending on the runner type:
+        - Command runner: PUGMARK_SESSION_METADATA environment variable (JSON-encoded)
+        - HTTP runner: X-Pugmark-Session-Metadata header (JSON-encoded)
+        - Function runner: Passed directly to the Session constructor
+
+        Returns:
+            Dictionary of metadata key-value pairs, empty dict if none set
+        """
+        return self._metadata
+
+    @contextmanager
+    def traced_execution(self) -> Iterator[None]:
+        """
+        Context manager that creates a tracing span for handler execution.
+
+        When OpenTelemetry is available, this creates a "pugmark.handler" span with
+        session context attributes. The span uses CONSUMER kind to indicate that
+        this is processing a message/event from the parent process.
+
+        Attributes included in the span:
+        - pugmark.session: The session ID
+        - pugmark.agent: The agent ID (if available)
+        - pugmark.snapshot: The snapshot version (if available)
+
+        If OpenTelemetry is not available, this is a no-op.
+
+        Yields:
+            None
+        """
+        tracer = _get_tracer()
+        if tracer is None:
+            yield
+            return
+
+        attributes: Dict[str, Any] = {}
+        if self._session_id:
+            attributes["pugmark.session"] = self._session_id
+        if self._agent:
+            attributes["pugmark.agent"] = self._agent
+        if self._snapshot is not None:
+            attributes["pugmark.snapshot"] = self._snapshot
+
+        with tracer.start_as_current_span(
+            "pugmark.handler",
+            kind=otel_trace.SpanKind.CONSUMER,
+            attributes=attributes,
+        ):
+            yield
+
+    @contextmanager
+    def open(
+        self,
+        name: str,
+        content_type: Optional[str] = None,
+        content_encoding: Optional[str] = None,
+        schema_url: Optional[str] = None,
+        metadata: Optional[Dict[str, str]] = None,
+    ) -> Iterator[IO[bytes]]:
+        """
+        Context manager for opening an object by name.
+
+        Args:
+            name: The name of the object to open
+            content_type: Optional content type for the object
+            content_encoding: Optional content encoding for the object
+            schema_url: URL of the schema associated with this object
+            metadata: Optional metadata dictionary for the object
+
+        Yields:
+            A temporary file-like object containing the object's data
+        """
+        # Try to load the object (this will check memory namespaces too)
+        obj = self.load(name)
+        if obj:
+            body = obj.body()
+            data = body
+        else:
+            body = b""
+            data = b""
+
+        temp = tempfile.NamedTemporaryFile(delete=True, mode="w+b")
+        try:
+            temp.write(data)
+            temp.seek(0)
+            yield temp
+            temp.flush()
+            temp.seek(0)
+            data = temp.read()
+        finally:
+            temp.close()
+
+        if data != body:
+            output = Output(
+                data=data,
+                name=name or "",
+                content_type=content_type or "",
+                content_encoding=content_encoding or "",
+                schema_url=schema_url or "",
+                metadata=metadata or {},
+            )
+            obj = LocalObject(output)
+            self.store(obj)
+
+
+# Default session instance for backward compatibility
+_default_session = Session()
+
+
+# Legacy global functions that use the default session
+stdin = _default_session.reader
+stdout = _default_session.writer
+
+
+def default_session() -> Session:
+    """
+    Get the default session instance used by global functions.
+
+    This allows access to the underlying session for advanced use cases
+    like accessing the session's state directly or replacing the default
+    session for testing purposes.
+
+    Returns:
+        The default Session instance
+    """
+    return _default_session
+
+
+def events() -> Iterator[Event]:
+    """
+    Return an iterator over events from the default session.
+
+    This is a convenience function that provides access to the event stream
+    without needing to manage a Session instance directly.
+
+    Example usage:
+        import pugmark
+
+        for event in pugmark.events():
+            match event:
+                case pugmark.StartEvent():
+                    print("Session started")
+                case pugmark.PushEvent() as push:
+                    print(f"Object pushed: {push.object.name()}")
+
+    Returns:
+        Iterator yielding Event instances from the default session
+    """
+    return _default_session.events()
+
+
+def read() -> Iterator[Object]:
+    """
+    Convenience function that returns an iterator using the default session.
+
+    This is the most common way to read Input records in subprocess implementations.
+
+    Returns:
+        Iterator yielding Object instances from stdin
+    """
+    return _default_session.read()
+
+
+def write(
+    value: Any,
+    *,
+    name: str = "",
+    content_type: Optional[str] = None,
+    content_encoding: str = "",
+    schema_url: str = "",
+    metadata: Optional[Dict[str, str]] = None,
+) -> None:
+    """
+    Convenience function that writes a value using the default session.
+
+    ``value`` may be a :class:`pugmark.Object` (written as-is) or a raw Python
+    value (str, bytes, dict, list, or any JSON-serializable scalar). Content
+    type is auto-detected from the value's type unless explicitly provided.
+
+    Examples::
+
+        pugmark.write("hello")                            # text/plain
+        pugmark.write({"answer": 42}, name="reply")       # application/json
+        pugmark.write(b"\\x00\\x01", name="bin")            # application/octet-stream
+    """
+    _default_session.write(
+        value,
+        name=name,
+        content_type=content_type,
+        content_encoding=content_encoding,
+        schema_url=schema_url,
+        metadata=metadata,
+    )
+
+
+def write_text(
+    text: str,
+    *,
+    name: str = "",
+    content_type: str = "text/plain",
+    metadata: Optional[Dict[str, str]] = None,
+) -> None:
+    """Write a string using the default session."""
+    _default_session.write_text(text, name=name, content_type=content_type, metadata=metadata)
+
+
+def write_json(
+    data: Any,
+    *,
+    name: str = "",
+    content_type: str = "application/json",
+    metadata: Optional[Dict[str, str]] = None,
+) -> None:
+    """Write a JSON-serializable value using the default session."""
+    _default_session.write_json(data, name=name, content_type=content_type, metadata=metadata)
+
+
+def write_bytes(
+    data: Union[bytes, bytearray],
+    *,
+    name: str = "",
+    content_type: str = "application/octet-stream",
+    metadata: Optional[Dict[str, str]] = None,
+) -> None:
+    """Write raw bytes using the default session."""
+    _default_session.write_bytes(data, name=name, content_type=content_type, metadata=metadata)
+
+
+def write_remote(
+    uri: str,
+    *,
+    name: str = "",
+    content_type: str = "",
+    content_encoding: str = "",
+    schema_url: str = "",
+    metadata: Optional[Dict[str, str]] = None,
+    session: str = "",
+) -> None:
+    """
+    Write a remote object reference using the default session.
+
+    Instead of sending the object's bytes inline, pugmark will fetch the object
+    at ``uri`` (e.g. ``s3://bucket/path/to/object``), compute its SHA-256, and
+    store it in the session.
+
+    Args:
+        uri: Remote object URI (e.g. ``s3://bucket/path/to/object``)
+        name: Optional name for the object
+        content_type: Optional MIME type override (defaults to remote value)
+        content_encoding: Optional encoding override (defaults to remote value)
+        schema_url: Optional schema URL
+        metadata: Optional metadata key-value pairs
+        session: Optional target session ID for cross-session writes
+    """
+    _default_session.write_remote(
+        uri,
+        name=name,
+        content_type=content_type,
+        content_encoding=content_encoding,
+        schema_url=schema_url,
+        metadata=metadata,
+        session=session,
+    )
+
+
+def state() -> State:
+    """
+    Get the current state of objects read from stdin using the default session.
+
+    Returns:
+        The State instance containing all objects seen so far
+    """
+    return _default_session.state()
+
+
+def load(name: str) -> Optional[Object]:
+    """
+    Load an object by name from the state using the default session.
+
+    Args:
+        name: The name of the object to retrieve
+
+    Returns:
+        The Object instance, or None if not found
+    """
+    return _default_session.load(name)
+
+
+def load_text(name: str) -> Optional[str]:
+    """Load a named object and decode it as UTF-8 text."""
+    return _default_session.load_text(name)
+
+
+def load_json(name: str) -> Any:
+    """Load a named object and decode it as JSON. Returns None if absent."""
+    return _default_session.load_json(name)
+
+
+def load_bytes(name: str) -> Optional[bytes]:
+    """Load a named object and return its raw bytes."""
+    return _default_session.load_bytes(name)
+
+
+def store(object: Object) -> None:
+    """
+    Store an object in the state using the default session.
+
+    Args:
+        object: The Object instance to store
+    """
+    _default_session.store(object)
+
+
+@contextmanager
+def open(
+    name: str,
+    content_type: Optional[str] = None,
+    content_encoding: Optional[str] = None,
+    schema_url: Optional[str] = None,
+    metadata: Optional[Dict[str, str]] = None,
+) -> Iterator[IO[bytes]]:
+    """
+    Context manager for opening an object by name using the default session.
+
+    Args:
+        name: The name of the object to open
+        content_type: Optional content type for the object
+        content_encoding: Optional content encoding for the object
+        schema_url: URL of the schema associated with this object
+        metadata: Optional metadata dictionary for the object
+
+    Yields:
+        A temporary file-like object containing the object's data
+    """
+    with _default_session.open(name, content_type, content_encoding, schema_url, metadata) as f:
+        yield f
+
+
+def pause(reason: str = "") -> None:
+    """
+    Pause the execution of the current Pugmark subprocess using the default session.
+
+    Args:
+        reason: An optional string explaining why execution is being paused
+    """
+    _default_session.pause(reason)
+
+
+def sleep(seconds: float, reason: str = "") -> None:
+    """
+    Sleep for a specified number of seconds using the default session.
+
+    Args:
+        seconds: Number of seconds to sleep
+        reason: Optional reason for the pause
+    """
+    _default_session.sleep(seconds, reason)
+
+
+def session_id() -> Optional[str]:
+    """
+    Get the pugmark session ID.
+
+    This function returns the session ID from the PUGMARK_SESSION_ID environment
+    variable, which is automatically set when running inside a pugmark subprocess
+    (via `pugmark run` or `pugmark exec`). Returns None if not running under pugmark.
+
+    Returns:
+        The session ID string if available, None otherwise
+
+    Example:
+        import pugmark
+
+        sid = pugmark.session_id()
+        if sid:
+            print(f"Running in session: {sid}")
+    """
+    return _default_session.session_id()
+
+
+def agent() -> Optional[str]:
+    """
+    Get the agent identity.
+
+    This function returns the agent ID from the PUGMARK_AGENT_ID environment
+    variable, which is automatically set when running inside a pugmark subprocess.
+    This enables multi-agent routing where handlers can route
+    requests to different implementations based on the agent.
+
+    Returns:
+        The agent identity string if available, None otherwise
+
+    Example:
+        import pugmark
+
+        for event in pugmark.events():
+            agent_id = pugmark.agent()
+            if agent_id == "claude":
+                handle_claude(event)
+            elif agent_id == "gpt4":
+                handle_gpt4(event)
+    """
+    return _default_session.agent()
+
+
+def metadata() -> Dict[str, str]:
+    """
+    Get the session metadata from the manifest.
+
+    This function returns metadata passed to the pugmark runner. The metadata
+    is passed via different mechanisms depending on the runner type:
+    - Command runner: PUGMARK_SESSION_METADATA environment variable (JSON-encoded)
+    - HTTP runner: X-Pugmark-Session-Metadata header (JSON-encoded)
+
+    Returns:
+        Dictionary of metadata key-value pairs, empty dict if none set
+
+    Example:
+        import pugmark
+
+        meta = pugmark.metadata()
+        user_id = meta.get("user-id")
+        tenant = meta.get("tenant")
+    """
+    return _default_session.metadata()
+
+
+@contextmanager
+def traced_execution() -> Iterator[None]:
+    """
+    Context manager that creates a tracing span for default session handler execution.
+
+    This is a convenience function that wraps the default session's traced_execution
+    method, creating a "pugmark.handler" span with session context attributes when
+    OpenTelemetry is available.
+
+    Example:
+        import pugmark
+
+        with pugmark.traced_execution():
+            for event in pugmark.events():
+                process(event)
+
+    Yields:
+        None
+    """
+    with _default_session.traced_execution():
+        yield