Persist provider profile dispatch leases

codex-rs/app-server/src/bespoke_event_handling.rs

@@ -2844,6 +2844,7 @@ mod tests {
                 schedule_id: "schedule-1".to_string(),
                 run_id: "run-1".to_string(),
                 lease_id: "lease-1".to_string(),
+                usage_profile_lease: None,
                 goal_id: None,
                 state_db,
             },

codex-rs/app-server/src/request_processors/thread_lifecycle.rs

@@ -325,20 +325,31 @@ pub(super) async fn ensure_listener_task_running(
                         event.msg,
                         EventMsg::TurnComplete(_) | EventMsg::TurnAborted(_) | EventMsg::Error(_)
                     );
-                    let (raw_events_enabled, tracked_scheduled_run, turn_error) = {
+                    let (
+                        raw_events_enabled,
+                        tracked_scheduled_run,
+                        tracked_usage_profile_lease,
+                        turn_error,
+                    ) = {
                         let mut thread_state = thread_state.lock().await;
                         thread_state.track_current_turn_event(&event.id, &event.msg);
                         let tracked_scheduled_run = if terminal_event {
                             thread_state.take_scheduled_run(&event.id)
                         } else {
                             None
                         };
+                        let tracked_usage_profile_lease = if terminal_event {
+                            thread_state.take_usage_profile_turn_lease(&event.id)
+                        } else {
+                            None
+                        };
                         let turn_error = terminal_event
                             .then(|| thread_state.turn_summary.last_error.clone())
                             .flatten();
                         (
                             thread_state.experimental_raw_events,
                             tracked_scheduled_run,
+                            tracked_usage_profile_lease,
                             turn_error,
                         )
                     };
@@ -418,6 +429,9 @@ pub(super) async fn ensure_listener_task_running(
                         )
                         .await;
                     }
+                    if let Some(usage_profile_lease) = tracked_usage_profile_lease {
+                        finish_usage_profile_turn_lease(usage_profile_lease).await;
+                    }
                 }
                 unloading_watchers_open = unloading_state.wait_for_unloading_trigger() => {
                     if !unloading_watchers_open {
@@ -455,15 +469,33 @@ pub(super) async fn ensure_listener_task_running(
             }
         }
 
-        let mut thread_state = thread_state.lock().await;
-        if thread_state.listener_generation == listener_generation {
-            thread_state_manager.unregister_listener_command_tx(conversation_id);
-            thread_state.clear_listener();
+        let abandoned_usage_profile_leases = {
+            let mut thread_state = thread_state.lock().await;
+            if thread_state.listener_generation == listener_generation {
+                thread_state_manager.unregister_listener_command_tx(conversation_id);
+                let leases = thread_state.drain_usage_profile_turn_leases();
+                thread_state.clear_listener();
+                leases
+            } else {
+                Vec::new()
+            }
+        };
+        for lease in abandoned_usage_profile_leases {
+            finish_usage_profile_turn_lease(lease).await;
         }
     });
     Ok(())
 }
 
+async fn finish_usage_profile_turn_lease(tracked: crate::thread_state::UsageProfileTurnLease) {
+    tracked.heartbeat_cancel.cancel();
+    super::usage_profile_broker::release_dispatch_auth_profile_lease(
+        &tracked.state_db,
+        &tracked.lease,
+    )
+    .await;
+}
+
 async fn heartbeat_local_active_session(
     listener_task_context: &ListenerTaskContext,
     conversation_id: ThreadId,