feat: preserve Context subtype and enforce input==output context type

docs/examples/context/custom_compactor.py

@@ -44,7 +44,7 @@ def compact(self, ctx, *, backend=None):
         items = ctx.as_list()
         if len(items) <= 1:
             return ctx
-        return _rebuild_chat_context(items[1:], compactor=ctx._compactor)
+        return _rebuild_chat_context(items[1:], source=ctx, compactor=ctx._compactor)
 
 
 def pattern_1_wired_into_context():

mellea/core/__init__.py

@@ -23,6 +23,7 @@
     ComputedModelOutputThunk,
     Context,
     ContextTurn,
+    ContextTypeMismatchError,
     GenerateLog,
     GenerateType,
     GenerationMetadata,
@@ -76,6 +77,7 @@ def __getattr__(name: str) -> object:
     "ComputedModelOutputThunk",
     "Context",
     "ContextTurn",
+    "ContextTypeMismatchError",
     "Formatter",
     "GenerateLog",
     "GenerateType",

mellea/core/base.py

@@ -1679,6 +1679,32 @@ class ContextTurn:
 ContextT = TypeVar("ContextT", bound="Context")
 
 
+class ContextTypeMismatchError(TypeError):
+    """Raised when a function returns a different `Context` subtype than it was given.
+
+    Mellea's convention is that the context type flowing out of a function equals
+    the context type flowing in (see issue #1522). This error enforces that
+    invariant. It is raised by the functional layer when a backend or sampling
+    strategy produces a context whose type differs from the input context's type,
+    unless the caller opted in to a deliberate type change via
+    `allow_context_type_change=True`.
+
+    Args:
+        input_type (type): The type of the context passed into the function.
+        output_type (type): The type of the context the function produced.
+    """
+
+    def __init__(self, input_type: type, output_type: type) -> None:
+        """Build the error message from the mismatched input and output context types."""
+        super().__init__(
+            f"Context type changed during generation: input was "
+            f"{input_type.__name__!r} but output is {output_type.__name__!r}. "
+            "Mellea functions must return the same Context subtype they were "
+            "given. If this change is deliberate (e.g. switching the context type "
+            "associated with a session), pass allow_context_type_change=True."
+        )
+
+
 class Context(abc.ABC):
     """A `Context` is used to track the state of a `MelleaSession`.
 
@@ -1691,12 +1717,23 @@ class Context(abc.ABC):
         node_data (Span | None): The data associated with this context node,
             or `None` for the root node.
         is_chat_context (bool): Whether this context operates in chat (multi-turn) mode.
+
+    Class Attributes:
+        _propagated_fields: Instance-attribute names copied from the source node
+            onto every node built by `from_previous()` (and, for `ChatContext`,
+            `_make_root()` / `_rebuild_chat_context()`). Because those factories
+            build via `cls.__new__(cls)` and never re-run `__init__`, a subclass
+            that stores state in its constructor must register the attribute here
+            or it is lost on the next `add()`. Empty on the base `Context`;
+            subclasses override with their own tuple (extend the parent's rather
+            than replacing it, e.g. `(*Context._propagated_fields, "my_field")`).
     """
 
     _previous: Context | None
     _data: Span | None
     _is_root: bool
     _is_chat_context: bool = True
+    _propagated_fields: tuple[str, ...] = ()
 
     def __init__(self) -> None:
         """Constructs a new root context with no content."""
@@ -1710,23 +1747,41 @@ def __init__(self) -> None:
     def from_previous(cls: type[ContextT], previous: Context, data: Span) -> ContextT:
         """Constructs a new context node linked to an existing context node.
 
+        The node is built with `cls.__new__(cls)` and its linked-list fields are
+        set directly, rather than by calling `cls()`. This deliberately does not
+        re-run the subclass `__init__`, so a subclass with required constructor
+        arguments (e.g. `def __init__(self, tag: str)`) still works when `add()`
+        reaches this factory — calling `cls()` would raise `TypeError`. Because
+        `__init__` is skipped, any subclass state that would otherwise be set
+        there must be registered in `_propagated_fields`; every such attribute is
+        copied from `previous` onto the new node here, so subclasses that carry
+        configuration keep it across `add()` without relying on `__init__`
+        re-running.
+
         Args:
             previous (Context): The existing context to extend.
             data (Span): The component, content block, or model output to associate with the new node.
 
         Returns:
             ContextT: A new context instance whose `previous_node` is `previous`.
+
+        Raises:
+            AssertionError: If `previous` is not a `Context`, or if `data` is `None`.
         """
         assert isinstance(previous, Context), (
             "Cannot create a new context from a non-Context object."
         )
         assert data is not None, "Cannot create a new context from None data."
 
