use async_trait::async_trait;
use lash_core::{
ToolArgumentProjectionPolicy, ToolAvailabilityConfig, ToolCall, ToolContract, ToolControl,
ToolDefinition, ToolManifest, ToolProvider, ToolResult, ToolScheduling,
};
use serde_json::{Value, json};
use std::sync::Arc;
use crate::projection::RlmSeed;
pub(crate) struct RlmControlToolsProvider;
#[async_trait]
impl ToolProvider for RlmControlToolsProvider {
fn tool_manifests(&self) -> Vec<ToolManifest> {
vec![continue_as_tool_definition().manifest()]
}
fn resolve_contract(&self, name: &str) -> Option<Arc<ToolContract>> {
(name == "continue_as").then(|| Arc::new(continue_as_tool_definition().contract()))
}
async fn execute(&self, call: ToolCall<'_>) -> ToolResult {
let result = match call.name {
"continue_as" => continue_as_switch_frame(call.args),
_ => return ToolResult::err_fmt(format_args!("Unknown tool: {}", call.name)),
};
finalise_tool_result(result)
}
}
pub fn continue_as_tool_definition() -> ToolDefinition {
ToolDefinition::raw(
"tool:continue_as",
"continue_as",
"Tail-call into a fresh RLM AgentFrame inside the current session with a clean window.\n\nThe new frame inherits **nothing** implicitly — no variables or message history. Pass everything it needs via `seed: { name: value, ... }`. Seed values copied from read-only values stay read-only in the new frame; computed expressions become writable variables.\n\n- Use when the current trajectory is stale, dominated by failed attempts, or the context budget is tight.\n- Treat `control.continue_as(...)` as a terminal control action: make it the last meaningful statement in the lashlang block, and do not call `submit` or perform more work after it.\n- `task` packs the concrete goal, constraints, and next steps the new frame must act on.\n- `seed` packs the concrete state (paths, facts already learned, partial results, read-only values) the new frame needs in scope; leave bulky raw output behind.\n- `continue_as` only changes the active AgentFrame. It does not start, transfer, list, cancel, or otherwise manage processes.",
continue_as_input_schema(),
continue_as_output_schema(),
)
.with_examples(vec![
r#"await control.continue_as({ task: "continue the audit from the summarized findings", seed: { problem: input.prompt, findings: findings } })?"#.into(),
])
.with_agent_surface(lash_core::ToolAgentSurface::new(["control"], "continue_as"))
.with_argument_projection(ToolArgumentProjectionPolicy::preserve_projected_refs_in_field(
"seed",
))
.with_availability(ToolAvailabilityConfig::callable())
.with_scheduling(ToolScheduling::Parallel)
}
fn continue_as_output_schema() -> Value {
json!({
"type": "object",
"properties": {
"ok": { "type": "boolean" },
"frame_id": { "type": "string" },
"task": { "type": "string" },
"seed_keys": {
"type": "array",
"items": { "type": "string" }
},
"seed_count": { "type": "integer", "minimum": 0 }
},
"required": [
"ok",
"frame_id",
"task",
"seed_keys",
"seed_count"
],
"additionalProperties": false
})
}
pub fn continue_as_input_schema() -> Value {
json!({
"type": "object",
"properties": {
"task": {
"type": "string",
"description": "Task for the new AgentFrame."
},
"seed": {
"type": "object",
"additionalProperties": true,
"description": "Optional record/dict of concrete state for the new AgentFrame."
}
},
"required": ["task"],
"additionalProperties": false
})
}
fn continue_as_switch_frame(args: &Value) -> Result<ContinueAsResult, String> {
let task = required_string(args, "task")?;
let seed = RlmSeed::from_tool_args(args).map_err(|err| format!("continue_as {err}"))?;
let mut seed_keys = seed
.globals
.keys()
.cloned()
.chain(seed.projected.entries.iter().map(|(name, _)| name.clone()))
.collect::<Vec<_>>();
seed_keys.sort();
let seed_count = seed_keys.len();
let frame_id = uuid::Uuid::new_v4().to_string();
let initial_nodes = crate::rlm_seed_initial_nodes(seed);
let initial_nodes = initial_nodes
.into_iter()
.map(|node| {
serde_json::to_value(node)
.map_err(|err| format!("failed to encode continue_as frame seed node: {err}"))
})
.collect::<Result<Vec<_>, _>>()?;
Ok(ContinueAsResult {
value: json!({
"ok": true,
"frame_id": frame_id.clone(),
"task": task.clone(),
"seed_keys": seed_keys,
"seed_count": seed_count,
}),
control: ToolControl::SwitchAgentFrame {
frame_id,
initial_nodes,
task: Some(task),
},
})
}
fn required_string(args: &Value, key: &str) -> Result<String, String> {
args.get(key)
.and_then(Value::as_str)
.map(str::trim)
.filter(|value| !value.is_empty())
.map(ToOwned::to_owned)
.ok_or_else(|| format!("missing required parameter: {key}"))
}
struct ContinueAsResult {
value: Value,
control: ToolControl,
}
fn finalise_tool_result(result: Result<ContinueAsResult, String>) -> ToolResult {
match result {
Ok(result) => ToolResult::ok(result.value).with_control(result.control),
Err(err) => ToolResult::err(json!(err)),
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::projection::{decode_rlm_protocol_event, rlm_protocol_event};
use std::sync::{Arc, Mutex};
use lash_core::plugin::runtime_host::{
SessionGraphService, SessionLifecycleService, SessionStateService,
};
use lash_core::plugin::{PluginError, SessionHandle};
use lash_core::runtime::RuntimeSessionState;
use lash_core::{
SessionAppendNode, SessionCreateRequest, SessionPolicy, SessionSnapshot, ToolProvider,
};
use lash_rlm_types::{RlmProtocolEvent, RlmTermination};
fn model_spec(model: &str) -> lash_core::ModelSpec {
lash_core::ModelSpec::from_token_limits(model, None, 200_000, None)
.expect("valid test model spec")
}
#[test]
fn continue_as_contract_documents_switch_result() {
let definition = continue_as_tool_definition();
assert_eq!(
definition.contract.output_schema["required"],
json!(["ok", "frame_id", "task", "seed_keys", "seed_count"])
);
let rendered = definition.compact_contract().render_signature();
assert!(rendered.contains("frame_id"), "{rendered}");
assert!(!rendered.contains("handle_count"), "{rendered}");
assert!(!rendered.contains("projected_count"), "{rendered}");
assert!(!rendered.contains("global_count"), "{rendered}");
}
#[derive(Default)]
struct BatonManager {
snapshot: RuntimeSessionState,
created: Mutex<Vec<SessionCreateRequest>>,
closed: Mutex<Vec<String>>,
}
#[test]
fn continue_as_tool_definition_preserves_projected_seed_refs_by_metadata() {
assert_eq!(
continue_as_tool_definition().manifest.argument_projection,
ToolArgumentProjectionPolicy::preserve_projected_refs_in_field("seed")
);
}
#[async_trait]
impl SessionStateService for BatonManager {
async fn snapshot_current(&self) -> Result<SessionSnapshot, PluginError> {
Ok(self.snapshot.to_snapshot())
}
async fn snapshot_session(
&self,
_session_id: &str,
) -> Result<SessionSnapshot, PluginError> {
Ok(self.snapshot.to_snapshot())
}
async fn tool_catalog(
&self,
_session_id: &str,
) -> Result<Vec<serde_json::Value>, PluginError> {
Ok(Vec::new())
}
}
#[async_trait]
impl SessionLifecycleService for BatonManager {
async fn create_session(
&self,
request: SessionCreateRequest,
) -> Result<SessionHandle, PluginError> {
self.created.lock().expect("created").push(request.clone());
Ok(SessionHandle {
session_id: request.session_id.unwrap_or_else(|| "child".to_string()),
parent_session_id: request.relation.parent_session_id().map(ToOwned::to_owned),
policy: request.policy.unwrap_or_default(),
})
}
async fn close_session(&self, session_id: &str) -> Result<(), PluginError> {
self.closed
.lock()
.expect("closed")
.push(session_id.to_string());
Ok(())
}
}
#[async_trait]
impl SessionGraphService for BatonManager {}
#[async_trait]
impl lash_core::ProcessService for BatonManager {
async fn start(
&self,
_session_id: &str,
_registration: lash_core::ProcessRegistration,
_options: lash_core::ProcessStartOptions,
_scope: lash_core::ProcessOpScope<'_>,
) -> Result<lash_core::ProcessRecord, PluginError> {
Err(PluginError::Session(
"process starts are unavailable in this test".to_string(),
))
}
async fn await_process(
&self,
_process_id: &str,
_scope: lash_core::ProcessOpScope<'_>,
) -> Result<lash_core::ProcessAwaitOutput, PluginError> {
Err(PluginError::Session(
"process awaiting is unavailable in this test".to_string(),
))
}
async fn list_visible(
&self,
_session_id: &str,
_mode: lash_core::ProcessListMode,
_scope: lash_core::ProcessOpScope<'_>,
) -> Result<Vec<lash_core::runtime::ProcessHandleGrantEntry>, PluginError> {
Ok(Vec::new())
}
async fn validate_visible(
&self,
_session_id: &str,
_handle_ids: &[String],
_scope: lash_core::ProcessOpScope<'_>,
) -> Result<(), PluginError> {
Err(PluginError::Session(
"continue_as must not validate process handles".to_string(),
))
}
async fn cancel(
&self,
_session_id: &str,
_process_id: &str,
_scope: lash_core::ProcessOpScope<'_>,
) -> Result<lash_core::ProcessRecord, PluginError> {
Err(PluginError::Session(
"process cancellation is unavailable in this test".to_string(),
))
}
async fn signal(
&self,
_session_id: &str,
_process_id: &str,
_signal_id: String,
_payload: serde_json::Value,
_scope: lash_core::ProcessOpScope<'_>,
) -> Result<lash_core::ProcessEvent, PluginError> {
Err(PluginError::Session(
"process signalling is unavailable in this test".to_string(),
))
}
async fn transfer(
&self,
_from_session_id: &str,
_to_session_id: &str,
_process_ids: Vec<String>,
_scope: lash_core::ProcessOpScope<'_>,
) -> Result<(), PluginError> {
Err(PluginError::Session(
"continue_as must not transfer process handles".to_string(),
))
}
async fn cancel_unreferenced(
&self,
_session_id: &str,
_keep_process_ids: Vec<String>,
_scope: lash_core::ProcessOpScope<'_>,
) -> Result<Vec<lash_core::ProcessRecord>, PluginError> {
Err(PluginError::Session(
"continue_as must not cancel process handles".to_string(),
))
}
}
async fn run_continue_as(
provider: &RlmControlToolsProvider,
manager: Arc<BatonManager>,
args: &Value,
) -> ToolResult {
let sessions: Arc<dyn SessionStateService> = manager.clone();
let session_lifecycle: Arc<dyn SessionLifecycleService> = manager.clone();
let session_graph: Arc<dyn SessionGraphService> = manager.clone();
let processes: Arc<dyn lash_core::ProcessService> = manager;
let context = lash_core::ToolContext::__for_testing(
"test-session".to_string(),
sessions,
session_lifecycle,
session_graph,
processes,
Arc::new(lash_core::InMemoryAttachmentStore::new()),
lash_core::DirectCompletionClient::from_fn(|_, _| {
Err(lash_core::PluginError::Session(
"direct completions are unavailable in continue_as tests".to_string(),
))
}),
Some("continue-as-test".to_string()),
);
provider
.execute(lash_core::ToolCall {
name: "continue_as",
args,
context: &context,
progress: None,
})
.await
}
#[test]
fn rlm_control_definitions_include_continue_as_only() {
let provider = RlmControlToolsProvider;
let names = provider
.tool_manifests()
.into_iter()
.map(|tool| tool.name)
.collect::<Vec<_>>();
assert_eq!(names, vec!["continue_as"]);
}
#[tokio::test]
async fn continue_as_creates_empty_rlm_frame_with_seed_and_task() {
let mut session_graph = lash_core::SessionGraph::default();
session_graph.append_protocol_event(rlm_protocol_event(RlmProtocolEvent::RlmGlobalsPatch(
lash_rlm_types::RlmGlobalsPatchPluginBody {
set_default: serde_json::Map::from_iter([("diary".to_string(), json!([]))]),
},
)));
let manager = Arc::new(BatonManager {
snapshot: RuntimeSessionState {
policy: SessionPolicy {
model: model_spec("model"),
..SessionPolicy::default()
},
protocol_turn_options: lash_core::ProtocolTurnOptions::typed(
RlmTermination::SubmitRequired {
schema: Some(json!({
"type": "object",
"properties": { "answer": { "type": "string" } },
"required": ["answer"]
})),
},
)
.expect("valid rlm turn options"),
session_graph,
..RuntimeSessionState::default()
},
created: Mutex::new(Vec::new()),
..BatonManager::default()
});
let provider = RlmControlToolsProvider;
let args = json!({
"task": "finish from here",
"seed": { "x": 1, "query": "original" }
});
let result = run_continue_as(&provider, manager.clone(), &args).await;
assert!(result.is_success(), "{:?}", result.value_for_projection());
let value = result.value_for_projection();
assert!(value.get("frame_id").and_then(Value::as_str).is_some());
assert_eq!(value.get("seed_keys"), Some(&json!(["query", "x"])));
assert_eq!(value.get("seed_count"), Some(&json!(2)));
assert!(value.get("projected_count").is_none());
assert!(value.get("global_count").is_none());
let Some(ToolControl::SwitchAgentFrame {
frame_id,
initial_nodes,
task,
}) = result.as_output().control.as_ref()
else {
panic!("expected frame switch control");
};
assert_eq!(
value.get("frame_id").and_then(Value::as_str),
Some(frame_id.as_str())
);
assert_eq!(task.as_deref(), Some("finish from here"));
assert_eq!(initial_nodes.len(), 1);
let node = serde_json::from_value::<SessionAppendNode>(initial_nodes[0].clone())
.expect("decode initial node");
let SessionAppendNode::ProtocolEvent {
event: protocol_event,
..
} = node
else {
panic!("expected seed globals event");
};
let Some(RlmProtocolEvent::RlmSeed(seed)) = decode_rlm_protocol_event(&protocol_event)
else {
panic!("expected RlmSeed");
};
assert_eq!(seed.globals["x"], json!(1));
assert_eq!(seed.globals["query"], json!("original"));
assert!(seed.projected.is_empty());
assert!(manager.created.lock().expect("created").is_empty());
}
#[tokio::test]
async fn continue_as_routes_projected_entries_and_globals_to_one_seed_event() {
// Mixed seed: `proj` was a projected source on the parent (encoded with
// the canonical `__projected__` JSON wrapper), `glob` was a regular
// global. The new frame receives both through one durable RLM seed event.
let manager = Arc::new(BatonManager {
snapshot: RuntimeSessionState {
policy: SessionPolicy {
model: model_spec("model"),
..SessionPolicy::default()
},
..RuntimeSessionState::default()
},
created: Mutex::new(Vec::new()),
..BatonManager::default()
});
let provider = RlmControlToolsProvider;
let args = json!({
"task": "finish from here",
"seed": {
"proj": { "__projected__": "carry-over" },
"glob": 7
}
});
let result = run_continue_as(&provider, manager.clone(), &args).await;
assert!(result.is_success(), "{:?}", result.value_for_projection());
let value = result.value_for_projection();
assert_eq!(value.get("seed_keys"), Some(&json!(["glob", "proj"])));
assert_eq!(value.get("seed_count"), Some(&json!(2)));
assert!(value.get("projected_count").is_none());
assert!(value.get("global_count").is_none());
let Some(ToolControl::SwitchAgentFrame { initial_nodes, .. }) =
result.as_output().control.as_ref()
else {
panic!("expected frame switch control");
};
assert_eq!(initial_nodes.len(), 1);
let node = serde_json::from_value::<SessionAppendNode>(initial_nodes[0].clone())
.expect("decode initial node");
let SessionAppendNode::ProtocolEvent {
event: protocol_event,
..
} = node
else {
panic!("expected seed globals event");
};
let Some(RlmProtocolEvent::RlmSeed(seed)) = decode_rlm_protocol_event(&protocol_event)
else {
panic!("expected RlmSeed");
};
assert_eq!(seed.globals.len(), 1, "only `glob` should land as a global");
assert_eq!(seed.globals["glob"], json!(7));
assert!(!seed.globals.contains_key("proj"));
assert_eq!(seed.projected.entries.len(), 1);
assert_eq!(seed.projected.entries[0].0, "proj");
assert_eq!(seed.projected.entries[0].1, json!("carry-over"));
assert!(manager.created.lock().expect("created").is_empty());
}
#[tokio::test]
async fn continue_as_preserves_process_shaped_seed_without_process_control() {
let manager = Arc::new(BatonManager {
snapshot: RuntimeSessionState {
policy: SessionPolicy {
model: model_spec("model"),
..SessionPolicy::default()
},
..RuntimeSessionState::default()
},
created: Mutex::new(Vec::new()),
..BatonManager::default()
});
let provider = RlmControlToolsProvider;
let args = json!({
"task": "continue with background work",
"seed": {
"one": { "__handle__": "process", "id": "h1", "tool": "slow" },
"nested": [{ "h": { "__handle__": "process", "id": "h2", "tool": "slow" } }]
}
});
let result = run_continue_as(&provider, manager.clone(), &args).await;
assert!(result.is_success(), "{:?}", result.value_for_projection());
let Some(ToolControl::SwitchAgentFrame { initial_nodes, .. }) =
result.as_output().control.as_ref()
else {
panic!("expected frame switch control");
};
let node = serde_json::from_value::<SessionAppendNode>(initial_nodes[0].clone())
.expect("decode initial node");
let SessionAppendNode::ProtocolEvent {
event: protocol_event,
..
} = node
else {
panic!("expected seed globals event");
};
let Some(RlmProtocolEvent::RlmSeed(seed)) = decode_rlm_protocol_event(&protocol_event)
else {
panic!("expected RlmSeed");
};
assert_eq!(
seed.globals["one"],
json!({ "__handle__": "process", "id": "h1", "tool": "slow" })
);
assert_eq!(
seed.globals["nested"],
json!([{ "h": { "__handle__": "process", "id": "h2", "tool": "slow" } }])
);
}
#[tokio::test]
async fn continue_as_does_not_validate_unknown_seed_handles() {
let manager = Arc::new(BatonManager {
snapshot: RuntimeSessionState {
policy: SessionPolicy {
model: model_spec("model"),
..SessionPolicy::default()
},
..RuntimeSessionState::default()
},
created: Mutex::new(Vec::new()),
..BatonManager::default()
});
let provider = RlmControlToolsProvider;
let args = json!({
"task": "continue",
"seed": { "h": { "__handle__": "process", "id": "missing", "tool": "slow" } }
});
let result = run_continue_as(&provider, manager.clone(), &args).await;
assert!(result.is_success(), "{:?}", result.value_for_projection());
assert!(manager.created.lock().expect("created").is_empty());
}
}