fix(sandbox): bound phase-one memory prompts to model context

[GautamSharma99]

Summary

This pull request fixes sandbox-memory phase-one requests that could exceed the selected model's context window even after the rollout was truncated.

Previously, every phase-one extraction allowed up to 150,000 estimated rollout tokens regardless of the configured model. That limit excluded the sizeable extraction instructions, terminal metadata, user-message wrapper, structured-output schema, provider framing, and model output.

The new budget resolution keeps one provider-neutral path:

• An explicit MemoryGenerateConfig.phase_one_model_context_window_tokens override takes precedence, which supports custom model objects and unknown provider model names.
• Otherwise, known string model names reuse CompactionModelInfo, the sandbox subsystem's existing context-window source of truth.
• Unknown models without an override preserve the released 150,000-token rollout fallback.
• The complete estimated phase-one input is limited to 70% of the resolved context window. Instructions, output schema, terminal metadata, and prompt-wrapper overhead are counted before the remaining budget is assigned to rollout content.
• The released 150,000-token rollout limit remains a maximum, so known large-context models do not unexpectedly receive larger prompts.
• Configured context windows that cannot fit the fixed phase-one request overhead fail before model invocation.

Phase one now uses the base Agent because it only transforms supplied rollout text into structured output and does not need sandbox tools or filesystem instructions. Phase two retains its existing sandbox-agent behavior.

The public field is optional and appended to MemoryGenerateConfig, preserving existing positional argument meaning. English documentation explains automatic lookup, the custom-model override, conservative token estimation, and fallback semantics.

Test plan

• Added regressions for known-model context lookup, explicit override precedence, unknown-model fallback, invalid non-positive configuration, fixed-overhead rejection, and complete-input budget enforcement.
• Added an end-to-end sandbox-memory regression that records the actual phase-one model request and verifies instructions, user input, and structured-output schema remain within 70% of the configured context.
• uv run pytest -q tests/sandbox/test_memory.py (74 passed)
• make build-docs
• bash .agents/skills/code-change-verification/scripts/run.sh
  • make format
  • make lint
  • make typecheck
  • make tests

Issue number

Closes #4058

Checks

• I've added new tests, if relevant
• I've run .agents/skills/code-change-verification/scripts/run.sh
• I've confirmed all verification steps pass
• If using Codex, I've run /review before submitting this PR

[seratch]

Before we can merge this, please replace the current approx_token_count() safety check. It computes ceil(UTF-8 bytes / 4), which is an average heuristic rather than a conservative upper bound. Token-dense JSON or code can therefore satisfy the new 70% assertion while the actual request still exceeds the model's input/context limit. The current end-to-end test verifies the same estimator used by the implementation, so it cannot catch that case.

Please use accounting that cannot undercount supported OpenAI input, add an independent token-dense regression that fails with the current implementation, and update the documentation to state the resulting guarantee accurately. The remaining model lookup, override, fallback, and compatibility shape can stay as proposed.

[GautamSharma99]

Thanks for the review @seratch — both points are addressed in the follow-up commit.

Non-undercounting token accounting. I replaced the ceil(UTF-8 bytes / 4) heuristic with a new openai_token_count() helper:

• When tiktoken is available it counts with the model's own encoding (falling back to o200k_base for unrecognized names), which is exact for supported OpenAI input and therefore cannot undercount.
• When tiktoken is not installed it returns a conservative UTF-8 byte-length upper bound (conservative_token_count()), which also cannot undercount because OpenAI's byte-level BPE never produces more tokens than bytes.

Both the fixed overhead (instructions + structured-output schema) and the per-iteration prompt check use this helper, so token-dense JSON/code can no longer slip past the 70% budget. Rollout content is truncated by bytes and each iteration re-measures the rendered prompt, converging via the prompt's measured byte-per-token density so dense content isn't discarded wholesale. The model lookup, override, fallback, and compatibility shape are unchanged, as you suggested.

