agent_client_protocol/jsonrpc/
close.rs

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

agent_client_protocol/jsonrpc/
close.rs
