agent_client_protocol/jsonrpc/
close.rs

1//! Connection-close handlers.
2
3use std::fmt::Debug;
4use std::future::Future;
5use std::marker::PhantomData;
6
7use crate::{
8    ConnectionTo,
9    jsonrpc::{ConnectionContext, RawConnectionContext, connection_context},
10    role::Role,
11};
12
13/// A handler that runs after the incoming transport reaches clean EOF.
14///
15/// Close handlers are composed by [`Builder::on_close`](crate::Builder::on_close)
16/// and run sequentially in registration order. Unlike
17/// [`RunWithConnectionTo`](crate::RunWithConnectionTo), they are a distinct
18/// connection-lifecycle phase rather than concurrent background work.
19pub trait HandleConnectionClose<Counterpart: Role>: Send {
20    /// Run this handler during clean incoming-transport shutdown.
21    fn handle_connection_close(
22        self,
23        connection: ConnectionTo<Counterpart>,
24    ) -> impl Future<Output = Result<(), crate::Error>> + Send;
25}
26
27/// A close handler that does nothing.
28#[derive(Debug, Default)]
29pub struct NullClose;
30
31impl<Counterpart: Role> HandleConnectionClose<Counterpart> for NullClose {
32    fn handle_connection_close(
33        self,
34        _connection: ConnectionTo<Counterpart>,
35    ) -> impl Future<Output = Result<(), crate::Error>> + Send {
36        std::future::ready(Ok(()))
37    }
38}
39
40pub(crate) struct CloseCallback<F, Context = RawConnectionContext> {
41    callback: F,
42    context: PhantomData<fn() -> Context>,
43}
44
45impl<F, Context> CloseCallback<F, Context> {
46    pub(crate) fn new(callback: F) -> Self {
47        Self {
48            callback,
49            context: PhantomData,
50        }
51    }
52}
53
54impl<F, Context> Debug for CloseCallback<F, Context> {
55    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
56        formatter
57            .debug_struct("CloseCallback")
58            .finish_non_exhaustive()
59    }
60}
61
62impl<Counterpart, F, Fut, Context> HandleConnectionClose<Counterpart> for CloseCallback<F, Context>
63where
64    Counterpart: Role,
65    Context: ConnectionContext,
66    F: FnOnce(Context::Connection<Counterpart>) -> Fut + Send,
67    Fut: Future<Output = Result<(), crate::Error>> + Send,
68{
69    async fn handle_connection_close(
70        self,
71        connection: ConnectionTo<Counterpart>,
72    ) -> Result<(), crate::Error> {
73        let result = (self.callback)(connection_context::from_raw::<Context, _>(connection)).await;
74        if let Err(error) = &result {
75            tracing::warn!(?error, "Connection close callback failed");
76        }
77        result
78    }
79}
80
81#[derive(Debug)]
82pub(crate) struct ChainedClose<A, B> {
83    first: A,
84    second: B,
85}
86
87impl<A, B> ChainedClose<A, B> {
88    pub(crate) fn new(first: A, second: B) -> Self {
89        Self { first, second }
90    }
91}
92
93impl<Counterpart, A, B> HandleConnectionClose<Counterpart> for ChainedClose<A, B>
94where
95    Counterpart: Role,
96    A: HandleConnectionClose<Counterpart>,
97    B: HandleConnectionClose<Counterpart>,
98{
99    async fn handle_connection_close(
100        self,
101        connection: ConnectionTo<Counterpart>,
102    ) -> Result<(), crate::Error> {
103        // Box each side to keep deeply composed close chains from producing
104        // correspondingly deep connection-driver futures.
105        let first = Box::pin(self.first.handle_connection_close(connection.clone())).await;
106        let second = Box::pin(self.second.handle_connection_close(connection)).await;
107        first.and(second)
108    }
109}
agent_client_protocol/jsonrpc/close.rs

agent_client_protocol/jsonrpc/
close.rs
