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

[AngeloDanducci]

Pull Request

Issue

Fixes #1522

Description

preserve Context subtype and enforce input==output context type

Testing

• Tests added to the respective file if code was changed
• New code has 100% coverage if code was added
• Ensure existing tests and github automation passes (a maintainer will kick off the github automation when the rest of the PR is populated)

Attribution

• AI coding assistants used

---

Adding a new component, requirement, sampling strategy, or tool?

If your PR adds or modifies one of the types below, check the matching box. A checklist of type-specific review items will be posted as a comment.

• Component
• Requirement
• Sampling Strategy
• Tool

NOTE: Please ensure you have an issue that has been acknowledged by a core contributor and routed you to open a pull request against this repository. Otherwise, please open an issue before continuing with this pull request.

[jakelorocco]

looks good; I have a few concerns / thoughts that I think would be worth addressing in this PR so that we can get a better surface for the context typing.

[jakelorocco]

Can you please add new tests for this as well?

[jakelorocco]

I looked at the implementation for our existing contexts and I think they will require fixes here as well. It looks like several of their methods use the named class constructor instead of self, etc... which will cause subclasses to fail.

[planetf1]

Rechecked the constructor paths at b4ca7bce. Using type(self) fixes bare-subclass demotion, but subclasses with required constructor arguments now fail: ChatContext._make_root() calls type(self)(), while both built-in add() paths reach Context.from_previous() → cls(). A subclass with __init__(self, tag: str) raises TypeError on its first add(); ChatContext also fails through new_instance() during reset/model binding.

Please either preserve subtype state without re-running subclass initialisers, including the Context.from_previous() path, or explicitly require subclasses to be constructible with no arguments. Add a required-argument-subclass regression test for the chosen contract.

[planetf1]

I found a few API, typing, and lifecycle issues inline. The new runtime subtype coverage and the checks over every sampling context are both useful additions.

[planetf1]

One lifecycle ordering issue: the context guard runs after component_post_success. When it rejects new_ctx, the exception handler then emits component_post_error, so plugins see both terminal events. I reproduced ['success', 'ContextTypeMismatchError']; the tracing plugin closes the action span as successful before the error hook runs. Could we validate the returned context(s) before constructing the success payload, so a rejected call emits only the error path?

[AngeloDanducci]

Yes, done.

[planetf1]

This escape hatch is runtime-correct but statically unsound. With allow_context_type_change=True, a ChatContext can become a SimpleContext, yet this cast and all public overloads still return ContextT. That also leaves MelleaSession[ChatContext].ctx typed as ChatContext after the switch, so session.ctx.model_id type-checks even though the runtime object has no such attribute. Could the Literal[True] overloads widen to Context, with sessions that permit switching represented as MelleaSession[Context]?

[AngeloDanducci]

Widened this, could use another set of eyes.

[planetf1]

Rechecked at b4ca7bce. The literal-True widening is in place for the existing functions, but three gaps remain:

1. achat has no overloads, so allow_context_type_change=True still returns tuple[Message, ChatContext] rather than widening the context. Please mirror chat and add a typing assertion.
2. MelleaSession(..., allow_context_type_change=True) remains MelleaSession[ChatContext], unlike start_session, which widens to MelleaSession[Context].
3. A runtime bool matches neither Literal overload. This leaves nine suppressed internal calls inferred as Any, and makes start_session(ctx=..., allow_context_type_change=flag) a public call-overload error. Add bool → Context fallback overloads for each applicable non-SamplingResult return shape and the ctx-supplied start_session forms.

[planetf1]

Two cases remain here, with different practical weight.

Required for this PR: the context_type and default-context overloads accept any bool but always return a narrowly typed session:

session = start_session(
    context_type="chat",
    allow_context_type_change=True,
)
reveal_type(session.ctx)  # ChatContext

The session is allowed to switch to another context subtype, so this should be MelleaSession[Context]. This is a normal public construction path and directly within the PR's typing scope. Could these overloads also distinguish Literal[False], Literal[True], and a runtime bool, with matching assert_type checks?

A separate low-likelihood, non-blocking edge is that allow_context_type_change remains writable after construction. Mypy accepts changing it to True on an existing MelleaSession[ChatContext], after which the runtime context can change subtype while session.ctx remains statically ChatContext. Nothing in-tree does this, but making the setting read-only would close that gap.