-        x = cls()
+        x = cls.__new__(cls)
         x._previous = previous
         x._data = data
         x._is_root = False
         x._is_chat_context = previous._is_chat_context
+        # Skipping `__init__` (above) means subclass-owned state would be lost;
+        # copy every registered field from the source node so it survives.
+        for field_name in cls._propagated_fields:
+            setattr(x, field_name, getattr(previous, field_name))
         return x
 
     @classmethod
@@ -1746,8 +1801,16 @@ def new_instance(self) -> Context:
         configuration (e.g. `ChatContext` with `model_id` and `window_size`)
         should override this to propagate their config into the fresh instance.
 
+        The base signature returns `Context` (not `Self`) so that existing typed
+        third-party subclasses whose override is annotated `-> Context` — the
+        previous base contract, which the docstring invites — continue to satisfy
+        mypy's override check. Built-in contexts narrow the return to their own
+        type on their concrete overrides (e.g. `ChatContext.new_instance` returns
+        `ChatContext`), preserving subtype inference for callers.
+
         Returns:
-            Context: A freshly initialised root context of the same type.
+            Context: A freshly initialised root context of the same runtime type.
+            Concrete built-in subclasses narrow this to their own type.
         """
         return self.reset_to_new()
 
@@ -1880,11 +1943,20 @@ def last_turn(self) -> ContextTurn | None:
     def add(self, c: Span) -> Context:
         """Returns a new context obtained by appending `c` to this context.
 
+        The abstract signature returns the base `Context` so that existing typed
+        third-party subclasses whose override is annotated `-> Context` continue
+        to satisfy mypy's override check (changing this to `Self` would be a
+        breaking API change for them). Built-in contexts narrow the return to
+        `Self` on their concrete overrides, so `ChatContext.add(...)` statically
+        yields a `ChatContext` and a subclass yields its own type.
+
         Args:
             c (Span): The component, content block, or model output to add to the context.
 
         Returns:
-            Context: A new context node with `c` as its data and this context as its previous node.
+            Context: A new context node of the same runtime type with `c` as its
+            data and this context as its previous node. Concrete built-in
+            subclasses narrow this to `Self`.
         """
         # something along ....from_previous(self, c)
         ...

mellea/stdlib/context/chat.py

@@ -5,7 +5,7 @@
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Self, cast
 
 from mellea.backends.context_lengths import get_context_length
 from mellea.backends.model_ids import ModelIdentifier
@@ -67,12 +67,14 @@ class ChatContext(Context):
             and uses it as the token budget.
 
     Class Attributes:
-        _propagated_fields: Instance-attribute names copied by `add()` and
-            `_make_root()` into every descendant node. Add new `ChatContext`
-            fields here so they propagate automatically.
+        _propagated_fields: Instance-attribute names copied into every descendant
+            node built via `__new__` (by `Context.from_previous`, `_make_root()`,
+            and `_rebuild_chat_context()`). Extends the base `Context` tuple; add
+            new `ChatContext` fields here so they propagate automatically.
     """
 
     _propagated_fields: tuple[str, ...] = (
+        *Context._propagated_fields,
         "_compactor",
         "_token_context_length_limit",
         "_model_id",
@@ -117,8 +119,17 @@ def model_id(self) -> str | ModelIdentifier | None:
         return self._model_id
 
     def _make_root(self, model_id: str | ModelIdentifier | None) -> ChatContext:
-        """Return a new empty root `ChatContext`, propagating all `_propagated_fields` then binding `model_id`."""
-        new = ChatContext()
+        """Return a new empty root `ChatContext`, propagating all `_propagated_fields` then binding `model_id`.
+
+        Builds via `type(self).__new__(...)` (not `type(self)()`), so a subclass
+        gets an instance of itself back rather than being demoted to
+        `ChatContext`, and a subclass with required constructor arguments does not
+        raise `TypeError` here (the subclass `__init__` is deliberately not
+        re-run). Subclass configuration travels via `_propagated_fields`.
+        """
+        cls = type(self)
+        new = cls.__new__(cls)
+        Context.__init__(new)
         for field in self._propagated_fields:
             setattr(new, field, getattr(self, field))
         # Override whatever _propagated_fields copied for _model_id: the caller
@@ -153,36 +164,59 @@ def _bind_model(self, model_id: str | ModelIdentifier) -> ChatContext:
             )
         return self._make_root(model_id)
 
