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