Skip to main content

leviath_runtime/
custom_region.rs

1//! Runtime side of script-backed custom regions (`RegionKind::Custom`).
2//!
3//! The scripting layer ([`leviath_scripting::region_hook`]) compiles and runs
4//! the hooks; this module owns everything on the runtime side of that JSON
5//! boundary: building the `ctx` objects a hook receives, interpreting each
6//! hook's returned value, and every fallback. The contract is that a hook
7//! failure can never fail an inference or lose a write:
8//!
9//! - `render` failure → the region renders as a Temporary-style block
10//!   (`[{name}]:\n…`, `Never` cache hint) and a warning names the script.
11//! - `on_write` failure → the entry is accepted unchanged.
12//! - `on_overflow` failure or invalid indices → oldest-first eviction.
13
14use 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/// Stage-level metadata threaded into `render(ctx)`. `Default` (empty name,
20/// zero iterations, empty model) is used by callers with no stage context -
21/// the transition-choice request and plain `assemble()` in tests.
22#[derive(Debug, Clone, Default)]
23pub struct AssembleMeta {
24    /// Current stage name (`AgentState::current_stage`).
25    pub stage_name: String,
26    /// Inference count within the current stage (`StageProgress::iterations`).
27    pub stage_iterations: usize,
28    /// Model id serving this stage.
29    pub model: String,
30    /// Hash of the system prefix the *previous* request sent, when there was
31    /// one.
32    ///
33    /// Only the message breakpoint reads it, and only to decide whether writing
34    /// that cache entry can pay for itself. `None` means nobody is tracking (a
35    /// first request, or a caller with no state to keep), and the breakpoint is
36    /// placed as it always was.
37    pub previous_system_hash: Option<u64>,
38}
39
40/// What `on_write` decided about an incoming entry.
41pub(crate) enum OnWriteOutcome {
42    /// Store the entry with this (possibly replaced) content and token count.
43    Accept(String, usize),
44    /// The script declined the entry; report success to the writer.
45    Drop,
46}
47
48/// Serialize one region entry for a hook ctx. Typed metadata crosses as plain
49/// data so a script can key decisions off tool ids/names, but it is read-only:
50/// hooks return instructions (drop indices, replacement text), never entries.
51fn 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
86/// The `region` sub-object shared by all three hook ctx shapes.
87fn 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
96/// The Temporary-style block a custom region falls back to whenever its hook
97/// can't run or misbehaves - identical to the plain-`assemble` arm, so the
98/// region is never silently dropped.
99fn 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
112/// Render a custom region through its script, appending the results to the
113/// caller's block/message accumulators. Any failure falls back to the
114/// Temporary-style block with a warning; success is followed by a warn-only
115/// token re-check against the region's budget (no truncation - the opt-in
116/// exact-token preflight remains the hard guard).
117/// What one custom region is rendered from.
118///
119/// The window figures travel with the region rather than the accumulators
120/// because a render hook is told how full the window is so it can decide how
121/// much to emit - they describe the input, not where the output lands.
122pub(crate) struct RegionRender<'a> {
123    /// The region being rendered.
124    pub region: &'a Region,
125    /// Its render hook, when it declares one.
126    pub script: Option<&'a Arc<RegionScript>>,
127    /// Whether the region persists across stage transitions.
128    pub persistent: bool,
129    /// Stage metadata the hook sees.
130    pub meta: &'a AssembleMeta,
131    /// How full the window is right now.
132    pub window_current: usize,
133    /// How full it may get.
134    pub window_max: usize,
135}
136
137/// Where a rendered region's output is appended.
138///
139/// The caller owns both accumulators and interleaves several regions into them,
140/// so they are borrowed together rather than returned.
141pub(crate) struct RenderSink<'a> {
142    /// System blocks, for regions that render into the system prompt.
143    pub system_blocks: &'a mut Vec<leviath_providers::SystemBlock>,
144    /// Messages, for regions that render into the conversation.
145    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        // No compiled script on the window (plain `assemble()` callers, or a
163        // spawn path that skipped resolution). Same shape as a hook failure.
164        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
253/// Interpret `render`'s returned value: a string (one system block) or an
254/// object with optional `system` (string or array of strings) and `messages`
255/// (array of message objects). Strict on shape - any surprise is an `Err`,
256/// which the caller turns into the fallback block.
257fn 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    // A persistent region's rendered output is expected stable → cacheable.
268    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
329/// Build one provider message from a script-emitted message object. Three
330/// accepted shapes, constructed with the same wire types the built-in
331/// SlidingWindow arm emits (so a Rhai recreation of it is byte-identical):
332///
333/// - `{ role, content }` - plain text, role `user` or `assistant`
334/// - `{ role: "assistant", content?, tool_calls: [{id, name, arguments}] }`
335/// - `{ role: "user", tool_results: [{tool_call_id, content, is_error?}] }`
336fn 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
444/// Run `on_write` for an entry headed into a custom region. Failure of any
445/// kind accepts the entry unchanged - a script bug must not lose writes.
446pub(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
493/// Ask `on_overflow` which entries to drop, validate the answer, and apply it.
494/// Returns the tokens freed (0 when the hook is absent, fails, or returns an
495/// invalid/empty answer - callers fall back to oldest-first for the rest).
496pub(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    // Descending order keeps earlier indices valid while removing.
531    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
539/// Validate an `on_overflow` return value into a sorted, deduped index list.
540/// `None` when the shape is wrong or any index is out of range.
541fn 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    // ─── entry_to_json ───────────────────────────────────────────────────
620
621    #[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    // ─── render: happy paths ─────────────────────────────────────────────
661
662    #[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(&region, 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(&region, 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        // A script that recreates the SlidingWindow wire shapes: assistant
680        // text+tool_use, then a user tool_result message - built-ins parity.
681        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(&region, 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        // Assert the wire shapes through serde - no enum destructuring, so
708        // there are no never-taken match arms for the coverage gate.
709        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(&region, 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(&region, 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(&region, 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        // Budget is 1000 tokens; the script emits ~2000 tokens of output. The
756        // result is kept (warn-only per the design), not truncated.
757        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(&region, Some(&script(src)), false);
760        assert_eq!(blocks.len(), 1);
761        assert_eq!(blocks[0].text.len(), 8000);
762    }
763
764    // ─── render: fallbacks ───────────────────────────────────────────────
765
766    #[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(&region, 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(&region, 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(&region, 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(&region, 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(&region, 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        // The invalid-shape fallback has nothing to fall back TO when the
839        // region is empty - no block at all.
840        let region = region_with(&[]);
841        let (blocks, messages) = render(&region, 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(&region, 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(&region, 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        // The kind string the script sees matches the entry's typed kind.
873        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(&region, 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        // Indexing a missing key yields Null, so this covers absent-or-null
914        // without a short-circuit branch the coverage gate can't see taken.
915        assert_eq!(blocks[0]["thought_signature"], serde_json::Value::Null);
916    }
917
918    // ─── on_write ────────────────────────────────────────────────────────
919
920    /// Run `apply_on_write` and collapse the outcome to an Option - both
921    /// enum arms are exercised across this suite (through this one shared
922    /// mapping), so it has no never-taken branch.
923    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), &region, 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    // ─── on_overflow / apply_overflow ────────────────────────────────────
979
980    #[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        // Duplicate + unordered indices are deduped and applied safely.
988        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}