1use std::sync::Arc;
15
16use leviath_core::{EntryKind, Region, RegionEntry};
17use leviath_scripting::region_hook::{RegionScript, run_on_overflow, run_on_write, run_render};
18
19#[derive(Debug, Clone, Default)]
23pub struct AssembleMeta {
24 pub stage_name: String,
26 pub stage_iterations: usize,
28 pub model: String,
30 pub previous_system_hash: Option<u64>,
38}
39
40pub(crate) enum OnWriteOutcome {
42 Accept(String, usize),
44 Drop,
46}
47
48fn entry_to_json(entry: &RegionEntry) -> serde_json::Value {
52 let mut obj = serde_json::json!({
53 "content": entry.content,
54 "tokens": entry.tokens,
55 "timestamp": entry.timestamp,
56 "key": entry.key,
57 });
58 let (kind, extra) = match &entry.kind {
59 EntryKind::Text => ("text", None),
60 EntryKind::UserMessage => ("user_message", None),
61 EntryKind::AssistantTurn { tool_calls } => (
62 "assistant_turn",
63 Some((
64 "tool_calls",
65 serde_json::to_value(tool_calls).unwrap_or_default(),
66 )),
67 ),
68 EntryKind::ToolResult {
69 tool_call_id,
70 tool_name,
71 is_error,
72 } => {
73 obj["tool_call_id"] = serde_json::json!(tool_call_id);
74 obj["tool_name"] = serde_json::json!(tool_name);
75 obj["is_error"] = serde_json::json!(is_error);
76 ("tool_result", None)
77 }
78 };
79 obj["kind"] = serde_json::json!(kind);
80 if let Some((k, v)) = extra {
81 obj[k] = v;
82 }
83 obj
84}
85
86fn region_to_json(region: &Region) -> serde_json::Value {
88 serde_json::json!({
89 "name": region.name,
90 "budget": region.max_tokens,
91 "current_tokens": region.current_tokens,
92 "entry_count": region.content.len(),
93 })
94}
95
96fn fallback_block(region: &Region) -> leviath_providers::SystemBlock {
100 let text = region
101 .content
102 .iter()
103 .map(|e| e.content.as_str())
104 .collect::<Vec<_>>()
105 .join("\n\n");
106 leviath_providers::SystemBlock {
107 text: format!("[{}]:\n{}", region.name, text),
108 cache_hint: leviath_core::CacheHint::Never,
109 }
110}
111
112pub(crate) struct RegionRender<'a> {
123 pub region: &'a Region,
125 pub script: Option<&'a Arc<RegionScript>>,
127 pub persistent: bool,
129 pub meta: &'a AssembleMeta,
131 pub window_current: usize,
133 pub window_max: usize,
135}
136
137pub(crate) struct RenderSink<'a> {
142 pub system_blocks: &'a mut Vec<leviath_providers::SystemBlock>,
144 pub messages: &'a mut Vec<leviath_providers::Message>,
146}
147
148pub(crate) fn render_custom_region(render: RegionRender<'_>, out: RenderSink<'_>) {
149 let RegionRender {
150 region,
151 script,
152 persistent,
153 meta,
154 window_current,
155 window_max,
156 } = render;
157 let RenderSink {
158 system_blocks,
159 messages,
160 } = out;
161 let Some(script) = script else {
162 if !region.content.is_empty() {
165 tracing::warn!(
166 region = %region.name,
167 "custom region has no compiled script; rendering fallback block"
168 );
169 system_blocks.push(fallback_block(region));
170 }
171 return;
172 };
173
174 let ctx = serde_json::json!({
175 "region": region_to_json(region),
176 "entries": region.content.iter().map(entry_to_json).collect::<Vec<_>>(),
177 "stage_name": meta.stage_name,
178 "stage_iterations": meta.stage_iterations,
179 "model": meta.model,
180 "window": { "total_tokens": window_current, "max_tokens": window_max },
181 });
182
183 let rendered = match run_render(script, ctx) {
184 Ok(value) => value,
185 Err(e) => {
186 tracing::warn!(
187 region = %region.name,
188 script = %script.path,
189 error = %e,
190 "custom region render failed; using fallback block"
191 );
192 if !region.content.is_empty() {
193 system_blocks.push(fallback_block(region));
194 }
195 return;
196 }
197 };
198
199 match parse_render_output(&rendered, persistent) {
200 Ok((blocks, msgs)) => {
201 let emitted_tokens: usize = blocks
202 .iter()
203 .map(|b| leviath_core::estimate_tokens(&b.text))
204 .chain(msgs.iter().map(|m| {
205 match &m.content {
206 leviath_providers::MessageContent::Text(t) => {
207 leviath_core::estimate_tokens(t)
208 }
209 leviath_providers::MessageContent::Blocks(bs) => bs
210 .iter()
211 .map(|b| match b {
212 leviath_providers::ContentBlock::Text { text } => {
213 leviath_core::estimate_tokens(text)
214 }
215 leviath_providers::ContentBlock::ToolUse { input, .. } => {
216 leviath_core::estimate_tokens(&input.to_string())
217 }
218 leviath_providers::ContentBlock::ToolResult { content, .. } => {
219 leviath_core::estimate_tokens(content)
220 }
221 })
222 .sum(),
223 }
224 }))
225 .sum();
226 if emitted_tokens > region.max_tokens {
227 tracing::warn!(
228 region = %region.name,
229 script = %script.path,
230 emitted_tokens,
231 budget = region.max_tokens,
232 "custom region render exceeds its budget; sending anyway \
233 (enable exact_token_counting for a hard guard)"
234 );
235 }
236 system_blocks.extend(blocks);
237 messages.extend(msgs);
238 }
239 Err(reason) => {
240 tracing::warn!(
241 region = %region.name,
242 script = %script.path,
243 reason = %reason,
244 "custom region render returned an invalid shape; using fallback block"
245 );
246 if !region.content.is_empty() {
247 system_blocks.push(fallback_block(region));
248 }
249 }
250 }
251}
252
253fn parse_render_output(
258 value: &serde_json::Value,
259 persistent: bool,
260) -> Result<
261 (
262 Vec<leviath_providers::SystemBlock>,
263 Vec<leviath_providers::Message>,
264 ),
265 String,
266> {
267 let hint = if persistent {
269 leviath_core::CacheHint::Always
270 } else {
271 leviath_core::CacheHint::UntilChanged
272 };
273 let block = |text: &str| leviath_providers::SystemBlock {
274 text: text.to_string(),
275 cache_hint: hint,
276 };
277
278 match value {
279 serde_json::Value::String(s) => {
280 let blocks = if s.is_empty() { vec![] } else { vec![block(s)] };
281 Ok((blocks, vec![]))
282 }
283 serde_json::Value::Object(obj) => {
284 let mut blocks = Vec::new();
285 match obj.get("system") {
286 None | Some(serde_json::Value::Null) => {}
287 Some(serde_json::Value::String(s)) => {
288 if !s.is_empty() {
289 blocks.push(block(s));
290 }
291 }
292 Some(serde_json::Value::Array(items)) => {
293 for item in items {
294 match item {
295 serde_json::Value::String(s) if !s.is_empty() => blocks.push(block(s)),
296 serde_json::Value::String(_) => {}
297 other => {
298 return Err(format!(
299 "system array items must be strings, found {other}"
300 ));
301 }
302 }
303 }
304 }
305 Some(other) => {
306 return Err(format!(
307 "system must be a string or array of strings, found {other}"
308 ));
309 }
310 }
311 let mut messages = Vec::new();
312 match obj.get("messages") {
313 None | Some(serde_json::Value::Null) => {}
314 Some(serde_json::Value::Array(items)) => {
315 for item in items {
316 messages.push(message_from_json(item)?);
317 }
318 }
319 Some(other) => return Err(format!("messages must be an array, found {other}")),
320 }
321 Ok((blocks, messages))
322 }
323 other => Err(format!(
324 "render must return a string or #{{ system, messages }} map, found {other}"
325 )),
326 }
327}
328
329fn message_from_json(value: &serde_json::Value) -> Result<leviath_providers::Message, String> {
337 let obj = value
338 .as_object()
339 .ok_or_else(|| format!("each message must be a map, found {value}"))?;
340 let role = obj
341 .get("role")
342 .and_then(|r| r.as_str())
343 .ok_or("each message needs a role of \"user\" or \"assistant\"")?;
344 if role != "user" && role != "assistant" {
345 return Err(format!(
346 "message role must be user or assistant, found {role}"
347 ));
348 }
349
350 let content_str = match obj.get("content") {
351 None | Some(serde_json::Value::Null) => None,
352 Some(serde_json::Value::String(s)) => Some(s.clone()),
353 Some(other) => return Err(format!("message content must be a string, found {other}")),
354 };
355
356 if let Some(calls) = obj.get("tool_calls") {
357 if role != "assistant" {
358 return Err("tool_calls are only valid on an assistant message".to_string());
359 }
360 let calls = calls
361 .as_array()
362 .ok_or_else(|| format!("tool_calls must be an array, found {calls}"))?;
363 let mut blocks = Vec::new();
364 if let Some(text) = content_str.filter(|s| !s.is_empty()) {
365 blocks.push(leviath_providers::ContentBlock::Text { text });
366 }
367 for call in calls {
368 let call = call
369 .as_object()
370 .ok_or_else(|| format!("each tool_call must be a map, found {call}"))?;
371 let id = call
372 .get("id")
373 .and_then(|v| v.as_str())
374 .ok_or("each tool_call needs a string id")?;
375 let name = call
376 .get("name")
377 .and_then(|v| v.as_str())
378 .ok_or("each tool_call needs a string name")?;
379 blocks.push(leviath_providers::ContentBlock::ToolUse {
380 id: id.to_string(),
381 name: name.to_string(),
382 input: call
383 .get("arguments")
384 .cloned()
385 .unwrap_or(serde_json::Value::Object(Default::default())),
386 thought_signature: call
387 .get("thought_signature")
388 .and_then(|v| v.as_str())
389 .map(String::from),
390 });
391 }
392 return Ok(leviath_providers::Message {
393 role: "assistant".to_string(),
394 content: leviath_providers::MessageContent::Blocks(blocks),
395 cache_breakpoint: false,
396 });
397 }
398
399 if let Some(results) = obj.get("tool_results") {
400 if role != "user" {
401 return Err("tool_results are only valid on a user message".to_string());
402 }
403 let results = results
404 .as_array()
405 .ok_or_else(|| format!("tool_results must be an array, found {results}"))?;
406 let mut blocks = Vec::new();
407 for result in results {
408 let result = result
409 .as_object()
410 .ok_or_else(|| format!("each tool_result must be a map, found {result}"))?;
411 let id = result
412 .get("tool_call_id")
413 .and_then(|v| v.as_str())
414 .ok_or("each tool_result needs a string tool_call_id")?;
415 let content = result
416 .get("content")
417 .and_then(|v| v.as_str())
418 .ok_or("each tool_result needs string content")?;
419 blocks.push(leviath_providers::ContentBlock::ToolResult {
420 tool_use_id: id.to_string(),
421 content: content.to_string(),
422 is_error: result
423 .get("is_error")
424 .and_then(|v| v.as_bool())
425 .unwrap_or(false),
426 });
427 }
428 return Ok(leviath_providers::Message {
429 role: "user".to_string(),
430 content: leviath_providers::MessageContent::Blocks(blocks),
431 cache_breakpoint: false,
432 });
433 }
434
435 let content =
436 content_str.ok_or("a message without tool_calls/tool_results needs string content")?;
437 Ok(leviath_providers::Message {
438 role: role.to_string(),
439 content: content.into(),
440 cache_breakpoint: false,
441 })
442}
443
444pub(crate) fn apply_on_write(
447 script: &RegionScript,
448 region: &Region,
449 content: String,
450 tokens: usize,
451 kind: &EntryKind,
452) -> OnWriteOutcome {
453 let kind_str = match kind {
454 EntryKind::Text => "text",
455 EntryKind::UserMessage => "user_message",
456 EntryKind::AssistantTurn { .. } => "assistant_turn",
457 EntryKind::ToolResult { .. } => "tool_result",
458 };
459 let ctx = serde_json::json!({
460 "region": region_to_json(region),
461 "entry": { "content": content, "kind": kind_str, "tokens": tokens },
462 });
463 match run_on_write(script, ctx) {
464 Ok(serde_json::Value::String(replacement)) => {
465 let tokens = leviath_core::estimate_tokens(&replacement);
466 OnWriteOutcome::Accept(replacement, tokens)
467 }
468 Ok(serde_json::Value::Bool(false)) => OnWriteOutcome::Drop,
469 Ok(serde_json::Value::Bool(true)) | Ok(serde_json::Value::Null) => {
470 OnWriteOutcome::Accept(content, tokens)
471 }
472 Ok(other) => {
473 tracing::warn!(
474 region = %region.name,
475 script = %script.path,
476 returned = %other,
477 "on_write must return a string, true/false, or unit; accepting entry unchanged"
478 );
479 OnWriteOutcome::Accept(content, tokens)
480 }
481 Err(e) => {
482 tracing::warn!(
483 region = %region.name,
484 script = %script.path,
485 error = %e,
486 "on_write failed; accepting entry unchanged"
487 );
488 OnWriteOutcome::Accept(content, tokens)
489 }
490 }
491}
492
493pub(crate) fn apply_overflow(
497 script: &RegionScript,
498 region: &mut Region,
499 needed_tokens: usize,
500) -> usize {
501 let ctx = serde_json::json!({
502 "region": region_to_json(region),
503 "entries": region.content.iter().map(entry_to_json).collect::<Vec<_>>(),
504 "needed_tokens": needed_tokens,
505 });
506 let value = match run_on_overflow(script, ctx) {
507 Ok(v) => v,
508 Err(e) => {
509 tracing::warn!(
510 region = %region.name,
511 script = %script.path,
512 error = %e,
513 "on_overflow failed; falling back to oldest-first eviction"
514 );
515 return 0;
516 }
517 };
518 let Some(indices) = valid_drop_indices(&value, region.content.len()) else {
519 tracing::warn!(
520 region = %region.name,
521 script = %script.path,
522 returned = %value,
523 "on_overflow must return an array of in-range entry indices; \
524 falling back to oldest-first eviction"
525 );
526 return 0;
527 };
528
529 let mut freed = 0;
530 for index in indices.into_iter().rev() {
532 let entry = region.content.remove(index);
533 freed += entry.tokens;
534 }
535 region.current_tokens = region.current_tokens.saturating_sub(freed);
536 freed
537}
538
539fn valid_drop_indices(value: &serde_json::Value, len: usize) -> Option<Vec<usize>> {
542 let items = value.as_array()?;
543 let mut indices = Vec::with_capacity(items.len());
544 for item in items {
545 let index = item.as_u64()? as usize;
546 if index >= len {
547 return None;
548 }
549 indices.push(index);
550 }
551 indices.sort_unstable();
552 indices.dedup();
553 Some(indices)
554}
555
556#[cfg(test)]
557mod tests {
558 use super::*;
559 use crate::test_support::with_tracing;
560 use leviath_core::RegionKind;
561 use leviath_scripting::region_hook::compile;
562 use serde_json::json;
563
564 fn script(src: &str) -> Arc<RegionScript> {
565 Arc::new(compile("test.rhai", src).unwrap())
566 }
567
568 fn region_with(entries: &[(&str, EntryKind)]) -> Region {
569 let mut region = Region::new(
570 "brain".to_string(),
571 RegionKind::Custom {
572 script: "test.rhai".to_string(),
573 persistent: false,
574 },
575 1000,
576 );
577 for (content, kind) in entries {
578 region
579 .add_typed_entry(content.to_string(), 10, kind.clone())
580 .unwrap();
581 }
582 region
583 }
584
585 fn render(
586 region: &Region,
587 script: Option<&Arc<RegionScript>>,
588 persistent: bool,
589 ) -> (
590 Vec<leviath_providers::SystemBlock>,
591 Vec<leviath_providers::Message>,
592 ) {
593 let mut blocks = Vec::new();
594 let mut messages = Vec::new();
595 with_tracing(|| {
596 render_custom_region(
597 RegionRender {
598 region,
599 script,
600 persistent,
601 meta: &AssembleMeta {
602 stage_name: "plan".to_string(),
603 stage_iterations: 2,
604 model: "m1".to_string(),
605 previous_system_hash: None,
606 },
607 window_current: 50,
608 window_max: 2000,
609 },
610 RenderSink {
611 system_blocks: &mut blocks,
612 messages: &mut messages,
613 },
614 )
615 });
616 (blocks, messages)
617 }
618
619 #[test]
622 fn entry_to_json_serializes_all_kinds() {
623 let mut region = region_with(&[
624 ("plain", EntryKind::Text),
625 ("hi", EntryKind::UserMessage),
626 (
627 "calling",
628 EntryKind::AssistantTurn {
629 tool_calls: vec![leviath_core::SerializedToolCall {
630 id: "c1".to_string(),
631 name: "shell".to_string(),
632 arguments: json!({"command": "ls"}),
633 thought_signature: None,
634 }],
635 },
636 ),
637 (
638 "result",
639 EntryKind::ToolResult {
640 tool_call_id: "c1".to_string(),
641 tool_name: "shell".to_string(),
642 is_error: true,
643 },
644 ),
645 ]);
646 region.content[0].key = Some("k".to_string());
647
648 let entries: Vec<_> = region.content.iter().map(entry_to_json).collect();
649 assert_eq!(entries[0]["kind"], json!("text"));
650 assert_eq!(entries[0]["key"], json!("k"));
651 assert_eq!(entries[0]["tokens"], json!(10));
652 assert_eq!(entries[1]["kind"], json!("user_message"));
653 assert_eq!(entries[2]["kind"], json!("assistant_turn"));
654 assert_eq!(entries[2]["tool_calls"][0]["id"], json!("c1"));
655 assert_eq!(entries[3]["kind"], json!("tool_result"));
656 assert_eq!(entries[3]["tool_call_id"], json!("c1"));
657 assert_eq!(entries[3]["is_error"], json!(true));
658 }
659
660 #[test]
663 fn render_string_becomes_one_block_with_persistence_hint() {
664 let region = region_with(&[("x", EntryKind::Text)]);
665 let s = script("fn render(ctx) { `<${ctx.region.name}>` }");
666
667 let (blocks, messages) = render(®ion, Some(&s), false);
668 assert_eq!(blocks.len(), 1);
669 assert_eq!(blocks[0].text, "<brain>");
670 assert_eq!(blocks[0].cache_hint, leviath_core::CacheHint::UntilChanged);
671 assert!(messages.is_empty());
672
673 let (blocks, _) = render(®ion, Some(&s), true);
674 assert_eq!(blocks[0].cache_hint, leviath_core::CacheHint::Always);
675 }
676
677 #[test]
678 fn render_map_emits_system_array_and_typed_messages() {
679 let src = r#"
682 fn render(ctx) {
683 #{
684 system: ["s1", "", "s2"],
685 messages: [
686 #{ role: "user", content: "hello" },
687 #{ role: "assistant", content: "thinking", tool_calls: [
688 #{ id: "c1", name: "shell", arguments: #{ command: "ls" } },
689 ] },
690 #{ role: "user", tool_results: [
691 #{ tool_call_id: "c1", content: "file_a", is_error: false },
692 ] },
693 ],
694 }
695 }
696 "#;
697 let region = region_with(&[("x", EntryKind::Text)]);
698 let (blocks, messages) = render(®ion, Some(&script(src)), false);
699
700 assert_eq!(
701 blocks.iter().map(|b| b.text.as_str()).collect::<Vec<_>>(),
702 vec!["s1", "s2"],
703 "empty system strings are skipped"
704 );
705 assert_eq!(messages.len(), 3);
706 assert_eq!(messages[0].role, "user");
707 let assistant = serde_json::to_value(&messages[1].content).unwrap();
710 assert_eq!(assistant[0], json!({ "type": "text", "text": "thinking" }));
711 assert_eq!(assistant[1]["type"], json!("tool_use"));
712 assert_eq!(assistant[1]["id"], json!("c1"));
713 assert_eq!(assistant[1]["name"], json!("shell"));
714 let results = serde_json::to_value(&messages[2].content).unwrap();
715 assert_eq!(results[0]["type"], json!("tool_result"));
716 assert_eq!(results[0]["tool_use_id"], json!("c1"));
717 assert_eq!(results[0]["content"], json!("file_a"));
718 assert_eq!(results[0]["is_error"], json!(false));
719 }
720
721 #[test]
722 fn render_map_accepts_single_system_string_and_null_fields() {
723 let src = r#"fn render(ctx) { #{ system: "solo", messages: () } }"#;
724 let region = region_with(&[("x", EntryKind::Text)]);
725 let (blocks, messages) = render(®ion, Some(&script(src)), false);
726 assert_eq!(blocks.len(), 1);
727 assert_eq!(blocks[0].text, "solo");
728 assert!(messages.is_empty());
729 }
730
731 #[test]
732 fn render_empty_map_and_empty_string_emit_nothing() {
733 let region = region_with(&[("x", EntryKind::Text)]);
734 for src in ["fn render(ctx) { #{} }", "fn render(ctx) { \"\" }"] {
735 let (blocks, messages) = render(®ion, Some(&script(src)), false);
736 assert!(blocks.is_empty(), "src: {src}");
737 assert!(messages.is_empty());
738 }
739 }
740
741 #[test]
742 fn render_sees_stage_meta_and_window_fields() {
743 let src = r#"
744 fn render(ctx) {
745 `${ctx.stage_name}|${ctx.stage_iterations}|${ctx.model}|${ctx.window.total_tokens}|${ctx.window.max_tokens}`
746 }
747 "#;
748 let region = region_with(&[("x", EntryKind::Text)]);
749 let (blocks, _) = render(®ion, Some(&script(src)), false);
750 assert_eq!(blocks[0].text, "plan|2|m1|50|2000");
751 }
752
753 #[test]
754 fn render_over_budget_warns_but_still_emits() {
755 let src = r#"fn render(ctx) { let s = "x"; s.pad(8000, 'x'); s }"#;
758 let region = region_with(&[("x", EntryKind::Text)]);
759 let (blocks, _) = render(®ion, Some(&script(src)), false);
760 assert_eq!(blocks.len(), 1);
761 assert_eq!(blocks[0].text.len(), 8000);
762 }
763
764 #[test]
767 fn render_missing_script_falls_back_to_temporary_style() {
768 let region = region_with(&[("a", EntryKind::Text), ("b", EntryKind::Text)]);
769 let (blocks, messages) = render(®ion, None, false);
770 assert_eq!(blocks.len(), 1);
771 assert_eq!(blocks[0].text, "[brain]:\na\n\nb");
772 assert_eq!(blocks[0].cache_hint, leviath_core::CacheHint::Never);
773 assert!(messages.is_empty());
774 }
775
776 #[test]
777 fn render_missing_script_on_empty_region_emits_nothing() {
778 let region = region_with(&[]);
779 let (blocks, messages) = render(®ion, None, false);
780 assert!(blocks.is_empty());
781 assert!(messages.is_empty());
782 }
783
784 #[test]
785 fn render_runtime_error_falls_back() {
786 let region = region_with(&[("kept", EntryKind::Text)]);
787 let s = script("fn render(ctx) { throw \"broken\" }");
788 let (blocks, _) = render(®ion, Some(&s), false);
789 assert_eq!(blocks.len(), 1);
790 assert_eq!(blocks[0].text, "[brain]:\nkept");
791 assert_eq!(blocks[0].cache_hint, leviath_core::CacheHint::Never);
792 }
793
794 #[test]
795 fn render_error_on_empty_region_emits_nothing() {
796 let region = region_with(&[]);
797 let s = script("fn render(ctx) { throw \"broken\" }");
798 let (blocks, _) = render(®ion, Some(&s), false);
799 assert!(blocks.is_empty());
800 }
801
802 #[test]
803 fn render_invalid_shapes_fall_back() {
804 let region = region_with(&[("kept", EntryKind::Text)]);
805 for src in [
806 "fn render(ctx) { 42 }",
807 "fn render(ctx) { true }",
808 "fn render(ctx) { [1, 2] }",
809 "fn render(ctx) { }",
810 "fn render(ctx) { #{ system: 42 } }",
811 "fn render(ctx) { #{ system: [1] } }",
812 "fn render(ctx) { #{ messages: \"not an array\" } }",
813 "fn render(ctx) { #{ messages: [42] } }",
814 "fn render(ctx) { #{ messages: [#{ content: \"no role\" }] } }",
815 "fn render(ctx) { #{ messages: [#{ role: \"system\", content: \"bad role\" }] } }",
816 "fn render(ctx) { #{ messages: [#{ role: \"user\", content: 42 }] } }",
817 "fn render(ctx) { #{ messages: [#{ role: \"user\" }] } }",
818 "fn render(ctx) { #{ messages: [#{ role: \"user\", tool_calls: [] }] } }",
819 "fn render(ctx) { #{ messages: [#{ role: \"assistant\", tool_calls: 42 }] } }",
820 "fn render(ctx) { #{ messages: [#{ role: \"assistant\", tool_calls: [42] }] } }",
821 "fn render(ctx) { #{ messages: [#{ role: \"assistant\", tool_calls: [#{ name: \"n\" }] }] } }",
822 "fn render(ctx) { #{ messages: [#{ role: \"assistant\", tool_calls: [#{ id: \"i\" }] }] } }",
823 "fn render(ctx) { #{ messages: [#{ role: \"assistant\", tool_results: [] }] } }",
824 "fn render(ctx) { #{ messages: [#{ role: \"user\", tool_results: 42 }] } }",
825 "fn render(ctx) { #{ messages: [#{ role: \"user\", tool_results: [42] }] } }",
826 "fn render(ctx) { #{ messages: [#{ role: \"user\", tool_results: [#{ content: \"c\" }] }] } }",
827 "fn render(ctx) { #{ messages: [#{ role: \"user\", tool_results: [#{ tool_call_id: \"i\" }] }] } }",
828 ] {
829 let (blocks, messages) = render(®ion, Some(&script(src)), false);
830 assert_eq!(blocks.len(), 1, "src must fall back: {src}");
831 assert_eq!(blocks[0].text, "[brain]:\nkept", "src: {src}");
832 assert!(messages.is_empty(), "src: {src}");
833 }
834 }
835
836 #[test]
837 fn render_invalid_shape_on_empty_region_emits_nothing() {
838 let region = region_with(&[]);
841 let (blocks, messages) = render(®ion, Some(&script("fn render(ctx) { 42 }")), false);
842 assert!(blocks.is_empty());
843 assert!(messages.is_empty());
844 }
845
846 #[test]
847 fn render_empty_single_system_string_is_skipped() {
848 let src = r#"fn render(ctx) { #{ system: "" } }"#;
849 let region = region_with(&[("x", EntryKind::Text)]);
850 let (blocks, messages) = render(®ion, Some(&script(src)), false);
851 assert!(blocks.is_empty());
852 assert!(messages.is_empty());
853 }
854
855 #[test]
856 fn render_tool_call_passes_thought_signature_through() {
857 let src = r#"
858 fn render(ctx) {
859 #{ messages: [#{ role: "assistant", tool_calls: [
860 #{ id: "c", name: "n", thought_signature: "sig123" },
861 ] }] }
862 }
863 "#;
864 let region = region_with(&[("x", EntryKind::Text)]);
865 let (_, messages) = render(®ion, Some(&script(src)), false);
866 let blocks = serde_json::to_value(&messages[0].content).unwrap();
867 assert_eq!(blocks[0]["thought_signature"], json!("sig123"));
868 }
869
870 #[test]
871 fn on_write_ctx_reports_every_entry_kind() {
872 let src = r#"
874 fn render(ctx) { "" }
875 fn on_write(ctx) { ctx.entry.kind }
876 "#;
877 for (kind, expected) in [
878 (EntryKind::Text, "text"),
879 (EntryKind::UserMessage, "user_message"),
880 (
881 EntryKind::AssistantTurn { tool_calls: vec![] },
882 "assistant_turn",
883 ),
884 (
885 EntryKind::ToolResult {
886 tool_call_id: "c".to_string(),
887 tool_name: "t".to_string(),
888 is_error: false,
889 },
890 "tool_result",
891 ),
892 ] {
893 let replaced = on_write_kind(src, "x", &kind);
894 assert_eq!(
895 replaced.map(|(content, _)| content),
896 Some(expected.to_string())
897 );
898 }
899 }
900
901 #[test]
902 fn render_assistant_tool_call_defaults_arguments_and_signature() {
903 let src = r#"
904 fn render(ctx) {
905 #{ messages: [#{ role: "assistant", tool_calls: [#{ id: "c", name: "n" }] }] }
906 }
907 "#;
908 let region = region_with(&[("x", EntryKind::Text)]);
909 let (_, messages) = render(®ion, Some(&script(src)), false);
910 let blocks = serde_json::to_value(&messages[0].content).unwrap();
911 assert_eq!(blocks[0]["type"], json!("tool_use"));
912 assert_eq!(blocks[0]["input"], json!({}));
913 assert_eq!(blocks[0]["thought_signature"], serde_json::Value::Null);
916 }
917
918 fn on_write_kind(src: &str, content: &str, kind: &EntryKind) -> Option<(String, usize)> {
924 let region = region_with(&[]);
925 let outcome =
926 with_tracing(|| apply_on_write(&script(src), ®ion, content.to_string(), 5, kind));
927 match outcome {
928 OnWriteOutcome::Accept(content, tokens) => Some((content, tokens)),
929 OnWriteOutcome::Drop => None,
930 }
931 }
932
933 fn on_write_of(src: &str, content: &str) -> Option<(String, usize)> {
934 on_write_kind(src, content, &EntryKind::Text)
935 }
936
937 #[test]
938 fn on_write_replaces_accepts_and_drops() {
939 let replaced = on_write_of(
940 "fn render(ctx) { \"\" }\nfn on_write(ctx) { ctx.entry.content.to_upper() }",
941 "hi",
942 );
943 assert_eq!(
944 replaced,
945 Some(("HI".to_string(), leviath_core::estimate_tokens("HI")))
946 );
947
948 for accept_body in ["true", ""] {
949 let src = format!("fn render(ctx) {{ \"\" }}\nfn on_write(ctx) {{ {accept_body} }}");
950 assert_eq!(
951 on_write_of(&src, "orig"),
952 Some(("orig".to_string(), 5)),
953 "body {accept_body:?} accepts unchanged with original tokens"
954 );
955 }
956
957 assert_eq!(
958 on_write_of("fn render(ctx) { \"\" }\nfn on_write(ctx) { false }", "x"),
959 None,
960 "false drops the entry"
961 );
962 }
963
964 #[test]
965 fn on_write_invalid_return_and_error_accept_unchanged() {
966 for src in [
967 "fn render(ctx) { \"\" }\nfn on_write(ctx) { 42 }",
968 "fn render(ctx) { \"\" }\nfn on_write(ctx) { throw \"bad\" }",
969 ] {
970 assert_eq!(
971 on_write_of(src, "keep"),
972 Some(("keep".to_string(), 5)),
973 "src: {src}"
974 );
975 }
976 }
977
978 #[test]
981 fn apply_overflow_drops_chosen_indices() {
982 let mut region = region_with(&[
983 ("a", EntryKind::Text),
984 ("b", EntryKind::Text),
985 ("c", EntryKind::Text),
986 ]);
987 let s = script("fn render(ctx) { \"\" }\nfn on_overflow(ctx) { [2, 0, 2] }");
989 let freed = with_tracing(|| apply_overflow(&s, &mut region, 15));
990 assert_eq!(freed, 20);
991 assert_eq!(region.content.len(), 1);
992 assert_eq!(region.content[0].content, "b");
993 assert_eq!(region.current_tokens, 10);
994 }
995
996 #[test]
997 fn apply_overflow_error_and_invalid_shapes_free_nothing() {
998 for src in [
999 "fn render(ctx) { \"\" }\nfn on_overflow(ctx) { throw \"bad\" }",
1000 "fn render(ctx) { \"\" }\nfn on_overflow(ctx) { \"not an array\" }",
1001 "fn render(ctx) { \"\" }\nfn on_overflow(ctx) { [\"x\"] }",
1002 "fn render(ctx) { \"\" }\nfn on_overflow(ctx) { [99] }",
1003 ] {
1004 let mut region = region_with(&[("a", EntryKind::Text)]);
1005 let freed = with_tracing(|| apply_overflow(&script(src), &mut region, 5));
1006 assert_eq!(freed, 0, "src: {src}");
1007 assert_eq!(region.content.len(), 1, "content untouched: {src}");
1008 }
1009 }
1010
1011 #[test]
1012 fn overflow_ctx_carries_needed_tokens_and_entries() {
1013 let src = r#"
1014 fn render(ctx) { "" }
1015 fn on_overflow(ctx) {
1016 if ctx.needed_tokens == 7 && ctx.entries.len() == 2 { [0] } else { [] }
1017 }
1018 "#;
1019 let mut region = region_with(&[("a", EntryKind::Text), ("b", EntryKind::Text)]);
1020 let freed = with_tracing(|| apply_overflow(&script(src), &mut region, 7));
1021 assert_eq!(freed, 10);
1022 }
1023}