use std::sync::Arc;
use std::time::{Duration, Instant};
use rhai::{Dynamic, Engine, EvalAltResult, Scope, AST};
use crate::data_context::{DataContext, DataDep};
use crate::segments::{RenderContext, RenderResult, Segment, SegmentError};
use super::ctx_mirror::build_ctx;
use super::engine::{
set_current_plugin_id, set_render_deadline, DeadlineAbortMarker, DEFAULT_RENDER_DEADLINE_MS,
};
use super::output::validate_return;
use super::registry::CompiledPlugin;
struct RenderState;
impl RenderState {
fn install(plugin_id: &str, deadline: Instant) -> Self {
debug_assert!(
super::engine::render_deadline_snapshot().is_none(),
"RENDER_DEADLINE leaked from a prior render"
);
debug_assert!(
super::engine::current_plugin_id_snapshot().is_none(),
"CURRENT_PLUGIN_ID leaked from a prior render"
);
set_render_deadline(Some(deadline));
set_current_plugin_id(Some(plugin_id));
Self
}
}
impl Drop for RenderState {
fn drop(&mut self) {
set_render_deadline(None);
set_current_plugin_id(None);
}
}
pub struct RhaiSegment {
id: String,
ast: AST,
engine: Arc<Engine>,
config: Dynamic,
declared_deps: &'static [DataDep],
}
impl RhaiSegment {
#[must_use]
pub fn from_compiled(plugin: CompiledPlugin, engine: Arc<Engine>, config: Dynamic) -> Self {
let declared_deps: &'static [DataDep] = Vec::leak(plugin.declared_deps);
Self {
id: plugin.id,
ast: plugin.ast,
engine,
config,
declared_deps,
}
}
#[must_use]
pub fn id(&self) -> &str {
&self.id
}
fn classify_render_error(&self, err: Box<EvalAltResult>) -> SegmentError {
if let EvalAltResult::ErrorTerminated(token, _) = err.as_ref() {
if token.is::<DeadlineAbortMarker>() {
return SegmentError::new(format!(
"plugin `{}` exceeded the {}ms render deadline",
self.id, DEFAULT_RENDER_DEADLINE_MS
));
}
}
SegmentError::new(format!("plugin `{}` render failed: {err}", self.id))
}
}
impl Segment for RhaiSegment {
fn render(&self, ctx: &DataContext, rc: &RenderContext) -> RenderResult {
let mirror = build_ctx(ctx, rc, self.declared_deps, self.config.clone());
let deadline = Instant::now() + Duration::from_millis(DEFAULT_RENDER_DEADLINE_MS);
let _state = RenderState::install(&self.id, deadline);
let mut scope = Scope::new();
let returned: Dynamic = self
.engine
.call_fn(&mut scope, &self.ast, "render", (mirror,))
.map_err(|e| self.classify_render_error(e))?;
validate_return(returned, &self.id).map_err(|e| {
SegmentError::new(format!(
"plugin `{}` returned malformed shape: {e}",
self.id
))
})
}
fn data_deps(&self) -> &'static [DataDep] {
self.declared_deps
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::input::{ModelInfo, StatusContext, Tool, WorkspaceInfo};
use crate::plugins::build_engine;
use crate::plugins::registry::PluginRegistry;
use std::fs;
use std::path::PathBuf;
use std::sync::Arc;
use tempfile::TempDir;
fn minimal_status() -> StatusContext {
StatusContext {
tool: Tool::ClaudeCode,
model: Some(ModelInfo {
display_name: "Sonnet".to_string(),
}),
workspace: Some(WorkspaceInfo {
project_dir: PathBuf::from("/repo"),
git_worktree: None,
}),
context_window: None,
cost: None,
effort: None,
vim: None,
output_style: None,
agent_name: None,
version: None,
raw: Arc::new(serde_json::json!({})),
}
}
fn load_single(
dir: &tempfile::TempDir,
name: &str,
src: &str,
) -> (CompiledPlugin, Arc<Engine>) {
fs::write(dir.path().join(name), src).expect("write plugin");
let engine = build_engine();
let registry =
PluginRegistry::load_with_xdg(&[dir.path().to_path_buf()], None, &engine, &[]);
assert!(
registry.load_errors().is_empty(),
"unexpected load errors: {:?}",
registry.load_errors()
);
let plugin = registry
.into_plugins()
.into_iter()
.next()
.expect("plugin loaded");
(plugin, engine)
}
#[test]
fn plugin_can_read_terminal_width_from_ctx_render() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"tw.rhai",
r#"
const ID = "tw";
fn render(ctx) {
#{ runs: [#{ text: `${ctx.render.terminal_width}` }] }
}
"#,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(137);
let rendered = seg.render(&dc, &rc).unwrap().expect("rendered");
assert_eq!(rendered.text(), "137");
}
#[test]
fn plugin_returning_unit_hides_segment() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"hide.rhai",
r#"
const ID = "hide";
fn render(ctx) { () }
"#,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(80);
assert_eq!(seg.render(&dc, &rc).unwrap(), None);
}
#[test]
fn plugin_returning_single_run_renders() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"simple.rhai",
r#"
const ID = "simple";
fn render(ctx) {
#{ runs: [#{ text: "hello" }] }
}
"#,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(80);
let rendered = seg.render(&dc, &rc).unwrap().expect("rendered");
assert_eq!(rendered.text(), "hello");
}
#[test]
fn plugin_sees_status_fields_via_ctx() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"model_echo.rhai",
r#"
const ID = "model_echo";
fn render(ctx) {
#{ runs: [#{ text: ctx.status.model.display_name }] }
}
"#,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(80);
let rendered = seg.render(&dc, &rc).unwrap().expect("rendered");
assert_eq!(rendered.text(), "Sonnet");
}
#[test]
fn plugin_receives_config_passed_in() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"cfg.rhai",
r#"
const ID = "cfg";
fn render(ctx) {
#{ runs: [#{ text: ctx.config.label }] }
}
"#,
);
let mut config = rhai::Map::new();
config.insert("label".into(), Dynamic::from("configured".to_string()));
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::from_map(config));
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(80);
let rendered = seg.render(&dc, &rc).unwrap().expect("rendered");
assert_eq!(rendered.text(), "configured");
}
#[test]
fn plugin_can_read_ctx_env_from_rhai_side() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"env.rhai",
r#"
const ID = "env_probe";
fn render(ctx) {
let term = ctx.env.TERM;
let label = if term == () { "unset" } else { "set" };
#{ runs: [#{ text: label }] }
}
"#,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(80);
let rendered = seg.render(&dc, &rc).unwrap().expect("rendered");
assert!(rendered.text() == "set" || rendered.text() == "unset");
}
#[test]
fn declared_deps_surface_via_trait() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"deps.rhai",
r#"// @data_deps = ["usage"]
const ID = "deps";
fn render(ctx) { () }
"#,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
assert!(seg.data_deps().contains(&DataDep::Status));
assert!(seg.data_deps().contains(&DataDep::Usage));
}
#[test]
fn plugin_runtime_error_maps_to_segment_error() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"boom.rhai",
r#"
const ID = "boom";
fn render(ctx) {
let n = 1 / 0;
#{ runs: [#{ text: `${n}` }] }
}
"#,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(80);
let err = seg.render(&dc, &rc).unwrap_err();
assert!(err.message.contains("boom"), "message: {}", err.message);
assert!(
err.message.contains("render failed"),
"non-deadline failures must use the generic branch: {}",
err.message
);
assert!(
!err.message.contains("deadline"),
"non-deadline failures must NOT be relabeled as a timeout: {}",
err.message
);
}
#[test]
fn plugin_throw_cannot_impersonate_deadline_abort() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"fake.rhai",
r##"
const ID = "fake_deadline";
fn render(ctx) {
throw "linesmith:render-deadline-exceeded";
}
"##,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(80);
let err = seg.render(&dc, &rc).unwrap_err();
assert!(
err.message.contains("render failed"),
"throw must use the generic branch: {}",
err.message
);
assert!(
!err.message.contains("exceeded the"),
"thrown payload must not impersonate the host deadline message: {}",
err.message
);
}
#[test]
fn render_state_drop_clears_thread_locals() {
use crate::plugins::engine::{current_plugin_id_snapshot, render_deadline_snapshot};
{
let _state =
RenderState::install("guard_test", Instant::now() + Duration::from_secs(60));
assert!(render_deadline_snapshot().is_some());
assert_eq!(current_plugin_id_snapshot().as_deref(), Some("guard_test"));
}
assert!(render_deadline_snapshot().is_none(), "deadline leaked");
assert!(current_plugin_id_snapshot().is_none(), "plugin id leaked");
}
#[test]
fn plugin_returning_malformed_shape_maps_to_segment_error() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"bad.rhai",
r#"
const ID = "bad_shape";
fn render(ctx) { 42 }
"#,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(80);
let err = seg.render(&dc, &rc).unwrap_err();
assert!(
err.message.contains("bad_shape"),
"message: {}",
err.message
);
assert!(
err.message.to_lowercase().contains("malformed") || err.message.contains("must return"),
"message: {}",
err.message
);
}
#[test]
fn deadline_abort_surfaces_clear_segment_error() {
use crate::plugins::engine::set_render_deadline;
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) =
load_single(&tmp, "x.rhai", r#"const ID = "x"; fn render(ctx) { () }"#);
set_render_deadline(Some(Instant::now()));
let err = engine.eval::<()>("loop {}").unwrap_err();
set_render_deadline(None);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let segment_err = seg.classify_render_error(err);
assert!(
segment_err.message.contains("deadline"),
"deadline aborts should name the timeout: {}",
segment_err.message
);
}
#[test]
fn operation_limit_kills_infinite_loop_without_hang() {
let tmp = TempDir::new().expect("tempdir");
let (plugin, engine) = load_single(
&tmp,
"loop.rhai",
r#"
const ID = "loop";
fn render(ctx) {
loop {}
}
"#,
);
let seg = RhaiSegment::from_compiled(plugin, engine, Dynamic::UNIT);
let dc = DataContext::new(minimal_status());
let rc = RenderContext::new(80);
let err = seg.render(&dc, &rc).unwrap_err();
assert!(
err.message.to_lowercase().contains("operation") || err.message.contains("loop"),
"message: {}",
err.message
);
}
}