Inject history notes hints into context windows (#40539)

## What changed

- Register the history-notes extension as a context contributor and fetch `thread_hint` from the configured history-notes backend.
- Add non-empty hints of at most 4,000 bytes to the context-window prompt, and omit hints when the request fails or the response is missing, empty, or oversized.
- Verify that app-server sessions use the native backend when `use_history_notes_extension` is enabled and retain the MCP bridge behavior when it is disabled.

## Testing

- Cover hint injection, request context, response validation, backend failures, and disabled-extension behavior in history-notes and app-server tests.

GitOrigin-RevId: 344a7bb79d962ea8b7c969f25b67a79423fc7fcc
This commit is contained in:
pmccrary-oai
2026-08-25 04:05:38 +00:00
committed by copyberry
parent 3a469a297d
commit 2e4675919e
3 changed files with 249 additions and 20 deletions
@@ -12,10 +12,13 @@ use codex_config::types::AuthCredentialsStoreMode;
use core_test_support::load_default_config_for_test;
use core_test_support::responses;
use pretty_assertions::assert_eq;
use serde_json::Value;
use serde_json::json;
use tempfile::TempDir;
use test_case::test_case;
use tokio::time::timeout;
use wiremock::Mock;
use wiremock::Request;
use wiremock::ResponseTemplate;
use wiremock::matchers::method;
use wiremock::matchers::path;
@@ -25,8 +28,20 @@ use super::analytics::wait_for_matching_analytics_event;
const DEFAULT_READ_TIMEOUT: Duration = Duration::from_secs(30);
const THREAD_HINT: &str =
"Recent notes (up to 5, most-recent first):\n- /root/notes/latest.md (2 lines, 14 UTF-8 bytes)";
const BRIDGE_HINT: &str = "unstructured notes/thread_hint fixture result";
#[test_case(true, 200, THREAD_HINT; "native_hint")]
#[test_case(true, 200, ""; "no_notes")]
#[test_case(true, 503, THREAD_HINT; "native_failure_does_not_use_bridge")]
#[test_case(false, 200, THREAD_HINT; "bridge_hint")]
#[tokio::test]
async fn app_server_registers_history_and_notes_tools_for_token_budget_threads() -> Result<()> {
async fn app_server_uses_configured_notes_backend_for_context_window_hints(
use_history_notes_extension: bool,
hint_status: u16,
hint_text: &str,
) -> Result<()> {
let server = responses::start_mock_server().await;
let response_mock = responses::mount_sse_once(
&server,
@@ -37,6 +52,66 @@ async fn app_server_registers_history_and_notes_tools_for_token_budget_threads()
]),
)
.await;
// HTTP MCP stays on the app-server host, including with a remote executor.
Mock::given(method("POST"))
.and(path("/mcp"))
.respond_with(|request: &Request| {
let body: Value = request.body_json().expect("MCP JSON-RPC request");
let Some(id) = body.get("id") else {
return ResponseTemplate::new(202);
};
let result = match body["method"].as_str() {
Some("initialize") => json!({
"protocolVersion": body["params"]["protocolVersion"],
"capabilities": {"tools": {}},
"serverInfo": {"name": "notes-fixture", "version": "1.0.0"},
}),
Some("tools/list") => json!({
"tools": [{
"name": "thread_hint",
"inputSchema": {"type": "object", "properties": {}},
"annotations": {"readOnlyHint": true},
"_meta": {"ui": {"visibility": []}},
}],
}),
Some("tools/call") => {
assert_eq!(body["params"]["name"], "thread_hint");
let thread_id = body["params"]["_meta"]["threadId"]
.as_str()
.expect("threadId metadata");
json!({"content": [
{"type": "text", "text": format!("manual history hint for thread {thread_id}")},
{"type": "text", "text": BRIDGE_HINT},
]})
}
_ => {
return ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0",
"id": id,
"error": {"code": -32601, "message": "Method not found"},
}));
}
};
ResponseTemplate::new(200).set_body_json(json!({
"jsonrpc": "2.0", "id": id, "result": result,
}))
})
.mount(&server)
.await;
Mock::given(method("GET"))
.and(path("/mcp"))
.respond_with(ResponseTemplate::new(405))
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/backend-api/codex/alpha/notes/v2/thread_hint"))
.respond_with(
ResponseTemplate::new(hint_status).set_body_json(serde_json::json!({
"text": hint_text
})),
)
.mount(&server)
.await;
let codex_home = TempDir::new()?;
MockResponsesConfig::new(&server.uri())
@@ -44,9 +119,10 @@ async fn app_server_registers_history_and_notes_tools_for_token_budget_threads()
.with_provider_name("OpenAI")
.with_provider_base_url(&format!("{}/backend-api/codex", server.uri()))
.with_provider_config("supports_websockets = false\nrequires_openai_auth = true")
.with_extra_config(
"[features.token_budget]\nenabled = true\nuse_history_notes_extension = true",
)
.with_extra_config(&format!(
"[features.token_budget]\nenabled = true\nuse_history_notes_extension = {use_history_notes_extension}\n\n[mcp_servers.notes]\nurl = \"{}/mcp\"\nstartup_timeout_sec = 10\n",
server.uri(),
))
.write(codex_home.path())?;
write_chatgpt_auth(
codex_home.path(),
@@ -77,22 +153,69 @@ async fn app_server_registers_history_and_notes_tools_for_token_budget_threads()
.await??;
let request = response_mock.single_request();
for (namespace, tool_name) in [
("history", "list_windows"),
("history", "list_items"),
("history", "read_item"),
("history", "search_contents"),
("notes", "list_files_by_prefix"),
("notes", "read_file"),
("notes", "search_contents"),
("notes", "append_to_file"),
("notes", "write_file"),
] {
assert!(
request.tool_by_name(namespace, tool_name).is_some(),
"app-server should expose {namespace}.{tool_name} to the model"
);
if use_history_notes_extension {
for (namespace, tool_name) in [
("history", "list_windows"),
("history", "list_items"),
("history", "read_item"),
("history", "search_contents"),
("notes", "list_files_by_prefix"),
("notes", "read_file"),
("notes", "search_contents"),
("notes", "append_to_file"),
("notes", "write_file"),
] {
assert!(
request.tool_by_name(namespace, tool_name).is_some(),
"app-server should expose {namespace}.{tool_name} to the model"
);
}
}
assert!(request.tool_by_name("notes", "thread_hint").is_none());
let input = request.input();
let bridge_hint_present = request
.message_input_texts("developer")
.iter()
.any(|text| text.contains(BRIDGE_HINT));
assert_eq!(bridge_hint_present, !use_history_notes_extension);
let requests = server.received_requests().await.expect("recorded requests");
let native_requests = requests
.iter()
.filter(|request| request.url.path() == "/backend-api/codex/alpha/notes/v2/thread_hint")
.count();
assert_eq!(native_requests, usize::from(use_history_notes_extension));
let bridge_calls = requests
.iter()
.filter(|request| request.url.path() == "/mcp")
.filter(|request| {
request
.body_json::<Value>()
.is_ok_and(|body| body["method"] == "tools/call")
})
.count();
assert_eq!(bridge_calls, usize::from(!use_history_notes_extension));
let developer_messages = request.message_input_texts("developer");
let context_window = developer_messages
.iter()
.find(|text| text.contains("<context_window>"))
.expect("context-window developer message");
let window_body = context_window
.split_once("<context_window>")
.expect("opening tag")
.1
.split_once("</context_window>")
.expect("closing tag")
.0;
assert_eq!(
window_body.contains(THREAD_HINT),
use_history_notes_extension && hint_status == 200 && !hint_text.is_empty(),
);
assert!(!input.iter().any(|item| {
item["type"] == "function_call"
&& item["namespace"] == "notes"
&& item["name"] == "thread_hint"
}));
Ok(())
}
@@ -2,9 +2,13 @@ use std::sync::Arc;
use codex_core::config::Config;
use codex_extension_api::ConfigContributor;
use codex_extension_api::ContentItemKind;
use codex_extension_api::ContextContributor;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionFuture;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::PromptFragment;
use codex_extension_api::PromptSlot;
use codex_extension_api::ThreadLifecycleContributor;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::ToolCall;
@@ -13,11 +17,15 @@ use codex_extension_api::ToolExecutor;
use codex_login::AuthManager;
use codex_model_provider::create_model_provider;
use codex_protocol::AgentPath;
use serde_json::json;
use crate::backend::HistoryNotesBackend;
use crate::tools::HistoryNotesAction;
use crate::tools::HistoryNotesTool;
// Bound model context even when note paths are unusually long.
const MAX_THREAD_HINT_BYTES: usize = 4_000;
struct HistoryNotesExtension {
auth_manager: Arc<AuthManager>,
}
@@ -82,6 +90,46 @@ impl ConfigContributor<Config> for HistoryNotesExtension {
}
}
impl ContextContributor for HistoryNotesExtension {
fn contribute_thread_context<'a>(
&'a self,
session_store: &'a ExtensionData,
thread_store: &'a ExtensionData,
) -> ExtensionFuture<'a, Vec<PromptFragment>> {
Box::pin(async move {
let Some(config) = thread_store.get::<HistoryNotesExtensionConfig>() else {
return Vec::new();
};
let Some(identity) = thread_store.get::<HistoryNotesAgentIdentity>() else {
return Vec::new();
};
let Ok(result) = config
.backend
.call(
"alpha/notes/v2/thread_hint",
session_store.level_id(),
&identity.agent_name,
json!({}),
)
.await
else {
return Vec::new();
};
let Some(text) = result.get("text").and_then(serde_json::Value::as_str) else {
return Vec::new();
};
if text.is_empty() || text.len() > MAX_THREAD_HINT_BYTES {
return Vec::new();
}
vec![PromptFragment::new(
PromptSlot::ContextWindow,
text,
ContentItemKind("notes.thread_hint".to_string()),
)]
})
}
}
impl ToolContributor for HistoryNotesExtension {
fn tools(
&self,
@@ -114,5 +162,6 @@ pub fn install(registry: &mut ExtensionRegistryBuilder<Config>, auth_manager: Ar
let extension = Arc::new(HistoryNotesExtension { auth_manager });
registry.thread_lifecycle_contributor(extension.clone());
registry.config_contributor(extension.clone());
registry.prompt_contributor(extension.clone());
registry.tool_contributor(extension);
}
@@ -4,11 +4,14 @@ use codex_config::LoaderOverrides;
use codex_core::config::Config;
use codex_core::config::ConfigBuilder;
use codex_core::config::TokenBudgetConfig;
use codex_extension_api::ContentItemKind;
use codex_extension_api::ConversationHistory;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionRegistry;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_extension_api::NoopTurnItemEmitter;
use codex_extension_api::PromptFragment;
use codex_extension_api::PromptSlot;
use codex_extension_api::ThreadStartInput;
use codex_extension_api::ToolCall;
use codex_extension_api::ToolCallSource;
@@ -40,6 +43,7 @@ use wiremock::matchers::method;
use wiremock::matchers::path;
type TestResult = Result<(), Box<dyn std::error::Error>>;
const THREAD_HINT: &str = "Recent notes (up to 5, most-recent first):\n- /root/worker/notes/latest.md (2 lines, 14 UTF-8 bytes)";
#[tokio::test]
async fn installed_extension_exposes_and_invokes_history_notes_tools() -> TestResult {
@@ -52,6 +56,12 @@ async fn installed_extension_exposes_and_invokes_history_notes_tools() -> TestRe
})))
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path("/backend-api/codex/alpha/notes/v2/thread_hint"))
.and(header("x-openai-actor-authorization", "actor-biscuit"))
.respond_with(ResponseTemplate::new(200).set_body_json(json!({"text": THREAD_HINT})))
.mount(&server)
.await;
let codex_home = TempDir::new()?;
let mut config = ConfigBuilder::default()
.loader_overrides(LoaderOverrides::without_managed_config_for_tests())
@@ -142,8 +152,20 @@ async fn installed_extension_exposes_and_invokes_history_notes_tools() -> TestRe
)
);
let hints = registry.context_contributors()[0]
.contribute_thread_context(&session_store, &thread_store)
.await;
assert_eq!(
hints,
vec![PromptFragment::new(
PromptSlot::ContextWindow,
THREAD_HINT,
ContentItemKind("notes.thread_hint".to_string()),
)]
);
let requests = server.received_requests().await.expect("recorded requests");
assert_eq!(requests.len(), 1);
assert_eq!(requests.len(), 2);
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&requests[0].body)?,
json!({
@@ -155,12 +177,47 @@ async fn installed_extension_exposes_and_invokes_history_notes_tools() -> TestRe
})
);
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&requests[1].body)?,
json!({
"context": {
"session_id": "session-123",
"current_agent_name": "/root/worker",
}
})
);
for result in [
json!({"text": ""}),
json!({"text": "x".repeat(4_001)}),
json!({"encrypted_output": "old-backend-hint"}),
] {
server.reset().await;
Mock::given(method("POST"))
.and(path("/backend-api/codex/alpha/notes/v2/thread_hint"))
.respond_with(ResponseTemplate::new(200).set_body_json(result))
.mount(&server)
.await;
assert!(
registry.context_contributors()[0]
.contribute_thread_context(&session_store, &thread_store)
.await
.is_empty()
);
}
let mut disabled_config = config.clone();
disabled_config.token_budget = None;
for contributor in registry.config_contributors() {
contributor.on_config_changed(&session_store, &thread_store, &config, &disabled_config);
}
assert!(exposed_tools(&registry, &session_store, &thread_store).is_empty());
assert!(
registry.context_contributors()[0]
.contribute_thread_context(&session_store, &thread_store)
.await
.is_empty()
);
Ok(())
}