refactor(worker): canonical runtime API on WorkerRuntime (Finding 4 Part B)

[slin1237]

Summary

Part B of the Worker trait trim outlined in .claude/plans/2026-04-14-worker-module-followup-cleanup.md (Finding 4, categories B + C). Follows the same Option 2 principle that landed in Part A (#1127): canonical implementation on the underlying struct, trait surface preserved so callers don't churn, only delete truly dead methods.

Part A consolidated the metadata delegates onto WorkerMetadata. Part B does the equivalent consolidation for runtime state, so the atomic bookkeeping (orderings, post-increment adjustments, saturating decrements) lives in exactly one place: impl WorkerRuntime.

What changed

All edits live in model_gateway/src/worker/worker.rs:

• Delete dead Worker::reset_load. grep -rn reset_load confirmed 0 callers anywhere in the tree (trait default was a no-op {}, and the only override was in BasicWorker). Dead code removed, not deprecated.
• Promote WorkerRuntime's lifecycle methods to pub. status, set_status, revision, bump_revision were all private on impl WorkerRuntime — promoted to pub fn so callers can reach them through a single accessor without bouncing through the Worker trait.
• Add canonical counter methods on impl WorkerRuntime as the single source of truth for runtime state:
  • Health counters: consecutive_failures_increment/_reset, consecutive_successes_increment/_reset, total_pending_probes/_increment/_reset
  • Load counter: load, increment_load, new try_decrement_load() -> bool (returns false on underflow so callers can warn)
  • Routing-key load: routing_key_load, increment_routing_key_load, decrement_routing_key_load
  • Processed counter: processed_requests, increment_processed
• Rewrite the 15 counter methods on impl Worker for BasicWorker so each is a one-line forward of the form self.runtime.load().method(). BasicWorker-specific side effects (update_running_requests_metrics() after every load mutation, Metrics::set_worker_health after set_status) stay at the forwarding layer — only the atomic bookkeeping moves to WorkerRuntime.

Not touched (already Option 2-compliant):

• Category C (circuit breaker) — the six BasicWorker wrappers (circuit_breaker_state, circuit_breaker_can_execute, record_circuit_breaker_outcome, plus the is_available / record_outcome default impls and the resilience accessor) already forward to CircuitBreaker::state(), .can_execute(), .record_outcome() which own the canonical implementation in worker/circuit_breaker.rs. No changes needed.
• The GrpcWorker Worker impl in bindings/golang/src/policy.rs — this one carries its own raw atomics with comments saying the counters are intentionally no-ops ("FFI workers don't run the state machine — these counters are unused"). Consolidating it onto Arc<WorkerRuntime> is a bigger cross-crate change and belongs to a separate PR. This PR leaves it alone and verifies it still compiles against the trimmed trait.

Why

Before this PR, every counter mutation on BasicWorker went through a ~6-line block that reached into WorkerRuntime's internal atomic fields directly (e.g. self.runtime.load().consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1). Consequences:

1. The struct fields had to be pub for the trait impl to touch them, which leaked implementation details.
2. Each atomic ordering choice (AcqRel for health counters, Relaxed for load/processed, Release/Acquire for status) was duplicated at every call site. If we ever need to reason about memory ordering or benchmark alternatives, we'd have to hunt 15+ call sites.
3. The post-increment + 1 and the saturating fetch_update idiom lived inside the trait impl.

Collapsing the implementation onto WorkerRuntime keeps precisely one copy of each atomic ordering choice, and lets future Worker implementations (e.g. the FFI GrpcWorker) adopt the canonical runtime wholesale without duplicating the memory-ordering rules.

reset_load had no callers anywhere — keeping it in the trait just added noise. Deleting dead methods keeps the trait definition honest.

How — Option 2 principle

Trait methods stay (so callers still write worker.load(), worker.increment_processed(), etc. — zero churn across routers/ and workflow/). The impl on BasicWorker becomes one-line forwards. The canonical code lives in exactly one place: impl WorkerRuntime.

Before / after on a representative method (consecutive_failures_increment):

// Before — BasicWorker reaches into WorkerRuntime internals
fn consecutive_failures_increment(&self) -> usize {
    self.runtime
        .load()
        .consecutive_failures
        .fetch_add(1, Ordering::AcqRel)
        + 1
}

// After — canonical impl on WorkerRuntime, BasicWorker forwards
fn consecutive_failures_increment(&self) -> usize {
    self.runtime.load().consecutive_failures_increment()
}

decrement_load is the one method with non-trivial BasicWorker-side behavior (it logs when the counter is already zero). The new split is:

// On WorkerRuntime — pure atomic bookkeeping
pub fn try_decrement_load(&self) -> bool {
    self.load_counter
        .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |c| c.checked_sub(1))
        .is_ok()
}

