use std::{collections::BTreeMap, sync::Arc, time::Duration};
use runtime_extension_api::{
AgentDefinition, ContextBudgetFrame, ContextEvent, ContextFrame, ContextInvalidation,
ContextItem, ContextResource, ContextSensitivity, ContextSlot, ContextToolFrame,
ExtensionContext,
};
use runtime_kernel_api::{KernelEvent, KernelEventSink, KernelFailure};
use runtime_types::ExecutionId;
pub(crate) struct ContextHookState {
resources: BTreeMap<(ContextSlot, String), ContextResource>,
pending_invalidations: ContextInvalidation,
budget: Option<ContextBudgetFrame>,
latest_tool: Option<ContextToolFrame>,
transcript_summary: String,
}
impl ContextHookState {
pub(crate) fn new(resources: Vec<ContextResource>) -> Result<Self, KernelFailure> {
let mut indexed = BTreeMap::new();
for resource in resources {
validate_resource(&resource)?;
indexed.insert((resource.slot, resource.id.clone()), resource);
}
Ok(Self {
resources: indexed,
pending_invalidations: ContextInvalidation::NONE,
budget: None,
latest_tool: None,
transcript_summary: String::new(),
})
}
pub(crate) fn observe_model_input(
&mut self,
budget: ContextBudgetFrame,
transcript_summary: String,
) {
self.budget = Some(budget);
self.transcript_summary = transcript_summary;
}
#[allow(clippy::too_many_arguments)]
pub(crate) async fn run(
&mut self,
definition: &AgentDefinition,
extension: &ExtensionContext,
event: ContextEvent,
execution_id: &ExecutionId,
turn: usize,
user_prompt: &str,
visible_tools: &[String],
events: &Arc<dyn KernelEventSink>,
) -> Result<(), KernelFailure> {
if let ContextEvent::AfterTool { invalidation, .. } = &event {
self.pending_invalidations = self.pending_invalidations.union(*invalidation);
}
if let ContextEvent::AfterTool {
tool_name, failed, ..
} = &event
{
self.latest_tool = Some(ContextToolFrame {
name: tool_name.clone(),
failed: *failed,
});
}
for hook in &definition.context_hooks {
let frame = ContextFrame {
execution_id: execution_id.clone(),
turn,
user_prompt: user_prompt.to_string(),
visible_tools: visible_tools.to_vec(),
pending_invalidations: self.pending_invalidations,
resources: self.resources.values().cloned().collect(),
budget: self.budget.clone(),
latest_tool: self.latest_tool.clone(),
transcript_summary: self.transcript_summary.clone(),
};
let result = tokio::time::timeout(
Duration::from_secs(5),
hook.apply(extension, &event, &frame),
)
.await;
match result {
Ok(Ok(patch)) => {
self.pending_invalidations =
self.pending_invalidations.union(patch.invalidations);
for removal in patch.removals {
self.resources.remove(&(removal.slot, removal.id));
}
for resource in patch.resources {
validate_resource(&resource)?;
let key = (resource.slot, resource.id.clone());
if self
.resources
.get(&key)
.is_some_and(|current| current.digest == resource.digest)
{
continue;
}
self.resources.insert(key, resource);
}
}
Ok(Err(error)) => {
events
.emit(KernelEvent::Warning {
code: "CONTEXT_HOOK_FAILED".into(),
message: format!("{}: {error}", hook.name()),
})
.await?;
}
Err(_) => {
events
.emit(KernelEvent::Warning {
code: "CONTEXT_HOOK_TIMEOUT".into(),
message: hook.name().to_string(),
})
.await?;
}
}
}
if matches!(event, ContextEvent::BeforeModel { .. }) {
self.pending_invalidations = ContextInvalidation::NONE;
}
Ok(())
}
pub(crate) fn context_items(&self) -> Vec<ContextItem> {
self.resources
.values()
.filter(|resource| {
resource.sensitivity != ContextSensitivity::Sensitive
&& resource.slot != ContextSlot::Checkpoint
})
.map(|resource| ContextItem {
source: format!("{}:{}@{}", resource.source, resource.id, resource.version),
content: resource.content.clone(),
priority: slot_priority(resource.slot),
required: resource.required,
})
.collect()
}
pub(crate) fn checkpoint_candidate(&self) -> Option<&str> {
self.resources
.iter()
.rev()
.find(|((slot, _), resource)| {
*slot == ContextSlot::Checkpoint
&& resource.sensitivity != ContextSensitivity::Sensitive
})
.map(|(_, resource)| resource.content.as_str())
}
}
fn validate_resource(resource: &ContextResource) -> Result<(), KernelFailure> {
if resource.id.trim().is_empty()
|| resource.id.len() > 256
|| resource.source.trim().is_empty()
|| resource.version.trim().is_empty()
|| resource.digest.trim().is_empty()
{
return Err(KernelFailure::new(
"CONTEXT_RESOURCE_INVALID",
"Context Hook resource requires bounded id, source, and version",
));
}
if resource.required && resource.sensitivity == ContextSensitivity::Sensitive {
return Err(KernelFailure::new(
"CONTEXT_RESOURCE_SENSITIVE_REQUIRED",
"required sensitive Context Hook resource cannot be projected to the model",
));
}
Ok(())
}
fn slot_priority(slot: ContextSlot) -> i32 {
match slot {
ContextSlot::System => 1_000,
ContextSlot::Instructions => 900,
ContextSlot::Checkpoint => 850,
ContextSlot::ToolCatalog => 800,
ContextSlot::Skills => 700,
ContextSlot::Memory => 600,
ContextSlot::Plan => 500,
ContextSlot::RuntimeState => 400,
ContextSlot::History => 350,
ContextSlot::Observations => 300,
}
}