1use std::{sync::Arc, time::Duration};
2
3use async_trait::async_trait;
4use runtime_extension_api::{
5 CompletionMode, ContextBudgetFrame, ContextEvent, ExtensionContext, ObserverEvent,
6};
7use runtime_kernel_api::{
8 KernelEvent, KernelEventSink, KernelFailure, KernelOutcome, KernelSpec, RuntimeKernel,
9};
10use runtime_ports::{
11 ContextAuditDerivation, ContextAuditSnapshot, ContextAuditStage, ContextDerivationKind,
12 ConversationMessage, ExecutionSessions, ModelContent, ModelInvocationAudit,
13 ModelInvocationAuditState, ModelMessage, ModelRequest, ModelRole, ModelToolDefinition,
14 OperationContext, TokenUsage,
15};
16use sha2::{Digest, Sha256};
17
18use crate::{
19 compaction::{CompactionController, has_unclosed_tool_transaction},
20 context::ContextManager,
21 context_hooks::ContextHookState,
22 context_measurement::{ContextMeasurement, model_request_digest},
23 failure_fold::fold_repeated_failures,
24 tool_runtime::ToolRuntime,
25};
26
27#[derive(Debug, Default)]
28pub struct AgentKernel;
29
30#[async_trait]
31impl RuntimeKernel for AgentKernel {
32 async fn execute(
33 &self,
34 operation: OperationContext,
35 spec: KernelSpec,
36 sessions: ExecutionSessions,
37 events: Arc<dyn KernelEventSink>,
38 ) -> Result<KernelOutcome, KernelFailure> {
39 validate_spec(&spec)?;
40 operation.remaining().map_err(port_error)?;
41 events.emit(KernelEvent::Started).await?;
42 observe(
43 &spec,
44 ObserverEvent::ExecutionStarted {
45 execution_id: spec.execution_id.clone(),
46 },
47 )
48 .await;
49
50 let resources = sessions
51 .resources
52 .load(&operation)
53 .await
54 .map_err(port_error)?;
55 let resource_version = resources.digest.clone();
56 let extension_context = ExtensionContext {
57 execution_id: spec.execution_id.clone(),
58 deadline: operation.deadline,
59 cancellation: operation.cancellation.clone(),
60 };
61 let mut context_resources = resources
62 .instructions
63 .into_iter()
64 .enumerate()
65 .map(|(index, content)| runtime_extension_api::ContextResource {
66 id: format!("instruction-{index}"),
67 slot: runtime_extension_api::ContextSlot::Instructions,
68 source: "resource-session".into(),
69 estimated_tokens: estimated_text_tokens(&content),
70 digest: content_digest(&content),
71 content,
72 required: true,
73 sensitivity: runtime_extension_api::ContextSensitivity::Internal,
74 version: resource_version.clone(),
75 })
76 .collect::<Vec<_>>();
77 context_resources.extend(resources.skills.into_iter().map(|(name, content)| {
78 runtime_extension_api::ContextResource {
79 id: name,
80 slot: runtime_extension_api::ContextSlot::Skills,
81 source: "resource-session".into(),
82 estimated_tokens: estimated_text_tokens(&content),
83 digest: content_digest(&content),
84 content,
85 required: false,
86 sensitivity: runtime_extension_api::ContextSensitivity::Internal,
87 version: resource_version.clone(),
88 }
89 }));
90 let mut legacy_source_index = 0usize;
91 for source in &spec.package.context_sources {
92 match tokio::time::timeout(Duration::from_secs(5), source.load(&extension_context))
93 .await
94 {
95 Ok(Ok(items)) => {
96 for item in items {
97 legacy_source_index = legacy_source_index.saturating_add(1);
98 context_resources.push(runtime_extension_api::ContextResource {
99 id: format!("legacy-context-{legacy_source_index}"),
100 slot: runtime_extension_api::ContextSlot::Instructions,
101 source: item.source,
102 estimated_tokens: estimated_text_tokens(&item.content),
103 digest: content_digest(&item.content),
104 content: item.content,
105 required: item.required,
106 sensitivity: runtime_extension_api::ContextSensitivity::Internal,
107 version: "legacy-v1".into(),
108 });
109 }
110 }
111 Ok(Err(error)) => {
112 events
113 .emit(KernelEvent::Warning {
114 code: "CONTEXT_SOURCE_FAILED".into(),
115 message: format!("{}: {error}", source.name()),
116 })
117 .await?
118 }
119 Err(_) => {
120 events
121 .emit(KernelEvent::Warning {
122 code: "CONTEXT_SOURCE_TIMEOUT".into(),
123 message: source.name().to_string(),
124 })
125 .await?
126 }
127 }
128 }
129 let history = sessions
130 .conversation
131 .load_recent(&operation, spec.limits.context.max_messages)
132 .await
133 .map_err(port_error)?;
134 let context = ContextManager::new(spec.limits.context.max_messages)?;
135 let tool_runtime = ToolRuntime::new(
136 spec.package.clone(),
137 spec.limits.max_tool_output_bytes,
138 &spec.granted_capabilities,
139 )?;
140 let tools = tool_runtime.model_definitions();
141 let visible_tools = tools
142 .iter()
143 .map(|tool| tool.name.clone())
144 .collect::<Vec<_>>();
145 let mut hook_state = ContextHookState::new(context_resources)?;
146 hook_state
147 .run(
148 &spec.package,
149 &extension_context,
150 ContextEvent::ExecutionStart,
151 &spec.execution_id,
152 0,
153 &spec.user_prompt,
154 &visible_tools,
155 &events,
156 )
157 .await?;
158 let assembled_items = hook_state.context_items();
159 let original_messages = raw_context_candidates(
160 &spec.package.system_prompt,
161 &assembled_items,
162 &history,
163 &spec.user_prompt,
164 );
165 let original_items_count = original_messages.len().saturating_add(tools.len());
166 let original_request_digest = model_request_digest(
167 &original_messages,
168 &tools,
169 Some(spec.limits.context.reserved_output_tokens),
170 );
171 let mut messages = context.assemble(
172 &spec.package.system_prompt,
173 assembled_items,
174 history,
175 &spec.user_prompt,
176 )?;
177 let original_snapshot_id = format!("{}-context-original", spec.execution_id);
178 let mut compaction = CompactionController::new(spec.limits.context.clone());
179 let original_measurement = compaction.measurement(&original_messages, &tools);
180 let original_tokens = original_measurement.total_tokens;
181 sessions
182 .context_audit
183 .snapshot(
184 &operation,
185 ContextAuditSnapshot {
186 id: original_snapshot_id.clone(),
187 stage: ContextAuditStage::Original,
188 purpose: "execution context candidates".into(),
189 messages: original_messages,
190 tools: tools.clone(),
191 protocol_overhead_tokens: original_measurement.protocol_overhead_tokens,
192 reserved_output_tokens: original_measurement.reserved_output_tokens,
193 estimated_tokens: original_tokens,
194 request_digest: original_request_digest,
195 },
196 )
197 .await
198 .map_err(port_error)?;
199 sessions
200 .conversation
201 .append(
202 &operation,
203 &[ConversationMessage {
204 role: ModelRole::User,
205 text: spec.user_prompt.clone(),
206 }],
207 )
208 .await
209 .map_err(port_error)?;
210
211 let mut total_usage = TokenUsage::zero();
212 let mut tool_calls = 0usize;
213
214 for turn in 1..=spec.limits.max_model_turns {
215 operation.remaining().map_err(port_error)?;
216 let current_measurement = compaction.measurement(&messages, &tools);
217 hook_state.observe_model_input(
218 budget_frame(current_measurement, spec.limits.context.max_input_tokens),
219 normalized_transcript_summary(&messages),
220 );
221 hook_state
222 .run(
223 &spec.package,
224 &extension_context,
225 ContextEvent::BeforeModel { turn },
226 &spec.execution_id,
227 turn,
228 &spec.user_prompt,
229 &visible_tools,
230 &events,
231 )
232 .await?;
233 let refreshed_items = hook_state.context_items();
234 context.refresh_system(&mut messages, &spec.package.system_prompt, refreshed_items)?;
235
236 fold_repeated_failures(&mut messages);
237 let initial_decision = compaction.decide(&messages, &tools, &spec.user_prompt, turn)?;
238 let unclosed_transaction = has_unclosed_tool_transaction(&messages);
239 if initial_decision.is_some() && unclosed_transaction {
240 events
241 .emit(KernelEvent::Warning {
242 code: "CONTEXT_COMPACTION_DEFERRED".into(),
243 message: "an unclosed ToolUse/ToolResult transaction is resident".into(),
244 })
245 .await?;
246 let measurement = compaction.measurement(&messages, &tools);
247 if measurement.total_tokens > spec.limits.context.max_input_tokens {
248 return finish_failure(
249 &spec,
250 "CONTEXT_LIMIT_EXCEEDED",
251 "context exceeds its hard limit while a Tool transaction is open",
252 )
253 .await;
254 }
255 } else if initial_decision.is_some() {
256 let measurement = compaction.measurement(&messages, &tools);
257 hook_state.observe_model_input(
258 budget_frame(measurement, spec.limits.context.max_input_tokens),
259 normalized_transcript_summary(&messages),
260 );
261 hook_state
262 .run(
263 &spec.package,
264 &extension_context,
265 ContextEvent::BeforeCompaction { turn },
266 &spec.execution_id,
267 turn,
268 &spec.user_prompt,
269 &visible_tools,
270 &events,
271 )
272 .await?;
273 let refreshed_items = hook_state.context_items();
274 context.refresh_system(
275 &mut messages,
276 &spec.package.system_prompt,
277 refreshed_items,
278 )?;
279 fold_repeated_failures(&mut messages);
280 }
281
282 let assembled_snapshot_id = format!("{}-context-turn-{turn}", spec.execution_id);
283 let assembled_measurement = compaction.measurement(&messages, &tools);
284 sessions
285 .context_audit
286 .snapshot(
287 &operation,
288 context_audit_snapshot(
289 assembled_snapshot_id.clone(),
290 ContextAuditStage::Assembled,
291 format!("model request turn {turn} before compaction"),
292 &messages,
293 &tools,
294 assembled_measurement,
295 spec.limits.context.reserved_output_tokens,
296 ),
297 )
298 .await
299 .map_err(port_error)?;
300 if turn == 1 {
301 sessions
302 .context_audit
303 .derivation(
304 &operation,
305 ContextAuditDerivation {
306 id: format!("{}-context-assemble", spec.execution_id),
307 source_snapshot_id: original_snapshot_id.clone(),
308 target_snapshot_id: assembled_snapshot_id.clone(),
309 kind: ContextDerivationKind::Assemble,
310 input_tokens: original_tokens,
311 output_tokens: assembled_measurement.total_tokens,
312 input_items: original_items_count,
313 output_items: messages.len().saturating_add(tools.len()),
314 },
315 )
316 .await
317 .map_err(port_error)?;
318 }
319 let mut request_snapshot_id = assembled_snapshot_id.clone();
320 if !unclosed_transaction
321 && let Some(decision) =
322 compaction.decide(&messages, &tools, &spec.user_prompt, turn)?
323 {
324 let checkpoint_candidate = hook_state.checkpoint_candidate().map(str::to_owned);
325 let outcome = compaction.compact(
326 &mut messages,
327 &tools,
328 &spec.user_prompt,
329 checkpoint_candidate.as_deref(),
330 turn,
331 )?;
332 let measurement = compaction.measurement(&messages, &tools);
333 hook_state.observe_model_input(
334 budget_frame(measurement, spec.limits.context.max_input_tokens),
335 normalized_transcript_summary(&messages),
336 );
337 hook_state
338 .run(
339 &spec.package,
340 &extension_context,
341 ContextEvent::AfterCompaction { turn },
342 &spec.execution_id,
343 turn,
344 &spec.user_prompt,
345 &visible_tools,
346 &events,
347 )
348 .await?;
349 let refreshed_items = hook_state.context_items();
350 context.refresh_system(
351 &mut messages,
352 &spec.package.system_prompt,
353 refreshed_items,
354 )?;
355 let compressed_measurement = compaction.measurement(&messages, &tools);
356 request_snapshot_id =
357 format!("{}-context-turn-{turn}-compressed", spec.execution_id);
358 sessions
359 .context_audit
360 .snapshot(
361 &operation,
362 context_audit_snapshot(
363 request_snapshot_id.clone(),
364 ContextAuditStage::Compressed,
365 format!("model request turn {turn} after checkpoint compaction"),
366 &messages,
367 &tools,
368 compressed_measurement,
369 spec.limits.context.reserved_output_tokens,
370 ),
371 )
372 .await
373 .map_err(port_error)?;
374 sessions
375 .context_audit
376 .derivation(
377 &operation,
378 ContextAuditDerivation {
379 id: format!("{}-context-compact-{turn}", spec.execution_id),
380 source_snapshot_id: assembled_snapshot_id,
381 target_snapshot_id: request_snapshot_id.clone(),
382 kind: ContextDerivationKind::Compress,
383 input_tokens: outcome.before.total_tokens,
384 output_tokens: compressed_measurement.total_tokens,
385 input_items: messages
386 .len()
387 .saturating_add(outcome.removed_messages)
388 .saturating_add(tools.len()),
389 output_items: messages.len().saturating_add(tools.len()),
390 },
391 )
392 .await
393 .map_err(port_error)?;
394 events
395 .emit(KernelEvent::Warning {
396 code: "CONTEXT_COMPACTED".into(),
397 message: format!(
398 "reason={:?}; before={}; after={}; target={}; reclaimable={}; fallback={}",
399 decision.reason,
400 outcome.before.total_tokens,
401 compressed_measurement.total_tokens,
402 decision.target_tokens,
403 decision.reclaimable_tokens,
404 outcome.used_fallback
405 ),
406 })
407 .await?;
408 }
409 let invocation_id = format!("{}-model-turn-{turn}", spec.execution_id);
410 sessions
411 .context_audit
412 .model_invocation(
413 &operation,
414 ModelInvocationAudit {
415 id: invocation_id.clone(),
416 request_snapshot_id: request_snapshot_id.clone(),
417 model: spec.model.clone(),
418 state: ModelInvocationAuditState::Prepared,
419 usage: TokenUsage::default(),
420 error_code: None,
421 },
422 )
423 .await
424 .map_err(port_error)?;
425 events.emit(KernelEvent::ModelStarted { turn }).await?;
426 let response = sessions
427 .model
428 .invoke(
429 &operation,
430 ModelRequest {
431 invocation_id: invocation_id.clone(),
432 model: spec.model.clone(),
433 messages: messages.clone(),
434 tools: tools.clone(),
435 max_output_tokens: Some(spec.limits.context.reserved_output_tokens),
436 },
437 )
438 .await;
439 let response = match response {
440 Ok(response) => response,
441 Err(error) => {
442 let _ = sessions
443 .context_audit
444 .model_invocation(
445 &operation,
446 ModelInvocationAudit {
447 id: invocation_id,
448 request_snapshot_id,
449 model: spec.model.clone(),
450 state: ModelInvocationAuditState::Failed,
451 usage: TokenUsage::default(),
452 error_code: Some(error.code.clone()),
453 },
454 )
455 .await;
456 return Err(port_error(error));
457 }
458 };
459 sessions
460 .context_audit
461 .model_invocation(
462 &operation,
463 ModelInvocationAudit {
464 id: invocation_id,
465 request_snapshot_id,
466 model: spec.model.clone(),
467 state: ModelInvocationAuditState::Succeeded,
468 usage: response.usage.clone(),
469 error_code: None,
470 },
471 )
472 .await
473 .map_err(port_error)?;
474 total_usage.add_assign(&response.usage);
475 events.emit(KernelEvent::ModelCompleted { turn }).await?;
476 observe(
477 &spec,
478 ObserverEvent::ModelCompleted {
479 execution_id: spec.execution_id.clone(),
480 turn,
481 },
482 )
483 .await;
484
485 if !response.finish.is_complete() {
486 return finish_failure(
487 &spec,
488 "MODEL_INCOMPLETE",
489 format!("model stopped with {:?}", response.finish),
490 )
491 .await;
492 }
493 let answer = response
494 .output
495 .iter()
496 .filter_map(|part| match part {
497 ModelContent::Text { text } => Some(text.as_str()),
498 _ => None,
499 })
500 .collect::<Vec<_>>()
501 .join("");
502 let calls = response
503 .output
504 .iter()
505 .filter_map(|part| match part {
506 ModelContent::ToolUse {
507 id,
508 name,
509 arguments,
510 } => Some((id.clone(), name.clone(), arguments.clone())),
511 _ => None,
512 })
513 .collect::<Vec<_>>();
514 messages.push(ModelMessage {
515 role: ModelRole::Assistant,
516 content: response.output.clone(),
517 });
518
519 if calls.is_empty() {
520 if answer.trim().is_empty() {
521 return finish_failure(
522 &spec,
523 "MODEL_EMPTY",
524 "model returned no text or tool calls",
525 )
526 .await;
527 }
528 if let CompletionMode::RequiredTool(name) = &spec.package.completion_mode {
529 return finish_failure(
530 &spec,
531 "COMPLETION_TOOL_REQUIRED",
532 format!("agent must finish with {name}"),
533 )
534 .await;
535 }
536 sessions
537 .conversation
538 .append(
539 &operation,
540 &[ConversationMessage {
541 role: ModelRole::Assistant,
542 text: answer.clone(),
543 }],
544 )
545 .await
546 .map_err(port_error)?;
547 observe(
548 &spec,
549 ObserverEvent::ExecutionFinished {
550 execution_id: spec.execution_id.clone(),
551 succeeded: true,
552 },
553 )
554 .await;
555 return Ok(KernelOutcome {
556 answer,
557 model_turns: turn,
558 tool_calls,
559 usage: total_usage,
560 });
561 }
562
563 for (call_id, name, arguments) in calls {
564 tool_calls += 1;
565 if tool_calls > spec.limits.max_tool_calls {
566 return finish_failure(
567 &spec,
568 "TOOL_BUDGET_EXHAUSTED",
569 "maximum tool calls exceeded",
570 )
571 .await;
572 }
573 let outcome = tool_runtime
574 .execute(&operation, &sessions, &events, &call_id, &name, arguments)
575 .await?;
576 let invalidation = outcome.context_invalidation;
577 messages.push(ModelMessage {
578 role: ModelRole::Tool,
579 content: vec![ModelContent::ToolResult {
580 tool_use_id: call_id,
581 name: name.clone(),
582 content: outcome.content.clone(),
583 failed: outcome.failed,
584 }],
585 });
586 let measurement = compaction.measurement(&messages, &tools);
587 hook_state.observe_model_input(
588 budget_frame(measurement, spec.limits.context.max_input_tokens),
589 normalized_transcript_summary(&messages),
590 );
591 hook_state
592 .run(
593 &spec.package,
594 &extension_context,
595 ContextEvent::AfterTool {
596 turn,
597 tool_name: name,
598 failed: outcome.failed,
599 invalidation,
600 },
601 &spec.execution_id,
602 turn,
603 &spec.user_prompt,
604 &visible_tools,
605 &events,
606 )
607 .await?;
608 if let Some(answer) = outcome.completion {
609 sessions
610 .conversation
611 .append(
612 &operation,
613 &[ConversationMessage {
614 role: ModelRole::Assistant,
615 text: answer.clone(),
616 }],
617 )
618 .await
619 .map_err(port_error)?;
620 observe(
621 &spec,
622 ObserverEvent::ExecutionFinished {
623 execution_id: spec.execution_id.clone(),
624 succeeded: true,
625 },
626 )
627 .await;
628 return Ok(KernelOutcome {
629 answer,
630 model_turns: turn,
631 tool_calls,
632 usage: total_usage,
633 });
634 }
635 }
636 }
637 finish_failure(
638 &spec,
639 "MODEL_TURN_BUDGET_EXHAUSTED",
640 "maximum model turns exceeded",
641 )
642 .await
643 }
644}
645
646#[allow(clippy::too_many_arguments)]
647fn context_audit_snapshot(
648 id: String,
649 stage: ContextAuditStage,
650 purpose: String,
651 messages: &[ModelMessage],
652 tools: &[ModelToolDefinition],
653 measurement: ContextMeasurement,
654 max_output_tokens: u32,
655) -> ContextAuditSnapshot {
656 ContextAuditSnapshot {
657 id,
658 stage,
659 purpose,
660 messages: messages.to_vec(),
661 tools: tools.to_vec(),
662 protocol_overhead_tokens: measurement.protocol_overhead_tokens,
663 reserved_output_tokens: measurement.reserved_output_tokens,
664 estimated_tokens: measurement.total_tokens,
665 request_digest: model_request_digest(messages, tools, Some(max_output_tokens)),
666 }
667}
668
669fn budget_frame(measurement: ContextMeasurement, max_input_tokens: u64) -> ContextBudgetFrame {
670 ContextBudgetFrame {
671 max_input_tokens,
672 message_tokens: measurement.message_tokens,
673 tool_schema_tokens: measurement.tool_schema_tokens,
674 protocol_overhead_tokens: measurement.protocol_overhead_tokens,
675 reserved_output_tokens: measurement.reserved_output_tokens,
676 total_tokens: measurement.total_tokens,
677 }
678}
679
680fn estimated_text_tokens(content: &str) -> u64 {
681 (content.len() as u64).saturating_add(3) / 4
682}
683
684fn content_digest(content: &str) -> String {
685 format!("sha256:{:x}", Sha256::digest(content.as_bytes()))
686}
687
688fn normalized_transcript_summary(messages: &[ModelMessage]) -> String {
691 messages
692 .iter()
693 .enumerate()
694 .rev()
695 .take(12)
696 .map(|(index, message)| {
697 let content = message
698 .content
699 .iter()
700 .map(|content| match content {
701 ModelContent::Text { text } => format!("text(chars={})", text.chars().count()),
702 ModelContent::ToolUse { name, .. } => format!("tool_use(name={name})"),
703 ModelContent::ToolResult {
704 name,
705 content,
706 failed,
707 ..
708 } => format!(
709 "tool_result(name={name},failed={failed},chars={})",
710 content.chars().count()
711 ),
712 })
713 .collect::<Vec<_>>()
714 .join(",");
715 format!("message[{index}] role={:?} {content}", message.role)
716 })
717 .collect::<Vec<_>>()
718 .into_iter()
719 .rev()
720 .collect::<Vec<_>>()
721 .join("\n")
722}
723
724fn raw_context_candidates(
725 system_prompt: &str,
726 context_items: &[runtime_extension_api::ContextItem],
727 history: &[ConversationMessage],
728 user_prompt: &str,
729) -> Vec<ModelMessage> {
730 let mut system = system_prompt.trim().to_string();
731 for item in context_items {
732 if !item.content.trim().is_empty() {
733 system.push_str("\n\n<context source=\"");
734 system.push_str(&item.source);
735 system.push_str("\">\n");
736 system.push_str(item.content.trim());
737 system.push_str("\n</context>");
738 }
739 }
740 let mut messages = Vec::with_capacity(history.len() + 2);
741 messages.push(ModelMessage::text(ModelRole::System, system));
742 messages.extend(
743 history
744 .iter()
745 .map(|message| ModelMessage::text(message.role, message.text.clone())),
746 );
747 messages.push(ModelMessage::text(ModelRole::User, user_prompt));
748 messages
749}
750
751fn validate_spec(spec: &KernelSpec) -> Result<(), KernelFailure> {
752 if spec.user_prompt.trim().is_empty() {
753 return Err(KernelFailure::new(
754 "PROMPT_EMPTY",
755 "user prompt must not be empty",
756 ));
757 }
758 if spec.model.trim().is_empty() {
759 return Err(KernelFailure::new(
760 "MODEL_EMPTY",
761 "model binding must not be empty",
762 ));
763 }
764 if spec.limits.max_model_turns == 0
765 || spec.limits.max_tool_calls == 0
766 || spec.limits.max_tool_output_bytes == 0
767 || spec.limits.context.max_messages < 2
768 || spec.limits.context.max_input_tokens < 1_024
769 || u64::from(spec.limits.context.reserved_output_tokens)
770 >= spec.limits.context.max_input_tokens
771 || spec.limits.context.compaction.target_percent
772 >= spec.limits.context.compaction.trigger_percent
773 || spec.limits.context.compaction.trigger_percent
774 >= spec.limits.context.compaction.emergency_percent
775 || spec.limits.context.compaction.emergency_percent > 100
776 {
777 return Err(KernelFailure::new(
778 "KERNEL_LIMIT_INVALID",
779 "kernel limits must be positive",
780 ));
781 }
782 Ok(())
783}
784
785async fn observe(spec: &KernelSpec, event: ObserverEvent) {
786 for observer in &spec.package.observers {
787 let _ = tokio::time::timeout(Duration::from_secs(2), observer.observe(event.clone())).await;
788 }
789}
790
791async fn finish_failure<T>(
792 spec: &KernelSpec,
793 code: &str,
794 message: impl Into<String>,
795) -> Result<T, KernelFailure> {
796 observe(
797 spec,
798 ObserverEvent::ExecutionFinished {
799 execution_id: spec.execution_id.clone(),
800 succeeded: false,
801 },
802 )
803 .await;
804 Err(KernelFailure::new(code, message))
805}
806
807fn port_error(error: runtime_ports::PortFailure) -> KernelFailure {
808 KernelFailure {
809 code: error.code,
810 message: error.message,
811 retryable: error.retryable,
812 commit: error.commit,
813 }
814}
815
816#[cfg(test)]
817mod tests {
818 use std::sync::Arc;
819 use std::time::{Duration, Instant};
820
821 use agent_runtime_testkit::{
822 CompleteTool, MemoryConversationSession, RecordingContextAuditSession,
823 RecordingKernelEvents, ScriptedModelSession, fixture_package, sessions,
824 sessions_with_context_audit,
825 };
826 use async_trait::async_trait;
827 use runtime_extension_api::{
828 AgentPackage, CompletionMode, ContextFrame, ContextHook, ContextInvalidation, ContextPatch,
829 ContextResource, ContextSensitivity, ContextSlot, ExtensionError, SideEffect, Tool,
830 ToolContext, ToolDescriptor, ToolOutcome,
831 };
832 use runtime_kernel_api::{ContextLimits, KernelLimits, RuntimeKernel};
833 use runtime_ports::{
834 ContextAuditStage, ModelContent, ModelFinish, ModelInvocationAuditState, ModelResponse,
835 OperationContext, TokenUsage,
836 };
837 use runtime_types::{ExecutionId, OperationId};
838 use serde_json::json;
839 use tokio_util::sync::CancellationToken;
840
841 use super::*;
842
843 fn operation() -> OperationContext {
844 OperationContext {
845 id: OperationId::random(),
846 execution_id: ExecutionId::random(),
847 deadline: Instant::now() + Duration::from_secs(10),
848 cancellation: CancellationToken::new(),
849 }
850 }
851
852 fn spec(operation: &OperationContext) -> KernelSpec {
853 KernelSpec {
854 execution_id: operation.execution_id.clone(),
855 conversation_id: None,
856 user_prompt: "finish the task".into(),
857 model: "fixture-model".into(),
858 package: fixture_package(CompletionMode::RequiredTool("runtime.complete".into())),
859 granted_capabilities: std::collections::BTreeSet::new(),
860 limits: KernelLimits {
861 max_model_turns: 4,
862 max_tool_calls: 4,
863 context: ContextLimits {
864 max_messages: 16,
865 max_input_tokens: 16 * 1024,
866 reserved_output_tokens: 1_024,
867 ..ContextLimits::default()
868 },
869 max_tool_output_bytes: 4096,
870 },
871 }
872 }
873
874 struct RefreshContextTool;
875
876 #[async_trait]
877 impl Tool for RefreshContextTool {
878 fn descriptor(&self) -> ToolDescriptor {
879 ToolDescriptor {
880 name: "context.refresh".into(),
881 description: "Refresh dynamic memory context.".into(),
882 input_schema: json!({"type":"object","additionalProperties":false}),
883 required_capabilities: Vec::new(),
884 side_effect: SideEffect::ReadOnly,
885 }
886 }
887
888 async fn invoke(
889 &self,
890 _: ToolContext,
891 _: serde_json::Value,
892 ) -> Result<ToolOutcome, ExtensionError> {
893 Ok(ToolOutcome::success("refreshed")
894 .with_context_invalidation(ContextInvalidation::MEMORY))
895 }
896 }
897
898 struct RefreshMemoryHook;
899
900 #[async_trait]
901 impl ContextHook for RefreshMemoryHook {
902 fn name(&self) -> &str {
903 "test.refresh-memory"
904 }
905
906 async fn apply(
907 &self,
908 _: &ExtensionContext,
909 event: &ContextEvent,
910 frame: &ContextFrame,
911 ) -> Result<ContextPatch, ExtensionError> {
912 if matches!(event, ContextEvent::BeforeModel { .. })
913 && frame
914 .pending_invalidations
915 .contains(ContextInvalidation::MEMORY)
916 {
917 return Ok(ContextPatch {
918 resources: vec![ContextResource {
919 id: "working-memory".into(),
920 slot: ContextSlot::Memory,
921 source: "test-hook".into(),
922 content: "dynamic-memory-v2".into(),
923 required: false,
924 sensitivity: ContextSensitivity::Internal,
925 version: "2".into(),
926 digest: "dynamic-memory-v2".into(),
927 estimated_tokens: 8,
928 }],
929 removals: Vec::new(),
930 invalidations: ContextInvalidation::NONE,
931 });
932 }
933 Ok(ContextPatch::default())
934 }
935 }
936
937 #[tokio::test]
938 async fn model_tool_completion_closes_the_loop() {
939 let model = Arc::new(ScriptedModelSession::new(vec![ModelResponse {
940 output: vec![ModelContent::ToolUse {
941 id: "call-1".into(),
942 name: "runtime.complete".into(),
943 arguments: json!({"answer":"done"}),
944 }],
945 finish: ModelFinish::ToolCalls,
946 usage: TokenUsage {
947 input_tokens: Some(10),
948 output_tokens: Some(2),
949 },
950 }]));
951 let operation = operation();
952 let outcome = AgentKernel
953 .execute(
954 operation.clone(),
955 spec(&operation),
956 sessions(model.clone()),
957 Arc::new(RecordingKernelEvents::default()),
958 )
959 .await
960 .unwrap();
961 assert_eq!(outcome.answer, "done");
962 assert_eq!(outcome.tool_calls, 1);
963 assert_eq!(model.requests().len(), 1);
964 }
965
966 #[tokio::test]
967 async fn tool_invalidation_refreshes_context_before_the_next_model_turn() {
968 let model = Arc::new(ScriptedModelSession::new(vec![
969 ModelResponse {
970 output: vec![ModelContent::ToolUse {
971 id: "call-refresh".into(),
972 name: "context.refresh".into(),
973 arguments: json!({}),
974 }],
975 finish: ModelFinish::ToolCalls,
976 usage: TokenUsage::default(),
977 },
978 ModelResponse {
979 output: vec![ModelContent::ToolUse {
980 id: "call-complete".into(),
981 name: "runtime.complete".into(),
982 arguments: json!({"answer":"done"}),
983 }],
984 finish: ModelFinish::ToolCalls,
985 usage: TokenUsage::default(),
986 },
987 ]));
988 let operation = operation();
989 let mut spec = spec(&operation);
990 let fixture = fixture_package(CompletionMode::RequiredTool("runtime.complete".into()));
991 spec.package = Arc::new(AgentPackage {
992 manifest: fixture.manifest.clone(),
993 system_prompt: fixture.system_prompt.clone(),
994 tools: vec![Arc::new(CompleteTool), Arc::new(RefreshContextTool)],
995 context_sources: Vec::new(),
996 context_hooks: vec![Arc::new(RefreshMemoryHook)],
997 guards: Vec::new(),
998 observers: Vec::new(),
999 completion_mode: CompletionMode::RequiredTool("runtime.complete".into()),
1000 });
1001
1002 AgentKernel
1003 .execute(
1004 operation.clone(),
1005 spec,
1006 sessions(model.clone()),
1007 Arc::new(RecordingKernelEvents::default()),
1008 )
1009 .await
1010 .unwrap();
1011
1012 let requests = model.requests();
1013 assert_eq!(requests.len(), 2);
1014 let second_system = serde_json::to_string(&requests[1].messages[0]).unwrap();
1015 assert!(second_system.contains("dynamic-memory-v2"));
1016 }
1017
1018 #[tokio::test]
1019 async fn records_exact_model_context_and_invocation_lifecycle() {
1020 let model = Arc::new(ScriptedModelSession::new(vec![ModelResponse {
1021 output: vec![ModelContent::ToolUse {
1022 id: "call-1".into(),
1023 name: "runtime.complete".into(),
1024 arguments: json!({"answer":"done"}),
1025 }],
1026 finish: ModelFinish::ToolCalls,
1027 usage: TokenUsage {
1028 input_tokens: Some(11),
1029 output_tokens: Some(3),
1030 },
1031 }]));
1032 let audit = Arc::new(RecordingContextAuditSession::default());
1033 let operation = operation();
1034 AgentKernel
1035 .execute(
1036 operation.clone(),
1037 spec(&operation),
1038 sessions_with_context_audit(model.clone(), audit.clone()),
1039 Arc::new(RecordingKernelEvents::default()),
1040 )
1041 .await
1042 .unwrap();
1043
1044 let snapshots = audit.snapshots.lock().unwrap();
1045 assert_eq!(snapshots.len(), 2);
1046 assert_eq!(snapshots[0].stage, ContextAuditStage::Original);
1047 assert_eq!(snapshots[1].stage, ContextAuditStage::Assembled);
1048 assert_eq!(snapshots[1].messages.len(), 2);
1049 let request = &model.requests()[0];
1050 assert_eq!(snapshots[1].tools, request.tools);
1051 assert_eq!(
1052 snapshots[1].reserved_output_tokens,
1053 u64::from(request.max_output_tokens.unwrap())
1054 );
1055 assert_eq!(
1056 snapshots[1].request_digest,
1057 model_request_digest(&request.messages, &request.tools, request.max_output_tokens)
1058 );
1059 drop(snapshots);
1060
1061 let derivations = audit.derivations.lock().unwrap();
1062 assert_eq!(derivations.len(), 1);
1063 assert_eq!(
1064 derivations[0].source_snapshot_id,
1065 format!("{}-context-original", operation.execution_id)
1066 );
1067 drop(derivations);
1068
1069 let invocations = audit.invocations.lock().unwrap();
1070 assert_eq!(invocations.len(), 2);
1071 assert_eq!(invocations[0].state, ModelInvocationAuditState::Prepared);
1072 assert_eq!(invocations[1].state, ModelInvocationAuditState::Succeeded);
1073 assert_eq!(invocations[1].usage.input_tokens, Some(11));
1074 }
1075
1076 struct RecordingCompactionHook {
1077 events: std::sync::Mutex<Vec<&'static str>>,
1078 }
1079
1080 #[async_trait]
1081 impl ContextHook for RecordingCompactionHook {
1082 fn name(&self) -> &str {
1083 "test.compaction-lifecycle"
1084 }
1085
1086 async fn apply(
1087 &self,
1088 _: &ExtensionContext,
1089 event: &ContextEvent,
1090 _: &ContextFrame,
1091 ) -> Result<ContextPatch, ExtensionError> {
1092 let name = match event {
1093 ContextEvent::BeforeCompaction { .. } => Some("before"),
1094 ContextEvent::AfterCompaction { .. } => Some("after"),
1095 _ => None,
1096 };
1097 if let Some(name) = name {
1098 self.events.lock().unwrap().push(name);
1099 }
1100 Ok(ContextPatch::default())
1101 }
1102 }
1103
1104 #[tokio::test]
1105 async fn compaction_preserves_resident_set_and_audits_lifecycle() {
1106 let model = Arc::new(ScriptedModelSession::new(vec![ModelResponse {
1107 output: vec![ModelContent::ToolUse {
1108 id: "call-complete".into(),
1109 name: "runtime.complete".into(),
1110 arguments: json!({"answer":"done"}),
1111 }],
1112 finish: ModelFinish::ToolCalls,
1113 usage: TokenUsage::default(),
1114 }]));
1115 let hook = Arc::new(RecordingCompactionHook {
1116 events: std::sync::Mutex::new(Vec::new()),
1117 });
1118 let operation = operation();
1119 let mut spec = spec(&operation);
1120 spec.limits.context.max_input_tokens = 2_000;
1121 spec.limits.context.reserved_output_tokens = 100;
1122 spec.limits.context.protocol_overhead_tokens = 50;
1123 spec.limits.context.compaction.trigger_percent = 50;
1124 spec.limits.context.compaction.target_percent = 30;
1125 spec.limits.context.compaction.emergency_percent = 90;
1126 let fixture = fixture_package(CompletionMode::RequiredTool("runtime.complete".into()));
1127 spec.package = Arc::new(AgentPackage {
1128 manifest: fixture.manifest.clone(),
1129 system_prompt: fixture.system_prompt.clone(),
1130 tools: fixture.tools.clone(),
1131 context_sources: fixture.context_sources.clone(),
1132 context_hooks: vec![hook.clone()],
1133 guards: fixture.guards.clone(),
1134 observers: fixture.observers.clone(),
1135 completion_mode: fixture.completion_mode.clone(),
1136 });
1137
1138 let audit = Arc::new(RecordingContextAuditSession::default());
1139 let mut execution_sessions = sessions_with_context_audit(model.clone(), audit.clone());
1140 execution_sessions.conversation = Arc::new(MemoryConversationSession::with_messages(vec![
1141 ConversationMessage {
1142 role: ModelRole::User,
1143 text: "old investigation".repeat(200),
1144 },
1145 ConversationMessage {
1146 role: ModelRole::Assistant,
1147 text: "old evidence".repeat(200),
1148 },
1149 ]));
1150
1151 AgentKernel
1152 .execute(
1153 operation,
1154 spec,
1155 execution_sessions,
1156 Arc::new(RecordingKernelEvents::default()),
1157 )
1158 .await
1159 .unwrap();
1160
1161 assert_eq!(*hook.events.lock().unwrap(), vec!["before", "after"]);
1162 let requests = model.requests();
1163 assert_eq!(requests.len(), 1);
1164 assert_eq!(requests[0].messages.len(), 3);
1165 let encoded = serde_json::to_string(&requests[0].messages).unwrap();
1166 assert!(encoded.contains("finish the task"));
1167 assert!(encoded.contains("changed_unverified") || encoded.contains("investigating"));
1168
1169 let snapshots = audit.snapshots.lock().unwrap();
1170 assert_eq!(
1171 snapshots.last().unwrap().stage,
1172 ContextAuditStage::Compressed
1173 );
1174 assert_eq!(
1175 snapshots.last().unwrap().request_digest,
1176 model_request_digest(
1177 &requests[0].messages,
1178 &requests[0].tools,
1179 requests[0].max_output_tokens,
1180 )
1181 );
1182 let derivations = audit.derivations.lock().unwrap();
1183 assert!(
1184 derivations
1185 .iter()
1186 .any(|derivation| derivation.kind == ContextDerivationKind::Compress)
1187 );
1188 }
1189
1190 #[tokio::test]
1191 async fn required_completion_rejects_plain_model_stop() {
1192 let model = Arc::new(ScriptedModelSession::new(vec![ModelResponse {
1193 output: vec![ModelContent::Text {
1194 text: "premature".into(),
1195 }],
1196 finish: ModelFinish::Stop,
1197 usage: TokenUsage::default(),
1198 }]));
1199 let operation = operation();
1200 let failure = AgentKernel
1201 .execute(
1202 operation.clone(),
1203 spec(&operation),
1204 sessions(model),
1205 Arc::new(RecordingKernelEvents::default()),
1206 )
1207 .await
1208 .unwrap_err();
1209 assert_eq!(failure.code, "COMPLETION_TOOL_REQUIRED");
1210 }
1211}