fix(credit-router): define register/release_child_routing methods (latent AttributeError on FORK DAGs)

[ajcasagrande]

Summary

• Fixes a latent AttributeError in main's DAG runtime: BranchOrchestrator calls self._sticky_router.register_child_routing(...) and release_child_routing(...) at five sites (_branch_orchestrator_spawn.py:197,257, branch_orchestrator.py:391,455,467) but those methods were never defined on StickyCreditRouter. Test suites mock them so CI is green — the moment a real FORK-mode DAG runs in prod, the orchestrator raises AttributeError at the first child dispatch.
• Refactors _sticky_sessions from dict[str, str] to dict[str, _StickyEntry] so the parent's sticky binding can outlive its own final turn while DAG children pinning to its worker are still in-flight (ref_count + parent_final_seen lifecycle).
• Adds a subagent_type: str \| None field on ConversationBranchInfo (SPAWN-only) so loaders can pass through upstream-recorded values like "Explore" or "Codex Subagent" verbatim without enum-narrowing.
• Optimizes send_credit: the dominant single-turn-root-no-forks workload (most AIPerf benchmarks) no longer pays an extra dict.get per credit — the final-turn lifecycle gates on sticky_entry is not None directly instead of falling back to a second dict lookup.

Hot-path impact

Hot-path performance was a deliberate constraint (per the StickyCreditRouter docstring: "Methods are intentionally large/inlined to avoid function call overhead"). Per-credit cost analysis:

• Sticky-hit common case: +1 attribute deref (sticky_entry.worker_id), +1 or-eval on routing key. ~15-25ns. Below noise at typical AIPerf throughputs.
• Single-turn root no-forks (dominant non-DAG workload): no _StickyEntry allocation, no second dict.get (thanks to the optimization). Effectively unchanged from main.
• Session creation (cold-ish): _StickyEntry(slots=True) allocation instead of dict[k] = str. ~200ns added, once per multi-turn session.
• Atomicity invariant preserved: zero new awaits in send_credit.

Lifecycle (new semantics)

                                       parent_correlation_id?  is_final_turn?  has_forks?
                                       ───────────────────────  ──────────────  ──────────
Single-turn root, no forks              None                    True            False     → no entry created, no allocation
Single-turn root, declares spawns       None                    True            True      → entry created + survives (register_child_routing finds it)
Multi-turn root (any non-final turn)    None                    False           any       → entry created/reused
Multi-turn root, final turn, no forks   None                    True            False     → entry decremented + popped
DAG child credit                        set                     any             any       → pins to parent's worker via routing_key, lifecycle block skipped

register_child_routing(parent_corr) increments refcount; logs a warning if the parent has no entry (graceful degradation — child routes via least-loaded, loses prefix-cache locality but stays correct).

release_child_routing(parent_corr) decrements refcount; evicts the entry only when both ref_count <= 0 AND parent_final_seen.

subagent_type field

Added because the Weka kv-cache-tester trace schema already carries subagent_type: str per recorded subagent (real corpus values: "Explore", "Codex Subagent", etc. — free-form, set by upstream Claude Code proxy). Typing the new aiperf-side field as str \| None rather than a narrow enum lets loaders pass these through verbatim without filter/map logic and without rejecting new agent types as corpora evolve.