Independent regression. test_render_phase_one_prompt_bounds_token_dense_input_by_real_tokens builds token-dense JSON and asserts the real token count of the rendered request (measured with tiktoken directly, not the implementation's estimator) stays within budget. It fails under the previous bytes/4 accounting — that estimator produced a prompt of ~7,363 real tokens against a 4,000-token budget — and passes now. I also added unit tests for the new counters, including the tiktoken-absent fallback.

Docs. I removed the docs/sandbox/memory.md change so it can land after the feature is released. The resulting guarantee is now stated accurately in the phase_one_model_context_window_tokens docstring instead.

make format, make lint, make typecheck, and make tests all pass.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 3abddc6e80

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[seratch]

Thanks for the follow-up. The token-dense regression fixes the original bytes/4 blind spot for recognized OpenAI models, but the new tiktoken path introduces two blockers.

phase_one_model also accepts Model objects and unknown/custom provider names; both are currently counted with o200k_base, although the Model interface provides no tokenizer contract. That makes the "never undercounts" guarantee false for the custom-model override this PR adds. Also, tiktoken is not a core dependency, and get_encoding() can synchronously download BPE data on a cache miss; an encoding load failure currently propagates instead of using the conservative fallback.

Please narrow this to deterministic UTF-8 byte-length upper-bound accounting inside phase one, remove the guessed/optional tiktoken path and its shared exports, and retain a dense-input regression for that bound. This preserves the provider-neutral contract without introducing environment-dependent I/O. After that, this should be ready for another review.

[GautamSharma99]

Thanks — addressed in 8e77c8b. I removed the tiktoken path entirely and narrowed phase one to deterministic UTF-8 byte-length upper-bound accounting.

Both Codex blockers + the tokenizer-contract concern. phase_one_model accepts custom Model objects and unknown provider names, and the Model interface provides no tokenizer contract, so counting those with o200k_base could undercount a denser tokenizer — the exact override case this PR adds. Byte length is a provider-neutral upper bound (a byte-level BPE tokenizer never emits more tokens than UTF-8 bytes), so it never undercounts any provider, with no tiktoken dependency and no synchronous get_encoding() BPE download. openai_token_count(), the tiktoken encoding helper, and their shared exports are gone; token_truncation.py and its __init__ are back to their released state, and the byte-length helper is private to phase_one.py.

150k rollout ceiling (Codex comment 2). The ceiling is now applied as a byte cap on the measured rollout (min(150_000, budget - overhead) in the same byte unit), so token-dense content can no longer retain well over 150k rollout tokens up to a larger context budget. Added test_render_phase_one_prompt_caps_rollout_at_ceiling_for_large_context to lock this in.

Tests. The dense-input regression and the end-to-end test now assert the deterministic byte-length bound instead of a tokenizer, keeping the check provider-neutral and free of environment-dependent I/O. The config docstring states the guarantee.

make format, make lint, make typecheck, and make tests all pass.

[seratch]

Thanks, the latest revision fixes the tiktoken and custom-model tokenizer blockers, and all checks are green. One behavior regression remains.

In v0.19.1, _PHASE_ONE_ROLLOUT_TOKEN_LIMIT = 150_000 is passed to TruncationPolicy.tokens(), corresponding to a 600,000-byte truncation budget. The new code reinterprets the same constant as 150,000 bytes, so large-context models retain only one quarter of the previous rollout evidence even when the context budget has ample room. The new large-context test locks in that reduction.

Please preserve the released 600,000-byte-equivalent rollout capacity while keeping the new whole-request byte bound as the safety check, and add a regression proving that intermediate-sized rollout content is retained when the context budget permits it. Also narrow the comments and docstring: UTF-8 bytes are an upper bound for byte-level BPE tokenizers, not an arbitrary custom Model contract. After those focused changes, this should be ready to approve.

[seratch]

After reviewing the issue, I came to think we don't need to make substantial changes for the purpose.