Support per-tool MCP output limits (#41421)

## What changed

- Add a positive `output_token_limit` setting to each entry under an MCP server's `tools` configuration.
- Apply the most restrictive limit when plugin and user policies overlap, while keeping approval policy independent.
- Carry the effective MCP output budget in conversation history so tool output, post-tool hook responses, and resumed sessions use the same truncation limit.

## Testing

- Cover configuration parsing, serialization, schema validation, and plugin policy merging.
- Cover MCP output below and above the configured limit, post-tool hook responses, and session resume.

GitOrigin-RevId: d0beb4fca9ba6055d9e1d31c137373b465d50d61
This commit is contained in:
pakrym-oai authored and copyberry committed 2026-08-28 21:40:23 +00:00
1 parent 03147407e3
commit f742dabc6f
38 files changed
+487 -224

No files matched your search

+2 -3
View File
@@ -72,6 +72,7 @@ impl Session {
&& matches!(&item, ResponseItem::Message { role, .. } if role == "developer"))
.then_some(CodexHarnessMetadata {
client_authored: true,
..Default::default()
});
ResponseItemEnvelope { item, metadata }
@@ -82,9 +83,7 @@ impl Session {
turn_context: &TurnContext,
items: Vec<ResponseItemEnvelope>,
) {
if !self.enabled(Feature::RetainClientDeveloperMessages)
|| items.iter().all(|item| item.metadata.is_none())
{
if items.iter().all(|item| item.metadata.is_none()) {
let items = items
.into_iter()
.map(ResponseItemEnvelope::into_item)
+1
View File
@@ -401,6 +401,7 @@ mod tests {
item: ResponseItem::Other,
metadata: Some(CodexHarnessMetadata {
client_authored: true,
..Default::default()
}),
});
assert!(serde_json::to_value(annotated).is_err());
+7 -10
View File
@@ -40,6 +40,7 @@ use crate::session::step_context::StepContext;
use crate::session::turn_context::TurnContext;
use crate::skills::emit_explicit_skill_invocations;
use crate::stream_events_utils::HandleOutputCtx;
use crate::stream_events_utils::InFlightFuture;
use crate::stream_events_utils::TurnItemContributorPolicy;
use crate::stream_events_utils::finalize_non_tool_response_item;
use crate::stream_events_utils::handle_non_tool_response_item;
@@ -90,7 +91,6 @@ use codex_protocol::items::build_hook_prompt_message;
use codex_protocol::models::BaseInstructions;
use codex_protocol::models::ContentItem;
use codex_protocol::models::MessagePhase;
use codex_protocol::models::ResponseInputItem;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::AgentMessageContentDeltaEvent;
use codex_protocol::protocol::AgentReasoningSectionBreakEvent;
@@ -124,7 +124,6 @@ use codex_utils_stream_parser::AssistantTextStreamParser;
use codex_utils_stream_parser::ProposedPlanSegment;
use codex_utils_stream_parser::extract_proposed_plan_text;
use codex_utils_stream_parser::strip_citations;
use futures::future::BoxFuture;
use futures::prelude::*;
use futures::stream::FuturesOrdered;
use tokio_util::sync::CancellationToken;
@@ -2154,22 +2153,21 @@ async fn handle_assistant_item_done_in_plan_mode(
#[instrument(level = "trace", skip_all)]
async fn drain_in_flight(
in_flight: &mut FuturesOrdered<BoxFuture<'static, CodexResult<ResponseInputItem>>>,
in_flight: &mut FuturesOrdered<InFlightFuture<'static>>,
sess: Arc<Session>,
turn_context: Arc<TurnContext>,
) -> CodexResult<()> {
while let Some(res) = in_flight.next().await {
match res {
Ok(response_input) => {
let response_item = response_input.into();
sess.record_conversation_items(&turn_context, std::slice::from_ref(&response_item))
.await;
Ok(envelope) => {
mark_thread_memory_mode_polluted_if_external_context(
sess.as_ref(),
turn_context.as_ref(),
&response_item,
&envelope.item,
)
.await;
sess.record_annotated_conversation_items(&turn_context, vec![envelope])
.await;
}
Err(err) => {
error_or_panic(format!("in-flight tool future failed during drain: {err}"));
@@ -2247,8 +2245,7 @@ async fn try_run_sampling_request(
.instrument(trace_span!("stream_request"))
.or_cancel(&cancellation_token)
.await??;
let mut in_flight: FuturesOrdered<BoxFuture<'static, CodexResult<ResponseInputItem>>> =
FuturesOrdered::new();
let mut in_flight: FuturesOrdered<InFlightFuture<'static>> = FuturesOrdered::new();
let mut needs_follow_up = false;
let mut last_agent_message: Option<String> = None;
let mut active_item: Option<TurnItem> = None;