The field validator still enforces the SPAWN-only invariant (FORK children inherit the parent's role, so they have no separate classification).

Test plan

• tests/unit/credit/test_sticky_router.py::TestStickyCreditRouterChildRoutingRefcount (6 tests): register/release refcount semantics, eviction conditions, missing-parent warning, no-op-on-unknown-parent.
• tests/unit/credit/test_sticky_router.py::TestStickyCreditRouterSendCreditFinalTurnLifecycle (4 tests): single-turn root no-forks (no entry created), final turn with forks (entry survives for register_child_routing), DAG-child final turn (parent untouched), multi-turn final turn (entry popped on no-forks).
• tests/unit/common/models/test_branch_model.py::TestConversationBranchInfoSubagentType (6 tests): SPAWN accepts "Explore" and arbitrary "Codex Subagent", FORK rejects non-None, round-trip.
• Full unit suite: 12,689 passed, 79 skipped, 0 failed.
• Component integration (DAG + sticky-routing): test_dag_adversarial_timing_modes, test_dag_combined_pathology, test_sticky_routing — 64 passed.
• pre-commit run --files <touched> clean.
• Real FORK-DAG benchmark run (any reviewer with hardware): verify the orchestrator no longer raises AttributeError on child dispatch.

Guardrails

StickyCreditRouter.send_credit grandfathered in tools/ergonomics_baseline.json (function-size 90 > 80) and tools/ruff_baseline.json (C901 cyclomatic 12 > 10, PLR0912 branches 13 > 12). The hot path is intentionally inlined per the class docstring; extracting helpers would add per-credit function-call overhead. Baseline additions restricted to exactly these two new entries; unrelated stale entries left untouched.

Related branches

• ajc/dag4 (1 commit ahead of main, dated May) had an earlier version of this fix bundled with the Weka loader stack and a narrow SubagentType enum — appears abandoned, no merges from main since.
• ajc/dag5 (17 commits ahead of main, actively being worked) is the successor of PR feat(dag): conversation DAG benchmarks (dag_jsonl) — FORK / SPAWN / f… #891 and still carries the same latent AttributeError. Happy to retarget this PR onto dag5 if you'd prefer to fold the fix into that review cycle.

🤖 Generated with Claude Code

Summary by CodeRabbit

• New Features

  • Added support for optional subagent_type specification in conversation branches for improved branch configuration.
  • Enhanced routing system with reference counting to better manage child session lifecycles in multi-branch conversations.
  • Introduced methods for managing child session routing registration and release.

• Tests

  • Added validation tests for branch configuration constraints.
  • Expanded routing tests for complex session lifecycle scenarios.

[github-actions]

Try out this PR

Quick install:

pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@2c53ea483b3d9e48fd006a82e68385a56108f5a6

Recommended with virtual environment (using uv):

uv venv --python 3.12 && source .venv/bin/activate
uv pip install --upgrade --force-reinstall git+https://github.com/ai-dynamo/aiperf.git@2c53ea483b3d9e48fd006a82e68385a56108f5a6

---

Last updated for commit: 2c53ea4 • Browse code

[coderabbitai]

Walkthrough

This PR introduces subagent-type metadata to conversation branching with SPAWN-mode validation, and refactors sticky-session state management to track refcounts and parent final-turn bookkeeping, enabling safe child routing past parent completion through new lifecycle methods.

Changes

Sticky Router Enhancement and Branch Model Metadata

Layer / File(s)	Summary
Branch model subagent_type field and validation src/aiperf/common/models/branch.py, tests/unit/common/models/test_branch_model.py	ConversationBranchInfo adds optional subagent_type field with field validator rejecting non-None values when mode is FORK. Test suite validates default handling, string acceptance, rejection for FORK, and model round-tripping.
Sticky entry data structure src/aiperf/credit/sticky_router.py	Introduces _StickyEntry dataclass to track sticky-session ownership, ref count, and parent final-turn state for lifecycle-based eviction.
Sticky router storage refactoring src/aiperf/credit/sticky_router.py, tests/unit/credit/test_sticky_router.py, tests/unit/timing/test_race_conditions.py	Replaces routing_key -> worker_id mapping with routing_key -> _StickyEntry. Updates routing lookup to derive keys from parent or session correlation ID. Test assertions and setup across multiple scenarios updated to reference _StickyEntry.worker_id.
Sticky router lifecycle and child routing methods src/aiperf/credit/sticky_router.py	send_credit now creates/binds entries for non-final or fork scenarios, sets parent_final_seen and decrements refcount on final turns, and evicts only when refcount reaches zero with no forks. New register_child_routing and release_child_routing methods manage refcount for parent correlation IDs, evicting after parent final turn is observed.
Comprehensive sticky router refcount and final-turn lifecycle tests tests/unit/credit/test_sticky_router.py	New test suite covering _StickyEntry refcount lifecycle for child routing including eviction conditions, and detailed send_credit final-turn paths for single-turn roots, fork scenarios, DAG child short-circuits, and multi-turn sessions.
Code quality baseline updates tools/ergonomics_baseline.json, tools/ruff_baseline.json	Ergonomics and Ruff baselines record violations for expanded StickyCreditRouter.send_credit method (function-size, C901, PLR0912).

---

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

---

Poem

🐰 A branch now knows its type so fine,
Sticky routers count with designs divine,
Entries dance through parent and child,
Refcounts keep sessions reconciled,
DAG routing, finally unfiled!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

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

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title accurately and specifically describes the main fix: defining missing register/release_child_routing methods to address a latent AttributeError in FORK DAGs.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

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

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

🧹 Nitpick comments (1)

tests/unit/credit/test_sticky_router.py (1)

48-48: 💤 Low value

Consider using next(iter()) for single-element access.

Static analysis suggests this minor optimization to avoid creating an intermediate list.

Suggested change

-        assert list(router._sticky_sessions.values())[0].worker_id == "worker-2"
+        assert next(iter(router._sticky_sessions.values())).worker_id == "worker-2"

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tests/unit/credit/test_sticky_router.py` at line 48, Replace the intermediate
list creation when asserting the single sticky session: instead of converting
router._sticky_sessions.values() to a list and indexing [0], use
next(iter(router._sticky_sessions.values())) to retrieve the first (and only)
element; update the assertion that checks its worker_id (the existing assertion
referencing router._sticky_sessions.values() and worker_id == "worker-2") to use
next(iter(...)) for more efficient single-element access.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tests/unit/credit/test_sticky_router.py`:
- Line 48: Replace the intermediate list creation when asserting the single
sticky session: instead of converting router._sticky_sessions.values() to a list
and indexing [0], use next(iter(router._sticky_sessions.values())) to retrieve
the first (and only) element; update the assertion that checks its worker_id
(the existing assertion referencing router._sticky_sessions.values() and
worker_id == "worker-2") to use next(iter(...)) for more efficient
single-element access.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 4cae323f-a088-49ab-8bf3-c12659dac25f

📥 Commits

Reviewing files that changed from the base of the PR and between bb6f421 and 541b93f.

📒 Files selected for processing (7)

• src/aiperf/common/models/branch.py
• src/aiperf/credit/sticky_router.py
• tests/unit/common/models/test_branch_model.py
• tests/unit/credit/test_sticky_router.py
• tests/unit/timing/test_race_conditions.py
• tools/ergonomics_baseline.json
• tools/ruff_baseline.json

[codecov]

Codecov Report

❌ Patch coverage is 96.15385% with 2 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/aiperf/credit/sticky_router.py	95.55%	0 Missing and 2 partials ⚠️

📢 Thoughts on this report? Let us know!

[dynamo-ops]

SPAWN child credits also carry parent_correlation_id, so this unconditionally keys them under the parent while the final-turn cleanup skips all parent_correlation_id credits, leaking multi-turn SPAWN sticky sessions and forcing sibling SPAWN children onto one worker. Fix: only use parent_correlation_id and the refcount lifecycle for credit.branch_mode == ConversationBranchMode.FORK; route and clean up SPAWN children by their own x_correlation_id.