Delete written managed tool replies so the reply limit counts only live calls

workspaces/runtime/packages/runtime-services/src/account_worker.rs

@@ -3406,6 +3406,7 @@ impl Worker {
             let monitor_dispatch_wake = self.bridge.monitor_dispatch_wake();
             let store = Arc::clone(&self.bridge.store);
             let managed_tool_runtime = self.bridge.managed_tool_runtime();
+            let journal_scope = self.journal_scope(actor_id);
             self.managed_callback_cursors
                 .insert((actor.clone(), actor_epoch, cursor.0));
             self.managed_callbacks.spawn(async move {
@@ -3440,6 +3441,10 @@ impl Worker {
                         wake,
                         monitor_dispatch_wake,
                         managed_runtime: managed_tool_runtime,
+                        journal: callback::CallbackJournal {
+                            scope: journal_scope,
+                            cursor: cursor.0,
+                        },
                     },
                 )
                 .await;

workspaces/runtime/packages/runtime-services/src/account_worker/callback.rs

@@ -13,6 +13,12 @@ pub(super) struct ManagedResponseDependencies {
     pub wake: Arc<Notify>,
     pub monitor_dispatch_wake: Arc<Notify>,
     pub managed_runtime: Option<Arc<crate::provider_callbacks::ManagedToolRuntime>>,
+    /// The journal position whose callback write answers this tool call.
+    pub journal: CallbackJournal,
+}
+pub(super) struct CallbackJournal {
+    pub scope: String,
+    pub cursor: u64,
 }
 pub fn dispatch(_request: &Frame) -> CallbackDecision {
     CallbackDecision::Unsupported
@@ -185,7 +191,27 @@ pub(super) async fn managed_response(
             )
             .await
             {
-                Ok(crate::provider_callbacks::CallbackDecision::Answer(reply)) => reply.body,
+                Ok(crate::provider_callbacks::CallbackDecision::Answer(reply)) => {
+                    let key = crate::provider_callbacks::reply_key(&reply.identity);
+                    let journal = &dependencies.journal;
+                    // The answer still goes out without the link. Its prepared
+                    // row then waits for the settled-turn sweep.
+                    if let Err(error) = super::store::link_callback_reply(
+                        store,
+                        &journal.scope,
+                        generation,
+                        journal.cursor,
+                        key,
+                    )
+                    .await
+                    {
+                        eprintln!(
+                            "managed tool reply was not linked to its callback; account={account} generation={generation} cursor={}: {error}",
+                            journal.cursor
+                        );
+                    }
+                    reply.body
+                }
                 Ok(
                     crate::provider_callbacks::CallbackDecision::Pending(_)
                     | crate::provider_callbacks::CallbackDecision::Settled(_),

workspaces/runtime/packages/runtime-services/src/account_worker/store.rs

@@ -2150,9 +2150,61 @@ fn settle_callback_in_transaction(
         row["ready"] = json!(true);
         replace(tx, &operation, "settled", row)?;
     }
+    // Every caller holds this cursor's written callback proof.
+    let callback = callback_operation(journal_scope, generation, cursor);
+    if let Some(reply) = receipt(tx, &callback, CALLBACK_REPLY_PHASE)? {
+        let key = reply.as_str().ok_or_else(|| invalid(&callback))?;
+        crate::provider_callbacks::delete_worker_written_reply(tx, key)
+            .map_err(|error| callback_store_error(error, &callback))?;
+        delete_receipt(tx, &callback, CALLBACK_REPLY_PHASE)?;
+    }
     compact(tx, journal_scope, generation)
 }
 
+/// The receipt phase that names the provider callback reply answering a
+/// callback cursor (value: that reply's receipt key).
+const CALLBACK_REPLY_PHASE: &str = "reply";
+
+fn callback_operation(journal_scope: &str, generation: u64, cursor: u64) -> String {
+    format!("callback:{journal_scope}:{generation}:{cursor}")
+}
+
+fn callback_store_error(
+    error: crate::provider_callbacks::CallbackError,
+    operation: &str,
+) -> WorkerError {
+    match error {
+        crate::provider_callbacks::CallbackError::Store(error) => error.into(),
+        _ => invalid(operation),
+    }
+}
+
+/// Link a managed tool reply to the callback cursor that writes it, so the
+/// written proof of that cursor deletes the reply. A cursor that already has
+/// its proof deletes the reply at once; the worker then only acknowledges it.
+pub(super) async fn link_callback_reply(
+    store: &RuntimeStore,
+    journal_scope: &str,
+    generation: u64,
+    cursor: u64,
+    reply_key: String,
+) -> Result<(), WorkerError> {
+    let journal_scope = journal_scope.to_owned();
+    store
+        .transaction_with(move |tx| {
+            let callback = callback_operation(&journal_scope, generation, cursor);
+            if cursor <= checkpoint(tx, &journal_scope, generation)?
+                || receipt(tx, &callback, "written")? == Some(json!(true))
+            {
+                crate::provider_callbacks::delete_worker_written_reply(tx, &reply_key)
+                    .map_err(|error| callback_store_error(error, &callback))?;
+                return Ok(());
+            }
+            replace(tx, &callback, CALLBACK_REPLY_PHASE, json!(reply_key))
+        })
+        .await
+}
+
 // Python codex_streaming.FLUSH_SECONDS is .75. The additional byte threshold
 // bounds a durable batch when producers outpace that timer; answers are uncut.
 pub(super) const STREAM_FLUSH_BYTES: usize = 64 * 1024;

workspaces/runtime/packages/runtime-services/src/account_worker/tests.rs

@@ -17,6 +17,7 @@ mod stale_disconnect_recovery;
 mod start_throttle;
 mod steer;
 mod submission;
+mod tool_reply_retirement;
 use super::*;
 use crate::event_write::{AnnotationWrite, annotate};
 use std::{

workspaces/runtime/packages/runtime-services/src/account_worker/tests/tool_reply_retirement.rs

@@ -0,0 +1,203 @@
+//! A managed tool reply that the worker wrote leaves no outstanding provider
+//! callback row, so the account's reply limit counts only real in-flight calls.
+use super::*;
+
+/// The account's reply limit (provider_callbacks/model.rs MAX_PENDING).
+const REPLY_LIMIT: usize = 128;
+
+struct ImmediateStatus;
+
+impl crate::provider_callbacks::tools::StatusWatchPort for ImmediateStatus {
+    fn status<'a>(
+        &'a self,
+        _store: &'a RuntimeStore,
+        _actor: &'a crate::managed_host_tools::Actor,
+        _args: &'a Value,
+        _now: f64,
+    ) -> futures_util::future::BoxFuture<'a, Result<Value, crate::managed_host_tools::Error>> {
+        Box::pin(async { Ok(json!({"status":"answered"})) })
+    }
+
+    fn rules<'a>(
+        &'a self,
+        _store: &'a RuntimeStore,
+        _actor: &'a crate::managed_host_tools::Actor,
+        _args: &'a Value,
+        _request_id: &'a str,
+        _now: f64,
+    ) -> futures_util::future::BoxFuture<
+        'a,
+        Result<
+            crate::provider_callbacks::tools::RulesActionResult,
+            crate::managed_host_tools::Error,
+        >,
+    > {
+        Box::pin(async {
+            Ok(crate::provider_callbacks::tools::RulesActionResult {
+                value: Value::Null,
+                cancel_monitor_ids: Vec::new(),
+            })
+        })
+    }
+}
+
+/// (outstanding replies, all provider callback rows, reply links)
+async fn reply_rows(state: &State) -> (i64, i64, i64) {
+    state
+        .bridge
+        .run(|store| async move {
+            let connection = store.read()?;
+            let count = |sql: &str| connection.query_row(sql, [], |row| row.get::<_, i64>(0));
+            Ok((
+                count(
+                    "SELECT count(*) FROM runtime_operation_receipts WHERE signature='provider-callback' AND json_extract(result,'$.delivery') IN ('prepared','submitted','unknown')",
+                )?,
+                count(
+                    "SELECT count(*) FROM runtime_operation_receipts WHERE signature='provider-callback'",
+                )?,
+                count(
+                    "SELECT count(*) FROM runtime_operation_receipts WHERE id LIKE 'account-worker:callback:%:reply'",
+                )?,
+            ))
+        })
+        .await
+        .unwrap()
+}
+
+async fn running_tool_turn(state: &State, peer: &mut Peer) -> AccountHandle {
+    let worker = peer.worker(state.bridge.clone()).await;
+    let start = state.start(peer, &worker).await;
+    peer.accepted(&start).await;
+    let agent = state.agent.clone();
+    state
+        .bridge
+        .run(move |store| async move {
+            store.transaction_with::<_, StoreError>(move |tx| {
+                tx.execute(
+                    "UPDATE runtime_agents SET record=json_set(record,'$.startAttempt.connectionId','1') WHERE id=?1",
+                    [agent],
+                )?;
+                Ok(())
+            }).await?;
+            Ok(())
+        })
+        .await
+        .unwrap();
+    state.bridge.set_managed_tool_runtime(
+        crate::provider_callbacks::ManagedToolRuntime::from_ports_for_test(
+            crate::provider_callbacks::ManagedToolRuntimePorts {
+                status_watch: Some(Arc::new(ImmediateStatus)),
+                child_creation: None,
+                monitor_create: None,
+                interrupt: None,
+                monitor_input: None,
+                monitor_cancel: None,
+                work_action: None,
+                late: Default::default(),
+            },
+        ),
+    );
+    worker
+}
+
+async fn answered_tool_call(peer: &mut Peer, index: usize) {
+    let id = format!("tool-call-{index}");
+    let cursor = peer
+        .frame(json!({
+            "id":id,
+            "method":"item/tool/call",
+            "params":{
+                "threadId":"fake-thread-0",
+                "turnId":"turn",
+                "callId":format!("status-{index}"),
+                "tool":"orchestration_status",
+                "arguments":{}
+            }
+        }))
+        .await;
+    let callback = peer.next_write().await;
+    assert_eq!(callback.value["id"], id.as_str());
+    assert_eq!(
+        callback.value["result"]["success"], true,
+        "tool call {index} must succeed: {}",
+        callback.value
+    );
+    callback
+        .receipt
+        .send(CommandReceipt::Written {
+            remote_id: None,
+            duplicate: false,
+        })
+        .ok()
+        .expect("callback write receipt receiver");
+    peer.acknowledged(cursor).await;
+}
+
+#[tokio::test]
+async fn written_tool_replies_leave_no_outstanding_reply_past_the_account_limit() {
+    let state = State::new().await;
+    let mut peer = Peer::new();
+    let mut worker = running_tool_turn(&state, &mut peer).await;
+    // The call after the limit must succeed: no written reply stays outstanding.
+    for index in 0..=REPLY_LIMIT {
+        answered_tool_call(&mut peer, index).await;
+    }
+    assert_eq!(
+        reply_rows(&state).await,
+        (0, 0, 0),
+        "written tool calls leave no reply row or link"
+    );
+    worker.shutdown().await.unwrap();
+}
+
+#[tokio::test]
+async fn leaked_replies_of_a_settled_turn_do_not_block_a_new_tool_call() {
+    let state = State::new().await;
+    let mut peer = Peer::new();
+    let mut worker = running_tool_turn(&state, &mut peer).await;
+    let agent = state.agent.clone();
+    state
+        .bridge
+        .run(move |store| async move {
+            store
+                .transaction_with::<_, StoreError>(move |tx| {
+                    for index in 0..REPLY_LIMIT {
+                        let identity = crate::provider_callbacks::ReplyIdentity {
+                            context: crate::provider_callbacks::CallbackContext {
+                                agent: agent.clone(),
+                                account: "default".into(),
+                                connection: "1".into(),
+                                generation: 1,
+                                epoch: 0,
+                                thread: "fake-thread-0".into(),
+                                turn: "earlier-turn".into(),
+                            },
+                            rpc_id: studio_runtime_providers::RequestId::Number(index as u64),
+                            attempt: format!("attempt-{index}"),
+                        };
+                        let value = json!({
+                            "identity":identity, "signature":"signature",
+                            "body":{"result":{"success":true,"contentItems":[]}},
+                            "delivery":"prepared", "toolRequest":format!("leaked-{index}"),
+                        });
+                        tx.execute(
+                            "INSERT INTO runtime_operation_receipts(id,signature,result) VALUES(?1,'provider-callback',?2)",
+                            [crate::provider_callbacks::reply_key(&identity), value.to_string()],
+                        )?;
+                    }
+                    Ok(())
+                })
+                .await?;
+            Ok(())
+        })
+        .await
+        .unwrap();
+    assert_eq!(
+        reply_rows(&state).await.0,
+        REPLY_LIMIT as i64,
+        "the account starts at its reply limit"
+    );
+    answered_tool_call(&mut peer, 0).await;
+    assert_eq!(reply_rows(&state).await, (0, 0, 0));
+    worker.shutdown().await.unwrap();
+}

workspaces/runtime/packages/runtime-services/src/provider_callbacks.rs

@@ -31,6 +31,7 @@ pub use storage::{
     DeliveryOutcome, bind_journal, delivery, prepared_replies, retire_checkpointed,
     submission_ready,
 };
+pub(crate) use storage::{delete_worker_written_reply, reply_key};
 use studio_runtime_providers::{Frame, ResponseBody};
 use studio_runtime_store::RuntimeStore;