Skip to content

(MOT-4014) fix: guarantee context assembly stays within budget - #495

Merged
ytallo merged 1 commit into
mainfrom
fix/context-assembly-budget
Jul 14, 2026
Merged

(MOT-4014) fix: guarantee context assembly stays within budget#495
ytallo merged 1 commit into
mainfrom
fix/context-assembly-budget

Conversation

@ytallo

@ytallo ytallo commented Jul 14, 2026

Copy link
Copy Markdown
Contributor

Problem

Context assembly could compact repeatedly and still return a successful response whose estimated token count exceeded the model's usable input budget. A multi-megabyte function result in the latest turn was protected by normal pruning and tail-selection rules, so it survived every compaction attempt.

The harness also added tool schemas, system aids, hook messages, response formats, and provider options after assembly. When context-manager was unavailable or returned an unusable result, the harness could fall back to raw history and call the provider without a complete final preflight.

Solution

  • Enforce token_count <= usable as a hard postcondition for every successful context::assemble response.
  • Add an emergency reduction pass that replaces oversized function results—even recent or normally protected ones—with bounded transcript references carrying call identity, size, estimated tokens, SHA-256, preview, and retrieval guidance.
  • Count complete tool schemas and caller-supplied request overhead throughout assembly.
  • Return stable context/overflow and harness/context_overflow failures when a safe request cannot be produced.
  • Resolve deterministic system text, invocation schemas, response format, provider options, and framing overhead before assembly.
  • Recount the final request after pre-generation hooks and call/result repair, before creating an assistant entry or invoking router::chat.
  • Remove the raw-history fallback and fail closed when context budgeting is unavailable.
  • Persist the correct applied.tokens_before compaction metric.

Impact

A 4.98 MB latest function result now becomes a bounded reference and the turn can continue without breaking call/result pairing. Irreducible requests fail locally with a structured context-overflow action and never reach the provider.

Validation

  • context-manager: full cargo test, including 84 BDD scenarios and 749 steps.
  • harness: full cargo test, including 220 unit tests plus manifest and schema coverage.
  • cargo clippy --all-targets -- -D warnings for both packages.
  • Updated and verified the context::assemble wire-schema golden.
  • git diff --check.

Fixes MOT-4014

Summary by CodeRabbit

  • New Features

    • Added support for tool schemas and provider request overhead in context budgeting.
    • Added token-counting support for model requests.
    • Oversized function results can now be replaced with bounded transcript references while preserving retrieval details.
  • Bug Fixes

    • Context assembly now guarantees results fit within the usable token budget.
    • Added clear context/overflow and harness/context_overflow errors when requests cannot safely fit.
    • Prevented over-budget requests from being sent after final request modifications.

@vercel

vercel Bot commented Jul 14, 2026

Copy link
Copy Markdown

The latest updates on your projects. Learn more about Vercel for GitHub.

Project Deployment Actions Updated (UTC)
workers Ready Ready Preview, Comment Jul 14, 2026 3:55pm
workers-tech-spec Ready Ready Preview, Comment Jul 14, 2026 3:55pm

Request Review

@coderabbitai

coderabbitai Bot commented Jul 14, 2026

Copy link
Copy Markdown

Review Change Stack

Caution

Review failed

The pull request is closed.

ℹ️ Recent review info
⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: ce08ccb7-ff5f-445e-a4f0-5df9214e0529

📥 Commits

Reviewing files that changed from the base of the PR and between 9f30a6c and 67f38e7.

📒 Files selected for processing (15)
  • context-manager/README.md
  • context-manager/src/core/prune.rs
  • context-manager/src/error.rs
  • context-manager/src/functions/assemble.rs
  • context-manager/tests/features/assemble.feature
  • context-manager/tests/features/engine_roundtrip.feature
  • context-manager/tests/golden/schemas/context.assemble.json
  • context-manager/tests/steps/call_steps.rs
  • context-manager/tests/steps/invariant_steps.rs
  • harness/src/clients/context.rs
  • harness/src/clients/mod.rs
  • harness/src/error.rs
  • harness/src/turn_loop.rs
  • tech-specs/2026-06-agentic/context-manager.md
  • tech-specs/2026-06-agentic/harness.md

📝 Walkthrough

Walkthrough

context::assemble now accounts for tools and request overhead, guarantees bounded successful output, applies emergency reduction to oversized function results, and returns stable overflow errors. The harness requires context services, recounts final requests after hooks, and prevents provider calls when budgets are exceeded.

Changes

Context budgeting and fail-closed generation

Layer / File(s) Summary
Budget contracts and API shapes
context-manager/src/error.rs, context-manager/src/functions/assemble.rs, context-manager/tests/golden/schemas/..., tech-specs/2026-06-agentic/..., context-manager/README.md
Defines tool and request-overhead inputs, token_count <= usable, emergency reduction semantics, and stable context overflow errors.
Emergency function-result reduction
context-manager/src/core/prune.rs
Replaces oversized function results with bounded deterministic references while preserving pairing metadata and tests the reduction behavior.
Assembly enforcement and feature coverage
context-manager/src/functions/assemble.rs, context-manager/tests/features/..., context-manager/tests/steps/...
Includes tools and overhead in counting, applies emergency reduction, fails irreducible requests with context/overflow, and updates feature coverage.
Harness context client contract
harness/src/clients/context.rs, harness/src/clients/mod.rs, harness/src/error.rs
Makes context assembly required, adds token counting and budget fields, and introduces harness/context_overflow.
Turn-loop final budget enforcement
harness/src/turn_loop.rs, tech-specs/2026-06-agentic/harness.md
Passes explicit assembly inputs, recounts the final request after hooks and repairs, and blocks generation when it exceeds the usable budget.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant TurnLoop
  participant ContextClient
  participant ContextManager
  participant LLMRouter
  TurnLoop->>ContextClient: assemble messages, tools, and request overhead
  ContextClient->>ContextManager: context::assemble
  ContextManager->>ContextManager: prune, compact, and emergency-reduce function results
  ContextManager-->>ContextClient: bounded context or context/overflow
  ContextClient-->>TurnLoop: assembled context and usable budget
  TurnLoop->>ContextClient: count final messages and tools after hooks
  ContextClient->>ContextManager: context::count-tokens
  ContextManager-->>TurnLoop: final token count
  TurnLoop->>LLMRouter: router::chat only when final request fits
Loading

Possibly related PRs

  • iii-hq/workers#477: Both changes modify the durable turn-loop flow in harness/src/turn_loop.rs.

Suggested reviewers: sergiofilhowz

Poem

I’m a rabbit guarding tokens tight,
Trimming tails that sprawl in flight.
Tools and overhead join the queue,
Overflow keeps unsafe calls from you.
A hashed hint, a bounded trail—
The model hops within the rail.

✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch fix/context-assembly-budget

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

Comment @coderabbitai help to get the list of available commands.

@github-actions

Copy link
Copy Markdown
Contributor

skill-check — worker

0 verified, 42 skipped (no docs/).

Layer Result
structure
vale
ai
render

Four for four. Nicely done.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant