use std::{sync::Arc, time::Duration};
use async_trait::async_trait;
use runtime_extension_api::{
CompletionMode, ContextBudgetFrame, ContextEvent, ExtensionContext, ObserverEvent,
};
use runtime_kernel_api::{
KernelEvent, KernelEventSink, KernelFailure, KernelOutcome, KernelSpec, RuntimeKernel,
};
use runtime_ports::{
ContextAuditDerivation, ContextAuditSnapshot, ContextAuditStage, ContextDerivationKind,
ConversationMessage, ExecutionSessions, ModelContent, ModelInvocationAudit,
ModelInvocationAuditState, ModelMessage, ModelRequest, ModelRole, ModelToolDefinition,
OperationContext, TokenUsage,
};
use sha2::{Digest, Sha256};
use crate::{
compaction::{CompactionController, has_unclosed_tool_transaction},
context::ContextManager,
context_hooks::ContextHookState,
context_measurement::{ContextMeasurement, model_request_digest},
failure_fold::fold_repeated_failures,
tool_runtime::ToolRuntime,
};
#[derive(Debug, Default)]
pub struct AgentKernel;
#[async_trait]
impl RuntimeKernel for AgentKernel {
async fn execute(
&self,
operation: OperationContext,
spec: KernelSpec,
sessions: ExecutionSessions,
events: Arc<dyn KernelEventSink>,
) -> Result<KernelOutcome, KernelFailure> {
validate_spec(&spec)?;
operation.remaining().map_err(port_error)?;
events.emit(KernelEvent::Started).await?;
observe(
&spec,
ObserverEvent::ExecutionStarted {
execution_id: spec.execution_id.clone(),
},
)
.await;
let resources = sessions
.resources
.load(&operation)
.await
.map_err(port_error)?;
let resource_version = resources.digest.clone();
let extension_context = ExtensionContext {
execution_id: spec.execution_id.clone(),
deadline: operation.deadline,
cancellation: operation.cancellation.clone(),
};
let mut context_resources = resources
.instructions
.into_iter()
.enumerate()
.map(|(index, content)| runtime_extension_api::ContextResource {
id: format!("instruction-{index}"),
slot: runtime_extension_api::ContextSlot::Instructions,
source: "resource-session".into(),
estimated_tokens: estimated_text_tokens(&content),
digest: content_digest(&content),
content,
required: true,
sensitivity: runtime_extension_api::ContextSensitivity::Internal,
version: resource_version.clone(),
})
.collect::<Vec<_>>();
context_resources.extend(resources.skills.into_iter().map(|(name, content)| {
runtime_extension_api::ContextResource {
id: name,
slot: runtime_extension_api::ContextSlot::Skills,
source: "resource-session".into(),
estimated_tokens: estimated_text_tokens(&content),
digest: content_digest(&content),
content,
required: false,
sensitivity: runtime_extension_api::ContextSensitivity::Internal,
version: resource_version.clone(),
}
}));
let mut legacy_source_index = 0usize;
for source in &spec.package.context_sources {
match tokio::time::timeout(Duration::from_secs(5), source.load(&extension_context))
.await
{
Ok(Ok(items)) => {
for item in items {
legacy_source_index = legacy_source_index.saturating_add(1);
context_resources.push(runtime_extension_api::ContextResource {
id: format!("legacy-context-{legacy_source_index}"),
slot: runtime_extension_api::ContextSlot::Instructions,
source: item.source,
estimated_tokens: estimated_text_tokens(&item.content),
digest: content_digest(&item.content),
content: item.content,
required: item.required,
sensitivity: runtime_extension_api::ContextSensitivity::Internal,
version: "legacy-v1".into(),
});
}
}
Ok(Err(error)) => {
events
.emit(KernelEvent::Warning {
code: "CONTEXT_SOURCE_FAILED".into(),
message: format!("{}: {error}", source.name()),
})
.await?
}
Err(_) => {
events
.emit(KernelEvent::Warning {
code: "CONTEXT_SOURCE_TIMEOUT".into(),
message: source.name().to_string(),
})
.await?
}
}
}
let history = sessions
.conversation
.load_recent(&operation, spec.limits.context.max_messages)
.await
.map_err(port_error)?;
let context = ContextManager::new(spec.limits.context.max_messages)?;
let tool_runtime = ToolRuntime::new(
spec.package.clone(),
spec.limits.max_tool_output_bytes,
&spec.granted_capabilities,
)?;
let tools = tool_runtime.model_definitions();
let visible_tools = tools
.iter()
.map(|tool| tool.name.clone())
.collect::<Vec<_>>();
let mut hook_state = ContextHookState::new(context_resources)?;
hook_state
.run(
&spec.package,
&extension_context,
ContextEvent::ExecutionStart,
&spec.execution_id,
0,
&spec.user_prompt,
&visible_tools,
&events,
)
.await?;
let assembled_items = hook_state.context_items();
let original_messages = raw_context_candidates(
&spec.package.system_prompt,
&assembled_items,
&history,
&spec.user_prompt,
);
let original_items_count = original_messages.len().saturating_add(tools.len());
let original_request_digest = model_request_digest(
&original_messages,
&tools,
Some(spec.limits.context.reserved_output_tokens),
);
let mut messages = context.assemble(
&spec.package.system_prompt,
assembled_items,
history,
&spec.user_prompt,
)?;
let original_snapshot_id = format!("{}-context-original", spec.execution_id);
let mut compaction = CompactionController::new(spec.limits.context.clone());
let original_measurement = compaction.measurement(&original_messages, &tools);
let original_tokens = original_measurement.total_tokens;
sessions
.context_audit
.snapshot(
&operation,
ContextAuditSnapshot {
id: original_snapshot_id.clone(),
stage: ContextAuditStage::Original,
purpose: "execution context candidates".into(),
messages: original_messages,
tools: tools.clone(),
protocol_overhead_tokens: original_measurement.protocol_overhead_tokens,
reserved_output_tokens: original_measurement.reserved_output_tokens,
estimated_tokens: original_tokens,
request_digest: original_request_digest,
},
)
.await
.map_err(port_error)?;
sessions
.conversation
.append(
&operation,
&[ConversationMessage {
role: ModelRole::User,
text: spec.user_prompt.clone(),
}],
)
.await
.map_err(port_error)?;
let mut total_usage = TokenUsage::zero();
let mut tool_calls = 0usize;
for turn in 1..=spec.limits.max_model_turns {
operation.remaining().map_err(port_error)?;
let current_measurement = compaction.measurement(&messages, &tools);
hook_state.observe_model_input(
budget_frame(current_measurement, spec.limits.context.max_input_tokens),
normalized_transcript_summary(&messages),
);
hook_state
.run(
&spec.package,
&extension_context,
ContextEvent::BeforeModel { turn },
&spec.execution_id,
turn,
&spec.user_prompt,
&visible_tools,
&events,
)
.await?;
let refreshed_items = hook_state.context_items();
context.refresh_system(&mut messages, &spec.package.system_prompt, refreshed_items)?;
fold_repeated_failures(&mut messages);
let initial_decision = compaction.decide(&messages, &tools, &spec.user_prompt, turn)?;
let unclosed_transaction = has_unclosed_tool_transaction(&messages);
if initial_decision.is_some() && unclosed_transaction {
events
.emit(KernelEvent::Warning {
code: "CONTEXT_COMPACTION_DEFERRED".into(),
message: "an unclosed ToolUse/ToolResult transaction is resident".into(),
})
.await?;
let measurement = compaction.measurement(&messages, &tools);
if measurement.total_tokens > spec.limits.context.max_input_tokens {
return finish_failure(
&spec,
"CONTEXT_LIMIT_EXCEEDED",
"context exceeds its hard limit while a Tool transaction is open",
)
.await;
}
} else if initial_decision.is_some() {
let measurement = compaction.measurement(&messages, &tools);
hook_state.observe_model_input(
budget_frame(measurement, spec.limits.context.max_input_tokens),
normalized_transcript_summary(&messages),
);
hook_state
.run(
&spec.package,
&extension_context,
ContextEvent::BeforeCompaction { turn },
&spec.execution_id,
turn,
&spec.user_prompt,
&visible_tools,
&events,
)
.await?;
let refreshed_items = hook_state.context_items();
context.refresh_system(
&mut messages,
&spec.package.system_prompt,
refreshed_items,
)?;
fold_repeated_failures(&mut messages);
}
let assembled_snapshot_id = format!("{}-context-turn-{turn}", spec.execution_id);
let assembled_measurement = compaction.measurement(&messages, &tools);
sessions
.context_audit
.snapshot(
&operation,
context_audit_snapshot(
assembled_snapshot_id.clone(),
ContextAuditStage::Assembled,
format!("model request turn {turn} before compaction"),
&messages,
&tools,
assembled_measurement,
spec.limits.context.reserved_output_tokens,
),
)
.await
.map_err(port_error)?;
if turn == 1 {
sessions
.context_audit
.derivation(
&operation,
ContextAuditDerivation {
id: format!("{}-context-assemble", spec.execution_id),
source_snapshot_id: original_snapshot_id.clone(),
target_snapshot_id: assembled_snapshot_id.clone(),
kind: ContextDerivationKind::Assemble,
input_tokens: original_tokens,
output_tokens: assembled_measurement.total_tokens,
input_items: original_items_count,
output_items: messages.len().saturating_add(tools.len()),
},
)
.await
.map_err(port_error)?;
}
let mut request_snapshot_id = assembled_snapshot_id.clone();
if !unclosed_transaction
&& let Some(decision) =
compaction.decide(&messages, &tools, &spec.user_prompt, turn)?
{
let checkpoint_candidate = hook_state.checkpoint_candidate().map(str::to_owned);
let outcome = compaction.compact(
&mut messages,
&tools,
&spec.user_prompt,
checkpoint_candidate.as_deref(),
turn,
)?;
let measurement = compaction.measurement(&messages, &tools);
hook_state.observe_model_input(
budget_frame(measurement, spec.limits.context.max_input_tokens),
normalized_transcript_summary(&messages),
);
hook_state
.run(
&spec.package,
&extension_context,
ContextEvent::AfterCompaction { turn },
&spec.execution_id,
turn,
&spec.user_prompt,
&visible_tools,
&events,
)
.await?;
let refreshed_items = hook_state.context_items();
context.refresh_system(
&mut messages,
&spec.package.system_prompt,
refreshed_items,
)?;
let compressed_measurement = compaction.measurement(&messages, &tools);
request_snapshot_id =
format!("{}-context-turn-{turn}-compressed", spec.execution_id);
sessions
.context_audit
.snapshot(
&operation,
context_audit_snapshot(
request_snapshot_id.clone(),
ContextAuditStage::Compressed,
format!("model request turn {turn} after checkpoint compaction"),
&messages,
&tools,
compressed_measurement,
spec.limits.context.reserved_output_tokens,
),
)
.await
.map_err(port_error)?;
sessions
.context_audit
.derivation(
&operation,
ContextAuditDerivation {
id: format!("{}-context-compact-{turn}", spec.execution_id),
source_snapshot_id: assembled_snapshot_id,
target_snapshot_id: request_snapshot_id.clone(),
kind: ContextDerivationKind::Compress,
input_tokens: outcome.before.total_tokens,
output_tokens: compressed_measurement.total_tokens,
input_items: messages
.len()
.saturating_add(outcome.removed_messages)
.saturating_add(tools.len()),
output_items: messages.len().saturating_add(tools.len()),
},
)
.await
.map_err(port_error)?;
events
.emit(KernelEvent::Warning {
code: "CONTEXT_COMPACTED".into(),
message: format!(
"reason={:?}; before={}; after={}; target={}; reclaimable={}; fallback={}",
decision.reason,
outcome.before.total_tokens,
compressed_measurement.total_tokens,
decision.target_tokens,
decision.reclaimable_tokens,
outcome.used_fallback
),
})
.await?;
}
let invocation_id = format!("{}-model-turn-{turn}", spec.execution_id);
sessions
.context_audit
.model_invocation(
&operation,
ModelInvocationAudit {
id: invocation_id.clone(),
request_snapshot_id: request_snapshot_id.clone(),
model: spec.model.clone(),
state: ModelInvocationAuditState::Prepared,
usage: TokenUsage::default(),
error_code: None,
},
)
.await
.map_err(port_error)?;
events.emit(KernelEvent::ModelStarted { turn }).await?;
let response = sessions
.model
.invoke(
&operation,
ModelRequest {
invocation_id: invocation_id.clone(),
model: spec.model.clone(),
messages: messages.clone(),
tools: tools.clone(),
max_output_tokens: Some(spec.limits.context.reserved_output_tokens),
},
)
.await;
let response = match response {
Ok(response) => response,
Err(error) => {
let _ = sessions
.context_audit
.model_invocation(
&operation,
ModelInvocationAudit {
id: invocation_id,
request_snapshot_id,
model: spec.model.clone(),
state: ModelInvocationAuditState::Failed,
usage: TokenUsage::default(),
error_code: Some(error.code.clone()),
},
)
.await;
return Err(port_error(error));
}
};
sessions
.context_audit
.model_invocation(
&operation,
ModelInvocationAudit {
id: invocation_id,
request_snapshot_id,
model: spec.model.clone(),
state: ModelInvocationAuditState::Succeeded,
usage: response.usage.clone(),
error_code: None,
},
)
.await
.map_err(port_error)?;
total_usage.add_assign(&response.usage);
events.emit(KernelEvent::ModelCompleted { turn }).await?;
observe(
&spec,
ObserverEvent::ModelCompleted {
execution_id: spec.execution_id.clone(),
turn,
},
)
.await;
if !response.finish.is_complete() {
return finish_failure(
&spec,
"MODEL_INCOMPLETE",
format!("model stopped with {:?}", response.finish),
)
.await;
}
let answer = response
.output
.iter()
.filter_map(|part| match part {
ModelContent::Text { text } => Some(text.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("");
let calls = response
.output
.iter()
.filter_map(|part| match part {
ModelContent::ToolUse {
id,
name,
arguments,
} => Some((id.clone(), name.clone(), arguments.clone())),
_ => None,
})
.collect::<Vec<_>>();
messages.push(ModelMessage {
role: ModelRole::Assistant,
content: response.output.clone(),
});
if calls.is_empty() {
if answer.trim().is_empty() {
return finish_failure(
&spec,
"MODEL_EMPTY",
"model returned no text or tool calls",
)
.await;
}
if let CompletionMode::RequiredTool(name) = &spec.package.completion_mode {
return finish_failure(
&spec,
"COMPLETION_TOOL_REQUIRED",
format!("agent must finish with {name}"),
)
.await;
}
sessions
.conversation
.append(
&operation,
&[ConversationMessage {
role: ModelRole::Assistant,
text: answer.clone(),
}],
)
.await
.map_err(port_error)?;
observe(
&spec,
ObserverEvent::ExecutionFinished {
execution_id: spec.execution_id.clone(),
succeeded: true,
},
)
.await;
return Ok(KernelOutcome {
answer,
model_turns: turn,
tool_calls,
usage: total_usage,
});
}
for (call_id, name, arguments) in calls {
tool_calls += 1;
if tool_calls > spec.limits.max_tool_calls {
return finish_failure(
&spec,
"TOOL_BUDGET_EXHAUSTED",
"maximum tool calls exceeded",
)
.await;
}
let outcome = tool_runtime
.execute(&operation, &sessions, &events, &call_id, &name, arguments)
.await?;
let invalidation = outcome.context_invalidation;
messages.push(ModelMessage {
role: ModelRole::Tool,
content: vec![ModelContent::ToolResult {
tool_use_id: call_id,
name: name.clone(),
content: outcome.content.clone(),
failed: outcome.failed,
}],
});
let measurement = compaction.measurement(&messages, &tools);
hook_state.observe_model_input(
budget_frame(measurement, spec.limits.context.max_input_tokens),
normalized_transcript_summary(&messages),
);
hook_state
.run(
&spec.package,
&extension_context,
ContextEvent::AfterTool {
turn,
tool_name: name,
failed: outcome.failed,
invalidation,
},
&spec.execution_id,
turn,
&spec.user_prompt,
&visible_tools,
&events,
)
.await?;
if let Some(answer) = outcome.completion {
sessions
.conversation
.append(
&operation,
&[ConversationMessage {
role: ModelRole::Assistant,
text: answer.clone(),
}],
)
.await
.map_err(port_error)?;
observe(
&spec,
ObserverEvent::ExecutionFinished {
execution_id: spec.execution_id.clone(),
succeeded: true,
},
)
.await;
return Ok(KernelOutcome {
answer,
model_turns: turn,
tool_calls,
usage: total_usage,
});
}
}
}
finish_failure(
&spec,
"MODEL_TURN_BUDGET_EXHAUSTED",
"maximum model turns exceeded",
)
.await
}
}
#[allow(clippy::too_many_arguments)]
fn context_audit_snapshot(
id: String,
stage: ContextAuditStage,
purpose: String,
messages: &[ModelMessage],
tools: &[ModelToolDefinition],
measurement: ContextMeasurement,
max_output_tokens: u32,
) -> ContextAuditSnapshot {
ContextAuditSnapshot {
id,
stage,
purpose,
messages: messages.to_vec(),
tools: tools.to_vec(),
protocol_overhead_tokens: measurement.protocol_overhead_tokens,
reserved_output_tokens: measurement.reserved_output_tokens,
estimated_tokens: measurement.total_tokens,
request_digest: model_request_digest(messages, tools, Some(max_output_tokens)),
}
}
fn budget_frame(measurement: ContextMeasurement, max_input_tokens: u64) -> ContextBudgetFrame {
ContextBudgetFrame {
max_input_tokens,
message_tokens: measurement.message_tokens,
tool_schema_tokens: measurement.tool_schema_tokens,
protocol_overhead_tokens: measurement.protocol_overhead_tokens,
reserved_output_tokens: measurement.reserved_output_tokens,
total_tokens: measurement.total_tokens,
}
}
fn estimated_text_tokens(content: &str) -> u64 {
(content.len() as u64).saturating_add(3) / 4
}
fn content_digest(content: &str) -> String {
format!("sha256:{:x}", Sha256::digest(content.as_bytes()))
}
fn normalized_transcript_summary(messages: &[ModelMessage]) -> String {
messages
.iter()
.enumerate()
.rev()
.take(12)
.map(|(index, message)| {
let content = message
.content
.iter()
.map(|content| match content {
ModelContent::Text { text } => format!("text(chars={})", text.chars().count()),
ModelContent::ToolUse { name, .. } => format!("tool_use(name={name})"),
ModelContent::ToolResult {
name,
content,
failed,
..
} => format!(
"tool_result(name={name},failed={failed},chars={})",
content.chars().count()
),
})
.collect::<Vec<_>>()
.join(",");
format!("message[{index}] role={:?} {content}", message.role)
})
.collect::<Vec<_>>()
.into_iter()
.rev()
.collect::<Vec<_>>()
.join("\n")
}
fn raw_context_candidates(
system_prompt: &str,
context_items: &[runtime_extension_api::ContextItem],
history: &[ConversationMessage],
user_prompt: &str,
) -> Vec<ModelMessage> {
let mut system = system_prompt.trim().to_string();
for item in context_items {
if !item.content.trim().is_empty() {
system.push_str("\n\n<context source=\"");
system.push_str(&item.source);
system.push_str("\">\n");
system.push_str(item.content.trim());
system.push_str("\n</context>");
}
}
let mut messages = Vec::with_capacity(history.len() + 2);
messages.push(ModelMessage::text(ModelRole::System, system));
messages.extend(
history
.iter()
.map(|message| ModelMessage::text(message.role, message.text.clone())),
);
messages.push(ModelMessage::text(ModelRole::User, user_prompt));
messages
}
fn validate_spec(spec: &KernelSpec) -> Result<(), KernelFailure> {
if spec.user_prompt.trim().is_empty() {
return Err(KernelFailure::new(
"PROMPT_EMPTY",
"user prompt must not be empty",
));
}
if spec.model.trim().is_empty() {
return Err(KernelFailure::new(
"MODEL_EMPTY",
"model binding must not be empty",
));
}
if spec.limits.max_model_turns == 0
|| spec.limits.max_tool_calls == 0
|| spec.limits.max_tool_output_bytes == 0
|| spec.limits.context.max_messages < 2
|| spec.limits.context.max_input_tokens < 1_024
|| u64::from(spec.limits.context.reserved_output_tokens)
>= spec.limits.context.max_input_tokens
|| spec.limits.context.compaction.target_percent
>= spec.limits.context.compaction.trigger_percent
|| spec.limits.context.compaction.trigger_percent
>= spec.limits.context.compaction.emergency_percent
|| spec.limits.context.compaction.emergency_percent > 100
{
return Err(KernelFailure::new(
"KERNEL_LIMIT_INVALID",
"kernel limits must be positive",
));
}
Ok(())
}
async fn observe(spec: &KernelSpec, event: ObserverEvent) {
for observer in &spec.package.observers {
let _ = tokio::time::timeout(Duration::from_secs(2), observer.observe(event.clone())).await;
}
}
async fn finish_failure<T>(
spec: &KernelSpec,
code: &str,
message: impl Into<String>,
) -> Result<T, KernelFailure> {
observe(
spec,
ObserverEvent::ExecutionFinished {
execution_id: spec.execution_id.clone(),
succeeded: false,
},
)
.await;
Err(KernelFailure::new(code, message))
}
fn port_error(error: runtime_ports::PortFailure) -> KernelFailure {
KernelFailure {
code: error.code,
message: error.message,
retryable: error.retryable,
commit: error.commit,
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use std::time::{Duration, Instant};
use agent_runtime_testkit::{
CompleteTool, MemoryConversationSession, RecordingContextAuditSession,
RecordingKernelEvents, ScriptedModelSession, fixture_package, sessions,
sessions_with_context_audit,
};
use async_trait::async_trait;
use runtime_extension_api::{
AgentPackage, CompletionMode, ContextFrame, ContextHook, ContextInvalidation, ContextPatch,
ContextResource, ContextSensitivity, ContextSlot, ExtensionError, SideEffect, Tool,
ToolContext, ToolDescriptor, ToolOutcome,
};
use runtime_kernel_api::{ContextLimits, KernelLimits, RuntimeKernel};
use runtime_ports::{
ContextAuditStage, ModelContent, ModelFinish, ModelInvocationAuditState, ModelResponse,
OperationContext, TokenUsage,
};
use runtime_types::{ExecutionId, OperationId};
use serde_json::json;
use tokio_util::sync::CancellationToken;
use super::*;
fn operation() -> OperationContext {
OperationContext {
id: OperationId::random(),
execution_id: ExecutionId::random(),
deadline: Instant::now() + Duration::from_secs(10),
cancellation: CancellationToken::new(),
}
}
fn spec(operation: &OperationContext) -> KernelSpec {
KernelSpec {
execution_id: operation.execution_id.clone(),
conversation_id: None,
user_prompt: "finish the task".into(),
model: "fixture-model".into(),
package: fixture_package(CompletionMode::RequiredTool("runtime.complete".into())),
granted_capabilities: std::collections::BTreeSet::new(),
limits: KernelLimits {
max_model_turns: 4,
max_tool_calls: 4,
context: ContextLimits {
max_messages: 16,
max_input_tokens: 16 * 1024,
reserved_output_tokens: 1_024,
..ContextLimits::default()
},
max_tool_output_bytes: 4096,
},
}
}
struct RefreshContextTool;
#[async_trait]
impl Tool for RefreshContextTool {
fn descriptor(&self) -> ToolDescriptor {
ToolDescriptor {
name: "context.refresh".into(),
description: "Refresh dynamic memory context.".into(),
input_schema: json!({"type":"object","additionalProperties":false}),
required_capabilities: Vec::new(),
side_effect: SideEffect::ReadOnly,
}
}
async fn invoke(
&self,
_: ToolContext,
_: serde_json::Value,
) -> Result<ToolOutcome, ExtensionError> {
Ok(ToolOutcome::success("refreshed")
.with_context_invalidation(ContextInvalidation::MEMORY))
}
}
struct RefreshMemoryHook;
#[async_trait]
impl ContextHook for RefreshMemoryHook {
fn name(&self) -> &str {
"test.refresh-memory"
}
async fn apply(
&self,
_: &ExtensionContext,
event: &ContextEvent,
frame: &ContextFrame,
) -> Result<ContextPatch, ExtensionError> {
if matches!(event, ContextEvent::BeforeModel { .. })
&& frame
.pending_invalidations
.contains(ContextInvalidation::MEMORY)
{
return Ok(ContextPatch {
resources: vec![ContextResource {
id: "working-memory".into(),
slot: ContextSlot::Memory,
source: "test-hook".into(),
content: "dynamic-memory-v2".into(),
required: false,
sensitivity: ContextSensitivity::Internal,
version: "2".into(),
digest: "dynamic-memory-v2".into(),
estimated_tokens: 8,
}],
removals: Vec::new(),
invalidations: ContextInvalidation::NONE,
});
}
Ok(ContextPatch::default())
}
}
#[tokio::test]
async fn model_tool_completion_closes_the_loop() {
let model = Arc::new(ScriptedModelSession::new(vec![ModelResponse {
output: vec![ModelContent::ToolUse {
id: "call-1".into(),
name: "runtime.complete".into(),
arguments: json!({"answer":"done"}),
}],
finish: ModelFinish::ToolCalls,
usage: TokenUsage {
input_tokens: Some(10),
output_tokens: Some(2),
},
}]));
let operation = operation();
let outcome = AgentKernel
.execute(
operation.clone(),
spec(&operation),
sessions(model.clone()),
Arc::new(RecordingKernelEvents::default()),
)
.await
.unwrap();
assert_eq!(outcome.answer, "done");
assert_eq!(outcome.tool_calls, 1);
assert_eq!(model.requests().len(), 1);
}
#[tokio::test]
async fn tool_invalidation_refreshes_context_before_the_next_model_turn() {
let model = Arc::new(ScriptedModelSession::new(vec![
ModelResponse {
output: vec![ModelContent::ToolUse {
id: "call-refresh".into(),
name: "context.refresh".into(),
arguments: json!({}),
}],
finish: ModelFinish::ToolCalls,
usage: TokenUsage::default(),
},
ModelResponse {
output: vec![ModelContent::ToolUse {
id: "call-complete".into(),
name: "runtime.complete".into(),
arguments: json!({"answer":"done"}),
}],
finish: ModelFinish::ToolCalls,
usage: TokenUsage::default(),
},
]));
let operation = operation();
let mut spec = spec(&operation);
let fixture = fixture_package(CompletionMode::RequiredTool("runtime.complete".into()));
spec.package = Arc::new(AgentPackage {
manifest: fixture.manifest.clone(),
system_prompt: fixture.system_prompt.clone(),
tools: vec![Arc::new(CompleteTool), Arc::new(RefreshContextTool)],
context_sources: Vec::new(),
context_hooks: vec![Arc::new(RefreshMemoryHook)],
guards: Vec::new(),
observers: Vec::new(),
completion_mode: CompletionMode::RequiredTool("runtime.complete".into()),
});
AgentKernel
.execute(
operation.clone(),
spec,
sessions(model.clone()),
Arc::new(RecordingKernelEvents::default()),
)
.await
.unwrap();
let requests = model.requests();
assert_eq!(requests.len(), 2);
let second_system = serde_json::to_string(&requests[1].messages[0]).unwrap();
assert!(second_system.contains("dynamic-memory-v2"));
}
#[tokio::test]
async fn records_exact_model_context_and_invocation_lifecycle() {
let model = Arc::new(ScriptedModelSession::new(vec![ModelResponse {
output: vec![ModelContent::ToolUse {
id: "call-1".into(),
name: "runtime.complete".into(),
arguments: json!({"answer":"done"}),
}],
finish: ModelFinish::ToolCalls,
usage: TokenUsage {
input_tokens: Some(11),
output_tokens: Some(3),
},
}]));
let audit = Arc::new(RecordingContextAuditSession::default());
let operation = operation();
AgentKernel
.execute(
operation.clone(),
spec(&operation),
sessions_with_context_audit(model.clone(), audit.clone()),
Arc::new(RecordingKernelEvents::default()),
)
.await
.unwrap();
let snapshots = audit.snapshots.lock().unwrap();
assert_eq!(snapshots.len(), 2);
assert_eq!(snapshots[0].stage, ContextAuditStage::Original);
assert_eq!(snapshots[1].stage, ContextAuditStage::Assembled);
assert_eq!(snapshots[1].messages.len(), 2);
let request = &model.requests()[0];
assert_eq!(snapshots[1].tools, request.tools);
assert_eq!(
snapshots[1].reserved_output_tokens,
u64::from(request.max_output_tokens.unwrap())
);
assert_eq!(
snapshots[1].request_digest,
model_request_digest(&request.messages, &request.tools, request.max_output_tokens)
);
drop(snapshots);
let derivations = audit.derivations.lock().unwrap();
assert_eq!(derivations.len(), 1);
assert_eq!(
derivations[0].source_snapshot_id,
format!("{}-context-original", operation.execution_id)
);
drop(derivations);
let invocations = audit.invocations.lock().unwrap();
assert_eq!(invocations.len(), 2);
assert_eq!(invocations[0].state, ModelInvocationAuditState::Prepared);
assert_eq!(invocations[1].state, ModelInvocationAuditState::Succeeded);
assert_eq!(invocations[1].usage.input_tokens, Some(11));
}
struct RecordingCompactionHook {
events: std::sync::Mutex<Vec<&'static str>>,
}
#[async_trait]
impl ContextHook for RecordingCompactionHook {
fn name(&self) -> &str {
"test.compaction-lifecycle"
}
async fn apply(
&self,
_: &ExtensionContext,
event: &ContextEvent,
_: &ContextFrame,
) -> Result<ContextPatch, ExtensionError> {
let name = match event {
ContextEvent::BeforeCompaction { .. } => Some("before"),
ContextEvent::AfterCompaction { .. } => Some("after"),
_ => None,
};
if let Some(name) = name {
self.events.lock().unwrap().push(name);
}
Ok(ContextPatch::default())
}
}
#[tokio::test]
async fn compaction_preserves_resident_set_and_audits_lifecycle() {
let model = Arc::new(ScriptedModelSession::new(vec![ModelResponse {
output: vec![ModelContent::ToolUse {
id: "call-complete".into(),
name: "runtime.complete".into(),
arguments: json!({"answer":"done"}),
}],
finish: ModelFinish::ToolCalls,
usage: TokenUsage::default(),
}]));
let hook = Arc::new(RecordingCompactionHook {
events: std::sync::Mutex::new(Vec::new()),
});
let operation = operation();
let mut spec = spec(&operation);
spec.limits.context.max_input_tokens = 2_000;
spec.limits.context.reserved_output_tokens = 100;
spec.limits.context.protocol_overhead_tokens = 50;
spec.limits.context.compaction.trigger_percent = 50;
spec.limits.context.compaction.target_percent = 30;
spec.limits.context.compaction.emergency_percent = 90;
let fixture = fixture_package(CompletionMode::RequiredTool("runtime.complete".into()));
spec.package = Arc::new(AgentPackage {
manifest: fixture.manifest.clone(),
system_prompt: fixture.system_prompt.clone(),
tools: fixture.tools.clone(),
context_sources: fixture.context_sources.clone(),
context_hooks: vec![hook.clone()],
guards: fixture.guards.clone(),
observers: fixture.observers.clone(),
completion_mode: fixture.completion_mode.clone(),
});
let audit = Arc::new(RecordingContextAuditSession::default());
let mut execution_sessions = sessions_with_context_audit(model.clone(), audit.clone());
execution_sessions.conversation = Arc::new(MemoryConversationSession::with_messages(vec![
ConversationMessage {
role: ModelRole::User,
text: "old investigation".repeat(200),
},
ConversationMessage {
role: ModelRole::Assistant,
text: "old evidence".repeat(200),
},
]));
AgentKernel
.execute(
operation,
spec,
execution_sessions,
Arc::new(RecordingKernelEvents::default()),
)
.await
.unwrap();
assert_eq!(*hook.events.lock().unwrap(), vec!["before", "after"]);
let requests = model.requests();
assert_eq!(requests.len(), 1);
assert_eq!(requests[0].messages.len(), 3);
let encoded = serde_json::to_string(&requests[0].messages).unwrap();
assert!(encoded.contains("finish the task"));
assert!(encoded.contains("changed_unverified") || encoded.contains("investigating"));
let snapshots = audit.snapshots.lock().unwrap();
assert_eq!(
snapshots.last().unwrap().stage,
ContextAuditStage::Compressed
);
assert_eq!(
snapshots.last().unwrap().request_digest,
model_request_digest(
&requests[0].messages,
&requests[0].tools,
requests[0].max_output_tokens,
)
);
let derivations = audit.derivations.lock().unwrap();
assert!(
derivations
.iter()
.any(|derivation| derivation.kind == ContextDerivationKind::Compress)
);
}
#[tokio::test]
async fn required_completion_rejects_plain_model_stop() {
let model = Arc::new(ScriptedModelSession::new(vec![ModelResponse {
output: vec![ModelContent::Text {
text: "premature".into(),
}],
finish: ModelFinish::Stop,
usage: TokenUsage::default(),
}]));
let operation = operation();
let failure = AgentKernel
.execute(
operation.clone(),
spec(&operation),
sessions(model),
Arc::new(RecordingKernelEvents::default()),
)
.await
.unwrap_err();
assert_eq!(failure.code, "COMPLETION_TOOL_REQUIRED");
}
}