Skip to main content

objectiveai_sdk/agent/
continuation.rs

1use base64::Engine;
2use schemars::JsonSchema;
3use serde::{Deserialize, Serialize};
4
5/// Continuation state for multi-turn agent completions.
6///
7/// Returned in the final streaming chunk and in unary responses.
8/// Pass it back in the next request to continue the conversation.
9/// Serialized as base64-encoded JSON.
10#[derive(Debug, Clone, Serialize, Deserialize, JsonSchema)]
11#[serde(untagged)]
12#[schemars(rename = "agent.Continuation")]
13pub enum Continuation {
14    #[schemars(title = "Openrouter")]
15    Openrouter(super::openrouter::Continuation),
16    #[schemars(title = "ClaudeAgentSdk")]
17    ClaudeAgentSdk(super::claude_agent_sdk::Continuation),
18    #[schemars(title = "CodexSdk")]
19    CodexSdk(super::codex_sdk::Continuation),
20    #[schemars(title = "Mock")]
21    Mock(super::mock::Continuation),
22}
23
24impl From<super::openrouter::Continuation> for Continuation {
25    fn from(inner: super::openrouter::Continuation) -> Self {
26        Self::Openrouter(inner)
27    }
28}
29
30impl From<super::claude_agent_sdk::Continuation> for Continuation {
31    fn from(inner: super::claude_agent_sdk::Continuation) -> Self {
32        Self::ClaudeAgentSdk(inner)
33    }
34}
35
36impl From<super::codex_sdk::Continuation> for Continuation {
37    fn from(inner: super::codex_sdk::Continuation) -> Self {
38        Self::CodexSdk(inner)
39    }
40}
41
42impl From<super::mock::Continuation> for Continuation {
43    fn from(inner: super::mock::Continuation) -> Self {
44        Self::Mock(inner)
45    }
46}
47
48impl Continuation {
49    /// Returns the MCP sessions map for this continuation.
50    pub fn mcp_sessions(&self) -> &indexmap::IndexMap<String, String> {
51        match self {
52            Self::Openrouter(c) => &c.mcp_sessions,
53            Self::ClaudeAgentSdk(c) => &c.mcp_sessions,
54            Self::CodexSdk(c) => &c.mcp_sessions,
55            Self::Mock(c) => &c.mcp_sessions,
56        }
57    }
58
59    /// Per-agent reverse-attach `ws_session_id` baked into the
60    /// `client_objectiveai_mcp` proxy URL path segment, if this
61    /// agent used `client_objectiveai_mcp` in a prior turn. The WS
62    /// streaming handler reads this on resume so the agent's MCP
63    /// proxy URL stays valid against the new WS reverse channel.
64    pub fn ws_session_id(&self) -> Option<&str> {
65        match self {
66            Self::Openrouter(c) => c.ws_session_id.as_deref(),
67            Self::ClaudeAgentSdk(c) => c.ws_session_id.as_deref(),
68            Self::CodexSdk(c) => c.ws_session_id.as_deref(),
69            Self::Mock(c) => c.ws_session_id.as_deref(),
70        }
71    }
72
73    /// Stamps the per-agent reverse-attach `ws_session_id` on the
74    /// outgoing continuation. Called by the agent client after it
75    /// resolves the id (incoming continuation's id or freshly minted).
76    pub fn set_ws_session_id(&mut self, id: Option<String>) {
77        match self {
78            Self::Openrouter(c) => c.ws_session_id = id,
79            Self::ClaudeAgentSdk(c) => c.ws_session_id = id,
80            Self::CodexSdk(c) => c.ws_session_id = id,
81            Self::Mock(c) => c.ws_session_id = id,
82        }
83    }
84
85    /// Full slash-separated lineage of the agent this continuation
86    /// belongs to. See per-upstream struct docs for the semantic.
87    /// Empty string if the continuation was minted before this field
88    /// existed (pre-field tokens deserialize the default).
89    pub fn agent_instance_hierarchy(&self) -> &str {
90        match self {
91            Self::Openrouter(c) => c.agent_instance_hierarchy.as_str(),
92            Self::ClaudeAgentSdk(c) => c.agent_instance_hierarchy.as_str(),
93            Self::CodexSdk(c) => c.agent_instance_hierarchy.as_str(),
94            Self::Mock(c) => c.agent_instance_hierarchy.as_str(),
95        }
96    }
97
98    /// Stamps `agent_instance_hierarchy` on the outgoing continuation.
99    /// Mirrors [`Self::set_ws_session_id`]: the api server calls this
100    /// just before returning the wire continuation so the next round
101    /// inherits the same lineage.
102    pub fn set_agent_instance_hierarchy(&mut self, id: String) {
103        match self {
104            Self::Openrouter(c) => c.agent_instance_hierarchy = id,
105            Self::ClaudeAgentSdk(c) => c.agent_instance_hierarchy = id,
106            Self::CodexSdk(c) => c.agent_instance_hierarchy = id,
107            Self::Mock(c) => c.agent_instance_hierarchy = id,
108        }
109    }
110
111    /// Returns the upstream type for this continuation.
112    pub fn upstream(&self) -> super::Upstream {
113        match self {
114            Self::Openrouter(_) => super::Upstream::Openrouter,
115            Self::ClaudeAgentSdk(_) => super::Upstream::ClaudeAgentSdk,
116            Self::CodexSdk(_) => super::Upstream::CodexSdk,
117            Self::Mock(_) => super::Upstream::Mock,
118        }
119    }
120
121    /// Serializes the continuation to a base64-encoded string.
122    pub fn to_string(&self) -> String {
123        let json = serde_json::to_string(self).unwrap();
124        base64::engine::general_purpose::STANDARD.encode(json)
125    }
126
127    /// Attempts to deserialize a continuation from a base64-encoded string.
128    pub fn try_from_string(s: &str) -> Option<Self> {
129        let json = base64::engine::general_purpose::STANDARD.decode(s).ok()?;
130        let continuation = serde_json::from_slice(&json).ok()?;
131        Some(continuation)
132    }
133}