Files
optimclaw/crates/ironclaw_engine/src/executor/orchestrator.rs
T
[email protected]andClaude Opus 4.6 6d6b7fab0e feat(ui): show tool arguments in CLI and gateway
Add params_summary to ActionExecuted/ActionFailed events so the CLI
and gateway can display what tools are doing:

  ● http(https://api.github.com/repos/nearai/ironclaw/issues)
  ● web_search(latest AI news)
  ● memory_read(HEARTBEAT.md)

The summarize_params() helper extracts the most relevant argument
per tool type (URL for http, query for search, path for memory, etc.)
and truncates to 80 chars. Sensitive params are not included.

Router forwards the summary in both StatusUpdate (CLI/REPL) and
AppEvent (web gateway SSE) display names.

Co-Authored-By: Claude Opus 4.6 (1M context) <[email protected]>
2026-03-27 23:28:28 -07:00

1545 lines
54 KiB
Rust

//! Python orchestrator — the self-modifiable execution loop.
//!
//! Replaces the Rust `ExecutionLoop::run()` with versioned Python code
//! executed via Monty. The orchestrator is the "glue layer" between the
//! LLM and tools — tool dispatch, output formatting, state management,
//! truncation — all in Python, patchable by the self-improvement Mission.
//!
//! Host functions exposed to the orchestrator Python:
//! - `__llm_complete__` — make an LLM call
//! - `__execute_code_step__` — run user CodeAct code in a nested Monty VM
//! - `__execute_action__` — execute a single tool action
//! - `__check_signals__` — poll for stop/inject signals
//! - `__emit_event__` — broadcast a ThreadEvent
//! - `__add_message__` — append a message to the thread
//! - `__save_checkpoint__` — persist thread state
//! - `__transition_to__` — change thread state (validated)
//! - `__retrieve_docs__` — query memory docs
//! - `__check_budget__` — remaining tokens/time/USD
//! - `__get_actions__` — available tool definitions
use std::sync::Arc;
use std::collections::HashMap;
use monty::{
ExtFunctionResult, LimitedTracker, MontyObject, MontyRun, NameLookupResult, PrintWriter,
ResourceLimits, RunProgress,
};
use tracing::{debug, warn};
use crate::capability::lease::LeaseManager;
use crate::capability::policy::PolicyEngine;
use crate::memory::RetrievalEngine;
use crate::runtime::messaging::{SignalReceiver, ThreadOutcome, ThreadSignal};
use crate::traits::effect::{EffectExecutor, ThreadExecutionContext};
use crate::traits::llm::{LlmBackend, LlmCallConfig};
use crate::traits::store::Store;
use crate::types::error::EngineError;
use crate::types::event::{EventKind, ThreadEvent, summarize_params};
use crate::types::message::ThreadMessage;
use crate::types::project::ProjectId;
use crate::types::step::{StepId, TokenUsage};
use crate::types::thread::Thread;
use super::scripting::{execute_code, json_to_monty, monty_to_json, monty_to_string};
/// The compiled-in default orchestrator (v0).
const DEFAULT_ORCHESTRATOR: &str = include_str!("../../orchestrator/default.py");
/// Well-known title for orchestrator code in the Store.
pub const ORCHESTRATOR_TITLE: &str = "orchestrator:main";
/// Well-known tag for orchestrator code docs.
pub const ORCHESTRATOR_TAG: &str = "orchestrator_code";
/// Result of running the orchestrator.
pub struct OrchestratorResult {
/// The thread outcome parsed from the orchestrator's return value.
pub outcome: ThreadOutcome,
/// Total tokens used by LLM calls within the orchestrator.
pub tokens_used: TokenUsage,
}
/// Resource limits for the orchestrator VM.
fn orchestrator_limits() -> ResourceLimits {
ResourceLimits::new()
.max_duration(std::time::Duration::from_secs(300)) // 5 min (longer than user code)
.max_allocations(5_000_000)
.max_memory(128 * 1024 * 1024) // 128 MB
}
/// Maximum consecutive failures before auto-rollback.
const MAX_FAILURES_BEFORE_ROLLBACK: u64 = 3;
/// Well-known title for orchestrator failure tracking.
const FAILURE_TRACKER_TITLE: &str = "orchestrator:failures";
/// Load orchestrator code: runtime version from Store, or compiled-in default.
///
/// Checks the failure tracker — if the latest version has >= 3 consecutive
/// failures, falls back to the previous version (or compiled-in default).
pub async fn load_orchestrator(
store: Option<&Arc<dyn Store>>,
project_id: ProjectId,
) -> (String, u64) {
let Some(store) = store else {
debug!("using compiled-in default orchestrator (v0, no store)");
return (DEFAULT_ORCHESTRATOR.to_string(), 0);
};
let docs = match store.list_memory_docs(project_id).await {
Ok(d) => d,
Err(_) => {
debug!("using compiled-in default orchestrator (v0, store error)");
return (DEFAULT_ORCHESTRATOR.to_string(), 0);
}
};
// Find all orchestrator versions, sorted by version number descending
let mut versions: Vec<_> = docs
.iter()
.filter(|d| d.title == ORCHESTRATOR_TITLE && d.tags.contains(&ORCHESTRATOR_TAG.to_string()))
.collect();
versions.sort_by(|a, b| {
let va = a
.metadata
.get("version")
.and_then(|v| v.as_u64())
.unwrap_or(0);
let vb = b
.metadata
.get("version")
.and_then(|v| v.as_u64())
.unwrap_or(0);
vb.cmp(&va) // descending
});
if versions.is_empty() {
debug!("using compiled-in default orchestrator (v0)");
return (DEFAULT_ORCHESTRATOR.to_string(), 0);
}
// Check failure count for the latest version
let failures = load_failure_count(&docs);
for doc in &versions {
let version = doc
.metadata
.get("version")
.and_then(|v| v.as_u64())
.unwrap_or(1);
// Skip versions with too many failures (only check the latest)
if version
== versions[0]
.metadata
.get("version")
.and_then(|v| v.as_u64())
.unwrap_or(1)
&& failures >= MAX_FAILURES_BEFORE_ROLLBACK
{
warn!(
version,
failures, "orchestrator version has too many failures, skipping"
);
continue;
}
debug!(version, "loaded runtime orchestrator");
return (doc.content.clone(), version);
}
// All versions failed — fall back to compiled-in default
debug!("all orchestrator versions failed, using compiled-in default (v0)");
(DEFAULT_ORCHESTRATOR.to_string(), 0)
}
/// Record a failure for the current orchestrator version.
pub async fn record_orchestrator_failure(
store: &Arc<dyn Store>,
project_id: ProjectId,
version: u64,
) {
use crate::types::memory::{DocType, MemoryDoc};
let docs = store.list_memory_docs(project_id).await.unwrap_or_default();
let existing = docs.iter().find(|d| d.title == FAILURE_TRACKER_TITLE);
let mut tracker = if let Some(doc) = existing {
doc.clone()
} else {
MemoryDoc::new(project_id, DocType::Note, FAILURE_TRACKER_TITLE, "")
.with_tags(vec!["orchestrator_meta".to_string()])
};
// Store failure count as JSON in content: {"version": N, "count": M}
let current: serde_json::Value =
serde_json::from_str(&tracker.content).unwrap_or(serde_json::json!({}));
let current_version = current.get("version").and_then(|v| v.as_u64()).unwrap_or(0);
let current_count = current.get("count").and_then(|v| v.as_u64()).unwrap_or(0);
let new_count = if current_version == version {
current_count + 1
} else {
1 // new version, reset count
};
tracker.content = serde_json::json!({
"version": version,
"count": new_count,
})
.to_string();
tracker.updated_at = chrono::Utc::now();
if let Err(e) = store.save_memory_doc(&tracker).await {
warn!("failed to save orchestrator failure tracker: {e}");
}
debug!(version, count = new_count, "recorded orchestrator failure");
}
/// Reset the failure counter (called after successful execution).
pub async fn reset_orchestrator_failures(store: &Arc<dyn Store>, project_id: ProjectId) {
let docs = store.list_memory_docs(project_id).await.unwrap_or_default();
let existing = docs.iter().find(|d| d.title == FAILURE_TRACKER_TITLE);
if let Some(doc) = existing {
let mut tracker = doc.clone();
tracker.content = serde_json::json!({"version": 0, "count": 0}).to_string();
tracker.updated_at = chrono::Utc::now();
let _ = store.save_memory_doc(&tracker).await;
}
}
/// Load failure count for the latest orchestrator version.
fn load_failure_count(docs: &[crate::types::memory::MemoryDoc]) -> u64 {
docs.iter()
.find(|d| d.title == FAILURE_TRACKER_TITLE)
.and_then(|d| serde_json::from_str::<serde_json::Value>(&d.content).ok())
.and_then(|v| v.get("count").and_then(|c| c.as_u64()))
.unwrap_or(0)
}
/// Execute the orchestrator Python code with host function dispatch.
///
/// This is the core function that replaces `ExecutionLoop::run()`'s inner loop.
/// The orchestrator Python calls host functions via Monty's suspension mechanism,
/// and this function handles each suspension by delegating to the appropriate
/// Rust implementation.
#[allow(clippy::too_many_arguments)]
pub async fn execute_orchestrator(
code: &str,
thread: &mut Thread,
llm: &Arc<dyn LlmBackend>,
effects: &Arc<dyn EffectExecutor>,
leases: &Arc<LeaseManager>,
policy: &Arc<PolicyEngine>,
signal_rx: &mut SignalReceiver,
event_tx: Option<&tokio::sync::broadcast::Sender<ThreadEvent>>,
retrieval: Option<&RetrievalEngine>,
store: Option<&Arc<dyn Store>>,
persisted_state: &serde_json::Value,
) -> Result<OrchestratorResult, EngineError> {
let mut total_tokens = TokenUsage::default();
// Build context variables for the orchestrator
let (input_names, input_values) = build_orchestrator_inputs(thread, persisted_state);
// Parse and compile
let runner = match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
MontyRun::new(code.to_string(), "orchestrator.py", input_names)
})) {
Ok(Ok(runner)) => runner,
Ok(Err(e)) => {
return Err(EngineError::Effect {
reason: format!("Orchestrator parse error: {e}"),
});
}
Err(_) => {
return Err(EngineError::Effect {
reason: "Monty VM panicked during orchestrator parsing".into(),
});
}
};
// Start execution
let mut stdout = String::new();
let tracker = LimitedTracker::new(orchestrator_limits());
let run_result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
runner.start(input_values, tracker, PrintWriter::Collect(&mut stdout))
}));
let mut progress = match run_result {
Ok(Ok(p)) => p,
Ok(Err(e)) => {
return Err(EngineError::Effect {
reason: format!("Orchestrator runtime error: {e}"),
});
}
Err(_) => {
return Err(EngineError::Effect {
reason: "Monty VM panicked during orchestrator start".into(),
});
}
};
// Drive the orchestrator dispatch loop
let mut final_result: Option<serde_json::Value> = None;
loop {
match progress {
RunProgress::Complete(obj) => {
// Use FINAL result if set, otherwise fall back to VM return value
let result = if let Some(ref fr) = final_result {
fr.clone()
} else {
monty_to_json(&obj)
};
return Ok(OrchestratorResult {
outcome: parse_outcome(&result),
tokens_used: total_tokens,
});
}
RunProgress::FunctionCall(call) => {
let action_name = call.function_name.clone();
let args = &call.args;
let kwargs = &call.kwargs;
debug!(action = %action_name, "orchestrator: host function call");
let ext_result = match action_name.as_str() {
// FINAL(result) — orchestrator returns its outcome
"FINAL" => {
let val = args.first().map(monty_to_json).unwrap_or_default();
final_result = Some(val);
ExtFunctionResult::Return(MontyObject::None)
}
// __llm_complete__(messages, actions, config)
"__llm_complete__" => {
handle_llm_complete(
args,
kwargs,
thread,
llm,
effects,
leases,
&mut total_tokens,
)
.await
}
// __execute_code_step__(code, state)
"__execute_code_step__" => {
handle_execute_code_step(
args, kwargs, thread, llm, effects, leases, policy, event_tx,
)
.await
}
// __execute_action__(name, params, call_id=...)
"__execute_action__" => {
handle_execute_action(
args, kwargs, thread, effects, leases, policy, event_tx,
)
.await
}
// __check_signals__()
"__check_signals__" => handle_check_signals(signal_rx),
// __emit_event__(kind, **data)
"__emit_event__" => handle_emit_event(args, kwargs, thread, event_tx),
// __add_message__(role, content)
"__add_message__" => handle_add_message(args, kwargs, thread),
// __save_checkpoint__(state, counters)
"__save_checkpoint__" => handle_save_checkpoint(args, kwargs, thread),
// __transition_to__(state, reason)
"__transition_to__" => handle_transition_to(args, kwargs, thread),
// __retrieve_docs__(goal, max_docs)
"__retrieve_docs__" => {
handle_retrieve_docs(args, kwargs, thread, retrieval).await
}
// __check_budget__()"
"__check_budget__" => handle_check_budget(thread),
// __get_actions__()
"__get_actions__" => handle_get_actions(thread, effects, leases).await,
// __list_skills__(max_candidates, max_tokens)
"__list_skills__" => handle_list_skills(args, thread, store).await,
// __record_skill_usage__(doc_id, success)
"__record_skill_usage__" => handle_record_skill_usage(args, store).await,
// Unknown — let Monty resolve it (user-defined functions, builtins)
other => ExtFunctionResult::NotFound(other.to_string()),
};
// Resume the orchestrator VM
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
call.resume(ext_result, PrintWriter::Collect(&mut stdout))
})) {
Ok(Ok(p)) => progress = p,
Ok(Err(e)) => {
return Err(EngineError::Effect {
reason: format!("Orchestrator error after resume: {e}"),
});
}
Err(_) => {
return Err(EngineError::Effect {
reason: "Monty VM panicked during orchestrator resume".into(),
});
}
}
// If FINAL was called, the VM should complete on next iteration
if final_result.is_some() {
continue;
}
}
RunProgress::NameLookup(lookup) => {
// Undefined variable — resume with NameError
let name = lookup.name.clone();
debug!(name = %name, "orchestrator: unresolved name");
match std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
lookup.resume(
NameLookupResult::Undefined,
PrintWriter::Collect(&mut stdout),
)
})) {
Ok(Ok(p)) => progress = p,
Ok(Err(e)) => {
return Err(EngineError::Effect {
reason: format!("Orchestrator NameError '{name}': {e}"),
});
}
Err(_) => {
return Err(EngineError::Effect {
reason: format!("Monty panic on NameLookup '{name}'"),
});
}
}
}
RunProgress::OsCall(_) => {
return Err(EngineError::Effect {
reason: "Orchestrator attempted OS call (blocked)".into(),
});
}
RunProgress::ResolveFutures(_) => {
return Err(EngineError::Effect {
reason: "Orchestrator attempted async (not supported)".into(),
});
}
}
}
}
// ── Host function handlers ──────────────────────────────────
/// Handle `__llm_complete__(messages, actions, config)`.
///
/// Calls the LLM and returns the response as a dict:
/// `{type: "text"|"code"|"actions", content/code/calls: ..., usage: {...}}`
///
/// For `ActionCalls` responses, the assistant message with structured action_calls
/// is added directly to the thread (not by Python) so the LLM backend can convert
/// them to the provider-specific tool_calls format on the next call.
async fn handle_llm_complete(
_args: &[MontyObject],
_kwargs: &[(MontyObject, MontyObject)],
thread: &mut Thread,
llm: &Arc<dyn LlmBackend>,
effects: &Arc<dyn EffectExecutor>,
leases: &Arc<LeaseManager>,
total_tokens: &mut TokenUsage,
) -> ExtFunctionResult {
use crate::types::step::LlmResponse;
// Build messages from thread (the orchestrator's __add_message__ calls
// have already populated thread.messages)
let active_leases = leases.active_for_thread(thread.id).await;
let actions = effects
.available_actions(&active_leases)
.await
.unwrap_or_default();
let config = LlmCallConfig {
max_tokens: None,
temperature: None,
force_text: false,
depth: thread.config.depth,
metadata: HashMap::new(),
};
match llm.complete(&thread.messages, &actions, &config).await {
Ok(output) => {
total_tokens.input_tokens += output.usage.input_tokens;
total_tokens.output_tokens += output.usage.output_tokens;
let usage = serde_json::json!({
"input_tokens": output.usage.input_tokens,
"output_tokens": output.usage.output_tokens,
});
let result = match output.response {
LlmResponse::Text(text) => {
serde_json::json!({"type": "text", "content": text, "usage": usage})
}
LlmResponse::Code { code, .. } => {
serde_json::json!({"type": "code", "code": code, "usage": usage})
}
LlmResponse::ActionCalls { calls, content } => {
// Add the assistant message with structured action_calls so the
// LLM backend sees proper tool_calls on the next round-trip.
// Python must NOT call __add_message__("assistant_actions", ...) —
// the message is already on the thread.
thread.add_message(ThreadMessage::assistant_with_actions(
content,
calls.clone(),
));
let calls_json: Vec<serde_json::Value> = calls
.iter()
.map(|c| {
serde_json::json!({
"name": c.action_name,
"call_id": c.id,
"params": c.parameters,
})
})
.collect();
serde_json::json!({"type": "actions", "calls": calls_json, "usage": usage})
}
};
ExtFunctionResult::Return(json_to_monty(&result))
}
Err(e) => ExtFunctionResult::Error(monty::MontyException::new(
monty::ExcType::RuntimeError,
Some(format!("LLM call failed: {e}")),
)),
}
}
/// Handle `__execute_code_step__(code, state)`.
///
/// Runs user CodeAct code in a nested Monty VM with full tool dispatch.
/// Returns a dict with stdout, return_value, action_results, etc.
#[allow(clippy::too_many_arguments)]
async fn handle_execute_code_step(
args: &[MontyObject],
_kwargs: &[(MontyObject, MontyObject)],
thread: &mut Thread,
llm: &Arc<dyn LlmBackend>,
effects: &Arc<dyn EffectExecutor>,
leases: &Arc<LeaseManager>,
policy: &Arc<PolicyEngine>,
event_tx: Option<&tokio::sync::broadcast::Sender<ThreadEvent>>,
) -> ExtFunctionResult {
let code = match args.first() {
Some(obj) => monty_to_string(obj),
None => {
return ExtFunctionResult::Error(monty::MontyException::new(
monty::ExcType::TypeError,
Some("__execute_code_step__ requires a code string".into()),
));
}
};
let state = args
.get(1)
.map(monty_to_json)
.unwrap_or(serde_json::json!({}));
let exec_ctx = ThreadExecutionContext {
thread_id: thread.id,
thread_type: thread.thread_type,
project_id: thread.project_id,
user_id: "orchestrator".into(),
step_id: StepId::new(),
};
// Run user code in a nested Monty VM (same pattern as rlm_query)
match Box::pin(execute_code(
&code,
thread,
llm,
effects,
leases,
policy,
&exec_ctx,
&[],
&state,
))
.await
{
Ok(result) => {
// Broadcast events from code execution to the thread and event channel.
// Without this, ActionExecuted events from CodeAct tool calls are lost
// and never appear in traces.
for event_kind in &result.events {
let event = ThreadEvent::new(thread.id, event_kind.clone());
if let Some(tx) = event_tx {
let _ = tx.send(event.clone());
}
thread.events.push(event);
}
thread.updated_at = chrono::Utc::now();
let action_results: Vec<serde_json::Value> = result
.action_results
.iter()
.map(|r| {
serde_json::json!({
"action_name": r.action_name,
"output": r.output,
"is_error": r.is_error,
"duration_ms": r.duration.as_millis(),
})
})
.collect();
let result_json = serde_json::json!({
"return_value": result.return_value,
"stdout": result.stdout,
"action_results": action_results,
"final_answer": result.final_answer,
"had_error": result.had_error,
"need_approval": result.need_approval.as_ref().map(|na| {
match na {
ThreadOutcome::NeedApproval { action_name, call_id, parameters } => {
serde_json::json!({
"action_name": action_name,
"call_id": call_id,
"parameters": parameters,
})
}
_ => serde_json::Value::Null,
}
}),
});
ExtFunctionResult::Return(json_to_monty(&result_json))
}
Err(e) => ExtFunctionResult::Error(monty::MontyException::new(
monty::ExcType::RuntimeError,
Some(format!("Code execution failed: {e}")),
)),
}
}
/// Handle `__execute_action__(name, params, call_id=...)`.
///
/// Single source of truth for action execution. Performs:
/// 1. Lease lookup
/// 2. Policy check
/// 3. Lease consumption
/// 4. Action execution via EffectExecutor
/// 5. Event emission (ActionExecuted/ActionFailed)
/// 6. Message addition (ActionResult with correct call_id)
///
/// Python only needs to check the returned `need_approval` flag.
async fn handle_execute_action(
args: &[MontyObject],
kwargs: &[(MontyObject, MontyObject)],
thread: &mut Thread,
effects: &Arc<dyn EffectExecutor>,
leases: &Arc<LeaseManager>,
policy: &Arc<PolicyEngine>,
event_tx: Option<&tokio::sync::broadcast::Sender<ThreadEvent>>,
) -> ExtFunctionResult {
let name = match extract_string_arg(args, kwargs, "name", 0) {
Some(n) => n,
None => {
return ExtFunctionResult::Error(monty::MontyException::new(
monty::ExcType::TypeError,
Some("__execute_action__ requires a name argument".into()),
));
}
};
let params = args
.get(1)
.map(monty_to_json)
.unwrap_or(serde_json::json!({}));
let call_id = extract_string_kwarg(kwargs, "call_id").unwrap_or_default();
let exec_ctx = ThreadExecutionContext {
thread_id: thread.id,
thread_type: thread.thread_type,
project_id: thread.project_id,
user_id: "orchestrator".into(),
step_id: StepId::new(),
};
// Helper: emit event and add ActionResult message to thread
let emit_and_record = |thread: &mut Thread,
event_tx: Option<&tokio::sync::broadcast::Sender<ThreadEvent>>,
event_kind: EventKind,
call_id: &str,
action_name: &str,
output: &serde_json::Value| {
let event = ThreadEvent::new(thread.id, event_kind);
if let Some(tx) = event_tx {
let _ = tx.send(event.clone());
}
thread.events.push(event);
thread.updated_at = chrono::Utc::now();
thread.add_message(ThreadMessage::action_result(
call_id,
action_name,
output.to_string(),
));
};
// 1. Find lease for this action
let lease = match leases.find_lease_for_action(thread.id, &name).await {
Some(l) => l,
None => {
let error = format!("No lease for action '{name}'");
let output = serde_json::json!({"error": &error});
emit_and_record(
thread,
event_tx,
EventKind::ActionFailed {
step_id: exec_ctx.step_id,
action_name: name.clone(),
call_id: call_id.clone(),
error,
params_summary: None,
},
&call_id,
&name,
&output,
);
let result = serde_json::json!({
"output": output,
"is_error": true,
});
return ExtFunctionResult::Return(json_to_monty(&result));
}
};
// 2. Check policy
let action_def = effects
.available_actions(std::slice::from_ref(&lease))
.await
.ok()
.and_then(|actions| actions.into_iter().find(|a| a.name == name));
if let Some(ref ad) = action_def {
match policy.evaluate(ad, &lease, &[]) {
crate::capability::policy::PolicyDecision::Deny { reason } => {
let output = serde_json::json!({"error": format!("Denied: {reason}")});
emit_and_record(
thread,
event_tx,
EventKind::ActionFailed {
step_id: exec_ctx.step_id,
action_name: name.clone(),
call_id: call_id.clone(),
error: reason,
params_summary: None,
},
&call_id,
&name,
&output,
);
let result = serde_json::json!({
"output": output,
"is_error": true,
});
return ExtFunctionResult::Return(json_to_monty(&result));
}
crate::capability::policy::PolicyDecision::RequireApproval { .. } => {
let result = serde_json::json!({
"need_approval": true,
"action_name": name,
});
return ExtFunctionResult::Return(json_to_monty(&result));
}
crate::capability::policy::PolicyDecision::Allow => {}
}
}
// 3. Consume a lease use
if let Err(e) = leases.consume_use(lease.id).await {
debug!(error = %e, "lease consumption failed (non-fatal)");
}
// 4. Execute
let ps = summarize_params(&name, &params);
match effects
.execute_action(&name, params, &lease, &exec_ctx)
.await
{
Ok(r) => {
emit_and_record(
thread,
event_tx,
EventKind::ActionExecuted {
step_id: exec_ctx.step_id,
action_name: name.clone(),
call_id: call_id.clone(),
duration_ms: r.duration.as_millis() as u64,
params_summary: ps.clone(),
},
&call_id,
&name,
&r.output,
);
let result = serde_json::json!({
"action_name": r.action_name,
"output": r.output,
"is_error": r.is_error,
"duration_ms": r.duration.as_millis(),
});
ExtFunctionResult::Return(json_to_monty(&result))
}
Err(e) => {
let output = serde_json::json!({"error": e.to_string()});
emit_and_record(
thread,
event_tx,
EventKind::ActionFailed {
step_id: exec_ctx.step_id,
action_name: name.clone(),
call_id: call_id.clone(),
error: e.to_string(),
params_summary: ps,
},
&call_id,
&name,
&output,
);
let result = serde_json::json!({
"output": output,
"is_error": true,
});
ExtFunctionResult::Return(json_to_monty(&result))
}
}
}
/// Handle `__check_signals__()`.
fn handle_check_signals(signal_rx: &mut SignalReceiver) -> ExtFunctionResult {
match signal_rx.try_recv() {
Ok(ThreadSignal::Stop) | Ok(ThreadSignal::Suspend) => {
ExtFunctionResult::Return(MontyObject::String("stop".into()))
}
Ok(ThreadSignal::InjectMessage(msg)) => {
let result = serde_json::json!({"inject": msg.content});
ExtFunctionResult::Return(json_to_monty(&result))
}
Ok(ThreadSignal::Resume) | Ok(ThreadSignal::ChildCompleted { .. }) => {
ExtFunctionResult::Return(MontyObject::None)
}
Err(_) => ExtFunctionResult::Return(MontyObject::None),
}
}
/// Handle `__emit_event__(kind, **data)`.
fn handle_emit_event(
args: &[MontyObject],
kwargs: &[(MontyObject, MontyObject)],
thread: &mut Thread,
event_tx: Option<&tokio::sync::broadcast::Sender<ThreadEvent>>,
) -> ExtFunctionResult {
let kind_str = args.first().map(monty_to_string).unwrap_or_default();
let kind = match kind_str.as_str() {
"step_started" => {
let _step = extract_u64_kwarg(kwargs, "step").unwrap_or(0);
EventKind::StepStarted {
step_id: StepId::new(),
}
}
"step_completed" => {
let input = extract_u64_kwarg(kwargs, "input_tokens").unwrap_or(0);
let output = extract_u64_kwarg(kwargs, "output_tokens").unwrap_or(0);
// Increment step count (mirrors the old Rust loop's step_count += 1)
thread.step_count += 1;
// Track token usage
thread.total_tokens_used += input + output;
EventKind::StepCompleted {
step_id: StepId::new(),
tokens: TokenUsage {
input_tokens: input,
output_tokens: output,
..Default::default()
},
}
}
"action_executed" => {
let action_name = extract_string_kwarg(kwargs, "action_name").unwrap_or_default();
let call_id = extract_string_kwarg(kwargs, "call_id").unwrap_or_default();
EventKind::ActionExecuted {
step_id: StepId::new(),
action_name,
call_id,
duration_ms: 0,
params_summary: None,
}
}
"action_failed" => {
let action_name = extract_string_kwarg(kwargs, "action_name").unwrap_or_default();
let call_id = extract_string_kwarg(kwargs, "call_id").unwrap_or_default();
let error = extract_string_kwarg(kwargs, "error").unwrap_or_default();
EventKind::ActionFailed {
step_id: StepId::new(),
action_name,
call_id,
error,
params_summary: None,
}
}
"skill_activated" => {
let names_str = extract_string_kwarg(kwargs, "skill_names").unwrap_or_default();
let skill_names: Vec<String> = names_str
.split(',')
.map(|s| s.trim().to_string())
.filter(|s| !s.is_empty())
.collect();
EventKind::SkillActivated { skill_names }
}
_ => {
debug!(kind = %kind_str, "orchestrator: unknown event kind, skipping");
return ExtFunctionResult::Return(MontyObject::None);
}
};
let event = ThreadEvent::new(thread.id, kind);
if let Some(tx) = event_tx {
let _ = tx.send(event.clone());
}
thread.events.push(event);
thread.updated_at = chrono::Utc::now();
ExtFunctionResult::Return(MontyObject::None)
}
/// Handle `__add_message__(role, content)`.
///
/// ActionResult messages are NOT added here — they are handled by
/// `__execute_action__` which is the single source of truth for
/// action execution, event emission, and message recording.
fn handle_add_message(
args: &[MontyObject],
_kwargs: &[(MontyObject, MontyObject)],
thread: &mut Thread,
) -> ExtFunctionResult {
let role = args.first().map(monty_to_string).unwrap_or_default();
let content = args.get(1).map(monty_to_string).unwrap_or_default();
match role.as_str() {
"user" => thread.add_message(ThreadMessage::user(&content)),
"assistant" => thread.add_message(ThreadMessage::assistant(&content)),
"system" => thread.add_message(ThreadMessage::system(&content)),
"system_append" => {
// Append to existing system message (for doc injection)
if let Some(msg) = thread
.messages
.iter_mut()
.find(|m| m.role == crate::types::message::MessageRole::System)
{
msg.content.push_str("\n\n");
msg.content.push_str(&content);
}
}
_ => {
thread.add_message(ThreadMessage::user(&content));
}
}
ExtFunctionResult::Return(MontyObject::None)
}
/// Handle `__save_checkpoint__(state, counters)`.
fn handle_save_checkpoint(
args: &[MontyObject],
_kwargs: &[(MontyObject, MontyObject)],
thread: &mut Thread,
) -> ExtFunctionResult {
let state = args
.first()
.map(monty_to_json)
.unwrap_or(serde_json::json!({}));
let counters = args
.get(1)
.map(monty_to_json)
.unwrap_or(serde_json::json!({}));
if let Some(metadata) = thread.metadata.as_object_mut() {
metadata.insert(
"runtime_checkpoint".into(),
serde_json::json!({
"persisted_state": state,
"nudge_count": counters.get("nudge_count").and_then(|v| v.as_u64()).unwrap_or(0),
"consecutive_errors": counters.get("consecutive_errors").and_then(|v| v.as_u64()).unwrap_or(0),
"compaction_count": counters.get("compaction_count").and_then(|v| v.as_u64()).unwrap_or(0),
}),
);
}
thread.updated_at = chrono::Utc::now();
ExtFunctionResult::Return(MontyObject::None)
}
/// Handle `__transition_to__(state, reason)`.
fn handle_transition_to(
args: &[MontyObject],
_kwargs: &[(MontyObject, MontyObject)],
thread: &mut Thread,
) -> ExtFunctionResult {
let state_str = args.first().map(monty_to_string).unwrap_or_default();
let reason = args.get(1).map(monty_to_string);
let target = match state_str.as_str() {
"running" => crate::types::thread::ThreadState::Running,
"completed" => crate::types::thread::ThreadState::Completed,
"failed" => crate::types::thread::ThreadState::Failed,
"waiting" => crate::types::thread::ThreadState::Waiting,
"suspended" => crate::types::thread::ThreadState::Suspended,
other => {
return ExtFunctionResult::Error(monty::MontyException::new(
monty::ExcType::ValueError,
Some(format!("Unknown thread state: {other}")),
));
}
};
match thread.transition_to(target, reason) {
Ok(()) => ExtFunctionResult::Return(MontyObject::None),
Err(e) => ExtFunctionResult::Error(monty::MontyException::new(
monty::ExcType::RuntimeError,
Some(format!("State transition failed: {e}")),
)),
}
}
/// Handle `__retrieve_docs__(goal, max_docs)`.
async fn handle_retrieve_docs(
args: &[MontyObject],
_kwargs: &[(MontyObject, MontyObject)],
thread: &Thread,
retrieval: Option<&RetrievalEngine>,
) -> ExtFunctionResult {
let retrieval = match retrieval {
Some(r) => r,
None => return ExtFunctionResult::Return(json_to_monty(&serde_json::json!([]))),
};
let goal = args.first().map(monty_to_string).unwrap_or_default();
let max_docs = args
.get(1)
.and_then(|v| match v {
MontyObject::Int(i) => Some(*i as usize),
_ => None,
})
.unwrap_or(5);
match retrieval
.retrieve_context(thread.project_id, &goal, max_docs)
.await
{
Ok(docs) => {
let docs_json: Vec<serde_json::Value> = docs
.iter()
.map(|d| {
serde_json::json!({
"type": format!("{:?}", d.doc_type),
"title": d.title,
"content": d.content,
})
})
.collect();
ExtFunctionResult::Return(json_to_monty(&serde_json::json!(docs_json)))
}
Err(e) => {
warn!("retrieve_docs failed: {e}");
ExtFunctionResult::Return(json_to_monty(&serde_json::json!([])))
}
}
}
/// Handle `__check_budget__()`.
fn handle_check_budget(thread: &Thread) -> ExtFunctionResult {
let tokens_remaining = thread
.config
.max_tokens_total
.map(|max| max.saturating_sub(thread.total_tokens_used))
.unwrap_or(u64::MAX);
let time_remaining_ms = thread
.config
.max_duration
.map(|dur| {
let elapsed = chrono::Utc::now()
.signed_duration_since(thread.created_at)
.num_milliseconds()
.max(0) as u64;
dur.as_millis() as u64 - elapsed.min(dur.as_millis() as u64)
})
.unwrap_or(u64::MAX);
let usd_remaining = thread
.config
.max_budget_usd
.map(|max| (max - thread.total_cost_usd).max(0.0));
let result = serde_json::json!({
"tokens_remaining": tokens_remaining,
"time_remaining_ms": time_remaining_ms,
"usd_remaining": usd_remaining,
});
ExtFunctionResult::Return(json_to_monty(&result))
}
/// Handle `__get_actions__()`.
async fn handle_get_actions(
thread: &Thread,
effects: &Arc<dyn EffectExecutor>,
leases: &Arc<LeaseManager>,
) -> ExtFunctionResult {
let active_leases = leases.active_for_thread(thread.id).await;
match effects.available_actions(&active_leases).await {
Ok(actions) => {
let actions_json: Vec<serde_json::Value> = actions
.iter()
.map(|a| {
serde_json::json!({
"name": a.name,
"description": a.description,
"params": a.parameters_schema,
})
})
.collect();
ExtFunctionResult::Return(json_to_monty(&serde_json::json!(actions_json)))
}
Err(e) => {
warn!("get_actions failed: {e}");
ExtFunctionResult::Return(json_to_monty(&serde_json::json!([])))
}
}
}
/// Handle `__list_skills__()`.
///
/// Loads all `DocType::Skill` MemoryDocs from the project and returns them
/// as a list of Python dicts. The Python orchestrator handles scoring,
/// selection, and injection — Rust just provides data access.
async fn handle_list_skills(
_args: &[MontyObject],
thread: &Thread,
store: Option<&Arc<dyn Store>>,
) -> ExtFunctionResult {
let Some(store) = store else {
return ExtFunctionResult::Return(json_to_monty(&serde_json::json!([])));
};
let docs = match store.list_memory_docs(thread.project_id).await {
Ok(docs) => docs,
Err(e) => {
debug!("__list_skills__: failed to load docs: {e}");
return ExtFunctionResult::Return(json_to_monty(&serde_json::json!([])));
}
};
let skills: Vec<serde_json::Value> = docs
.into_iter()
.filter(|d| d.doc_type == crate::types::memory::DocType::Skill)
.map(|d| {
serde_json::json!({
"doc_id": d.id.0.to_string(),
"title": d.title,
"content": d.content,
"metadata": d.metadata,
})
})
.collect();
ExtFunctionResult::Return(json_to_monty(&serde_json::json!(skills)))
}
/// Handle `__record_skill_usage__(doc_id, success)`.
///
/// Records that a skill was used in this thread. Called by the Python
/// orchestrator after skill-assisted execution completes.
async fn handle_record_skill_usage(
args: &[MontyObject],
store: Option<&Arc<dyn Store>>,
) -> ExtFunctionResult {
let Some(store) = store else {
return ExtFunctionResult::Return(MontyObject::None);
};
let doc_id_str = args.first().map(monty_to_string).unwrap_or_default();
let success = args
.get(1)
.map(|o| matches!(o, MontyObject::Bool(true)))
.unwrap_or(false);
let Ok(uuid) = uuid::Uuid::parse_str(&doc_id_str) else {
debug!("__record_skill_usage__: invalid doc_id: {doc_id_str}");
return ExtFunctionResult::Return(MontyObject::None);
};
let tracker = crate::capability::skill_tracker::SkillTracker::new(Arc::clone(store));
if let Err(e) = tracker
.record_usage(crate::types::memory::DocId(uuid), success)
.await
{
debug!("__record_skill_usage__: failed: {e}");
}
ExtFunctionResult::Return(MontyObject::None)
}
// ── Helpers ─────────────────────────────────────────────────
/// Build the context variables injected into the orchestrator Python.
fn build_orchestrator_inputs(
thread: &Thread,
persisted_state: &serde_json::Value,
) -> (Vec<String>, Vec<MontyObject>) {
let names = vec![
"context".into(),
"goal".into(),
"actions".into(),
"state".into(),
"config".into(),
];
// Build context (message history)
let context: Vec<serde_json::Value> = thread
.messages
.iter()
.map(|m| {
serde_json::json!({
"role": format!("{:?}", m.role),
"content": m.content,
})
})
.collect();
// Build config
let config = serde_json::json!({
"max_iterations": thread.config.max_iterations,
"max_tool_intent_nudges": thread.config.max_tool_intent_nudges,
"enable_tool_intent_nudge": thread.config.enable_tool_intent_nudge,
"max_consecutive_errors": thread.config.max_consecutive_errors,
"max_tokens_total": thread.config.max_tokens_total,
"max_budget_usd": thread.config.max_budget_usd,
"model_context_limit": thread.config.model_context_limit,
"enable_compaction": thread.config.enable_compaction,
"depth": thread.config.depth,
"max_depth": thread.config.max_depth,
"step_count": thread.step_count,
});
let values = vec![
json_to_monty(&serde_json::json!(context)),
MontyObject::String(thread.goal.clone()),
json_to_monty(&serde_json::json!([])), // actions loaded dynamically via __get_actions__
json_to_monty(persisted_state),
json_to_monty(&config),
];
(names, values)
}
/// Parse the orchestrator's return value into a ThreadOutcome.
fn parse_outcome(result: &serde_json::Value) -> ThreadOutcome {
let outcome = result
.get("outcome")
.and_then(|v| v.as_str())
.unwrap_or("completed");
match outcome {
"completed" => ThreadOutcome::Completed {
response: result
.get("response")
.and_then(|v| v.as_str())
.map(String::from),
},
"stopped" => ThreadOutcome::Stopped,
"max_iterations" => ThreadOutcome::MaxIterations,
"failed" => ThreadOutcome::Failed {
error: result
.get("error")
.and_then(|v| v.as_str())
.unwrap_or("unknown error")
.to_string(),
},
"need_approval" => ThreadOutcome::NeedApproval {
action_name: result
.get("action_name")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
call_id: result
.get("call_id")
.and_then(|v| v.as_str())
.unwrap_or("")
.to_string(),
parameters: result
.get("parameters")
.cloned()
.unwrap_or(serde_json::json!({})),
},
_ => ThreadOutcome::Completed { response: None },
}
}
fn extract_string_arg(
args: &[MontyObject],
kwargs: &[(MontyObject, MontyObject)],
name: &str,
position: usize,
) -> Option<String> {
for (k, v) in kwargs {
if let MontyObject::String(key) = k
&& key == name
{
return Some(monty_to_string(v));
}
}
args.get(position).map(monty_to_string)
}
fn extract_string_kwarg(kwargs: &[(MontyObject, MontyObject)], name: &str) -> Option<String> {
for (k, v) in kwargs {
if let MontyObject::String(key) = k
&& key == name
{
return Some(monty_to_string(v));
}
}
None
}
fn extract_u64_kwarg(kwargs: &[(MontyObject, MontyObject)], name: &str) -> Option<u64> {
for (k, v) in kwargs {
if let MontyObject::String(key) = k
&& key == name
&& let MontyObject::Int(i) = v
{
return Some(*i as u64);
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use crate::types::memory::{DocType, MemoryDoc};
use crate::types::project::ProjectId;
#[tokio::test]
async fn load_orchestrator_without_store_returns_default() {
let (code, version) = load_orchestrator(None, ProjectId::new()).await;
assert_eq!(version, 0);
assert!(code.contains("run_loop"));
assert!(code.contains("__llm_complete__"));
}
#[tokio::test]
async fn load_orchestrator_with_runtime_version() {
let project_id = ProjectId::new();
let mut doc = MemoryDoc::new(
project_id,
DocType::Note,
ORCHESTRATOR_TITLE,
"custom_orchestrator_code()",
)
.with_tags(vec![ORCHESTRATOR_TAG.to_string()]);
doc.metadata = serde_json::json!({"version": 1});
let store = Arc::new(crate::tests::InMemoryStore::with_docs(vec![doc]));
let (code, version) = load_orchestrator(Some(&(store as Arc<dyn Store>)), project_id).await;
assert_eq!(version, 1);
assert!(code.contains("custom_orchestrator_code"));
}
#[tokio::test]
async fn load_orchestrator_picks_highest_version() {
let project_id = ProjectId::new();
let mut doc_v1 = MemoryDoc::new(project_id, DocType::Note, ORCHESTRATOR_TITLE, "v1_code()")
.with_tags(vec![ORCHESTRATOR_TAG.to_string()]);
doc_v1.metadata = serde_json::json!({"version": 1});
let mut doc_v3 = MemoryDoc::new(project_id, DocType::Note, ORCHESTRATOR_TITLE, "v3_code()")
.with_tags(vec![ORCHESTRATOR_TAG.to_string()]);
doc_v3.metadata = serde_json::json!({"version": 3});
let mut doc_v2 = MemoryDoc::new(project_id, DocType::Note, ORCHESTRATOR_TITLE, "v2_code()")
.with_tags(vec![ORCHESTRATOR_TAG.to_string()]);
doc_v2.metadata = serde_json::json!({"version": 2});
let store = Arc::new(crate::tests::InMemoryStore::with_docs(vec![
doc_v1, doc_v3, doc_v2,
]));
let (code, version) = load_orchestrator(Some(&(store as Arc<dyn Store>)), project_id).await;
assert_eq!(version, 3);
assert!(code.contains("v3_code"));
}
#[tokio::test]
async fn rollback_after_max_failures() {
let project_id = ProjectId::new();
// Create v2 orchestrator
let mut doc_v2 =
MemoryDoc::new(project_id, DocType::Note, ORCHESTRATOR_TITLE, "v2_buggy()")
.with_tags(vec![ORCHESTRATOR_TAG.to_string()]);
doc_v2.metadata = serde_json::json!({"version": 2});
// Create v1 orchestrator (fallback)
let mut doc_v1 =
MemoryDoc::new(project_id, DocType::Note, ORCHESTRATOR_TITLE, "v1_stable()")
.with_tags(vec![ORCHESTRATOR_TAG.to_string()]);
doc_v1.metadata = serde_json::json!({"version": 1});
// Create failure tracker showing v2 has 3 failures
let tracker = MemoryDoc::new(
project_id,
DocType::Note,
FAILURE_TRACKER_TITLE,
r#"{"version": 2, "count": 3}"#,
)
.with_tags(vec!["orchestrator_meta".to_string()]);
let store = Arc::new(crate::tests::InMemoryStore::with_docs(vec![
doc_v2, doc_v1, tracker,
]));
let (code, version) = load_orchestrator(Some(&(store as Arc<dyn Store>)), project_id).await;
// Should skip v2 (too many failures) and load v1
assert_eq!(version, 1);
assert!(code.contains("v1_stable"));
}
#[tokio::test]
async fn rollback_to_default_when_all_versions_fail() {
let project_id = ProjectId::new();
// Single version with 3 failures
let mut doc_v1 =
MemoryDoc::new(project_id, DocType::Note, ORCHESTRATOR_TITLE, "v1_broken()")
.with_tags(vec![ORCHESTRATOR_TAG.to_string()]);
doc_v1.metadata = serde_json::json!({"version": 1});
let tracker = MemoryDoc::new(
project_id,
DocType::Note,
FAILURE_TRACKER_TITLE,
r#"{"version": 1, "count": 5}"#,
)
.with_tags(vec!["orchestrator_meta".to_string()]);
let store = Arc::new(crate::tests::InMemoryStore::with_docs(vec![
doc_v1, tracker,
]));
let (code, version) = load_orchestrator(Some(&(store as Arc<dyn Store>)), project_id).await;
// Should fall back to compiled-in default (v0)
assert_eq!(version, 0);
assert!(code.contains("run_loop"));
}
#[tokio::test]
async fn record_and_reset_failures() {
let project_id = ProjectId::new();
let store: Arc<dyn Store> = Arc::new(crate::tests::InMemoryStore::with_docs(vec![]));
// Record 3 failures
record_orchestrator_failure(&store, project_id, 2).await;
record_orchestrator_failure(&store, project_id, 2).await;
record_orchestrator_failure(&store, project_id, 2).await;
let docs = store.list_memory_docs(project_id).await.unwrap();
let count = load_failure_count(&docs);
assert_eq!(count, 3);
// Reset
reset_orchestrator_failures(&store, project_id).await;
let docs = store.list_memory_docs(project_id).await.unwrap();
let count = load_failure_count(&docs);
assert_eq!(count, 0);
}
#[tokio::test]
async fn failure_count_resets_on_new_version() {
let project_id = ProjectId::new();
let store: Arc<dyn Store> = Arc::new(crate::tests::InMemoryStore::with_docs(vec![]));
// Record failures for version 1
record_orchestrator_failure(&store, project_id, 1).await;
record_orchestrator_failure(&store, project_id, 1).await;
// Switch to version 2 — count should reset to 1
record_orchestrator_failure(&store, project_id, 2).await;
let docs = store.list_memory_docs(project_id).await.unwrap();
let count = load_failure_count(&docs);
assert_eq!(count, 1);
}
#[test]
fn parse_outcome_completed() {
let result = serde_json::json!({"outcome": "completed", "response": "Hello!"});
let outcome = parse_outcome(&result);
assert!(matches!(outcome, ThreadOutcome::Completed { response: Some(r) } if r == "Hello!"));
}
#[test]
fn parse_outcome_failed() {
let result = serde_json::json!({"outcome": "failed", "error": "boom"});
let outcome = parse_outcome(&result);
assert!(matches!(outcome, ThreadOutcome::Failed { error } if error == "boom"));
}
#[test]
fn parse_outcome_need_approval() {
let result = serde_json::json!({
"outcome": "need_approval",
"action_name": "shell",
"call_id": "abc",
"parameters": {"cmd": "rm -rf /"}
});
let outcome = parse_outcome(&result);
assert!(
matches!(outcome, ThreadOutcome::NeedApproval { action_name, .. } if action_name == "shell")
);
}
#[test]
fn parse_outcome_max_iterations() {
let result = serde_json::json!({"outcome": "max_iterations"});
let outcome = parse_outcome(&result);
assert!(matches!(outcome, ThreadOutcome::MaxIterations));
}
#[test]
fn parse_outcome_stopped() {
let result = serde_json::json!({"outcome": "stopped"});
let outcome = parse_outcome(&result);
assert!(matches!(outcome, ThreadOutcome::Stopped));
}
}