Record additional context counts, sizes, and truncation metrics (#58144)

## What changed

Add `codex.additional_context` to count submitted entries, including unchanged entries, and `codex.additional_context.value_bytes` to measure UTF-8 value sizes before and after the existing truncation limit. Sizes exclude message wrappers and include the truncation marker after truncation.

Tag metrics with context kind, truncation status, and an explicitly supplied name when present; never use opaque source keys as tags. Record metrics through the current turn's telemetry.

## Testing

Add integration coverage for message and tool-output formats, repeated submissions, truncation boundaries, empty and multibyte values, optional name tags, and turn-specific model attribution.

GitOrigin-RevId: 463f257af27be58c372970fac853b85fde495dd5
This commit is contained in:
Alex Kira
2026-10-29 18:24:34 +00:00
committed by copyberry
parent f32b6cfd2b
commit 2ca71f1683
6 changed files with 281 additions and 5 deletions
@@ -7,7 +7,7 @@ use codex_utils_string::truncate_middle_with_token_budget;
use crate::ContextualUserFragment;
const MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS: usize = 1_000;
pub const MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS: usize = 1_000;
const MAX_DERIVED_TOOL_NAME_CHARS: usize = 128;
const ADDITIONAL_CONTEXT_TOOL_NAMESPACE: &str = "agent";
const ADDITIONAL_CONTEXT_END_MARKER_SUFFIX: &str = ">";
+1
View File
@@ -8,6 +8,7 @@ mod recap_prompt;
pub use additional_context::AdditionalContextDeveloperFragment;
pub use additional_context::AdditionalContextToolOutput;
pub use additional_context::AdditionalContextUserFragment;
pub use additional_context::MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS;
pub use annotated_content::AnnotatedContent;
pub use annotated_content::message_from_parts;
pub use annotated_content::set_annotated_content;
+11 -4
View File
@@ -441,7 +441,8 @@ async fn start_or_steer(
session
.maybe_emit_model_warnings_for_turn(turn_context.as_ref())
.await;
let mut task_input = merge_additional_context_input(session, additional_context).await;
let mut task_input =
merge_additional_context_input(session, &turn_context, additional_context).await;
if has_explicit_input {
task_input.push(
pending_turn_input(
@@ -591,7 +592,8 @@ async fn start_if_idle(
.maybe_emit_model_warnings_for_turn(turn_context.as_ref())
.await;
let mut task_input = merge_additional_context_input(session, additional_context).await;
let mut task_input =
merge_additional_context_input(session, &turn_context, additional_context).await;
let sender_messages = capture_sender_user_messages(session, &mut input).await;
match kind {
TurnStartKind::User => {
@@ -910,7 +912,9 @@ impl Session {
{
return Err(NotSubmittedReason::ActiveTurnOutputSchemaMismatch);
}
let mut pending_input = merge_additional_context_input(self, additional_context).await;
let mut pending_input =
merge_additional_context_input(self, &active_task.turn_context, additional_context)
.await;
active_task
.turn_context
@@ -953,6 +957,7 @@ impl Session {
async fn merge_additional_context_input(
session: &Session,
turn_context: &TurnContext,
additional_context: BTreeMap<String, AdditionalContextEntry>,
) -> Vec<TurnInput> {
let format = if session.enabled(Feature::UntrustedAsToolOutput) {
@@ -962,7 +967,9 @@ async fn merge_additional_context_input(
};
let additional_context_input = {
let mut state = session.state.lock().await;
state.additional_context.merge(additional_context, format)
state
.additional_context
.merge(additional_context, format, &turn_context.session_telemetry)
};
additional_context_input
.into_iter()
@@ -6,9 +6,15 @@ use crate::context::AdditionalContextDeveloperFragment;
use crate::context::AdditionalContextToolOutput;
use crate::context::AdditionalContextUserFragment;
use crate::context::ContextualUserFragment;
use codex_context_fragments::MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS;
use codex_otel::ADDITIONAL_CONTEXT_COUNT_METRIC;
use codex_otel::ADDITIONAL_CONTEXT_VALUE_BYTES_BUCKETS;
use codex_otel::ADDITIONAL_CONTEXT_VALUE_BYTES_METRIC;
use codex_otel::SessionTelemetry;
use codex_protocol::models::ResponseItem;
use codex_protocol::protocol::AdditionalContextEntry;
use codex_protocol::protocol::AdditionalContextKind;
use codex_utils_string::truncate_middle_with_token_budget;
pub(crate) enum AdditionalContextFormat {
Messages,
@@ -45,7 +51,11 @@ impl AdditionalContextStore {
&mut self,
values: BTreeMap<String, AdditionalContextEntry>,
format: AdditionalContextFormat,
telemetry: &SessionTelemetry,
) -> Vec<ResponseItem> {
for entry in values.values() {
Self::record_metrics(entry, telemetry);
}
let values =
values
.into_iter()
@@ -87,4 +97,32 @@ impl AdditionalContextStore {
self.values = values;
fragments
}
fn record_metrics(entry: &AdditionalContextEntry, telemetry: &SessionTelemetry) {
let (bounded_value, original_tokens) =
truncate_middle_with_token_budget(&entry.value, MAX_ADDITIONAL_CONTEXT_VALUE_TOKENS);
let kind = match entry.kind {
AdditionalContextKind::Untrusted => "untrusted",
AdditionalContextKind::Application => "application",
};
let truncated = original_tokens.is_some().to_string();
let mut tags = vec![("kind", kind), ("truncated", truncated.as_str())];
if let Some(name) = entry.name.as_deref() {
tags.push(("name", name));
}
telemetry.counter(ADDITIONAL_CONTEXT_COUNT_METRIC, /*inc*/ 1, &tags);
for (stage, bytes) in [
("before", entry.value.len()),
("after", bounded_value.len()),
] {
let mut size_tags = tags.clone();
size_tags.push(("stage", stage));
telemetry.histogram_with_boundaries(
ADDITIONAL_CONTEXT_VALUE_BYTES_METRIC,
i64::try_from(bytes).unwrap_or(i64::MAX),
ADDITIONAL_CONTEXT_VALUE_BYTES_BUCKETS,
&size_tags,
);
}
}
}
@@ -1,12 +1,17 @@
use anyhow::Result;
use codex_core::TurnInputRequest;
use codex_features::Feature;
use codex_otel::ADDITIONAL_CONTEXT_COUNT_METRIC;
use codex_otel::ADDITIONAL_CONTEXT_VALUE_BYTES_METRIC;
use codex_otel::MetricsClient;
use codex_otel::MetricsConfig;
use codex_protocol::error::CodexErrorDetails;
use codex_protocol::items::TurnItem;
use codex_protocol::protocol::AdditionalContextEntry;
use codex_protocol::protocol::AdditionalContextKind;
use codex_protocol::protocol::EventMsg;
use codex_protocol::protocol::ItemCompletedEvent;
use codex_protocol::protocol::ThreadSettingsOverrides;
use codex_protocol::user_input::UserInput;
use core_test_support::context_snapshot;
use core_test_support::context_snapshot::ContextSnapshotOptions;
@@ -18,6 +23,10 @@ use core_test_support::responses::start_mock_server;
use core_test_support::skip_if_no_network;
use core_test_support::test_codex::test_codex;
use core_test_support::wait_for_event_match;
use opentelemetry_sdk::metrics::InMemoryMetricExporter;
use opentelemetry_sdk::metrics::data::AggregatedMetrics;
use opentelemetry_sdk::metrics::data::MetricData;
use opentelemetry_sdk::metrics::data::ScopeMetrics;
use pretty_assertions::assert_eq;
use serde_json::json;
use std::collections::BTreeMap;
@@ -705,6 +714,205 @@ async fn additional_context_values_are_truncated_before_model_input(
Ok(())
}
#[test_case::test_case(false; "messages")]
#[test_case::test_case(true; "tool_outputs")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn additional_context_metrics_record_submitted_sizes_and_truncation(
tool_outputs: bool,
) -> Result<()> {
skip_if_no_network!(Ok(()));
let server = start_mock_server().await;
let metrics = MetricsClient::new(
MetricsConfig::in_memory(
"test",
"additional-context",
env!("CARGO_PKG_VERSION"),
InMemoryMetricExporter::default(),
)
.with_runtime_reader(),
)?;
let test = test_codex()
.with_model("gpt-6-sol")
.with_config(move |config| {
config
.features
.set_enabled(Feature::UntrustedAsToolOutput, tool_outputs)
.expect("test config should allow feature update");
})
.build_with_auto_env(&server)
.await?;
let mut options = test.start_thread_options().await;
options.thread_extension_init.insert(metrics.clone());
let thread = test.thread_manager.start_thread(options).await?.thread;
let application = AdditionalContextKind::Application;
let untrusted = AdditionalContextKind::Untrusted;
let cases = [
(
"application_limit",
application,
None,
"a".repeat(4_000),
false,
),
(
"application_over",
application,
None,
"a".repeat(4_001),
true,
),
(
"opaque_named_limit",
untrusted,
Some("browser_limit"),
"b".repeat(4_000),
false,
),
(
"opaque_named_over",
untrusted,
Some("browser_over"),
"b".repeat(4_001),
true,
),
// Size is measured in UTF-8 bytes; the unnamed entry must not acquire a key-derived label.
("opaque_unnamed", untrusted, None, "é".repeat(2_001), true),
("opaque_empty", untrusted, None, String::new(), false),
];
let values = cases
.iter()
.map(|(key, kind, name, value, _)| {
(
key.to_string(),
AdditionalContextEntry {
name: name.map(str::to_owned),
kind: *kind,
value: value.clone(),
},
)
})
.collect::<BTreeMap<_, _>>();
for turn in 0..2 {
let request = mount_sse_once(&server, sse(vec![ev_completed("context-metrics")])).await;
thread
.start_or_steer_turn(
TurnInputRequest::user_input(vec![UserInput::Text {
text: format!("inspect context {turn}"),
text_elements: Vec::new(),
}])
.with_additional_context(values.clone())
.with_thread_settings(ThreadSettingsOverrides {
model: Some("gpt-5.5".to_string()),
..Default::default()
}),
)
.await?;
wait_for_event_match(&thread, |event| {
matches!(event, EventMsg::TurnComplete(_)).then_some(())
})
.await;
let request = request.single_request();
assert_eq!(request.body_json()["model"], "gpt-5.5");
// Repeated submissions are measured even though retained context is deduplicated.
assert_eq!(
request.inputs_of_type("function_call_output").len(),
if tool_outputs { 4 } else { 0 }
);
for (key, kind, _, _, _) in &cases {
let role = match kind {
AdditionalContextKind::Application => "developer",
AdditionalContextKind::Untrusted if tool_outputs => continue,
AdditionalContextKind::Untrusted => "user",
};
assert_eq!(
request
.message_input_texts(role)
.iter()
.filter(|text| text.contains(key))
.count(),
1,
);
}
}
let context_attribute = |attribute: &opentelemetry::KeyValue| {
(!matches!(
attribute.key.as_str(),
"auth_mode" | "session_source" | "originator" | "service_name" | "app.version"
))
.then(|| (attribute.key.to_string(), attribute.value.to_string()))
};
let snapshot = metrics.snapshot()?;
let mut counts = BTreeMap::new();
let mut sizes = BTreeMap::new();
for metric in snapshot.scope_metrics().flat_map(ScopeMetrics::metrics) {
match metric.name() {
ADDITIONAL_CONTEXT_COUNT_METRIC => {
let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else {
panic!("additional context count must be a counter");
};
for point in sum.data_points() {
counts.insert(
point
.attributes()
.filter_map(context_attribute)
.collect::<BTreeMap<_, _>>(),
point.value(),
);
}
}
ADDITIONAL_CONTEXT_VALUE_BYTES_METRIC => {
let AggregatedMetrics::F64(MetricData::Histogram(histogram)) = metric.data() else {
panic!("additional context size must be a histogram");
};
for point in histogram.data_points() {
sizes.insert(
point
.attributes()
.filter_map(context_attribute)
.collect::<BTreeMap<_, _>>(),
(point.count(), point.sum()),
);
}
}
_ => {}
}
}
let mut expected_counts = BTreeMap::new();
let mut expected_sizes = BTreeMap::new();
for (_, kind, name, value, truncated) in cases {
let kind = match kind {
AdditionalContextKind::Application => "application",
AdditionalContextKind::Untrusted => "untrusted",
};
let mut tags = BTreeMap::from([
("kind".to_string(), kind.to_string()),
("model".to_string(), "gpt-5.5".to_string()),
("truncated".to_string(), truncated.to_string()),
]);
if let Some(name) = name {
tags.insert("name".to_string(), name.to_string());
}
expected_counts.insert(tags.clone(), 2);
// The marker is outside the retained-text budget, so a truncated value can grow.
let after_bytes = if truncated {
4_000 + "…1 tokens truncated…".len()
} else {
value.len()
};
for (stage, bytes) in [("before", value.len()), ("after", after_bytes)] {
let mut tags = tags.clone();
tags.insert("stage".to_string(), stage.to_string());
expected_sizes.insert(tags, (2, 2.0 * bytes as f64));
}
}
assert_eq!(counts, expected_counts);
assert_eq!(sizes, expected_sizes);
metrics.shutdown()?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn additional_context_rejects_invalid_names_before_admitting_input() -> Result<()> {
skip_if_no_network!(Ok(()));
+22
View File
@@ -70,6 +70,28 @@ pub const STARTUP_PREWARM_DURATION_METRIC: &str = "codex.startup_prewarm.duratio
pub const STARTUP_PREWARM_AGE_AT_FIRST_TURN_METRIC: &str =
"codex.startup_prewarm.age_at_first_turn_ms";
pub const THREAD_STARTED_METRIC: &str = "codex.thread.started";
/// Submitted additional-context entries, including unchanged entries, by kind and truncation status.
/// The name tag is included only when explicitly supplied; opaque source keys are never tags.
pub const ADDITIONAL_CONTEXT_COUNT_METRIC: &str = "codex.additional_context";
/// UTF-8 value bytes before/after truncation, excluding message wrappers; after includes the marker.
pub const ADDITIONAL_CONTEXT_VALUE_BYTES_METRIC: &str = "codex.additional_context.value_bytes";
pub const ADDITIONAL_CONTEXT_VALUE_BYTES_BUCKETS: &[f64] = &[
0.0,
256.0,
512.0,
1_000.0,
2_000.0,
4_000.0,
8_000.0,
16_000.0,
32_000.0,
64_000.0,
128_000.0,
256_000.0,
512_000.0,
1_024_000.0,
4_096_000.0,
];
pub const THREAD_SKILLS_ENABLED_TOTAL_METRIC: &str = "codex.thread.skills.enabled_total";
pub const THREAD_SKILLS_KEPT_TOTAL_METRIC: &str = "codex.thread.skills.kept_total";
pub const THREAD_SKILLS_DESCRIPTION_TRUNCATED_CHARS_METRIC: &str =