objectiveai_mcp_proxy/
dropper.rs

1//! Caller-driven forceful teardown + permanent ban of a response id.
2//!
3//! The embedder (the API) constructs a [`Dropper`], passes a clone into
4//! [`crate::setup`], and keeps a clone to invoke. Calling
5//! [`Dropper::drop`] for a `response_id`:
6//!   1. **forever-bans** it (a concurrency-safe set checked on every
7//!      `initialize`), then
8//!   2. tears down any live session for it.
9//!
10//! Banning before teardown — paired with `handle_initialize` inserting
11//! before it checks the ban — guarantees no `initialize` racing a `drop`
12//! can leave a session alive (see the module's teardown / the
13//! `handle_initialize` checks).
14//!
15//! The teardown context (the [`SessionManager`] and the optional
16//! [`ReverseChannel`]) only exists once `setup` has built it, after the
17//! caller already holds the `Dropper`. It is therefore late-bound via a
18//! `OnceLock`, wired exactly once by `setup` ([`Dropper::wire`]).
19
20use std::sync::{Arc, OnceLock};
21
22use dashmap::DashMap;
23
24use crate::ReverseChannel;
25use crate::session_manager::SessionManager;
26
27/// Cheaply-cloneable handle for banning + dropping response ids. Shared
28/// between the caller (who invokes [`Self::drop`]) and the proxy (which
29/// checks [`Self::is_banned`] on every `initialize`).
30#[derive(Clone, Default)]
31pub struct Dropper {
32    inner: Arc<Inner>,
33}
34
35#[derive(Default)]
36struct Inner {
37    /// Forever-ban set. A response id present here is never (re)admitted.
38    banned: DashMap<String, ()>,
39    /// Teardown context, late-bound by [`Dropper::wire`] from `setup`.
40    ctx: OnceLock<TeardownCtx>,
41}
42
43struct TeardownCtx {
44    sessions: Arc<SessionManager>,
45    reverse_channel: Option<ReverseChannel>,
46}
47
48impl std::fmt::Debug for Dropper {
49    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
50        f.debug_struct("Dropper").finish_non_exhaustive()
51    }
52}
53
54impl Dropper {
55    pub fn new() -> Self {
56        Self::default()
57    }
58
59    /// Forever-ban `response_id`, then tear down any live session for it.
60    ///
61    /// Bans FIRST so an `initialize` for the same id that runs
62    /// concurrently observes the ban (its post-insert check then tears
63    /// the session down itself). Teardown is idempotent, so a double
64    /// teardown across the race is harmless.
65    pub async fn drop(&self, response_id: String) {
66        self.inner.banned.insert(response_id.clone(), ());
67        if let Some(ctx) = self.inner.ctx.get() {
68            teardown(&response_id, &ctx.sessions, ctx.reverse_channel.as_ref()).await;
69        }
70    }
71
72    /// Whether `response_id` has been banned. Checked on every
73    /// `initialize` (before reuse/connect, and again after insert).
74    pub(crate) fn is_banned(&self, response_id: &str) -> bool {
75        self.inner.banned.contains_key(response_id)
76    }
77
78    /// Wire the proxy's teardown context into the dropper. Called once by
79    /// `setup` after it has built the session manager + reverse channel.
80    /// Idempotent (first write wins).
81    pub(crate) fn wire(
82        &self,
83        sessions: Arc<SessionManager>,
84        reverse_channel: Option<ReverseChannel>,
85    ) {
86        let _ = self.inner.ctx.set(TeardownCtx {
87            sessions,
88            reverse_channel,
89        });
90    }
91}
92
93/// Forceful teardown of one response id's session.
94///
95/// Removes the session from the registry (dropping it closes the
96/// proxy-side upstream connections), and — only if at least one upstream
97/// is a `client_objectiveai_mcp` (ws) upstream — fires the reverse-channel
98/// [`Drop`](objectiveai_sdk::client_objectiveai_mcp::server_request::Payload::Drop)
99/// so the CLI tears down its side (connections + plugin subprocesses).
100/// The `Drop` result is discarded. Idempotent: an absent id is a no-op,
101/// and the CLI's `drop` handler is itself idempotent.
102pub(crate) async fn teardown(
103    response_id: &str,
104    sessions: &SessionManager,
105    reverse_channel: Option<&ReverseChannel>,
106) {
107    let Some(session) = sessions.remove(response_id) else {
108        return;
109    };
110    let uses_objectiveai_mcp = session.connections.values().any(|u| u.is_ws());
111    if uses_objectiveai_mcp {
112        if let Some(rc) = reverse_channel {
113            rc.drop_response(response_id.to_string()).await;
114        }
115    }
116    // `session` (Arc<Session>) drops at end of scope → proxy-side upstream
117    // connections close.
118}
objectiveai_mcp_proxy/dropper.rs

objectiveai_mcp_proxy/
dropper.rs
