1use std::collections::{BTreeMap, BTreeSet};
2use std::sync::{Arc, LazyLock, OnceLock};
3
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6
7use crate::tool_annotations::{SideEffectLevel, ToolAnnotations, ToolKind};
8use crate::value::{VmDictExt, VmError, VmValue};
9
10use super::manifest::{BindingManifestEntry, BindingPolicyStatus};
11
12pub const COMPOSITION_STATE_SCHEMA_VERSION: u32 = 1;
13pub const COMPOSITION_STATE_CAPABILITY: &str = "composition_state";
14
15const STATE_BINDING_NAME: &str = "state";
16const DEFAULT_MAX_VALUE_BYTES: u64 = 16 * 1024;
17const DEFAULT_MAX_TOTAL_BYTES: u64 = 64 * 1024;
18const DEFAULT_MAX_KEYS: u64 = 64;
19
20#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
21#[serde(default, deny_unknown_fields)]
22pub struct CompositionStateBinding {
23 pub schema_version: u32,
24 pub max_value_bytes: u64,
25 pub max_total_bytes: u64,
26 pub max_keys: u64,
27}
28
29impl Default for CompositionStateBinding {
30 fn default() -> Self {
31 Self {
32 schema_version: COMPOSITION_STATE_SCHEMA_VERSION,
33 max_value_bytes: DEFAULT_MAX_VALUE_BYTES,
34 max_total_bytes: DEFAULT_MAX_TOTAL_BYTES,
35 max_keys: DEFAULT_MAX_KEYS,
36 }
37 }
38}
39
40impl CompositionStateBinding {
41 pub fn validate(&self) -> Result<(), CompositionStateError> {
42 if self.schema_version != COMPOSITION_STATE_SCHEMA_VERSION {
43 return Err(CompositionStateError::invalid_limits(format!(
44 "unsupported state schema_version={} (expected {})",
45 self.schema_version, COMPOSITION_STATE_SCHEMA_VERSION
46 )));
47 }
48 if self.max_value_bytes == 0 {
49 return Err(CompositionStateError::invalid_limits(
50 "max_value_bytes must be greater than zero",
51 ));
52 }
53 if self.max_total_bytes == 0 {
54 return Err(CompositionStateError::invalid_limits(
55 "max_total_bytes must be greater than zero",
56 ));
57 }
58 if self.max_keys == 0 {
59 return Err(CompositionStateError::invalid_limits(
60 "max_keys must be greater than zero",
61 ));
62 }
63 if self.max_value_bytes > self.max_total_bytes {
64 return Err(CompositionStateError::invalid_limits(
65 "max_value_bytes cannot exceed max_total_bytes",
66 ));
67 }
68 Ok(())
69 }
70}
71
72#[derive(Clone, Copy, Debug, Eq, PartialEq, Serialize, Deserialize)]
73#[serde(rename_all = "snake_case")]
74pub enum CompositionStateErrorCode {
75 InvalidLimits,
76 SessionRequired,
77 InvalidKey,
78 NonJsonValue,
79 ValueTooLarge,
80 TotalTooLarge,
81 TooManyKeys,
82}
83
84impl CompositionStateErrorCode {
85 pub fn as_str(self) -> &'static str {
86 match self {
87 Self::InvalidLimits => "invalid_limits",
88 Self::SessionRequired => "session_required",
89 Self::InvalidKey => "invalid_key",
90 Self::NonJsonValue => "non_json_value",
91 Self::ValueTooLarge => "value_too_large",
92 Self::TotalTooLarge => "total_too_large",
93 Self::TooManyKeys => "too_many_keys",
94 }
95 }
96}
97
98#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
99pub struct CompositionStateError {
100 pub code: CompositionStateErrorCode,
101 pub operation: String,
102 pub message: String,
103 #[serde(skip_serializing_if = "Option::is_none")]
104 pub key: Option<String>,
105 #[serde(skip_serializing_if = "Option::is_none")]
106 pub limit: Option<u64>,
107 #[serde(skip_serializing_if = "Option::is_none")]
108 pub actual: Option<u64>,
109}
110
111impl CompositionStateError {
112 pub fn invalid_limits(message: impl Into<String>) -> Self {
113 Self {
114 code: CompositionStateErrorCode::InvalidLimits,
115 operation: "configure".to_string(),
116 message: message.into(),
117 key: None,
118 limit: None,
119 actual: None,
120 }
121 }
122
123 pub fn session_required(operation: &str) -> Self {
124 Self {
125 code: CompositionStateErrorCode::SessionRequired,
126 operation: operation.to_string(),
127 message: "state operations require a non-empty composition session_id".to_string(),
128 key: None,
129 limit: None,
130 actual: None,
131 }
132 }
133
134 pub fn invalid_key(operation: &str, message: impl Into<String>) -> Self {
135 Self {
136 code: CompositionStateErrorCode::InvalidKey,
137 operation: operation.to_string(),
138 message: message.into(),
139 key: None,
140 limit: None,
141 actual: None,
142 }
143 }
144
145 pub fn non_json(operation: &str, key: Option<&str>, message: impl Into<String>) -> Self {
146 Self {
147 code: CompositionStateErrorCode::NonJsonValue,
148 operation: operation.to_string(),
149 message: message.into(),
150 key: key.map(ToOwned::to_owned),
151 limit: None,
152 actual: None,
153 }
154 }
155
156 fn for_key(
157 code: CompositionStateErrorCode,
158 operation: &str,
159 key: &str,
160 message: impl Into<String>,
161 ) -> Self {
162 Self {
163 code,
164 operation: operation.to_string(),
165 message: message.into(),
166 key: Some(key.to_string()),
167 limit: None,
168 actual: None,
169 }
170 }
171
172 fn with_bound(mut self, limit: u64, actual: u64) -> Self {
173 self.limit = Some(limit);
174 self.actual = Some(actual);
175 self
176 }
177
178 pub fn to_value(&self) -> Value {
179 serde_json::to_value(self).unwrap_or_else(|_| {
180 serde_json::json!({
181 "code": self.code.as_str(),
182 "operation": self.operation,
183 "message": self.message,
184 })
185 })
186 }
187
188 pub fn into_vm_error(self) -> VmError {
189 let mut fields = crate::value::DictMap::new();
190 fields.put_str("type", "composition_state_error");
191 fields.put_str("code", self.code.as_str());
192 fields.put_str("operation", &self.operation);
193 fields.put_str("message", &self.message);
194 if let Some(key) = self.key {
195 fields.put_str("key", key);
196 }
197 if let Some(limit) = self.limit {
198 fields.put(
199 "limit",
200 VmValue::Int(i64::try_from(limit).unwrap_or(i64::MAX)),
201 );
202 }
203 if let Some(actual) = self.actual {
204 fields.put(
205 "actual",
206 VmValue::Int(i64::try_from(actual).unwrap_or(i64::MAX)),
207 );
208 }
209 VmError::Thrown(VmValue::dict(fields))
210 }
211}
212
213#[derive(Clone, Debug, Eq, PartialEq, Ord, PartialOrd)]
214pub struct CompositionStateScope {
215 session_id: String,
216 tool_window_id: String,
217}
218
219impl CompositionStateScope {
220 pub fn new(session_id: String, tool_window_id: String) -> Self {
221 Self {
222 session_id,
223 tool_window_id,
224 }
225 }
226
227 pub fn tool_window_id(&self) -> &str {
228 &self.tool_window_id
229 }
230}
231
232#[derive(Default)]
233struct StateWindow {
234 values: BTreeMap<String, Value>,
235 total_bytes: u64,
236}
237
238static STATE_WINDOWS: LazyLock<parking_lot::Mutex<BTreeMap<CompositionStateScope, StateWindow>>> =
239 LazyLock::new(|| parking_lot::Mutex::new(BTreeMap::new()));
240
241fn ensure_cleanup_hook() {
242 static REGISTRATION: OnceLock<crate::llm::SessionCloseHookRegistration> = OnceLock::new();
243 REGISTRATION.get_or_init(|| {
244 crate::llm::register_session_close_hook(Arc::new(|session_id| {
245 STATE_WINDOWS
246 .lock()
247 .retain(|scope, _| scope.session_id != session_id);
248 }))
249 });
250}
251
252pub fn execute(
253 scope: &CompositionStateScope,
254 limits: &CompositionStateBinding,
255 operation: &str,
256 key: Option<&str>,
257 value: Option<Value>,
258) -> Result<Value, CompositionStateError> {
259 ensure_cleanup_hook();
260 limits.validate()?;
261 match operation {
262 "get" => {
263 let key = validate_key(operation, key)?;
264 Ok(STATE_WINDOWS
265 .lock()
266 .get(scope)
267 .and_then(|window| window.values.get(key))
268 .cloned()
269 .unwrap_or(Value::Null))
270 }
271 "list" => Ok(Value::Array(
272 STATE_WINDOWS
273 .lock()
274 .get(scope)
275 .map(|window| window.values.keys().cloned().map(Value::String).collect())
276 .unwrap_or_default(),
277 )),
278 "put" => {
279 let key = validate_key(operation, key)?;
280 let value = value.ok_or_else(|| {
281 CompositionStateError::for_key(
282 CompositionStateErrorCode::NonJsonValue,
283 operation,
284 key,
285 "state.put requires a JSON value",
286 )
287 })?;
288 put(scope, limits, key, value)?;
289 Ok(Value::Null)
290 }
291 "delete" => {
292 let key = validate_key(operation, key)?;
293 Ok(Value::Bool(delete(scope, key)))
294 }
295 _ => Err(CompositionStateError::invalid_key(
296 operation,
297 format!("unknown state operation '{operation}'"),
298 )),
299 }
300}
301
302fn validate_key<'a>(
303 operation: &str,
304 key: Option<&'a str>,
305) -> Result<&'a str, CompositionStateError> {
306 let key = key.unwrap_or_default();
307 if key.is_empty() {
308 return Err(CompositionStateError::for_key(
309 CompositionStateErrorCode::InvalidKey,
310 operation,
311 key,
312 "state keys must be non-empty strings",
313 ));
314 }
315 Ok(key)
316}
317
318fn put(
319 scope: &CompositionStateScope,
320 limits: &CompositionStateBinding,
321 key: &str,
322 value: Value,
323) -> Result<(), CompositionStateError> {
324 let value_bytes = u64::try_from(
325 serde_json::to_vec(&value)
326 .map_err(|error| {
327 CompositionStateError::for_key(
328 CompositionStateErrorCode::NonJsonValue,
329 "put",
330 key,
331 format!("state value is not valid JSON: {error}"),
332 )
333 })?
334 .len(),
335 )
336 .unwrap_or(u64::MAX);
337 if value_bytes > limits.max_value_bytes {
338 return Err(CompositionStateError::for_key(
339 CompositionStateErrorCode::ValueTooLarge,
340 "put",
341 key,
342 format!(
343 "state value exceeds max_value_bytes={}",
344 limits.max_value_bytes
345 ),
346 )
347 .with_bound(limits.max_value_bytes, value_bytes));
348 }
349
350 let key_bytes = u64::try_from(key.len()).unwrap_or(u64::MAX);
351 let mut windows = STATE_WINDOWS.lock();
352 let window = windows.get(scope);
353 let old_value = window.and_then(|window| window.values.get(key));
354 let old_bytes = old_value
355 .and_then(|old| serde_json::to_vec(old).ok())
356 .and_then(|old| u64::try_from(old.len()).ok())
357 .map(|old| old.saturating_add(key_bytes))
358 .unwrap_or(0);
359 let key_count = window
360 .map(|window| u64::try_from(window.values.len()).unwrap_or(u64::MAX))
361 .unwrap_or(0);
362 if old_value.is_none() && key_count >= limits.max_keys {
363 return Err(CompositionStateError::for_key(
364 CompositionStateErrorCode::TooManyKeys,
365 "put",
366 key,
367 format!("state store exceeds max_keys={}", limits.max_keys),
368 )
369 .with_bound(limits.max_keys, key_count.saturating_add(1)));
370 }
371 let new_total = window
372 .map(|window| window.total_bytes)
373 .unwrap_or(0)
374 .saturating_sub(old_bytes)
375 .saturating_add(key_bytes)
376 .saturating_add(value_bytes);
377 if new_total > limits.max_total_bytes {
378 return Err(CompositionStateError::for_key(
379 CompositionStateErrorCode::TotalTooLarge,
380 "put",
381 key,
382 format!(
383 "state store exceeds max_total_bytes={}",
384 limits.max_total_bytes
385 ),
386 )
387 .with_bound(limits.max_total_bytes, new_total));
388 }
389 let window = windows.entry(scope.clone()).or_default();
390 window.values.insert(key.to_string(), value);
391 window.total_bytes = new_total;
392 Ok(())
393}
394
395fn delete(scope: &CompositionStateScope, key: &str) -> bool {
396 let mut windows = STATE_WINDOWS.lock();
397 let Some(window) = windows.get_mut(scope) else {
398 return false;
399 };
400 let Some(value) = window.values.remove(key) else {
401 return false;
402 };
403 let value_bytes = serde_json::to_vec(&value)
404 .ok()
405 .and_then(|bytes| u64::try_from(bytes.len()).ok())
406 .unwrap_or(0);
407 window.total_bytes = window
408 .total_bytes
409 .saturating_sub(u64::try_from(key.len()).unwrap_or(u64::MAX))
410 .saturating_sub(value_bytes);
411 if window.values.is_empty() {
412 windows.remove(scope);
413 }
414 true
415}
416
417pub fn binding_entry(operation: &str) -> BindingManifestEntry {
418 let writes = matches!(operation, "put" | "delete");
419 let capability = if writes { "write" } else { "read" };
420 let mut capabilities = BTreeMap::new();
421 capabilities.insert(
422 COMPOSITION_STATE_CAPABILITY.to_string(),
423 vec![capability.to_string()],
424 );
425 let annotations = ToolAnnotations {
426 kind: if writes {
427 ToolKind::Edit
428 } else {
429 ToolKind::Read
430 },
431 side_effect_level: if writes {
432 SideEffectLevel::WorkspaceWrite
433 } else {
434 SideEffectLevel::ReadOnly
435 },
436 capabilities: capabilities.clone(),
437 inline_result: true,
438 ..ToolAnnotations::default()
439 };
440 BindingManifestEntry {
441 name: format!("{STATE_BINDING_NAME}.{operation}"),
442 binding: format!("{STATE_BINDING_NAME}.{operation}"),
443 namespace: Some(STATE_BINDING_NAME.to_string()),
444 description: Some(format!("Session-scoped composition state {operation}")),
445 input_schema: state_input_schema(operation),
446 output_schema: Some(state_output_schema(operation)),
447 annotations,
448 side_effect_level: if writes {
449 SideEffectLevel::WorkspaceWrite
450 } else {
451 SideEffectLevel::ReadOnly
452 },
453 capabilities,
454 source: COMPOSITION_STATE_CAPABILITY.to_string(),
455 policy: BindingPolicyStatus::default(),
456 metadata: serde_json::json!({
457 "internal": true,
458 "state_operation": operation,
459 }),
460 ..BindingManifestEntry::default()
461 }
462}
463
464fn state_input_schema(operation: &str) -> Value {
465 match operation {
466 "put" => serde_json::json!({
467 "type": "object",
468 "required": ["key", "value"],
469 "properties": {
470 "key": {"type": "string"},
471 "value": {},
472 },
473 }),
474 "get" | "delete" => serde_json::json!({
475 "type": "object",
476 "required": ["key"],
477 "properties": {"key": {"type": "string"}},
478 }),
479 _ => serde_json::json!({"type": "object", "properties": {}}),
480 }
481}
482
483fn state_output_schema(operation: &str) -> Value {
484 match operation {
485 "list" => serde_json::json!({"type": "array", "items": {"type": "string"}}),
486 "delete" => serde_json::json!({"type": "boolean"}),
487 _ => serde_json::json!({}),
488 }
489}
490
491pub fn operation_names() -> BTreeSet<String> {
492 ["get", "put", "list", "delete"]
493 .into_iter()
494 .map(|operation| format!("{STATE_BINDING_NAME}.{operation}"))
495 .collect()
496}
497
498pub fn harn_runtime_source() -> &'static str {
499 "const state = {\n\
500 _namespace: \"composition_state\",\n\
501 get: { key -> __composition_state(\"get\", key) },\n\
502 put: { key, value -> __composition_state(\"put\", key, value) },\n\
503 list: { -> __composition_state(\"list\") },\n\
504 delete: { key -> __composition_state(\"delete\", key) },\n\
505 }\n"
506}
507
508pub fn harn_api_source() -> &'static str {
509 "type CompositionState = {\n\
510 _namespace: string,\n\
511 get: fn(string) -> JsonValue,\n\
512 put: fn(string, JsonValue) -> nil,\n\
513 list: fn() -> list<string>,\n\
514 delete: fn(string) -> bool,\n\
515 }\n\
516 const state: CompositionState = {\n\
517 _namespace: \"composition_state\",\n\
518 get: { key -> __composition_state(\"get\", key) },\n\
519 put: { key, value -> __composition_state(\"put\", key, value) },\n\
520 list: { -> __composition_state(\"list\") },\n\
521 delete: { key -> __composition_state(\"delete\", key) },\n\
522 }\n\n"
523}
524
525pub fn typescript_api_source() -> &'static str {
526 "export interface CompositionState {\n\
527 get(key: string): Promise<JsonValue>;\n\
528 put(key: string, value: JsonValue): Promise<void>;\n\
529 list(): Promise<string[]>;\n\
530 delete(key: string): Promise<boolean>;\n\
531 }\n\
532 export declare const state: CompositionState;\n\n"
533}