Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
96 changes: 85 additions & 11 deletions crates/core/src/api/registry.rs
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
use crate::api::runtime::{
EventSanitizeFn, LlmConditionalFn, LlmExecutionFn, LlmRequestInterceptFn, LlmSanitizeRequestFn,
LlmSanitizeResponseFn, LlmStreamExecutionFn, ToolConditionalFn, ToolExecutionFn,
ToolInterceptFn, ToolSanitizeFn,
ToolExecutionFrameFn, ToolInterceptFn, ToolSanitizeFn,
};
use crate::api::runtime::{current_scope_stack, global_context};
use crate::api::shared::ensure_runtime_owner;
Expand Down Expand Up @@ -203,12 +203,10 @@ macro_rules! global_intercept_registry_api {
};
}

macro_rules! global_execution_registry_api {
macro_rules! global_execution_registry_register_api {
(
$(#[$register_meta:meta])*
$register_name:ident,
$(#[$deregister_meta:meta])*
$deregister_name:ident,
$field:ident,
$fn_type:ty
) => {
Expand All @@ -233,10 +231,18 @@ macro_rules! global_execution_registry_api {
.map_err(|error| FlowError::Internal(error.to_string()))?;
state
.$field
.register(ExecutionIntercept::new(name, priority, callable))
.register(ExecutionIntercept::new(name, priority, callable.into()))
.map_err(FlowError::AlreadyExists)
}
};
}

macro_rules! global_execution_registry_deregister_api {
(
$(#[$deregister_meta:meta])*
$deregister_name:ident,
$field:ident
) => {
$(#[$deregister_meta])*
///
/// # Parameters
Expand All @@ -259,6 +265,29 @@ macro_rules! global_execution_registry_api {
};
}

macro_rules! global_execution_registry_api {
(
$(#[$register_meta:meta])*
$register_name:ident,
$(#[$deregister_meta:meta])*
$deregister_name:ident,
$field:ident,
$fn_type:ty
) => {
global_execution_registry_register_api!(
$(#[$register_meta])*
$register_name,
$field,
$fn_type
);
global_execution_registry_deregister_api!(
$(#[$deregister_meta])*
$deregister_name,
$field
);
};
}

macro_rules! scope_guardrail_registry_api {
(
$(#[$register_meta:meta])*
Expand Down Expand Up @@ -400,12 +429,10 @@ macro_rules! scope_intercept_registry_api {
};
}

macro_rules! scope_execution_registry_api {
macro_rules! scope_execution_registry_register_api {
(
$(#[$register_meta:meta])*
$register_name:ident,
$(#[$deregister_meta:meta])*
$deregister_name:ident,
$field:ident,
$fn_type:ty
) => {
Expand Down Expand Up @@ -438,10 +465,18 @@ macro_rules! scope_execution_registry_api {
.ok_or_else(|| FlowError::NotFound(format!("scope {scope_uuid} not found")))?;
registries
.$field
.register(ExecutionIntercept::new(name, priority, callable))
.register(ExecutionIntercept::new(name, priority, callable.into()))
.map_err(FlowError::AlreadyExists)
}
};
}

macro_rules! scope_execution_registry_deregister_api {
(
$(#[$deregister_meta:meta])*
$deregister_name:ident,
$field:ident
) => {
$(#[$deregister_meta])*
///
/// # Parameters
Expand All @@ -467,6 +502,29 @@ macro_rules! scope_execution_registry_api {
};
}

macro_rules! scope_execution_registry_api {
(
$(#[$register_meta:meta])*
$register_name:ident,
$(#[$deregister_meta:meta])*
$deregister_name:ident,
$field:ident,
$fn_type:ty
) => {
scope_execution_registry_register_api!(
$(#[$register_meta])*
$register_name,
$field,
$fn_type
);
scope_execution_registry_deregister_api!(
$(#[$deregister_meta])*
$deregister_name,
$field
);
};
}

global_guardrail_registry_api!(
/// Register a global mark event sanitizer.
register_mark_sanitize_guardrail,
Expand Down Expand Up @@ -535,14 +593,22 @@ global_intercept_registry_api!(
global_execution_registry_api!(
/// Register a global tool execution intercept.
/// Execution intercepts can wrap or replace the tool callback. Each
/// callback returns a canonical tool execution outcome, while its
/// callback returns Relay's tool execution outcome wrapper, while its
/// continuation resolves to the raw downstream result JSON.
register_tool_execution_intercept,
/// Deregister a global tool execution intercept.
deregister_tool_execution_intercept,
tool_execution_intercepts,
ToolExecutionFn
);
global_execution_registry_register_api!(
/// Register a global annotation-aware tool execution intercept.
/// Frame intercepts share the existing tool execution registry, namespace,
/// and priority order with raw-JSON intercepts.
register_tool_execution_frame_intercept,
tool_execution_intercepts,
ToolExecutionFrameFn
);

global_guardrail_registry_api!(
/// Register a global LLM sanitize-request guardrail.
Expand Down Expand Up @@ -669,14 +735,22 @@ scope_intercept_registry_api!(
scope_execution_registry_api!(
/// Register a scope-local tool execution intercept.
/// Execution intercepts can wrap or replace the tool callback inside the
/// owning scope. Each callback returns a canonical tool execution outcome,
/// owning scope. Each callback returns Relay's tool execution outcome wrapper,
/// while its continuation resolves to the raw downstream result JSON.
scope_register_tool_execution_intercept,
/// Deregister a scope-local tool execution intercept.
scope_deregister_tool_execution_intercept,
tool_execution_intercepts,
ToolExecutionFn
);
scope_execution_registry_register_api!(
/// Register a scope-local annotation-aware tool execution intercept.
/// Frame intercepts share the existing tool execution registry, namespace,
/// and priority order with raw-JSON intercepts.
scope_register_tool_execution_frame_intercept,
tool_execution_intercepts,
ToolExecutionFrameFn
);

scope_guardrail_registry_api!(
/// Register a scope-local LLM sanitize-request guardrail.
Expand Down
3 changes: 2 additions & 1 deletion crates/core/src/api/runtime.rs
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,8 @@ pub use callbacks::{
LlmRequestInterceptFn, LlmSanitizeRequestContext, LlmSanitizeRequestFn,
LlmSanitizeResponseContext, LlmSanitizeResponseFn, LlmStreamExecutionFn,
LlmStreamExecutionNextFn, LlmStreamInner, ToolConditionalFn, ToolExecutionFn,
ToolExecutionNextFn, ToolInterceptFn, ToolSanitizeFn,
ToolExecutionFrameFn, ToolExecutionFrameNextFn, ToolExecutionNextFn, ToolInterceptFn,
ToolSanitizeFn,
};
#[doc(hidden)]
pub use continuation_context::MiddlewareContinuationContext;
Expand Down
63 changes: 57 additions & 6 deletions crates/core/src/api/runtime/callbacks.rs
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,9 @@ use tokio_stream::Stream;

use crate::api::event::{Event, EventSanitizeFields};
use crate::api::llm::{LlmRequest, LlmRequestInterceptOutcome};
use crate::api::tool::ToolExecutionInterceptOutcome;
use crate::api::tool::{
ToolExecutionFrame, ToolExecutionFrameOutcome, ToolExecutionInterceptOutcome,
};
use crate::codec::request::AnnotatedLlmRequest;
use crate::codec::traits::{LlmCodec, LlmResponseCodec};
use crate::error::Result;
Expand Down Expand Up @@ -128,8 +130,9 @@ pub type ToolExecutionNextFn =
/// - Third argument: Continuation for the remaining execution chain.
///
/// # Returns
/// A future resolving to the canonical tool execution outcome, containing the
/// tool result and any pending lifecycle marks produced by this intercept.
/// A future resolving to Relay's execution outcome wrapper, containing the
/// harness-owned tool result and any pending lifecycle marks produced by this
/// intercept.
///
/// # Errors
/// The future resolves to an error when the intercept or remaining execution
Expand All @@ -144,13 +147,61 @@ pub type ToolExecutionFn = Arc<
+ Sync,
>;

/// Internal continuation carrying both a tool result and accumulated marks.
pub(crate) type ToolExecutionOutcomeNextFn = Arc<
dyn Fn(Json) -> Pin<Box<dyn Future<Output = Result<ToolExecutionInterceptOutcome>> + Send>>
/// Annotation-aware continuation invoked by tool execution frame intercepts.
///
/// The continuation exposes the raw downstream result together with its
/// optional opaque annotation. Relay retains downstream pending marks
/// internally, matching [`ToolExecutionNextFn`].
pub type ToolExecutionFrameNextFn = Arc<
dyn Fn(Json) -> Pin<Box<dyn Future<Output = Result<ToolExecutionFrame>> + Send>> + Send + Sync,
>;

/// Annotation-aware tool execution intercept.
///
/// This callback participates in the same priority-ordered chain as
/// [`ToolExecutionFn`], but its continuation and outcome carry a
/// [`ToolExecutionFrame`].
pub type ToolExecutionFrameFn = Arc<
dyn Fn(
&str,
Json,
ToolExecutionFrameNextFn,
) -> Pin<Box<dyn Future<Output = Result<ToolExecutionFrameOutcome>> + Send>>
+ Send
+ Sync,
>;

/// Internal continuation carrying a tool frame and accumulated marks.
pub(crate) type ToolExecutionFrameOutcomeNextFn = Arc<
dyn Fn(Json) -> Pin<Box<dyn Future<Output = Result<ToolExecutionFrameOutcome>> + Send>>
+ Send
+ Sync,
>;

/// One registry payload for legacy and annotation-aware tool intercepts.
///
/// Keeping both callback forms in this enum preserves a single namespace and
/// priority order rather than creating a second middleware chain.
#[derive(Clone)]
pub(crate) enum ToolExecutionCallable {
/// Existing raw-JSON execution intercept.
Legacy(ToolExecutionFn),
/// Annotation-aware frame execution intercept.
Frame(ToolExecutionFrameFn),
}

impl From<ToolExecutionFn> for ToolExecutionCallable {
fn from(value: ToolExecutionFn) -> Self {
Self::Legacy(value)
}
}

impl From<ToolExecutionFrameFn> for ToolExecutionCallable {
fn from(value: ToolExecutionFrameFn) -> Self {
Self::Frame(value)
}
}

/// Relay's built-in LLM codec identities.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BuiltinLlmCodec {
Expand Down
Loading