use std::pin::Pin;
use std::sync::Arc;
use std::time::Instant;
use serde::{Deserialize, Serialize};
use tokio::sync::{mpsc, oneshot};
use tokio_util::sync::CancellationToken;
use crate::{RuntimeError, RuntimeErrorCode};
use super::super::envelope::{RuntimeEffectEnvelope, RuntimeEffectOutcome};
use super::{RuntimeEffectControllerError, RuntimeEffectLocalExecutor};
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum ExecutionScope {
Turn {
session_id: String,
turn_id: String,
},
Process {
process_id: String,
},
QueueDrain {
session_id: String,
drain_id: String,
},
SessionDelete {
session_id: String,
},
RuntimeOperation {
operation_id: String,
},
}
impl ExecutionScope {
pub fn turn(session_id: impl Into<String>, turn_id: impl Into<String>) -> Self {
Self::Turn {
session_id: session_id.into(),
turn_id: turn_id.into(),
}
}
pub fn process(process_id: impl Into<String>) -> Self {
Self::Process {
process_id: process_id.into(),
}
}
pub fn queue_drain(session_id: impl Into<String>, drain_id: impl Into<String>) -> Self {
Self::QueueDrain {
session_id: session_id.into(),
drain_id: drain_id.into(),
}
}
pub fn session_delete(session_id: impl Into<String>) -> Self {
Self::SessionDelete {
session_id: session_id.into(),
}
}
pub fn runtime_operation(operation_id: impl Into<String>) -> Self {
Self::RuntimeOperation {
operation_id: operation_id.into(),
}
}
pub fn id(&self) -> &str {
match self {
Self::Turn { turn_id, .. } => turn_id,
Self::Process { process_id } => process_id,
Self::QueueDrain { drain_id, .. } => drain_id,
Self::SessionDelete { session_id } => session_id,
Self::RuntimeOperation { operation_id } => operation_id,
}
}
pub fn session_id(&self) -> Option<&str> {
match self {
Self::Turn { session_id, .. }
| Self::QueueDrain { session_id, .. }
| Self::SessionDelete { session_id } => Some(session_id),
Self::Process { .. } | Self::RuntimeOperation { .. } => None,
}
}
pub fn turn_id(&self) -> Option<&str> {
match self {
Self::Turn { turn_id, .. } => Some(turn_id),
_ => None,
}
}
pub fn validates_turn_trace_id(&self) -> bool {
matches!(self, Self::Turn { .. })
}
pub(crate) fn validate(&self) -> Result<(), RuntimeError> {
let missing = match self {
Self::Turn {
session_id,
turn_id,
} => session_id.trim().is_empty() || turn_id.trim().is_empty(),
Self::Process { process_id } => process_id.trim().is_empty(),
Self::QueueDrain {
session_id,
drain_id,
} => session_id.trim().is_empty() || drain_id.trim().is_empty(),
Self::SessionDelete { session_id } => session_id.trim().is_empty(),
Self::RuntimeOperation { operation_id } => operation_id.trim().is_empty(),
};
if missing {
return Err(RuntimeError::new(
RuntimeErrorCode::MissingExecutionScopeId,
"execution scopes require non-empty stable ids",
));
}
Ok(())
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum AwaitEventWaitIdentity {
ToolCompletion {
tool_call_id: String,
},
ProcessSignal {
process_id: String,
signal_name: String,
ordinal: u64,
},
TurnCancelGate,
TurnTerminal,
Custom {
key: String,
},
}
impl AwaitEventWaitIdentity {
pub fn tool_completion(tool_call_id: impl Into<String>) -> Self {
Self::ToolCompletion {
tool_call_id: tool_call_id.into(),
}
}
pub fn process_signal(
process_id: impl Into<String>,
signal_name: impl Into<String>,
ordinal: u64,
) -> Self {
Self::ProcessSignal {
process_id: process_id.into(),
signal_name: signal_name.into(),
ordinal,
}
}
pub(in crate::runtime::effect) fn validate(&self) -> Result<(), RuntimeError> {
let invalid = match self {
Self::ToolCompletion { tool_call_id } => tool_call_id.trim().is_empty(),
Self::ProcessSignal {
process_id,
signal_name,
ordinal,
} => process_id.trim().is_empty() || signal_name.trim().is_empty() || *ordinal == 0,
Self::TurnCancelGate | Self::TurnTerminal => false,
Self::Custom { key } => key.trim().is_empty(),
};
if invalid {
return Err(RuntimeError::new(
"invalid_await_event_wait_identity",
"await-event wait identity requires non-empty stable ids",
));
}
Ok(())
}
pub fn is_turn_control(&self) -> bool {
matches!(self, Self::TurnCancelGate | Self::TurnTerminal)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct AwaitEventKey {
pub scope: ExecutionScope,
pub wait: AwaitEventWaitIdentity,
pub key_id: String,
pub signature: String,
}
impl AwaitEventKey {
pub fn promise_key(&self) -> String {
format!("lash-await-event:{}", self.key_id)
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
pub struct ExternalCompletionError {
pub code: String,
pub message: String,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub raw: Option<serde_json::Value>,
}
impl ExternalCompletionError {
pub fn new(code: impl Into<String>, message: impl Into<String>) -> Self {
Self {
code: code.into(),
message: message.into(),
raw: None,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "status", content = "payload", rename_all = "snake_case")]
pub enum Resolution {
Ok(serde_json::Value),
Err(ExternalCompletionError),
Timeout,
Cancelled,
}
#[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "status", rename_all = "snake_case")]
pub enum ResolveOutcome {
Accepted,
AlreadyResolved { terminal: Resolution },
UnknownOrRevoked,
}
pub(super) enum ScopedEffectControllerInner<'run> {
Borrowed(&'run dyn RuntimeEffectController),
Shared(Arc<dyn RuntimeEffectController>),
}
impl Clone for ScopedEffectControllerInner<'_> {
fn clone(&self) -> Self {
match self {
Self::Borrowed(controller) => Self::Borrowed(*controller),
Self::Shared(controller) => Self::Shared(Arc::clone(controller)),
}
}
}
#[derive(Clone)]
pub struct ScopedEffectController<'run> {
pub(super) controller: ScopedEffectControllerInner<'run>,
pub(super) scope: ExecutionScope,
}
impl<'run> ScopedEffectController<'run> {
pub fn borrowed(
controller: &'run dyn RuntimeEffectController,
scope: ExecutionScope,
) -> Result<Self, RuntimeError> {
scope.validate()?;
Ok(Self {
controller: ScopedEffectControllerInner::Borrowed(controller),
scope,
})
}
pub fn shared(
controller: Arc<dyn RuntimeEffectController>,
scope: ExecutionScope,
) -> Result<Self, RuntimeError> {
scope.validate()?;
Ok(Self {
controller: ScopedEffectControllerInner::Shared(controller),
scope,
})
}
pub fn controller(&self) -> &dyn RuntimeEffectController {
match &self.controller {
ScopedEffectControllerInner::Borrowed(controller) => *controller,
ScopedEffectControllerInner::Shared(controller) => controller.as_ref(),
}
}
pub fn execution_scope(&self) -> &ExecutionScope {
&self.scope
}
pub fn scope_id(&self) -> &str {
self.scope.id()
}
pub fn turn_id(&self) -> Option<&str> {
self.scope.turn_id()
}
pub(crate) fn to_static(&self) -> Option<ScopedEffectController<'static>> {
let ScopedEffectControllerInner::Shared(controller) = &self.controller else {
return None;
};
Some(ScopedEffectController {
controller: ScopedEffectControllerInner::Shared(Arc::clone(controller)),
scope: self.scope.clone(),
})
}
}
type EffectControllerTaskFuture<'run> = Pin<Box<dyn Future<Output = ()> + Send + 'run>>;
pub(crate) enum EffectControllerTaskRequest {
Execute {
envelope: Box<RuntimeEffectEnvelope>,
local_executor: Box<RuntimeEffectLocalExecutor<'static>>,
response: oneshot::Sender<Result<RuntimeEffectOutcome, RuntimeEffectControllerError>>,
},
AwaitEventKey {
scope: ExecutionScope,
wait: AwaitEventWaitIdentity,
response: oneshot::Sender<Result<AwaitEventKey, RuntimeError>>,
},
ResolveAwaitEvent {
key: AwaitEventKey,
resolution: Resolution,
response: oneshot::Sender<Result<ResolveOutcome, RuntimeError>>,
},
}
impl EffectControllerTaskRequest {
fn into_future<'run>(
self,
controller: &'run dyn RuntimeEffectController,
) -> EffectControllerTaskFuture<'run> {
match self {
Self::Execute {
envelope,
local_executor,
response,
} => Box::pin(async move {
let _ = response.send(controller.execute_effect(*envelope, *local_executor).await);
}),
Self::AwaitEventKey {
scope,
wait,
response,
} => Box::pin(async move {
let _ = response.send(controller.await_event_key(&scope, wait).await);
}),
Self::ResolveAwaitEvent {
key,
resolution,
response,
} => Box::pin(async move {
let _ = response.send(controller.resolve_await_event(&key, resolution).await);
}),
}
}
}
pub(super) struct RemoteLocalExecutionRequest {
pub(super) envelope: RuntimeEffectEnvelope,
pub(super) response:
oneshot::Sender<Result<RuntimeEffectOutcome, RuntimeEffectControllerError>>,
}
#[derive(Clone)]
pub(crate) struct EffectTaskController {
requests: mpsc::UnboundedSender<EffectControllerTaskRequest>,
durability_tier: crate::DurabilityTier,
allows_process_lifetime_completion_keys: bool,
supports_concurrent_effects: bool,
}
impl EffectTaskController {
pub(crate) fn scoped(
controller: &dyn RuntimeEffectController,
scope: ExecutionScope,
) -> Result<
(
ScopedEffectController<'static>,
mpsc::UnboundedReceiver<EffectControllerTaskRequest>,
),
RuntimeError,
> {
let (requests, request_rx) = mpsc::unbounded_channel();
let proxy = Self {
requests,
durability_tier: controller.durability_tier(),
allows_process_lifetime_completion_keys: controller
.allows_process_lifetime_completion_keys(),
supports_concurrent_effects: controller.supports_concurrent_effects(),
};
Ok((
ScopedEffectController::shared(Arc::new(proxy), scope)?,
request_rx,
))
}
}
#[async_trait::async_trait]
impl AwaitEventResolver for EffectTaskController {
fn durability_tier(&self) -> crate::DurabilityTier {
self.durability_tier
}
fn allows_process_lifetime_completion_keys(&self) -> bool {
self.allows_process_lifetime_completion_keys
}
async fn await_event_key(
&self,
scope: &ExecutionScope,
wait: AwaitEventWaitIdentity,
) -> Result<AwaitEventKey, RuntimeError> {
let (response_tx, response_rx) = oneshot::channel();
self.requests
.send(EffectControllerTaskRequest::AwaitEventKey {
scope: scope.clone(),
wait,
response: response_tx,
})
.map_err(|_| {
RuntimeError::new(
"runtime_effect_controller_task_closed",
"await-event key controller task is no longer running",
)
})?;
response_rx.await.map_err(|_| {
RuntimeError::new(
"runtime_effect_controller_task_closed",
"await-event key controller response was dropped",
)
})?
}
async fn resolve_await_event(
&self,
key: &AwaitEventKey,
resolution: Resolution,
) -> Result<ResolveOutcome, RuntimeError> {
let (response_tx, response_rx) = oneshot::channel();
self.requests
.send(EffectControllerTaskRequest::ResolveAwaitEvent {
key: key.clone(),
resolution,
response: response_tx,
})
.map_err(|_| {
RuntimeError::new(
"runtime_effect_controller_task_closed",
"await-event resolution controller task is no longer running",
)
})?;
response_rx.await.map_err(|_| {
RuntimeError::new(
"runtime_effect_controller_task_closed",
"await-event resolution controller response was dropped",
)
})?
}
}
#[async_trait::async_trait]
impl RuntimeEffectController for EffectTaskController {
fn supports_concurrent_effects(&self) -> bool {
self.supports_concurrent_effects
}
async fn execute_effect(
&self,
envelope: RuntimeEffectEnvelope,
local_executor: RuntimeEffectLocalExecutor<'_>,
) -> Result<RuntimeEffectOutcome, RuntimeEffectControllerError> {
let (local_executor, mut local_execution) = local_executor.into_remote_execution();
let (response_tx, response_rx) = oneshot::channel();
self.requests
.send(EffectControllerTaskRequest::Execute {
envelope: Box::new(envelope),
local_executor: Box::new(local_executor),
response: response_tx,
})
.map_err(|_| {
RuntimeEffectControllerError::new(
"runtime_effect_controller_task_closed",
"effect controller task is no longer running",
)
})?;
tokio::pin!(response_rx);
loop {
tokio::select! {
response = &mut response_rx => {
return response.map_err(|_| {
RuntimeEffectControllerError::new(
"runtime_effect_controller_task_closed",
"effect controller response was dropped",
)
})?;
}
request = async {
match local_execution.as_mut() {
Some((_, requests)) => requests.recv().await,
None => std::future::pending().await,
}
} => {
let Some(request) = request else {
local_execution = None;
continue;
};
let Some((executor, _)) = local_execution.take() else {
unreachable!("local execution request requires a local executor");
};
let result = executor.execute_forwarded(request.envelope).await;
let _ = request.response.send(result);
}
}
}
}
}
pub(crate) async fn drive_effect_controller_task(
controller: &dyn RuntimeEffectController,
envelope: RuntimeEffectEnvelope,
local_executor: RuntimeEffectLocalExecutor<'static>,
mut requests: mpsc::UnboundedReceiver<EffectControllerTaskRequest>,
) -> Result<RuntimeEffectOutcome, RuntimeEffectControllerError> {
let (root_tx, root_rx) = oneshot::channel();
let root = EffectControllerTaskRequest::Execute {
envelope: Box::new(envelope),
local_executor: Box::new(local_executor),
response: root_tx,
};
let mut stack = vec![root.into_future(controller)];
let mut requests_open = true;
tokio::pin!(root_rx);
loop {
let Some(active) = stack.last_mut() else {
return root_rx.await.map_err(|_| {
RuntimeEffectControllerError::new(
"runtime_effect_controller_task_closed",
"root effect controller response was dropped",
)
})?;
};
tokio::select! {
biased;
response = &mut root_rx => {
return response.map_err(|_| {
RuntimeEffectControllerError::new(
"runtime_effect_controller_task_closed",
"root effect controller response was dropped",
)
})?;
}
() = active => {
stack.pop();
}
request = async {
if requests_open {
requests.recv().await
} else {
std::future::pending().await
}
} => {
match request {
Some(request) => stack.push(request.into_future(controller)),
None => requests_open = false,
}
}
}
}
}
#[async_trait::async_trait]
pub trait AwaitEventResolver: Send + Sync {
fn durability_tier(&self) -> crate::DurabilityTier {
crate::DurabilityTier::Inline
}
fn allows_process_lifetime_completion_keys(&self) -> bool {
self.durability_tier() == crate::DurabilityTier::Durable
}
async fn await_event_key(
&self,
_scope: &ExecutionScope,
_wait: AwaitEventWaitIdentity,
) -> Result<AwaitEventKey, RuntimeError> {
Err(RuntimeError::new(
"await_event_unsupported",
"this effect boundary does not support await-event keys",
))
}
async fn resolve_await_event(
&self,
_key: &AwaitEventKey,
_resolution: Resolution,
) -> Result<ResolveOutcome, RuntimeError> {
Ok(ResolveOutcome::UnknownOrRevoked)
}
async fn peek_await_event(
&self,
_key: &AwaitEventKey,
) -> Result<Option<Resolution>, RuntimeError> {
Err(RuntimeError::new(
"await_event_unsupported",
"this effect boundary does not support await-event reads",
))
}
async fn await_await_event(
&self,
_key: &AwaitEventKey,
_cancel: CancellationToken,
_deadline: Option<Instant>,
) -> Result<Resolution, RuntimeError> {
Err(RuntimeError::new(
"await_event_unsupported",
"this effect boundary does not support await-event waits",
))
}
async fn revoke_await_events_for_session(&self, _session_id: &str) -> Result<(), RuntimeError> {
Err(RuntimeError::new(
"await_event_unsupported",
"this effect boundary does not support revoking await-event waits",
))
}
async fn cancel_await_events_for_session(&self, _session_id: &str) -> Result<(), RuntimeError> {
Err(RuntimeError::new(
"await_event_cancel_unsupported",
"this effect boundary does not support cancelling durable waits",
))
}
}
#[async_trait::async_trait]
pub trait EffectHost: AwaitEventResolver {
fn scoped<'run>(
&'run self,
scope: ExecutionScope,
) -> Result<ScopedEffectController<'run>, RuntimeError>;
fn scoped_static(
&self,
_scope: ExecutionScope,
) -> Result<Option<ScopedEffectController<'static>>, RuntimeError> {
Ok(None)
}
}
#[async_trait::async_trait]
pub trait RuntimeEffectController: AwaitEventResolver {
fn wants_segment_boundary(&self, _progress: &SegmentProgress) -> Option<BoundaryReason> {
None
}
fn supports_concurrent_effects(&self) -> bool {
true
}
async fn execute_effect(
&self,
envelope: RuntimeEffectEnvelope,
local_executor: RuntimeEffectLocalExecutor<'_>,
) -> Result<RuntimeEffectOutcome, RuntimeEffectControllerError>;
}
#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
pub struct SegmentProgress {
pub effects_executed: u64,
pub journaled_bytes_estimate: Option<u64>,
}
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum BoundaryReason {
JournalBudget,
DurationCap,
}
#[derive(Clone)]
pub(crate) enum RuntimeEffectControllerHandle<'run> {
Borrowed(ScopedEffectController<'run>),
#[cfg(any(test, feature = "testing"))]
Shared {
controller: Arc<dyn RuntimeEffectController>,
scope: ExecutionScope,
},
}
impl<'run> RuntimeEffectControllerHandle<'run> {
pub(crate) fn borrowed(scoped: ScopedEffectController<'run>) -> Self {
Self::Borrowed(scoped)
}
#[cfg(any(test, feature = "testing"))]
pub(crate) fn shared(controller: Arc<dyn RuntimeEffectController>) -> Self {
Self::Shared {
controller,
scope: ExecutionScope::runtime_operation("test-runtime-effect-controller"),
}
}
pub(crate) fn controller(&self) -> &dyn RuntimeEffectController {
match self {
Self::Borrowed(scoped) => scoped.controller(),
#[cfg(any(test, feature = "testing"))]
Self::Shared { controller, .. } => controller.as_ref(),
}
}
pub(crate) fn scoped(&self) -> ScopedEffectController<'_> {
match self {
Self::Borrowed(scoped) => scoped.clone(),
#[cfg(any(test, feature = "testing"))]
Self::Shared { controller, scope } => {
ScopedEffectController::shared(Arc::clone(controller), scope.clone())
.expect("runtime effect controller handle carries a valid scope")
}
}
}
pub(crate) fn clone_scoped(&self) -> RuntimeEffectControllerHandle<'run> {
self.clone()
}
pub(crate) fn to_static(&self) -> Option<RuntimeEffectControllerHandle<'static>> {
match self {
Self::Borrowed(scoped) => scoped
.to_static()
.map(RuntimeEffectControllerHandle::Borrowed),
#[cfg(any(test, feature = "testing"))]
Self::Shared { controller, scope } => Some(RuntimeEffectControllerHandle::Shared {
controller: Arc::clone(controller),
scope: scope.clone(),
}),
}
}
}