@@ -22,6 +22,7 @@ use codex_protocol::models::ReasoningItemReasoningSummary;
22
22
use codex_protocol:: models:: WebSearchAction ;
23
23
use core_test_support:: load_default_config_for_test;
24
24
use core_test_support:: load_sse_fixture_with_id;
25
+ use core_test_support:: responses;
25
26
use core_test_support:: wait_for_event;
26
27
use futures:: StreamExt ;
27
28
use serde_json:: json;
@@ -776,6 +777,100 @@ async fn azure_responses_request_includes_store_and_reasoning_ids() {
776
777
assert_eq ! ( body[ "input" ] [ 5 ] [ "id" ] . as_str( ) , Some ( "custom-tool-id" ) ) ;
777
778
}
778
779
780
+ #[ tokio:: test( flavor = "multi_thread" , worker_threads = 2 ) ]
781
+ async fn token_count_includes_rate_limits_snapshot ( ) {
782
+ let server = MockServer :: start ( ) . await ;
783
+
784
+ let sse_body = responses:: sse ( vec ! [ responses:: ev_completed_with_tokens( "resp_rate" , 123 ) ] ) ;
785
+
786
+ let response = ResponseTemplate :: new ( 200 )
787
+ . insert_header ( "content-type" , "text/event-stream" )
788
+ . insert_header ( "x-codex-primary-used-percent" , "12.5" )
789
+ . insert_header ( "x-codex-protection-used-percent" , "40.0" )
790
+ . insert_header ( "x-codex-primary-over-protection-limit-percent" , "75.0" )
791
+ . insert_header ( "x-codex-primary-window-minutes" , "10" )
792
+ . insert_header ( "x-codex-protection-window-minutes" , "60" )
793
+ . set_body_raw ( sse_body, "text/event-stream" ) ;
794
+
795
+ Mock :: given ( method ( "POST" ) )
796
+ . and ( path ( "/v1/responses" ) )
797
+ . respond_with ( response)
798
+ . expect ( 1 )
799
+ . mount ( & server)
800
+ . await ;
801
+
802
+ let mut provider = built_in_model_providers ( ) [ "openai" ] . clone ( ) ;
803
+ provider. base_url = Some ( format ! ( "{}/v1" , server. uri( ) ) ) ;
804
+
805
+ let home = TempDir :: new ( ) . unwrap ( ) ;
806
+ let mut config = load_default_config_for_test ( & home) ;
807
+ config. model_provider = provider;
808
+
809
+ let conversation_manager = ConversationManager :: with_auth ( CodexAuth :: from_api_key ( "test" ) ) ;
810
+ let codex = conversation_manager
811
+ . new_conversation ( config)
812
+ . await
813
+ . expect ( "create conversation" )
814
+ . conversation ;
815
+
816
+ codex
817
+ . submit ( Op :: UserInput {
818
+ items : vec ! [ InputItem :: Text {
819
+ text: "hello" . into( ) ,
820
+ } ] ,
821
+ } )
822
+ . await
823
+ . unwrap ( ) ;
824
+
825
+ let token_event = wait_for_event ( & codex, |msg| matches ! ( msg, EventMsg :: TokenCount ( _) ) ) . await ;
826
+ let final_payload = match token_event {
827
+ EventMsg :: TokenCount ( ev) => ev,
828
+ _ => unreachable ! ( ) ,
829
+ } ;
830
+ // Assert full JSON for the final token count event (usage + rate limits)
831
+ let final_json = serde_json:: to_value ( & final_payload) . unwrap ( ) ;
832
+ pretty_assertions:: assert_eq!(
833
+ final_json,
834
+ json!( {
835
+ "info" : {
836
+ "total_token_usage" : {
837
+ "input_tokens" : 123 ,
838
+ "cached_input_tokens" : 0 ,
839
+ "output_tokens" : 0 ,
840
+ "reasoning_output_tokens" : 0 ,
841
+ "total_tokens" : 123
842
+ } ,
843
+ "last_token_usage" : {
844
+ "input_tokens" : 123 ,
845
+ "cached_input_tokens" : 0 ,
846
+ "output_tokens" : 0 ,
847
+ "reasoning_output_tokens" : 0 ,
848
+ "total_tokens" : 123
849
+ } ,
850
+ // Default model is gpt-5 in tests → 272000 context window
851
+ "model_context_window" : 272000
852
+ } ,
853
+ "rate_limits" : {
854
+ "primary_used_percent" : 12.5 ,
855
+ "weekly_used_percent" : 40.0 ,
856
+ "primary_to_weekly_ratio_percent" : 75.0 ,
857
+ "primary_window_minutes" : 10 ,
858
+ "weekly_window_minutes" : 60
859
+ }
860
+ } )
861
+ ) ;
862
+ let usage = final_payload
863
+ . info
864
+ . expect ( "token usage info should be recorded after completion" ) ;
865
+ assert_eq ! ( usage. total_token_usage. total_tokens, 123 ) ;
866
+ let final_snapshot = final_payload
867
+ . rate_limits
868
+ . expect ( "latest rate limit snapshot should be retained" ) ;
869
+ assert_eq ! ( final_snapshot. primary_used_percent, 12.5 ) ;
870
+
871
+ wait_for_event ( & codex, |msg| matches ! ( msg, EventMsg :: TaskComplete ( _) ) ) . await ;
872
+ }
873
+
779
874
#[ tokio:: test( flavor = "multi_thread" , worker_threads = 2 ) ]
780
875
async fn azure_overrides_assign_properties_used_for_responses_url ( ) {
781
876
let existing_env_var_with_random_value = if cfg ! ( windows) { "USERNAME" } else { "USER" } ;
0 commit comments