-    def new_instance(self) -> ChatContext:
+    def new_instance(self) -> Self:
         """Return a new empty root `ChatContext`, preserving compactor, token budget, and `model_id`.
 
         Use this instead of `reset_to_new()` when you need to preserve the
         configuration of an existing instance. `reset_to_new()` is a classmethod
         that returns a bare `ChatContext()` with no configuration.
 
+        Narrows the base `Context.new_instance()` return (`Context`) to `Self`, so
+        a subclass gets an instance of its own type back. `_make_root` builds via
+        `type(self).__new__(...)`, so the runtime guarantee matches this static
+        type.
+
         Returns:
-            ChatContext: A fresh root context with the same compactor,
-            `token_context_length_limit`, and `model_id` as this instance, but
-            no history.
+            Self: A fresh root context of the same concrete subtype with the same
+            compactor, `token_context_length_limit`, and `model_id` as this
+            instance, but no history.
         """
-        return self._make_root(self._model_id)
+        return cast("Self", self._make_root(self._model_id))
 
-    def add(self, c: Span) -> ChatContext:
+    def add(self, c: Span) -> Self:
         """Append `c` and run the compactor; return the resulting context.
 
         Args:
             c (Span): The component, content
                 block, or model output to append.
 
         Returns:
-            ChatContext: A new `ChatContext` carrying the same configuration.
+            Self: A new context of the same concrete subtype carrying the same
+            configuration. Returning `Self` (not the hard-coded `ChatContext`)
+            keeps a subclass statically its own type, matching the runtime
+            `type(self)` construction below.
         """
-        new = ChatContext.from_previous(self, c)
-        for field in self._propagated_fields:
-            setattr(new, field, getattr(self, field))
+        # `type(self)`, not `ChatContext`, so a subclass gets an instance of
+        # itself back rather than being silently demoted to `ChatContext`.
+        # Typed as `ChatContext` (not `Self`) because `compact()` below returns
+        # `ChatContext`; the final `cast` re-narrows to `Self` for the return.
+        # `from_previous` already copies `_propagated_fields` from `self` onto the
+        # new node, so no explicit copy is needed here.
+        new: ChatContext = type(self).from_previous(self, c)
         if self._compactor is not None:
             new = self._compactor.compact(new)
-        return new
+            # The built-in compactors rebuild via `type(ctx)`, so they preserve
+            # the concrete subtype. A *custom* `InlineCompactor` need not — its
+            # `compact()` may return a plain `ChatContext`, which would make the
+            # `Self` cast below unsound (issue #1522). If the returned type was
+            # demoted, rebuild the compacted history back into `type(self)`,
+            # re-copying `_propagated_fields` so subclass-owned state survives.
+            if type(new) is not type(self):
+                new = _rebuild_chat_context(new.as_list(), source=self)
+        # `new` is now guaranteed to be a `type(self)` instance (either the
+        # compactor preserved it, or the rebuild above restored it), so the cast
+        # informs the checker of what the runtime guarantees.
+        return cast("Self", new)
 
     def view_for_generation(self) -> list[Span] | None:
         """Return the components to forward to the model.
@@ -275,9 +309,11 @@ def _as_list_token_budget(self, token_budget: int) -> list[Span]:
 def _rebuild_chat_context(
     components: list[Span],
     *,
+    source: ChatContext,
     compactor: InlineCompactor | None = None,
     token_context_length_limit: int | None = None,
     model_id: str | ModelIdentifier | None = None,
+    cls: type[ChatContext] | None = None,
 ) -> ChatContext:
     """Build a fresh `ChatContext` linked-list without triggering compaction.
 
