[Proposal] Extend context-editing spill protection to subagents — lossy admission truncation is their only context defense today

[subaoyan16]

Summary

The orchestrator's middleware stack has a lossless context-defense layer that subagents don't get. Long-running specialists (deep web_crawler chains, multi-page google_search, etc.) currently fall back to character-level truncation in admit_langchain_messages, which is lossy, can split AIMessage.tool_calls / ToolMessage pairs, and leaves nothing recoverable. I'd like to propose a small (~200 line) PR that gives every subagent the same spill-based context editing the main agent already has, and I'm volunteering to implement it.

Current behavior

Main agent (build in surfsense_backend/app/agents/chat/multi_agent_chat/main_agent/middleware/stack.py):

• build_context_editing_mw(...) installs SpillingContextEditingMiddleware (main_agent/middleware/context_editing/middleware.py). When context exceeds 0.55 × max_input_tokens (see context_editing/builder.py), older ToolMessage contents are persisted to the tool_output_spills table first and replaced with a spill_<uuid> placeholder; the main agent can read them back via read_run / search_run.
• SurfSenseCompactionMiddleware (app/agents/chat/shared/middleware/compaction.py) additionally summarizes older history.

Subagents (compiled in main_agent/middleware/checkpointed_subagent_middleware/middleware.py::_compile_one from the spec built in subagents/shared/subagent_builder.py::pack_subagent):

• Middleware list is: todos, citation, retry, fallback, model_call_limit, tool_call_limit, per-subagent permission, PatchToolCallsMiddleware. No context editing, no compaction.
• Subagents share the same ChatLiteLLM instance, so every subagent model call still passes through ChatLiteLLM._admit_messages → admit_langchain_messages (app/services/context_admission.py). Admission is their only context defense.

Failure modes (observed consequences of the asymmetry)

1. Lossy, unrecoverable truncation. When a long subagent run crosses the admission budget, tool outputs are truncated character-by-character (trim_messages_to_fit_context, priority order tool/assistant → user with <document>). Unlike main-agent spills, the original content is gone forever.
2. Pair-splitting. Admission truncates message content independently, so a ToolMessage can be cut while its originating AIMessage.tool_calls args remain intact — or vice versa. Some providers reject these malformed sequences; in the extreme, preserve_protected_content=True leads to ContextOverflowError after emptying every removable message.
3. (Out of scope, listed for context) Subagent checkpoints ({parent_thread}::task:{tool_call_id}) grow without any cleanup policy.

Concrete scenario: a web_crawler subagent crawling a long article chain accumulates several large ToolMessages mid-run. Today: admission hard-truncates them, the subagent continues with corrupted context, and the parent receives a final answer built on truncated data with no way to inspect what was lost.

Proposed change

Reuse the existing, battle-tested factory for subagents — no new abstractions:

1. Thread max_input_tokens and workspace_id through subagent_dependencies in main_agent/middleware/stack.py (build_main_agent_deepagent_middleware).
2. In subagents/shared/subagent_builder.py::pack_subagent, build one SpillingContextEditingMiddleware instance per subagent via the existing build_context_editing_mw(flags, max_input_tokens, tools=<subagent's tools>, workspace_id) and insert it after tool_call_limit (mirroring the main-agent stack, where context editing precedes permission). The factory already returns None when the flag is off or max_input_tokens is missing, so behavior stays identical by default.

Sketch:

# pack_subagent(), middleware assembly section
context_editing = build_context_editing_mw(
    flags=dependencies["flags"],
    max_input_tokens=dependencies.get("max_input_tokens"),
    tools=tools,
    workspace_id=dependencies.get("workspace_id"),
)
if context_editing is not None:
    prepended.append(context_editing)

# stack.py, subagent_dependencies construction
"max_input_tokens": max_input_tokens,
"workspace_id": workspace_id,

Spill rows land under the subagent's thread_id ({parent}::task:{tool_call_id}, set in subagents/shared/invocation.py::subagent_invoke_config), so per-invocation isolation and deterministic uuid5 spill ids come for free.

Why a per-subagent instance (not the shared stack)

build_subagent_middleware_stack returns one dict shared by all subagents, and batch task(tasks=[...]) fan-out runs several subagents concurrently. SpillingContextEditingMiddleware buffers pending_spills per instance (SpillToBackendEdit.pending_spills, drained in awrap_model_call), so a single shared instance would let concurrent runs drain each other's rows. Building it inside pack_subagent (one call per subagent spec) avoids the problem entirely.

Gating & scope

• Gated by the existing enable_context_editing flag — off by default keeps current behavior unchanged; on, subagents become symmetric with the main agent.
• In scope: spill-based context editing for subagents + tests.
• Out of scope (possible follow-ups, happy to discuss): subagent compaction (summarization offload path conflicts with the shared StateBackend), subagent checkpoint pruning, unifying the three budget thresholds (compaction / context editing / admission) under resolve_max_input_tokens.

Test plan

1. Assembly matrix: flag × max_input_tokens (4 combinations) → middleware present/absent as expected.
2. Per-instance isolation: two pack_subagent calls yield distinct SpillingContextEditingMiddleware instances.
3. Trigger behavior: oversized tool-output message sequence → rows in tool_output_spills with *::task:* thread ids, placeholders substituted, protected tool names (prune_tool_names.py) exempt.
4. Existing main-agent tests must stay green (main path untouched).

Questions

1. Would you be interested in accepting a PR along these lines?
2. Is anything already planned for subagent context management (public board / plans/ may not show internal tracks)?
3. Preference on gating: reuse enable_context_editing, or introduce a separate subagent-scoped flag?

Thanks for the great architecture — happy to adjust the proposal to fit your roadmap.