mirror of
https://github.com/openai/codex.git
synced 2026-11-03 19:53:25 +00:00
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:
@@ -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(®istry, &session_store, &thread_store).is_empty());
|
||||
assert!(
|
||||
registry.context_contributors()[0]
|
||||
.contribute_thread_context(&session_store, &thread_store)
|
||||
.await
|
||||
.is_empty()
|
||||
);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user