agent_client_protocol/jsonrpc/
run.rs

1//! Run trait for background tasks that run alongside a connection.
2//!
3//! Run implementations are composable background tasks that run while a connection is active.
4//! They're used for things like MCP tool handlers that need to receive calls through
5//! channels and invoke user-provided closures.
6
7use std::future::Future;
8use std::marker::PhantomData;
9
10use crate::{
11    ConnectionTo,
12    jsonrpc::{ConnectionContext, RawConnectionContext, connection_context},
13    role::Role,
14};
15
16/// A background task that runs alongside a connection.
17///
18/// `RunIn<R>` means "run in the context of being role R". The task receives
19/// a `ConnectionTo<R::Counterpart>` for communicating with the other side.
20///
21/// Implementations are composed using [`ChainRun`] and run in parallel
22/// when the connection is active.
23pub trait RunWithConnectionTo<Counterpart: Role>: Send {
24    /// Run this task to completion.
25    fn run_with_connection_to(
26        self,
27        cx: ConnectionTo<Counterpart>,
28    ) -> impl Future<Output = Result<(), crate::Error>> + Send;
29}
30
31/// A no-op RunIn that completes immediately.
32#[derive(Debug, Default)]
33pub struct NullRun;
34
35impl<Counterpart: Role> RunWithConnectionTo<Counterpart> for NullRun {
36    fn run_with_connection_to(
37        self,
38        _cx: ConnectionTo<Counterpart>,
39    ) -> impl Future<Output = Result<(), crate::Error>> + Send {
40        std::future::ready(Ok(()))
41    }
42}
43
44/// Chains two RunIn implementations to run in parallel.
45#[derive(Debug)]
46pub struct ChainRun<A, B> {
47    a: A,
48    b: B,
49}
50
51impl<A, B> ChainRun<A, B> {
52    /// Create a new chained RunIn from two RunIn implementations.
53    pub fn new(a: A, b: B) -> Self {
54        Self { a, b }
55    }
56}
57
58impl<Counterpart: Role, A, B> RunWithConnectionTo<Counterpart> for ChainRun<A, B>
59where
60    A: RunWithConnectionTo<Counterpart>,
61    B: RunWithConnectionTo<Counterpart>,
62{
63    async fn run_with_connection_to(
64        self,
65        cx: ConnectionTo<Counterpart>,
66    ) -> Result<(), crate::Error> {
67        // Box the futures to avoid stack overflow with deeply nested RunIn chains
68        let a_fut = Box::pin(self.a.run_with_connection_to(cx.clone()));
69        let b_fut = Box::pin(self.b.run_with_connection_to(cx.clone()));
70        let ((), ()) = futures::future::try_join(a_fut, b_fut).await?;
71        Ok(())
72    }
73}
74
75/// A RunIn created from a closure via [`with_spawned`](crate::Builder::with_spawned).
76pub struct SpawnedRun<F, Context = RawConnectionContext> {
77    task_fn: F,
78    location: &'static std::panic::Location<'static>,
79    context: PhantomData<fn() -> Context>,
80}
81
82impl<F, Context> SpawnedRun<F, Context> {
83    /// Create a new spawned RunIn from a closure.
84    pub fn new(location: &'static std::panic::Location<'static>, task_fn: F) -> Self {
85        Self {
86            task_fn,
87            location,
88            context: PhantomData,
89        }
90    }
91}
92
93impl<Counterpart, F, Fut, Context> RunWithConnectionTo<Counterpart> for SpawnedRun<F, Context>
94where
95    Counterpart: Role,
96    Context: ConnectionContext,
97    F: FnOnce(Context::Connection<Counterpart>) -> Fut + Send,
98    Fut: Future<Output = Result<(), crate::Error>> + Send,
99{
100    async fn run_with_connection_to(
101        self,
102        connection: ConnectionTo<Counterpart>,
103    ) -> Result<(), crate::Error> {
104        let location = self.location;
105        (self.task_fn)(connection_context::from_raw::<Context, _>(connection))
106            .await
107            .map_err(|err| {
108                let data = err.data.clone();
109                err.data(serde_json::json!({
110                    "spawned_at": format!("{}:{}:{}", location.file(), location.line(), location.column()),
111                    "data": data,
112                }))
113            })
114    }
115}
agent_client_protocol/jsonrpc/run.rs

agent_client_protocol/jsonrpc/
run.rs
