starweaver_context/
context_protocol.rs1use 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#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
13pub struct AgentInfo {
14 pub agent_id: String,
16 pub agent_name: String,
18 #[serde(default, skip_serializing_if = "Option::is_none")]
20 pub parent_agent_id: Option<String>,
21}
22
23impl AgentInfo {
24 #[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 #[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
42pub type DeferredToolMetadata = Metadata;
44
45pub type WrapperMetadata = Metadata;
47
48#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
50pub struct ContextLifecycleState {
51 #[serde(default)]
53 pub entered: bool,
54 #[serde(default)]
56 pub stream_queue_enabled: bool,
57 #[serde(default)]
59 pub compact_depth: u32,
60}
61
62impl ContextLifecycleState {
63 #[must_use]
65 pub fn is_default(&self) -> bool {
66 self == &Self::default()
67 }
68}
69
70#[derive(Clone, Debug, Deserialize, Eq, PartialEq, Serialize)]
72pub struct ToolIdWrapper {
73 #[serde(default = "default_tool_id_prefix")]
75 pub prefix: String,
76 #[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 pub fn clear(&mut self) {
93 self.tool_call_maps.clear();
94 }
95
96 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 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 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 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 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 #[must_use]
185 pub fn is_empty(&self) -> bool {
186 self.tool_call_maps.is_empty()
187 }
188}
189
190#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
192pub struct AgentStreamQueueRegistry {
193 #[serde(default, skip_serializing_if = "Vec::is_empty")]
195 pub queues: Vec<String>,
196}
197
198impl AgentStreamQueueRegistry {
199 #[must_use]
201 pub const fn is_empty(&self) -> bool {
202 self.queues.is_empty()
203 }
204}
205
206#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
208pub struct ToolSearchState {
209 #[serde(default, skip_serializing_if = "Vec::is_empty")]
211 pub loaded_tools: Vec<String>,
212 #[serde(default, skip_serializing_if = "Vec::is_empty")]
214 pub loaded_namespaces: Vec<String>,
215}
216
217impl ToolSearchState {
218 #[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#[derive(Clone, Debug, Default, Deserialize, Eq, PartialEq, Serialize)]
227pub struct ToolSearchInvalidation {
228 #[serde(default, skip_serializing_if = "Vec::is_empty")]
230 pub removed_tools: Vec<String>,
231 #[serde(default, skip_serializing_if = "Vec::is_empty")]
233 pub removed_namespaces: Vec<String>,
234}
235
236impl ToolSearchInvalidation {
237 #[must_use]
239 pub const fn is_empty(&self) -> bool {
240 self.removed_tools.is_empty() && self.removed_namespaces.is_empty()
241 }
242}
243
244pub type ModelWrapperMetadata = Value;
246
247fn default_tool_id_prefix() -> String {
248 "sw-tool-".to_string()
249}