// On BasicWorker — logging + metric export stay here
fn decrement_load(&self) {
    if !self.runtime.load().try_decrement_load() {
        tracing::warn!(
            worker_url = %self.metadata.spec.url,
            "Attempted to decrement load counter that is already at 0"
        );
    }
    self.update_running_requests_metrics();
}

Test plan

• cargo check -p smg — clean.
• cargo clippy -p smg --all-targets -- -D warnings — clean.
• cargo clippy -p smg-golang --all-targets -- -D warnings — clean (GrpcWorker in bindings/golang still compiles against the trimmed trait).
• cargo check -p smg-python — clean.
• cargo test -p smg --lib — 538 passed; 0 failed; 4 ignored (no test changes were required — the existing worker counter tests in worker::worker::tests exercise the new forwarding code path transparently through the same trait methods).

Refs: .claude/plans/2026-04-14-worker-module-followup-cleanup.md (Finding 4 Part B)

Summary by CodeRabbit

• Refactor
  • Enhanced worker runtime infrastructure with improved health-check monitoring and state tracking capabilities
  • Improved load management with safer operations to prevent counter overflow
  • Expanded internal counters for tracking worker failures, successes, and pending operations

[claude]

Clean refactoring — reviewed all changes. Atomic orderings preserved, return values match, side effects (metrics updates) stay at the BasicWorker forwarding layer, reset_load removal is safe (zero callers confirmed). Fields are now properly encapsulated behind WorkerRuntime's public API. LGTM.

[coderabbitai]

📝 Walkthrough

Walkthrough

The PR refactors WorkerRuntime by exposing previously internal methods and adding new public APIs for health-check counters, probe tracking, and load management. Removed reset_load from the Worker trait and updated BasicWorker to delegate counter operations to new WorkerRuntime methods with improved safety.

Changes

Cohort / File(s)

Summary

WorkerRuntime API Refactoring & Health Counters model_gateway/src/worker/worker.rs

Exposed WorkerRuntime lifecycle accessors (status, set_status, revision, bump_revision) as public methods. Added new public counter APIs for health tracking (consecutive_failures_increment/reset, consecutive_successes_increment/reset), probe management (total_pending_probes with increment/reset variants), load operations (load, increment_load, try_decrement_load() with saturation safety), and routing-key load tracking. Removed reset_load from Worker trait. Updated BasicWorker to delegate counter operations to new WorkerRuntime methods and use try_decrement_load() for checked decrement behavior.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~22 minutes

Possibly related PRs

• fix(realtime-api): worker health tracking in websocket session #725 — Adds worker.record_outcome(success) call that directly integrates with the newly exposed WorkerRuntime health-tracking and consecutive counter APIs.
• refactor(anthropic): simplify worker lifecycle in Anthropic router #476 — Removes calls to per-worker load and circuit-breaker functions alongside LoadTrackingStream, touching the same worker load/counter interaction surface.
• refactor(worker): replace healthy AtomicBool with status AtomicU8 #1101 — Replaces AtomicBool health with status AtomicU8 and introduces status()/set_status() methods, which this PR also exposes publicly.

Suggested reviewers

• CatherineSue
• key4ng
• whybeyoung

Poem

🐰 The worker now shows its inner light,
Counters exposed, APIs shining bright!
Health tracking blooms, no secrets to hide,
Safe load decrement—with checked stride.
A refactored grace, the trait's stride aligned! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 5.41% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'refactor(worker): canonical runtime API on WorkerRuntime (Finding 4 Part B)' accurately reflects the main objective: consolidating runtime-state bookkeeping onto WorkerRuntime with a canonical public API.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch refactor/worker-runtime-canonical

---

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Warning

Gemini is experiencing higher than usual traffic and was unable to create the review. Please try again in a few hours by commenting /gemini review.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@model_gateway/src/worker/worker.rs`:
- Around line 594-596: The new public mutator set_status(WorkerStatus) (and the
other public mutators around 642-654) introduce a second write path that
bypasses BasicWorker's forwarding layer and skips Metrics::set_worker_health,
update_running_requests_metrics(), and decrement_load() underflow checks; change
these methods to non-public (remove pub) or otherwise restrict visibility so
callers must go through BasicWorker's runtime forwarding API, and if internal
updates are needed ensure they invoke the same helpers
(Metrics::set_worker_health, update_running_requests_metrics, decrement_load)
used by the forwarding layer so metrics and runtime state remain consistent.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 78c85359-fae4-4fdb-b496-db3f08aa0fd6

📥 Commits

Reviewing files that changed from the base of the PR and between 9e9c0c3 and c04cac3.

📒 Files selected for processing (1)

• model_gateway/src/worker/worker.rs