conversation.rs 46 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401
  1. use std::collections::BTreeMap;
  2. use std::fmt::{Display, Formatter};
  3. use serde_json::{Map, Value};
  4. use telemetry::SessionTracer;
  5. use crate::compact::{
  6. compact_session, estimate_session_tokens, CompactionConfig, CompactionResult,
  7. };
  8. use crate::config::RuntimeFeatureConfig;
  9. use crate::hooks::{HookAbortSignal, HookProgressReporter, HookRunResult, HookRunner};
  10. use crate::permissions::{
  11. PermissionContext, PermissionOutcome, PermissionPolicy, PermissionPrompter,
  12. };
  13. use crate::session::{ContentBlock, ConversationMessage, Session};
  14. use crate::usage::{TokenUsage, UsageTracker};
  15. const DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD: u32 = 100_000;
  16. const AUTO_COMPACTION_THRESHOLD_ENV_VAR: &str = "CLAUDE_CODE_AUTO_COMPACT_INPUT_TOKENS";
  17. #[derive(Debug, Clone, PartialEq, Eq)]
  18. pub struct ApiRequest {
  19. pub system_prompt: Vec<String>,
  20. pub messages: Vec<ConversationMessage>,
  21. }
  22. #[derive(Debug, Clone, PartialEq, Eq)]
  23. pub enum AssistantEvent {
  24. TextDelta(String),
  25. ToolUse {
  26. id: String,
  27. name: String,
  28. input: String,
  29. },
  30. Usage(TokenUsage),
  31. PromptCache(PromptCacheEvent),
  32. MessageStop,
  33. }
  34. #[derive(Debug, Clone, PartialEq, Eq)]
  35. pub struct PromptCacheEvent {
  36. pub unexpected: bool,
  37. pub reason: String,
  38. pub previous_cache_read_input_tokens: u32,
  39. pub current_cache_read_input_tokens: u32,
  40. pub token_drop: u32,
  41. }
  42. pub trait ApiClient {
  43. fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError>;
  44. }
  45. pub trait ToolExecutor {
  46. fn execute(&mut self, tool_name: &str, input: &str) -> Result<String, ToolError>;
  47. }
  48. #[derive(Debug, Clone, PartialEq, Eq)]
  49. pub struct ToolError {
  50. message: String,
  51. }
  52. impl ToolError {
  53. #[must_use]
  54. pub fn new(message: impl Into<String>) -> Self {
  55. Self {
  56. message: message.into(),
  57. }
  58. }
  59. }
  60. impl Display for ToolError {
  61. fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
  62. write!(f, "{}", self.message)
  63. }
  64. }
  65. impl std::error::Error for ToolError {}
  66. #[derive(Debug, Clone, PartialEq, Eq)]
  67. pub struct RuntimeError {
  68. message: String,
  69. }
  70. impl RuntimeError {
  71. #[must_use]
  72. pub fn new(message: impl Into<String>) -> Self {
  73. Self {
  74. message: message.into(),
  75. }
  76. }
  77. }
  78. impl Display for RuntimeError {
  79. fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
  80. write!(f, "{}", self.message)
  81. }
  82. }
  83. impl std::error::Error for RuntimeError {}
  84. #[derive(Debug, Clone, PartialEq, Eq)]
  85. pub struct TurnSummary {
  86. pub assistant_messages: Vec<ConversationMessage>,
  87. pub tool_results: Vec<ConversationMessage>,
  88. pub prompt_cache_events: Vec<PromptCacheEvent>,
  89. pub iterations: usize,
  90. pub usage: TokenUsage,
  91. pub auto_compaction: Option<AutoCompactionEvent>,
  92. }
  93. #[derive(Debug, Clone, Copy, PartialEq, Eq)]
  94. pub struct AutoCompactionEvent {
  95. pub removed_message_count: usize,
  96. }
  97. pub struct ConversationRuntime<C, T> {
  98. session: Session,
  99. api_client: C,
  100. tool_executor: T,
  101. permission_policy: PermissionPolicy,
  102. system_prompt: Vec<String>,
  103. max_iterations: usize,
  104. usage_tracker: UsageTracker,
  105. hook_runner: HookRunner,
  106. auto_compaction_input_tokens_threshold: u32,
  107. hook_abort_signal: HookAbortSignal,
  108. hook_progress_reporter: Option<Box<dyn HookProgressReporter>>,
  109. session_tracer: Option<SessionTracer>,
  110. }
  111. impl<C, T> ConversationRuntime<C, T>
  112. where
  113. C: ApiClient,
  114. T: ToolExecutor,
  115. {
  116. #[must_use]
  117. pub fn new(
  118. session: Session,
  119. api_client: C,
  120. tool_executor: T,
  121. permission_policy: PermissionPolicy,
  122. system_prompt: Vec<String>,
  123. ) -> Self {
  124. Self::new_with_features(
  125. session,
  126. api_client,
  127. tool_executor,
  128. permission_policy,
  129. system_prompt,
  130. &RuntimeFeatureConfig::default(),
  131. )
  132. }
  133. #[must_use]
  134. #[allow(clippy::needless_pass_by_value)]
  135. pub fn new_with_features(
  136. session: Session,
  137. api_client: C,
  138. tool_executor: T,
  139. permission_policy: PermissionPolicy,
  140. system_prompt: Vec<String>,
  141. feature_config: &RuntimeFeatureConfig,
  142. ) -> Self {
  143. let usage_tracker = UsageTracker::from_session(&session);
  144. Self {
  145. session,
  146. api_client,
  147. tool_executor,
  148. permission_policy,
  149. system_prompt,
  150. max_iterations: usize::MAX,
  151. usage_tracker,
  152. hook_runner: HookRunner::from_feature_config(feature_config),
  153. auto_compaction_input_tokens_threshold: auto_compaction_threshold_from_env(),
  154. hook_abort_signal: HookAbortSignal::default(),
  155. hook_progress_reporter: None,
  156. session_tracer: None,
  157. }
  158. }
  159. #[must_use]
  160. pub fn with_max_iterations(mut self, max_iterations: usize) -> Self {
  161. self.max_iterations = max_iterations;
  162. self
  163. }
  164. #[must_use]
  165. pub fn with_auto_compaction_input_tokens_threshold(mut self, threshold: u32) -> Self {
  166. self.auto_compaction_input_tokens_threshold = threshold;
  167. self
  168. }
  169. #[must_use]
  170. pub fn with_hook_abort_signal(mut self, hook_abort_signal: HookAbortSignal) -> Self {
  171. self.hook_abort_signal = hook_abort_signal;
  172. self
  173. }
  174. #[must_use]
  175. pub fn with_hook_progress_reporter(
  176. mut self,
  177. hook_progress_reporter: Box<dyn HookProgressReporter>,
  178. ) -> Self {
  179. self.hook_progress_reporter = Some(hook_progress_reporter);
  180. self
  181. }
  182. #[must_use]
  183. pub fn with_session_tracer(mut self, session_tracer: SessionTracer) -> Self {
  184. self.session_tracer = Some(session_tracer);
  185. self
  186. }
  187. fn run_pre_tool_use_hook(&mut self, tool_name: &str, input: &str) -> HookRunResult {
  188. if let Some(reporter) = self.hook_progress_reporter.as_mut() {
  189. self.hook_runner.run_pre_tool_use_with_context(
  190. tool_name,
  191. input,
  192. Some(&self.hook_abort_signal),
  193. Some(reporter.as_mut()),
  194. )
  195. } else {
  196. self.hook_runner.run_pre_tool_use_with_context(
  197. tool_name,
  198. input,
  199. Some(&self.hook_abort_signal),
  200. None,
  201. )
  202. }
  203. }
  204. fn run_post_tool_use_hook(
  205. &mut self,
  206. tool_name: &str,
  207. input: &str,
  208. output: &str,
  209. is_error: bool,
  210. ) -> HookRunResult {
  211. if let Some(reporter) = self.hook_progress_reporter.as_mut() {
  212. self.hook_runner.run_post_tool_use_with_context(
  213. tool_name,
  214. input,
  215. output,
  216. is_error,
  217. Some(&self.hook_abort_signal),
  218. Some(reporter.as_mut()),
  219. )
  220. } else {
  221. self.hook_runner.run_post_tool_use_with_context(
  222. tool_name,
  223. input,
  224. output,
  225. is_error,
  226. Some(&self.hook_abort_signal),
  227. None,
  228. )
  229. }
  230. }
  231. fn run_post_tool_use_failure_hook(
  232. &mut self,
  233. tool_name: &str,
  234. input: &str,
  235. output: &str,
  236. ) -> HookRunResult {
  237. if let Some(reporter) = self.hook_progress_reporter.as_mut() {
  238. self.hook_runner.run_post_tool_use_failure_with_context(
  239. tool_name,
  240. input,
  241. output,
  242. Some(&self.hook_abort_signal),
  243. Some(reporter.as_mut()),
  244. )
  245. } else {
  246. self.hook_runner.run_post_tool_use_failure_with_context(
  247. tool_name,
  248. input,
  249. output,
  250. Some(&self.hook_abort_signal),
  251. None,
  252. )
  253. }
  254. }
  255. #[allow(clippy::too_many_lines)]
  256. pub fn run_turn(
  257. &mut self,
  258. user_input: impl Into<String>,
  259. mut prompter: Option<&mut dyn PermissionPrompter>,
  260. ) -> Result<TurnSummary, RuntimeError> {
  261. let user_input = user_input.into();
  262. self.record_turn_started(&user_input);
  263. self.session
  264. .push_user_text(user_input)
  265. .map_err(|error| RuntimeError::new(error.to_string()))?;
  266. let mut assistant_messages = Vec::new();
  267. let mut tool_results = Vec::new();
  268. let mut prompt_cache_events = Vec::new();
  269. let mut iterations = 0;
  270. loop {
  271. iterations += 1;
  272. if iterations > self.max_iterations {
  273. let error = RuntimeError::new(
  274. "conversation loop exceeded the maximum number of iterations",
  275. );
  276. self.record_turn_failed(iterations, &error);
  277. return Err(error);
  278. }
  279. let request = ApiRequest {
  280. system_prompt: self.system_prompt.clone(),
  281. messages: self.session.messages.clone(),
  282. };
  283. let events = match self.api_client.stream(request) {
  284. Ok(events) => events,
  285. Err(error) => {
  286. self.record_turn_failed(iterations, &error);
  287. return Err(error);
  288. }
  289. };
  290. let (assistant_message, usage, turn_prompt_cache_events) =
  291. match build_assistant_message(events) {
  292. Ok(result) => result,
  293. Err(error) => {
  294. self.record_turn_failed(iterations, &error);
  295. return Err(error);
  296. }
  297. };
  298. if let Some(usage) = usage {
  299. self.usage_tracker.record(usage);
  300. }
  301. prompt_cache_events.extend(turn_prompt_cache_events);
  302. let pending_tool_uses = assistant_message
  303. .blocks
  304. .iter()
  305. .filter_map(|block| match block {
  306. ContentBlock::ToolUse { id, name, input } => {
  307. Some((id.clone(), name.clone(), input.clone()))
  308. }
  309. _ => None,
  310. })
  311. .collect::<Vec<_>>();
  312. self.record_assistant_iteration(
  313. iterations,
  314. &assistant_message,
  315. pending_tool_uses.len(),
  316. );
  317. self.session
  318. .push_message(assistant_message.clone())
  319. .map_err(|error| RuntimeError::new(error.to_string()))?;
  320. assistant_messages.push(assistant_message);
  321. if pending_tool_uses.is_empty() {
  322. break;
  323. }
  324. for (tool_use_id, tool_name, input) in pending_tool_uses {
  325. let pre_hook_result = self.run_pre_tool_use_hook(&tool_name, &input);
  326. let effective_input = pre_hook_result
  327. .updated_input()
  328. .map_or_else(|| input.clone(), ToOwned::to_owned);
  329. let permission_context = PermissionContext::new(
  330. pre_hook_result.permission_override(),
  331. pre_hook_result.permission_reason().map(ToOwned::to_owned),
  332. );
  333. let permission_outcome = if pre_hook_result.is_cancelled() {
  334. PermissionOutcome::Deny {
  335. reason: format_hook_message(
  336. &pre_hook_result,
  337. &format!("PreToolUse hook cancelled tool `{tool_name}`"),
  338. ),
  339. }
  340. } else if pre_hook_result.is_denied() {
  341. PermissionOutcome::Deny {
  342. reason: format_hook_message(
  343. &pre_hook_result,
  344. &format!("PreToolUse hook denied tool `{tool_name}`"),
  345. ),
  346. }
  347. } else if let Some(prompt) = prompter.as_mut() {
  348. self.permission_policy.authorize_with_context(
  349. &tool_name,
  350. &effective_input,
  351. &permission_context,
  352. Some(*prompt),
  353. )
  354. } else {
  355. self.permission_policy.authorize_with_context(
  356. &tool_name,
  357. &effective_input,
  358. &permission_context,
  359. None,
  360. )
  361. };
  362. let result_message = match permission_outcome {
  363. PermissionOutcome::Allow => {
  364. self.record_tool_started(iterations, &tool_name);
  365. let (mut output, mut is_error) =
  366. match self.tool_executor.execute(&tool_name, &effective_input) {
  367. Ok(output) => (output, false),
  368. Err(error) => (error.to_string(), true),
  369. };
  370. output = merge_hook_feedback(pre_hook_result.messages(), output, false);
  371. let post_hook_result = if is_error {
  372. self.run_post_tool_use_failure_hook(
  373. &tool_name,
  374. &effective_input,
  375. &output,
  376. )
  377. } else {
  378. self.run_post_tool_use_hook(
  379. &tool_name,
  380. &effective_input,
  381. &output,
  382. false,
  383. )
  384. };
  385. if post_hook_result.is_denied() || post_hook_result.is_cancelled() {
  386. is_error = true;
  387. }
  388. output = merge_hook_feedback(
  389. post_hook_result.messages(),
  390. output,
  391. post_hook_result.is_denied() || post_hook_result.is_cancelled(),
  392. );
  393. ConversationMessage::tool_result(tool_use_id, tool_name, output, is_error)
  394. }
  395. PermissionOutcome::Deny { reason } => ConversationMessage::tool_result(
  396. tool_use_id,
  397. tool_name,
  398. merge_hook_feedback(pre_hook_result.messages(), reason, true),
  399. true,
  400. ),
  401. };
  402. self.session
  403. .push_message(result_message.clone())
  404. .map_err(|error| RuntimeError::new(error.to_string()))?;
  405. self.record_tool_finished(iterations, &result_message);
  406. tool_results.push(result_message);
  407. }
  408. }
  409. let auto_compaction = self.maybe_auto_compact();
  410. let summary = TurnSummary {
  411. assistant_messages,
  412. tool_results,
  413. prompt_cache_events,
  414. iterations,
  415. usage: self.usage_tracker.cumulative_usage(),
  416. auto_compaction,
  417. };
  418. self.record_turn_completed(&summary);
  419. Ok(summary)
  420. }
  421. #[must_use]
  422. pub fn compact(&self, config: CompactionConfig) -> CompactionResult {
  423. compact_session(&self.session, config)
  424. }
  425. #[must_use]
  426. pub fn estimated_tokens(&self) -> usize {
  427. estimate_session_tokens(&self.session)
  428. }
  429. #[must_use]
  430. pub fn usage(&self) -> &UsageTracker {
  431. &self.usage_tracker
  432. }
  433. #[must_use]
  434. pub fn session(&self) -> &Session {
  435. &self.session
  436. }
  437. #[must_use]
  438. pub fn fork_session(&self, branch_name: Option<String>) -> Session {
  439. self.session.fork(branch_name)
  440. }
  441. #[must_use]
  442. pub fn into_session(self) -> Session {
  443. self.session
  444. }
  445. fn maybe_auto_compact(&mut self) -> Option<AutoCompactionEvent> {
  446. if self.usage_tracker.cumulative_usage().input_tokens
  447. < self.auto_compaction_input_tokens_threshold
  448. {
  449. return None;
  450. }
  451. let result = compact_session(
  452. &self.session,
  453. CompactionConfig {
  454. max_estimated_tokens: 0,
  455. ..CompactionConfig::default()
  456. },
  457. );
  458. if result.removed_message_count == 0 {
  459. return None;
  460. }
  461. self.session = result.compacted_session;
  462. Some(AutoCompactionEvent {
  463. removed_message_count: result.removed_message_count,
  464. })
  465. }
  466. fn record_turn_started(&self, user_input: &str) {
  467. let Some(session_tracer) = &self.session_tracer else {
  468. return;
  469. };
  470. let mut attributes = Map::new();
  471. attributes.insert(
  472. "user_input".to_string(),
  473. Value::String(user_input.to_string()),
  474. );
  475. session_tracer.record("turn_started", attributes);
  476. }
  477. fn record_assistant_iteration(
  478. &self,
  479. iteration: usize,
  480. assistant_message: &ConversationMessage,
  481. pending_tool_use_count: usize,
  482. ) {
  483. let Some(session_tracer) = &self.session_tracer else {
  484. return;
  485. };
  486. let mut attributes = Map::new();
  487. attributes.insert("iteration".to_string(), Value::from(iteration as u64));
  488. attributes.insert(
  489. "assistant_blocks".to_string(),
  490. Value::from(assistant_message.blocks.len() as u64),
  491. );
  492. attributes.insert(
  493. "pending_tool_use_count".to_string(),
  494. Value::from(pending_tool_use_count as u64),
  495. );
  496. session_tracer.record("assistant_iteration_completed", attributes);
  497. }
  498. fn record_tool_started(&self, iteration: usize, tool_name: &str) {
  499. let Some(session_tracer) = &self.session_tracer else {
  500. return;
  501. };
  502. let mut attributes = Map::new();
  503. attributes.insert("iteration".to_string(), Value::from(iteration as u64));
  504. attributes.insert(
  505. "tool_name".to_string(),
  506. Value::String(tool_name.to_string()),
  507. );
  508. session_tracer.record("tool_execution_started", attributes);
  509. }
  510. fn record_tool_finished(&self, iteration: usize, result_message: &ConversationMessage) {
  511. let Some(session_tracer) = &self.session_tracer else {
  512. return;
  513. };
  514. let Some(ContentBlock::ToolResult {
  515. tool_name,
  516. is_error,
  517. ..
  518. }) = result_message.blocks.first()
  519. else {
  520. return;
  521. };
  522. let mut attributes = Map::new();
  523. attributes.insert("iteration".to_string(), Value::from(iteration as u64));
  524. attributes.insert("tool_name".to_string(), Value::String(tool_name.clone()));
  525. attributes.insert("is_error".to_string(), Value::Bool(*is_error));
  526. session_tracer.record("tool_execution_finished", attributes);
  527. }
  528. fn record_turn_completed(&self, summary: &TurnSummary) {
  529. let Some(session_tracer) = &self.session_tracer else {
  530. return;
  531. };
  532. let mut attributes = Map::new();
  533. attributes.insert(
  534. "iterations".to_string(),
  535. Value::from(summary.iterations as u64),
  536. );
  537. attributes.insert(
  538. "assistant_messages".to_string(),
  539. Value::from(summary.assistant_messages.len() as u64),
  540. );
  541. attributes.insert(
  542. "tool_results".to_string(),
  543. Value::from(summary.tool_results.len() as u64),
  544. );
  545. attributes.insert(
  546. "prompt_cache_events".to_string(),
  547. Value::from(summary.prompt_cache_events.len() as u64),
  548. );
  549. session_tracer.record("turn_completed", attributes);
  550. }
  551. fn record_turn_failed(&self, iteration: usize, error: &RuntimeError) {
  552. let Some(session_tracer) = &self.session_tracer else {
  553. return;
  554. };
  555. let mut attributes = Map::new();
  556. attributes.insert("iteration".to_string(), Value::from(iteration as u64));
  557. attributes.insert("error".to_string(), Value::String(error.to_string()));
  558. session_tracer.record("turn_failed", attributes);
  559. }
  560. }
  561. #[must_use]
  562. pub fn auto_compaction_threshold_from_env() -> u32 {
  563. parse_auto_compaction_threshold(
  564. std::env::var(AUTO_COMPACTION_THRESHOLD_ENV_VAR)
  565. .ok()
  566. .as_deref(),
  567. )
  568. }
  569. #[must_use]
  570. fn parse_auto_compaction_threshold(value: Option<&str>) -> u32 {
  571. value
  572. .and_then(|raw| raw.trim().parse::<u32>().ok())
  573. .filter(|threshold| *threshold > 0)
  574. .unwrap_or(DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD)
  575. }
  576. fn build_assistant_message(
  577. events: Vec<AssistantEvent>,
  578. ) -> Result<
  579. (
  580. ConversationMessage,
  581. Option<TokenUsage>,
  582. Vec<PromptCacheEvent>,
  583. ),
  584. RuntimeError,
  585. > {
  586. let mut text = String::new();
  587. let mut blocks = Vec::new();
  588. let mut prompt_cache_events = Vec::new();
  589. let mut finished = false;
  590. let mut usage = None;
  591. for event in events {
  592. match event {
  593. AssistantEvent::TextDelta(delta) => text.push_str(&delta),
  594. AssistantEvent::ToolUse { id, name, input } => {
  595. flush_text_block(&mut text, &mut blocks);
  596. blocks.push(ContentBlock::ToolUse { id, name, input });
  597. }
  598. AssistantEvent::Usage(value) => usage = Some(value),
  599. AssistantEvent::PromptCache(event) => prompt_cache_events.push(event),
  600. AssistantEvent::MessageStop => {
  601. finished = true;
  602. }
  603. }
  604. }
  605. flush_text_block(&mut text, &mut blocks);
  606. if !finished {
  607. return Err(RuntimeError::new(
  608. "assistant stream ended without a message stop event",
  609. ));
  610. }
  611. if blocks.is_empty() {
  612. return Err(RuntimeError::new("assistant stream produced no content"));
  613. }
  614. Ok((
  615. ConversationMessage::assistant_with_usage(blocks, usage),
  616. usage,
  617. prompt_cache_events,
  618. ))
  619. }
  620. fn flush_text_block(text: &mut String, blocks: &mut Vec<ContentBlock>) {
  621. if !text.is_empty() {
  622. blocks.push(ContentBlock::Text {
  623. text: std::mem::take(text),
  624. });
  625. }
  626. }
  627. fn format_hook_message(result: &HookRunResult, fallback: &str) -> String {
  628. if result.messages().is_empty() {
  629. fallback.to_string()
  630. } else {
  631. result.messages().join("\n")
  632. }
  633. }
  634. fn merge_hook_feedback(messages: &[String], output: String, denied: bool) -> String {
  635. if messages.is_empty() {
  636. return output;
  637. }
  638. let mut sections = Vec::new();
  639. if !output.trim().is_empty() {
  640. sections.push(output);
  641. }
  642. let label = if denied {
  643. "Hook feedback (denied)"
  644. } else {
  645. "Hook feedback"
  646. };
  647. sections.push(format!("{label}:\n{}", messages.join("\n")));
  648. sections.join("\n\n")
  649. }
  650. type ToolHandler = Box<dyn FnMut(&str) -> Result<String, ToolError>>;
  651. #[derive(Default)]
  652. pub struct StaticToolExecutor {
  653. handlers: BTreeMap<String, ToolHandler>,
  654. }
  655. impl StaticToolExecutor {
  656. #[must_use]
  657. pub fn new() -> Self {
  658. Self::default()
  659. }
  660. #[must_use]
  661. pub fn register(
  662. mut self,
  663. tool_name: impl Into<String>,
  664. handler: impl FnMut(&str) -> Result<String, ToolError> + 'static,
  665. ) -> Self {
  666. self.handlers.insert(tool_name.into(), Box::new(handler));
  667. self
  668. }
  669. }
  670. impl ToolExecutor for StaticToolExecutor {
  671. fn execute(&mut self, tool_name: &str, input: &str) -> Result<String, ToolError> {
  672. self.handlers
  673. .get_mut(tool_name)
  674. .ok_or_else(|| ToolError::new(format!("unknown tool: {tool_name}")))?(input)
  675. }
  676. }
  677. #[cfg(test)]
  678. mod tests {
  679. use super::{
  680. parse_auto_compaction_threshold, ApiClient, ApiRequest, AssistantEvent,
  681. AutoCompactionEvent, ConversationRuntime, PromptCacheEvent, RuntimeError,
  682. StaticToolExecutor, DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD,
  683. };
  684. use crate::compact::CompactionConfig;
  685. use crate::config::{RuntimeFeatureConfig, RuntimeHookConfig};
  686. use crate::permissions::{
  687. PermissionMode, PermissionPolicy, PermissionPromptDecision, PermissionPrompter,
  688. PermissionRequest,
  689. };
  690. use crate::prompt::{ProjectContext, SystemPromptBuilder};
  691. use crate::session::{ContentBlock, MessageRole, Session};
  692. use crate::usage::TokenUsage;
  693. use std::fs;
  694. use std::path::PathBuf;
  695. use std::sync::Arc;
  696. use std::time::{SystemTime, UNIX_EPOCH};
  697. use telemetry::{MemoryTelemetrySink, SessionTracer, TelemetryEvent};
  698. struct ScriptedApiClient {
  699. call_count: usize,
  700. }
  701. impl ApiClient for ScriptedApiClient {
  702. fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> {
  703. self.call_count += 1;
  704. match self.call_count {
  705. 1 => {
  706. assert!(request
  707. .messages
  708. .iter()
  709. .any(|message| message.role == MessageRole::User));
  710. Ok(vec![
  711. AssistantEvent::TextDelta("Let me calculate that.".to_string()),
  712. AssistantEvent::ToolUse {
  713. id: "tool-1".to_string(),
  714. name: "add".to_string(),
  715. input: "2,2".to_string(),
  716. },
  717. AssistantEvent::Usage(TokenUsage {
  718. input_tokens: 20,
  719. output_tokens: 6,
  720. cache_creation_input_tokens: 1,
  721. cache_read_input_tokens: 2,
  722. }),
  723. AssistantEvent::MessageStop,
  724. ])
  725. }
  726. 2 => {
  727. let last_message = request
  728. .messages
  729. .last()
  730. .expect("tool result should be present");
  731. assert_eq!(last_message.role, MessageRole::Tool);
  732. Ok(vec![
  733. AssistantEvent::TextDelta("The answer is 4.".to_string()),
  734. AssistantEvent::Usage(TokenUsage {
  735. input_tokens: 24,
  736. output_tokens: 4,
  737. cache_creation_input_tokens: 1,
  738. cache_read_input_tokens: 3,
  739. }),
  740. AssistantEvent::PromptCache(PromptCacheEvent {
  741. unexpected: true,
  742. reason:
  743. "cache read tokens dropped while prompt fingerprint remained stable"
  744. .to_string(),
  745. previous_cache_read_input_tokens: 6_000,
  746. current_cache_read_input_tokens: 1_000,
  747. token_drop: 5_000,
  748. }),
  749. AssistantEvent::MessageStop,
  750. ])
  751. }
  752. _ => Err(RuntimeError::new("unexpected extra API call")),
  753. }
  754. }
  755. }
  756. struct PromptAllowOnce;
  757. impl PermissionPrompter for PromptAllowOnce {
  758. fn decide(&mut self, request: &PermissionRequest) -> PermissionPromptDecision {
  759. assert_eq!(request.tool_name, "add");
  760. PermissionPromptDecision::Allow
  761. }
  762. }
  763. #[test]
  764. fn runs_user_to_tool_to_result_loop_end_to_end_and_tracks_usage() {
  765. let api_client = ScriptedApiClient { call_count: 0 };
  766. let tool_executor = StaticToolExecutor::new().register("add", |input| {
  767. let total = input
  768. .split(',')
  769. .map(|part| part.parse::<i32>().expect("input must be valid integer"))
  770. .sum::<i32>();
  771. Ok(total.to_string())
  772. });
  773. let permission_policy = PermissionPolicy::new(PermissionMode::WorkspaceWrite);
  774. let system_prompt = SystemPromptBuilder::new()
  775. .with_project_context(ProjectContext {
  776. cwd: PathBuf::from("/tmp/project"),
  777. current_date: "2026-03-31".to_string(),
  778. git_status: None,
  779. git_diff: None,
  780. instruction_files: Vec::new(),
  781. })
  782. .with_os("linux", "6.8")
  783. .build();
  784. let mut runtime = ConversationRuntime::new(
  785. Session::new(),
  786. api_client,
  787. tool_executor,
  788. permission_policy,
  789. system_prompt,
  790. );
  791. let summary = runtime
  792. .run_turn("what is 2 + 2?", Some(&mut PromptAllowOnce))
  793. .expect("conversation loop should succeed");
  794. assert_eq!(summary.iterations, 2);
  795. assert_eq!(summary.assistant_messages.len(), 2);
  796. assert_eq!(summary.tool_results.len(), 1);
  797. assert_eq!(summary.prompt_cache_events.len(), 1);
  798. assert_eq!(runtime.session().messages.len(), 4);
  799. assert_eq!(summary.usage.output_tokens, 10);
  800. assert_eq!(summary.auto_compaction, None);
  801. assert!(matches!(
  802. runtime.session().messages[1].blocks[1],
  803. ContentBlock::ToolUse { .. }
  804. ));
  805. assert!(matches!(
  806. runtime.session().messages[2].blocks[0],
  807. ContentBlock::ToolResult {
  808. is_error: false,
  809. ..
  810. }
  811. ));
  812. }
  813. #[test]
  814. fn records_runtime_session_trace_events() {
  815. let sink = Arc::new(MemoryTelemetrySink::default());
  816. let tracer = SessionTracer::new("session-runtime", sink.clone());
  817. let mut runtime = ConversationRuntime::new(
  818. Session::new(),
  819. ScriptedApiClient { call_count: 0 },
  820. StaticToolExecutor::new().register("add", |_input| Ok("4".to_string())),
  821. PermissionPolicy::new(PermissionMode::WorkspaceWrite),
  822. vec!["system".to_string()],
  823. )
  824. .with_session_tracer(tracer);
  825. runtime
  826. .run_turn("what is 2 + 2?", Some(&mut PromptAllowOnce))
  827. .expect("conversation loop should succeed");
  828. let events = sink.events();
  829. let trace_names = events
  830. .iter()
  831. .filter_map(|event| match event {
  832. TelemetryEvent::SessionTrace(trace) => Some(trace.name.as_str()),
  833. _ => None,
  834. })
  835. .collect::<Vec<_>>();
  836. assert!(trace_names.contains(&"turn_started"));
  837. assert!(trace_names.contains(&"assistant_iteration_completed"));
  838. assert!(trace_names.contains(&"tool_execution_started"));
  839. assert!(trace_names.contains(&"tool_execution_finished"));
  840. assert!(trace_names.contains(&"turn_completed"));
  841. }
  842. #[test]
  843. fn records_denied_tool_results_when_prompt_rejects() {
  844. struct RejectPrompter;
  845. impl PermissionPrompter for RejectPrompter {
  846. fn decide(&mut self, _request: &PermissionRequest) -> PermissionPromptDecision {
  847. PermissionPromptDecision::Deny {
  848. reason: "not now".to_string(),
  849. }
  850. }
  851. }
  852. struct SingleCallApiClient;
  853. impl ApiClient for SingleCallApiClient {
  854. fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> {
  855. if request
  856. .messages
  857. .iter()
  858. .any(|message| message.role == MessageRole::Tool)
  859. {
  860. return Ok(vec![
  861. AssistantEvent::TextDelta("I could not use the tool.".to_string()),
  862. AssistantEvent::MessageStop,
  863. ]);
  864. }
  865. Ok(vec![
  866. AssistantEvent::ToolUse {
  867. id: "tool-1".to_string(),
  868. name: "blocked".to_string(),
  869. input: "secret".to_string(),
  870. },
  871. AssistantEvent::MessageStop,
  872. ])
  873. }
  874. }
  875. let mut runtime = ConversationRuntime::new(
  876. Session::new(),
  877. SingleCallApiClient,
  878. StaticToolExecutor::new(),
  879. PermissionPolicy::new(PermissionMode::WorkspaceWrite),
  880. vec!["system".to_string()],
  881. );
  882. let summary = runtime
  883. .run_turn("use the tool", Some(&mut RejectPrompter))
  884. .expect("conversation should continue after denied tool");
  885. assert_eq!(summary.tool_results.len(), 1);
  886. assert!(matches!(
  887. &summary.tool_results[0].blocks[0],
  888. ContentBlock::ToolResult { is_error: true, output, .. } if output == "not now"
  889. ));
  890. }
  891. #[test]
  892. fn denies_tool_use_when_pre_tool_hook_blocks() {
  893. struct SingleCallApiClient;
  894. impl ApiClient for SingleCallApiClient {
  895. fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> {
  896. if request
  897. .messages
  898. .iter()
  899. .any(|message| message.role == MessageRole::Tool)
  900. {
  901. return Ok(vec![
  902. AssistantEvent::TextDelta("blocked".to_string()),
  903. AssistantEvent::MessageStop,
  904. ]);
  905. }
  906. Ok(vec![
  907. AssistantEvent::ToolUse {
  908. id: "tool-1".to_string(),
  909. name: "blocked".to_string(),
  910. input: r#"{"path":"secret.txt"}"#.to_string(),
  911. },
  912. AssistantEvent::MessageStop,
  913. ])
  914. }
  915. }
  916. let mut runtime = ConversationRuntime::new_with_features(
  917. Session::new(),
  918. SingleCallApiClient,
  919. StaticToolExecutor::new().register("blocked", |_input| {
  920. panic!("tool should not execute when hook denies")
  921. }),
  922. PermissionPolicy::new(PermissionMode::DangerFullAccess),
  923. vec!["system".to_string()],
  924. &RuntimeFeatureConfig::default().with_hooks(RuntimeHookConfig::new(
  925. vec![shell_snippet("printf 'blocked by hook'; exit 2")],
  926. Vec::new(),
  927. Vec::new(),
  928. )),
  929. );
  930. let summary = runtime
  931. .run_turn("use the tool", None)
  932. .expect("conversation should continue after hook denial");
  933. assert_eq!(summary.tool_results.len(), 1);
  934. let ContentBlock::ToolResult {
  935. is_error, output, ..
  936. } = &summary.tool_results[0].blocks[0]
  937. else {
  938. panic!("expected tool result block");
  939. };
  940. assert!(
  941. *is_error,
  942. "hook denial should produce an error result: {output}"
  943. );
  944. assert!(
  945. output.contains("denied tool") || output.contains("blocked by hook"),
  946. "unexpected hook denial output: {output:?}"
  947. );
  948. }
  949. #[test]
  950. fn appends_post_tool_hook_feedback_to_tool_result() {
  951. struct TwoCallApiClient {
  952. calls: usize,
  953. }
  954. impl ApiClient for TwoCallApiClient {
  955. fn stream(&mut self, request: ApiRequest) -> Result<Vec<AssistantEvent>, RuntimeError> {
  956. self.calls += 1;
  957. match self.calls {
  958. 1 => Ok(vec![
  959. AssistantEvent::ToolUse {
  960. id: "tool-1".to_string(),
  961. name: "add".to_string(),
  962. input: r#"{"lhs":2,"rhs":2}"#.to_string(),
  963. },
  964. AssistantEvent::MessageStop,
  965. ]),
  966. 2 => {
  967. assert!(request
  968. .messages
  969. .iter()
  970. .any(|message| message.role == MessageRole::Tool));
  971. Ok(vec![
  972. AssistantEvent::TextDelta("done".to_string()),
  973. AssistantEvent::MessageStop,
  974. ])
  975. }
  976. _ => Err(RuntimeError::new("unexpected extra API call")),
  977. }
  978. }
  979. }
  980. let mut runtime = ConversationRuntime::new_with_features(
  981. Session::new(),
  982. TwoCallApiClient { calls: 0 },
  983. StaticToolExecutor::new().register("add", |_input| Ok("4".to_string())),
  984. PermissionPolicy::new(PermissionMode::DangerFullAccess),
  985. vec!["system".to_string()],
  986. &RuntimeFeatureConfig::default().with_hooks(RuntimeHookConfig::new(
  987. vec![shell_snippet("printf 'pre hook ran'")],
  988. vec![shell_snippet("printf 'post hook ran'")],
  989. Vec::new(),
  990. )),
  991. );
  992. let summary = runtime
  993. .run_turn("use add", None)
  994. .expect("tool loop succeeds");
  995. assert_eq!(summary.tool_results.len(), 1);
  996. let ContentBlock::ToolResult {
  997. is_error, output, ..
  998. } = &summary.tool_results[0].blocks[0]
  999. else {
  1000. panic!("expected tool result block");
  1001. };
  1002. assert!(
  1003. !*is_error,
  1004. "post hook should preserve non-error result: {output:?}"
  1005. );
  1006. assert!(
  1007. output.contains('4'),
  1008. "tool output missing value: {output:?}"
  1009. );
  1010. assert!(
  1011. output.contains("pre hook ran"),
  1012. "tool output missing pre hook feedback: {output:?}"
  1013. );
  1014. assert!(
  1015. output.contains("post hook ran"),
  1016. "tool output missing post hook feedback: {output:?}"
  1017. );
  1018. }
  1019. #[test]
  1020. fn reconstructs_usage_tracker_from_restored_session() {
  1021. struct SimpleApi;
  1022. impl ApiClient for SimpleApi {
  1023. fn stream(
  1024. &mut self,
  1025. _request: ApiRequest,
  1026. ) -> Result<Vec<AssistantEvent>, RuntimeError> {
  1027. Ok(vec![
  1028. AssistantEvent::TextDelta("done".to_string()),
  1029. AssistantEvent::MessageStop,
  1030. ])
  1031. }
  1032. }
  1033. let mut session = Session::new();
  1034. session
  1035. .messages
  1036. .push(crate::session::ConversationMessage::assistant_with_usage(
  1037. vec![ContentBlock::Text {
  1038. text: "earlier".to_string(),
  1039. }],
  1040. Some(TokenUsage {
  1041. input_tokens: 11,
  1042. output_tokens: 7,
  1043. cache_creation_input_tokens: 2,
  1044. cache_read_input_tokens: 1,
  1045. }),
  1046. ));
  1047. let runtime = ConversationRuntime::new(
  1048. session,
  1049. SimpleApi,
  1050. StaticToolExecutor::new(),
  1051. PermissionPolicy::new(PermissionMode::DangerFullAccess),
  1052. vec!["system".to_string()],
  1053. );
  1054. assert_eq!(runtime.usage().turns(), 1);
  1055. assert_eq!(runtime.usage().cumulative_usage().total_tokens(), 21);
  1056. }
  1057. #[test]
  1058. fn compacts_session_after_turns() {
  1059. struct SimpleApi;
  1060. impl ApiClient for SimpleApi {
  1061. fn stream(
  1062. &mut self,
  1063. _request: ApiRequest,
  1064. ) -> Result<Vec<AssistantEvent>, RuntimeError> {
  1065. Ok(vec![
  1066. AssistantEvent::TextDelta("done".to_string()),
  1067. AssistantEvent::MessageStop,
  1068. ])
  1069. }
  1070. }
  1071. let mut runtime = ConversationRuntime::new(
  1072. Session::new(),
  1073. SimpleApi,
  1074. StaticToolExecutor::new(),
  1075. PermissionPolicy::new(PermissionMode::DangerFullAccess),
  1076. vec!["system".to_string()],
  1077. );
  1078. runtime.run_turn("a", None).expect("turn a");
  1079. runtime.run_turn("b", None).expect("turn b");
  1080. runtime.run_turn("c", None).expect("turn c");
  1081. let result = runtime.compact(CompactionConfig {
  1082. preserve_recent_messages: 2,
  1083. max_estimated_tokens: 1,
  1084. });
  1085. assert!(result.summary.contains("Conversation summary"));
  1086. assert_eq!(
  1087. result.compacted_session.messages[0].role,
  1088. MessageRole::System
  1089. );
  1090. assert_eq!(
  1091. result.compacted_session.session_id,
  1092. runtime.session().session_id
  1093. );
  1094. assert!(result.compacted_session.compaction.is_some());
  1095. }
  1096. #[test]
  1097. fn persists_conversation_turn_messages_to_jsonl_session() {
  1098. struct SimpleApi;
  1099. impl ApiClient for SimpleApi {
  1100. fn stream(
  1101. &mut self,
  1102. _request: ApiRequest,
  1103. ) -> Result<Vec<AssistantEvent>, RuntimeError> {
  1104. Ok(vec![
  1105. AssistantEvent::TextDelta("done".to_string()),
  1106. AssistantEvent::MessageStop,
  1107. ])
  1108. }
  1109. }
  1110. let path = temp_session_path("persisted-turn");
  1111. let session = Session::new().with_persistence_path(path.clone());
  1112. let mut runtime = ConversationRuntime::new(
  1113. session,
  1114. SimpleApi,
  1115. StaticToolExecutor::new(),
  1116. PermissionPolicy::new(PermissionMode::DangerFullAccess),
  1117. vec!["system".to_string()],
  1118. );
  1119. runtime
  1120. .run_turn("persist this turn", None)
  1121. .expect("turn should succeed");
  1122. let restored = Session::load_from_path(&path).expect("persisted session should reload");
  1123. fs::remove_file(&path).expect("temp session file should be removable");
  1124. assert_eq!(restored.messages.len(), 2);
  1125. assert_eq!(restored.messages[0].role, MessageRole::User);
  1126. assert_eq!(restored.messages[1].role, MessageRole::Assistant);
  1127. assert_eq!(restored.session_id, runtime.session().session_id);
  1128. }
  1129. #[test]
  1130. fn forks_runtime_session_without_mutating_original() {
  1131. let mut session = Session::new();
  1132. session
  1133. .push_user_text("branch me")
  1134. .expect("message should append");
  1135. let runtime = ConversationRuntime::new(
  1136. session.clone(),
  1137. ScriptedApiClient { call_count: 0 },
  1138. StaticToolExecutor::new(),
  1139. PermissionPolicy::new(PermissionMode::DangerFullAccess),
  1140. vec!["system".to_string()],
  1141. );
  1142. let forked = runtime.fork_session(Some("alt-path".to_string()));
  1143. assert_eq!(forked.messages, session.messages);
  1144. assert_ne!(forked.session_id, session.session_id);
  1145. assert_eq!(
  1146. forked
  1147. .fork
  1148. .as_ref()
  1149. .map(|fork| (fork.parent_session_id.as_str(), fork.branch_name.as_deref())),
  1150. Some((session.session_id.as_str(), Some("alt-path")))
  1151. );
  1152. assert!(runtime.session().fork.is_none());
  1153. }
  1154. fn temp_session_path(label: &str) -> PathBuf {
  1155. let nanos = SystemTime::now()
  1156. .duration_since(UNIX_EPOCH)
  1157. .expect("system time should be after epoch")
  1158. .as_nanos();
  1159. std::env::temp_dir().join(format!("runtime-conversation-{label}-{nanos}.json"))
  1160. }
  1161. #[cfg(windows)]
  1162. fn shell_snippet(script: &str) -> String {
  1163. script.replace('\'', "\"")
  1164. }
  1165. #[cfg(not(windows))]
  1166. fn shell_snippet(script: &str) -> String {
  1167. script.to_string()
  1168. }
  1169. #[test]
  1170. fn auto_compacts_when_cumulative_input_threshold_is_crossed() {
  1171. struct SimpleApi;
  1172. impl ApiClient for SimpleApi {
  1173. fn stream(
  1174. &mut self,
  1175. _request: ApiRequest,
  1176. ) -> Result<Vec<AssistantEvent>, RuntimeError> {
  1177. Ok(vec![
  1178. AssistantEvent::TextDelta("done".to_string()),
  1179. AssistantEvent::Usage(TokenUsage {
  1180. input_tokens: 120_000,
  1181. output_tokens: 4,
  1182. cache_creation_input_tokens: 0,
  1183. cache_read_input_tokens: 0,
  1184. }),
  1185. AssistantEvent::MessageStop,
  1186. ])
  1187. }
  1188. }
  1189. let mut session = Session::new();
  1190. session.messages = vec![
  1191. crate::session::ConversationMessage::user_text("one"),
  1192. crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
  1193. text: "two".to_string(),
  1194. }]),
  1195. crate::session::ConversationMessage::user_text("three"),
  1196. crate::session::ConversationMessage::assistant(vec![ContentBlock::Text {
  1197. text: "four".to_string(),
  1198. }]),
  1199. ];
  1200. let mut runtime = ConversationRuntime::new(
  1201. session,
  1202. SimpleApi,
  1203. StaticToolExecutor::new(),
  1204. PermissionPolicy::new(PermissionMode::DangerFullAccess),
  1205. vec!["system".to_string()],
  1206. )
  1207. .with_auto_compaction_input_tokens_threshold(100_000);
  1208. let summary = runtime
  1209. .run_turn("trigger", None)
  1210. .expect("turn should succeed");
  1211. assert_eq!(
  1212. summary.auto_compaction,
  1213. Some(AutoCompactionEvent {
  1214. removed_message_count: 2,
  1215. })
  1216. );
  1217. assert_eq!(runtime.session().messages[0].role, MessageRole::System);
  1218. }
  1219. #[test]
  1220. fn skips_auto_compaction_below_threshold() {
  1221. struct SimpleApi;
  1222. impl ApiClient for SimpleApi {
  1223. fn stream(
  1224. &mut self,
  1225. _request: ApiRequest,
  1226. ) -> Result<Vec<AssistantEvent>, RuntimeError> {
  1227. Ok(vec![
  1228. AssistantEvent::TextDelta("done".to_string()),
  1229. AssistantEvent::Usage(TokenUsage {
  1230. input_tokens: 99_999,
  1231. output_tokens: 4,
  1232. cache_creation_input_tokens: 0,
  1233. cache_read_input_tokens: 0,
  1234. }),
  1235. AssistantEvent::MessageStop,
  1236. ])
  1237. }
  1238. }
  1239. let mut runtime = ConversationRuntime::new(
  1240. Session::new(),
  1241. SimpleApi,
  1242. StaticToolExecutor::new(),
  1243. PermissionPolicy::new(PermissionMode::DangerFullAccess),
  1244. vec!["system".to_string()],
  1245. )
  1246. .with_auto_compaction_input_tokens_threshold(100_000);
  1247. let summary = runtime
  1248. .run_turn("trigger", None)
  1249. .expect("turn should succeed");
  1250. assert_eq!(summary.auto_compaction, None);
  1251. assert_eq!(runtime.session().messages.len(), 2);
  1252. }
  1253. #[test]
  1254. fn auto_compaction_threshold_defaults_and_parses_values() {
  1255. assert_eq!(
  1256. parse_auto_compaction_threshold(None),
  1257. DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD
  1258. );
  1259. assert_eq!(parse_auto_compaction_threshold(Some("4321")), 4321);
  1260. assert_eq!(
  1261. parse_auto_compaction_threshold(Some("not-a-number")),
  1262. DEFAULT_AUTO_COMPACTION_INPUT_TOKENS_THRESHOLD
  1263. );
  1264. }
  1265. }