use std::time::Duration;
#[cfg(feature = "python")]
use pyo3::prelude::*;
#[cfg(feature = "python")]
use tokio::sync::{mpsc, oneshot};
#[cfg(feature = "python")]
use crate::error::Error;
#[cfg(feature = "native")]
pub mod native;
#[cfg(feature = "native")]
pub use native::NativeRuntime;
pub(crate) mod bridge_state;
mod config;
#[cfg(feature = "python")]
pub(crate) mod command_loop;
#[cfg(feature = "python")]
pub(crate) mod ffi_dispatch;
#[cfg(feature = "python")]
mod handlers;
#[cfg(feature = "python")]
pub(crate) mod py_scripts;
#[cfg(feature = "python")]
pub(crate) mod streaming;
#[cfg(feature = "python")]
pub(crate) mod venv;
#[cfg(test)]
#[cfg(feature = "python")]
mod tests;
#[cfg(feature = "python")]
pub(crate) use bridge_state::AgentId;
#[cfg(test)]
pub(crate) use bridge_state::set_agent_conversation_id;
pub(crate) use bridge_state::{
AgentBridgeState, bridge_state, initializing_hook_runners, next_agent_id,
};
pub use config::{BackendLogLevel, RuntimeConfig};
#[cfg(feature = "python")]
pub(crate) use ffi_dispatch::{
dispatch_rust_hook, dispatch_rust_policy_confirm, dispatch_rust_tool,
};
pub const DEFAULT_INTER_AGENT_DELAY: Duration = Duration::from_millis(500);
const DEFAULT_CHANNEL_CAPACITY: usize = 64;
const DEFAULT_SHUTDOWN_TIMEOUT: Duration = Duration::from_secs(10);
#[cfg(feature = "python")]
pub(crate) enum PyCommand {
CreateAgent {
agent_id: u64,
config_json: String,
reply: oneshot::Sender<Result<(AgentId, Vec<handlers::agent::RawToolInfo>), Error>>,
},
Chat {
agent_id: AgentId,
prompt: String,
reply: oneshot::Sender<Result<crate::streaming::ChatResponseHandle, Error>>,
},
ShutdownAgent {
agent_id: AgentId,
reply: oneshot::Sender<Result<(), Error>>,
},
Cancel {
agent_id: AgentId,
reply: oneshot::Sender<Result<(), Error>>,
},
WaitForIdle {
agent_id: AgentId,
reply: oneshot::Sender<Result<(), Error>>,
},
Send {
agent_id: AgentId,
prompt: String,
reply: oneshot::Sender<Result<(), Error>>,
},
SignalIdle {
agent_id: AgentId,
reply: oneshot::Sender<Result<(), Error>>,
},
WaitForWakeup {
agent_id: AgentId,
timeout_secs: f64,
reply: oneshot::Sender<Result<bool, Error>>,
},
Shutdown,
GetHistory {
agent_id: AgentId,
reply: oneshot::Sender<Result<Vec<crate::types::ConversationMessage>, Error>>,
},
GetTurnCount {
agent_id: AgentId,
reply: oneshot::Sender<Result<u32, Error>>,
},
GetActiveAgentCount {
reply: oneshot::Sender<Result<usize, Error>>,
},
GetTotalUsage {
agent_id: AgentId,
reply: oneshot::Sender<Result<crate::types::UsageMetadata, Error>>,
},
GetLastTurnUsage {
agent_id: AgentId,
reply: oneshot::Sender<Result<crate::types::UsageMetadata, Error>>,
},
ClearHistory {
agent_id: AgentId,
reply: oneshot::Sender<Result<(), Error>>,
},
GetCompactionIndices {
agent_id: AgentId,
reply: oneshot::Sender<Result<Vec<u32>, Error>>,
},
GetLastResponse {
agent_id: AgentId,
reply: oneshot::Sender<Result<Option<String>, Error>>,
},
Delete {
agent_id: AgentId,
reply: oneshot::Sender<Result<(), Error>>,
},
Disconnect {
agent_id: AgentId,
reply: oneshot::Sender<Result<(), Error>>,
},
IsIdle {
agent_id: AgentId,
reply: oneshot::Sender<Result<bool, Error>>,
},
}
#[cfg(feature = "python")]
pub struct PythonRuntime {
cmd_tx: Option<mpsc::Sender<PyCommand>>,
thread: Option<std::thread::JoinHandle<()>>,
config: RuntimeConfig,
}
#[cfg(feature = "python")]
impl std::fmt::Debug for PythonRuntime {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("PythonRuntime")
.field("config", &self.config)
.field(
"thread_running",
&self.thread.as_ref().is_some_and(|t| !t.is_finished()),
)
.finish_non_exhaustive()
}
}
#[cfg(feature = "python")]
impl PythonRuntime {
pub fn new(config: RuntimeConfig) -> Result<Self, Error> {
let (cmd_tx, cmd_rx) = mpsc::channel(config.channel_capacity);
let thread_config = config.clone();
let thread = std::thread::Builder::new()
.name("agy-bridge-python-runtime".into())
.spawn(move || {
python_thread_main(cmd_rx, &thread_config);
})
.map_err(|e| Error::BackendError {
message: format!("Failed to spawn Python runtime thread: {e}"),
})?;
Ok(Self {
cmd_tx: Some(cmd_tx),
thread: Some(thread),
config,
})
}
async fn send_command<T>(
&self,
operation: &str,
build_cmd: impl FnOnce(oneshot::Sender<Result<T, Error>>) -> PyCommand,
) -> Result<T, Error> {
let Some(ref tx) = self.cmd_tx else {
return Err(Error::ChannelClosed {
message: format!("Python runtime thread is shut down (sending {operation})"),
});
};
let (reply_tx, reply_rx) = oneshot::channel();
let cmd = build_cmd(reply_tx);
tx.send(cmd).await.map_err(|e| Error::ChannelClosed {
message: format!("Python runtime thread has exited (sending {operation}): {e}"),
})?;
let result = reply_rx.await.map_err(|e| Error::ChannelClosed {
message: format!("Reply channel dropped for {operation}: {e}"),
})??;
Ok(result)
}
pub(crate) async fn active_agent_count(&self) -> Result<usize, Error> {
self.send_command("active_agent_count", |reply| {
PyCommand::GetActiveAgentCount { reply }
})
.await
}
pub async fn shutdown(mut self) -> Result<(), Error> {
if let Some(tx) = self.cmd_tx.take()
&& let Err(e) = tx.send(PyCommand::Shutdown).await
{
tracing::warn!("Shutdown command send failed (thread may already be exiting): {e}");
}
let Some(thread) = self.thread.take() else {
tracing::warn!("PythonRuntime::shutdown() called but thread handle already taken");
return Ok(());
};
let shutdown_timeout = self.config.shutdown_timeout;
let join_result = tokio::time::timeout(
shutdown_timeout,
tokio::task::spawn_blocking(move || thread.join()),
)
.await;
match join_result {
Ok(Ok(Ok(()))) => {
tracing::info!("Python runtime thread joined successfully");
Ok(())
}
Ok(Ok(Err(panic_payload))) => {
let panic_msg = panic_payload.downcast_ref::<&str>().map_or_else(
|| {
panic_payload
.downcast_ref::<String>()
.map_or_else(|| format!("{panic_payload:?}"), Clone::clone)
},
|s| (*s).to_string(),
);
tracing::error!(
panic_message = %panic_msg,
"Python runtime thread panicked during shutdown"
);
Err(Error::BackendError {
message: format!("Python runtime thread panicked during shutdown: {panic_msg}"),
})
}
Ok(Err(join_err)) => {
tracing::error!("spawn_blocking join error: {join_err}");
Err(Error::BackendError {
message: format!("Failed to join Python thread: {join_err}"),
})
}
Err(_elapsed) => {
tracing::error!(
timeout_secs = shutdown_timeout.as_secs(),
"Python runtime thread did not exit within shutdown timeout"
);
Err(Error::Timeout {
duration: shutdown_timeout,
operation: "PythonRuntime::shutdown (thread join)".to_string(),
})
}
}
}
}
#[cfg(feature = "python")]
impl Drop for PythonRuntime {
fn drop(&mut self) {
let Some(tx) = self.cmd_tx.take() else {
return;
};
let Some(thread) = self.thread.take() else {
return;
};
if !thread.is_finished()
&& let Err(e) = tx.try_send(PyCommand::Shutdown)
{
tracing::debug!(
error = %e,
"PythonRuntime::drop: could not eagerly signal shutdown; \
relying on channel close"
);
}
drop(tx);
let deadline = std::time::Instant::now() + self.config.shutdown_timeout;
while !thread.is_finished() && std::time::Instant::now() < deadline {
std::thread::sleep(std::time::Duration::from_millis(5));
}
if thread.is_finished() {
if thread.join().is_err() {
tracing::error!("Python runtime thread panicked during drop cleanup");
} else {
tracing::debug!("Python runtime thread joined cleanly on drop");
}
} else {
tracing::warn!(
"Python runtime thread still running after shutdown timeout during drop — \
detaching; agent cleanup will complete asynchronously"
);
}
}
}
#[cfg(feature = "python")]
fn python_thread_main(cmd_rx: mpsc::Receiver<PyCommand>, config: &RuntimeConfig) {
Python::initialize();
Python::attach(|py| {
if let Err(e) = venv::configure_python_sys_path(py) {
tracing::error!(
error = %e,
"Failed to configure Python sys.path in runtime thread — \
venv imports will likely fail"
);
}
});
if let Err(e) = run_live_thread(cmd_rx, config) {
tracing::error!(error = %e, "Python runtime thread failed");
}
tracing::info!("Python runtime thread exiting");
}
#[cfg(feature = "python")]
fn run_live_thread(cmd_rx: mpsc::Receiver<PyCommand>, config: &RuntimeConfig) -> Result<(), Error> {
Python::attach(|py| {
let asyncio = py.import("asyncio").map_err(|e| Error::BackendError {
message: format!("Failed to import asyncio: {e}"),
})?;
let event_loop =
asyncio
.call_method0("new_event_loop")
.map_err(|e| Error::BackendError {
message: format!("Failed to create new asyncio event loop: {e}"),
})?;
asyncio
.call_method1("set_event_loop", (&event_loop,))
.map_err(|e| Error::BackendError {
message: format!("Failed to set asyncio event loop: {e}"),
})?;
let sys = py.import("sys").map_err(|e| Error::BackendError {
message: format!("Failed to import sys: {e}"),
})?;
let sys_modules = sys.getattr("modules").map_err(|e| Error::BackendError {
message: format!("Failed to get sys.modules: {e}"),
})?;
let globals_mod = if sys_modules
.contains(command_loop::AGY_BRIDGE_GLOBALS_MODULE)
.map_err(|e| Error::BackendError {
message: format!("Failed to check sys.modules: {e}"),
})? {
sys_modules
.get_item(command_loop::AGY_BRIDGE_GLOBALS_MODULE)
.map_err(|e| Error::BackendError {
message: format!("Failed to get _agy_bridge_globals: {e}"),
})?
} else {
let types = py.import("types").map_err(|e| Error::BackendError {
message: format!("Failed to import types: {e}"),
})?;
let module = types
.getattr("ModuleType")
.map_err(|e| Error::BackendError {
message: format!("Failed to get ModuleType: {e}"),
})?
.call1((command_loop::AGY_BRIDGE_GLOBALS_MODULE,))
.map_err(|e| Error::BackendError {
message: format!("Failed to create ModuleType: {e}"),
})?;
sys_modules
.set_item(command_loop::AGY_BRIDGE_GLOBALS_MODULE, &module)
.map_err(|e| Error::BackendError {
message: format!("Failed to register _agy_bridge_globals: {e}"),
})?;
module
};
globals_mod
.setattr("EVENT_LOOP", &event_loop)
.map_err(|e| Error::BackendError {
message: format!("Failed to set EVENT_LOOP in globals: {e}"),
})?;
register_thread_event_loop(py, &globals_mod, &event_loop)?;
tracing::info!("Python asyncio event loop created on runtime thread");
let inter_agent_delay = config.inter_agent_delay;
let stream_limits = streaming::StreamLimits::from_config(config);
let event_loop_obj = event_loop.clone().unbind();
let run_fut =
pyo3_async_runtimes::tokio::run_until_complete(event_loop.clone(), async move {
command_loop::run_async_command_loop(
event_loop_obj,
cmd_rx,
inter_agent_delay,
stream_limits,
)
.await
});
unregister_thread_event_loop(py, &globals_mod);
if let Err(e) = run_fut {
if let Err(close_err) = event_loop.call_method0("close") {
tracing::warn!("Failed to close asyncio event loop: {close_err}");
}
return Err(Error::BackendError {
message: format!("Python runtime command loop failed: {e}"),
});
}
if let Err(e) = event_loop.call_method0("close") {
tracing::warn!("Failed to close asyncio event loop: {e}");
}
Ok(())
})
}
#[cfg(feature = "python")]
fn register_thread_event_loop(
py: Python<'_>,
globals_mod: &Bound<'_, PyAny>,
event_loop: &Bound<'_, PyAny>,
) -> Result<(), Error> {
let threading = py.import("threading").map_err(|e| Error::BackendError {
message: format!("Failed to import threading for event-loop registration: {e}"),
})?;
let thread_id = threading
.call_method0("get_ident")
.map_err(|e| Error::BackendError {
message: format!("Failed to read threading.get_ident(): {e}"),
})?;
let loops = if globals_mod
.hasattr("EVENT_LOOPS")
.map_err(|e| Error::BackendError {
message: format!("Failed to check for EVENT_LOOPS attribute: {e}"),
})? {
globals_mod
.getattr("EVENT_LOOPS")
.map_err(|e| Error::BackendError {
message: format!("Failed to get EVENT_LOOPS map: {e}"),
})?
} else {
let dict = pyo3::types::PyDict::new(py).into_any();
globals_mod
.setattr("EVENT_LOOPS", &dict)
.map_err(|e| Error::BackendError {
message: format!("Failed to create EVENT_LOOPS map: {e}"),
})?;
dict
};
loops
.set_item(thread_id, event_loop)
.map_err(|e| Error::BackendError {
message: format!("Failed to register runtime event loop by thread id: {e}"),
})?;
Ok(())
}
#[cfg(feature = "python")]
fn unregister_thread_event_loop(py: Python<'_>, globals_mod: &Bound<'_, PyAny>) {
let unregister_res = (|| -> PyResult<()> {
let threading = py.import("threading")?;
let thread_id = threading.call_method0("get_ident")?;
if globals_mod.hasattr("EVENT_LOOPS")? {
let loops = globals_mod.getattr("EVENT_LOOPS")?;
let dict = loops.cast::<pyo3::types::PyDict>()?;
dict.del_item(thread_id)?;
}
Ok(())
})();
if let Err(e) = unregister_res {
tracing::debug!(error = %e, "Failed to unregister thread event loop on teardown");
}
}
#[cfg(feature = "python")]
fn compute_active_builtins(
config: &crate::config::AgentConfig,
) -> Vec<crate::config::BuiltinTools> {
let Some(caps) = config.capabilities.as_ref() else {
return crate::config::BuiltinTools::all_tools().to_vec();
};
if let Some(enabled) = caps.enabled_tools.as_ref() {
return enabled.clone();
}
if let Some(disabled) = caps.disabled_tools.as_ref() {
return crate::config::BuiltinTools::all_tools()
.iter()
.filter(|t| !disabled.contains(t))
.copied()
.collect();
}
crate::config::BuiltinTools::all_tools().to_vec()
}
#[cfg(feature = "python")]
impl crate::agent::Runtime for PythonRuntime {
async fn create_agent(
&self,
agent_id: u64,
config: crate::config::AgentConfig,
) -> Result<(crate::agent::AgentId, Vec<crate::tools::AvailableTool>), Error> {
if let Some(save_dir) = config.save_dir.as_ref()
&& let Err(e) = std::fs::create_dir_all(save_dir)
{
tracing::warn!(
save_dir = %save_dir.display(),
error = ?e,
"Failed to create save_dir; conversation state may not persist \
and resume may fail with \"conversation not found\""
);
}
let config_json = {
let mut val = serde_json::to_value(&config).map_err(|e| Error::BackendError {
message: format!("Failed to serialize AgentConfig: {e}"),
})?;
if let serde_json::Value::Object(ref mut map) = val {
map.insert(
"_backend_log_level".to_owned(),
serde_json::Value::String(self.config.backend_log_level.as_str().to_owned()),
);
}
serde_json::to_string(&val).map_err(|e| Error::BackendError {
message: format!("Failed to re-serialize config JSON: {e}"),
})?
};
let custom_tool_names: std::collections::HashSet<String> =
config.tools.iter().map(|t| t.name.clone()).collect();
let (raw_id, raw_tools) = self
.send_command("create_agent", |reply| PyCommand::CreateAgent {
agent_id,
config_json,
reply,
})
.await?;
let active_builtins = compute_active_builtins(&config);
let builtin_names: std::collections::HashSet<&str> = active_builtins
.iter()
.map(crate::config::BuiltinTools::as_sdk_name)
.collect();
let mut available_tools: Vec<crate::tools::AvailableTool> = raw_tools
.into_iter()
.filter(|raw| !builtin_names.contains(raw.name.as_str()))
.map(|raw| {
let source = if custom_tool_names.contains(&raw.name) {
crate::tools::ToolSource::Custom
} else {
crate::tools::ToolSource::Mcp
};
crate::tools::AvailableTool {
name: raw.name,
description: raw.description,
parameter_schema: raw.parameter_schema,
source,
}
})
.collect();
for builtin in active_builtins {
available_tools.push(crate::tools::AvailableTool {
name: builtin.as_sdk_name().to_owned(),
description: builtin.description().to_owned(),
parameter_schema: serde_json::Value::Null,
source: crate::tools::ToolSource::Builtin,
});
}
tracing::info!(
agent_id = raw_id.0,
tool_count = available_tools.len(),
tools = ?available_tools.iter().map(|t| format!("{t}")).collect::<Vec<_>>(),
"Agent created with available tools"
);
Ok((raw_id.0, available_tools))
}
async fn chat(
&self,
agent_id: crate::agent::AgentId,
content: &crate::content::Content,
) -> Result<crate::streaming::ChatResponseHandle, Error> {
let prompt = match content {
crate::content::Content::Text { text } => text.clone(),
other => crate::content::content_to_json(other)?,
};
self.send_command("chat", |reply| PyCommand::Chat {
agent_id: AgentId(agent_id),
prompt,
reply,
})
.await
}
async fn shutdown_agent(&self, agent_id: crate::agent::AgentId) -> Result<(), Error> {
self.send_command("shutdown_agent", |reply| PyCommand::ShutdownAgent {
agent_id: AgentId(agent_id),
reply,
})
.await
}
fn try_shutdown_agent(&self, agent_id: crate::agent::AgentId) {
if let Some(ref tx) = self.cmd_tx {
let (reply, _) = oneshot::channel();
if let Err(e) = tx.try_send(PyCommand::ShutdownAgent {
agent_id: AgentId(agent_id),
reply,
}) {
tracing::debug!(
agent_id = agent_id,
error = %e,
"try_shutdown_agent: channel send failed (runtime may already be gone)"
);
}
}
}
async fn cancel(&self, agent_id: crate::agent::AgentId) -> Result<(), Error> {
self.send_command("cancel", |reply| PyCommand::Cancel {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn wait_for_idle(&self, agent_id: crate::agent::AgentId) -> Result<(), Error> {
self.send_command("wait_for_idle", |reply| PyCommand::WaitForIdle {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn send(
&self,
agent_id: crate::agent::AgentId,
content: &crate::content::Content,
) -> Result<(), Error> {
let prompt = match content {
crate::content::Content::Text { text } => text.clone(),
other => crate::content::content_to_json(other)?,
};
self.send_command("send", |reply| PyCommand::Send {
agent_id: AgentId(agent_id),
prompt,
reply,
})
.await
}
async fn signal_idle(&self, agent_id: crate::agent::AgentId) -> Result<(), Error> {
self.send_command("signal_idle", |reply| PyCommand::SignalIdle {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn wait_for_wakeup(
&self,
agent_id: crate::agent::AgentId,
timeout: std::time::Duration,
) -> Result<bool, Error> {
self.send_command("wait_for_wakeup", |reply| PyCommand::WaitForWakeup {
agent_id: AgentId(agent_id),
timeout_secs: timeout.as_secs_f64(),
reply,
})
.await
}
async fn history(
&self,
agent_id: crate::agent::AgentId,
) -> Result<Vec<crate::types::ConversationMessage>, Error> {
self.send_command("get_history", |reply| PyCommand::GetHistory {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn turn_count(&self, agent_id: crate::agent::AgentId) -> Result<u32, Error> {
self.send_command("get_turn_count", |reply| PyCommand::GetTurnCount {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn total_usage(
&self,
agent_id: crate::agent::AgentId,
) -> Result<crate::types::UsageMetadata, Error> {
self.send_command("get_total_usage", |reply| PyCommand::GetTotalUsage {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn last_turn_usage(
&self,
agent_id: crate::agent::AgentId,
) -> Result<crate::types::UsageMetadata, Error> {
self.send_command("get_last_turn_usage", |reply| PyCommand::GetLastTurnUsage {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn clear_history(&self, agent_id: crate::agent::AgentId) -> Result<(), Error> {
self.send_command("clear_history", |reply| PyCommand::ClearHistory {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn compaction_indices(&self, agent_id: crate::agent::AgentId) -> Result<Vec<u32>, Error> {
self.send_command("compaction_indices", |reply| {
PyCommand::GetCompactionIndices {
agent_id: AgentId(agent_id),
reply,
}
})
.await
}
async fn last_response(
&self,
agent_id: crate::agent::AgentId,
) -> Result<Option<String>, Error> {
self.send_command("last_response", |reply| PyCommand::GetLastResponse {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn delete(&self, agent_id: crate::agent::AgentId) -> Result<(), Error> {
self.send_command("delete", |reply| PyCommand::Delete {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn disconnect(&self, agent_id: crate::agent::AgentId) -> Result<(), Error> {
self.send_command("disconnect", |reply| PyCommand::Disconnect {
agent_id: AgentId(agent_id),
reply,
})
.await
}
async fn is_idle(&self, agent_id: crate::agent::AgentId) -> Result<bool, Error> {
self.send_command("is_idle", |reply| PyCommand::IsIdle {
agent_id: AgentId(agent_id),
reply,
})
.await
}
}