1use std::cmp::Ordering;
2use std::collections::{BTreeSet, HashSet};
3use std::fmt;
4use std::panic::{catch_unwind, AssertUnwindSafe};
5use std::sync::Arc;
6
7use serde::{Deserialize, Serialize};
8use serde_json::{json, Value};
9
10use crate::types::{CompletionReason, CycleRecord, Message, Metadata, TaskTokenUsage};
11
12pub const AFTER_CYCLE_CONTROL_STATE_KEY: &str = "_vv_agent_after_cycle_control";
13pub const AFTER_CYCLE_CONTROL_SCHEMA: &str = "vv-agent.after-cycle-control.v1";
14pub const MAX_STEERING_MESSAGES: usize = 32;
15pub const MAX_STEERING_MESSAGE_UTF8_BYTES: usize = 16_384;
16pub const MAX_TOTAL_STEERING_UTF8_BYTES: usize = 65_536;
17pub const MAX_DISALLOW_TOOLS: usize = 1_024;
18pub const MAX_TOOL_NAME_UTF8_BYTES: usize = 256;
19pub const MAX_STOP_CODE_ASCII_BYTES: usize = 128;
20pub const MAX_STOP_MESSAGE_UTF8_BYTES: usize = 4_096;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
23#[serde(rename_all = "snake_case")]
24pub enum AfterCycleAction {
25 Continue,
26 Steer,
27 StopNonSuccess,
28}
29
30#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
31#[serde(rename_all = "snake_case")]
32pub enum NativeCycleOutcomeKind {
33 Continue,
34 Completed,
35 WaitUser,
36 MaxCycles,
37}
38
39#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
40pub struct NativeCycleOutcome {
41 pub kind: NativeCycleOutcomeKind,
42 pub completion_reason: Option<CompletionReason>,
43 pub completion_tool_name: Option<String>,
44 pub steer_allowed: bool,
45}
46
47impl NativeCycleOutcome {
48 pub fn continuing() -> Self {
49 Self {
50 kind: NativeCycleOutcomeKind::Continue,
51 completion_reason: None,
52 completion_tool_name: None,
53 steer_allowed: true,
54 }
55 }
56}
57
58#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
59pub struct AfterCycleSnapshot {
60 pub task_id: String,
61 pub cycle_index: u32,
62 pub max_cycles: u32,
63 pub remaining_cycles: u32,
64 pub cycle: CycleRecord,
65 pub messages: Vec<Message>,
66 pub shared_state: Metadata,
67 pub cumulative_token_usage: TaskTokenUsage,
68 pub available_tool_names: Vec<String>,
69 pub disallowed_tool_names: Vec<String>,
70 pub native_outcome: NativeCycleOutcome,
71}
72
73impl AfterCycleSnapshot {
74 #[allow(clippy::too_many_arguments)]
75 pub fn capture(
76 task_id: impl Into<String>,
77 cycle_index: u32,
78 max_cycles: u32,
79 cycle: &CycleRecord,
80 messages: &[Message],
81 shared_state: &Metadata,
82 cumulative_token_usage: TaskTokenUsage,
83 available_tool_names: Vec<String>,
84 disallowed_tool_names: Vec<String>,
85 native_outcome: NativeCycleOutcome,
86 ) -> Self {
87 Self {
88 task_id: task_id.into(),
89 cycle_index,
90 max_cycles,
91 remaining_cycles: max_cycles.saturating_sub(cycle_index),
92 cycle: cycle.clone(),
93 messages: messages.to_vec(),
94 shared_state: shared_state.clone(),
95 cumulative_token_usage,
96 available_tool_names,
97 disallowed_tool_names,
98 native_outcome,
99 }
100 }
101}
102
103#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
104#[serde(deny_unknown_fields)]
105pub struct AfterCycleStop {
106 pub code: String,
107 pub message: String,
108}
109
110#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
111#[serde(deny_unknown_fields)]
112pub struct AfterCycleDecision {
113 pub action: AfterCycleAction,
114 pub steering_messages: Vec<String>,
115 pub disallow_tools: Vec<String>,
116 pub stop: Option<AfterCycleStop>,
117}
118
119impl Default for AfterCycleDecision {
120 fn default() -> Self {
121 Self::continue_run()
122 }
123}
124
125impl AfterCycleDecision {
126 pub fn continue_run() -> Self {
127 Self {
128 action: AfterCycleAction::Continue,
129 steering_messages: Vec::new(),
130 disallow_tools: Vec::new(),
131 stop: None,
132 }
133 }
134
135 pub fn continue_with_disallowed_tools(
136 tools: impl IntoIterator<Item = impl Into<String>>,
137 ) -> Result<Self, AfterCycleDecisionError> {
138 let decision = Self {
139 disallow_tools: tools.into_iter().map(Into::into).collect(),
140 ..Self::continue_run()
141 };
142 decision.validate()?;
143 Ok(decision)
144 }
145
146 pub fn steer(
147 messages: impl IntoIterator<Item = impl Into<String>>,
148 ) -> Result<Self, AfterCycleDecisionError> {
149 Self::steer_with_disallowed_tools(messages, Vec::<String>::new())
150 }
151
152 pub fn steer_with_disallowed_tools(
153 messages: impl IntoIterator<Item = impl Into<String>>,
154 tools: impl IntoIterator<Item = impl Into<String>>,
155 ) -> Result<Self, AfterCycleDecisionError> {
156 let decision = Self {
157 action: AfterCycleAction::Steer,
158 steering_messages: messages.into_iter().map(Into::into).collect(),
159 disallow_tools: tools.into_iter().map(Into::into).collect(),
160 stop: None,
161 };
162 decision.validate()?;
163 Ok(decision)
164 }
165
166 pub fn stop_non_success(
167 code: impl Into<String>,
168 message: impl Into<String>,
169 ) -> Result<Self, AfterCycleDecisionError> {
170 let decision = Self {
171 action: AfterCycleAction::StopNonSuccess,
172 steering_messages: Vec::new(),
173 disallow_tools: Vec::new(),
174 stop: Some(AfterCycleStop {
175 code: code.into(),
176 message: message.into(),
177 }),
178 };
179 decision.validate()?;
180 Ok(decision)
181 }
182
183 pub fn validate(&self) -> Result<(), AfterCycleDecisionError> {
184 if self.steering_messages.len() > MAX_STEERING_MESSAGES {
185 return Err(AfterCycleDecisionError::new(
186 "after-cycle steering message count exceeds the limit",
187 ));
188 }
189 let mut total_bytes = 0_usize;
190 for message in &self.steering_messages {
191 validate_bounded_text(
192 message,
193 "after-cycle steering message",
194 MAX_STEERING_MESSAGE_UTF8_BYTES,
195 )?;
196 total_bytes = total_bytes.saturating_add(message.len());
197 }
198 if total_bytes > MAX_TOTAL_STEERING_UTF8_BYTES {
199 return Err(AfterCycleDecisionError::new(
200 "after-cycle steering messages exceed the total byte limit",
201 ));
202 }
203 if self.disallow_tools.len() > MAX_DISALLOW_TOOLS {
204 return Err(AfterCycleDecisionError::new(
205 "after-cycle disallowed tool count exceeds the limit",
206 ));
207 }
208 let mut seen = HashSet::new();
209 for tool_name in &self.disallow_tools {
210 validate_bounded_text(
211 tool_name,
212 "after-cycle disallowed tool name",
213 MAX_TOOL_NAME_UTF8_BYTES,
214 )?;
215 if !seen.insert(tool_name) {
216 return Err(AfterCycleDecisionError::new(
217 "after-cycle disallowed tools must be unique",
218 ));
219 }
220 }
221 match self.action {
222 AfterCycleAction::Continue => {
223 if !self.steering_messages.is_empty() || self.stop.is_some() {
224 return Err(AfterCycleDecisionError::new(
225 "continue cannot include steering messages or a stop payload",
226 ));
227 }
228 }
229 AfterCycleAction::Steer => {
230 if self.steering_messages.is_empty() || self.stop.is_some() {
231 return Err(AfterCycleDecisionError::new(
232 "steer requires messages and cannot include a stop payload",
233 ));
234 }
235 }
236 AfterCycleAction::StopNonSuccess => {
237 if !self.steering_messages.is_empty()
238 || !self.disallow_tools.is_empty()
239 || self.stop.is_none()
240 {
241 return Err(AfterCycleDecisionError::new(
242 "stop_non_success requires only a typed stop payload",
243 ));
244 }
245 let stop = self.stop.as_ref().expect("checked stop payload");
246 validate_stop_code(&stop.code)?;
247 validate_bounded_text(
248 &stop.message,
249 "after-cycle stop message",
250 MAX_STOP_MESSAGE_UTF8_BYTES,
251 )?;
252 }
253 }
254 Ok(())
255 }
256}
257
258#[derive(Debug, Clone, PartialEq, Eq)]
259pub struct AfterCycleDecisionError {
260 message: String,
261}
262
263impl AfterCycleDecisionError {
264 fn new(message: impl Into<String>) -> Self {
265 Self {
266 message: message.into(),
267 }
268 }
269}
270
271impl fmt::Display for AfterCycleDecisionError {
272 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
273 formatter.write_str(&self.message)
274 }
275}
276
277impl std::error::Error for AfterCycleDecisionError {}
278
279pub trait AfterCycleHook: Send + Sync {
280 fn after_cycle(
281 &self,
282 snapshot: &AfterCycleSnapshot,
283 ) -> Result<Option<AfterCycleDecision>, String>;
284}
285
286impl<F> AfterCycleHook for F
287where
288 F: Fn(&AfterCycleSnapshot) -> Result<Option<AfterCycleDecision>, String> + Send + Sync,
289{
290 fn after_cycle(
291 &self,
292 snapshot: &AfterCycleSnapshot,
293 ) -> Result<Option<AfterCycleDecision>, String> {
294 self(snapshot)
295 }
296}
297
298#[derive(Debug, Clone, PartialEq, Eq)]
299pub struct AfterCycleHookError {
300 pub code: &'static str,
301 message: String,
302}
303
304impl AfterCycleHookError {
305 fn new(code: &'static str, message: impl Into<String>) -> Self {
306 Self {
307 code,
308 message: message.into(),
309 }
310 }
311}
312
313impl fmt::Display for AfterCycleHookError {
314 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
315 formatter.write_str(&self.message)
316 }
317}
318
319impl std::error::Error for AfterCycleHookError {}
320
321#[derive(Clone, Default)]
322pub struct AfterCycleHookManager {
323 hooks: Vec<Arc<dyn AfterCycleHook>>,
324}
325
326impl AfterCycleHookManager {
327 pub fn new(hooks: Vec<Arc<dyn AfterCycleHook>>) -> Self {
328 Self { hooks }
329 }
330
331 pub fn has_hooks(&self) -> bool {
332 !self.hooks.is_empty()
333 }
334
335 pub fn apply(
336 &self,
337 snapshot: &AfterCycleSnapshot,
338 ) -> Result<AfterCycleDecision, AfterCycleHookError> {
339 let mut steering_messages = Vec::new();
340 let mut disallow_tools = Vec::new();
341 let mut seen_tools = HashSet::new();
342 for hook in &self.hooks {
343 let outcome = catch_unwind(AssertUnwindSafe(|| hook.after_cycle(snapshot)))
344 .map_err(|_| {
345 AfterCycleHookError::new("after_cycle_hook_failed", "after-cycle hook panicked")
346 })?
347 .map_err(|error| {
348 AfterCycleHookError::new(
349 "after_cycle_hook_failed",
350 format!("after-cycle hook failed: {error}"),
351 )
352 })?;
353 let Some(decision) = outcome else {
354 continue;
355 };
356 decision.validate().map_err(|error| {
357 AfterCycleHookError::new(
358 "after_cycle_decision_invalid",
359 format!("after-cycle hook returned an invalid decision: {error}"),
360 )
361 })?;
362 for tool_name in &decision.disallow_tools {
363 if seen_tools.insert(tool_name.clone()) {
364 disallow_tools.push(tool_name.clone());
365 }
366 }
367 if decision.action == AfterCycleAction::StopNonSuccess {
368 return Ok(decision);
369 }
370 if decision.action == AfterCycleAction::Steer {
371 steering_messages.extend(decision.steering_messages);
372 }
373 }
374 let composed = if steering_messages.is_empty() {
375 AfterCycleDecision::continue_with_disallowed_tools(disallow_tools)
376 } else {
377 AfterCycleDecision::steer_with_disallowed_tools(steering_messages, disallow_tools)
378 };
379 composed.map_err(|error| {
380 AfterCycleHookError::new(
381 "after_cycle_decision_invalid",
382 format!("composed after-cycle decision is invalid: {error}"),
383 )
384 })
385 }
386}
387
388pub fn read_after_cycle_disallowed_tools(
389 shared_state: &Metadata,
390) -> Result<Vec<String>, AfterCycleHookError> {
391 let Some(raw) = shared_state.get(AFTER_CYCLE_CONTROL_STATE_KEY) else {
392 return Ok(Vec::new());
393 };
394 let object = raw
395 .as_object()
396 .ok_or_else(|| control_state_error("after-cycle control state must be an object"))?;
397 let expected = BTreeSet::from(["schema_version", "disallowed_tools"]);
398 let actual = object.keys().map(String::as_str).collect::<BTreeSet<_>>();
399 if actual != expected {
400 return Err(control_state_error(
401 "after-cycle control state has missing or unknown fields",
402 ));
403 }
404 if object.get("schema_version").and_then(Value::as_str) != Some(AFTER_CYCLE_CONTROL_SCHEMA) {
405 return Err(control_state_error(
406 "after-cycle control state schema is unsupported",
407 ));
408 }
409 let values = object
410 .get("disallowed_tools")
411 .and_then(Value::as_array)
412 .ok_or_else(|| {
413 control_state_error("after-cycle control disallowed_tools must be an array")
414 })?
415 .iter()
416 .map(|value| {
417 value.as_str().map(str::to_string).ok_or_else(|| {
418 control_state_error("after-cycle control disallowed_tools must contain strings")
419 })
420 })
421 .collect::<Result<Vec<_>, _>>()?;
422 AfterCycleDecision::continue_with_disallowed_tools(values.clone())
423 .map_err(|error| control_state_error(error.to_string()))?;
424 let mut ordered = values.clone();
425 ordered.sort_by(|left, right| utf16_cmp(left, right));
426 if values != ordered {
427 return Err(control_state_error(
428 "after-cycle control disallowed_tools must be sorted and unique",
429 ));
430 }
431 Ok(ordered)
432}
433
434pub fn persist_after_cycle_disallowed_tools(
435 shared_state: &mut Metadata,
436 additional_tools: &[String],
437) -> Result<Vec<String>, AfterCycleHookError> {
438 let mut values = read_after_cycle_disallowed_tools(shared_state)?;
439 values.extend(additional_tools.iter().cloned());
440 values.sort_by(|left, right| utf16_cmp(left, right));
441 values.dedup();
442 if values.is_empty() {
443 return Ok(values);
444 }
445 AfterCycleDecision::continue_with_disallowed_tools(values.clone())
446 .map_err(|error| control_state_error(error.to_string()))?;
447 shared_state.insert(
448 AFTER_CYCLE_CONTROL_STATE_KEY.to_string(),
449 json!({
450 "schema_version": AFTER_CYCLE_CONTROL_SCHEMA,
451 "disallowed_tools": values,
452 }),
453 );
454 Ok(values)
455}
456
457pub(crate) fn utf16_cmp(left: &str, right: &str) -> Ordering {
458 left.encode_utf16().cmp(right.encode_utf16())
459}
460
461fn validate_bounded_text(
462 value: &str,
463 field_name: &str,
464 max_bytes: usize,
465) -> Result<(), AfterCycleDecisionError> {
466 if value.trim().is_empty() {
467 return Err(AfterCycleDecisionError::new(format!(
468 "{field_name} must be a non-empty string"
469 )));
470 }
471 if value.len() > max_bytes {
472 return Err(AfterCycleDecisionError::new(format!(
473 "{field_name} exceeds {max_bytes} UTF-8 bytes"
474 )));
475 }
476 Ok(())
477}
478
479fn validate_stop_code(code: &str) -> Result<(), AfterCycleDecisionError> {
480 let valid = code.is_ascii()
481 && !code.is_empty()
482 && code.len() <= MAX_STOP_CODE_ASCII_BYTES
483 && code.as_bytes()[0].is_ascii_lowercase()
484 && code.bytes().all(|byte| {
485 byte.is_ascii_lowercase() || byte.is_ascii_digit() || b"_.-".contains(&byte)
486 });
487 if !valid {
488 return Err(AfterCycleDecisionError::new(
489 "after-cycle stop code is invalid",
490 ));
491 }
492 Ok(())
493}
494
495fn control_state_error(message: impl Into<String>) -> AfterCycleHookError {
496 AfterCycleHookError::new("after_cycle_control_state_invalid", message)
497}