better example

This commit is contained in:
jif-oai
2026-05-08 12:09:26 +02:00
parent aa9bb90611
commit 1378f42e04
2 changed files with 57 additions and 104 deletions

View File

@@ -16,8 +16,3 @@ codex-protocol = { workspace = true }
codex-tools = { workspace = true }
serde_json = { workspace = true }
thiserror = { workspace = true }
tokio = { version = "1.49.0", features = ["rt", "rt-multi-thread", "macros"] }
[dev-dependencies]
codex-git-attribution = { workspace = true }
codex-memories = { workspace = true }

View File

@@ -1,110 +1,68 @@
#[path = "enabled_extensions/shared_state_extension.rs"]
mod shared_state_extension;
use std::sync::Arc;
use codex_extension_api::ExtensionData;
use codex_extension_api::ExtensionRegistryBuilder;
use codex_git_attribution as git_attribution;
use codex_memories::MemoriesExtension;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use codex_memories::citation_output::MemoryState;
use codex_protocol::items::{TurnItem, UserMessageItem};
use shared_state_extension::SharedStateExtension;
use shared_state_extension::recorded_style_contributions;
use shared_state_extension::recorded_usage_contributions;
#[tokio::main]
async fn main() {
// 1. Build the concrete extension values owned by the host.
let memories = Arc::new(MemoriesExtension::new(
Some("Use memories when they help answer the user.".to_string()),
std::env::temp_dir().join("codex-memories-example"),
));
let git_attribution = git_attribution::extension();
fn main() {
// 1. Build the extension value owned by the host.
let extension = Arc::new(SharedStateExtension);
// 2. Install those extensions into one registry for the runtime context type
// this host exposes to contributors.
let registry = ExtensionRegistryBuilder::<ctx::RuntimeContext>::new()
.with_extension(git_attribution)
.with_extension(memories)
// 2. Install it into the registry for the context type this host exposes.
let registry = ExtensionRegistryBuilder::<()>::new()
.with_extension(extension)
.build();
// 3. Build the runtime facts contributors can inspect when they decide
// whether they are active for this request.
//
// Ideally, this is instead the TurnContext or the Config or so
let runtime = ctx::RuntimeContext {
commit_attribution: None,
memory_tool_enabled: true,
use_memories: true,
};
// 3. The host decides which stores are shared.
let session_store = ExtensionData::new();
let first_thread_store = ExtensionData::new();
let second_thread_store = ExtensionData::new();
// 4. Build the host-owned stores used by the active contribution families.
let session_data = ExtensionData::new();
let thread_data = ExtensionData::new();
let turn_data = ExtensionData::new();
// 4. Reusing the same session store shares session state across threads.
let first_thread_fragments = contribute_prompt(&registry, &session_store, &first_thread_store);
contribute_prompt(&registry, &session_store, &first_thread_store);
contribute_prompt(&registry, &session_store, &second_thread_store);
// 5. Invoke whichever contribution families this insertion point needs.
let prompt_fragments = registry
println!("first prompt fragments: {}", first_thread_fragments.len());
println!(
"session style contributions: {}",
recorded_style_contributions(&session_store)
);
println!(
"session usage contributions: {}",
recorded_usage_contributions(&session_store)
);
println!(
"first thread style contributions: {}",
recorded_style_contributions(&first_thread_store)
);
println!(
"first thread usage contributions: {}",
recorded_usage_contributions(&first_thread_store)
);
println!(
"second thread style contributions: {}",
recorded_style_contributions(&second_thread_store)
);
println!(
"second thread usage contributions: {}",
recorded_usage_contributions(&second_thread_store)
);
}
fn contribute_prompt(
registry: &codex_extension_api::ExtensionRegistry<()>,
session_store: &ExtensionData,
thread_store: &ExtensionData,
) -> Vec<codex_extension_api::PromptFragment> {
registry
.context_contributors()
.iter()
.flat_map(|contributor| contributor.contribute(&runtime, &session_data, &thread_data))
.collect::<Vec<_>>();
let tools = registry
.tool_contributors()
.iter()
.flat_map(|contributor| contributor.tools(&runtime, &thread_data))
.collect::<Vec<_>>();
let tools = registry
.tool_contributors()
.iter()
.flat_map(|contributor| contributor.tools(&runtime, &thread_data))
.collect::<Vec<_>>();
let mut item =TurnItem::UserMessage(UserMessageItem {
content: vec!(),
id: String::new()
});
for contributor in registry.turn_item_contributors() {
let _ = contributor.contribute(&runtime, &thread_data, &turn_data, &mut item).await;
}
for contributor in registry.turn_item_contributors() {
let _ = contributor.contribute(&runtime, &thread_data, &turn_data, &mut item).await;
}
let session_state = session_data.get_or_init::<MemoryState>(Default::default);
let thread_state = thread_data.get_or_init::<MemoryState>(Default::default);
let turn_state = turn_data.get_or_init::<MemoryState>(Default::default);
println!("Session: {} (expected 1)", session_state.counter.load(Ordering::Relaxed));
println!("Thread: {} (expected 3)", thread_state.counter.load(Ordering::Relaxed));
println!("Turn: {} (expected 2)", turn_state.counter.load(Ordering::Relaxed));
println!("prompt fragments: {}", prompt_fragments.len());
println!("native tools: {}", tools.len());
}
mod ctx {
use codex_git_attribution::GitAttributionContext;
use codex_memories::ctx::MemoriesContext;
/// Host-owned projection exposed to extension contributors.
#[derive(Clone)]
pub struct RuntimeContext {
pub commit_attribution: Option<String>,
pub memory_tool_enabled: bool,
pub use_memories: bool,
}
impl GitAttributionContext for RuntimeContext {
fn commit_attribution(&self) -> Option<&str> {
self.commit_attribution.as_deref()
}
}
impl MemoriesContext for RuntimeContext {
fn memory_tool_enabled(&self) -> bool {
self.memory_tool_enabled
}
fn use_memories(&self) -> bool {
self.use_memories
}
}
.flat_map(|contributor| contributor.contribute(&(), session_store, thread_store))
.collect()
}