Skip to main content

rpi_agent/
tool_program.rs

1//! Host-owned, invocation-local nested execution. No model turns or transcript events.
2use crate::tool_program_observation::{
3    self, ToolProgramChild, ToolProgramObservation, ToolProgramObserver,
4};
5use crate::{
6    AgentContext, AgentEmitter, AgentEvent, AgentLoopConfig, TextContentOrImage, ToolExecutionMode,
7};
8use futures::{future::BoxFuture, FutureExt};
9use rpi_ai::types::{AssistantMessage, ToolCall, ToolCallType};
10use serde_json::{json, Value};
11use std::{
12    collections::HashMap,
13    future::Future,
14    sync::{
15        atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
16        Arc, Mutex, Weak,
17    },
18    time::Duration,
19};
20use tokio::sync::{Notify, RwLock, Semaphore};
21use tokio_util::sync::CancellationToken;
22
23#[path = "tool_program_accounting.rs"]
24mod accounting;
25#[path = "tool_program_results.rs"]
26mod results;
27
28/// Native content composition. References select only content, never child metadata.
29#[derive(Clone, Debug)]
30pub enum ToolProgramContentPart {
31    /// Adapter-validated owned blocks; plugin frames remain limited to 1 MiB.
32    Inline(Vec<TextContentOrImage>),
33    /// Invocation-owned native content. None selects all; an empty selection
34    /// selects none. Ordered indices may repeat, but must remain in bounds.
35    Reference {
36        id: String,
37        indices: Option<Vec<usize>>,
38    },
39}
40
41const MAX_BYTES: usize = 1024 * 1024;
42static NEXT: AtomicU64 = AtomicU64::new(1);
43tokio::task_local! { static CURRENT: Arc<ToolProgram>; static PARENT: NestedParent; }
44
45/// Identity only the native execution driver can install; JSON is never trusted.
46#[derive(Clone)]
47pub struct NestedParent {
48    pub root_id: String,
49    pub root_name: String,
50    pub parent_id: String,
51}
52pub fn nested_parent() -> Option<NestedParent> {
53    PARENT.try_with(Clone::clone).ok()
54}
55pub fn current() -> Option<Arc<ToolProgram>> {
56    CURRENT.try_with(Clone::clone).ok()
57}
58pub(crate) async fn scope<F: Future>(capability: Arc<ToolProgram>, future: F) -> F::Output {
59    CURRENT.scope(capability, future).await
60}
61pub(crate) struct InvocationLifetime(pub CancellationToken);
62impl Drop for InvocationLifetime {
63    fn drop(&mut self) {
64        self.0.cancel();
65    }
66}
67
68struct Request {
69    signal: CancellationToken,
70    result: Option<Value>,
71}
72#[derive(Default)]
73struct State {
74    requests: HashMap<String, Request>,
75    accounting: accounting::Accounting,
76    results: results::Store,
77}
78/// Captured from the effective executing loop, never from a root registry.
79pub struct ToolProgram {
80    context: Arc<AgentContext>,
81    assistant: Arc<AssistantMessage>,
82    config: AgentLoopConfig,
83    ancestors: Vec<String>,
84    identity: NestedParent,
85    child_scope: Option<Arc<crate::tool_execution_scope::ToolExecutionScope>>,
86    observer: Option<Arc<dyn ToolProgramObserver>>,
87    starts: Arc<AtomicUsize>,
88    state: Mutex<State>,
89    serial: RwLock<()>,
90    active: Arc<Semaphore>,
91    tree_active: Arc<Semaphore>,
92    result_bytes: Arc<Semaphore>,
93    ordered: bool,
94    terminated: AtomicBool,
95    progress: Arc<Progress>,
96    pub(crate) signal: CancellationToken,
97}
98impl ToolProgram {
99    pub(crate) fn capture(
100        context: &Arc<AgentContext>,
101        assistant: &Arc<AssistantMessage>,
102        call: &ToolCall,
103        config: &AgentLoopConfig,
104        signal: CancellationToken,
105        on_update: Arc<dyn Fn(crate::ToolResultPartial) + Send + Sync>,
106        emit: &Arc<dyn AgentEmitter>,
107    ) -> Arc<Self> {
108        // Closing this composition must not cancel its containing native tool:
109        // wrappers may still need that tool's signal for result postprocessing.
110        // Caller cancellation still propagates down into the program and children.
111        let signal = signal.child_token();
112        let previous = current();
113        let mut ancestors = previous
114            .as_ref()
115            .map(|p| p.ancestors.clone())
116            .unwrap_or_default();
117        ancestors.push(call.name.clone());
118        let identity = previous
119            .as_ref()
120            .map(|p| NestedParent {
121                root_id: p.identity.root_id.clone(),
122                root_name: p.identity.root_name.clone(),
123                parent_id: call.id.clone(),
124            })
125            .unwrap_or_else(|| NestedParent {
126                root_id: call.id.clone(),
127                root_name: call.name.clone(),
128                parent_id: call.id.clone(),
129            });
130        Arc::new(Self {
131            context: context.clone(),
132            assistant: assistant.clone(),
133            config: config.clone(),
134            ancestors,
135            identity,
136            child_scope: crate::tool_execution_scope::current_scope(),
137            // Nested execution inherits only the actual containing capture.
138            observer: match &previous {
139                Some(parent) => parent.observer.clone(),
140                None => emit.tool_program_observer(),
141            },
142            starts: previous
143                .as_ref()
144                .map(|p| p.starts.clone())
145                .unwrap_or_default(),
146            tree_active: previous
147                .as_ref()
148                .map(|p| p.tree_active.clone())
149                .unwrap_or_else(|| Arc::new(Semaphore::new(8))),
150            result_bytes: previous
151                .as_ref()
152                .map(|p| p.result_bytes.clone())
153                .unwrap_or_else(|| Arc::new(Semaphore::new(results::TREE_BYTES))),
154            ordered: config.tool_execution == ToolExecutionMode::Sequential
155                || context.tools.iter().any(|t| {
156                    t.schema().name == call.name
157                        && t.execution_mode() == ToolExecutionMode::Sequential
158                }),
159            state: Mutex::default(),
160            serial: RwLock::new(()),
161            active: Arc::new(Semaphore::new(8)),
162            terminated: AtomicBool::new(false),
163            progress: previous
164                .as_ref()
165                .map(|p| p.progress.clone())
166                .unwrap_or_else(|| {
167                    Arc::new(Progress {
168                        callback: on_update,
169                        signal: signal.clone(),
170                        budget: Mutex::default(),
171                    })
172                }),
173            signal,
174        })
175    }
176    /// Effective names, including this tool, for the existing names-only action.
177    pub fn active_tool_names(&self) -> Vec<String> {
178        self.context
179            .tools
180            .iter()
181            .map(|t| t.schema().name.clone())
182            .collect()
183    }
184    /// Cancellation stops admission. Cleanup/poll remain available to the binding.
185    pub fn cancel_all(&self) {
186        self.signal.cancel();
187    }
188    /// Observe driver-local cancellation, without retaining a finished invocation.
189    pub fn bind_signal(self: &Arc<Self>, signal: CancellationToken) {
190        let weak = Arc::downgrade(self);
191        let finished = self.signal.clone();
192        tokio::spawn(async move {
193            tokio::select! {
194                biased;
195                _ = signal.cancelled() => {
196                    if let Some(program) = weak.upgrade() { program.cancel_all(); }
197                }
198                _ = finished.cancelled() => {}
199            }
200        });
201    }
202    pub(crate) async fn finish(&self) {
203        self.cancel_all();
204        // An outer Sequential call cannot finish ahead of its descendants.
205        let _idle = self.active.acquire_many(8).await;
206        self.state.lock().unwrap().results.clear();
207    }
208    pub(crate) fn merge_accounting(
209        &self,
210        result: &mut crate::AgentToolResult,
211    ) -> Result<(), String> {
212        self.state.lock().unwrap().accounting.merge(result)
213    }
214    pub fn handle(self: &Arc<Self>, args: Value) -> BoxFuture<'static, Result<Value, String>> {
215        let this = self.clone();
216        Box::pin(async move { this.handle_inner(args).await })
217    }
218    async fn handle_inner(self: &Arc<Self>, args: Value) -> Result<Value, String> {
219        bounded(&args)?;
220        let op = args["op"].as_str().ok_or("op is required")?;
221        if matches!(op, "list" | "start")
222            && (self.signal.is_cancelled() || self.terminated.load(Ordering::SeqCst))
223        {
224            return Err("parent invocation is no longer active".into());
225        }
226        match op {
227            "list" => {
228                let offset = integer(&args, "offset", 0, 0, 4096)? as usize;
229                let limit = integer(&args, "limit", 32, 1, 128)? as usize;
230                let tools: Vec<_> = self
231                    .context
232                    .tools
233                    .iter()
234                    .filter(|t| !self.ancestors.contains(&t.schema().name))
235                    .collect();
236                if tools.len() > 4096 {
237                    return Err("tool catalog exceeds 4096 entries".into());
238                }
239                // Bound the complete catalog as well as each page.
240                let all:Vec<_> = tools.iter().map(|t|json!({"name":t.schema().name,"description":t.schema().description,"parameters":t.schema().parameters})).collect();
241                bounded(&json!(all))?;
242                let end = offset.saturating_add(limit).min(all.len());
243                let page = all.get(offset..end).unwrap_or(&[]);
244                Ok(
245                    json!({"protocolVersion":1,"capabilities":{"resultEnvelope":1,"hostAccounting":1,"parentProgress":1,"losslessValues":1,"resultReferences":1,"contentAttachments":1},"tools":page,"nextOffset":if end<all.len(){Some(end)}else{None}}),
246                )
247            }
248            "read_result" => self.read_result(&args),
249            "measure_result_content" => self.measure_result_content(&args),
250            "release_result" => self.release_result(&args),
251            "start" => {
252                let reference = match args.get("resultMode") {
253                    None => false,
254                    Some(Value::String(mode)) if mode == "reference" => true,
255                    _ => return Err("resultMode must be reference when provided".into()),
256                };
257                // The executing parent cannot hold its children's gate. It
258                // must instead exclude peers in its own containing batch.
259                // Inductively, this also orders independent program groups.
260                if !self.ordered {
261                    return Err(
262                        "ToolProgram start requires a Sequential executing tool or loop".into(),
263                    );
264                }
265                let name = args["name"].as_str().ok_or("name is required")?.to_owned();
266                if self.ancestors.contains(&name) || self.ancestors.len() >= 16 {
267                    return Err("recursive/ancestor tool invocation is forbidden".into());
268                }
269                let tool = self
270                    .context
271                    .tools
272                    .iter()
273                    .find(|t| t.schema().name == name)
274                    .ok_or("tool is not in the effective active subset")?;
275                let sequential = self.config.tool_execution == ToolExecutionMode::Sequential
276                    || tool.execution_mode() == ToolExecutionMode::Sequential;
277                let arguments = args
278                    .get("arguments")
279                    .filter(|v| v.is_object())
280                    .ok_or("arguments must be an object")?
281                    .clone();
282                let timeout = integer(&args, "timeoutMs", 30000, 1, 120000)?;
283                let permit = self
284                    .active
285                    .clone()
286                    .try_acquire_owned()
287                    .map_err(|_| "ToolProgram active request limit exceeded")?;
288                let tree_permit = self
289                    .tree_active
290                    .clone()
291                    .try_acquire_owned()
292                    .map_err(|_| "ToolProgram invocation-tree active request limit exceeded")?;
293                let signal = self.signal.child_token();
294                let id = format!("tp-{}", NEXT.fetch_add(1, Ordering::Relaxed));
295                let ordinal;
296                {
297                    let mut state = self.state.lock().unwrap();
298                    if state.requests.len() >= 16 {
299                        return Err("ToolProgram retained/active request limit exceeded".into());
300                    }
301                    #[allow(deprecated)] // Keep the workspace Rust 1.88 MSRV.
302                    let previous = self
303                        .starts
304                        .fetch_update(Ordering::SeqCst, Ordering::SeqCst, |n| {
305                            (n < 128).then_some(n + 1)
306                        })
307                        .map_err(|_| "ToolProgram start budget exceeded")?;
308                    ordinal = previous + 1;
309                    if self.signal.is_cancelled() || self.terminated.load(Ordering::SeqCst) {
310                        return Err("parent invocation is no longer active".into());
311                    }
312                    state.requests.insert(
313                        id.clone(),
314                        Request {
315                            signal: signal.clone(),
316                            result: None,
317                        },
318                    );
319                }
320                let this = self.clone();
321                let request_id = id.clone();
322                let progress_name = name.clone();
323                let updates = Arc::new(AtomicUsize::new(0));
324                self.report(&id, &progress_name, "started", 0, &signal);
325                tokio::spawn(async move {
326                    let _permit = permit;
327                    let _tree_permit = tree_permit;
328                    // One deadline covers native work, callback retirement and writer close.
329                    let deadline = tokio::time::Instant::now() + Duration::from_millis(timeout);
330                    let update_gate = this
331                        .observer
332                        .as_ref()
333                        .map(|_| Arc::new(UpdateGate::default()));
334                    // Strong writer ownership survives the work future's unwind/drain.
335                    let mut observation: Option<Arc<dyn ToolProgramObservation>> = None;
336                    let observation_error = this
337                        .observer
338                        .as_ref()
339                        .map(|_| Arc::new(Mutex::new(None::<String>)));
340                    let publish = |result: Result<Value, String>| {
341                        this.report(
342                            &request_id,
343                            &progress_name,
344                            if result.is_ok() {
345                                if result.as_ref().is_ok_and(|r| r["isError"] == true) {
346                                    "failed"
347                                } else {
348                                    "completed"
349                                }
350                            } else {
351                                "failed"
352                            },
353                            updates.load(Ordering::Relaxed),
354                            &signal,
355                        );
356                        let reply = result.unwrap_or_else(
357                            |error| json!({"requestId":request_id,"status":"failed","error":error}),
358                        );
359                        if let Some(request) =
360                            this.state.lock().unwrap().requests.get_mut(&request_id)
361                        {
362                            request.result = Some(reply);
363                        }
364                    };
365                    let (mut result, mut published) = {
366                        let work = async {
367                            // The parent never takes this lock: Sequential programs cannot self-deadlock.
368                            let _write;
369                            let _read;
370                            if sequential {
371                                _write = Some(this.serial.write().await);
372                                _read = None;
373                            } else {
374                                _read = Some(this.serial.read().await);
375                                _write = None;
376                            }
377                            if signal.is_cancelled() {
378                                return Err("nested invocation cancelled".to_owned());
379                            }
380                            if let Some(observer) = &this.observer {
381                                let child = ToolProgramChild {
382                                    root_id: this.identity.root_id.clone(),
383                                    root_name: this.identity.root_name.clone(),
384                                    parent_id: this.identity.parent_id.clone(),
385                                    request_id: request_id.clone(),
386                                    name: name.clone(),
387                                    ordinal,
388                                };
389                                let writer = observer
390                                    .begin(child, &arguments, signal.clone())
391                                    .await
392                                    .inspect_err(|_| this.cancel_all())?;
393                                observation = Some(Arc::new(NativeObservation {
394                                    writer,
395                                    program: Arc::downgrade(&this),
396                                    error: observation_error.as_ref().unwrap().clone(),
397                                }));
398                                if signal.is_cancelled() {
399                                    return Err("nested invocation cancelled".into());
400                                }
401                            }
402                            let call = ToolCall {
403                                kind: ToolCallType,
404                                id: request_id.clone(),
405                                name,
406                                arguments,
407                                thought_signature: None,
408                                namespace: None,
409                                recorded: None,
410                            };
411                            let mut config = this.config.clone();
412                            config.signal = signal.clone();
413                            let emit: Arc<dyn AgentEmitter> = Arc::new(ProgressEmitter {
414                                program: Arc::downgrade(&this),
415                                request_id: request_id.clone(),
416                                name: progress_name.clone(),
417                                signal: signal.clone(),
418                                updates: updates.clone(),
419                                observation: observation.as_ref().map(Arc::downgrade),
420                                observation_error: observation_error.clone(),
421                                update_gate: update_gate.clone(),
422                            });
423                            let execute = PARENT.scope(
424                                this.identity.clone(),
425                                scope(
426                                    this.clone(),
427                                    crate::agent_loop::run_one_tool_call(
428                                        &this.context,
429                                        &this.assistant,
430                                        &call,
431                                        &config,
432                                        &emit,
433                                    ),
434                                ),
435                            );
436                            let execute = async {
437                                if let Some(scope) = &this.child_scope {
438                                    scope.run(execute).await
439                                } else {
440                                    execute.await
441                                }
442                            };
443                            let result = if let Some(observation) = &observation {
444                                tool_program_observation::scope(observation.clone(), execute).await
445                            } else {
446                                execute.await
447                            };
448                            let finalized = result.map_err(|e| e.to_string())?;
449                            let result = &finalized.result;
450                            // Exactly one retention per direct request, before any fallible frame encoding.
451                            {
452                                let mut state = this.state.lock().unwrap();
453                                state.accounting.retain(result);
454                                // Serialize terminal retention with start's last admission check.
455                                // Finalization, not fallible reply encoding, decides admission.
456                                if result.terminate {
457                                    this.terminated.store(true, Ordering::SeqCst);
458                                }
459                            }
460                            if let Some(gate) = &update_gate {
461                                gate.retire().await;
462                            }
463                            if let Some(observation) = &observation {
464                                observation
465                                    .finish(result, finalized.is_error, signal.is_cancelled())
466                                    .await
467                                    .inspect_err(|_| this.cancel_all())?;
468                            }
469                            if let Some(error) = observation_error
470                                .as_ref()
471                                .and_then(|e| e.lock().unwrap().clone())
472                            {
473                                return Err(error);
474                            }
475                            let reply = if reference {
476                                let encoded =
477                                    serde_json::to_string(result).map_err(|e| e.to_string())?;
478                                this.retain_result(
479                                    &request_id,
480                                    finalized.result,
481                                    finalized.is_error,
482                                    encoded,
483                                )?
484                            } else {
485                                let payload =
486                                    serde_json::to_value(result).map_err(|e| e.to_string())?;
487                                let reply = json!({"requestId":request_id,"status":"done","isError":finalized.is_error,"result":payload});
488                                bounded(&reply)?;
489                                reply
490                            };
491                            Ok::<_, String>(reply)
492                        };
493                        let work = std::panic::AssertUnwindSafe(work).catch_unwind();
494                        tokio::pin!(work);
495                        let (result, drain) = tokio::select! {
496                            biased;
497                            _ = signal.cancelled() => (Err("nested invocation cancelled".into()), true),
498                            _ = tokio::time::sleep_until(deadline) => (
499                                Err("nested invocation timed out; side effects may have occurred".into()), true),
500                            result = &mut work => (match result {
501                                Ok(reply) => reply,
502                                Err(_) => {
503                                    if this.observer.is_some() { this.cancel_all(); }
504                                    Err("nested invocation panicked".into())
505                                },
506                            }, false),
507                        };
508                        if drain {
509                            // Publish promptly, but retain native execution and permits.
510                            publish(result.clone());
511                            signal.cancel();
512                            if (&mut work).await.is_err() && this.observer.is_some() {
513                                // A writer may panic only after cancellation won
514                                // the select. It still fails program admission closed.
515                                this.cancel_all();
516                            }
517                        }
518                        (result, drain)
519                    };
520                    let reason = result.as_ref().err().cloned();
521                    // Own exactly one cleanup future. Cancellation/deadline may publish
522                    // failure, but never drop retirement/close or release either permit.
523                    let cleanup = async {
524                        if let Some(gate) = &update_gate {
525                            gate.retire().await;
526                        }
527                        // Native callbacks have retired, including their root reports.
528                        if let Some(observation) = observation {
529                            std::panic::AssertUnwindSafe(async {
530                                observation.close(reason.as_deref()).await
531                            })
532                            .catch_unwind()
533                            .await
534                            .unwrap_or_else(|_| Err("observation close panicked".into()))
535                        } else {
536                            Ok(())
537                        }
538                    };
539                    tokio::pin!(cleanup);
540                    let closed = if published {
541                        (&mut cleanup).await
542                    } else {
543                        let interruption = tokio::select! {
544                            biased;
545                            _ = signal.cancelled() => Some("nested invocation cancelled"),
546                            _ = tokio::time::sleep_until(deadline) => Some(
547                                "nested invocation timed out; side effects may have occurred"),
548                            closed = &mut cleanup => {
549                                if let Err(error) = closed {
550                                    this.cancel_all();
551                                    result = Err(error);
552                                }
553                                None
554                            }
555                        };
556                        if let Some(error) = interruption {
557                            publish(Err(error.into()));
558                            published = true;
559                            signal.cancel();
560                            (&mut cleanup).await
561                        } else {
562                            Ok(())
563                        }
564                    };
565                    if let Err(error) = closed {
566                        this.cancel_all();
567                        result = Err(error);
568                    }
569                    if !published {
570                        publish(result);
571                    }
572                    signal.cancel();
573                });
574                Ok(json!({"requestId":id,"status":"running"}))
575            }
576            "poll" | "cancel" | "dispose" => {
577                let id = args["requestId"].as_str().ok_or("requestId is required")?;
578                let mut state = self.state.lock().unwrap();
579                if op == "dispose" {
580                    if let Some(r) = state.requests.remove(id) {
581                        r.signal.cancel();
582                    }
583                    return Ok(json!({"requestId":id,"status":"disposed"}));
584                }
585                let request = state
586                    .requests
587                    .get_mut(id)
588                    .ok_or("unknown request for this invocation")?;
589                if op == "cancel" {
590                    request.signal.cancel();
591                    Ok(json!({"requestId":id,"status":"cancelled"}))
592                } else {
593                    Ok(request
594                        .result
595                        .clone()
596                        .unwrap_or_else(|| json!({"requestId":id,"status":"running"})))
597                }
598            }
599            _ => Err("unknown ToolProgram operation".into()),
600        }
601    }
602}
603fn bounded(v: &Value) -> Result<(), String> {
604    if serde_json::to_vec(v).map_err(|e| e.to_string())?.len() > MAX_BYTES {
605        Err("ToolProgram JSON exceeds 1 MiB".into())
606    } else {
607        Ok(())
608    }
609}
610fn integer(args: &Value, key: &str, default: u64, min: u64, max: u64) -> Result<u64, String> {
611    let value = match args.get(key) {
612        None => default,
613        Some(v) => v
614            .as_u64()
615            .ok_or_else(|| format!("{key} must be an integer"))?,
616    };
617    if !(min..=max).contains(&value) {
618        Err(format!("{key} must be between {min} and {max}"))
619    } else {
620        Ok(value)
621    }
622}
623// A prepared-barrier failure is an observation failure, not merely a tool
624// error: stop admission in the genuine containing program. No strong registry
625// reference is retained by the observation or by native callbacks.
626struct NativeObservation {
627    writer: Arc<dyn ToolProgramObservation>,
628    program: Weak<ToolProgram>,
629    error: Arc<Mutex<Option<String>>>,
630}
631impl ToolProgramObservation for NativeObservation {
632    fn id(&self) -> &str {
633        self.writer.id()
634    }
635    fn update(&self, partial: Arc<crate::AgentToolResult>) -> Result<(), String> {
636        self.writer.update(partial)
637    }
638    fn prepared<'a>(&'a self, arguments: &'a Value) -> BoxFuture<'a, Result<(), String>> {
639        Box::pin(async move {
640            // Construct the writer future inside the catch as well as polling it.
641            let result =
642                std::panic::AssertUnwindSafe(async { self.writer.prepared(arguments).await })
643                    .catch_unwind()
644                    .await
645                    .unwrap_or_else(|_| Err("observation prepared panicked".into()));
646            if let Err(error) = &result {
647                *self.error.lock().unwrap() = Some(error.clone());
648                if let Some(program) = self.program.upgrade() {
649                    program.cancel_all();
650                }
651            }
652            result
653        })
654    }
655    fn finish<'a>(
656        &'a self,
657        result: &'a crate::AgentToolResult,
658        is_error: bool,
659        cancelled: bool,
660    ) -> BoxFuture<'a, Result<(), String>> {
661        self.writer.finish(result, is_error, cancelled)
662    }
663    fn close<'a>(&'a self, reason: Option<&'a str>) -> BoxFuture<'a, Result<(), String>> {
664        self.writer.close(reason)
665    }
666}
667
668// A single atomic orders admission against retirement: once CLOSED is set,
669// no callback can increment the active count. Only observed requests allocate it.
670#[derive(Default)]
671struct UpdateGate {
672    state: AtomicUsize,
673    idle: Notify,
674}
675const UPDATES_CLOSED: usize = 1usize << (usize::BITS - 1);
676impl UpdateGate {
677    fn enter(&self) -> Option<UpdatePermit<'_>> {
678        #[allow(deprecated)] // Keep the workspace Rust 1.88 MSRV.
679        self.state
680            .fetch_update(Ordering::AcqRel, Ordering::Acquire, |state| {
681                (state & UPDATES_CLOSED == 0 && state < UPDATES_CLOSED - 1).then(|| state + 1)
682            })
683            .ok()
684            .map(|_| UpdatePermit(self))
685    }
686    async fn retire(&self) {
687        self.state.fetch_or(UPDATES_CLOSED, Ordering::AcqRel);
688        loop {
689            let idle = self.idle.notified();
690            tokio::pin!(idle);
691            // Register before checking the counter, so the last exiting callback
692            // cannot notify between the check and waiter registration.
693            idle.as_mut().enable();
694            if self.state.load(Ordering::Acquire) == UPDATES_CLOSED {
695                return;
696            }
697            idle.await;
698        }
699    }
700}
701struct UpdatePermit<'a>(&'a UpdateGate);
702impl Drop for UpdatePermit<'_> {
703    fn drop(&mut self) {
704        if self.0.state.fetch_sub(1, Ordering::AcqRel) & !UPDATES_CLOSED == 1 {
705            self.0.idle.notify_waiters();
706        }
707    }
708}
709
710struct ProgressEmitter {
711    // Native tools may retain their update callback in the program's tool registry.
712    // Such callbacks must not keep that registry (and themselves) alive.
713    program: Weak<ToolProgram>,
714    request_id: String,
715    name: String,
716    signal: CancellationToken,
717    updates: Arc<AtomicUsize>,
718    observation: Option<Weak<dyn ToolProgramObservation>>,
719    observation_error: Option<Arc<Mutex<Option<String>>>>,
720    update_gate: Option<Arc<UpdateGate>>,
721}
722impl AgentEmitter for ProgressEmitter {
723    fn emit(&self, event: AgentEvent) -> BoxFuture<'static, ()> {
724        self.try_emit(event);
725        Box::pin(async {})
726    }
727    fn try_emit(&self, event: AgentEvent) {
728        // Saved callbacks own only this standalone gate and weak runtime/writer
729        // references. Hold admission through writer acceptance AND root reporting.
730        let _admission = match &self.update_gate {
731            Some(gate) => match gate.enter() {
732                Some(permit) => Some(permit),
733                None => return,
734            },
735            None => None,
736        };
737        if let (AgentEvent::ToolExecutionUpdate { partial_result, .. }, Some(program)) =
738            (event, self.program.upgrade())
739        {
740            if !self.signal.is_cancelled() && !program.signal.is_cancelled() {
741                if let Some(observation) = self.observation.as_ref().and_then(Weak::upgrade) {
742                    let updated = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
743                        observation.update(partial_result)
744                    }))
745                    .unwrap_or_else(|_| Err("observation update panicked".into()));
746                    if let Err(error) = updated {
747                        if let Some(slot) = &self.observation_error {
748                            *slot.lock().unwrap() = Some(error);
749                        }
750                        program.cancel_all();
751                    }
752                }
753            }
754            let count = self
755                .updates
756                .fetch_add(1, Ordering::Relaxed)
757                .saturating_add(1);
758            program.report(
759                &self.request_id,
760                &self.name,
761                "progress",
762                count,
763                &self.signal,
764            );
765        }
766    }
767}
768
769/// Only lifecycle identity and a count cross the outer native update callback.
770struct Progress {
771    callback: Arc<dyn Fn(crate::ToolResultPartial) + Send + Sync>,
772    signal: CancellationToken,
773    budget: Mutex<(usize, usize)>,
774}
775impl ToolProgram {
776    fn report(&self, id: &str, name: &str, phase: &str, count: usize, signal: &CancellationToken) {
777        if signal.is_cancelled()
778            || self.signal.is_cancelled()
779            || self.progress.signal.is_cancelled()
780        {
781            return;
782        }
783        let short = |s: &str| s.chars().take(256).collect::<String>();
784        let details = json!({"toolProgram": {"rootId":short(&self.identity.root_id),
785            "parentId":short(&self.identity.parent_id), "requestId":short(id),
786            "tool":short(name), "phase":phase, "updateCount":count}});
787        let bytes = details.to_string().len();
788        {
789            let mut budget = self.progress.budget.lock().unwrap();
790            if budget.0 >= 256 || budget.1 + bytes > 64 * 1024 {
791                return;
792            }
793            budget.0 += 1;
794            budget.1 += bytes;
795        }
796        // An observer must never unwind past a cancelled worker's drain/permits.
797        let _ = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| {
798            (self.progress.callback)(crate::AgentToolResult {
799                details,
800                ..Default::default()
801            });
802        }));
803    }
804}