objectiveai_mcp/
agent_args_registry.rs

1//! Per-rmcp-session in-memory map of [`SessionState`]. Populated by
2//! [`crate::header_session_manager::HeaderSessionManager`] on every
3//! `initialize` (fresh, no-session-id POST, or lazy-rehydration of
4//! a reconnect with an id we haven't seen this process lifetime),
5//! consumed by every tool handler before dispatching to the
6//! executor.
7//!
8//! FULL-REPLACE semantics on every record: a reconnect with a new
9//! header set wholesale replaces the prior entry, and missing
10//! headers become `None` (filter values) / `true` (root) on the new
11//! struct. Mirrors the lifecycle of
12//! `objectiveai-mcp-proxy::Session::transient_headers` and the
13//! `SessionRegistry` in `psychological-operations-x-api-mcp`.
14//!
15//! In-memory only. A process restart silently flushes the map; the
16//! CLI re-sends the headers on its next request, and the
17//! lazy-rehydration path re-captures them.
18
19use std::collections::HashMap;
20use std::sync::Arc;
21
22use objectiveai_sdk::agent::ClientObjectiveaiMcpEntry;
23use objectiveai_sdk::cli::command::AgentArguments;
24use rmcp::transport::common::server_side_http::SessionId;
25use tokio::sync::RwLock;
26
27/// Per-session state recorded by the header parser. Wraps the
28/// legacy [`AgentArguments`] identity bag alongside the three
29/// optional `X-OBJECTIVEAI-MCP-*` filter values, so a single
30/// `(record, get, remove)` call cycle covers both. The wrapper
31/// shape keeps the registry's inner map single — one entry per
32/// session, projected by callers via `.args` or the `mcp_*`
33/// fields as needed.
34///
35/// `mcp_root` resolves to `true` when the
36/// `X-OBJECTIVEAI-MCP-ROOT` header is absent on connect — the
37/// header parser writes the resolved value here, so anything
38/// stored is the final per-session decision. `mcp_tools` /
39/// `mcp_plugins` are `None` when the corresponding header is
40/// absent ⇒ no filter at list time. `Some(vec![])` means
41/// "explicitly allow none."
42#[derive(Debug, Clone)]
43pub struct SessionState {
44    pub args: AgentArguments,
45    pub mcp_root: bool,
46    pub mcp_tools: Option<Vec<ClientObjectiveaiMcpEntry>>,
47    pub mcp_plugins: Option<Vec<ClientObjectiveaiMcpEntry>>,
48}
49
50/// Shared registry of per-session [`SessionState`]. Cheap to clone
51/// (the inner state is `Arc`'d).
52#[derive(Default, Debug, Clone)]
53pub struct AgentArgumentsRegistry {
54    inner: Arc<RwLock<HashMap<SessionId, Arc<SessionState>>>>,
55}
56
57impl AgentArgumentsRegistry {
58    pub fn new() -> Self {
59        Self::default()
60    }
61
62    /// Insert or FULL-REPLACE the entry for `id` with `state`. Any
63    /// prior state (with whatever fields it had) is discarded.
64    pub async fn record(&self, id: SessionId, state: Arc<SessionState>) {
65        self.inner.write().await.insert(id, state);
66    }
67
68    /// Look up the current state for `id`. Returns `None` if no
69    /// session has been registered under this id.
70    pub async fn get(&self, id: &SessionId) -> Option<Arc<SessionState>> {
71        self.inner.read().await.get(id).cloned()
72    }
73
74    /// Drop the entry for `id`. Returns the prior value if any.
75    pub async fn remove(&self, id: &SessionId) -> Option<Arc<SessionState>> {
76        self.inner.write().await.remove(id)
77    }
78}
objectiveai_mcp/agent_args_registry.rs

objectiveai_mcp/
agent_args_registry.rs