@@ -286,26 +322,72 @@ def _rebuild_chat_context(
     given the same configuration so the rebuilt context behaves identically to
     its source (e.g. token-budget views still apply).
 
+    Subclass state is preserved: every attribute named in the concrete class's
+    `_propagated_fields` is copied from `source` onto each rebuilt node — not
+    just the three built-in `ChatContext` fields — so a subclass that registers
+    its own field there keeps it across compaction rather than losing it (and
+    then raising `AttributeError` on the next `add`). The three built-in fields
+    can be overridden via the explicit `compactor` / `token_context_length_limit`
+    / `model_id` arguments; any that is left `None` falls back to `source`'s
+    value.
+
+    Note:
+        Nodes are constructed via `cls.__new__(cls)` and configured by copying
+        fields, so the subclass initializer is deliberately not re-run. A
+        subclass whose invariants live only in `__init__` (rather than in
+        `_propagated_fields`) will not have them re-established here; register
+        such state in `_propagated_fields` so it propagates.
+
+    Migration:
+        `source` is now a **required** keyword argument (it was absent before the
+        subclass-preservation change). Custom compactors that call this helper —
+        including the documented `custom_compactor` recipe — must pass
+        `source=ctx` so the rebuilt nodes inherit `type(ctx)` and its
+        `_propagated_fields`. It is required rather than optional because a
+        missing `source` would silently discard subclass identity and state (the
+        exact regression this argument fixes). The `compactor` /
+        `token_context_length_limit` / `model_id` arguments still default to
+        `None`; a `None` now inherits the corresponding value from `source`
+        rather than clearing the field.
+
     Args:
         components: Components to materialise as the new context, in order.
-        compactor: Compactor to attach to every node of the rebuilt context.
-        token_context_length_limit: Token budget to attach to every node.
-        model_id: Model identifier to attach to every node.
+        source: The context being rebuilt. Its `_propagated_fields` values are
+            copied onto every node so subclass-owned state survives.
+        compactor: Compactor to attach to every node; when `None`, `source`'s
+            compactor is used.
+        token_context_length_limit: Token budget to attach to every node; when
+            `None`, `source`'s value is used.
+        model_id: Model identifier to attach to every node; when `None`,
+            `source`'s value is used.
+        cls: The concrete `ChatContext` subtype to construct. Defaults to
+            `type(source)` so a subclassed context is rebuilt as its own type
+            rather than being demoted to `ChatContext`.
 
     Returns:
-        A new `ChatContext` whose linear history is exactly `components`.
+        A new context of type `cls` whose linear history is exactly `components`.
     """
+    target_cls = cls if cls is not None else type(source)
+    overrides = {
+        "_compactor": compactor,
+        "_token_context_length_limit": token_context_length_limit,
+        "_model_id": model_id,
+    }
 
     def _configure(node: ChatContext) -> None:
-        node._compactor = compactor
-        node._token_context_length_limit = token_context_length_limit
-        node._model_id = model_id
-
-    ctx: ChatContext = ChatContext.__new__(ChatContext)
+        # Copy every propagated field from the source so subclass-owned state
+        # survives compaction; explicit non-None overrides take precedence.
+        for field in source._propagated_fields:
+            setattr(node, field, getattr(source, field))
+        for field, value in overrides.items():
+            if value is not None:
+                setattr(node, field, value)
+
+    ctx: ChatContext = target_cls.__new__(target_cls)
     Context.__init__(ctx)
     _configure(ctx)
     for c in components:
-        new: ChatContext = ChatContext.__new__(ChatContext)
+        new: ChatContext = target_cls.__new__(target_cls)
         new._previous = ctx
         new._data = c
         new._is_root = False

mellea/stdlib/context/compactor.py

@@ -270,9 +270,11 @@ def compact(
         compacted = full[:pin_end] + keep_body
         return _rebuild_chat_context(
             compacted,
+            source=ctx,
             compactor=ctx._compactor,
             token_context_length_limit=ctx._token_context_length_limit,
             model_id=ctx._model_id,
+            cls=type(ctx),
         )
 
 
@@ -616,7 +618,9 @@ async def _async_compact(
         compacted = [*prefix, summary_message, *recent]
         return _rebuild_chat_context(
             compacted,
+            source=ctx,
             compactor=ctx._compactor,
             token_context_length_limit=ctx._token_context_length_limit,
             model_id=ctx._model_id,
+            cls=type(ctx),
         )