Skip to main content

starweaver_context/
context_protocol.rs

1//! Context protocol support types carried by [`crate::AgentContext`].
2
3use std::collections::BTreeMap;
4
5use serde::{Deserialize, Serialize};
6use serde_json::Value;
7use starweaver_core::Metadata;
8use starweaver_model::{ModelMessage, ModelRequestPart, ModelResponsePart};
9use uuid::Uuid;
10
11/// Metadata for a registered agent.
12#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
13pub struct AgentInfo {
14    /// Unique identifier for the agent.
15    pub agent_id: String,
16    /// Human-readable agent name.
17    pub agent_name: String,
18    /// Parent agent id, if this is a subagent.
19    #[serde(default, skip_serializing_if = "Option::is_none")]
20    pub parent_agent_id: Option<String>,
21}
22
23impl AgentInfo {
24    /// Create agent metadata.
25    #[must_use]
26    pub fn new(agent_id: impl Into<String>, agent_name: impl Into<String>) -> Self {
27        Self {
28            agent_id: agent_id.into(),
29            agent_name: agent_name.into(),
30            parent_agent_id: None,
31        }
32    }
33
34    /// Attach parent agent id.
35    #[must_use]
36    pub fn with_parent_agent_id(mut self, parent_agent_id: impl Into<String>) -> Self {
37        self.parent_agent_id = Some(parent_agent_id.into());
38        self
39    }
40}
41
42/// Metadata for one deferred tool call.
43pub type DeferredToolMetadata = Metadata;
44
45/// Runtime wrapper metadata passed to model/subagent wrapper equivalents.
46pub type WrapperMetadata = Metadata;
47
48/// Context lifecycle fields for enter, exit, streaming, and compaction state.
49#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
50pub struct ContextLifecycleState {
51    /// Whether the context has been entered.
52    #[serde(default)]
53    pub entered: bool,
54    /// Whether stream queue side-channel behavior is enabled.
55    #[serde(default)]
56    pub stream_queue_enabled: bool,
57    /// Current compact recursion depth.
58    #[serde(default)]
59    pub compact_depth: u32,
60}
61
62impl ContextLifecycleState {
63    /// Return whether this state has default values.
64    #[must_use]
65    pub fn is_default(&self) -> bool {
66        self == &Self::default()
67    }
68}
69
70/// Tool call ID normalizer for cross-provider tool call matching.
71#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
72pub struct ToolIdWrapper {
73    /// Prefix for normalized tool call IDs.
74    #[serde(default = "default_tool_id_prefix")]
75    pub prefix: String,
76    /// Original provider id to normalized id mapping.
77    #[serde(default, skip_serializing_if = "BTreeMap::is_empty")]
78    pub tool_call_maps: BTreeMap<String, String>,
79}
80
81impl Default for ToolIdWrapper {
82    fn default() -> Self {
83        Self {
84            prefix: default_tool_id_prefix(),
85            tool_call_maps: BTreeMap::new(),
86        }
87    }
88}
89
90impl ToolIdWrapper {
91    /// Clear cached mappings.
92    pub fn clear(&mut self) {
93        self.tool_call_maps.clear();
94    }
95
96    /// Normalize a tool call id.
97    pub fn upsert_tool_call_id(&mut self, tool_call_id: &str) -> String {
98        if tool_call_id.starts_with(&self.prefix) {
99            return tool_call_id.to_string();
100        }
101        if let Some(existing) = self.tool_call_maps.get(tool_call_id) {
102            return existing.clone();
103        }
104        let wrapped = format!("{}{}", self.prefix, Uuid::new_v4().simple());
105        self.tool_call_maps
106            .insert(tool_call_id.to_string(), wrapped.clone());
107        wrapped
108    }
109
110    /// Wrap all tool call IDs in a message history in place.
111    pub fn wrap_messages(&mut self, message_history: &mut [ModelMessage]) {
112        for message in message_history {
113            self.wrap_message(message);
114        }
115    }
116
117    /// Wrap all tool call IDs in one message in place.
118    pub fn wrap_message(&mut self, message: &mut ModelMessage) {
119        match message {
120            ModelMessage::Request(request) => {
121                for part in &mut request.parts {
122                    match part {
123                        ModelRequestPart::ToolReturn(tool_return) => {
124                            tool_return.tool_call_id =
125                                self.upsert_tool_call_id(&tool_return.tool_call_id);
126                        }
127                        ModelRequestPart::RetryPrompt { tool_call_id, .. } => {
128                            if let Some(id) = tool_call_id {
129                                *id = self.upsert_tool_call_id(id);
130                            }
131                        }
132                        ModelRequestPart::SystemPrompt { .. }
133                        | ModelRequestPart::UserPrompt { .. }
134                        | ModelRequestPart::Instruction { .. } => {}
135                    }
136                }
137            }
138            ModelMessage::Response(response) => {
139                for part in &mut response.parts {
140                    self.wrap_response_part(part);
141                }
142            }
143        }
144    }
145
146    /// Wrap all tool call IDs in one response part in place.
147    pub fn wrap_response_part(&mut self, part: &mut ModelResponsePart) {
148        match part {
149            ModelResponsePart::ToolCall(call)
150            | ModelResponsePart::ProviderToolCall { call, .. } => {
151                call.id = self.upsert_tool_call_id(&call.id);
152            }
153            ModelResponsePart::NativeToolCall { payload, .. }
154            | ModelResponsePart::NativeToolReturn { payload, .. }
155            | ModelResponsePart::ProviderOpaque { payload, .. } => {
156                self.wrap_provider_payload(payload);
157            }
158            ModelResponsePart::Text { .. }
159            | ModelResponsePart::ProviderText { .. }
160            | ModelResponsePart::Thinking { .. }
161            | ModelResponsePart::ProviderThinking { .. }
162            | ModelResponsePart::File { .. }
163            | ModelResponsePart::Compaction { .. } => {}
164        }
165    }
166
167    /// Wrap tool call IDs in provider-native payload objects when they use common id keys.
168    pub fn wrap_provider_payload(&mut self, payload: &mut Value) {
169        let Some(object) = payload.as_object_mut() else {
170            return;
171        };
172        for key in ["tool_call_id", "toolUseId", "tool_use_id", "call_id", "id"] {
173            let Some(value) = object.get_mut(key) else {
174                continue;
175            };
176            let Some(id) = value.as_str().map(ToString::to_string) else {
177                continue;
178            };
179            *value = Value::String(self.upsert_tool_call_id(&id));
180        }
181    }
182
183    /// Return whether no mappings are stored.
184    #[must_use]
185    pub fn is_empty(&self) -> bool {
186        self.tool_call_maps.is_empty()
187    }
188}
189
190/// Runtime-only stream queue registry placeholder.
191#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
192pub struct AgentStreamQueueRegistry {
193    /// Queue names/ids known to the context. Actual async queues live outside serializable state.
194    #[serde(default, skip_serializing_if = "Vec::is_empty")]
195    pub queues: Vec<String>,
196}
197
198impl AgentStreamQueueRegistry {
199    /// Return whether no queues are registered.
200    #[must_use]
201    pub const fn is_empty(&self) -> bool {
202        self.queues.is_empty()
203    }
204}
205
206/// Loaded tool-search state.
207#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
208pub struct ToolSearchState {
209    /// Tool names loaded via tool search.
210    #[serde(default, skip_serializing_if = "Vec::is_empty")]
211    pub loaded_tools: Vec<String>,
212    /// Namespace IDs loaded via tool search.
213    #[serde(default, skip_serializing_if = "Vec::is_empty")]
214    pub loaded_namespaces: Vec<String>,
215}
216
217impl ToolSearchState {
218    /// Return whether no tool-search state is present.
219    #[must_use]
220    pub const fn is_empty(&self) -> bool {
221        self.loaded_tools.is_empty() && self.loaded_namespaces.is_empty()
222    }
223}
224
225/// Removed tool-search state after host invalidation or refresh.
226#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
227pub struct ToolSearchInvalidation {
228    /// Tool names removed from loaded tool-search state.
229    #[serde(default, skip_serializing_if = "Vec::is_empty")]
230    pub removed_tools: Vec<String>,
231    /// Namespace IDs removed from loaded tool-search state.
232    #[serde(default, skip_serializing_if = "Vec::is_empty")]
233    pub removed_namespaces: Vec<String>,
234}
235
236impl ToolSearchInvalidation {
237    /// Return whether no loaded tool-search entries were removed.
238    #[must_use]
239    pub const fn is_empty(&self) -> bool {
240        self.removed_tools.is_empty() && self.removed_namespaces.is_empty()
241    }
242}
243
244/// Runtime model wrapper placeholder. Actual wrapper functions are crate-specific dependencies.
245pub type ModelWrapperMetadata = Value;
246
247fn default_tool_id_prefix() -> String {
248    "sw-tool-".to_string()
249}