diff --git a/Cargo.lock b/Cargo.lock index 715ca55..95fb304 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -204,6 +204,16 @@ version = "0.15.7" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + [[package]] name = "fallible-iterator" version = "0.3.0" @@ -1118,6 +1128,16 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + [[package]] name = "siphasher" version = "1.0.3" @@ -1314,6 +1334,7 @@ dependencies = [ "libc", "mio", "pin-project-lite", + "signal-hook-registry", "socket2", "tokio-macros", "windows-sys 0.61.2", diff --git a/Cargo.toml b/Cargo.toml index a509c31..b96b00e 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,6 +62,12 @@ default = [] sqlite = ["dep:rusqlite"] # Embedded Rhai-backed `.ragsh` REPL session runtime (`repl::session`). repl = ["dep:rhai"] +# Recursive-language-model runtime (`rlm`): a driver model writes code cells +# executed in a sandboxed interpreter (embedded Rhai, or an external Python / +# JavaScript process) whose only host surface is capability calls back into +# the registry (`llm`, `tool`, `agent`). `tokio/process` + `tokio/io-util` +# drive the external interpreter subprocesses. +rlm = ["dep:rhai", "tokio/process", "tokio/io-util"] # Builtin generic tools (`harness::tools`) kept out of the default dependency # graph so host applications can choose whether they want these implementations. tools = ["dep:chrono", "dep:chrono-tz"] @@ -78,3 +84,12 @@ futures = "0.3" # The OpenAI provider is always compiled, so these build with a plain # `cargo build --examples`; they only need a network call + `OPENAI_API_KEY` # at run time. + +# --- RLM examples (need the `rlm` feature) --- +[[example]] +name = "rlm_rhai" +required-features = ["rlm"] + +[[example]] +name = "rlm_python" +required-features = ["rlm"] diff --git a/docs/modules/rlm/README.md b/docs/modules/rlm/README.md new file mode 100644 index 0000000..59b2f88 --- /dev/null +++ b/docs/modules/rlm/README.md @@ -0,0 +1,179 @@ +# RLM — the recursive-language-model runtime (`src/rlm/`) + +The `rlm` module (Cargo feature `rlm`) turns a code sandbox into a +recursive-language-model harness: a **driver model** writes code cells, the +cells execute in a **sandboxed interpreter**, and the only host surface the +scripts see is a set of **capability calls back into the registry** — sub-LLM +queries, tools, and sub-agent delegation. Because scripts can call models +(and agents that call models), the loop is recursive by construction, which +is the execution pattern of the RLM literature this crate is architected +around (see the crate-level docs in `src/lib.rs`). + +```text + ┌──────────────────────────────────────────────┐ + task ───► │ RlmRunner: driver model ⇄ code cells ⇄ obs. │ ───► answer + └──────────────┬───────────────────────────────┘ + │ eval(code) + ┌──────────────▼───────────────┐ + │ RlmSession │ + │ ┌──────────┐ HostCall │ + │ │ RlmInter-│ ─────────────► │ RlmHost ──► CapabilityRegistry + │ │ preter │ ◄───────────── │ (policy, counters, depth) + │ └──────────┘ Value/error │ ├─ llm → ChatModel::invoke + └──────────────────────────────┘ ├─ tool → Tool::call + └─ agent → HarnessAgent::run +``` + +## The three layers + +| Layer | Type | What it is | +|---|---|---| +| Interpreter | `RlmInterpreter` (trait) | Pluggable cell execution: `eval_cell`, `set_variable`, `usage_guide`, `shutdown`. State persists across cells like a notebook. | +| Session | `RlmSession` | One interpreter bound to one `RlmHost`; the programmatic "interpreter as an API" surface. Enforces per-cell policy fail-closed. | +| Runner | `RlmRunner` | The model-driven loop: template → system prompt, fenced code cells in, observations out, stop on `final_answer(...)`. | + +Every layer is usable on its own: an embedder that wants its own loop (or a +notebook UI) drives `RlmSession::eval` directly and never touches the runner. + +## Interpreter backends + +- **`InterpreterSpec::Rhai`** (default) — the embedded Rhai engine. This is + the only *hermetic* sandbox: no filesystem, network, or process access + exists inside the engine; the registered capability closures are its whole + world. Bounded by `max_operations`, the cell deadline (`on_progress` + hook), and the blocking bridge around every capability call (the same + fail-closed adapter as `repl::session`). +- **`InterpreterSpec::Python { binary, args }`** — an external CPython child + process (`python3` by default). A bootstrap prelude is injected via `-c`; + cells run in a persistent exec namespace. +- **`InterpreterSpec::Javascript { binary, args }`** — a Node.js child + (`node` by default, prelude via `-e`); cells run in a persistent `vm` + context. +- **`InterpreterSpec::Command { binary, args }`** — any command that speaks + the wire protocol itself (a containerized runner, a jailed interpreter, a + different language). + +### The wire protocol (external backends) + +Line-delimited JSON on the child's stdin/stdout; calls are strictly +sequential so no correlation ids are needed. Child → host: `ready`, +`call {call: HostCall}`, `result {stdout, value, error}`, `var_set`. +Host → child: `eval {code}`, `set_var {name, value}`, +`call_result {ok, value|error}`, `shutdown`. `HostCall` is the shared, +serde-stable call shape (`{"capability": "llm"|"tool"|"agent"|"final_answer", ...}`). +Any runtime that speaks this protocol is a valid backend — that is the +extension point for other harnesses (e.g. openhuman) to bring their own +interpreter. + +### Sandboxing honesty + +The host enforces every policy limit fail-closed for **all** backends (a +child that exceeds its deadline or trips a bound is killed, not asked). But +an external child process has whatever OS access the environment grants it. +For untrusted driver models, either stay on the embedded Rhai backend or run +the external interpreter inside real isolation (container/jail/seccomp) via +`InterpreterSpec::Command`. + +## The capability surface inside scripts + +Identical semantics across languages (spelling varies per usage guide): + +- `llm(prompt)` / `llm({model, prompt, system})` → string — a sub-LLM call; + the unnamed form uses the session's default sub-model. +- `tool(name, args)` → tool result (raw JSON when the tool provides it). + Arguments are validated against the tool's schema at the host boundary. +- `agent(name, input)` → string — delegates to a registered + `HarnessAgent`; a full nested agent run with event fan-out and the shared + recursion-depth guard (`RunConfig::checked_child_depth`). +- `final_answer(text)` — ends the run with this answer. +- `print(...)` / `console.log(...)` — captured and echoed back to the + driver next turn, bounded by `max_output_bytes` (explicit truncation + marker). + +### Error contract + +Script-visible failures (unknown tool, schema mismatch, tool error, provider +error) surface *inside* the script as catchable exceptions (`RlmError` in +Python/JS, `try`/`catch` in Rhai) so the driving model can adapt — that +feedback loop is the point of an RLM. Policy violations (`LimitExceeded`, +`Timeout`, `Cancelled`, `SubAgentDepth`) are **fatal**: the cell aborts (and +an external child is killed) so scripts can never observe and route around +their own resource limits. `rlm::is_fatal` is the classifier. + +## Config-driven runs + +Everything a run needs is one serde document — the integration surface for +external harnesses: + +```json +{ + "interpreter": {"kind": "python", "binary": "/opt/venv/bin/python3"}, + "driver_model": "openai", + "sub_model": "openai", + "template": "context-explorer", + "policy": { + "max_cells": 8, "max_llm_calls": 16, "max_tool_calls": 64, + "max_agent_calls": 8, "max_depth": 8, + "max_script_bytes": 65536, "max_output_bytes": 262144, + "cell_timeout": 90000, "max_operations": 5000000 + } +} +``` + +`RlmConfig::from_json` → `RlmRunner::from_config(config, registry, state)` → +`runner.set_context(json)` (optional) → `runner.run(task)`. + +### Templates + +`TemplateSpec::Named` selects a built-in; `TemplateSpec::Inline` carries a +custom `RlmTemplate` in the config document. Placeholders `{{language}}`, +`{{usage}}`, `{{capabilities}}`, `{{limits}}` are substituted at run time +from the live session (so the prompt always reflects the actual registry and +policy). Built-ins: + +- `general` — solve the task with code; sub-LLMs for fuzzy subproblems. +- `context-explorer` — the RLM long-context pattern: material injected as + the `context` variable, probed programmatically, never printed whole. +- `orchestrator` — decompose and delegate to registered sub-agents. + +## Policy (all fail-closed) + +`RlmPolicy`: `max_cells`, `max_script_bytes`, `max_output_bytes`, +`max_llm_calls`, `max_tool_calls`, `max_agent_calls`, `max_depth`, +`cell_timeout` (ms in JSON), `max_operations` (Rhai only). Counters are +session-cumulative. `RlmCancelFlag` provides sticky external cancellation, +observed mid-script (Rhai `on_progress`), mid-call (blocking bridge), and +before each cell. + +## Runner loop details + +- The driver must reply with exactly one fenced code block; the first fence + is extracted regardless of its info string (models mislabel languages). +- A fence-less reply earns one nudge (it is often raw unfenced code); a + second consecutive fence-less reply is accepted as a prose answer + (`RlmStopReason::ModelAnswered`). +- Observations echo captured stdout, the cell value, and any script error. +- Stop reasons: `Answered` (a cell called `final_answer`), `ModelAnswered`, + `CellBudgetExhausted`. The full per-cell trajectory is returned in + `RlmOutcome::steps`. + +## Tests and examples + +- Unit: `src/rlm/test.rs` (config round-trips, templates, extraction, the + embedded backend, runner loop against `ScriptedModel`). +- E2E: `tests/e2e_rlm.rs` — sub-agent delegation from scripts, real + `python3`/`node` protocol round-trips (skipped when the binary is + missing), fail-closed timeout kill, config-driven runner. +- Live: `tests/live_rlm.rs` — network-gated on `OPENAI_API_KEY`. +- Examples: `examples/rlm_rhai.rs` (context-explorer over expense records), + `examples/rlm_python.rs` (tools + a real sub-agent, external Python). + Run with `cargo run --features rlm --example rlm_rhai`. + +## Relationship to `repl::session` + +`repl::session` is the `.ragsh` interactive orchestration surface: Rhai +only, richer built-ins (graph authoring, batching), REPL-first. `rlm` is the +model-*driven* counterpart: interpreter-pluggable, config-first, and shaped +for embedding in other harnesses. They share design DNA deliberately — the +fail-closed blocking bridge, reserved capability boundary, and policy +posture are the same pattern. diff --git a/examples/rlm_python.rs b/examples/rlm_python.rs new file mode 100644 index 0000000..13db721 --- /dev/null +++ b/examples/rlm_python.rs @@ -0,0 +1,135 @@ +//! A live RLM run over an **external Python interpreter**, with a real tool +//! and a real sub-agent registered. +//! +//! The driver model writes Python cells; the child `python3` process executes +//! them, and every capability call (`llm`, `tool`, `agent`, `final_answer`) +//! travels back over the wire protocol to the host, which enforces the policy +//! and lowers to the harness runtime. The `agent("summarizer", ...)` call +//! drives a *complete nested agent run* on its own OpenAI-backed harness — +//! scripts calling agents calling models. +//! +//! Run with (needs `OPENAI_API_KEY` and `python3` on PATH): +//! +//! ```text +//! cargo run --features rlm --example rlm_python +//! ``` + +use std::sync::Arc; + +use async_trait::async_trait; +use serde_json::json; + +use tinyagents::harness::providers::openai::OpenAiModel; +use tinyagents::harness::runtime::AgentHarness; +use tinyagents::harness::tool::{Tool, ToolCall, ToolResult, ToolSchema}; +use tinyagents::registry::CapabilityRegistry; +use tinyagents::rlm::{RlmConfig, RlmRunner}; +use tinyagents::{HarnessSubAgent, Result, SubAgent}; + +/// A deterministic metrics tool the script can query. +struct MetricsTool; + +#[async_trait] +impl Tool<()> for MetricsTool { + fn name(&self) -> &str { + "service_metrics" + } + + fn description(&self) -> &str { + "Returns latency and error-rate metrics for a named service." + } + + fn schema(&self) -> ToolSchema { + ToolSchema::new( + "service_metrics", + "Returns latency and error-rate metrics for a named service.", + json!({ + "type": "object", + "properties": { "service": { "type": "string" } }, + "required": ["service"], + "additionalProperties": false + }), + ) + } + + async fn call(&self, _state: &(), call: ToolCall) -> Result { + let service = call + .arguments + .get("service") + .and_then(|v| v.as_str()) + .unwrap_or("unknown"); + let metrics = match service { + "checkout" => json!({"p99_ms": 2140, "error_rate": 0.031, "deploys_today": 3}), + "search" => json!({"p99_ms": 180, "error_rate": 0.002, "deploys_today": 0}), + "auth" => json!({"p99_ms": 95, "error_rate": 0.001, "deploys_today": 1}), + other => json!({"error": format!("unknown service `{other}`")}), + }; + let mut result = ToolResult::text(call.id, call.name, metrics.to_string()); + result.raw = Some(metrics); + Ok(result) + } +} + +#[tokio::main] +async fn main() -> Result<()> { + dotenvy::dotenv().ok(); + + let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); + registry.register_model("openai", Arc::new(OpenAiModel::from_env()?))?; + registry.register_tool(Arc::new(MetricsTool))?; + + // A real sub-agent on its own harness: scripts delegate to it by name. + let mut child: AgentHarness<()> = AgentHarness::new(); + child + .register_model("openai", Arc::new(OpenAiModel::from_env()?)) + .set_default_model("openai"); + let summarizer = Arc::new( + SubAgent::new( + "summarizer", + "Writes a crisp two-sentence incident summary from raw findings.", + Arc::new(child), + ) + .with_system_prompt( + "You summarize incident findings for executives: two sentences, plain language, \ + lead with impact.", + ), + ); + registry.register_agent(Arc::new(HarnessSubAgent::new(summarizer)))?; + + let config = RlmConfig::from_json( + r#"{ + "interpreter": {"kind": "python"}, + "driver_model": "openai", + "template": "general", + "policy": {"max_cells": 8, "cell_timeout": 90000} + }"#, + )?; + + let mut runner = RlmRunner::from_config(config, Arc::new(registry), Arc::new(()))?; + let outcome = runner + .run( + "Check the service_metrics tool for the services checkout, search, and auth. \ + Identify which service looks unhealthy and why. Then delegate to the `summarizer` \ + agent to produce an executive summary of your findings, and return that summary \ + as the final answer.", + ) + .await?; + + for (i, step) in outcome.steps.iter().enumerate() { + println!("── cell {} ──\n{}\n", i + 1, step.code); + if !step.outcome.stdout.is_empty() { + println!("stdout:\n{}", step.outcome.stdout); + } + if let Some(error) = &step.outcome.error { + println!("error: {error}"); + } + } + println!("── answer ({:?}) ──", outcome.stop_reason); + println!("{}", outcome.answer.as_deref().unwrap_or("(none)")); + println!( + "\ndriver calls: {}, sub-llm calls: {}, tool calls: {}, agent calls: {}", + outcome.driver_calls, outcome.sub_llm_calls, outcome.tool_calls, outcome.agent_calls + ); + runner.shutdown().await?; + Ok(()) +} diff --git a/examples/rlm_rhai.rs b/examples/rlm_rhai.rs new file mode 100644 index 0000000..26a304a --- /dev/null +++ b/examples/rlm_rhai.rs @@ -0,0 +1,91 @@ +//! A live recursive-language-model run over the embedded Rhai sandbox. +//! +//! The classic RLM shape: a context too noisy to eyeball is injected into the +//! sandbox as the `context` variable, and the driver model must *probe it +//! with code* — slicing, filtering, and delegating fuzzy judgment on +//! individual entries to sub-LLM calls (`llm(...)`) — before answering with +//! `final_answer(...)`. +//! +//! Run with (needs `OPENAI_API_KEY`, optional `OPENAI_MODEL`): +//! +//! ```text +//! cargo run --features rlm --example rlm_rhai +//! ``` + +use std::sync::Arc; + +use serde_json::json; + +use tinyagents::Result; +use tinyagents::harness::providers::openai::OpenAiModel; +use tinyagents::registry::CapabilityRegistry; +use tinyagents::rlm::{RlmConfig, RlmRunner}; + +#[tokio::main] +async fn main() -> Result<()> { + dotenvy::dotenv().ok(); + let model = OpenAiModel::from_env()?; + + let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); + registry.register_model("openai", Arc::new(model))?; + + // The whole run is one JSON document — the same document an external + // harness could load from disk. + let config = RlmConfig::from_json( + r#"{ + "interpreter": {"kind": "rhai"}, + "driver_model": "openai", + "template": "context-explorer", + "policy": {"max_cells": 8, "max_llm_calls": 16, "cell_timeout": 90000} + }"#, + )?; + + let mut runner = RlmRunner::from_config(config, Arc::new(registry), Arc::new(()))?; + + // A synthetic "too big to read" context: expense records with three + // hidden anomalies among routine noise. + let mut records = Vec::new(); + for i in 0..200usize { + let team = ["platform", "growth", "ops"][i % 3]; + let amount = 40 + (i * 7) % 60; + records.push(json!({ + "id": i, + "team": team, + "amount_usd": amount, + "memo": format!("routine cloud spend, invoice {i}"), + })); + } + records[57] = json!({"id": 57, "team": "growth", "amount_usd": 18400, + "memo": "annual conference sponsorship paid twice, needs review"}); + records[121] = json!({"id": 121, "team": "ops", "amount_usd": 9750, + "memo": "emergency hardware replacement after flood damage"}); + records[188] = json!({"id": 188, "team": "platform", "amount_usd": 12300, + "memo": "contractor invoice with mismatched PO number"}); + runner.set_context(json!(records)).await?; + + println!("── system prompt ──\n{}\n", runner.system_prompt()); + + let outcome = runner + .run( + "The `context` variable holds 200 expense records. Find every anomalous record \ + (unusual amount or memo), and summarize each anomaly in one line.", + ) + .await?; + + for (i, step) in outcome.steps.iter().enumerate() { + println!("── cell {} ──\n{}\n", i + 1, step.code); + if !step.outcome.stdout.is_empty() { + println!("stdout:\n{}", step.outcome.stdout); + } + if let Some(error) = &step.outcome.error { + println!("error: {error}"); + } + } + println!("── answer ({:?}) ──", outcome.stop_reason); + println!("{}", outcome.answer.as_deref().unwrap_or("(none)")); + println!( + "\ndriver calls: {}, sub-llm calls: {}, tool calls: {}, agent calls: {}", + outcome.driver_calls, outcome.sub_llm_calls, outcome.tool_calls, outcome.agent_calls + ); + Ok(()) +} diff --git a/src/lib.rs b/src/lib.rs index 6efec5e..3299563 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -56,10 +56,14 @@ //! Hosted and local providers (OpenAI plus the OpenAI-compatible endpoints for //! Anthropic, Ollama, DeepSeek, Groq, xAI, OpenRouter, Together, and Mistral) //! are compiled in unconditionally alongside the offline, deterministic -//! [`harness::providers::MockModel`]. Two Cargo features gate optional, +//! [`harness::providers::MockModel`]. Three Cargo features gate optional, //! heavier dependencies instead: `sqlite` (embedded SQLite checkpointer, -//! [`graph::checkpoint::SqliteCheckpointer`]) and `repl` (embedded Rhai engine -//! powering the `.ragsh` session runtime, [`repl::session`]). +//! [`graph::checkpoint::SqliteCheckpointer`]), `repl` (embedded Rhai engine +//! powering the `.ragsh` session runtime, [`repl::session`]), and `rlm` (the +//! recursive-language-model runtime: a driver model writes code cells run in +//! a sandboxed interpreter — embedded Rhai or an external Python/JavaScript +//! process — whose only host surface is capability calls back into the +//! registry). //! //! ## Crate-root re-exports //! @@ -73,6 +77,8 @@ pub mod harness; pub mod language; pub mod registry; pub mod repl; +#[cfg(feature = "rlm")] +pub mod rlm; // --- Error: the crate-wide error type and `Result` alias --- pub use error::{Result, TinyAgentsError}; @@ -245,3 +251,15 @@ pub use repl::session::{ LanguageCompiler, ReplCallKind, ReplCallRecord, ReplCancelFlag, ReplCapabilities, ReplPolicy, ReplResult, ReplSession, ReplValue, ReplVariables, }; + +// --- RLM runtime (feature = "rlm") --- +// The recursive-language-model surface: a driver model writes code cells that +// run in a sandboxed interpreter (embedded Rhai or an external Python/Node +// process) whose only host surface is capability calls (`llm`, `tool`, +// `agent`) back into the registry. Config-driven end to end (`RlmConfig`). +#[cfg(feature = "rlm")] +pub use rlm::{ + CellOutcome, HostCall, InterpreterSpec, RlmCallKind, RlmCallRecord, RlmCancelFlag, RlmConfig, + RlmHost, RlmHostApi, RlmInterpreter, RlmOutcome, RlmPolicy, RlmRunner, RlmSession, RlmStep, + RlmStopReason, RlmTemplate, TemplateSpec, +}; diff --git a/src/rlm/host.rs b/src/rlm/host.rs new file mode 100644 index 0000000..406b2a1 --- /dev/null +++ b/src/rlm/host.rs @@ -0,0 +1,412 @@ +//! The host side of the RLM capability boundary. +//! +//! Every interpreter backend — embedded or external — funnels script +//! capability calls through one object-safe seam, [`RlmHostApi::handle`]. +//! [`RlmHost`] is the standard implementation: it resolves names through the +//! session's [`CapabilityRegistry`], enforces the [`RlmPolicy`] call and +//! recursion limits fail-closed, records an [`RlmCallRecord`] per call, and +//! lowers to the real harness runtime (`ChatModel::invoke`, `Tool::call`, +//! `HarnessAgent::run`). +//! +//! ## Fatal vs script-visible errors +//! +//! A capability call can fail two ways, and the distinction is the sandbox +//! contract: +//! +//! - **Script-visible** failures (unknown tool, tool returned an error, model +//! provider error) surface *inside* the script as a raised +//! exception/runtime error. The driving model observes them in the cell +//! outcome and may adapt — that feedback loop is the whole point of an RLM. +//! - **Fatal** failures ([`TinyAgentsError::LimitExceeded`], +//! [`Timeout`](TinyAgentsError::Timeout), +//! [`Cancelled`](TinyAgentsError::Cancelled), +//! [`SubAgentDepth`](TinyAgentsError::SubAgentDepth)) mean a policy bound +//! tripped; the cell is aborted (and an external interpreter's child +//! process killed) rather than letting the script observe and route around +//! its own resource limits. + +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use serde_json::{Value, json}; + +use super::types::{HostCall, RlmCallRecord, RlmCancelFlag, RlmPolicy}; +use crate::error::{Result, TinyAgentsError}; +use crate::graph::subagent_node::SubAgentInput; +use crate::harness::events::EventSink; +use crate::harness::message::Message; +use crate::harness::model::ModelRequest; +use crate::harness::tool::ToolCall; +use crate::registry::{CapabilityRegistry, ComponentKind}; + +/// Returns whether a capability error must abort the cell (policy bound +/// tripped) instead of surfacing inside the script. +pub fn is_fatal(err: &TinyAgentsError) -> bool { + matches!( + err, + TinyAgentsError::LimitExceeded(_) + | TinyAgentsError::Timeout(_) + | TinyAgentsError::Cancelled + | TinyAgentsError::SubAgentDepth(_) + ) +} + +/// A snapshot of the registered capability names a session may reach, +/// rendered into the driver prompt so the model knows what it can call. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct CapabilityListing { + /// Registered model names. + pub models: Vec, + /// Registered tool names (with descriptions when available). + pub tools: Vec<(String, String)>, + /// Registered agent names. + pub agents: Vec, +} + +/// The object-safe host seam every interpreter backend calls through. +/// +/// Implementations must be cheap to share (`Arc`): the +/// embedded Rhai engine clones the handle into each capability closure, and +/// the external-interpreter driver holds it across protocol frames. +#[async_trait] +pub trait RlmHostApi: Send + Sync { + /// Executes one capability call on behalf of a script. + /// + /// A `Ok(value)` is handed back to the script; an `Err` is surfaced into + /// the script when recoverable (see [`is_fatal`]) or aborts the cell when + /// fatal. + async fn handle(&self, call: HostCall) -> Result; + + /// The capability names available to scripts (for prompt rendering). + fn capabilities(&self) -> CapabilityListing; + + /// The wall-clock deadline of the cell currently being evaluated. + fn deadline(&self) -> Option; + + /// The session cancellation flag. + fn cancel_flag(&self) -> RlmCancelFlag; +} + +/// Session-cumulative counters enforced against [`RlmPolicy`]. +#[derive(Debug, Default, Clone, Copy)] +struct CallCounters { + llm: usize, + tool: usize, + agent: usize, +} + +/// Per-cell mutable buffers the session arms before evaluating a cell and +/// drains after. +#[derive(Debug, Default)] +struct CellBuffers { + deadline: Option, + calls: Vec, + final_answer: Option, +} + +/// The standard capability host, bound to a [`CapabilityRegistry`]. +pub struct RlmHost { + registry: Arc>, + state: Arc, + policy: RlmPolicy, + /// The model `llm(...)` reaches when the script names none. + default_model: Option, + /// The session's run depth — the parent depth for sub-agent runs. + run_depth: usize, + events: EventSink, + cancel: RlmCancelFlag, + counters: Mutex, + cell: Mutex, +} + +impl RlmHost { + /// Builds a host over a capability registry and application state. + pub fn new(registry: Arc>, state: Arc) -> Self { + Self { + registry, + state, + policy: RlmPolicy::default(), + default_model: None, + run_depth: 0, + events: EventSink::default(), + cancel: RlmCancelFlag::new(), + counters: Mutex::new(CallCounters::default()), + cell: Mutex::new(CellBuffers::default()), + } + } + + /// Sets the session policy. + pub fn with_policy(mut self, policy: RlmPolicy) -> Self { + self.policy = policy; + self + } + + /// Sets the default model `llm(...)` reaches when the script names none. + pub fn with_default_model(mut self, model: impl Into) -> Self { + self.default_model = Some(model.into()); + self + } + + /// Sets the run depth sub-agent calls recurse below. + pub fn with_run_depth(mut self, depth: usize) -> Self { + self.run_depth = depth; + self + } + + /// Installs an event sink child agent runs fan onto. + pub fn with_events(mut self, events: EventSink) -> Self { + self.events = events; + self + } + + /// Installs an external cancellation flag. + pub fn with_cancel_flag(mut self, cancel: RlmCancelFlag) -> Self { + self.cancel = cancel; + self + } + + /// The session policy. + pub fn policy(&self) -> &RlmPolicy { + &self.policy + } + + /// The application state capability calls run against. + pub fn app_state(&self) -> Arc { + self.state.clone() + } + + /// Arms the per-cell buffers before a cell is evaluated. + pub(super) fn begin_cell(&self) { + let mut cell = self.cell.lock().expect("cell buffers poisoned"); + cell.deadline = self.policy.cell_timeout.map(|t| Instant::now() + t); + cell.calls.clear(); + cell.final_answer = None; + } + + /// Drains the per-cell buffers after a cell finished: the recorded calls + /// and the final answer, if `final_answer(...)` was called. + pub(super) fn end_cell(&self) -> (Vec, Option) { + let mut cell = self.cell.lock().expect("cell buffers poisoned"); + cell.deadline = None; + (std::mem::take(&mut cell.calls), cell.final_answer.take()) + } + + fn record(&self, record: RlmCallRecord) { + self.cell + .lock() + .expect("cell buffers poisoned") + .calls + .push(record); + } + + fn bump(&self, call: &HostCall) -> Result<()> { + let mut counters = self.counters.lock().expect("counters poisoned"); + match call { + HostCall::Llm { .. } => { + if counters.llm >= self.policy.max_llm_calls { + return Err(TinyAgentsError::LimitExceeded(format!( + "llm call limit ({}) exceeded", + self.policy.max_llm_calls + ))); + } + counters.llm += 1; + } + HostCall::Tool { .. } => { + if counters.tool >= self.policy.max_tool_calls { + return Err(TinyAgentsError::LimitExceeded(format!( + "tool call limit ({}) exceeded", + self.policy.max_tool_calls + ))); + } + counters.tool += 1; + } + HostCall::Agent { .. } => { + if counters.agent >= self.policy.max_agent_calls { + return Err(TinyAgentsError::LimitExceeded(format!( + "agent call limit ({}) exceeded", + self.policy.max_agent_calls + ))); + } + counters.agent += 1; + } + HostCall::FinalAnswer { .. } => {} + } + Ok(()) + } + + /// Session-cumulative `(llm, tool, agent)` call counts. + pub fn call_counts(&self) -> (usize, usize, usize) { + let counters = self.counters.lock().expect("counters poisoned"); + (counters.llm, counters.tool, counters.agent) + } + + async fn handle_llm( + &self, + model: Option, + prompt: String, + system: Option, + ) -> Result { + let model_name = model + .or_else(|| self.default_model.clone()) + .ok_or_else(|| { + TinyAgentsError::Validation( + "llm: no model named and the session has no default model".to_string(), + ) + })?; + let model = self + .registry + .model(&model_name) + .ok_or_else(|| TinyAgentsError::ModelNotFound(model_name.clone()))?; + let mut messages = Vec::new(); + if let Some(system) = system { + messages.push(Message::system(system)); + } + messages.push(Message::user(prompt)); + // `model_name` is the registry name, not a provider model id; the + // resolved ChatModel carries its own provider configuration. + let request = ModelRequest { + messages, + ..Default::default() + }; + let start = Instant::now(); + let response = model.invoke(&self.state, request).await?; + let text = Message::Assistant(response.message).text(); + self.record(RlmCallRecord { + kind: super::types::RlmCallKind::Llm, + name: model_name, + detail: json!({ "chars": text.len() }), + elapsed: start.elapsed(), + }); + Ok(Value::String(text)) + } + + async fn handle_tool(&self, tool_name: String, arguments: Value) -> Result { + let tool = self + .registry + .tool(&tool_name) + .ok_or_else(|| TinyAgentsError::ToolNotFound(tool_name.clone()))?; + let call = ToolCall::new( + crate::harness::ids::new_call_id().as_str().to_string(), + tool_name.clone(), + arguments.clone(), + ); + // Validate against the tool's schema up front so the script gets a + // precise, catchable error instead of tool-dependent behavior. + tool.schema().validate_call(&call)?; + let start = Instant::now(); + let result = tool.call(&self.state, call).await?; + self.record(RlmCallRecord { + kind: super::types::RlmCallKind::Tool, + name: tool_name, + detail: json!({ "arguments": arguments }), + elapsed: start.elapsed(), + }); + if let Some(error) = result.error { + return Err(TinyAgentsError::Tool(error)); + } + match result.raw { + Some(raw) => Ok(raw), + None => Ok(Value::String(result.content)), + } + } + + async fn handle_agent( + &self, + agent_name: String, + input: String, + data: Option, + ) -> Result { + // Reuse the shared harness depth guard so RLM sub-runs stay in + // lock-step with SubAgent / SubAgentTool / the REPL. + crate::harness::context::RunConfig::checked_child_depth( + self.run_depth, + self.policy.max_depth, + )?; + let agent = self.registry.agent(&agent_name).ok_or_else(|| { + TinyAgentsError::Capability(format!("agent `{agent_name}` is not registered")) + })?; + let mut sub_input = SubAgentInput::prompt(input); + if let Some(data) = data { + sub_input = sub_input.with_data(data); + } + let start = Instant::now(); + let output = agent.run(sub_input, self.events.clone()).await?; + self.record(RlmCallRecord { + kind: super::types::RlmCallKind::Agent, + name: agent_name, + detail: json!({ + "model_calls": output.model_calls, + "tool_calls": output.tool_calls, + }), + elapsed: start.elapsed(), + }); + Ok(Value::String(output.text)) + } +} + +#[async_trait] +impl RlmHostApi for RlmHost { + async fn handle(&self, call: HostCall) -> Result { + if self.cancel.is_cancelled() { + return Err(TinyAgentsError::Cancelled); + } + if let Some(deadline) = self.deadline() + && Instant::now() >= deadline + { + return Err(TinyAgentsError::Timeout( + "rlm cell exceeded its wall-clock timeout".to_string(), + )); + } + self.bump(&call)?; + match call { + HostCall::Llm { + model, + prompt, + system, + } => self.handle_llm(model, prompt, system).await, + HostCall::Tool { tool, arguments } => self.handle_tool(tool, arguments).await, + HostCall::Agent { agent, input, data } => self.handle_agent(agent, input, data).await, + HostCall::FinalAnswer { answer } => { + let mut cell = self.cell.lock().expect("cell buffers poisoned"); + cell.final_answer = Some(answer); + cell.calls.push(RlmCallRecord { + kind: super::types::RlmCallKind::FinalAnswer, + name: "final_answer".to_string(), + detail: Value::Null, + elapsed: Duration::default(), + }); + Ok(Value::Null) + } + } + } + + fn capabilities(&self) -> CapabilityListing { + let tools = self + .registry + .names(ComponentKind::Tool) + .into_iter() + .map(|name| { + let description = self + .registry + .tool(&name) + .map(|tool| tool.description().to_string()) + .unwrap_or_default(); + (name, description) + }) + .collect(); + CapabilityListing { + models: self.registry.names(ComponentKind::Model), + tools, + agents: self.registry.names(ComponentKind::Agent), + } + } + + fn deadline(&self) -> Option { + self.cell.lock().expect("cell buffers poisoned").deadline + } + + fn cancel_flag(&self) -> RlmCancelFlag { + self.cancel.clone() + } +} diff --git a/src/rlm/interpreter/external.rs b/src/rlm/interpreter/external.rs new file mode 100644 index 0000000..fc5bd1c --- /dev/null +++ b/src/rlm/interpreter/external.rs @@ -0,0 +1,547 @@ +//! The external-process interpreter backend (Python, Node.js, or any command +//! speaking the wire protocol). +//! +//! ## Wire protocol +//! +//! Line-delimited JSON over the child's stdin/stdout. The child is a +//! long-lived REPL: globals persist across cells. Frames, host → child: +//! +//! - `{"op":"eval","code":"..."}` — evaluate one cell. +//! - `{"op":"set_var","name":"...","value":}` — set a global. +//! - `{"op":"call_result","ok":true,"value":}` / +//! `{"op":"call_result","ok":false,"error":"..."}` — the reply to a +//! capability call (script-visible failures arrive as `ok:false`; *fatal* +//! failures never get a reply — the host kills the child instead). +//! - `{"op":"shutdown"}` — exit cleanly. +//! +//! Child → host: +//! +//! - `{"op":"ready"}` — emitted once after bootstrap. +//! - `{"op":"call","call":{"capability":"llm"|"tool"|"agent"|"final_answer",...}}` +//! — a capability call (the `call` payload is a serialized +//! [`HostCall`]); the child blocks until the matching `call_result`. +//! - `{"op":"result","stdout":"...","value":,"error":}` — +//! the cell outcome. +//! - `{"op":"var_set"}` — acknowledges `set_var`. +//! +//! Calls are strictly sequential (one cell evaluates at a time and blocks on +//! each capability call), so frames need no correlation ids. +//! +//! ## Sandboxing honesty +//! +//! Unlike the embedded Rhai backend, a child process has whatever OS access +//! the embedder's environment grants it. The host still enforces every +//! [`RlmPolicy`](crate::rlm::RlmPolicy) limit fail-closed — a cell that +//! exceeds its deadline or trips a policy bound gets its child **killed**, +//! not asked nicely — but filesystem/network isolation for untrusted models +//! must come from the embedder (container, jail, seccomp, a locked-down +//! `InterpreterSpec::Command` runner). + +use std::process::Stdio; +use std::sync::{Arc, Mutex}; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use serde_json::{Value, json}; +use tokio::io::{AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader}; +use tokio::process::{Child, ChildStdin, ChildStdout}; + +use super::{CellEval, RlmInterpreter}; +use crate::error::{Result, TinyAgentsError}; +use crate::rlm::host::{RlmHostApi, is_fatal}; +use crate::rlm::types::HostCall; + +/// The Python bootstrap program (`python3 -u -c `), a minimal REPL +/// speaking the wire protocol. User prints are captured per cell; the +/// protocol channel is the real stdout, saved before any redirection. +const PYTHON_PRELUDE: &str = r#" +import ast, contextlib, io, json, sys, traceback +_out, _in = sys.stdout, sys.stdin + +def _send(obj): + _out.write(json.dumps(obj) + "\n"); _out.flush() + +class RlmError(Exception): + pass + +def _call(call): + _send({"op": "call", "call": call}) + while True: + line = _in.readline() + if not line: + sys.exit(0) + msg = json.loads(line) + if msg.get("op") == "call_result": + if msg.get("ok"): + return msg.get("value") + raise RlmError(msg.get("error") or "capability call failed") + +def llm(prompt, model=None, system=None): + if isinstance(prompt, dict): + model, system, prompt = prompt.get("model"), prompt.get("system"), prompt.get("prompt") + return _call({"capability": "llm", "prompt": prompt, "model": model, "system": system}) + +def tool(name, arguments=None): + return _call({"capability": "tool", "tool": name, "arguments": arguments}) + +def agent(name, input, data=None): + return _call({"capability": "agent", "agent": name, "input": str(input), "data": data}) + +def final_answer(answer): + _call({"capability": "final_answer", "answer": str(answer)}) + +_g = {"__name__": "__rlm__", "llm": llm, "tool": tool, "agent": agent, + "final_answer": final_answer, "RlmError": RlmError} + +def _eval(code): + buf = io.StringIO() + value, error = None, None + try: + with contextlib.redirect_stdout(buf): + tree = ast.parse(code, mode="exec") + if tree.body and isinstance(tree.body[-1], ast.Expr): + last = ast.Expression(tree.body[-1].value) + body = ast.Module(body=tree.body[:-1], type_ignores=[]) + exec(compile(body, "", "exec"), _g) + value = eval(compile(last, "", "eval"), _g) + else: + exec(compile(tree, "", "exec"), _g) + except Exception: + error = traceback.format_exc(limit=4) + try: + json.dumps(value) + except Exception: + value = repr(value) + _send({"op": "result", "stdout": buf.getvalue(), "value": value, "error": error}) + +_send({"op": "ready"}) +for _line in _in: + _msg = json.loads(_line) + _op = _msg.get("op") + if _op == "eval": + _eval(_msg.get("code") or "") + elif _op == "set_var": + _g[_msg["name"]] = _msg.get("value") + _send({"op": "var_set"}) + elif _op == "shutdown": + break +"#; + +/// The Node.js bootstrap program (`node -e `). Cells run in a +/// persistent `vm` context; `console.log` is captured per cell; capability +/// calls block on stdin with `fs.readSync`. +const JAVASCRIPT_PRELUDE: &str = r#" +const fs = require('fs'); +const vm = require('vm'); +let inbuf = Buffer.alloc(0); +function readLine() { + for (;;) { + const idx = inbuf.indexOf(10); + if (idx >= 0) { + const line = inbuf.slice(0, idx).toString('utf8'); + inbuf = inbuf.slice(idx + 1); + return line; + } + const chunk = Buffer.alloc(65536); + let n; + try { n = fs.readSync(0, chunk, 0, chunk.length, null); } + catch (e) { if (e.code === 'EAGAIN') continue; throw e; } + if (n === 0) process.exit(0); + inbuf = Buffer.concat([inbuf, chunk.slice(0, n)]); + } +} +function send(obj) { fs.writeSync(1, JSON.stringify(obj) + '\n'); } +class RlmError extends Error {} +function call(c) { + send({ op: 'call', call: c }); + for (;;) { + const msg = JSON.parse(readLine()); + if (msg.op === 'call_result') { + if (msg.ok) return msg.value === undefined ? null : msg.value; + throw new RlmError(msg.error || 'capability call failed'); + } + } +} +let stdoutBuf = ''; +function logLine(args) { + stdoutBuf += args.map(x => (typeof x === 'string' ? x : JSON.stringify(x))).join(' ') + '\n'; +} +const sandbox = { + llm: p => (typeof p === 'string' + ? call({ capability: 'llm', prompt: p, model: null, system: null }) + : call({ capability: 'llm', prompt: p.prompt, model: p.model || null, system: p.system || null })), + tool: (name, args) => call({ capability: 'tool', tool: name, arguments: args === undefined ? null : args }), + agent: (name, input) => call({ capability: 'agent', agent: name, input: String(input), data: null }), + final_answer: a => { call({ capability: 'final_answer', answer: String(a) }); }, + console: { log: (...xs) => logLine(xs), error: (...xs) => logLine(xs), info: (...xs) => logLine(xs) }, + RlmError, + JSON, Math, +}; +const ctx = vm.createContext(sandbox); +send({ op: 'ready' }); +for (;;) { + const msg = JSON.parse(readLine()); + if (msg.op === 'eval') { + stdoutBuf = ''; + let value = null, error = null; + try { + value = vm.runInContext(msg.code, ctx, { filename: '' }); + if (value === undefined) value = null; + if (typeof value === 'function') value = String(value); + try { JSON.stringify(value); } catch { value = String(value); } + } catch (e) { + error = e && e.stack ? String(e.stack).split('\n').slice(0, 4).join('\n') : String(e); + } + send({ op: 'result', stdout: stdoutBuf, value, error }); + } else if (msg.op === 'set_var') { + sandbox[msg.name] = msg.value; + send({ op: 'var_set' }); + } else if (msg.op === 'shutdown') { + process.exit(0); + } +} +"#; + +/// How long the child gets to print its `ready` frame after spawning. +const STARTUP_TIMEOUT: Duration = Duration::from_secs(15); + +/// The fallback bound for protocol exchanges when the session armed no cell +/// deadline (`RlmPolicy::cell_timeout: None`). +const DEFAULT_EXCHANGE_TIMEOUT: Duration = Duration::from_secs(300); + +/// A live child process with its protocol streams. +struct ChildHandle { + child: Child, + stdin: ChildStdin, + stdout: BufReader, + /// Rolling capture of the child's stderr, included in error messages. + stderr: Arc>, +} + +/// The external-process backend. See the [module docs](self). +pub struct ExternalInterpreter { + language: String, + binary: String, + args: Vec, + child: Option, +} + +impl ExternalInterpreter { + /// A CPython child (`binary` defaults to `python3`). + pub fn python(binary: Option<&str>, extra_args: Vec) -> Self { + let mut args = extra_args; + args.extend([ + "-u".to_string(), + "-c".to_string(), + PYTHON_PRELUDE.to_string(), + ]); + Self { + language: "python".to_string(), + binary: binary.unwrap_or("python3").to_string(), + args, + child: None, + } + } + + /// A Node.js child (`binary` defaults to `node`). + pub fn javascript(binary: Option<&str>, extra_args: Vec) -> Self { + let mut args = extra_args; + args.extend(["-e".to_string(), JAVASCRIPT_PRELUDE.to_string()]); + Self { + language: "javascript".to_string(), + binary: binary.unwrap_or("node").to_string(), + args, + child: None, + } + } + + /// An arbitrary command that speaks the wire protocol itself. Scripts are + /// assumed to be Python-flavored for prompt purposes. + pub fn command(binary: String, args: Vec) -> Self { + Self { + language: "python".to_string(), + binary, + args, + child: None, + } + } + + fn stderr_tail(&self) -> String { + self.child + .as_ref() + .map(|c| { + let text = c.stderr.lock().expect("stderr buffer poisoned"); + let tail: String = text.chars().rev().take(2000).collect(); + tail.chars().rev().collect() + }) + .unwrap_or_default() + } + + fn broken(&mut self, context: &str) -> TinyAgentsError { + let stderr = self.stderr_tail(); + self.child = None; + let mut message = format!("rlm external interpreter ({}): {context}", self.binary); + if !stderr.is_empty() { + message.push_str(&format!("; stderr tail: {stderr}")); + } + TinyAgentsError::Capability(message) + } + + async fn ensure_started(&mut self) -> Result<()> { + if self.child.is_some() { + return Ok(()); + } + let mut child = tokio::process::Command::new(&self.binary) + .args(&self.args) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::piped()) + .kill_on_drop(true) + .spawn() + .map_err(|err| { + TinyAgentsError::Capability(format!( + "rlm external interpreter: failed to spawn `{}`: {err}", + self.binary + )) + })?; + let stdin = child.stdin.take().expect("piped stdin"); + let stdout = BufReader::new(child.stdout.take().expect("piped stdout")); + let stderr_pipe = child.stderr.take().expect("piped stderr"); + let stderr = Arc::new(Mutex::new(String::new())); + let stderr_buf = stderr.clone(); + tokio::spawn(async move { + let mut reader = BufReader::new(stderr_pipe); + let mut buffer = [0u8; 4096]; + loop { + match reader.read(&mut buffer).await { + Ok(0) | Err(_) => break, + Ok(n) => { + let mut text = stderr_buf.lock().expect("stderr buffer poisoned"); + text.push_str(&String::from_utf8_lossy(&buffer[..n])); + // Keep only a bounded tail. + if text.len() > 16 * 1024 { + let cut = text.len() - 8 * 1024; + *text = text[cut..].to_string(); + } + } + } + } + }); + self.child = Some(ChildHandle { + child, + stdin, + stdout, + stderr, + }); + + // Wait for the bootstrap's `ready` frame. + match self.read_frame(Instant::now() + STARTUP_TIMEOUT).await { + Ok(frame) if frame.get("op").and_then(Value::as_str) == Some("ready") => Ok(()), + Ok(frame) => Err(self.broken(&format!("unexpected startup frame: {frame}"))), + Err(err) => { + self.kill().await; + Err(self.broken(&format!("did not become ready: {err}"))) + } + } + } + + async fn send_frame(&mut self, frame: Value) -> Result<()> { + let handle = self + .child + .as_mut() + .ok_or_else(|| TinyAgentsError::Capability("rlm interpreter not running".into()))?; + let mut line = frame.to_string(); + line.push('\n'); + if handle.stdin.write_all(line.as_bytes()).await.is_err() { + return Err(self.broken("stdin closed")); + } + let _ = handle.stdin.flush().await; + Ok(()) + } + + async fn read_frame(&mut self, deadline: Instant) -> Result { + let handle = self + .child + .as_mut() + .ok_or_else(|| TinyAgentsError::Capability("rlm interpreter not running".into()))?; + loop { + let remaining = deadline.saturating_duration_since(Instant::now()); + if remaining.is_zero() { + return Err(TinyAgentsError::Timeout( + "rlm cell exceeded its wall-clock timeout".to_string(), + )); + } + let mut line = String::new(); + let read = tokio::time::timeout(remaining, handle.stdout.read_line(&mut line)).await; + match read { + Err(_) => { + return Err(TinyAgentsError::Timeout( + "rlm cell exceeded its wall-clock timeout".to_string(), + )); + } + Ok(Err(err)) => return Err(self.broken(&format!("stdout read failed: {err}"))), + Ok(Ok(0)) => return Err(self.broken("exited unexpectedly")), + Ok(Ok(_)) => { + let trimmed = line.trim(); + if trimmed.is_empty() { + continue; + } + match serde_json::from_str::(trimmed) { + Ok(frame) => return Ok(frame), + // Non-protocol noise on stdout (a library printing + // directly to fd 1) is ignored rather than fatal. + Err(_) => continue, + } + } + } + } + } + + async fn kill(&mut self) { + if let Some(mut handle) = self.child.take() { + let _ = handle.child.kill().await; + } + } +} + +#[async_trait] +impl RlmInterpreter for ExternalInterpreter { + fn language(&self) -> &str { + &self.language + } + + fn usage_guide(&self) -> String { + match self.language.as_str() { + "javascript" => r#"Write JavaScript (Node vm context; globals persist across cells). Host functions: +- llm("prompt") or llm({model, prompt, system}) -> string // ask a sub-LLM +- tool("name", {arg: value, ...}) -> result // call a registered tool +- agent("name", "input") -> string // delegate to a sub-agent +- final_answer("...") // end the task with this answer +- console.log(x) // observe a value next turn +Capability failures throw RlmError; catch them with try/catch. The completion +value of the cell (its last expression) is echoed back to you."# + .to_string(), + _ => r#"Write Python (a persistent exec namespace; globals survive across cells). Host functions: +- llm("prompt") or llm(prompt, model=..., system=...) -> str # ask a sub-LLM +- tool("name", {"arg": value, ...}) -> result # call a registered tool +- agent("name", "input") -> str # delegate to a sub-agent +- final_answer("...") # end the task with this answer +- print(x) # observe a value next turn +Capability failures raise RlmError; catch them with try/except. If the last +statement of a cell is an expression, its value is echoed back to you."# + .to_string(), + } + } + + async fn set_variable(&mut self, name: &str, value: Value) -> Result<()> { + self.ensure_started().await?; + self.send_frame(json!({"op": "set_var", "name": name, "value": value})) + .await?; + let deadline = Instant::now() + STARTUP_TIMEOUT; + loop { + let frame = self.read_frame(deadline).await?; + if frame.get("op").and_then(Value::as_str) == Some("var_set") { + return Ok(()); + } + } + } + + async fn eval_cell(&mut self, code: &str, host: Arc) -> Result { + self.ensure_started().await?; + self.send_frame(json!({"op": "eval", "code": code})).await?; + let deadline = host + .deadline() + .unwrap_or_else(|| Instant::now() + DEFAULT_EXCHANGE_TIMEOUT); + + loop { + let frame = match self.read_frame(deadline).await { + Ok(frame) => frame, + Err(err) => { + // Fail closed: a cell that timed out (or broke the + // protocol) leaves the child in an unknown state. + self.kill().await; + return Err(err); + } + }; + match frame.get("op").and_then(Value::as_str) { + Some("call") => { + let call: HostCall = match serde_json::from_value( + frame.get("call").cloned().unwrap_or_default(), + ) { + Ok(call) => call, + Err(err) => { + self.send_frame(json!({ + "op": "call_result", + "ok": false, + "error": format!("malformed capability call: {err}"), + })) + .await?; + continue; + } + }; + let remaining = deadline.saturating_duration_since(Instant::now()); + let outcome = tokio::time::timeout(remaining, host.handle(call)).await; + match outcome { + Err(_) => { + self.kill().await; + return Err(TinyAgentsError::Timeout( + "rlm cell exceeded its wall-clock timeout".to_string(), + )); + } + Ok(Ok(value)) => { + self.send_frame( + json!({"op": "call_result", "ok": true, "value": value}), + ) + .await?; + } + Ok(Err(err)) if is_fatal(&err) => { + // Policy bound tripped: kill the child rather than + // letting the script observe its own limits. + self.kill().await; + return Err(err); + } + Ok(Err(err)) => { + self.send_frame(json!({ + "op": "call_result", + "ok": false, + "error": err.to_string(), + })) + .await?; + } + } + } + Some("result") => { + let value = frame.get("value").cloned().unwrap_or(Value::Null); + return Ok(CellEval { + stdout: frame + .get("stdout") + .and_then(Value::as_str) + .unwrap_or_default() + .to_string(), + value: (!value.is_null()).then_some(value), + error: frame + .get("error") + .and_then(Value::as_str) + .map(str::to_string), + }); + } + _ => continue, + } + } + } + + async fn shutdown(&mut self) -> Result<()> { + if self.child.is_some() { + let _ = self.send_frame(json!({"op": "shutdown"})).await; + self.kill().await; + } + Ok(()) + } +} + +impl Drop for ExternalInterpreter { + fn drop(&mut self) { + // `kill_on_drop(true)` reaps the child if shutdown was never called. + self.child = None; + } +} diff --git a/src/rlm/interpreter/mod.rs b/src/rlm/interpreter/mod.rs new file mode 100644 index 0000000..e6af212 --- /dev/null +++ b/src/rlm/interpreter/mod.rs @@ -0,0 +1,169 @@ +//! Interpreter backends for the RLM runtime. +//! +//! [`RlmInterpreter`] is the pluggable execution API: the session hands a +//! backend one code cell at a time plus the shared [`RlmHostApi`] handle, and +//! the backend returns a raw [`CellEval`]. State (variables, imports, +//! definitions) persists across cells within one backend instance, so the +//! driving model can build up a workspace incrementally like a notebook. +//! +//! Two backends ship built in: +//! +//! - [`rhai_cell::RhaiInterpreter`] — the embedded Rhai engine. Hermetic: no +//! filesystem, network, or process access; the registered capability +//! functions are its entire host surface. +//! - [`external::ExternalInterpreter`] — a child process (Python, Node, or +//! any command speaking the wire protocol). The binary is configuration, +//! so embedders choose the exact interpreter (virtualenv Python, Deno, a +//! containerized runner, …). +//! +//! Construct a backend from configuration with [`build_interpreter`]. + +pub mod external; +pub mod rhai_cell; + +use std::future::Future; +use std::sync::Arc; +use std::time::{Duration, Instant}; + +use async_trait::async_trait; +use serde_json::Value; + +use super::host::RlmHostApi; +use super::types::{InterpreterSpec, RlmCancelFlag}; +use crate::error::{Result, TinyAgentsError}; + +/// The raw output of evaluating one cell, before the session merges in the +/// host-side call records and final answer. +#[derive(Debug, Clone, Default, PartialEq)] +pub struct CellEval { + /// Captured print/console output. + pub stdout: String, + /// The cell's final expression value, if any. + pub value: Option, + /// A recoverable script error (exception / runtime error). + pub error: Option, +} + +/// A pluggable code-cell execution backend. +#[async_trait] +pub trait RlmInterpreter: Send { + /// The language cells are written in (`"rhai"`, `"python"`, …), used for + /// prompt rendering and code-fence extraction. + fn language(&self) -> &str; + + /// A prompt fragment teaching the driver model how to reach the host + /// capabilities *in this language* (function signatures, examples). + fn usage_guide(&self) -> String; + + /// Sets (or replaces) a global variable visible to subsequent cells. + /// + /// Used to inject the task context (`context`) without string-splicing + /// user data into script source. + async fn set_variable(&mut self, name: &str, value: Value) -> Result<()>; + + /// Evaluates one code cell against the host. + /// + /// Returns `Ok(CellEval)` for both success and *recoverable* script + /// errors (carried in [`CellEval::error`]); returns `Err` only for fatal + /// conditions (policy limits, timeout, cancellation, a dead backend). + async fn eval_cell(&mut self, code: &str, host: Arc) -> Result; + + /// Releases backend resources (kills a child process). Idempotent. + async fn shutdown(&mut self) -> Result<()>; +} + +/// Builds the interpreter backend described by an [`InterpreterSpec`]. +/// +/// `max_operations` bounds the embedded Rhai engine; external backends are +/// bounded by the cell deadline instead (a wedged child is killed). +pub fn build_interpreter( + spec: &InterpreterSpec, + max_operations: u64, +) -> Result> { + match spec { + InterpreterSpec::Rhai => Ok(Box::new(rhai_cell::RhaiInterpreter::new(max_operations))), + InterpreterSpec::Python { binary, args } => Ok(Box::new( + external::ExternalInterpreter::python(binary.as_deref(), args.clone()), + )), + InterpreterSpec::Javascript { binary, args } => Ok(Box::new( + external::ExternalInterpreter::javascript(binary.as_deref(), args.clone()), + )), + InterpreterSpec::Command { binary, args } => Ok(Box::new( + external::ExternalInterpreter::command(binary.clone(), args.clone()), + )), + } +} + +/// How often the watcher thread wakes to observe cancellation while a +/// capability call is in flight inside the blocking bridge. +const CANCEL_POLL_INTERVAL: Duration = Duration::from_millis(25); + +/// Why the watcher tripped a bounded bridge call. +enum BridgeStop { + Deadline, + Cancelled, +} + +/// Drives an async host-capability future to completion **synchronously**, +/// bounded by the cell deadline and the session cancel flag. +/// +/// This is the same fail-closed blocking bridge the `.ragsh` REPL uses (see +/// `repl::session::builtins`): the embedded Rhai engine is synchronous, so a +/// capability closure must block its thread — but never unboundedly. A +/// detached watcher thread races the future; if the deadline elapses or the +/// flag trips first, the future is dropped (cancelling the underlying +/// request) and a `Timeout` / `Cancelled` error is returned. +pub(super) fn bridge_block_on( + deadline: Option, + cancel: &RlmCancelFlag, + future: F, +) -> Result +where + F: Future>, +{ + if cancel.is_cancelled() { + return Err(TinyAgentsError::Cancelled); + } + if let Some(deadline) = deadline + && Instant::now() >= deadline + { + return Err(TinyAgentsError::Timeout( + "rlm cell deadline elapsed before a host capability call could start".to_string(), + )); + } + + let (tx, rx) = futures::channel::oneshot::channel::(); + let watcher_cancel = cancel.clone(); + std::thread::spawn(move || { + loop { + if tx.is_canceled() { + return; + } + if watcher_cancel.is_cancelled() { + let _ = tx.send(BridgeStop::Cancelled); + return; + } + match deadline { + Some(deadline) => { + let now = Instant::now(); + if now >= deadline { + let _ = tx.send(BridgeStop::Deadline); + return; + } + std::thread::sleep((deadline - now).min(CANCEL_POLL_INTERVAL)); + } + None => std::thread::sleep(CANCEL_POLL_INTERVAL), + } + } + }); + + match futures::executor::block_on(futures::future::select(Box::pin(future), rx)) { + futures::future::Either::Left((output, _watcher)) => output, + futures::future::Either::Right((stop, _fut)) => match stop { + Ok(BridgeStop::Cancelled) => Err(TinyAgentsError::Cancelled), + Ok(BridgeStop::Deadline) | Err(_) => Err(TinyAgentsError::Timeout( + "rlm cell deadline elapsed during a host capability call".to_string(), + )), + }, + } +} diff --git a/src/rlm/interpreter/rhai_cell.rs b/src/rlm/interpreter/rhai_cell.rs new file mode 100644 index 0000000..6dd2eeb --- /dev/null +++ b/src/rlm/interpreter/rhai_cell.rs @@ -0,0 +1,396 @@ +//! The embedded Rhai interpreter backend. +//! +//! The engine is built fresh for every cell over a **persistent scope**, so +//! variables survive across cells while the capability closures always see +//! the current cell's deadline. The engine has no filesystem, network, or +//! process access — the closures registered here are its entire host +//! surface, which is what makes this backend the hermetic default. +//! +//! Capability calls block the evaluating thread through the fail-closed +//! [`bridge_block_on`](super::bridge_block_on) adapter; the session drives +//! `eval_cell` inside [`tokio::task::spawn_blocking`], so blocking here never +//! starves the async runtime driving the underlying provider I/O. + +use std::sync::{Arc, Mutex}; +use std::time::Instant; + +use async_trait::async_trait; +use rhai::{Array, Dynamic, Engine, EvalAltResult, Map, Position, Scope}; +use serde_json::Value; + +use super::{CellEval, RlmInterpreter, bridge_block_on}; +use crate::error::{Result, TinyAgentsError}; +use crate::rlm::host::{RlmHostApi, is_fatal}; +use crate::rlm::types::HostCall; + +/// Sentinel runtime-error text `on_progress` terminates a cell with when the +/// wall-clock deadline elapses mid-script. +const DEADLINE_TOKEN: &str = "rlm cell exceeded its wall-clock timeout"; + +/// Sentinel runtime-error text `on_progress` terminates a cell with when the +/// session cancel flag trips mid-script. +const CANCELLED_TOKEN: &str = "rlm cell cancelled by host"; + +/// The embedded Rhai backend. See the [module docs](self). +pub struct RhaiInterpreter { + max_operations: u64, + scope: Scope<'static>, +} + +impl RhaiInterpreter { + /// Creates a backend bounded by `max_operations` Rhai operations per cell + /// (`0` means unlimited). + pub fn new(max_operations: u64) -> Self { + Self { + max_operations, + scope: Scope::new(), + } + } +} + +/// Shared per-cell buffers the capability closures write into. +#[derive(Default)] +struct CellState { + stdout: String, + /// A fatal host error (limit/timeout/cancel) stashed so `eval_cell` can + /// surface the precise crate error instead of its stringified form. + fatal: Option, +} + +type SharedCellState = Arc>; + +/// Dispatches one host call from a synchronous capability closure, blocking +/// through the bridge and splitting fatal from script-visible failures. +fn dispatch( + host: &Arc, + cell: &SharedCellState, + call: HostCall, +) -> std::result::Result> { + let deadline = host.deadline(); + let cancel = host.cancel_flag(); + match bridge_block_on(deadline, &cancel, host.handle(call)) { + Ok(value) => Ok(value), + Err(err) => { + let message = err.to_string(); + if is_fatal(&err) { + cell.lock().expect("cell state poisoned").fatal = Some(err); + } + Err(Box::new(EvalAltResult::ErrorRuntime( + Dynamic::from(message), + Position::NONE, + ))) + } + } +} + +// ── Dynamic ⇄ JSON conversion ─────────────────────────────────────────────── + +/// Converts a Rhai value into JSON. Opaque host types are stringified rather +/// than leaked. +pub(crate) fn dynamic_to_json(value: &Dynamic) -> Value { + if value.is_unit() { + Value::Null + } else if let Ok(b) = value.as_bool() { + Value::Bool(b) + } else if let Ok(i) = value.as_int() { + Value::from(i) + } else if let Ok(f) = value.as_float() { + serde_json::Number::from_f64(f) + .map(Value::Number) + .unwrap_or(Value::Null) + } else if let Some(s) = value.read_lock::() { + Value::String(s.to_string()) + } else if let Some(array) = value.read_lock::() { + Value::Array(array.iter().map(dynamic_to_json).collect()) + } else if let Some(map) = value.read_lock::() { + Value::Object( + map.iter() + .map(|(k, v)| (k.to_string(), dynamic_to_json(v))) + .collect(), + ) + } else { + Value::String(value.to_string()) + } +} + +/// Converts JSON into a Rhai value. +pub(crate) fn json_to_dynamic(value: &Value) -> Dynamic { + match value { + Value::Null => Dynamic::UNIT, + Value::Bool(b) => Dynamic::from(*b), + Value::Number(n) => { + if let Some(i) = n.as_i64() { + Dynamic::from(i) + } else { + Dynamic::from(n.as_f64().unwrap_or(0.0)) + } + } + Value::String(s) => Dynamic::from(s.clone()), + Value::Array(items) => { + Dynamic::from_array(items.iter().map(json_to_dynamic).collect::()) + } + Value::Object(map) => { + let mut out = Map::new(); + for (k, v) in map { + out.insert(k.clone().into(), json_to_dynamic(v)); + } + Dynamic::from_map(out) + } + } +} + +/// Builds the sandboxed engine for one cell, registering the capability +/// closures against `host` and the shared cell buffers. +fn build_engine(host: Arc, cell: SharedCellState, max_operations: u64) -> Engine { + let mut engine = Engine::new(); + engine.set_max_operations(max_operations); + + // Fail-closed mid-script enforcement: `on_progress` fires between Rhai + // statements/operations, catching runaway script loops with no host + // calls. In-flight capability calls are bounded separately by the + // blocking bridge. + let progress_host = host.clone(); + let progress_cell = cell.clone(); + engine.on_progress(move |_ops| { + if progress_host.cancel_flag().is_cancelled() { + return Some(Dynamic::from(CANCELLED_TOKEN.to_string())); + } + if progress_cell + .lock() + .expect("cell state poisoned") + .fatal + .is_some() + { + return Some(Dynamic::from(DEADLINE_TOKEN.to_string())); + } + match progress_host.deadline() { + Some(deadline) if Instant::now() >= deadline => { + Some(Dynamic::from(DEADLINE_TOKEN.to_string())) + } + _ => None, + } + }); + + // ── print / debug capture ── + let print_cell = cell.clone(); + engine.on_print(move |text| { + let mut state = print_cell.lock().expect("cell state poisoned"); + state.stdout.push_str(text); + state.stdout.push('\n'); + }); + let debug_cell = cell.clone(); + engine.on_debug(move |text, _source, _pos| { + let mut state = debug_cell.lock().expect("cell state poisoned"); + state.stdout.push_str(text); + state.stdout.push('\n'); + }); + + // ── llm(prompt) / llm(#{ model, prompt, system }) ── + let llm_host = host.clone(); + let llm_cell = cell.clone(); + engine.register_fn( + "llm", + move |prompt: &str| -> std::result::Result> { + let value = dispatch( + &llm_host, + &llm_cell, + HostCall::Llm { + model: None, + prompt: prompt.to_string(), + system: None, + }, + )?; + Ok(value.as_str().unwrap_or_default().to_string()) + }, + ); + let llm_map_host = host.clone(); + let llm_map_cell = cell.clone(); + engine.register_fn( + "llm", + move |params: Map| -> std::result::Result> { + let get = |key: &str| params.get(key).and_then(|d| d.clone().into_string().ok()); + let prompt = get("prompt").ok_or_else(|| { + Box::new(EvalAltResult::ErrorRuntime( + Dynamic::from("llm: missing `prompt`".to_string()), + Position::NONE, + )) + })?; + let value = dispatch( + &llm_map_host, + &llm_map_cell, + HostCall::Llm { + model: get("model"), + prompt, + system: get("system"), + }, + )?; + Ok(value.as_str().unwrap_or_default().to_string()) + }, + ); + + // ── tool(name) / tool(name, #{ ... }) ── + let tool_host = host.clone(); + let tool_cell = cell.clone(); + engine.register_fn( + "tool", + move |name: &str| -> std::result::Result> { + let value = dispatch( + &tool_host, + &tool_cell, + HostCall::Tool { + tool: name.to_string(), + arguments: Value::Null, + }, + )?; + Ok(json_to_dynamic(&value)) + }, + ); + let tool_args_host = host.clone(); + let tool_args_cell = cell.clone(); + engine.register_fn( + "tool", + move |name: &str, args: Map| -> std::result::Result> { + let arguments = dynamic_to_json(&Dynamic::from_map(args)); + let value = dispatch( + &tool_args_host, + &tool_args_cell, + HostCall::Tool { + tool: name.to_string(), + arguments, + }, + )?; + Ok(json_to_dynamic(&value)) + }, + ); + + // ── agent(name, input) ── + let agent_host = host.clone(); + let agent_cell = cell.clone(); + engine.register_fn( + "agent", + move |name: &str, input: &str| -> std::result::Result> { + let value = dispatch( + &agent_host, + &agent_cell, + HostCall::Agent { + agent: name.to_string(), + input: input.to_string(), + data: None, + }, + )?; + Ok(value.as_str().unwrap_or_default().to_string()) + }, + ); + + // ── final_answer(text) ── + let answer_host = host.clone(); + let answer_cell = cell.clone(); + engine.register_fn( + "final_answer", + move |text: &str| -> std::result::Result<(), Box> { + dispatch( + &answer_host, + &answer_cell, + HostCall::FinalAnswer { + answer: text.to_string(), + }, + )?; + Ok(()) + }, + ); + + engine +} + +#[async_trait] +impl RlmInterpreter for RhaiInterpreter { + fn language(&self) -> &str { + "rhai" + } + + fn usage_guide(&self) -> String { + r#"Write Rhai. Variables persist across cells. Host functions: +- llm(prompt) -> string // ask the default sub-LLM +- llm(#{ model: "name", prompt: "...", system: "..." }) -> string +- tool("name", #{ arg: value, ... }) -> result // call a registered tool +- agent("name", "input") -> string // delegate to a sub-agent +- final_answer("...") // end the task with this answer +- print(x) // observe a value in the next turn +Errors raised by capabilities can be caught with try/catch. The last +expression of a cell is echoed back to you as its value. + +Rhai syntax notes (Rhai is NOT JavaScript or Rust): +- There are NO tuples: return an array `[a, b]` or an object map `#{ k: v }`. +- Object maps are written `#{ name: value }` and indexed with `m.name` or `m["name"]`. +- Sub-arrays: `arr.extract(0..5)`; also `arr.len()`, `arr.push(x)`, `arr.filter(|x| ...)`, + `arr.map(|x| ...)`; loops: `for x in arr { ... }` and `for i in 0..n { ... }`. +- Strings concatenate with `+`; convert with `x.to_string()`; interpolate with `` `${x}` ``. +- `let` declares a mutable variable; statements end with `;`."# + .to_string() + } + + async fn set_variable(&mut self, name: &str, value: Value) -> Result<()> { + self.scope + .set_value(name.to_string(), json_to_dynamic(&value)); + Ok(()) + } + + async fn eval_cell(&mut self, code: &str, host: Arc) -> Result { + let cell: SharedCellState = Arc::new(Mutex::new(CellState::default())); + let engine = build_engine(host, cell.clone(), self.max_operations); + let mut scope = std::mem::take(&mut self.scope); + let code = code.to_string(); + + // Rhai is synchronous and the capability closures block through the + // bridge, so evaluate on the blocking pool to keep the async runtime + // (which drives the actual provider I/O) responsive. + let (scope_back, eval) = tokio::task::spawn_blocking(move || { + let result = engine.eval_with_scope::(&mut scope, &code); + (scope, result) + }) + .await + .map_err(|err| TinyAgentsError::Model(format!("rlm rhai eval task failed: {err}")))?; + self.scope = scope_back; + + let mut state = cell.lock().expect("cell state poisoned"); + if let Some(fatal) = state.fatal.take() { + return Err(fatal); + } + let stdout = std::mem::take(&mut state.stdout); + drop(state); + + match eval { + Ok(value) => { + let json = dynamic_to_json(&value); + Ok(CellEval { + stdout, + value: (!json.is_null()).then_some(json), + error: None, + }) + } + Err(err) => { + let message = err.to_string(); + if message.contains(DEADLINE_TOKEN) { + return Err(TinyAgentsError::Timeout(DEADLINE_TOKEN.to_string())); + } + if message.contains(CANCELLED_TOKEN) { + return Err(TinyAgentsError::Cancelled); + } + if matches!(*err, EvalAltResult::ErrorTooManyOperations(_)) { + return Err(TinyAgentsError::LimitExceeded( + "rlm cell exceeded the operation limit".to_string(), + )); + } + Ok(CellEval { + stdout, + value: None, + error: Some(message), + }) + } + } + } + + async fn shutdown(&mut self) -> Result<()> { + Ok(()) + } +} diff --git a/src/rlm/mod.rs b/src/rlm/mod.rs new file mode 100644 index 0000000..2024b48 --- /dev/null +++ b/src/rlm/mod.rs @@ -0,0 +1,60 @@ +//! Recursive-language-model (RLM) runtime: a driver model writes code cells +//! executed in a sandboxed interpreter whose only host surface is capability +//! calls back into the [`CapabilityRegistry`](crate::registry::CapabilityRegistry) +//! — sub-LLM queries, tools, and sub-agent delegation. Scripts can therefore +//! *recursively* call language models, which is what turns a code sandbox +//! into an RLM. +//! +//! ## The three layers +//! +//! 1. **[`RlmInterpreter`]** — the pluggable execution API ("the interpreter +//! exposed as an API"). Built-ins: the embedded Rhai engine (hermetic +//! sandbox, the default) and external Python / JavaScript processes +//! (binary + args are configuration; they speak a line-delimited JSON +//! wire protocol, see [`interpreter::external`]). +//! 2. **[`RlmSession`]** — one interpreter bound to one [`RlmHost`], with +//! every [`RlmPolicy`] limit enforced fail-closed per cell. Drive it +//! directly when embedding your own loop. +//! 3. **[`RlmRunner`]** — the model-driven loop: render a template into a +//! system prompt, let the driver model emit fenced code cells, execute, +//! feed observations back, stop on `final_answer(...)`. +//! +//! Everything a run needs is describable as one serde document +//! ([`RlmConfig`]), so external harnesses can define RLM behaviors as +//! configuration rather than code. +//! +//! ```no_run +//! use std::sync::Arc; +//! use tinyagents::registry::CapabilityRegistry; +//! use tinyagents::rlm::{RlmConfig, RlmRunner}; +//! +//! # async fn demo() -> tinyagents::Result<()> { +//! let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); +//! // registry.register_model("openai", Arc::new(model))?; … +//! let config = RlmConfig::from_json( +//! r#"{ "interpreter": {"kind": "rhai"}, "template": "general" }"#, +//! )?; +//! let mut runner = RlmRunner::from_config(config, Arc::new(registry), Arc::new(()))?; +//! let outcome = runner.run("How many primes are there below 1000?").await?; +//! println!("{:?}", outcome.answer); +//! # Ok(()) } +//! ``` + +mod host; +pub mod interpreter; +mod runner; +mod session; +pub mod templates; +mod types; + +#[cfg(test)] +mod test; + +pub use host::{CapabilityListing, RlmHost, RlmHostApi, is_fatal}; +pub use interpreter::{CellEval, RlmInterpreter, build_interpreter}; +pub use runner::{RlmRunner, extract_code_cell}; +pub use session::RlmSession; +pub use types::{ + CellOutcome, HostCall, InterpreterSpec, RlmCallKind, RlmCallRecord, RlmCancelFlag, RlmConfig, + RlmOutcome, RlmPolicy, RlmStep, RlmStopReason, RlmTemplate, TemplateSpec, +}; diff --git a/src/rlm/runner.rs b/src/rlm/runner.rs new file mode 100644 index 0000000..e862da3 --- /dev/null +++ b/src/rlm/runner.rs @@ -0,0 +1,240 @@ +//! The model-driven RLM loop: a driver model writes code cells, the session +//! executes them, and the observations flow back until the script (or the +//! model) produces a final answer. +//! +//! The loop deliberately drives [`ChatModel::invoke`] directly instead of +//! going through `AgentHarness`: the "tool" here is the whole sandboxed +//! interpreter, whose feedback protocol (code fence in, observation out) is +//! the RLM contract rather than a JSON tool call. + +use std::sync::Arc; + +use serde_json::Value; + +use super::host::{RlmHost, RlmHostApi}; +use super::session::RlmSession; +use super::templates; +use super::types::{RlmConfig, RlmOutcome, RlmStep, RlmStopReason}; +use crate::error::{Result, TinyAgentsError}; +use crate::harness::message::Message; +use crate::harness::model::ModelRequest; +use crate::registry::CapabilityRegistry; + +/// Extracts the first fenced code block from a model reply. +/// +/// Accepts ```` ``` ````, a bare ```` ``` ````, or any other fence +/// info string (models occasionally mislabel the language); returns `None` +/// when the reply contains no complete fence. +pub fn extract_code_cell(reply: &str) -> Option { + let fence_start = reply.find("```")?; + let after_fence = &reply[fence_start + 3..]; + let newline = after_fence.find('\n')?; + let body = &after_fence[newline + 1..]; + let fence_end = body.find("```")?; + let code = body[..fence_end].trim_end(); + (!code.trim().is_empty()).then(|| code.to_string()) +} + +/// Renders a cell outcome as the observation message fed back to the driver. +fn render_observation(outcome: &super::types::CellOutcome) -> String { + let mut out = String::new(); + if !outcome.stdout.is_empty() { + out.push_str("stdout:\n"); + out.push_str(&outcome.stdout); + if !outcome.stdout.ends_with('\n') { + out.push('\n'); + } + } + if let Some(value) = &outcome.value { + out.push_str(&format!("value: {value}\n")); + } + if let Some(error) = &outcome.error { + out.push_str(&format!("error: {error}\n")); + } + if out.is_empty() { + out.push_str("(cell produced no output)\n"); + } + out.push_str("Continue. Reply with the next code cell, or call final_answer(...) when done."); + out +} + +/// The model-driven RLM runner. Construct with [`RlmRunner::from_config`], +/// optionally inject a context, then [`run`](RlmRunner::run) a task. +pub struct RlmRunner { + registry: Arc>, + config: RlmConfig, + session: RlmSession, + driver_model: String, + system_prompt: String, +} + +impl RlmRunner { + /// Builds a runner from a config document, a capability registry, and the + /// application state capability calls run against. + pub fn from_config( + config: RlmConfig, + registry: Arc>, + state: Arc, + ) -> Result { + let driver_model = config + .driver_model + .clone() + .or_else(|| { + registry + .names(crate::registry::ComponentKind::Model) + .into_iter() + .next() + }) + .ok_or_else(|| { + TinyAgentsError::Validation( + "rlm: no driver model configured and no model registered".to_string(), + ) + })?; + let sub_model = config + .sub_model + .clone() + .unwrap_or_else(|| driver_model.clone()); + let host = Arc::new( + RlmHost::new(registry.clone(), state) + .with_policy(config.policy.clone()) + .with_default_model(sub_model), + ); + let session = RlmSession::new(&config.interpreter, host)?; + + let template = templates::resolve(&config.template)?; + let system_prompt = templates::render_system_prompt( + &template, + &session.language(), + &session.usage_guide(), + &session.host().capabilities(), + &config.policy, + ); + Ok(Self { + registry, + config, + session, + driver_model, + system_prompt, + }) + } + + /// The session, for injecting variables or inspecting call counts. + pub fn session_mut(&mut self) -> &mut RlmSession { + &mut self.session + } + + /// The rendered driver system prompt (for inspection/telemetry). + pub fn system_prompt(&self) -> &str { + &self.system_prompt + } + + /// Injects the task context as the `context` variable in the sandbox. + pub async fn set_context(&mut self, context: Value) -> Result<()> { + self.session.set_variable("context", context).await + } + + /// Runs the loop for one task until a final answer or the cell budget. + pub async fn run(&mut self, task: impl Into) -> Result { + let driver = self + .registry + .model(&self.driver_model) + .ok_or_else(|| TinyAgentsError::ModelNotFound(self.driver_model.clone()))?; + let state = self.session.host().app_state(); + + let mut messages = vec![ + Message::system(self.system_prompt.clone()), + Message::user(task.into()), + ]; + let mut steps: Vec = Vec::new(); + let mut driver_calls = 0usize; + // Set after a reply with no code fence: the driver gets one nudge to + // produce a cell before its prose is accepted as the answer. Models + // occasionally emit raw, unfenced code; without the nudge that code + // would be mistaken for a final answer. + let mut nudged = false; + + let outcome = loop { + if steps.len() >= self.config.policy.max_cells { + break RlmOutcome { + answer: None, + stop_reason: RlmStopReason::CellBudgetExhausted, + steps, + driver_calls, + sub_llm_calls: 0, + tool_calls: 0, + agent_calls: 0, + }; + } + + // `driver_model` is a *registry* name; the resolved ChatModel + // already knows its provider model id, so the request leaves + // `model` unset rather than leaking the registry name upstream. + let request = ModelRequest { + messages: messages.clone(), + ..Default::default() + }; + driver_calls += 1; + let response = driver.invoke(&state, request).await?; + let reply_text = Message::Assistant(response.message.clone()).text(); + messages.push(Message::Assistant(response.message)); + + let Some(code) = extract_code_cell(&reply_text) else { + if !nudged { + nudged = true; + messages.push(Message::user( + "Your reply contained no fenced code block, so nothing was executed. \ + Reply with exactly one fenced code block, or — if you are done — call \ + final_answer(...) from code. If you truly have nothing to run, repeat \ + your final answer in plain prose.", + )); + continue; + } + // Two fence-less replies in a row: accept the prose answer. + break RlmOutcome { + answer: Some(reply_text.trim().to_string()), + stop_reason: RlmStopReason::ModelAnswered, + steps, + driver_calls, + sub_llm_calls: 0, + tool_calls: 0, + agent_calls: 0, + }; + }; + nudged = false; + + let cell = self.session.eval(&code).await?; + let answered = cell.final_answer.clone(); + let observation = render_observation(&cell); + steps.push(RlmStep { + code, + outcome: cell, + }); + + if let Some(answer) = answered { + break RlmOutcome { + answer: Some(answer), + stop_reason: RlmStopReason::Answered, + steps, + driver_calls, + sub_llm_calls: 0, + tool_calls: 0, + agent_calls: 0, + }; + } + messages.push(Message::user(observation)); + }; + + let (llm, tool, agent) = self.session.host().call_counts(); + Ok(RlmOutcome { + sub_llm_calls: llm, + tool_calls: tool, + agent_calls: agent, + ..outcome + }) + } + + /// Releases interpreter resources. + pub async fn shutdown(&mut self) -> Result<()> { + self.session.shutdown().await + } +} diff --git a/src/rlm/session.rs b/src/rlm/session.rs new file mode 100644 index 0000000..ac3929c --- /dev/null +++ b/src/rlm/session.rs @@ -0,0 +1,140 @@ +//! An RLM session: one interpreter instance bound to one capability host. +//! +//! The session is the programmatic surface — "the interpreter exposed as an +//! API". Embedders that want direct cell execution (their own loop, a +//! notebook UI, a test) drive [`RlmSession::eval`] themselves; the +//! model-driven loop in [`super::runner`] is built on exactly this surface. + +use std::sync::Arc; +use std::time::Instant; + +use serde_json::Value; + +use super::host::{RlmHost, RlmHostApi}; +use super::interpreter::{RlmInterpreter, build_interpreter}; +use super::types::{CellOutcome, InterpreterSpec}; +use crate::error::{Result, TinyAgentsError}; + +/// Marker appended when captured output exceeds +/// [`RlmPolicy::max_output_bytes`] and is truncated. +const TRUNCATION_MARKER: &str = "\n… [output truncated by rlm policy]"; + +/// One sandboxed script workspace: a persistent interpreter plus the +/// capability host its cells call back into. +pub struct RlmSession { + interpreter: Box, + host: Arc>, + cells_run: usize, +} + +impl RlmSession { + /// Builds a session from an interpreter spec and a configured host. + pub fn new(spec: &InterpreterSpec, host: Arc>) -> Result { + let interpreter = build_interpreter(spec, host.policy().max_operations)?; + Ok(Self { + interpreter, + host, + cells_run: 0, + }) + } + + /// Builds a session over an already-constructed interpreter backend + /// (for custom [`RlmInterpreter`] implementations). + pub fn from_interpreter( + interpreter: Box, + host: Arc>, + ) -> Self { + Self { + interpreter, + host, + cells_run: 0, + } + } + + /// The capability host this session's cells call back into. + pub fn host(&self) -> &Arc> { + &self.host + } + + /// The language cells are written in. + pub fn language(&self) -> String { + self.interpreter.language().to_string() + } + + /// The interpreter-specific capability usage guide (prompt fragment). + pub fn usage_guide(&self) -> String { + self.interpreter.usage_guide() + } + + /// Number of cells evaluated so far. + pub fn cells_run(&self) -> usize { + self.cells_run + } + + /// Injects a global variable visible to subsequent cells — the safe way + /// to hand a task context to scripts without splicing it into source. + pub async fn set_variable(&mut self, name: &str, value: Value) -> Result<()> { + self.interpreter.set_variable(name, value).await + } + + /// Evaluates one code cell, enforcing the session policy fail-closed. + pub async fn eval(&mut self, code: &str) -> Result { + let policy = self.host.policy().clone(); + if self.host.cancel_flag().is_cancelled() { + return Err(TinyAgentsError::Cancelled); + } + if self.cells_run >= policy.max_cells { + return Err(TinyAgentsError::LimitExceeded(format!( + "cell limit ({}) exceeded", + policy.max_cells + ))); + } + if code.len() > policy.max_script_bytes { + return Err(TinyAgentsError::LimitExceeded(format!( + "cell source is {} bytes, over the {}-byte limit", + code.len(), + policy.max_script_bytes + ))); + } + self.cells_run += 1; + + self.host.begin_cell(); + let start = Instant::now(); + let eval = self + .interpreter + .eval_cell(code, self.host.clone() as Arc) + .await; + let (calls, final_answer) = self.host.end_cell(); + let mut eval = eval?; + + // Bound what flows back into the driver conversation. Truncation is + // explicit (marked) so the model knows it saw a prefix. + if eval.stdout.len() > policy.max_output_bytes { + eval.stdout.truncate(policy.max_output_bytes); + eval.stdout.push_str(TRUNCATION_MARKER); + } + if let Some(value) = &eval.value { + let rendered = value.to_string(); + if rendered.len() > policy.max_output_bytes { + let mut clipped = rendered; + clipped.truncate(policy.max_output_bytes); + clipped.push_str(TRUNCATION_MARKER); + eval.value = Some(Value::String(clipped)); + } + } + + Ok(CellOutcome { + stdout: eval.stdout, + value: eval.value, + error: eval.error, + calls, + final_answer, + elapsed: start.elapsed(), + }) + } + + /// Releases interpreter resources (kills an external child process). + pub async fn shutdown(&mut self) -> Result<()> { + self.interpreter.shutdown().await + } +} diff --git a/src/rlm/templates.rs b/src/rlm/templates.rs new file mode 100644 index 0000000..95f1f51 --- /dev/null +++ b/src/rlm/templates.rs @@ -0,0 +1,147 @@ +//! Built-in prompt templates for the RLM driver model, and the placeholder +//! renderer. +//! +//! A template is a plain [`RlmTemplate`] document, so external harnesses can +//! ship their own as JSON and select them through +//! [`TemplateSpec::Inline`](super::types::TemplateSpec). The built-ins cover +//! the three recurring RLM shapes: +//! +//! - **`general`** — solve a task with code, calling capabilities as needed. +//! - **`context-explorer`** — the recursive-LM pattern from the RLM +//! literature: a context too large to read at once is injected as the +//! `context` variable, and the model probes it programmatically (slice, +//! search, summarize slices via sub-LLM calls) instead of reading it all. +//! - **`orchestrator`** — decompose the task and delegate the pieces to +//! registered sub-agents, then synthesize. + +use super::host::CapabilityListing; +use super::types::{RlmPolicy, RlmTemplate, TemplateSpec}; +use crate::error::{Result, TinyAgentsError}; + +/// Shared preamble describing the cell loop contract. +const LOOP_CONTRACT: &str = r#"You are operating a sandboxed code notebook. Every reply MUST contain exactly +one fenced code block, opened with ```{{language}} and closed with ``` — code +outside a fence is never executed. The host executes the block and shows you +the captured output, the cell's value, and any error; variables persist +between cells. Keep cells small and observe intermediate results +instead of writing one giant script. When you have the answer, call +final_answer("...") from code. Never fabricate outputs you have not observed. + +{{usage}} + +Available capabilities: +{{capabilities}} + +Resource limits (exceeding them aborts the run): +{{limits}}"#; + +/// The built-in `general` template. +pub fn general() -> RlmTemplate { + RlmTemplate { + name: "general".to_string(), + system_prompt: format!( + "{LOOP_CONTRACT}\n\nSolve the user's task with code. Use sub-LLM calls for fuzzy \ + subproblems (summarization, extraction, judgment) and plain code for exact ones \ + (counting, filtering, arithmetic)." + ), + } +} + +/// The built-in `context-explorer` template (recursive-LM context probing). +pub fn context_explorer() -> RlmTemplate { + RlmTemplate { + name: "context-explorer".to_string(), + system_prompt: format!( + "{LOOP_CONTRACT}\n\nA variable named `context` holds material that is too large to \ + read in one glance. NEVER print all of it. Probe it programmatically: inspect its \ + length and structure first, then slice/search it, and delegate fuzzy analysis of \ + individual chunks to sub-LLM calls (`llm`). Combine the per-chunk findings with \ + code, then answer." + ), + } +} + +/// The built-in `orchestrator` template (sub-agent delegation). +pub fn orchestrator() -> RlmTemplate { + RlmTemplate { + name: "orchestrator".to_string(), + system_prompt: format!( + "{LOOP_CONTRACT}\n\nYou are an orchestrator. Decompose the task into independent \ + pieces, delegate each to the most suitable registered agent with `agent(name, \ + input)`, inspect their replies, iterate if a piece came back weak, and synthesize \ + the final answer yourself." + ), + } +} + +/// Resolves a [`TemplateSpec`] to a concrete template. +pub fn resolve(spec: &TemplateSpec) -> Result { + match spec { + TemplateSpec::Inline(template) => Ok(template.clone()), + TemplateSpec::Named(name) => match name.as_str() { + "general" => Ok(general()), + "context-explorer" => Ok(context_explorer()), + "orchestrator" => Ok(orchestrator()), + other => Err(TinyAgentsError::Validation(format!( + "unknown rlm template `{other}` (built-ins: general, context-explorer, \ + orchestrator)" + ))), + }, + } +} + +/// Renders a template's system prompt, substituting the documented +/// placeholders. +pub fn render_system_prompt( + template: &RlmTemplate, + language: &str, + usage: &str, + capabilities: &CapabilityListing, + policy: &RlmPolicy, +) -> String { + template + .system_prompt + .replace("{{language}}", language) + .replace("{{usage}}", usage) + .replace("{{capabilities}}", &render_capabilities(capabilities)) + .replace("{{limits}}", &render_limits(policy)) +} + +fn render_capabilities(listing: &CapabilityListing) -> String { + let mut out = String::new(); + if listing.models.is_empty() { + out.push_str("- models: (none registered)\n"); + } else { + out.push_str(&format!("- models: {}\n", listing.models.join(", "))); + } + if listing.tools.is_empty() { + out.push_str("- tools: (none registered)\n"); + } else { + out.push_str("- tools:\n"); + for (name, description) in &listing.tools { + if description.is_empty() { + out.push_str(&format!(" - {name}\n")); + } else { + out.push_str(&format!(" - {name}: {description}\n")); + } + } + } + if listing.agents.is_empty() { + out.push_str("- agents: (none registered)"); + } else { + out.push_str(&format!("- agents: {}", listing.agents.join(", "))); + } + out +} + +fn render_limits(policy: &RlmPolicy) -> String { + let timeout = policy + .cell_timeout + .map(|t| format!("{}s", t.as_secs())) + .unwrap_or_else(|| "none".to_string()); + format!( + "- max cells: {}\n- max sub-LLM calls: {}\n- max tool calls: {}\n- max agent calls: \ + {}\n- per-cell timeout: {timeout}", + policy.max_cells, policy.max_llm_calls, policy.max_tool_calls, policy.max_agent_calls + ) +} diff --git a/src/rlm/test.rs b/src/rlm/test.rs new file mode 100644 index 0000000..e84b3d9 --- /dev/null +++ b/src/rlm/test.rs @@ -0,0 +1,346 @@ +//! Module-local unit tests for the RLM runtime: config serialization, +//! template rendering, code-fence extraction, and the embedded-Rhai session +//! against deterministic capability doubles. + +use std::sync::Arc; + +use serde_json::json; + +use super::*; +use crate::harness::testkit::{FakeTool, ScriptedModel}; +use crate::registry::CapabilityRegistry; + +fn registry_with_mock(replies: Vec<&str>) -> Arc> { + let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); + registry + .register_model("mock", Arc::new(ScriptedModel::replies(replies))) + .expect("register model"); + registry + .register_tool(Arc::new(FakeTool::returning("echo", "echoed"))) + .expect("register tool"); + Arc::new(registry) +} + +fn rhai_session(registry: Arc>, policy: RlmPolicy) -> RlmSession<()> { + let host = Arc::new( + RlmHost::new(registry, Arc::new(())) + .with_policy(policy) + .with_default_model("mock"), + ); + RlmSession::new(&InterpreterSpec::Rhai, host).expect("build session") +} + +// ── Config round-trips ────────────────────────────────────────────────────── + +#[test] +fn config_round_trips_through_json() { + let config = RlmConfig { + interpreter: InterpreterSpec::Python { + binary: Some("python3".to_string()), + args: vec![], + }, + driver_model: Some("openai".to_string()), + sub_model: None, + policy: RlmPolicy::default(), + template: TemplateSpec::Named("context-explorer".to_string()), + }; + let json = config.to_json().expect("serialize"); + let back = RlmConfig::from_json(&json).expect("parse"); + assert_eq!(config, back); +} + +#[test] +fn minimal_config_document_parses_with_defaults() { + let config = RlmConfig::from_json(r#"{ "interpreter": {"kind": "rhai"} }"#).expect("parse"); + assert_eq!(config.interpreter, InterpreterSpec::Rhai); + assert_eq!(config.template, TemplateSpec::Named("general".to_string())); + assert_eq!(config.policy, RlmPolicy::default()); +} + +#[test] +fn host_call_wire_shape_is_stable() { + let call: HostCall = serde_json::from_value(json!({ + "capability": "llm", + "prompt": "hi", + "model": null, + "system": null, + })) + .expect("parse llm call"); + assert_eq!( + call, + HostCall::Llm { + model: None, + prompt: "hi".to_string(), + system: None + } + ); + let call: HostCall = serde_json::from_value(json!({ + "capability": "tool", + "tool": "echo", + })) + .expect("parse tool call without arguments"); + assert!(matches!(call, HostCall::Tool { arguments, .. } if arguments.is_null())); +} + +// ── Code-fence extraction ─────────────────────────────────────────────────── + +#[test] +fn extracts_fenced_code_and_rejects_prose() { + assert_eq!( + extract_code_cell("Let me try:\n```rhai\nlet x = 1;\nx\n```\nDone."), + Some("let x = 1;\nx".to_string()) + ); + assert_eq!( + extract_code_cell("```\nprint(1)\n```"), + Some("print(1)".to_string()) + ); + assert_eq!(extract_code_cell("no code here"), None); + assert_eq!(extract_code_cell("unterminated ```python\nprint(1)"), None); + assert_eq!(extract_code_cell("```rhai\n\n```"), None); +} + +// ── Template rendering ────────────────────────────────────────────────────── + +#[test] +fn renders_placeholders_into_the_system_prompt() { + let listing = CapabilityListing { + models: vec!["mock".to_string()], + tools: vec![("echo".to_string(), "Echoes.".to_string())], + agents: vec!["helper".to_string()], + }; + let prompt = templates::render_system_prompt( + &templates::general(), + "rhai", + "USAGE GUIDE", + &listing, + &RlmPolicy::default(), + ); + assert!(prompt.contains("```rhai")); + assert!(prompt.contains("USAGE GUIDE")); + assert!(prompt.contains("echo: Echoes.")); + assert!(prompt.contains("helper")); + assert!(!prompt.contains("{{")); +} + +#[test] +fn unknown_named_template_fails_closed() { + let err = templates::resolve(&TemplateSpec::Named("nope".to_string())); + assert!(err.is_err()); +} + +// ── Embedded Rhai session ─────────────────────────────────────────────────── + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rhai_cell_evaluates_and_persists_variables() { + let mut session = rhai_session(registry_with_mock(vec!["unused"]), RlmPolicy::default()); + let outcome = session.eval("let x = 21; x").await.expect("cell 1"); + assert_eq!(outcome.value, Some(json!(21))); + let outcome = session.eval("x * 2").await.expect("cell 2"); + assert_eq!(outcome.value, Some(json!(42))); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rhai_cell_calls_llm_tool_and_final_answer() { + let mut session = rhai_session( + registry_with_mock(vec!["sub-model says hi"]), + RlmPolicy::default(), + ); + let outcome = session + .eval( + r#" + let reply = llm("hello?"); + print(reply); + let echoed = tool("echo", #{ q: 7 }); + final_answer(reply); + "#, + ) + .await + .expect("cell"); + assert!(outcome.stdout.contains("sub-model says hi")); + assert_eq!(outcome.final_answer.as_deref(), Some("sub-model says hi")); + assert_eq!(outcome.calls.len(), 3); + assert_eq!(outcome.calls[0].kind, RlmCallKind::Llm); + assert_eq!(outcome.calls[1].kind, RlmCallKind::Tool); + assert_eq!(outcome.calls[2].kind, RlmCallKind::FinalAnswer); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn script_error_is_recoverable_not_fatal() { + let mut session = rhai_session(registry_with_mock(vec![]), RlmPolicy::default()); + let outcome = session.eval("this is not rhai ][").await.expect("cell"); + assert!(outcome.error.is_some()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn unknown_tool_error_is_catchable_in_script() { + let mut session = rhai_session(registry_with_mock(vec![]), RlmPolicy::default()); + let outcome = session + .eval(r#"try { tool("missing") } catch (e) { print("caught: " + e); } "ok""#) + .await + .expect("cell"); + assert!(outcome.stdout.contains("caught")); + assert_eq!(outcome.value, Some(json!("ok"))); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn llm_call_limit_is_fatal_and_aborts_the_cell() { + let policy = RlmPolicy { + max_llm_calls: 1, + ..RlmPolicy::default() + }; + let mut session = rhai_session(registry_with_mock(vec!["one", "two"]), policy); + let err = session + .eval(r#"llm("first"); llm("second")"#) + .await + .expect_err("limit must abort"); + assert!( + matches!(err, crate::error::TinyAgentsError::LimitExceeded(_)), + "got {err:?}" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn cell_budget_and_script_size_fail_closed() { + let policy = RlmPolicy { + max_cells: 1, + max_script_bytes: 16, + ..RlmPolicy::default() + }; + let mut session = rhai_session(registry_with_mock(vec![]), policy); + let err = session + .eval("1 + 1 + 1 + 1 + 1 + 1 + 1") + .await + .expect_err("script too large"); + assert!(matches!( + err, + crate::error::TinyAgentsError::LimitExceeded(_) + )); + session.eval("1").await.expect("first small cell"); + let err = session.eval("2").await.expect_err("cell budget"); + assert!(matches!( + err, + crate::error::TinyAgentsError::LimitExceeded(_) + )); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn oversized_stdout_is_truncated_with_a_marker() { + let policy = RlmPolicy { + max_output_bytes: 64, + ..RlmPolicy::default() + }; + let mut session = rhai_session(registry_with_mock(vec![]), policy); + let outcome = session + .eval(r#"for i in 0..100 { print("aaaaaaaaaaaaaaaaaaaaaaaa"); }"#) + .await + .expect("cell"); + assert!(outcome.stdout.len() < 200); + assert!(outcome.stdout.contains("truncated")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn context_variable_is_visible_to_scripts() { + let mut session = rhai_session(registry_with_mock(vec![]), RlmPolicy::default()); + session + .set_variable("context", json!({"words": ["alpha", "beta"]})) + .await + .expect("set context"); + let outcome = session.eval("context.words[1]").await.expect("cell"); + assert_eq!(outcome.value, Some(json!("beta"))); +} + +// ── The model-driven runner ───────────────────────────────────────────────── + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn runner_loops_until_final_answer() { + // Cell 1 computes and prints; cell 2 answers with the observed value. + let registry = registry_with_mock(vec![ + "Let me compute.\n```rhai\nlet x = 6 * 7;\nprint(x);\nx\n```", + "Now I know.\n```rhai\nfinal_answer(\"the answer is 42\")\n```", + ]); + let config = RlmConfig { + driver_model: Some("mock".to_string()), + ..RlmConfig::default() + }; + let mut runner = RlmRunner::from_config(config, registry, Arc::new(())).expect("build runner"); + let outcome = runner.run("multiply 6 by 7").await.expect("run"); + assert_eq!(outcome.answer.as_deref(), Some("the answer is 42")); + assert_eq!(outcome.stop_reason, RlmStopReason::Answered); + assert_eq!(outcome.steps.len(), 2); + assert_eq!(outcome.driver_calls, 2); + assert!(outcome.steps[0].outcome.stdout.contains("42")); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn runner_nudges_once_then_accepts_prose_as_the_answer() { + // A fence-less reply first earns a nudge (it may be unfenced code, not an + // answer); only a second fence-less reply is accepted as prose. + let registry = registry_with_mock(vec!["The answer is 4.", "The answer is 4."]); + let config = RlmConfig { + driver_model: Some("mock".to_string()), + ..RlmConfig::default() + }; + let mut runner = RlmRunner::from_config(config, registry, Arc::new(())).expect("build runner"); + let outcome = runner.run("what is 2+2?").await.expect("run"); + assert_eq!(outcome.answer.as_deref(), Some("The answer is 4.")); + assert_eq!(outcome.stop_reason, RlmStopReason::ModelAnswered); + assert_eq!(outcome.driver_calls, 2); + assert!(outcome.steps.is_empty()); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn runner_recovers_a_cell_after_a_nudge() { + // Unfenced code first (would previously have been mistaken for an + // answer), fenced after the nudge, then a final answer. + let registry = registry_with_mock(vec![ + "let x = 6 * 7; x", + "```rhai\nlet x = 6 * 7;\nx\n```", + "```rhai\nfinal_answer(\"42\")\n```", + ]); + let config = RlmConfig { + driver_model: Some("mock".to_string()), + ..RlmConfig::default() + }; + let mut runner = RlmRunner::from_config(config, registry, Arc::new(())).expect("build runner"); + let outcome = runner.run("multiply 6 by 7").await.expect("run"); + assert_eq!(outcome.answer.as_deref(), Some("42")); + assert_eq!(outcome.stop_reason, RlmStopReason::Answered); + assert_eq!(outcome.steps.len(), 2); + assert_eq!(outcome.driver_calls, 3); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn runner_stops_at_the_cell_budget() { + // The driver keeps emitting cells and never answers. + let cells: Vec<&str> = vec!["```rhai\n1\n```"; 4]; + let registry = registry_with_mock(cells); + let config = RlmConfig { + driver_model: Some("mock".to_string()), + policy: RlmPolicy { + max_cells: 2, + ..RlmPolicy::default() + }, + ..RlmConfig::default() + }; + let mut runner = RlmRunner::from_config(config, registry, Arc::new(())).expect("build runner"); + let outcome = runner.run("loop forever").await.expect("run"); + assert_eq!(outcome.answer, None); + assert_eq!(outcome.stop_reason, RlmStopReason::CellBudgetExhausted); + assert_eq!(outcome.steps.len(), 2); +} + +// ── Cancellation ──────────────────────────────────────────────────────────── + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn pre_cancelled_session_refuses_cells() { + let cancel = RlmCancelFlag::new(); + cancel.cancel(); + let host = Arc::new( + RlmHost::new(registry_with_mock(vec![]), Arc::new(())) + .with_default_model("mock") + .with_cancel_flag(cancel), + ); + let mut session = RlmSession::new(&InterpreterSpec::Rhai, host).expect("session"); + let err = session.eval("1").await.expect_err("must refuse"); + assert!(matches!(err, crate::error::TinyAgentsError::Cancelled)); +} diff --git a/src/rlm/types.rs b/src/rlm/types.rs new file mode 100644 index 0000000..7fc2303 --- /dev/null +++ b/src/rlm/types.rs @@ -0,0 +1,418 @@ +//! Public data types for the recursive-language-model (RLM) runtime. +//! +//! Everything here is `serde`-serializable on purpose: an RLM run is meant to +//! be **config-driven**, so an external harness (a CLI, a service, another +//! agent runtime) can describe an entire run — interpreter choice, resource +//! policy, prompt template — as a JSON document and hand it to +//! [`RlmConfig::from_json`]. +//! +//! Logic lives in the sibling modules: the host capability boundary in +//! [`super::host`], the interpreter backends in [`super::interpreter`], the +//! session in [`super::session`], and the model-driven loop in +//! [`super::runner`]. + +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use serde::{Deserialize, Serialize}; +use serde_json::Value; + +use crate::error::{Result, TinyAgentsError}; + +// ── Interpreter selection ─────────────────────────────────────────────────── + +/// Which interpreter executes the code cells of an RLM session. +/// +/// The embedded [`Rhai`](InterpreterSpec::Rhai) engine is the default and the +/// only *hermetically* sandboxed choice: it has no filesystem, network, or +/// process access — the capability functions registered by the host are its +/// entire world. External interpreters run as a child **process** provided by +/// the embedding application (the binary and args are configuration, exactly +/// so a harness can point at a virtualenv Python, a Deno binary, a container +/// entrypoint, …); they speak the line-delimited JSON wire protocol described +/// in [`super::interpreter::external`]. The host still enforces every +/// [`RlmPolicy`] limit fail-closed (killing the child on violation), but the +/// child process itself has whatever OS access the embedder's environment +/// grants it — isolate it externally (container, seccomp, jail) when running +/// untrusted models. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize, Default)] +#[serde(tag = "kind", rename_all = "snake_case")] +pub enum InterpreterSpec { + /// The embedded Rhai engine (feature-default, hermetic sandbox). + #[default] + Rhai, + /// An external CPython-compatible interpreter. + Python { + /// The interpreter binary (defaults to `python3` on `PATH`). + #[serde(default, skip_serializing_if = "Option::is_none")] + binary: Option, + /// Extra arguments placed before the bootstrap `-c` program. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + args: Vec, + }, + /// An external Node.js-compatible JavaScript interpreter. + Javascript { + /// The interpreter binary (defaults to `node` on `PATH`). + #[serde(default, skip_serializing_if = "Option::is_none")] + binary: Option, + /// Extra arguments placed before the bootstrap `-e` program. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + args: Vec, + }, + /// An arbitrary command that already speaks the RLM wire protocol on its + /// stdin/stdout (for embedders that ship their own runner, e.g. a + /// container image or a jailed interpreter). + Command { + /// The command binary. + binary: String, + /// Arguments passed verbatim. + #[serde(default, skip_serializing_if = "Vec::is_empty")] + args: Vec, + }, +} + +impl InterpreterSpec { + /// The language name scripts are written in, used in prompts and code + /// fence extraction (```rhai / ```python / ```javascript). + pub fn language(&self) -> &'static str { + match self { + InterpreterSpec::Rhai => "rhai", + InterpreterSpec::Python { .. } => "python", + InterpreterSpec::Javascript { .. } => "javascript", + InterpreterSpec::Command { .. } => "python", + } + } +} + +// ── Policy ────────────────────────────────────────────────────────────────── + +/// Resource limits bounding an RLM session and its model-driven loop. +/// +/// Every limit is enforced **fail closed**: exceeding a bound aborts the cell +/// (and, for an external interpreter, kills the child process) instead of +/// silently truncating or running unbounded work. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(default)] +pub struct RlmPolicy { + /// Maximum code cells one [`super::RlmRunner::run`] loop may execute. + pub max_cells: usize, + /// Maximum source size, in bytes, of a single cell. + pub max_script_bytes: usize, + /// Maximum captured stdout + value size, in bytes, per cell. + pub max_output_bytes: usize, + /// Maximum sub-LLM (`llm`) calls per session. + pub max_llm_calls: usize, + /// Maximum `tool` calls per session. + pub max_tool_calls: usize, + /// Maximum sub-agent (`agent`) calls per session. + pub max_agent_calls: usize, + /// Maximum recursion depth for sub-agent calls, enforced through the + /// shared harness guard + /// ([`RunConfig::checked_child_depth`](crate::harness::context::RunConfig::checked_child_depth)). + pub max_depth: usize, + /// Wall-clock timeout per cell (script + in-flight capability calls). + #[serde(with = "humantime_millis")] + pub cell_timeout: Option, + /// Maximum Rhai operations per cell (embedded interpreter only; `0` + /// means unlimited). + pub max_operations: u64, +} + +impl Default for RlmPolicy { + fn default() -> Self { + Self { + max_cells: 16, + max_script_bytes: 64 * 1024, + max_output_bytes: 256 * 1024, + max_llm_calls: 64, + max_tool_calls: 128, + max_agent_calls: 32, + max_depth: 8, + cell_timeout: Some(Duration::from_secs(120)), + max_operations: 5_000_000, + } + } +} + +/// Serializes the optional cell timeout as integer milliseconds so an RLM +/// config is a plain JSON document (`"cell_timeout": 120000`). +mod humantime_millis { + use std::time::Duration; + + use serde::{Deserialize, Deserializer, Serializer}; + + pub fn serialize( + value: &Option, + serializer: S, + ) -> Result { + match value { + Some(duration) => serializer.serialize_some(&(duration.as_millis() as u64)), + None => serializer.serialize_none(), + } + } + + pub fn deserialize<'de, D: Deserializer<'de>>( + deserializer: D, + ) -> Result, D::Error> { + Ok(Option::::deserialize(deserializer)?.map(Duration::from_millis)) + } +} + +// ── Cancellation ──────────────────────────────────────────────────────────── + +/// A shared, sticky cancellation flag for an RLM session (the same contract as +/// the REPL's flag: once cancelled, the session refuses further work until a +/// fresh flag is installed). +#[derive(Clone, Debug, Default)] +pub struct RlmCancelFlag(Arc); + +impl RlmCancelFlag { + /// Creates a fresh, un-cancelled flag. + pub fn new() -> Self { + Self(Arc::new(AtomicBool::new(false))) + } + + /// Requests cancellation; idempotent, observed by every clone. + pub fn cancel(&self) { + self.0.store(true, Ordering::SeqCst); + } + + /// Returns whether cancellation has been requested. + pub fn is_cancelled(&self) -> bool { + self.0.load(Ordering::SeqCst) + } +} + +// ── The host-call boundary ────────────────────────────────────────────────── + +/// One capability call a script makes back into the host. +/// +/// This is the **entire** host surface a sandboxed script sees, across every +/// interpreter backend: the embedded Rhai closures build these values +/// directly, and the external wire protocol carries them as the `call` field +/// of a `{"op":"call"}` frame. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +#[serde(tag = "capability", rename_all = "snake_case")] +pub enum HostCall { + /// A sub-LLM query (`llm(...)` in scripts). + Llm { + /// Registry model name; `None` selects the session's default model. + #[serde(default, skip_serializing_if = "Option::is_none")] + model: Option, + /// The user prompt. + prompt: String, + /// Optional system prompt. + #[serde(default, skip_serializing_if = "Option::is_none")] + system: Option, + }, + /// A tool invocation (`tool(name, args)` in scripts). + Tool { + /// Registry tool name. + tool: String, + /// JSON arguments matching the tool's schema. + #[serde(default)] + arguments: Value, + }, + /// A sub-agent delegation (`agent(name, input)` in scripts). + Agent { + /// Registry agent name. + agent: String, + /// The prompt the child run is seeded with. + input: String, + /// Optional structured side-channel payload. + #[serde(default, skip_serializing_if = "Option::is_none")] + data: Option, + }, + /// The script's final answer (`final_answer(text)`); ends the run loop. + FinalAnswer { + /// The answer text handed back to the caller. + answer: String, + }, +} + +impl HostCall { + /// The capability name used in call records and telemetry. + pub fn name(&self) -> String { + match self { + HostCall::Llm { model, .. } => model.clone().unwrap_or_else(|| "default".to_string()), + HostCall::Tool { tool, .. } => tool.clone(), + HostCall::Agent { agent, .. } => agent.clone(), + HostCall::FinalAnswer { .. } => "final_answer".to_string(), + } + } + + /// The record kind for this call. + pub fn kind(&self) -> RlmCallKind { + match self { + HostCall::Llm { .. } => RlmCallKind::Llm, + HostCall::Tool { .. } => RlmCallKind::Tool, + HostCall::Agent { .. } => RlmCallKind::Agent, + HostCall::FinalAnswer { .. } => RlmCallKind::FinalAnswer, + } + } +} + +/// The kind of capability an [`RlmCallRecord`] describes. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RlmCallKind { + /// A sub-LLM query. + Llm, + /// A tool invocation. + Tool, + /// A sub-agent delegation. + Agent, + /// The final answer. + FinalAnswer, +} + +/// A record of one capability call a cell performed. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct RlmCallRecord { + /// Which capability kind was invoked. + pub kind: RlmCallKind, + /// The capability name (model, tool, or agent registry name). + pub name: String, + /// Structured detail about the call (argument summary, sizes). + pub detail: Value, + /// Wall-clock time the call took. + pub elapsed: Duration, +} + +// ── Cell + run outcomes ───────────────────────────────────────────────────── + +/// The structured result of evaluating one code cell. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct CellOutcome { + /// Captured print/console output, bounded by + /// [`RlmPolicy::max_output_bytes`]. + pub stdout: String, + /// The cell's final expression value, if it produced one. + pub value: Option, + /// A script-level error (exception / runtime error), when the cell + /// failed *recoverably* — the driving model sees this and may adapt. + pub error: Option, + /// Capability calls recorded during the cell, in order. + pub calls: Vec, + /// The final answer, if the cell called `final_answer(...)`. + pub final_answer: Option, + /// Wall-clock time the cell took to evaluate. + pub elapsed: Duration, +} + +/// Why an [`RlmOutcome`] run loop stopped. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)] +#[serde(rename_all = "snake_case")] +pub enum RlmStopReason { + /// A cell called `final_answer(...)`. + Answered, + /// The driver model replied with prose and no code cell; the prose is + /// taken as the answer. + ModelAnswered, + /// The [`RlmPolicy::max_cells`] budget was exhausted without an answer. + CellBudgetExhausted, +} + +/// One executed step of the model-driven loop: the code the driver model +/// wrote and what evaluating it produced. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct RlmStep { + /// The code cell the driver model emitted. + pub code: String, + /// The evaluation outcome fed back to the model. + pub outcome: CellOutcome, +} + +/// The result of one complete [`super::RlmRunner::run`] loop. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct RlmOutcome { + /// The final answer, when the run produced one. + pub answer: Option, + /// Why the loop stopped. + pub stop_reason: RlmStopReason, + /// Every executed step, in order (the full trajectory). + pub steps: Vec, + /// Driver-model calls made by the loop itself (excludes sub-LLM calls + /// made *by scripts*, which are counted in [`RlmOutcome::sub_llm_calls`]). + pub driver_calls: usize, + /// Sub-LLM calls scripts made through the `llm` capability. + pub sub_llm_calls: usize, + /// Tool calls scripts made through the `tool` capability. + pub tool_calls: usize, + /// Sub-agent calls scripts made through the `agent` capability. + pub agent_calls: usize, +} + +// ── Templates ─────────────────────────────────────────────────────────────── + +/// A named prompt scaffold for the driver model. +/// +/// The `system_prompt` may reference these placeholders, substituted at run +/// time by [`super::templates::render_system_prompt`]: +/// +/// - `{{language}}` — the interpreter language (`rhai`, `python`, …) +/// - `{{usage}}` — the interpreter-specific capability usage guide +/// - `{{capabilities}}` — the live model/tool/agent registry listing +/// - `{{limits}}` — a human-readable summary of the [`RlmPolicy`] +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +pub struct RlmTemplate { + /// The template name (used by [`TemplateSpec::Named`]). + pub name: String, + /// The system prompt scaffold with `{{placeholder}}` slots. + pub system_prompt: String, +} + +/// How a config selects its prompt template: one of the built-in named +/// templates, or a fully inline scaffold. +#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)] +#[serde(untagged)] +pub enum TemplateSpec { + /// A built-in template by name (`"general"`, `"context-explorer"`, + /// `"orchestrator"`). + Named(String), + /// An inline template document. + Inline(RlmTemplate), +} + +impl Default for TemplateSpec { + fn default() -> Self { + TemplateSpec::Named("general".to_string()) + } +} + +// ── Config ────────────────────────────────────────────────────────────────── + +/// A complete, serializable description of an RLM run — the document an +/// external harness hands to [`super::RlmRunner::from_config`]. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Default)] +#[serde(default)] +pub struct RlmConfig { + /// Which interpreter executes code cells. + pub interpreter: InterpreterSpec, + /// The registry name of the driver model that writes cells; `None` + /// selects the registry default. + #[serde(skip_serializing_if = "Option::is_none")] + pub driver_model: Option, + /// The registry name of the default sub-LLM scripts reach with + /// `llm(...)` when they don't name a model; `None` falls back to the + /// driver model. + #[serde(skip_serializing_if = "Option::is_none")] + pub sub_model: Option, + /// Resource limits for the session. + pub policy: RlmPolicy, + /// The driver prompt template. + pub template: TemplateSpec, +} + +impl RlmConfig { + /// Parses a config from a JSON document. + pub fn from_json(json: &str) -> Result { + serde_json::from_str(json).map_err(TinyAgentsError::Serialization) + } + + /// Serializes this config to pretty JSON. + pub fn to_json(&self) -> Result { + serde_json::to_string_pretty(self).map_err(TinyAgentsError::Serialization) + } +} diff --git a/tests/e2e_rlm.rs b/tests/e2e_rlm.rs new file mode 100644 index 0000000..8fcc2a4 --- /dev/null +++ b/tests/e2e_rlm.rs @@ -0,0 +1,269 @@ +//! End-to-end tests for the `rlm` surface (feature = "rlm"). +//! +//! Deterministic tests drive the embedded Rhai backend and the model-driven +//! runner against testkit doubles. The external-interpreter tests run against +//! a real `python3` / `node` from `PATH` and **skip gracefully** (early +//! return) when the binary is missing, mirroring the `live_*.rs` gating +//! convention — no network or API key is needed either way. + +#![cfg(feature = "rlm")] + +use std::sync::Arc; + +use serde_json::json; +use tinyagents::harness::message::Message; +use tinyagents::harness::runtime::AgentHarness; +use tinyagents::harness::testkit::{FakeTool, ScriptedModel}; +use tinyagents::registry::CapabilityRegistry; +use tinyagents::rlm::{ + InterpreterSpec, RlmConfig, RlmHost, RlmPolicy, RlmRunner, RlmSession, RlmStopReason, +}; +use tinyagents::{HarnessSubAgent, SubAgent}; + +fn registry_with_doubles(replies: Vec<&str>) -> Arc> { + let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); + registry + .register_model("mock", Arc::new(ScriptedModel::replies(replies))) + .expect("register model"); + registry + .register_tool(Arc::new(FakeTool::returning("lookup", "tool-result-42"))) + .expect("register tool"); + Arc::new(registry) +} + +fn session_for(spec: &InterpreterSpec, registry: Arc>) -> RlmSession<()> { + let host = Arc::new( + RlmHost::new(registry, Arc::new(())) + .with_policy(RlmPolicy::default()) + .with_default_model("mock"), + ); + RlmSession::new(spec, host).expect("build session") +} + +fn binary_available(binary: &str) -> bool { + std::process::Command::new(binary) + .arg("--version") + .stdout(std::process::Stdio::null()) + .stderr(std::process::Stdio::null()) + .status() + .map(|status| status.success()) + .unwrap_or(false) +} + +// ── Sub-agent delegation from inside a script ─────────────────────────────── + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn rhai_script_delegates_to_a_registered_subagent() { + let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); + registry + .register_model("mock", Arc::new(ScriptedModel::replies(vec!["unused"]))) + .expect("register model"); + + // A real harness-backed sub-agent with its own scripted model. + let mut child_harness: AgentHarness<()> = AgentHarness::new(); + child_harness + .register_model( + "child-model", + Arc::new(ScriptedModel::replies(vec!["report from the child agent"])), + ) + .set_default_model("child-model"); + let subagent = Arc::new(SubAgent::new( + "researcher", + "Investigates a question and reports back.", + Arc::new(child_harness), + )); + registry + .register_agent(Arc::new(HarnessSubAgent::new(subagent))) + .expect("register agent"); + + let mut session = session_for(&InterpreterSpec::Rhai, Arc::new(registry)); + let outcome = session + .eval(r#"let report = agent("researcher", "investigate X"); final_answer(report)"#) + .await + .expect("cell"); + assert_eq!( + outcome.final_answer.as_deref(), + Some("report from the child agent") + ); +} + +// ── External Python interpreter ───────────────────────────────────────────── + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn python_interpreter_runs_cells_and_calls_capabilities() { + if !binary_available("python3") { + eprintln!("skipping: python3 is not on PATH"); + return; + } + let spec = InterpreterSpec::Python { + binary: None, + args: vec![], + }; + let mut session = session_for(&spec, registry_with_doubles(vec!["sub-llm reply"])); + + // Globals persist across cells; last expression is the value. + let outcome = session.eval("x = 40\nx + 2").await.expect("cell 1"); + assert_eq!(outcome.value, Some(json!(42))); + + // Context injection without source splicing. + session + .set_variable("context", json!({"items": [1, 2, 3]})) + .await + .expect("set context"); + let outcome = session + .eval("len(context['items']) + x") + .await + .expect("cell 2"); + assert_eq!(outcome.value, Some(json!(43))); + + // Capability calls: llm + tool + final_answer through the wire protocol. + let outcome = session + .eval( + "reply = llm('hello?')\nprint(reply)\nr = tool('lookup', {'q': 1})\nfinal_answer(reply)", + ) + .await + .expect("cell 3"); + assert!(outcome.stdout.contains("sub-llm reply")); + assert_eq!(outcome.final_answer.as_deref(), Some("sub-llm reply")); + + // Script exceptions are recoverable and RlmError is catchable. + let outcome = session + .eval("try:\n tool('missing')\nexcept RlmError as e:\n print('caught', e)\n'ok'") + .await + .expect("cell 4"); + assert!(outcome.stdout.contains("caught")); + assert_eq!(outcome.value, Some(json!("ok"))); + + session.shutdown().await.expect("shutdown"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn python_cell_timeout_kills_the_child_fail_closed() { + if !binary_available("python3") { + eprintln!("skipping: python3 is not on PATH"); + return; + } + let registry = registry_with_doubles(vec![]); + let host = Arc::new( + RlmHost::new(registry, Arc::new(())) + .with_policy(RlmPolicy { + cell_timeout: Some(std::time::Duration::from_millis(400)), + ..RlmPolicy::default() + }) + .with_default_model("mock"), + ); + let mut session = RlmSession::new( + &InterpreterSpec::Python { + binary: None, + args: vec![], + }, + host, + ) + .expect("session"); + let started = std::time::Instant::now(); + let err = session + .eval("import time\ntime.sleep(60)") + .await + .expect_err("must time out"); + assert!( + matches!(err, tinyagents::TinyAgentsError::Timeout(_)), + "got {err:?}" + ); + assert!(started.elapsed() < std::time::Duration::from_secs(10)); +} + +// ── External JavaScript interpreter ───────────────────────────────────────── + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn javascript_interpreter_runs_cells_and_calls_capabilities() { + if !binary_available("node") { + eprintln!("skipping: node is not on PATH"); + return; + } + let spec = InterpreterSpec::Javascript { + binary: None, + args: vec![], + }; + let mut session = session_for(&spec, registry_with_doubles(vec!["js sub-llm reply"])); + + let outcome = session.eval("let x = 40; x + 2").await.expect("cell 1"); + assert_eq!(outcome.value, Some(json!(42))); + + session + .set_variable("context", json!("needle in a haystack")) + .await + .expect("set context"); + let outcome = session + .eval("context.includes('needle')") + .await + .expect("cell 2"); + assert_eq!(outcome.value, Some(json!(true))); + + let outcome = session + .eval( + "const reply = llm('hello?'); console.log(reply); tool('lookup', {q: 1}); final_answer(reply)", + ) + .await + .expect("cell 3"); + assert!(outcome.stdout.contains("js sub-llm reply")); + assert_eq!(outcome.final_answer.as_deref(), Some("js sub-llm reply")); + + session.shutdown().await.expect("shutdown"); +} + +// ── Config-driven runner over an external interpreter ─────────────────────── + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn runner_drives_a_python_session_from_a_json_config() { + if !binary_available("python3") { + eprintln!("skipping: python3 is not on PATH"); + return; + } + let registry = registry_with_doubles(vec![ + // The scripted "driver model" writes a python cell, then answers. + "```python\ntotal = sum(range(10))\nprint(total)\ntotal\n```", + "```python\nfinal_answer(f'the sum is {total}')\n```", + ]); + let config = RlmConfig::from_json( + r#"{ + "interpreter": {"kind": "python"}, + "driver_model": "mock", + "template": "general", + "policy": {"max_cells": 4, "cell_timeout": 30000} + }"#, + ) + .expect("parse config"); + let mut runner = RlmRunner::from_config(config, registry, Arc::new(())).expect("runner"); + let outcome = runner.run("sum the numbers below 10").await.expect("run"); + assert_eq!(outcome.answer.as_deref(), Some("the sum is 45")); + assert_eq!(outcome.stop_reason, RlmStopReason::Answered); + assert_eq!(outcome.steps.len(), 2); + assert!(outcome.steps[0].outcome.stdout.contains("45")); + runner.shutdown().await.expect("shutdown"); +} + +// ── Prompt surface sanity ─────────────────────────────────────────────────── + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn system_prompt_lists_live_capabilities_and_language() { + let registry = registry_with_doubles(vec![]); + let config = RlmConfig { + driver_model: Some("mock".to_string()), + ..RlmConfig::default() + }; + let runner = RlmRunner::from_config(config, registry, Arc::new(())).expect("runner"); + let prompt = runner.system_prompt(); + assert!(prompt.contains("```rhai")); + assert!(prompt.contains("lookup")); + assert!(prompt.contains("final_answer")); +} + +// ── Messages are ordinary harness messages ────────────────────────────────── + +#[test] +fn observation_shapes_are_plain_messages() { + // Guard against accidental coupling: the runner speaks in ordinary + // Message values, so any ChatModel implementation can drive it. + let m = Message::user("observation"); + assert_eq!(m.text(), "observation"); +} diff --git a/tests/live_rlm.rs b/tests/live_rlm.rs new file mode 100644 index 0000000..21215be --- /dev/null +++ b/tests/live_rlm.rs @@ -0,0 +1,102 @@ +//! Live, network-gated RLM tests (feature = "rlm"). +//! +//! Skips gracefully (early return, not a panic) when `OPENAI_API_KEY` is not +//! set, following the `live_*.rs` convention. Assertions are structural only +//! (an answer was produced, cells actually executed) — never on exact model +//! prose. + +#![cfg(feature = "rlm")] + +use std::sync::Arc; + +use tinyagents::harness::providers::openai::OpenAiModel; +use tinyagents::registry::CapabilityRegistry; +use tinyagents::rlm::{RlmConfig, RlmRunner, RlmStopReason}; + +fn live_registry() -> Option>> { + let _ = dotenvy::dotenv(); + if std::env::var("OPENAI_API_KEY").is_err() { + eprintln!("skipping: OPENAI_API_KEY is not set"); + return None; + } + let mut registry: CapabilityRegistry<()> = CapabilityRegistry::new(); + registry + .register_model("openai", Arc::new(OpenAiModel::from_env().expect("model"))) + .expect("register model"); + Some(Arc::new(registry)) +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn live_rhai_rlm_computes_with_code() { + let Some(registry) = live_registry() else { + return; + }; + let config = RlmConfig::from_json( + r#"{ + "interpreter": {"kind": "rhai"}, + "driver_model": "openai", + "template": "general", + "policy": {"max_cells": 6, "cell_timeout": 60000} + }"#, + ) + .expect("config"); + let mut runner = RlmRunner::from_config(config, registry, Arc::new(())).expect("runner"); + let outcome = runner + .run( + "Compute the exact sum of the cubes of the integers from 1 to 25 with code, then \ + return just that number as the final answer.", + ) + .await + .expect("run"); + // 1³+…+25³ = (25·26/2)² = 105625. + let answer = outcome.answer.expect("an answer"); + assert!(answer.contains("105625"), "unexpected answer: {answer}"); + assert!( + !outcome.steps.is_empty(), + "the model must have executed at least one cell" + ); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +async fn live_python_rlm_probes_an_injected_context() { + if std::process::Command::new("python3") + .arg("--version") + .output() + .is_err() + { + eprintln!("skipping: python3 is not on PATH"); + return; + } + let Some(registry) = live_registry() else { + return; + }; + let config = RlmConfig::from_json( + r#"{ + "interpreter": {"kind": "python"}, + "driver_model": "openai", + "template": "context-explorer", + "policy": {"max_cells": 6, "cell_timeout": 60000} + }"#, + ) + .expect("config"); + let mut runner = RlmRunner::from_config(config, registry, Arc::new(())).expect("runner"); + + // A context with a needle the model must find programmatically. + let mut lines: Vec = (0..500) + .map(|i| format!("log line {i}: heartbeat ok")) + .collect(); + lines[317] = "log line 317: FATAL disk failure on node srv-42".to_string(); + runner + .set_context(serde_json::json!(lines.join("\n"))) + .await + .expect("set context"); + + let outcome = runner + .run("Exactly one log line in `context` is not a heartbeat. Which node failed?") + .await + .expect("run"); + let answer = outcome.answer.expect("an answer"); + assert!(answer.contains("srv-42"), "unexpected answer: {answer}"); + assert_ne!(outcome.stop_reason, RlmStopReason::CellBudgetExhausted); + runner.shutdown().await.expect("shutdown"); +}