codex_wrapper/
session.rs

1//! Stateful multi-turn session manager for the Codex CLI.
2//!
3//! [`Session`] wraps a [`Codex`] client and automatically threads
4//! conversation state across turns. The first call to [`send`](Session::send)
5//! dispatches via [`ExecCommand`]; subsequent calls use
6//! [`ExecResumeCommand`] with the captured `thread_id`.
7//!
8//! # Example
9//!
10//! ```no_run
11//! use std::sync::Arc;
12//! use codex_wrapper::{Codex, Session};
13//!
14//! # async fn example() -> codex_wrapper::Result<()> {
15//! let codex = Arc::new(Codex::builder().build()?);
16//! let mut session = Session::new(codex);
17//!
18//! let events = session.send("create a hello world program").await?;
19//! println!("turn 1: {} events", events.len());
20//!
21//! let events = session.send("now add error handling").await?;
22//! println!("turn 2: {} events, thread_id={:?}", events.len(), session.id());
23//! # Ok(())
24//! # }
25//! ```
26
27use std::sync::Arc;
28
29use crate::Codex;
30use crate::command::exec::{ExecCommand, ExecResumeCommand};
31use crate::error::{Error, Result};
32use crate::types::{JsonLineEvent, QueryResult, TokenUsage};
33
34/// A record of a single turn within a session.
35#[derive(Debug, Clone)]
36pub struct TurnRecord {
37    /// The typed result for this turn, assembled from its JSONL events.
38    ///
39    /// Holding the assembled [`QueryResult`] rather than the raw event vector
40    /// is what lets a session report usage: token counts live on the terminal
41    /// `turn.completed` event and are otherwise discarded.
42    pub result: QueryResult,
43}
44
45impl TurnRecord {
46    /// The parsed JSONL events returned by this turn.
47    #[must_use]
48    pub fn events(&self) -> &[JsonLineEvent] {
49        &self.result.events
50    }
51
52    /// Token counts for this turn, if the CLI reported them.
53    #[must_use]
54    pub fn usage(&self) -> Option<TokenUsage> {
55        self.result.usage
56    }
57}
58
59/// Stateful multi-turn session manager.
60///
61/// Wraps a [`Codex`] client and automatically threads conversation state
62/// across turns. On the first turn, an [`ExecCommand`] is used; on subsequent
63/// turns, an [`ExecResumeCommand`] resumes the session using the `thread_id`
64/// extracted from the JSONL event stream.
65///
66/// The `thread_id` is preserved even when a turn fails, as long as at least
67/// one event in the output carried it.
68///
69/// # Example
70///
71/// ```no_run
72/// use std::sync::Arc;
73/// use codex_wrapper::{Codex, Session};
74///
75/// # async fn example() -> codex_wrapper::Result<()> {
76/// let codex = Arc::new(Codex::builder().build()?);
77/// let mut session = Session::new(codex);
78///
79/// let events = session.send("summarize this repo").await?;
80/// assert!(session.id().is_some());
81/// assert_eq!(session.total_turns(), 1);
82///
83/// let events = session.send("now add more detail").await?;
84/// assert_eq!(session.total_turns(), 2);
85/// # Ok(())
86/// # }
87/// ```
88pub struct Session {
89    codex: Arc<Codex>,
90    thread_id: Option<String>,
91    history: Vec<TurnRecord>,
92    budget: Option<crate::budget::TokenBudget>,
93}
94
95impl Session {
96    /// Create a new session with no prior state.
97    ///
98    /// The first call to [`send`](Session::send) will use [`ExecCommand`].
99    pub fn new(codex: Arc<Codex>) -> Self {
100        Self {
101            codex,
102            thread_id: None,
103            history: Vec::new(),
104            budget: None,
105        }
106    }
107
108    /// Resume an existing session by its `thread_id`.
109    ///
110    /// The next call to [`send`](Session::send) will use
111    /// [`ExecResumeCommand`] with the provided ID.
112    pub fn resume(codex: Arc<Codex>, thread_id: impl Into<String>) -> Self {
113        Self {
114            codex,
115            thread_id: Some(thread_id.into()),
116            history: Vec::new(),
117            budget: None,
118        }
119    }
120
121    /// Attach a [`TokenBudget`](crate::TokenBudget).
122    ///
123    /// Every turn's reported usage is added to it, and each turn is refused
124    /// with [`Error::TokenBudgetExceeded`] once the ceiling is reached. The
125    /// check happens before a turn starts, so the ceiling can be overshot by
126    /// at most the turn that crosses it: usage is only known once spent.
127    ///
128    /// The budget is [`Clone`] over shared state, so one can span several
129    /// sessions.
130    ///
131    /// ```
132    /// use std::sync::Arc;
133    /// use codex_wrapper::{Codex, Session, TokenBudget};
134    ///
135    /// # fn example() -> codex_wrapper::Result<()> {
136    /// let codex = Arc::new(Codex::builder().build()?);
137    /// let budget = TokenBudget::builder().max_tokens(200_000).build();
138    /// let session = Session::new(codex).with_budget(budget.clone());
139    /// # let _ = session;
140    /// # Ok(())
141    /// # }
142    /// ```
143    #[must_use]
144    pub fn with_budget(mut self, budget: crate::budget::TokenBudget) -> Self {
145        self.budget = Some(budget);
146        self
147    }
148
149    /// The attached budget, if any.
150    #[must_use]
151    pub fn budget(&self) -> Option<&crate::budget::TokenBudget> {
152        self.budget.as_ref()
153    }
154
155    /// Send a prompt, automatically routing to `exec` or `exec resume`.
156    ///
157    /// On the first turn (no `thread_id`), dispatches via [`ExecCommand`].
158    /// On subsequent turns, dispatches via [`ExecResumeCommand`] with the
159    /// captured `thread_id`.
160    ///
161    /// Returns the parsed JSONL events for this turn.
162    pub async fn send(&mut self, prompt: impl Into<String>) -> Result<Vec<JsonLineEvent>> {
163        let prompt = prompt.into();
164
165        match &self.thread_id {
166            None => {
167                let cmd = ExecCommand::new(&prompt);
168                self.run_exec(cmd).await
169            }
170            Some(id) => {
171                let cmd = ExecResumeCommand::new()
172                    .session_id(id.clone())
173                    .prompt(prompt);
174                self.run_resume(cmd).await
175            }
176        }
177    }
178
179    /// Execute an [`ExecCommand`] with full control over its options.
180    ///
181    /// Use this when you need to configure model, sandbox, approval policy,
182    /// or other flags beyond what [`send`](Session::send) provides.
183    /// The session still captures the `thread_id` from the output.
184    pub async fn execute(&mut self, cmd: ExecCommand) -> Result<Vec<JsonLineEvent>> {
185        self.run_exec(cmd).await
186    }
187
188    /// Execute an [`ExecResumeCommand`] with full control over its options.
189    ///
190    /// Use this when you need to configure flags on the resume command
191    /// beyond what [`send`](Session::send) provides.
192    /// The session still captures the `thread_id` from the output.
193    pub async fn execute_resume(&mut self, cmd: ExecResumeCommand) -> Result<Vec<JsonLineEvent>> {
194        self.run_resume(cmd).await
195    }
196
197    /// Send a prompt, streaming events to `handler` as they arrive.
198    ///
199    /// The streaming equivalent of [`send`](Session::send): routes to `exec`
200    /// or `exec resume` the same way, and leaves the session in the same
201    /// state. Events are handed to `handler` as the CLI emits them and are
202    /// also retained, so `thread_id`, history, and cost are captured exactly
203    /// as they would be on the buffered path.
204    ///
205    /// Returns the full event stream for the turn.
206    ///
207    /// ```no_run
208    /// use std::sync::Arc;
209    /// use codex_wrapper::{Codex, Session};
210    ///
211    /// # async fn example() -> codex_wrapper::Result<()> {
212    /// let codex = Arc::new(Codex::builder().build()?);
213    /// let mut session = Session::new(codex);
214    ///
215    /// session
216    ///     .stream("summarize this repo", |event| {
217    ///         println!("{}", event.event_type);
218    ///     })
219    ///     .await?;
220    ///
221    /// assert_eq!(session.total_turns(), 1);
222    /// # Ok(())
223    /// # }
224    /// ```
225    pub async fn stream<F>(
226        &mut self,
227        prompt: impl Into<String>,
228        handler: F,
229    ) -> Result<Vec<JsonLineEvent>>
230    where
231        F: FnMut(JsonLineEvent),
232    {
233        let prompt = prompt.into();
234        match &self.thread_id {
235            None => {
236                self.stream_execute(ExecCommand::new(&prompt), handler)
237                    .await
238            }
239            Some(id) => {
240                let cmd = ExecResumeCommand::new()
241                    .session_id(id.clone())
242                    .prompt(prompt);
243                self.stream_execute_resume(cmd, handler).await
244            }
245        }
246    }
247
248    /// Stream an [`ExecCommand`] with full control over its options.
249    ///
250    /// The streaming equivalent of [`execute`](Session::execute).
251    pub async fn stream_execute<F>(
252        &mut self,
253        cmd: ExecCommand,
254        mut handler: F,
255    ) -> Result<Vec<JsonLineEvent>>
256    where
257        F: FnMut(JsonLineEvent),
258    {
259        self.check_budget()?;
260        let codex = Arc::clone(&self.codex);
261        let mut collected = Vec::new();
262        let outcome = cmd
263            .stream(&codex, |event| {
264                collected.push(event.clone());
265                handler(event);
266            })
267            .await;
268        self.finish_stream(collected, outcome)
269    }
270
271    /// Stream an [`ExecResumeCommand`] with full control over its options.
272    ///
273    /// The streaming equivalent of [`execute_resume`](Session::execute_resume).
274    pub async fn stream_execute_resume<F>(
275        &mut self,
276        cmd: ExecResumeCommand,
277        mut handler: F,
278    ) -> Result<Vec<JsonLineEvent>>
279    where
280        F: FnMut(JsonLineEvent),
281    {
282        self.check_budget()?;
283        let codex = Arc::clone(&self.codex);
284        let mut collected = Vec::new();
285        let outcome = cmd
286            .stream(&codex, |event| {
287                collected.push(event.clone());
288                handler(event);
289            })
290            .await;
291        self.finish_stream(collected, outcome)
292    }
293
294    /// Record a streamed turn, or salvage `thread_id` from a failed one.
295    ///
296    /// A stream that fails partway has still delivered real events, so the
297    /// `thread_id` is taken from what arrived. This mirrors the buffered
298    /// path's recovery, which re-parses stdout for the same reason. No turn is
299    /// recorded on failure: an incomplete stream has no terminal `completed`
300    /// event, so its usage would be missing and would silently undercount
301    /// [`total_tokens`](Self::total_tokens).
302    fn finish_stream(
303        &mut self,
304        collected: Vec<JsonLineEvent>,
305        outcome: Result<()>,
306    ) -> Result<Vec<JsonLineEvent>> {
307        match outcome {
308            Ok(()) => Ok(self.record_turn(collected)),
309            Err(e) => {
310                self.capture_thread_id(&collected);
311                Err(e)
312            }
313        }
314    }
315
316    /// Returns the `thread_id` captured from the most recent turn, if any.
317    #[must_use]
318    pub fn id(&self) -> Option<&str> {
319        self.thread_id.as_deref()
320    }
321
322    /// Total number of completed turns in this session.
323    #[must_use]
324    pub fn total_turns(&self) -> usize {
325        self.history.len()
326    }
327
328    /// Borrow the full turn history.
329    #[must_use]
330    pub fn history(&self) -> &[TurnRecord] {
331        &self.history
332    }
333
334    /// The typed result of the most recent completed turn.
335    #[must_use]
336    pub fn last_result(&self) -> Option<&QueryResult> {
337        self.history.last().map(|turn| &turn.result)
338    }
339
340    /// Sum of the per-turn token totals the CLI reported.
341    ///
342    /// Turns that reported no usage contribute nothing, so a total of `0` can
343    /// mean either "nothing was used" or "nothing was reported". Pair this
344    /// with [`turns_missing_usage`](Self::turns_missing_usage) to tell those
345    /// apart.
346    ///
347    /// There is no cost equivalent: the CLI reports tokens, not money. See
348    /// [`TokenUsage`].
349    #[must_use]
350    pub fn total_tokens(&self) -> u64 {
351        self.history
352            .iter()
353            .filter_map(|turn| turn.usage().and_then(|u| u.total()))
354            .sum()
355    }
356
357    /// How many completed turns reported no usable token total.
358    ///
359    /// Non-zero means [`total_tokens`](Self::total_tokens) is an undercount
360    /// rather than a full accounting.
361    #[must_use]
362    pub fn turns_missing_usage(&self) -> usize {
363        self.history
364            .iter()
365            .filter(|turn| turn.usage().and_then(|u| u.total()).is_none())
366            .count()
367    }
368
369    /// Capture state from a completed turn and record it.
370    ///
371    /// Shared by the buffered and streaming paths so a streaming turn leaves
372    /// the session in the same state a buffered one would.
373    fn record_turn(&mut self, events: Vec<JsonLineEvent>) -> Vec<JsonLineEvent> {
374        self.capture_thread_id(&events);
375        let result = QueryResult::from_events(events);
376        let events = result.events.clone();
377        if let Some(budget) = &self.budget {
378            // `None` when the turn reported no usage, which the budget counts
379            // separately rather than treating as zero consumption.
380            budget.record(result.usage.and_then(|usage| usage.total()));
381        }
382        self.history.push(TurnRecord { result });
383        events
384    }
385
386    /// Refuse a turn that would run past the budget.
387    ///
388    /// Checked before the turn rather than after, so the ceiling stops the
389    /// next run instead of being noticed once more has already been spent.
390    fn check_budget(&self) -> Result<()> {
391        match &self.budget {
392            Some(budget) => budget.check(),
393            None => Ok(()),
394        }
395    }
396
397    /// Run an [`ExecCommand`] and record the turn.
398    async fn run_exec(&mut self, cmd: ExecCommand) -> Result<Vec<JsonLineEvent>> {
399        self.check_budget()?;
400        match cmd.execute_json_lines(&self.codex).await {
401            Ok(events) => Ok(self.record_turn(events)),
402            Err(Error::CommandFailed {
403                stdout,
404                stderr,
405                exit_code,
406                command,
407                working_dir,
408            }) => {
409                self.try_capture_thread_id_from_stdout(&stdout);
410                Err(Error::CommandFailed {
411                    stdout,
412                    stderr,
413                    exit_code,
414                    command,
415                    working_dir,
416                })
417            }
418            Err(e) => Err(e),
419        }
420    }
421
422    /// Run an [`ExecResumeCommand`] and record the turn.
423    async fn run_resume(&mut self, cmd: ExecResumeCommand) -> Result<Vec<JsonLineEvent>> {
424        self.check_budget()?;
425        match cmd.execute_json_lines(&self.codex).await {
426            Ok(events) => Ok(self.record_turn(events)),
427            Err(Error::CommandFailed {
428                stdout,
429                stderr,
430                exit_code,
431                command,
432                working_dir,
433            }) => {
434                self.try_capture_thread_id_from_stdout(&stdout);
435                Err(Error::CommandFailed {
436                    stdout,
437                    stderr,
438                    exit_code,
439                    command,
440                    working_dir,
441                })
442            }
443            Err(e) => Err(e),
444        }
445    }
446
447    /// Extract `thread_id` from parsed events (first match wins).
448    fn capture_thread_id(&mut self, events: &[JsonLineEvent]) {
449        if let Some(id) = events.iter().find_map(|e| e.thread_id()) {
450            self.thread_id = Some(id.to_string());
451        }
452    }
453
454    /// Best-effort extraction of `thread_id` from raw stdout on error paths.
455    fn try_capture_thread_id_from_stdout(&mut self, stdout: &str) {
456        for line in stdout.lines() {
457            if let Ok(event) = serde_json::from_str::<JsonLineEvent>(line)
458                && let Some(id) = event.thread_id()
459            {
460                self.thread_id = Some(id.to_string());
461                return;
462            }
463        }
464    }
465}
466
467impl std::fmt::Debug for Session {
468    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
469        f.debug_struct("Session")
470            .field("thread_id", &self.thread_id)
471            .field("total_turns", &self.history.len())
472            .finish_non_exhaustive()
473    }
474}
475
476#[cfg(test)]
477mod tests {
478    use super::*;
479
480    fn test_codex() -> Arc<Codex> {
481        Arc::new(Codex::builder().binary("/usr/bin/false").build().unwrap())
482    }
483
484    #[test]
485    fn new_session_has_no_state() {
486        let session = Session::new(test_codex());
487        assert!(session.id().is_none());
488        assert_eq!(session.total_turns(), 0);
489        assert!(session.history().is_empty());
490    }
491
492    #[test]
493    fn resume_session_has_thread_id() {
494        let session = Session::resume(test_codex(), "thread_abc");
495        assert_eq!(session.id(), Some("thread_abc"));
496        assert_eq!(session.total_turns(), 0);
497    }
498
499    #[test]
500    fn capture_thread_id_from_events() {
501        let mut session = Session::new(test_codex());
502        let events: Vec<JsonLineEvent> = vec![
503            serde_json::from_str(r#"{"type":"message.created","role":"assistant"}"#).unwrap(),
504            serde_json::from_str(
505                r#"{"type":"thread.started","thread_id":"thread_xyz","session_id":"sess_1"}"#,
506            )
507            .unwrap(),
508        ];
509        session.capture_thread_id(&events);
510        assert_eq!(session.id(), Some("thread_xyz"));
511    }
512
513    #[test]
514    fn capture_thread_id_noop_when_absent() {
515        let mut session = Session::new(test_codex());
516        let events: Vec<JsonLineEvent> =
517            vec![serde_json::from_str(r#"{"type":"message.created"}"#).unwrap()];
518        session.capture_thread_id(&events);
519        assert!(session.id().is_none());
520    }
521
522    #[test]
523    fn try_capture_thread_id_from_stdout_parses_json() {
524        let mut session = Session::new(test_codex());
525        let stdout = r#"{"type":"thread.started","thread_id":"thread_err"}
526{"type":"error","message":"something went wrong"}"#;
527        session.try_capture_thread_id_from_stdout(stdout);
528        assert_eq!(session.id(), Some("thread_err"));
529    }
530
531    #[test]
532    fn try_capture_thread_id_from_stdout_ignores_garbage() {
533        let mut session = Session::new(test_codex());
534        session.try_capture_thread_id_from_stdout("not json\nalso not json");
535        assert!(session.id().is_none());
536    }
537
538    #[test]
539    fn debug_impl() {
540        let session = Session::resume(test_codex(), "thread_dbg");
541        let debug = format!("{session:?}");
542        assert!(debug.contains("thread_dbg"));
543        assert!(debug.contains("total_turns: 0"));
544    }
545
546    /// Build a session backed by the fake-codex script the streaming tests use.
547    ///
548    /// Unix-only: the fake CLI is a bash script, matching how
549    /// [`crate::streaming`] gates its own tests.
550    #[cfg(unix)]
551    fn streaming_session() -> Session {
552        let script = std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
553            .join("tests")
554            .join("fake-codex.sh");
555        Session::new(Arc::new(
556            Codex::builder()
557                .binary("/bin/bash")
558                .arg(script.to_str().unwrap())
559                .build()
560                .expect("bash must exist"),
561        ))
562    }
563
564    fn turn(json: &[&str]) -> Vec<JsonLineEvent> {
565        json.iter()
566            .map(|line| serde_json::from_str(line).unwrap())
567            .collect()
568    }
569
570    /// Build a completed turn reporting `total` tokens.
571    fn completed(total: u64) -> Vec<JsonLineEvent> {
572        turn(&[&format!(
573            r#"{{"type":"turn.completed","usage":{{"total_tokens":{total}}}}}"#
574        )])
575    }
576
577    #[test]
578    fn record_turn_captures_usage_and_thread_id() {
579        let mut session = Session::new(test_codex());
580        session.record_turn(turn(&[
581            r#"{"type":"thread.started","thread_id":"thread_1"}"#,
582            r#"{"type":"item.completed","item":{"item_type":"agent_message","text":"hi"}}"#,
583            r#"{"type":"turn.completed","usage":{"total_tokens":42}}"#,
584        ]));
585
586        assert_eq!(session.id(), Some("thread_1"));
587        assert_eq!(session.total_turns(), 1);
588        assert_eq!(session.last_result().unwrap().result, "hi");
589        assert_eq!(session.total_tokens(), 42);
590        assert_eq!(session.turns_missing_usage(), 0);
591    }
592
593    #[test]
594    fn total_tokens_sums_across_turns() {
595        let mut session = Session::new(test_codex());
596        for total in [10, 20, 30] {
597            session.record_turn(completed(total));
598        }
599        assert_eq!(session.total_turns(), 3);
600        assert_eq!(session.total_tokens(), 60);
601        assert_eq!(session.turns_missing_usage(), 0);
602    }
603
604    /// The CLI does not always report usage. A total that silently skipped
605    /// those turns would look authoritative while undercounting, so the count
606    /// of unreported turns is exposed alongside it.
607    #[test]
608    fn unreported_usage_is_counted_not_hidden() {
609        let mut session = Session::new(test_codex());
610        session.record_turn(completed(40));
611        session.record_turn(turn(&[r#"{"type":"turn.completed"}"#]));
612
613        assert_eq!(session.total_tokens(), 40);
614        assert_eq!(session.turns_missing_usage(), 1);
615        assert_eq!(session.total_turns(), 2);
616    }
617
618    #[test]
619    fn zero_usage_and_unreported_usage_are_distinguishable() {
620        let mut reported = Session::new(test_codex());
621        reported.record_turn(completed(0));
622
623        let mut unreported = Session::new(test_codex());
624        unreported.record_turn(turn(&[r#"{"type":"turn.completed"}"#]));
625
626        assert_eq!(reported.total_tokens(), unreported.total_tokens());
627        assert_eq!(reported.turns_missing_usage(), 0);
628        assert_eq!(unreported.turns_missing_usage(), 1);
629    }
630
631    #[test]
632    fn turn_record_exposes_events_and_usage() {
633        let mut session = Session::new(test_codex());
634        session.record_turn(turn(&[
635            r#"{"type":"turn.started"}"#,
636            r#"{"type":"turn.completed","usage":{"total_tokens":5}}"#,
637        ]));
638
639        let record = &session.history()[0];
640        assert_eq!(record.events().len(), 2);
641        assert_eq!(record.usage().unwrap().total(), Some(5));
642    }
643
644    #[test]
645    fn last_result_is_none_before_any_turn() {
646        let session = Session::new(test_codex());
647        assert!(session.last_result().is_none());
648        assert_eq!(session.total_tokens(), 0);
649        assert_eq!(session.turns_missing_usage(), 0);
650    }
651
652    #[tokio::test]
653    #[cfg(unix)]
654    async fn stream_delivers_events_and_records_the_turn() {
655        let mut session = streaming_session();
656        let mut seen = Vec::new();
657
658        let events = session
659            .stream("test prompt", |event| seen.push(event.event_type.clone()))
660            .await
661            .unwrap();
662
663        assert!(
664            seen.contains(&"turn.completed".to_string()),
665            "saw: {seen:?}"
666        );
667        assert_eq!(events.len(), seen.len(), "handler and return value agree");
668
669        // The streaming path must leave the same state a buffered turn would.
670        assert_eq!(session.total_turns(), 1);
671        assert_eq!(session.id(), Some("thread_test"));
672        assert_eq!(session.total_tokens(), 165);
673        assert_eq!(session.last_result().unwrap().result, "hello");
674    }
675
676    #[tokio::test]
677    #[cfg(unix)]
678    async fn streaming_turns_accumulate_like_buffered_ones() {
679        let mut session = streaming_session();
680        session.stream("first", |_| {}).await.unwrap();
681        // Second turn routes through exec resume, since thread_id is now set.
682        session.stream("second", |_| {}).await.unwrap();
683
684        assert_eq!(session.total_turns(), 2);
685        assert_eq!(session.total_tokens(), 330);
686        assert_eq!(session.turns_missing_usage(), 0);
687    }
688
689    // -----------------------------------------------------------------
690    // Budget (#60)
691    // -----------------------------------------------------------------
692
693    #[test]
694    fn a_session_without_a_budget_is_unaffected() {
695        let mut session = Session::new(test_codex());
696        session.record_turn(completed(500));
697        assert!(session.budget().is_none());
698        assert_eq!(session.total_tokens(), 500);
699    }
700
701    #[test]
702    fn turns_accumulate_into_the_attached_budget() {
703        let budget = crate::budget::TokenBudget::builder()
704            .max_tokens(1000)
705            .build();
706        let mut session = Session::new(test_codex()).with_budget(budget.clone());
707
708        session.record_turn(completed(300));
709        session.record_turn(completed(250));
710
711        assert_eq!(budget.total_tokens(), 550);
712        assert_eq!(session.total_tokens(), 550);
713    }
714
715    /// The ceiling stops the next turn rather than being noticed after more
716    /// has already been spent.
717    #[tokio::test]
718    #[cfg(unix)]
719    async fn a_turn_past_the_ceiling_is_refused_before_it_runs() {
720        let budget = crate::budget::TokenBudget::builder()
721            .max_tokens(100)
722            .build();
723        let mut session = streaming_session().with_budget(budget.clone());
724
725        // First turn runs and puts the budget over its ceiling: the fixture
726        // reports 165 tokens.
727        session.send("first").await.unwrap();
728        assert!(budget.total_tokens() >= 100);
729
730        let refused = session.send("second").await;
731        assert!(
732            matches!(refused, Err(Error::TokenBudgetExceeded { .. })),
733            "expected the second turn to be refused, got: {refused:?}"
734        );
735        // Refused before running, so no turn was recorded for it.
736        assert_eq!(session.total_turns(), 1);
737    }
738
739    /// A review reports usage as all zeros (#79), so a session of reviews
740    /// never advances a budget. Pinned rather than left as a surprise.
741    #[test]
742    fn an_all_zero_usage_turn_does_not_advance_the_budget() {
743        let budget = crate::budget::TokenBudget::builder()
744            .max_tokens(100)
745            .build();
746        let mut session = Session::new(test_codex()).with_budget(budget.clone());
747
748        for _ in 0..10 {
749            session.record_turn(completed(0));
750        }
751
752        assert_eq!(budget.total_tokens(), 0);
753        assert_eq!(
754            budget.turns_missing_usage(),
755            0,
756            "zero is reported, not missing"
757        );
758        assert!(budget.check().is_ok(), "a review-only session never stops");
759    }
760
761    /// A turn with no usage object must not read as zero consumption, which
762    /// would let an unmeasured session run forever against its ceiling.
763    #[test]
764    fn a_turn_with_no_usage_is_counted_as_unmeasured() {
765        let budget = crate::budget::TokenBudget::builder()
766            .max_tokens(100)
767            .build();
768        let mut session = Session::new(test_codex()).with_budget(budget.clone());
769
770        session.record_turn(turn(&[r#"{"type":"turn.completed"}"#]));
771
772        assert_eq!(budget.total_tokens(), 0);
773        assert_eq!(budget.turns_missing_usage(), 1);
774        assert_eq!(session.turns_missing_usage(), 1);
775    }
776}
codex_wrapper/session.rs

codex_wrapper/
session.rs