[planetf1]

Changing this public abstract method from Context to Self is a source-compatible runtime change, but it breaks typed third-party subclasses: an existing def add(...) -> Context now fails mypy's override check. Is that breaking API change intentional? If not, it may be safer to retain the abstract Context return and provide the narrower Self types on the concrete built-in contexts.

[AngeloDanducci]

should now narrow to self, the return while the abstract
signature stays Context for source-compat

[planetf1]

Preserving the subclass here bypasses its initializer and restores only the three built-in fields. That conflicts with _propagated_fields, whose class documentation invites subclasses to add state there. I reproduced a subclass-owned field disappearing after window compaction, then raising AttributeError on the next add(); a subclass with a required constructor argument already fails on the first add() through type(self)(). This needs a subclass-aware factory or clone/rebuild hook rather than cls.__new__().

[AngeloDanducci]

should now properly propagate

[planetf1]

The rebuild now preserves the concrete subtype, but the new compaction test only checks type identity. The original regression was a subclass-owned _propagated_fields value being discarded. Please extend it with a ChatContext subclass that registers an extra field, trigger compaction, and assert that field survives on the rebuilt context.

[planetf1]

The new runtime subclass checks are helpful. There is still a static half to the same contract: add() is annotated as returning ChatContext, and SimpleContext.add() returns SimpleContext. Mypy therefore widens a bare subclass at its first add(), despite the runtime object retaining its subtype. Could these use Self (and the compactor path preserve that type), with subclass assert_type checks added alongside the runtime tests?

[AngeloDanducci]

ChatContext.add and SimpleContext.add now return Self, and the
compactor path preserves it

[planetf1]

The built-in compactors preserve the concrete context type, but the InlineCompactor extension contract still permits a custom compactor to return a plain ChatContext:

class DemotingCompactor(InlineCompactor):
    def compact(self, ctx: ChatContext, *, backend=None) -> ChatContext:
        return ChatContext()

result = CustomContext(compactor=DemotingCompactor()).add(message)
reveal_type(result)  # CustomContext
print(type(result))  # ChatContext

I reproduced both results. This requires the less common combination of a custom ChatContext subtype and custom inline compactor, but both are supported extension points and add() now promises Self, so I think it is within this PR's scope. Could the compactor contract preserve its input subtype, or could add() validate the returned type before casting? A custom-compactor regression test would pin the guarantee.

[planetf1]

Separate source-compatibility issue on the sibling factory: Context.new_instance() now returns Self. A third-party override returning Context—the previous base contract, and one the docstring invites—now fails mypy’s override check.

Please keep the base return type as Context, matching add(). ChatContext.new_instance() can narrow its own return to Self, preserving subtype inference without breaking existing overrides.

[planetf1]

Please document the sampling-context validation scope. With return_sampling_results=False, only the chosen result_ctx is returned and checked; with True, every sample_contexts entry is returned and validated. That distinction is coherent, but it is surprising enough to deserve one sentence in the public API documentation.

[planetf1]

Was this compatibility break intentional? The documented custom-compactor recipe imports this underscore-private helper directly, so copied versions of the old call shape now fail because source is required; None for the configuration arguments also now inherits from source instead of clearing the field.

If intentional, please note the migration. If not, preserve the previous call shape with a source=None fallback.

[planetf1]

Re-review at b4ca7bce: the previous fixes hold. Requesting changes for the five remaining issues called out inline: the incomplete literal/runtime-bool overload contract, required-argument context subclasses, and new_instance() source compatibility. The sampling-scope and compactor-compatibility comments are non-blocking.

Required CI is green at this head.

[planetf1]

Re-review at a2f8a0e4: the five previously requested items are addressed. Requesting changes for two remaining in-scope type-safety gaps: start_session does not widen sessions created through context_type or the default context when switching is enabled, and a custom InlineCompactor can invalidate ChatContext.add() -> Self. The mutable session-flag case is low-likelihood and non-blocking; details inline.

Verified with full-repository mypy, Ruff, targeted context tests, and focused mypy/runtime reproducers.

[planetf1]

lgtm