gate4agent_shell_native/
provider_supervisor.rs

1use gate4agent_handle::{
2    ProviderCancellation, ProviderCancellationToken, ProviderInvocation,
3    ProviderObservationOutcome, ProviderObservationStatus, ProviderRuntimeError,
4    ProviderRuntimeHandle, ProviderRuntimeState, ProviderWork,
5};
6use gate4agent_tool_protocol::{
7    CapabilityObservation, CapabilityProviderDescriptor, CapabilityRequestKey,
8    InvocationCancelReason, ProviderBindingId, ToolCapabilityId, ToolFailure, ToolFailureKind,
9    ToolOperationId, ToolProviderId,
10};
11use std::collections::{BTreeMap, VecDeque};
12use std::sync::mpsc::TryRecvError;
13use std::time::{Duration, Instant};
14
15pub const MAX_PROVIDER_SUPERVISOR_EVENTS: usize = 1_024;
16pub const MAX_PROVIDER_SUPERVISOR_OPERATIONS: usize = 1_024;
17pub const MAX_PROVIDER_SUPERVISOR_TOMBSTONES: usize =
18    MAX_PROVIDER_SUPERVISOR_OPERATIONS * 2;
19pub const MAX_PROVIDER_SUPERVISOR_OUTCOMES_PER_TICK: usize = 64;
20pub const MAX_PROVIDER_SUPERVISOR_WORK_PER_TICK: usize = 64;
21pub const MAX_PROVIDER_STOP_SIGNAL_ATTEMPTS: u8 = 3;
22pub const MAX_PROVIDER_FORCE_STOP_ATTEMPTS: u8 = 3;
23pub const DEFAULT_PROVIDER_STOP_GRACE: Duration = Duration::from_secs(5);
24
25/// A reviewed provider adapter with fixed launch policy.
26///
27/// The supervisor deliberately offers no command, shell, filesystem, or
28/// working-directory launch surface. Implementations map the already-admitted
29/// invocation to one provider-specific operation.
30pub trait NativeProviderExecutor: Send {
31    fn start(
32        &mut self,
33        invocation: &ProviderInvocation,
34    ) -> Result<Box<dyn NativeProviderOperation>, ToolFailure>;
35}
36
37/// Provider-specific ownership of one physical operation.
38///
39/// `request_stop` asks for cooperative shutdown. `request_force_stop` must
40/// terminate the entire provider-owned resource tree after the bounded grace
41/// window. `try_wait` is the sole source of physical-exit truth and must
42/// return `Some` only after every owned OS resource has been reaped. A failed
43/// stop request does not release the implementer from continuing to service
44/// `try_wait`. Both stop methods must be idempotent.
45///
46/// Implementations must also make their own `Drop` fail-safe for live native
47/// resources. Dropping [`ProviderSupervisor`] makes one best-effort force-stop
48/// request, but is not the clean shutdown protocol and cannot prove reaping.
49pub trait NativeProviderOperation: Send {
50    fn request_stop(&mut self) -> Result<(), NativeProviderOperationError>;
51
52    fn request_force_stop(&mut self) -> Result<(), NativeProviderOperationError>;
53
54    fn try_wait(&mut self) -> Result<Option<NativeProviderExit>, NativeProviderOperationError>;
55
56    fn try_poll_result(
57        &mut self,
58    ) -> Result<NativeProviderResultPoll, NativeProviderOperationError>;
59}
60
61#[derive(Clone, Copy, Debug, Eq, PartialEq)]
62pub struct NativeProviderOperationError {
63    pub code: &'static str,
64}
65
66impl NativeProviderOperationError {
67    pub const fn new(code: &'static str) -> Self {
68        Self { code }
69    }
70}
71
72#[derive(Clone, Debug, Eq, PartialEq)]
73pub enum NativeProviderResultPoll {
74    Pending,
75    Ready(CapabilityObservation),
76    Closed,
77}
78
79#[derive(Clone, Copy, Debug, Eq, PartialEq)]
80pub struct NativeProviderExit {
81    pub success: bool,
82    pub code: Option<i32>,
83}
84
85#[derive(Clone, Debug, Eq, PartialEq)]
86pub struct ProviderOperationKey {
87    pub binding_id: ProviderBindingId,
88    pub operation_id: ToolOperationId,
89    pub request_key: CapabilityRequestKey,
90}
91
92#[derive(Clone, Debug, Eq, PartialEq)]
93pub struct PhysicalExitAck {
94    pub provider_id: ToolProviderId,
95    pub binding_id: ProviderBindingId,
96    pub operation: ProviderOperationKey,
97    pub exit: NativeProviderExit,
98    pub stop_signal_attempted: bool,
99    pub stop_signalled: bool,
100    pub force_stop_attempted: bool,
101    pub force_stop_signalled: bool,
102    pub stop_cause: Option<ProviderStopCause>,
103}
104
105#[derive(Clone, Copy, Debug, Eq, PartialEq)]
106pub enum ProviderStopCause {
107    Cancellation {
108        reason: InvocationCancelReason,
109    },
110    Retirement,
111    SupervisorFault,
112}
113
114#[derive(Clone, Copy, Debug, Eq, PartialEq)]
115pub enum ProviderSupervisorState {
116    Running,
117    Draining,
118    Closing,
119    Closed,
120    Faulted,
121}
122
123#[derive(Clone, Debug, Eq, PartialEq)]
124pub struct ProviderOperationSnapshot {
125    pub operation: ProviderOperationKey,
126    pub capability_id: ToolCapabilityId,
127    pub deadline_tick: u64,
128    pub stop_signal_attempted: bool,
129    pub stop_signalled: bool,
130    pub force_stop_attempted: bool,
131    pub force_stop_signalled: bool,
132    pub stop_cause: Option<ProviderStopCause>,
133    pub physical_exit: Option<NativeProviderExit>,
134    pub result_ready: bool,
135}
136
137#[derive(Clone, Debug, Eq, PartialEq)]
138pub struct ProviderSupervisorSnapshot {
139    pub provider_id: ToolProviderId,
140    pub binding_id: ProviderBindingId,
141    pub state: ProviderSupervisorState,
142    pub runtime_state: ProviderRuntimeState,
143    pub operations: Vec<ProviderOperationSnapshot>,
144    pub close_sequence: Option<u64>,
145    pub buffered_exit_acks: usize,
146    pub buffered_faults: usize,
147    pub dropped_faults: u64,
148    pub exit_ack_backpressured: bool,
149}
150
151#[derive(Clone, Copy, Debug, Eq, PartialEq)]
152pub enum ProviderSupervisorBuildError {
153    InvalidDescriptor,
154    ProviderMismatch,
155}
156
157#[derive(Clone, Copy, Debug, Eq, PartialEq)]
158pub enum ProviderSupervisorFaultKind {
159    InvalidExecutorFailure,
160    InvalidExecutorObservation,
161    StopSignalFailed,
162    ForceStopFailed,
163    WaitFailed,
164    ResultPollFailed,
165    OperationCapacityExceeded,
166    ConflictingInvocation,
167    UnknownCancellation,
168    CancellationIdentityMismatch,
169    ObservationOutcomeRejected,
170    ObservationOutcomeContractViolation,
171    ObservationOutcomeRuntimeClosed,
172    ObservationOutcomeMismatch,
173    TombstoneCapacityExceeded,
174    RuntimeQuiesceFailed(ProviderRuntimeError),
175    RuntimeCompletionFailed(ProviderRuntimeError),
176    RuntimeCloseFailed(ProviderRuntimeError),
177    FaultBufferOverflow { dropped: u64 },
178}
179
180#[derive(Clone, Debug, Eq, PartialEq)]
181pub struct ProviderSupervisorFault {
182    pub provider_id: ToolProviderId,
183    pub binding_id: ProviderBindingId,
184    pub operation: Option<ProviderOperationKey>,
185    pub kind: ProviderSupervisorFaultKind,
186    pub executor_error_code: Option<&'static str>,
187    pub blocks_detach: bool,
188}
189
190#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
191pub struct ProviderSupervisorTick {
192    pub work_received: usize,
193    pub observation_outcomes_received: usize,
194    pub operations_started: usize,
195    pub stop_signals: usize,
196    pub force_stop_signals: usize,
197    pub physical_exit_acks: usize,
198    pub observations_submitted: usize,
199    pub faults_recorded: usize,
200    pub close_requested: bool,
201}
202
203/// Owns native provider operations through physical exit and canonical
204/// observation disposition.
205///
206/// Clean shutdown requires `begin_retirement` followed by ticks until
207/// [`ProviderSupervisorState::Closed`]. `Drop` only sends one best-effort
208/// force-stop request to each still-live physical owner; it does not wait or
209/// prove that an OS resource was reaped.
210pub struct ProviderSupervisor {
211    descriptor: CapabilityProviderDescriptor,
212    owned: BTreeMap<ToolOperationId, OwnedProviderOperation>,
213    executor: Box<dyn NativeProviderExecutor>,
214    runtime: ProviderRuntimeHandle,
215    stop_grace: Duration,
216    observation_capacity: usize,
217    outstanding_observations: usize,
218    lifecycle: ProviderSupervisorState,
219    retirement_requested: bool,
220    detach_blocked: bool,
221    tombstones: VecDeque<CompletedOperationTombstone>,
222    close_sequence: Option<u64>,
223    exit_acks: VecDeque<PhysicalExitAck>,
224    faults: VecDeque<ProviderSupervisorFault>,
225    dropped_faults: u64,
226}
227
228enum OwnedProviderOperation {
229    Physical(OwnedPhysicalOperation),
230    PendingCompletion(PendingCompletion),
231    AwaitingObservationOutcome(AwaitingObservationOutcome),
232    AwaitingCancellation(AwaitingCancellation),
233}
234
235struct OwnedPhysicalOperation {
236    key: ProviderOperationKey,
237    capability_id: ToolCapabilityId,
238    deadline_tick: u64,
239    invocation: ProviderInvocation,
240    operation: Box<dyn NativeProviderOperation>,
241    stop_signal_attempted: bool,
242    stop_signal_attempts: u8,
243    stop_signalled: bool,
244    stop_signalled_at: Option<Instant>,
245    force_stop_attempts: u8,
246    force_stop_signalled: bool,
247    stop_cause: Option<ProviderStopCause>,
248    cancellation_received: bool,
249    result: Option<CapabilityObservation>,
250    result_closed: bool,
251    physical_exit: Option<NativeProviderExit>,
252    exit_ack_emitted: bool,
253}
254
255struct PendingCompletion {
256    key: ProviderOperationKey,
257    capability_id: ToolCapabilityId,
258    deadline_tick: u64,
259    invocation: ProviderInvocation,
260    observation: CapabilityObservation,
261}
262
263struct AwaitingObservationOutcome {
264    key: ProviderOperationKey,
265    capability_id: ToolCapabilityId,
266    deadline_tick: u64,
267    sequence: u64,
268    cancellation_token: ProviderCancellationToken,
269    cancellation_reason: Option<InvocationCancelReason>,
270}
271
272struct AwaitingCancellation {
273    key: ProviderOperationKey,
274    capability_id: ToolCapabilityId,
275    deadline_tick: u64,
276}
277
278struct CompletedOperationTombstone {
279    key: ProviderOperationKey,
280}
281
282impl ProviderSupervisor {
283    pub fn new(
284        descriptor: CapabilityProviderDescriptor,
285        runtime: ProviderRuntimeHandle,
286        executor: Box<dyn NativeProviderExecutor>,
287    ) -> Result<Self, ProviderSupervisorBuildError> {
288        Self::new_with_stop_grace(descriptor, runtime, executor, DEFAULT_PROVIDER_STOP_GRACE)
289    }
290
291    pub fn new_with_stop_grace(
292        descriptor: CapabilityProviderDescriptor,
293        runtime: ProviderRuntimeHandle,
294        executor: Box<dyn NativeProviderExecutor>,
295        stop_grace: Duration,
296    ) -> Result<Self, ProviderSupervisorBuildError> {
297        descriptor
298            .validate()
299            .map_err(|_| ProviderSupervisorBuildError::InvalidDescriptor)?;
300        if &descriptor.id != runtime.provider_id() {
301            return Err(ProviderSupervisorBuildError::ProviderMismatch);
302        }
303        let observation_capacity = runtime.effect_capacity();
304        Ok(Self {
305            descriptor,
306            owned: BTreeMap::new(),
307            executor,
308            runtime,
309            stop_grace: if stop_grace.is_zero() {
310                Duration::from_millis(1)
311            } else {
312                stop_grace
313            },
314            observation_capacity,
315            outstanding_observations: 0,
316            lifecycle: ProviderSupervisorState::Running,
317            retirement_requested: false,
318            detach_blocked: false,
319            tombstones: VecDeque::new(),
320            close_sequence: None,
321            exit_acks: VecDeque::new(),
322            faults: VecDeque::new(),
323            dropped_faults: 0,
324        })
325    }
326
327    pub fn descriptor(&self) -> &CapabilityProviderDescriptor {
328        &self.descriptor
329    }
330
331    pub fn begin_retirement(&mut self) -> Result<(), ProviderRuntimeError> {
332        if self.retirement_requested {
333            return Ok(());
334        }
335        self.retirement_requested = true;
336        if matches!(self.runtime.state(), ProviderRuntimeState::Closed) && self.owned.is_empty() {
337            self.detach_blocked = false;
338            self.lifecycle = ProviderSupervisorState::Closed;
339            return Ok(());
340        }
341        self.lifecycle = ProviderSupervisorState::Draining;
342        for owned in self.owned.values_mut() {
343            if let OwnedProviderOperation::Physical(owned) = owned {
344                if owned.physical_exit.is_none()
345                    && !owned.stop_signal_attempted
346                    && owned.stop_cause.is_none()
347                {
348                    owned.stop_cause = Some(ProviderStopCause::Retirement);
349                }
350            }
351        }
352        if matches!(self.runtime.state(), ProviderRuntimeState::Closed) {
353            return Ok(());
354        }
355        if matches!(self.runtime.state(), ProviderRuntimeState::Attaching)
356            && self.owned.is_empty()
357        {
358            let mut report = ProviderSupervisorTick::default();
359            return self.request_runtime_close(&mut report);
360        }
361        match self.runtime.begin_quiesce() {
362            Ok(()) => Ok(()),
363            Err(ProviderRuntimeError::Inactive)
364                if matches!(self.runtime.state(), ProviderRuntimeState::Closing) =>
365            {
366                Ok(())
367            }
368            Err(error) => {
369                let mut report = ProviderSupervisorTick::default();
370                self.record_blocking_fault(
371                    None,
372                    ProviderSupervisorFaultKind::RuntimeQuiesceFailed(error),
373                    None,
374                    &mut report,
375                );
376                Err(error)
377            }
378        }
379    }
380
381    pub fn snapshot(&self) -> ProviderSupervisorSnapshot {
382        ProviderSupervisorSnapshot {
383            provider_id: self.descriptor.id.clone(),
384            binding_id: self.runtime.binding_id(),
385            state: self.state(),
386            runtime_state: self.runtime.state(),
387            operations: self
388                .owned
389                .values()
390                .map(OwnedProviderOperation::snapshot)
391                .collect(),
392            close_sequence: self.close_sequence,
393            buffered_exit_acks: self.exit_acks.len(),
394            buffered_faults: self.faults.len() + usize::from(self.dropped_faults != 0),
395            dropped_faults: self.dropped_faults,
396            exit_ack_backpressured: self
397                .owned
398                .values()
399                .any(OwnedProviderOperation::exit_ack_backpressured),
400        }
401    }
402
403    pub fn state(&self) -> ProviderSupervisorState {
404        if self.detach_blocked {
405            ProviderSupervisorState::Faulted
406        } else {
407            self.lifecycle
408        }
409    }
410
411    pub fn drain_exit_acks(&mut self, limit: usize) -> Vec<PhysicalExitAck> {
412        drain_queue(&mut self.exit_acks, limit)
413    }
414
415    pub fn drain_faults(&mut self, limit: usize) -> Vec<ProviderSupervisorFault> {
416        if limit == 0 {
417            return Vec::new();
418        }
419        let mut drained = Vec::new();
420        if self.dropped_faults != 0 {
421            let dropped = std::mem::take(&mut self.dropped_faults);
422            drained.push(ProviderSupervisorFault {
423                provider_id: self.descriptor.id.clone(),
424                binding_id: self.runtime.binding_id(),
425                operation: None,
426                kind: ProviderSupervisorFaultKind::FaultBufferOverflow { dropped },
427                executor_error_code: None,
428                blocks_detach: false,
429            });
430        }
431        drained.extend(drain_queue(&mut self.faults, limit - drained.len()));
432        drained
433    }
434
435    pub fn tick(&mut self) -> ProviderSupervisorTick {
436        let mut report = ProviderSupervisorTick::default();
437        if matches!(self.lifecycle, ProviderSupervisorState::Closed) {
438            return report;
439        }
440
441        self.drain_observation_outcomes(&mut report);
442
443        let runtime_closed = matches!(self.runtime.state(), ProviderRuntimeState::Closed);
444        if runtime_closed {
445            if !self.owned.is_empty() && !self.detach_blocked {
446                self.record_blocking_fault(
447                    None,
448                    ProviderSupervisorFaultKind::RuntimeCompletionFailed(
449                        ProviderRuntimeError::Inactive,
450                    ),
451                    None,
452                    &mut report,
453                );
454            }
455            self.outstanding_observations = 0;
456        }
457
458        let operation_ids = self.owned.keys().copied().collect::<Vec<_>>();
459        for operation_id in operation_ids {
460            let Some(mut owned) = self.owned.remove(&operation_id) else {
461                continue;
462            };
463            if runtime_closed && !owned.has_physical_owner() {
464                continue;
465            }
466            if runtime_closed {
467                owned.fence_for_runtime_close(self.retirement_requested);
468            }
469            if let Some(owned) = self.tick_owned(owned, &mut report) {
470                self.owned.insert(operation_id, owned);
471            }
472        }
473
474        if runtime_closed {
475            if self.owned.is_empty() {
476                self.lifecycle = ProviderSupervisorState::Closed;
477                if self.retirement_requested {
478                    self.detach_blocked = false;
479                }
480            }
481            return report;
482        }
483
484        if self.retirement_requested {
485            self.lifecycle = if self.close_sequence.is_some() {
486                ProviderSupervisorState::Closing
487            } else {
488                ProviderSupervisorState::Draining
489            };
490            if !matches!(
491                self.runtime.state(),
492                ProviderRuntimeState::Attaching | ProviderRuntimeState::Closed
493            ) {
494                self.drain_work(&mut report);
495            }
496            if self.owned.is_empty() {
497                self.detach_blocked = false;
498                let _ = self.request_runtime_close(&mut report);
499            }
500            return report;
501        }
502
503        if !matches!(
504            self.runtime.state(),
505            ProviderRuntimeState::Attaching | ProviderRuntimeState::Closed
506        ) {
507            self.drain_work(&mut report);
508        }
509        report
510    }
511
512    fn drain_observation_outcomes(&mut self, report: &mut ProviderSupervisorTick) {
513        for _ in 0..MAX_PROVIDER_SUPERVISOR_OUTCOMES_PER_TICK {
514            match self.runtime.try_recv_observation_outcome() {
515                Ok(outcome) => {
516                    report.observation_outcomes_received += 1;
517                    self.handle_observation_outcome(outcome, report);
518                }
519                Err(TryRecvError::Empty | TryRecvError::Disconnected) => return,
520            }
521        }
522    }
523
524    fn handle_observation_outcome(
525        &mut self,
526        outcome: ProviderObservationOutcome,
527        report: &mut ProviderSupervisorTick,
528    ) {
529        if self.outstanding_observations == 0 {
530            self.record_blocking_fault(
531                None,
532                ProviderSupervisorFaultKind::ObservationOutcomeMismatch,
533                None,
534                report,
535            );
536        } else {
537            self.outstanding_observations -= 1;
538        }
539        let operation_id = outcome.operation_id;
540        let Some(mut owned) = self.owned.remove(&operation_id) else {
541            self.record_blocking_fault(
542                None,
543                ProviderSupervisorFaultKind::ObservationOutcomeMismatch,
544                None,
545                report,
546            );
547            return;
548        };
549        let OwnedProviderOperation::AwaitingObservationOutcome(awaiting) = owned else {
550            owned.fence_for_fault(self.retirement_requested);
551            let operation = owned.key().clone();
552            self.owned.insert(operation_id, owned);
553            self.record_blocking_fault(
554                Some(operation),
555                ProviderSupervisorFaultKind::ObservationOutcomeMismatch,
556                None,
557                report,
558            );
559            return;
560        };
561        if awaiting.sequence != outcome.sequence
562            || awaiting.key.operation_id != outcome.operation_id
563            || awaiting.key.request_key != outcome.request_key
564        {
565            let operation = awaiting.key.clone();
566            if awaiting.cancellation_reason.is_none()
567                && awaiting.cancellation_token.is_cancelled()
568            {
569                self.push_tombstone(awaiting.key, report);
570            }
571            self.record_blocking_fault(
572                Some(operation),
573                ProviderSupervisorFaultKind::ObservationOutcomeMismatch,
574                None,
575                report,
576            );
577            return;
578        }
579
580        match outcome.status {
581            ProviderObservationStatus::Applied | ProviderObservationStatus::Ignored { .. } => {
582                if awaiting.cancellation_reason.is_none()
583                    && awaiting.cancellation_token.is_cancelled()
584                {
585                    self.push_tombstone(awaiting.key, report);
586                }
587            }
588            status => {
589                let kind = match status {
590                    ProviderObservationStatus::Rejected(_) => {
591                        ProviderSupervisorFaultKind::ObservationOutcomeRejected
592                    }
593                    ProviderObservationStatus::ContractViolation => {
594                        ProviderSupervisorFaultKind::ObservationOutcomeContractViolation
595                    }
596                    ProviderObservationStatus::RuntimeClosed => {
597                        ProviderSupervisorFaultKind::ObservationOutcomeRuntimeClosed
598                    }
599                    ProviderObservationStatus::Applied
600                    | ProviderObservationStatus::Ignored { .. } => unreachable!(),
601                };
602                let operation = awaiting.key.clone();
603                if awaiting.cancellation_reason.is_none()
604                    && awaiting.cancellation_token.is_cancelled()
605                {
606                    self.push_tombstone(awaiting.key, report);
607                }
608                self.record_blocking_fault(Some(operation), kind, None, report);
609            }
610        }
611    }
612
613    fn tick_owned(
614        &mut self,
615        owned: OwnedProviderOperation,
616        report: &mut ProviderSupervisorTick,
617    ) -> Option<OwnedProviderOperation> {
618        match owned {
619            OwnedProviderOperation::Physical(mut owned) => {
620                if self.retirement_requested
621                    && owned.physical_exit.is_none()
622                    && !owned.stop_signal_attempted
623                    && owned.stop_cause.is_none()
624                {
625                    owned.stop_cause = Some(ProviderStopCause::Retirement);
626                }
627                if self.detach_blocked
628                    && owned.physical_exit.is_none()
629                    && !owned.stop_signal_attempted
630                    && owned.stop_cause.is_none()
631                {
632                    owned.stop_cause = Some(ProviderStopCause::SupervisorFault);
633                }
634                let cancellation_requested =
635                    owned.invocation.cancellation_token().is_cancelled();
636                let stop_required = self.retirement_requested
637                    || self.detach_blocked
638                    || cancellation_requested
639                    || owned.stop_cause.is_some();
640                if owned.physical_exit.is_none()
641                    && stop_required
642                    && !owned.stop_signalled
643                    && owned.stop_signal_attempts < MAX_PROVIDER_STOP_SIGNAL_ATTEMPTS
644                {
645                    owned.stop_signal_attempted = true;
646                    owned.stop_signal_attempts += 1;
647                    report.stop_signals += 1;
648                    match owned.operation.request_stop() {
649                        Ok(()) => {
650                            owned.stop_signalled = true;
651                            owned.stop_signalled_at = Some(Instant::now());
652                        }
653                        Err(error) => {
654                            if owned.stop_signal_attempts
655                                == MAX_PROVIDER_STOP_SIGNAL_ATTEMPTS
656                            {
657                                self.record_blocking_fault(
658                                    Some(owned.key.clone()),
659                                    ProviderSupervisorFaultKind::StopSignalFailed,
660                                    Some(error.code),
661                                    report,
662                                );
663                            } else {
664                                self.record_fault(
665                                    Some(owned.key.clone()),
666                                    ProviderSupervisorFaultKind::StopSignalFailed,
667                                    Some(error.code),
668                                    false,
669                                    report,
670                                );
671                            }
672                        }
673                    }
674                }
675                let force_stop_required = owned.physical_exit.is_none()
676                    && stop_required
677                    && !owned.force_stop_signalled
678                    && ((!owned.stop_signalled
679                        && owned.stop_signal_attempts
680                            == MAX_PROVIDER_STOP_SIGNAL_ATTEMPTS)
681                        || owned
682                            .stop_signalled_at
683                            .is_some_and(|started| started.elapsed() >= self.stop_grace));
684                if force_stop_required
685                    && owned.force_stop_attempts < MAX_PROVIDER_FORCE_STOP_ATTEMPTS
686                {
687                    owned.force_stop_attempts += 1;
688                    report.force_stop_signals += 1;
689                    match owned.operation.request_force_stop() {
690                        Ok(()) => owned.force_stop_signalled = true,
691                        Err(error) => {
692                            if owned.force_stop_attempts
693                                == MAX_PROVIDER_FORCE_STOP_ATTEMPTS
694                            {
695                                self.record_blocking_fault(
696                                    Some(owned.key.clone()),
697                                    ProviderSupervisorFaultKind::ForceStopFailed,
698                                    Some(error.code),
699                                    report,
700                                );
701                            } else {
702                                self.record_fault(
703                                    Some(owned.key.clone()),
704                                    ProviderSupervisorFaultKind::ForceStopFailed,
705                                    Some(error.code),
706                                    false,
707                                    report,
708                                );
709                            }
710                        }
711                    }
712                }
713
714                if !cancellation_requested
715                    && !owned.stop_signal_attempted
716                    && owned.stop_cause.is_none()
717                    && owned.result.is_none()
718                    && !owned.result_closed
719                {
720                    match owned.operation.try_poll_result() {
721                        Ok(NativeProviderResultPoll::Pending) => {}
722                        Ok(NativeProviderResultPoll::Closed) => owned.result_closed = true,
723                        Ok(NativeProviderResultPoll::Ready(observation)) => {
724                            if observation.validate().is_ok() {
725                                owned.result = Some(observation);
726                            } else {
727                                owned.result = Some(contract_violation_observation());
728                                self.record_fault(
729                                    Some(owned.key.clone()),
730                                    ProviderSupervisorFaultKind::InvalidExecutorObservation,
731                                    None,
732                                    false,
733                                    report,
734                                );
735                            }
736                        }
737                        Err(error) => {
738                            owned.result = Some(contract_violation_observation());
739                            owned.result_closed = true;
740                            if !self.retirement_requested && owned.stop_cause.is_none() {
741                                owned.stop_cause = Some(ProviderStopCause::SupervisorFault);
742                            }
743                            self.record_blocking_fault(
744                                Some(owned.key.clone()),
745                                ProviderSupervisorFaultKind::ResultPollFailed,
746                                Some(error.code),
747                                report,
748                            );
749                        }
750                    }
751                }
752
753                if owned.physical_exit.is_none() {
754                    match owned.operation.try_wait() {
755                        Ok(Some(exit)) => owned.physical_exit = Some(exit),
756                        Ok(None) => {}
757                        Err(error) => {
758                            if !self.retirement_requested
759                                && !owned.stop_signal_attempted
760                                && owned.stop_cause.is_none()
761                            {
762                                owned.stop_cause = Some(ProviderStopCause::SupervisorFault);
763                            }
764                            self.record_blocking_fault(
765                                Some(owned.key.clone()),
766                                ProviderSupervisorFaultKind::WaitFailed,
767                                Some(error.code),
768                                report,
769                            );
770                            return Some(OwnedProviderOperation::Physical(owned));
771                        }
772                    }
773                }
774
775                if let Some(exit) = owned.physical_exit {
776                    if !owned.exit_ack_emitted {
777                        if owned.stop_signal_attempted && owned.stop_cause.is_none() {
778                            return Some(OwnedProviderOperation::Physical(owned));
779                        }
780                        if self.exit_acks.len() >= MAX_PROVIDER_SUPERVISOR_EVENTS {
781                            return Some(OwnedProviderOperation::Physical(owned));
782                        }
783                        self.exit_acks.push_back(PhysicalExitAck {
784                            provider_id: self.descriptor.id.clone(),
785                            binding_id: self.runtime.binding_id(),
786                            operation: owned.key.clone(),
787                            exit,
788                            stop_signal_attempted: owned.stop_signal_attempted,
789                            stop_signalled: owned.stop_signalled,
790                            force_stop_attempted: owned.force_stop_attempts != 0,
791                            force_stop_signalled: owned.force_stop_signalled,
792                            stop_cause: owned.stop_cause,
793                        });
794                        owned.exit_ack_emitted = true;
795                        report.physical_exit_acks += 1;
796                    }
797
798                    if self.retirement_requested
799                        || self.detach_blocked
800                        || owned.cancellation_received
801                        || owned.stop_cause.is_some()
802                    {
803                        if !owned.cancellation_received {
804                            self.push_tombstone(owned.key, report);
805                        }
806                        return None;
807                    }
808                    if owned.result.is_none() && owned.result_closed {
809                        owned.result = Some(contract_violation_observation());
810                        self.record_fault(
811                            Some(owned.key.clone()),
812                            ProviderSupervisorFaultKind::InvalidExecutorObservation,
813                            None,
814                            false,
815                            report,
816                        );
817                    }
818                    if let Some(observation) = owned.result.take() {
819                        return Some(OwnedProviderOperation::PendingCompletion(
820                            PendingCompletion {
821                                key: owned.key,
822                                capability_id: owned.capability_id,
823                                deadline_tick: owned.deadline_tick,
824                                invocation: owned.invocation,
825                                observation,
826                            },
827                        ));
828                    }
829                }
830                Some(OwnedProviderOperation::Physical(owned))
831            }
832            OwnedProviderOperation::PendingCompletion(mut pending) => {
833                if self.retirement_requested {
834                    self.push_tombstone(pending.key, report);
835                    return None;
836                }
837                if pending.invocation.cancellation_token().is_cancelled() {
838                    return Some(OwnedProviderOperation::AwaitingCancellation(
839                        AwaitingCancellation {
840                            key: pending.key,
841                            capability_id: pending.capability_id,
842                            deadline_tick: pending.deadline_tick,
843                        },
844                    ));
845                }
846                if self.outstanding_observations >= self.observation_capacity {
847                    return Some(OwnedProviderOperation::PendingCompletion(pending));
848                }
849                let cancellation_token = pending.invocation.cancellation_token();
850                match self
851                    .runtime
852                    .try_complete(&mut pending.invocation, &pending.observation)
853                {
854                    Ok(sequence) => {
855                        self.outstanding_observations += 1;
856                        report.observations_submitted += 1;
857                        Some(OwnedProviderOperation::AwaitingObservationOutcome(
858                            AwaitingObservationOutcome {
859                                key: pending.key,
860                                capability_id: pending.capability_id,
861                                deadline_tick: pending.deadline_tick,
862                                sequence,
863                                cancellation_token,
864                                cancellation_reason: None,
865                            },
866                        ))
867                    }
868                    Err(ProviderRuntimeError::Full) => {
869                        Some(OwnedProviderOperation::PendingCompletion(pending))
870                    }
871                    Err(error) => {
872                        let operation = pending.key.clone();
873                        self.push_tombstone(pending.key, report);
874                        self.record_blocking_fault(
875                            Some(operation),
876                            ProviderSupervisorFaultKind::RuntimeCompletionFailed(error),
877                            None,
878                            report,
879                        );
880                        None
881                    }
882                }
883            }
884            OwnedProviderOperation::AwaitingObservationOutcome(awaiting) => {
885                Some(OwnedProviderOperation::AwaitingObservationOutcome(awaiting))
886            }
887            OwnedProviderOperation::AwaitingCancellation(awaiting) => {
888                if self.retirement_requested {
889                    self.push_tombstone(awaiting.key, report);
890                    None
891                } else {
892                    Some(OwnedProviderOperation::AwaitingCancellation(awaiting))
893                }
894            }
895        }
896    }
897
898    fn drain_work(&mut self, report: &mut ProviderSupervisorTick) {
899        for _ in 0..MAX_PROVIDER_SUPERVISOR_WORK_PER_TICK {
900            match self.runtime.try_recv() {
901                Ok(ProviderWork::Cancel(cancellation)) => {
902                    report.work_received += 1;
903                    self.handle_cancellation(cancellation, report);
904                }
905                Ok(ProviderWork::Invoke(invocation)) => {
906                    report.work_received += 1;
907                    if self.retirement_requested || self.detach_blocked {
908                        self.push_tombstone(operation_key(&invocation), report);
909                    } else {
910                        self.start_invocation(invocation, report);
911                    }
912                }
913                Err(TryRecvError::Empty | TryRecvError::Disconnected) => return,
914            }
915        }
916    }
917
918    fn start_invocation(
919        &mut self,
920        invocation: ProviderInvocation,
921        report: &mut ProviderSupervisorTick,
922    ) {
923        let key = operation_key(&invocation);
924        let capability_id = invocation.capability_id().clone();
925        let deadline_tick = invocation.deadline_tick();
926        if self.owned.contains_key(&key.operation_id)
927            || self
928                .tombstones
929                .iter()
930                .any(|tombstone| tombstone.key.operation_id == key.operation_id)
931        {
932            self.record_blocking_fault(
933                Some(key),
934                ProviderSupervisorFaultKind::ConflictingInvocation,
935                None,
936                report,
937            );
938            return;
939        }
940        if self.owned.len() >= MAX_PROVIDER_SUPERVISOR_OPERATIONS {
941            self.record_blocking_fault(
942                Some(key),
943                ProviderSupervisorFaultKind::OperationCapacityExceeded,
944                None,
945                report,
946            );
947            return;
948        }
949        if !self.descriptor.has_capability(&capability_id) {
950            self.record_fault(
951                Some(key.clone()),
952                ProviderSupervisorFaultKind::InvalidExecutorObservation,
953                None,
954                false,
955                report,
956            );
957            self.owned.insert(key.operation_id, OwnedProviderOperation::PendingCompletion(
958                PendingCompletion {
959                    key: key.clone(),
960                    capability_id,
961                    deadline_tick,
962                    invocation,
963                    observation: contract_violation_observation(),
964                },
965            ));
966            return;
967        }
968
969        if invocation.cancellation_token().is_cancelled() {
970            self.owned.insert(
971                key.operation_id,
972                OwnedProviderOperation::AwaitingCancellation(AwaitingCancellation {
973                    key,
974                    capability_id,
975                    deadline_tick,
976                }),
977            );
978            return;
979        }
980
981        match self.executor.start(&invocation) {
982            Ok(operation) => {
983                report.operations_started += 1;
984                self.owned.insert(key.operation_id, OwnedProviderOperation::Physical(OwnedPhysicalOperation {
985                    key: key.clone(),
986                    capability_id,
987                    deadline_tick,
988                    invocation,
989                    operation,
990                    stop_signal_attempted: false,
991                    stop_signal_attempts: 0,
992                    stop_signalled: false,
993                    stop_signalled_at: None,
994                    force_stop_attempts: 0,
995                    force_stop_signalled: false,
996                    stop_cause: None,
997                    cancellation_received: false,
998                    result: None,
999                    result_closed: false,
1000                    physical_exit: None,
1001                    exit_ack_emitted: false,
1002                }));
1003            }
1004            Err(failure) => {
1005                let observation = if failure.validate().is_ok() {
1006                    CapabilityObservation::Failed { failure }
1007                } else {
1008                    self.record_fault(
1009                        Some(key.clone()),
1010                        ProviderSupervisorFaultKind::InvalidExecutorFailure,
1011                        None,
1012                        false,
1013                        report,
1014                    );
1015                    contract_violation_observation()
1016                };
1017                self.owned.insert(key.operation_id, OwnedProviderOperation::PendingCompletion(
1018                    PendingCompletion {
1019                        key: key.clone(),
1020                        capability_id,
1021                        deadline_tick,
1022                        invocation,
1023                        observation,
1024                    },
1025                ));
1026            }
1027        }
1028    }
1029
1030    fn handle_cancellation(
1031        &mut self,
1032        cancellation: ProviderCancellation,
1033        report: &mut ProviderSupervisorTick,
1034    ) {
1035        let key = ProviderOperationKey {
1036            binding_id: cancellation.binding_id(),
1037            operation_id: cancellation.operation_id(),
1038            request_key: cancellation.request_key().clone(),
1039        };
1040        let Some(mut owned) = self.owned.remove(&key.operation_id) else {
1041            if let Some(index) = self
1042                .tombstones
1043                .iter()
1044                .position(|tombstone| tombstone.key.operation_id == key.operation_id)
1045            {
1046                if self.tombstones[index].key != key {
1047                    self.record_blocking_fault(
1048                        Some(key),
1049                        ProviderSupervisorFaultKind::CancellationIdentityMismatch,
1050                        None,
1051                        report,
1052                    );
1053                } else {
1054                    self.tombstones.remove(index);
1055                }
1056                return;
1057            }
1058            self.record_blocking_fault(
1059                Some(key),
1060                ProviderSupervisorFaultKind::UnknownCancellation,
1061                None,
1062                report,
1063            );
1064            return;
1065        };
1066        if owned.key() != &key {
1067            owned.fence_for_fault(self.retirement_requested);
1068            self.owned.insert(key.operation_id, owned);
1069            self.record_blocking_fault(
1070                Some(key),
1071                ProviderSupervisorFaultKind::CancellationIdentityMismatch,
1072                None,
1073                report,
1074            );
1075            return;
1076        }
1077
1078        match owned {
1079            OwnedProviderOperation::Physical(mut owned) => {
1080                let stop_cause = ProviderStopCause::Cancellation {
1081                    reason: cancellation.reason(),
1082                };
1083                owned.cancellation_received = true;
1084                if owned.stop_cause.is_none()
1085                    && (owned.physical_exit.is_none() || owned.stop_signal_attempted)
1086                {
1087                    owned.stop_cause = Some(stop_cause);
1088                }
1089                if owned.exit_ack_emitted {
1090                    if let Some(ack) = self
1091                        .exit_acks
1092                        .iter_mut()
1093                        .find(|ack| ack.operation == key)
1094                    {
1095                        if ack.stop_cause.is_none()
1096                            && owned.stop_cause == Some(stop_cause)
1097                        {
1098                            ack.stop_cause = owned.stop_cause;
1099                        }
1100                    }
1101                }
1102                if let Some(owned) =
1103                    self.tick_owned(OwnedProviderOperation::Physical(owned), report)
1104                {
1105                    self.owned.insert(key.operation_id, owned);
1106                }
1107            }
1108            OwnedProviderOperation::PendingCompletion(_)
1109            | OwnedProviderOperation::AwaitingCancellation(_) => {}
1110            OwnedProviderOperation::AwaitingObservationOutcome(mut awaiting) => {
1111                awaiting.cancellation_reason = Some(cancellation.reason());
1112                self.owned.insert(
1113                    key.operation_id,
1114                    OwnedProviderOperation::AwaitingObservationOutcome(awaiting),
1115                );
1116            }
1117        }
1118    }
1119
1120    fn push_tombstone(
1121        &mut self,
1122        key: ProviderOperationKey,
1123        report: &mut ProviderSupervisorTick,
1124    ) {
1125        if self.tombstones.len() >= MAX_PROVIDER_SUPERVISOR_TOMBSTONES {
1126            self.record_blocking_fault(
1127                Some(key),
1128                ProviderSupervisorFaultKind::TombstoneCapacityExceeded,
1129                None,
1130                report,
1131            );
1132            return;
1133        }
1134        self.tombstones.push_back(CompletedOperationTombstone { key });
1135    }
1136
1137    fn request_runtime_close(
1138        &mut self,
1139        report: &mut ProviderSupervisorTick,
1140    ) -> Result<(), ProviderRuntimeError> {
1141        match self.runtime.close() {
1142            Ok(sequence) => {
1143                self.close_sequence = Some(sequence);
1144                self.lifecycle = ProviderSupervisorState::Closing;
1145                report.close_requested = true;
1146                Ok(())
1147            }
1148            Err(ProviderRuntimeError::Full) => {
1149                self.lifecycle = ProviderSupervisorState::Closing;
1150                Ok(())
1151            }
1152            Err(ProviderRuntimeError::Inactive)
1153                if matches!(self.runtime.state(), ProviderRuntimeState::Closed) =>
1154            {
1155                self.lifecycle = ProviderSupervisorState::Closed;
1156                Ok(())
1157            }
1158            Err(error) => {
1159                self.record_blocking_fault(
1160                    None,
1161                    ProviderSupervisorFaultKind::RuntimeCloseFailed(error),
1162                    None,
1163                    report,
1164                );
1165                Err(error)
1166            }
1167        }
1168    }
1169
1170    fn record_blocking_fault(
1171        &mut self,
1172        operation: Option<ProviderOperationKey>,
1173        kind: ProviderSupervisorFaultKind,
1174        executor_error_code: Option<&'static str>,
1175        report: &mut ProviderSupervisorTick,
1176    ) {
1177        self.detach_blocked = true;
1178        for owned in self.owned.values_mut() {
1179            owned.fence_for_fault(self.retirement_requested);
1180        }
1181        let _ = self.runtime.begin_quiesce();
1182        self.record_fault(operation, kind, executor_error_code, true, report);
1183    }
1184
1185    fn record_fault(
1186        &mut self,
1187        operation: Option<ProviderOperationKey>,
1188        kind: ProviderSupervisorFaultKind,
1189        executor_error_code: Option<&'static str>,
1190        blocks_detach: bool,
1191        report: &mut ProviderSupervisorTick,
1192    ) {
1193        if self.faults.len() >= MAX_PROVIDER_SUPERVISOR_EVENTS {
1194            self.dropped_faults = self.dropped_faults.saturating_add(1);
1195            report.faults_recorded += 1;
1196            return;
1197        }
1198        self.faults.push_back(ProviderSupervisorFault {
1199            provider_id: self.descriptor.id.clone(),
1200            binding_id: self.runtime.binding_id(),
1201            operation,
1202            kind,
1203            executor_error_code,
1204            blocks_detach,
1205        });
1206        report.faults_recorded += 1;
1207    }
1208}
1209
1210impl OwnedProviderOperation {
1211    fn has_physical_owner(&self) -> bool {
1212        matches!(self, Self::Physical(_))
1213    }
1214
1215    fn fence_for_fault(&mut self, retirement_requested: bool) {
1216        if let Self::Physical(owned) = self {
1217            if retirement_requested {
1218                if owned.physical_exit.is_none()
1219                    && !owned.stop_signal_attempted
1220                    && owned.stop_cause.is_none()
1221                {
1222                    owned.stop_cause = Some(ProviderStopCause::Retirement);
1223                }
1224            } else if owned.physical_exit.is_none()
1225                && !owned.stop_signal_attempted
1226                && owned.stop_cause.is_none()
1227            {
1228                owned.stop_cause = Some(ProviderStopCause::SupervisorFault);
1229            }
1230        }
1231    }
1232
1233    fn fence_for_runtime_close(&mut self, retirement_requested: bool) {
1234        let Self::Physical(owned) = self else {
1235            return;
1236        };
1237        if owned.stop_cause.is_some() {
1238            return;
1239        }
1240        if owned.stop_signal_attempted {
1241            owned.stop_cause = Some(ProviderStopCause::SupervisorFault);
1242        } else if owned.physical_exit.is_none() {
1243            owned.stop_cause = Some(if retirement_requested {
1244                ProviderStopCause::Retirement
1245            } else {
1246                ProviderStopCause::SupervisorFault
1247            });
1248        }
1249    }
1250
1251    fn key(&self) -> &ProviderOperationKey {
1252        match self {
1253            Self::Physical(owned) => &owned.key,
1254            Self::PendingCompletion(pending) => &pending.key,
1255            Self::AwaitingObservationOutcome(awaiting) => &awaiting.key,
1256            Self::AwaitingCancellation(awaiting) => &awaiting.key,
1257        }
1258    }
1259
1260    fn snapshot(&self) -> ProviderOperationSnapshot {
1261        match self {
1262            Self::Physical(owned) => ProviderOperationSnapshot {
1263                operation: owned.key.clone(),
1264                capability_id: owned.capability_id.clone(),
1265                deadline_tick: owned.deadline_tick,
1266                stop_signal_attempted: owned.stop_signal_attempted,
1267                stop_signalled: owned.stop_signalled,
1268                force_stop_attempted: owned.force_stop_attempts != 0,
1269                force_stop_signalled: owned.force_stop_signalled,
1270                stop_cause: owned.stop_cause,
1271                physical_exit: owned.physical_exit,
1272                result_ready: owned.result.is_some(),
1273            },
1274            Self::PendingCompletion(pending) => ProviderOperationSnapshot {
1275                operation: pending.key.clone(),
1276                capability_id: pending.capability_id.clone(),
1277                deadline_tick: pending.deadline_tick,
1278                stop_signal_attempted: false,
1279                stop_signalled: false,
1280                force_stop_attempted: false,
1281                force_stop_signalled: false,
1282                stop_cause: None,
1283                physical_exit: None,
1284                result_ready: true,
1285            },
1286            Self::AwaitingObservationOutcome(awaiting) => ProviderOperationSnapshot {
1287                operation: awaiting.key.clone(),
1288                capability_id: awaiting.capability_id.clone(),
1289                deadline_tick: awaiting.deadline_tick,
1290                stop_signal_attempted: false,
1291                stop_signalled: false,
1292                force_stop_attempted: false,
1293                force_stop_signalled: false,
1294                stop_cause: awaiting
1295                    .cancellation_reason
1296                    .map(|reason| ProviderStopCause::Cancellation { reason }),
1297                physical_exit: None,
1298                result_ready: true,
1299            },
1300            Self::AwaitingCancellation(awaiting) => ProviderOperationSnapshot {
1301                operation: awaiting.key.clone(),
1302                capability_id: awaiting.capability_id.clone(),
1303                deadline_tick: awaiting.deadline_tick,
1304                stop_signal_attempted: false,
1305                stop_signalled: false,
1306                force_stop_attempted: false,
1307                force_stop_signalled: false,
1308                stop_cause: None,
1309                physical_exit: None,
1310                result_ready: false,
1311            },
1312        }
1313    }
1314
1315    fn exit_ack_backpressured(&self) -> bool {
1316        matches!(
1317            self,
1318            Self::Physical(OwnedPhysicalOperation {
1319                physical_exit: Some(_),
1320                exit_ack_emitted: false,
1321                ..
1322            })
1323        )
1324    }
1325}
1326
1327impl Drop for ProviderSupervisor {
1328    fn drop(&mut self) {
1329        for owned in self.owned.values_mut() {
1330            let OwnedProviderOperation::Physical(owned) = owned else {
1331                continue;
1332            };
1333            if owned.physical_exit.is_none() {
1334                owned.stop_signal_attempted = true;
1335                if owned.stop_cause.is_none() {
1336                    owned.stop_cause = Some(ProviderStopCause::Retirement);
1337                }
1338                let _ = owned.operation.request_force_stop();
1339            }
1340        }
1341    }
1342}
1343
1344fn operation_key(invocation: &ProviderInvocation) -> ProviderOperationKey {
1345    ProviderOperationKey {
1346        binding_id: invocation.binding_id(),
1347        operation_id: invocation.operation_id(),
1348        request_key: invocation.request_key().clone(),
1349    }
1350}
1351
1352fn contract_violation_observation() -> CapabilityObservation {
1353    CapabilityObservation::Failed {
1354        failure: ToolFailure {
1355            kind: ToolFailureKind::ProviderContractViolation,
1356            redacted_message: None,
1357        },
1358    }
1359}
1360
1361fn drain_queue<T>(queue: &mut VecDeque<T>, limit: usize) -> Vec<T> {
1362    let count = limit.min(queue.len());
1363    queue.drain(..count).collect()
1364}
gate4agent_shell_native/provider_supervisor.rs

gate4agent_shell_native/
provider_supervisor.rs
