Skip to content

feat(dispatch): wire context models into handler dispatch path [OMN-984] - #225

Merged
jonahgabriel merged 4 commits into
mainfrom
jonah/omn-984-wire-context-models-into-handler-dispatch-path-follow-up
Dec 20, 2025
Merged

jonahgabriel merged 4 commits into
mainfrom
jonah/omn-984-wire-context-models-into-handler-dispatch-path-follow-up

Conversation

@jonahgabriel

@jonahgabriel jonahgabriel commented Dec 20, 2025 •

Copy link
Copy Markdown
Contributor

Summary

Changes

New Protocol (protocols/handler/protocol_handler_context.py):

  • Defines common context interface: correlation_id, envelope_id, trace_id, span_id
  • All context models implement this protocol

MessageDispatchEngine Updates:

  • register_handler() now requires node_kind parameter
  • _build_handler_context() factory creates context based on node_kind
  • _execute_handler() passes context to handlers
  • RUNTIME_HOST rejected as handler node_kind

Context Types by Node Kind:

Node Kind Context Has now Has retry_attempt
EFFECT ModelEffectContext ✅ ✅
COMPUTE ModelComputeContext ❌ (pure) ❌
REDUCER ModelReducerContext ❌ (pure) ❌
ORCHESTRATOR ModelOrchestratorContext ✅ ❌

Breaking Changes

  • register_handler() now requires node_kind: EnumNodeKind parameter
  • Handlers must accept (envelope, context) signature

Test Plan

  • All 73 dispatch engine tests pass
  • New context injection tests verify correct context types
  • RUNTIME_HOST rejection test
  • Correlation ID preservation test
  • mypy strict mode passes on all new code

Summary by CodeRabbit

Release Notes

  • New Features

    • Handlers now receive structured execution context containing correlation_id, envelope_id, trace_id, and span_id identifiers.
    • Context attributes vary by handler type (effect, compute, reducer, orchestrator).
  • Breaking Changes

    • Handler signatures now require an execution context parameter alongside message envelopes.
    • Handler registration requires specifying the handler type during setup.

✏️ Tip: You can customize this high-level summary in your review settings.

Integrates context models from OMN-948 (PR #224) into MessageDispatchEngine
so handlers receive appropriate context during execution.

Changes:
- Add ProtocolHandlerContext protocol for common context interface
- Update register_handler() to require node_kind parameter
- Add _build_handler_context() factory for context creation
- Update _execute_handler() to pass context to handlers
- Reject RUNTIME_HOST as handler node_kind

Context types by node_kind:
- EFFECT: ModelEffectContext (with now, retry_attempt)
- COMPUTE: ModelComputeContext (no now - pure function)
- REDUCER: ModelReducerContext (no now - pure function)
- ORCHESTRATOR: ModelOrchestratorContext (with now)

Breaking change: register_handler() now requires node_kind,
handlers must accept (envelope, context) signature.
@coderabbitai

coderabbitai Bot commented Dec 20, 2025 •

Copy link
Copy Markdown
Contributor

Warning

Rate limit exceeded

@jonahgabriel has exceeded the limit for the number of commits or files that can be reviewed per hour. Please wait 12 minutes and 54 seconds before requesting another review.

⌛ How to resolve this issue?

After the wait time has elapsed, a review can be triggered using the @coderabbitai review command as a PR comment. Alternatively, push new commits to this PR.

We recommend that you space out your commits to avoid hitting the rate limit.

🚦 How do rate limits work?

CodeRabbit enforces hourly rate limits for each developer per organization.

Our paid plans have higher rate limits than the trial, open-source and free plans. In all cases, we re-allow further reviews after a brief timeout.

Please see our FAQ for further information.

📥 Commits

Reviewing files that changed from the base of the PR and between 04621bb and c03c6ea.

📒 Files selected for processing (1)
  • src/omnibase_core/runtime/message_dispatch_engine.py (9 hunks)

Walkthrough

The PR adds a new ProtocolHandlerContext protocol that standardizes handler execution context with correlation and tracing fields. The message dispatch engine is updated to construct node-specific context objects and pass them to handlers during invocation, enabling context propagation across EFFECT, COMPUTE, REDUCER, and ORCHESTRATOR node types.

Changes

Cohort / File(s) Summary
Protocol Handler Context Exports
src/omnibase_core/protocols/__init__.py, src/omnibase_core/protocols/handler/__init__.py
Re-exports ProtocolHandlerContext through package initializers to expose the new context protocol as part of the public protocols API.
Protocol Handler Context Definition
src/omnibase_core/protocols/handler/protocol_handler_context.py
Introduces ProtocolHandlerContext, a runtime-checkable protocol defining four read-only properties: correlation_id, envelope_id, trace_id, and span_id for standardizing handler execution context.
Dispatch Engine Context Integration
src/omnibase_core/runtime/message_dispatch_engine.py
Updates handler registration to require node_kind parameter, adds _build_handler_context method to construct node-specific contexts (ModelEffectContext, ModelComputeContext, ModelReducerContext, ModelOrchestratorContext), and modifies HandlerFunc signature to accept both envelope and context.
Handler Context Test Coverage
tests/unit/runtime/test_message_dispatch_engine.py
Extends test suite to verify context propagation for all node kinds, validate correlation/tracing field mapping from envelope to context, test node-kind-specific context attributes, and enforce RUNTIME_HOST rejection.

Sequence Diagram

sequenceDiagram
    participant Client
    participant DispatchEngine as Message<br/>Dispatch Engine
    participant ContextBuilder as Context<br/>Builder
    participant Handler

    Client->>DispatchEngine: dispatch(envelope)
    activate DispatchEngine
    DispatchEngine->>DispatchEngine: lookup handler + node_kind
    DispatchEngine->>ContextBuilder: _build_handler_context(envelope, node_kind)
    activate ContextBuilder
    ContextBuilder->>ContextBuilder: select context type<br/>(EFFECT/COMPUTE/REDUCER/ORCHESTRATOR)
    ContextBuilder->>ContextBuilder: map envelope fields<br/>(correlation_id, trace_id, etc.)
    ContextBuilder-->>DispatchEngine: context: ProtocolHandlerContext
    deactivate ContextBuilder
    DispatchEngine->>Handler: invoke(envelope, context)
    activate Handler
    Handler->>Handler: process with context
    Handler-->>DispatchEngine: result
    deactivate Handler
    DispatchEngine-->>Client: result
    deactivate DispatchEngine
Loading

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

  • message_dispatch_engine.py: Review signature changes to HandlerFunc and register_handler, validate the _build_handler_context logic for correct context type selection and field mapping based on node_kind, and ensure RUNTIME_HOST rejection is properly enforced.
  • test_message_dispatch_engine.py: Verify comprehensive coverage of context propagation across all node kinds, check that context-specific attributes (e.g., retry_attempt in EFFECT, absence of now in COMPUTE/REDUCER) are correctly validated, and confirm test assertions align with new dispatch semantics.

Possibly related PRs

  • PR #224: Introduces the node-specific context model types (ModelEffectContext, ModelOrchestratorContext, ModelReducerContext, ModelComputeContext) that this PR depends on for context instantiation.
  • PR #145: Modifies the protocols package exports; both PRs update the public API surface of the protocols module.
  • PR #223: Modifies the dispatch engine's handler invocation surface; both PRs alter HandlerFunc type and _execute_handler execution flow.

Poem

🐰 A context is born, so pristine and bright,
Correlation and traces, now woven just right,
Four flavors of handlers—Effect, Compute, Redux—
Each knows where it came from, no more handler flux! ✨

Pre-merge checks and finishing touches

❌ Failed checks (1 warning)
Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 52.43% which is insufficient. The required threshold is 80.00%. You can run @coderabbitai generate docstrings to improve docstring coverage.
✅ 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 accurately summarizes the main change: integrating context models into the handler dispatch path via the MessageDispatchEngine component.

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

@claude

claude Bot commented Dec 20, 2025

Copy link
Copy Markdown

PR Review: Context Models Integration into Message Dispatch

Overall Assessment

LGTM with minor suggestions ✅

This is an excellent implementation that cleanly integrates the node-specific context models from PR #224 into the message dispatch path. The design is architecturally sound, well-tested, and follows ONEX conventions.


Strengths 💪

1. Architectural Soundness

  • ✅ Protocol-driven design: ProtocolHandlerContext provides clean abstraction for type-safe context handling
  • ✅ Node-kind dispatch: Context selection based on EnumNodeKind aligns perfectly with ONEX four-node architecture
  • ✅ Purity enforcement: COMPUTE/REDUCER contexts correctly exclude now field (pure functions), while EFFECT/ORCHESTRATOR include it (time-dependent operations)
  • ✅ Breaking change clarity: Handler signature change is clearly documented and justified

2. Code Quality

  • ✅ Thread safety: All context models are frozen/immutable (per CLAUDE.md threading guidelines)
  • ✅ Type safety: Proper use of PEP 604 union syntax (UUID | None vs Optional[UUID])
  • ✅ Error handling: Comprehensive validation with structured ModelOnexError exceptions
  • ✅ Documentation: Excellent docstrings with examples, design rationale, and node-specific field tables

3. Test Coverage

  • ✅ 73 tests passing: All dispatch engine tests updated and passing
  • ✅ New test scenarios: Context injection, RUNTIME_HOST rejection, correlation ID preservation
  • ✅ Test pattern consistency: All test handlers updated to (envelope, context) signature
  • ✅ Both sync and async: Proper coverage of executor-based sync handler dispatch

4. ONEX Conventions Compliance

  • ✅ Follows CLAUDE.md: Protocol-driven DI, structured errors, node architecture patterns
  • ✅ Causality tracking: Proper propagation of correlation_id, envelope_id, trace_id, span_id
  • ✅ No backwards compatibility hacks: Clean breaking change (good per CLAUDE.md anti-patterns)

Suggestions for Improvement 🔍

1. Correlation ID Fallback (Minor)

Location: message_dispatch_engine.py:1337

Current implementation generates a fallback UUID when correlation_id is None. Consider logging a warning for observability, as missing correlation IDs could indicate upstream issues in production.

2. RUNTIME_HOST Rejection - Error Code Choice (Nitpick)

Location: message_dispatch_engine.py:1383

The defensive check uses INVALID_STATE, but INVALID_PARAMETER might be more semantically accurate. Low impact since this code path should never execute.

3. Test Coverage - Retry Attempt Mapping

I don't see explicit tests verifying that envelope.retry_count correctly maps to ModelEffectContext.retry_attempt. Consider adding a test case to verify this mapping.

4. Protocol Compliance Verification (Nice-to-Have)

Consider adding a runtime assertion in _build_handler_context to verify protocol compliance during development.


Security & Performance ✅

Security

  • ✅ No security concerns: Immutable contexts prevent state tampering
  • ✅ UUID generation: Uses secure uuid4() for fallback correlation IDs
  • ✅ No injection risks: Frozen models prevent field injection attacks

Performance

  • ✅ Minimal overhead: Context construction is lightweight
  • ✅ No blocking operations: Context building is synchronous and fast
  • ✅ Thread-safe: Frozen models avoid synchronization overhead
  • ✅ Async-friendly: Works seamlessly with existing async dispatch

Breaking Changes - Migration Path 📋

Impact: All existing handlers must update signature from handler(envelope) to handler(envelope, context).

Recommendations:

  1. ✅ Breaking change is documented in PR description
  2. Consider adding migration example to CHANGELOG showing before/after handler signature
  3. Consider grepping codebase for handler registrations to identify affected code

Final Verdict

This PR represents high-quality engineering:

  • Clean architecture aligned with ONEX principles
  • Comprehensive test coverage (73 tests, all scenarios covered)
  • Excellent documentation and type safety
  • Proper thread safety and immutability
  • Breaking change is justified and well-communicated

The minor suggestions above are optional improvements - this PR is ready to merge as-is.

Recommendation: APPROVE ✅


Testing Note: Verify all 73 dispatch engine tests pass in CI before merge.

Great work! 🎉

…ntext building [OMN-984]

- Add warning log when correlation_id is None and fallback UUID is generated
- Change error code from INVALID_STATE to INVALID_PARAMETER for RUNTIME_HOST rejection
- Add runtime assertion for protocol compliance verification (debug-only)
- Refactor _build_handler_context to use single return point
- Add test verifying retry_count → retry_attempt mapping for EFFECT handlers
@claude

claude Bot commented Dec 20, 2025

Copy link
Copy Markdown

Code Review: Context Models Integration (PR #225)

Summary

This PR successfully integrates node-specific context models into the message dispatch path. The implementation is well-architected, follows ONEX conventions, and includes comprehensive test coverage.


✅ Strengths

1. Excellent Architecture & Design

  • Protocol-driven approach: ProtocolHandlerContext provides a clean interface for type-safe context handling
  • Node-specific contexts: Correctly maps EnumNodeKind to appropriate context types (Effect, Compute, Reducer, Orchestrator)
  • Purity enforcement: COMPUTE and REDUCER contexts correctly omit now field to maintain deterministic behavior
  • Immutability: All context models use frozen=True, ensuring thread safety (critical per CLAUDE.md threading guidelines)

2. Strong Documentation

  • Comprehensive docstrings in protocol_handler_context.py (88 lines!)
  • Clear rationale for design decisions (e.g., why COMPUTE/REDUCER lack now)
  • Excellent inline comments in _build_handler_context() explaining context selection logic
  • PR description clearly documents breaking changes and migration path

3. Robust Error Handling

  • ✅ Validates node_kind parameter type with helpful error message
  • ✅ Rejects RUNTIME_HOST node kind with clear guidance (message_dispatch_engine.py:448-455)
  • ✅ Handles missing correlation_id gracefully with warning + fallback UUID (message_dispatch_engine.py:1337-1342)
  • ✅ Uses ModelOnexError with proper EnumCoreErrorCode (follows CLAUDE.md conventions)

4. Comprehensive Test Coverage

  • All 73 existing tests updated with new node_kind parameter
  • New tests verify context injection correctness
  • Tests cover edge cases (RUNTIME_HOST rejection, correlation ID preservation)
  • Follows project's pytest-xdist parallel testing strategy

5. Type Safety

  • @runtime_checkable protocol enables isinstance checks when needed
  • PEP 604 union syntax (UUID | None) used consistently ✅
  • Handler signature updated to require context: HandlerFunc = Callable[[ModelEventEnvelope[Any], ProtocolHandlerContext], Any | Awaitable[Any]]

🔍 Issues Found

CRITICAL: Missing from_attributes=True Verification

Per CLAUDE.md guidelines, value objects nested in Pydantic models used with pytest-xdist require from_attributes=True to prevent class identity issues in parallel execution.

Verification needed:

# All context models have this ✅
model_config = ConfigDict(frozen=True, extra="forbid", from_attributes=True)

I verified:

  • ✅ ModelComputeContext (line 78)
  • ✅ ModelEffectContext (line 92)
  • ✅ ModelReducerContext (line 70)
  • ✅ ModelOrchestratorContext (line 89-93)

Status: ✅ PASSED - All context models correctly implement from_attributes=True


💡 Suggestions for Improvement

1. Consider Adding Validation for retry_attempt Mapping (Low Priority)

In _build_handler_context(), the retry count is mapped from envelope.retry_count:

retry_attempt=envelope.retry_count,  # Line 1360

Suggestion: Add a comment or assertion that retry_count >= 0, since ModelEffectContext.retry_attempt has ge=0 constraint. This prevents potential Pydantic validation errors if envelope data is corrupted.

# Ensure retry_count is non-negative (ModelEffectContext requires ge=0)
retry_attempt = max(0, envelope.retry_count)

2. Potential Optimization: Cache Context Creation (Performance - Future)

_build_handler_context() creates new context instances on every dispatch. For high-throughput systems, consider caching contexts with identical IDs (correlation_id, envelope_id, trace_id, span_id).

Tradeoff: Adds complexity vs. minor GC pressure reduction. Recommendation: Profile first before optimizing.

3. Test Enhancement: Verify Context Fields Propagate Correctly (Medium Priority)

Current tests verify context type, but could strengthen by verifying field values:

async def test_context_fields_propagate_from_envelope():
    """Verify correlation_id, trace_id, span_id propagate to context."""
    test_envelope = ModelEventEnvelope(
        correlation_id=known_uuid,
        trace_id=known_trace_uuid,
        ...
    )
    
    # Dispatch and capture context in handler
    async def capturing_handler(envelope, context):
        assert context.correlation_id == known_uuid
        assert context.trace_id == known_trace_uuid
        ...

This would catch field mapping bugs early.

4. Documentation: Add Migration Guide Link (Low Priority)

Since this is a breaking change (handler signature changed), consider adding a link to migration docs in the PR description or in message_dispatch_engine.py docstring:

# Breaking Change (v0.4.0): Handlers now require (envelope, context) signature
# See: docs/guides/MIGRATING_TO_CONTEXT_HANDLERS.md

🔒 Security Review

✅ No security concerns identified:

  • Context models are immutable (prevents tampering)
  • No sensitive data exposed in context (only IDs)
  • Correlation ID fallback uses uuid4() (cryptographically random)
  • No injection vulnerabilities (all inputs validated via Pydantic)

⚡ Performance Considerations

✅ Efficient implementation:

  • Context creation is lightweight (no I/O, no blocking calls)
  • Sync handlers still run in executor (prevents blocking event loop)
  • Async handlers await directly (no unnecessary context switches)

Potential concern: The if __debug__ assertion at message_dispatch_engine.py:1398-1402 runs on every dispatch in development. This is acceptable for dev/test but ensure production runs with python -O to skip assertions.


📊 Test Coverage Assessment

✅ Excellent coverage (73 tests updated + new tests):

  • ✅ All four node kinds (EFFECT, COMPUTE, REDUCER, ORCHESTRATOR)
  • ✅ RUNTIME_HOST rejection
  • ✅ Sync and async handlers
  • ✅ Missing correlation_id fallback
  • ✅ Invalid node_kind parameter

Missing test: Handler receives correct retry_attempt value from envelope.retry_count. Suggested test:

async def test_effect_context_retry_attempt_from_envelope():
    envelope = create_test_envelope(retry_count=3)
    
    async def check_retry(env, ctx):
        assert isinstance(ctx, ModelEffectContext)
        assert ctx.retry_attempt == 3
    
    # ... register and dispatch

🎯 Final Verdict

Recommendation: ✅ APPROVE with minor suggestions

This is high-quality, production-ready code that:

  • Follows ONEX architecture principles
  • Adheres to CLAUDE.md conventions (Poetry, type safety, error handling, thread safety)
  • Includes comprehensive tests and documentation
  • Implements a clean, extensible design

The suggestions above are non-blocking improvements that can be addressed in follow-up PRs.

Blocking concerns: NONE ✅


📝 Action Items (Optional - Post-Merge)

  1. Consider adding migration guide document for handler signature changes
  2. Add test for retry_attempt field propagation
  3. Add test verifying all context fields propagate from envelope
  4. Profile context creation in high-throughput scenarios (if performance issues arise)

Great work on this PR! The context injection architecture is well-designed and sets a solid foundation for future enhancements. 🚀


Reviewed by: Claude Sonnet 4.5 (Automated Code Review)
Review Date: 2025-12-20
PR: #225 (OMN-984)

… tests [OMN-984]

- Add max(0, retry_count) validation to prevent Pydantic ge=0 constraint violation
- Add test verifying context fields propagate correctly from envelope
- Add test for null trace_id/span_id handling in context propagation
@claude

claude Bot commented Dec 20, 2025

Copy link
Copy Markdown

PR Review: feat(dispatch): wire context models into handler dispatch path [OMN-984]

Summary

This PR successfully integrates context models from PR #224 into the MessageDispatchEngine, enabling handlers to receive appropriate execution context based on their node_kind. The implementation is well-designed and follows ONEX architecture principles.


✅ Strengths

1. Excellent Protocol Design

The ProtocolHandlerContext protocol is well-crafted:

  • Clear separation of common fields (correlation_id, envelope_id, trace_id, span_id) vs node-specific fields
  • @runtime_checkable decorator enables structural typing
  • Comprehensive docstrings with design rationale and examples
  • Thread-safe by design (all implementations are frozen/immutable)

2. Type Safety & Validation

# Strong validation at registration time (lines 435-454)
if node_kind == EnumNodeKind.RUNTIME_HOST:
    raise ModelOnexError(
        message="Cannot register handler with node_kind RUNTIME_HOST...",
        error_code=EnumCoreErrorCode.INVALID_PARAMETER,
    )

# Runtime assertion during development (lines 1402-1406)
if __debug__:
    assert isinstance(context, ProtocolHandlerContext), ...

3. Defensive Programming

  • Line 1353: retry_attempt = max(0, envelope.retry_count) prevents Pydantic constraint violations from corrupted data
  • Lines 1337-1343: Fallback UUID generation for missing correlation_id with warning log
  • Single return point pattern in _build_handler_context() ensures consistent protocol assertion

4. Comprehensive Test Coverage

634+ lines of test additions covering:

  • Context injection for all 4 node kinds (EFFECT, COMPUTE, REDUCER, ORCHESTRATOR)
  • RUNTIME_HOST rejection
  • Correlation ID preservation and fallback generation
  • retry_count → retry_attempt mapping
  • Context field propagation from envelope
  • Null trace_id/span_id handling

5. Breaking Change Communication

Clear documentation of breaking changes:

  • register_handler() now requires node_kind parameter
  • Handler signature changed from (envelope) to (envelope, context)

🔍 Code Quality Observations

Minor: Potential Optimization Opportunity

File: message_dispatch_engine.py:1353

retry_attempt = max(0, envelope.retry_count)

Question: Is envelope.retry_count expected to be negative in production? If not, this defensive check is good but could be paired with a warning log (similar to the correlation_id fallback) to alert operators of data integrity issues.

Suggestion (optional):

retry_attempt = envelope.retry_count
if retry_attempt < 0:
    self._logger.warning(
        "Envelope has negative retry_count=%d (envelope_id=%s), clamping to 0",
        retry_attempt, envelope.envelope_id
    )
    retry_attempt = 0

Minor: Import Organization

File: message_dispatch_engine.py:90-106

The imports are getting dense. Consider grouping by category for better readability:

# Enums
from omnibase_core.enums.enum_core_error_code import EnumCoreErrorCode
from omnibase_core.enums.enum_dispatch_status import EnumDispatchStatus
from omnibase_core.enums.enum_execution_shape import EnumMessageCategory
from omnibase_core.enums.enum_node_kind import EnumNodeKind

# Context Models
from omnibase_core.models.compute.model_compute_context import ModelComputeContext
from omnibase_core.models.effect.model_effect_context import ModelEffectContext
from omnibase_core.models.reducer.model_reducer_context import ModelReducerContext
from omnibase_core.models.orchestrator.model_orchestrator_context import ModelOrchestratorContext

# Dispatch Models
from omnibase_core.models.dispatch.model_dispatch_metrics import ModelDispatchMetrics
...

🛡️ Security & Thread Safety

✅ Thread-Safe Design

All context models use ConfigDict(frozen=True):

  • ModelComputeContext
  • ModelEffectContext
  • ModelReducerContext
  • ModelOrchestratorContext

This ensures immutability after creation, making them safe for concurrent handler execution per CLAUDE.md guidelines.

✅ No Security Concerns

  • No credential handling or sensitive data exposure
  • Proper UUID generation for fallback scenarios
  • Error messages don't leak sensitive information

📊 Performance Considerations

✅ Minimal Overhead

  • Context creation is lightweight (Pydantic model instantiation)
  • No I/O operations in _build_handler_context()
  • __debug__ assertion only runs in development mode, not production (python -O disables)

✅ Executor Pattern Preserved

Line 1442: Sync handlers still run in executor to avoid blocking event loop

return await loop.run_in_executor(None, handler, envelope, context)

🧪 Test Quality

Excellent Coverage

The test suite demonstrates:

  1. All node kinds tested: EFFECT, COMPUTE, REDUCER, ORCHESTRATOR
  2. Edge cases covered: Missing correlation_id, negative retry_count, null trace IDs
  3. Protocol compliance: All test handlers use ProtocolHandlerContext type hints
  4. Integration patterns: Tests verify context fields propagate correctly from envelope

Test Statistics

  • 73 dispatch engine tests (as stated in PR description)
  • 634 lines of test additions
  • Tests follow repository conventions (pytest markers, async/await patterns)

📝 Documentation Quality

✅ Excellent Docstrings

  • ProtocolHandlerContext: Comprehensive module-level and class-level docs
  • _build_handler_context(): Clear context type mapping table
  • register_handler(): Updated with node_kind parameter documentation and examples

✅ Architecture Alignment

The PR correctly maps context types to node kinds per ONEX Four-Node Architecture:

Node Kind Context Has now Has retry_attempt Rationale
EFFECT ModelEffectContext ✅ ✅ Side effects need time & retry tracking
COMPUTE ModelComputeContext ❌ ❌ Pure functions (deterministic/cacheable)
REDUCER ModelReducerContext ❌ ❌ Event sourcing/replay safety
ORCHESTRATOR ModelOrchestratorContext ✅ ❌ Workflow timing/deadlines

This design enforces architectural purity at the type level.


🎯 Recommendations

✅ Ready to Merge

This PR is high-quality and ready for merge with no blocking issues. The optional suggestions above are for future consideration.

Post-Merge Actions

  1. Update migration guides: Document handler signature changes for consumers
  2. Monitor correlation_id warnings: Track fallback UUID generation frequency in production logs
  3. Consider adding metric: Track context creation time if performance becomes a concern

🔬 Mypy & Type Checking

The PR maintains strict type safety:

  • HandlerFunc type alias updated with ProtocolHandlerContext parameter
  • _build_handler_context() return type explicitly typed as ProtocolHandlerContext
  • All context models properly implement the protocol via structural typing

Per CLAUDE.md requirement: "Strict Type Checking: 100% mypy strict mode compliance (0 errors across 1865 source files)" - this PR adheres to that standard.


Final Verdict

LGTM! 🚀

This is a well-architected, thoroughly tested feature that strengthens the ONEX framework's type safety and handler execution model. The breaking changes are necessary and well-communicated. The code quality, documentation, and test coverage all meet the project's high standards.

Recommended Action: Approve and merge


Reviewed by: Claude Code (Sonnet 4.5)
Review Date: 2025-12-20

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
src/omnibase_core/runtime/message_dispatch_engine.py (1)

419-427: Stale error message references old handler signature.

The error message on line 424 still references the old handler signature without the context parameter:

Expected signature: (envelope: ModelEventEnvelope[Any]) -> Any

This should be updated to reflect the new signature with context.

🔎 Proposed fix
         if handler is None or not callable(handler):
             raise ModelOnexError(
                 message=(
                     f"Handler for '{handler_id}' must be callable (function or async function). "
                     f"Got {type(handler).__name__}. "
-                    f"Expected signature: (envelope: ModelEventEnvelope[Any]) -> Any"
+                    f"Expected signature: (envelope: ModelEventEnvelope[Any], context: ProtocolHandlerContext) -> Any"
                 ),
                 error_code=EnumCoreErrorCode.INVALID_PARAMETER,
             )
📜 Review details

Configuration used: defaults

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 006c459 and 04621bb.

📒 Files selected for processing (5)
  • src/omnibase_core/protocols/__init__.py (2 hunks)
  • src/omnibase_core/protocols/handler/__init__.py (1 hunks)
  • src/omnibase_core/protocols/handler/protocol_handler_context.py (1 hunks)
  • src/omnibase_core/runtime/message_dispatch_engine.py (9 hunks)
  • tests/unit/runtime/test_message_dispatch_engine.py (39 hunks)
🧰 Additional context used
📓 Path-based instructions (2)
**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

**/*.py: Always call super().init(container) in node class constructors to initialize base class boilerplate
Use ModelOnexError with EnumCoreErrorCode for all error handling instead of generic Exception
Use protocol names with container.get_service() for dependency injection, never use concrete class names
Use PEP 604 union syntax (str | None) instead of Optional or Union types
Use @standard_error_handling decorator from omnibase_core.decorators.error_handling to eliminate try/catch boilerplate
Do NOT share node instances across threads - use thread-local instances or separate instances per thread
ORCHESTRATOR nodes CANNOT return typed results - they can only emit events and intents, with result=None required
Use EnumNodeKind for architectural role classification (EFFECT, COMPUTE, REDUCER, ORCHESTRATOR, RUNTIME_HOST) and EnumNodeType for implementation type discovery
Use FileRegistry.load() or FileRegistry.load_all() for contract loading with proper error handling for FILE_NOT_FOUND, FILE_READ_ERROR, CONFIGURATION_PARSE_ERROR, CONTRACT_VALIDATION_ERROR, DUPLICATE_REGISTRATION, DIRECTORY_NOT_FOUND
Use ModelONEXContainer in node constructors for dependency injection, never use ModelContainer[T]
Use duck typing with protocols for service resolution - do not use isinstance checks

Files:

  • src/omnibase_core/protocols/handler/__init__.py
  • src/omnibase_core/protocols/handler/protocol_handler_context.py
  • src/omnibase_core/protocols/__init__.py
  • tests/unit/runtime/test_message_dispatch_engine.py
  • src/omnibase_core/runtime/message_dispatch_engine.py
tests/**/*.py

📄 CodeRabbit inference engine (CLAUDE.md)

Use pytest markers @pytest.mark.unit, @pytest.mark.integration, @pytest.mark.slow, @pytest.mark.smoke, @pytest.mark.performance to classify tests

Files:

  • tests/unit/runtime/test_message_dispatch_engine.py
🧠 Learnings (2)
📚 Learning: 2025-12-20T04:09:41.822Z
Learnt from: CR
Repo: OmniNode-ai/omnibase_core PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-20T04:09:41.822Z
Learning: Applies to **/*.py : Use EnumNodeKind for architectural role classification (EFFECT, COMPUTE, REDUCER, ORCHESTRATOR, RUNTIME_HOST) and EnumNodeType for implementation type discovery

Applied to files:

  • tests/unit/runtime/test_message_dispatch_engine.py
  • src/omnibase_core/runtime/message_dispatch_engine.py
📚 Learning: 2025-12-20T04:09:41.822Z
Learnt from: CR
Repo: OmniNode-ai/omnibase_core PR: 0
File: CLAUDE.md:0-0
Timestamp: 2025-12-20T04:09:41.822Z
Learning: Applies to **/*.py : Use ModelOnexError with EnumCoreErrorCode for all error handling instead of generic Exception

Applied to files:

  • tests/unit/runtime/test_message_dispatch_engine.py
  • src/omnibase_core/runtime/message_dispatch_engine.py
🧬 Code graph analysis (4)
src/omnibase_core/protocols/handler/__init__.py (2)
tests/unit/runtime/test_message_dispatch_engine.py (16)
  • handler (263-266)
  • handler (324-327)
  • handler (352-355)
  • handler (372-375)
  • handler (420-423)
  • handler (441-444)
  • handler (512-515)
  • handler (565-569)
  • handler (699-702)
  • handler (735-738)
  • handler (830-833)
  • handler (1082-1085)
  • handler (1189-1192)
  • handler (1298-1301)
  • handler (1482-1485)
  • handler (1528-1535)
src/omnibase_core/protocols/handler/protocol_handler_context.py (1)
  • ProtocolHandlerContext (34-85)
src/omnibase_core/protocols/handler/protocol_handler_context.py (1)
src/omnibase_core/models/metadata/model_metadata_field_info.py (1)
  • uuid (157-176)
src/omnibase_core/protocols/__init__.py (2)
tests/unit/runtime/test_message_dispatch_engine.py (16)
  • handler (263-266)
  • handler (324-327)
  • handler (352-355)
  • handler (372-375)
  • handler (420-423)
  • handler (441-444)
  • handler (512-515)
  • handler (565-569)
  • handler (699-702)
  • handler (735-738)
  • handler (830-833)
  • handler (1082-1085)
  • handler (1189-1192)
  • handler (1298-1301)
  • handler (1482-1485)
  • handler (1528-1535)
src/omnibase_core/protocols/handler/protocol_handler_context.py (1)
  • ProtocolHandlerContext (34-85)
src/omnibase_core/runtime/message_dispatch_engine.py (6)
src/omnibase_core/enums/enum_node_kind.py (1)
  • EnumNodeKind (13-118)
src/omnibase_core/models/compute/model_compute_context.py (1)
  • ModelComputeContext (37-104)
src/omnibase_core/models/effect/model_effect_context.py (1)
  • ModelEffectContext (51-123)
src/omnibase_core/models/orchestrator/model_orchestrator_context.py (1)
  • ModelOrchestratorContext (51-120)
src/omnibase_core/models/reducer/model_reducer_context.py (1)
  • ModelReducerContext (31-96)
src/omnibase_core/protocols/handler/protocol_handler_context.py (5)
  • ProtocolHandlerContext (34-85)
  • correlation_id (68-70)
  • envelope_id (73-75)
  • trace_id (78-80)
  • span_id (83-85)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)
  • GitHub Check: claude-review
🔇 Additional comments (13)
src/omnibase_core/protocols/handler/protocol_handler_context.py (1)

1-88: Well-designed protocol with clear documentation.

The ProtocolHandlerContext protocol is well-structured with:

  • Appropriate use of @runtime_checkable for duck typing support
  • PEP 604 union syntax (UUID | None) per coding guidelines
  • Comprehensive docstrings covering thread safety and design rationale
  • Clean, minimal interface with only the necessary causality tracking fields

The decision to keep node-specific fields (like now, retry_attempt) on concrete types rather than the protocol is sound for maintaining separation of concerns.

src/omnibase_core/protocols/__init__.py (1)

123-127: LGTM - Clean public API export.

The ProtocolHandlerContext is correctly imported from the handler subpackage and added to __all__. The placement follows the existing organizational pattern with a dedicated "Handler Module Exports" section.

Also applies to: 294-297

src/omnibase_core/protocols/handler/__init__.py (1)

1-11: LGTM - Standard package initializer.

Clean re-export pattern that makes ProtocolHandlerContext available at the package level.

tests/unit/runtime/test_message_dispatch_engine.py (5)

30-37: LGTM - Imports updated correctly.

New imports for EnumNodeKind and ProtocolHandlerContext are correctly added and follow the existing import organization.


2343-2392: Well-structured context injection tests.

The TestContextInjection class provides comprehensive coverage for the new context injection functionality. The test for reducer context correctly verifies that ModelReducerContext is received and that the now attribute is absent (preserving purity).


2443-2501: Critical test for retry_attempt mapping.

This test correctly validates the mapping from envelope.retry_count to ModelEffectContext.retry_attempt. The explicit assertion on line 2499-2500 is valuable for documenting this mapping behavior.


2641-2721: Comprehensive tracing field propagation test.

Excellent test coverage for verifying all causality tracking fields (correlation_id, trace_id, span_id, envelope_id) propagate correctly from the envelope to the handler context. The detailed assertion messages will be helpful for debugging failures.


2786-2805: RUNTIME_HOST rejection test.

Good boundary test verifying that RUNTIME_HOST is explicitly rejected with the correct error code (INVALID_PARAMETER). This aligns with the design decision that RUNTIME_HOST is for infrastructure, not message handlers.

src/omnibase_core/runtime/message_dispatch_engine.py (5)

90-107: LGTM - Imports organized correctly.

New imports for EnumNodeKind, all context model types, and ProtocolHandlerContext are properly added. The grouping follows the existing import organization pattern.


130-134: Breaking change: Handler signature now includes context.

The HandlerFunc type alias update reflects the new handler signature (envelope, context). This is a documented breaking change per PR objectives.


438-456: Proper validation for node_kind parameter.

The validation correctly:

  1. Checks that node_kind is an EnumNodeKind instance
  2. Explicitly rejects RUNTIME_HOST with a clear error message
  3. Uses ModelOnexError with EnumCoreErrorCode.INVALID_PARAMETER per coding guidelines

1303-1408: Well-implemented context factory method.

The _build_handler_context method correctly:

  1. Extracts causality IDs from the envelope with a fallback for missing correlation_id
  2. Protects against negative retry_count with max(0, envelope.retry_count)
  3. Creates the appropriate context model based on node_kind
  4. Includes a debug-only assertion for protocol compliance

The defensive logging for missing correlation_id (lines 1338-1343) helps identify upstream issues without failing the dispatch.


1410-1442: Context injection works for both sync and async handlers.

The updated _execute_handler correctly:

  1. Builds the context once via _build_handler_context
  2. Passes both envelope and context to async handlers
  3. Passes both arguments to sync handlers via run_in_executor

The sync handler path on line 1442 correctly passes both positional arguments to the executor.

…-984]

Replace silent max(0, retry_count) with explicit check and warning log
to alert operators of potential data corruption when envelope has
negative retry_count values.
@claude

claude Bot commented Dec 20, 2025

Copy link
Copy Markdown

Code Review - PR #225: Context Model Integration

Summary

This PR successfully integrates the context models from PR #224 into the message dispatch pipeline. The implementation is well-designed, thoroughly tested, and follows ONEX architectural principles. I recommend approval with minor suggestions below.

✅ Strengths

1. Excellent Protocol Design

  • ProtocolHandlerContext cleanly defines the common interface for all context types
  • Marked as @runtime_checkable for dynamic type validation
  • Well-documented with clear explanations of each field's purpose
  • Thread safety correctly documented (frozen models)

2. Strong Type Safety

  • Context type selection correctly maps node kinds to context models:
    • EFFECT → ModelEffectContext (with now, retry_attempt)
    • COMPUTE → ModelComputeContext (pure, no time)
    • REDUCER → ModelReducerContext (pure, no time)
    • ORCHESTRATOR → ModelOrchestratorContext (with now)
  • RUNTIME_HOST correctly rejected with clear error message
  • Runtime assertion in _build_handler_context provides development-time safety

3. Comprehensive Test Coverage

  • 73 total dispatch engine tests (up from 39)
  • Excellent context injection test suite covering:
    • Each node kind receives correct context type
    • Correlation ID preservation and propagation
    • Trace ID propagation (with and without values)
    • Retry attempt mapping from envelope.retry_count
    • RUNTIME_HOST rejection
  • Tests verify both positive and negative scenarios

4. Defensive Programming

  • Handles missing correlation_id with fallback UUID generation (lines 1336-1343)
  • Clamps negative retry_count to 0 (lines 1353-1361)
  • Clear warning logs for data anomalies
  • Protocol compliance assertion in debug mode (lines 1410-1414)

5. Breaking Change Handled Well

  • Clear documentation in PR description about signature changes
  • All existing tests updated to new (envelope, context) signature
  • Migration is straightforward for downstream consumers

🔍 Areas for Improvement

1. Missing Import in _build_handler_context

Location: src/omnibase_core/runtime/message_dispatch_engine.py:1343

The uuid4 function is used but not imported at the top of the file. This will cause a runtime error.

Fix:

from uuid import uuid4  # Add this to imports

Severity: ⚠️ HIGH - This is a runtime bug that will crash when correlation_id is None.

2. Inconsistent Error Handling for Invalid node_kind

Location: Lines 438-454 vs. 1400-1406

The registration validation checks isinstance(node_kind, EnumNodeKind), but this check will always pass in Python since all enum values are instances of their enum class. The check should validate against unexpected values.

Current:

if not isinstance(node_kind, EnumNodeKind):  # This won't catch invalid values
    raise ModelOnexError(...)

Suggestion:

# The check is redundant due to type hints, but if you want runtime validation:
if node_kind not in {EnumNodeKind.EFFECT, EnumNodeKind.COMPUTE, 
                      EnumNodeKind.REDUCER, EnumNodeKind.ORCHESTRATOR}:
    raise ModelOnexError(
        message=f"Invalid node_kind: {node_kind}. Valid values: EFFECT, COMPUTE, REDUCER, ORCHESTRATOR",
        error_code=EnumCoreErrorCode.INVALID_PARAMETER,
    )

However, since RUNTIME_HOST is already explicitly rejected on line 448, this check might be redundant. Consider removing the isinstance check entirely and relying on type hints + RUNTIME_HOST check.

Severity: 🟡 MEDIUM - Not a bug, but the validation logic could be clearer.

3. Thread Safety Documentation Needed

Location: MessageDispatchEngine._build_handler_context

While the context models themselves are frozen/immutable, the dispatch engine method accesses mutable state (entry.node_kind, envelope.*). The docstring should clarify thread safety assumptions.

Suggestion: Add to docstring:

Thread Safety:
    Safe to call concurrently. Reads immutable fields from entry and envelope.
    Returns a new frozen context model on each invocation.

Severity: 🟢 LOW - Documentation enhancement.

4. Minor: Debug Assertion Could Be More Specific

Location: Lines 1410-1414

The assertion message could include which protocol methods are missing if the check fails.

Enhancement:

if __debug__:
    # Verify protocol compliance
    required_attrs = ['correlation_id', 'envelope_id', 'trace_id', 'span_id']
    missing = [attr for attr in required_attrs if not hasattr(context, attr)]
    assert not missing, (
        f"Context {type(context).__name__} missing required protocol attributes: {missing}"
    )

Severity: 🟢 LOW - Nice-to-have improvement.

🔒 Security Review

  • ✅ No injection vulnerabilities
  • ✅ UUID generation is cryptographically secure (uuid4)
  • ✅ Frozen models prevent accidental mutation
  • ✅ No sensitive data logged in warning messages (only UUIDs)

📊 Performance Considerations

  • ✅ Context creation is lightweight (Pydantic model instantiation)
  • ✅ No database/network calls in context building
  • ✅ Single context instance per handler invocation (no unnecessary copies)
  • ⚠️ Minor: Generating fallback UUIDs for missing correlation_id adds overhead. Consider whether upstream should be fixed instead to always provide correlation_id.

🧪 Test Coverage Analysis

  • ✅ All four node kinds tested
  • ✅ Edge cases covered (missing correlation_id, negative retry_count, null trace IDs)
  • ✅ Protocol compliance verified
  • ✅ RUNTIME_HOST rejection verified
  • ✅ Sync and async handler execution paths tested

Missing Test: Consider adding a test for what happens when envelope.retry_count is negative (though this is tested implicitly via the clamping logic).

📝 Documentation Quality

  • ✅ Excellent docstrings in ProtocolHandlerContext
  • ✅ Clear examples in protocol docstring
  • ✅ Breaking changes well-documented in PR description
  • ✅ Context model differences clearly explained in table

🎯 Architectural Alignment

This PR strongly aligns with ONEX principles:

  • ✅ Protocol-driven design
  • ✅ Separation of concerns (pure compute vs. effectful operations)
  • ✅ Causality tracking through correlation IDs
  • ✅ Node kind distinctions enforced at runtime

📋 Checklist

  • ✅ Code follows ONEX conventions from CLAUDE.md
  • ✅ All tests pass (73 tests)
  • ✅ Type checking passes (mypy mentioned in test plan)
  • ✅ Breaking changes documented
  • ✅ Thread safety considered
  • ⚠️ Missing: Import for uuid4 (critical fix needed)

🎬 Recommendation

APPROVE after fixing the missing import. The missing uuid4 import is the only blocking issue. Once fixed:

# Verify the fix
poetry run mypy src/omnibase_core/runtime/message_dispatch_engine.py
poetry run pytest tests/unit/runtime/test_message_dispatch_engine.py -xvs

🙌 Great Work!

This is a well-crafted PR with excellent test coverage and clear documentation. The context injection pattern will be valuable for all downstream ONEX applications.


Review by Claude Code | Architectural alignment checked against CLAUDE.md

@claude

claude Bot commented Dec 20, 2025

Copy link
Copy Markdown

🔄 Correction to Previous Review

I need to correct an error in my previous review:

✅ uuid4 Import is Present

I incorrectly stated that uuid4 was missing from the imports. After re-checking line 85 of message_dispatch_engine.py, I can confirm:

from uuid import uuid4  # Line 85 - PRESENT

This is NOT a blocking issue. The import is correct.

Updated Recommendation

APPROVE - All critical issues from my previous review were incorrect. The code is production-ready as-is.

The "Areas for Improvement" I mentioned are all optional enhancements, not blockers:

  • The isinstance(node_kind, EnumNodeKind) check (MEDIUM) is a minor code style suggestion
  • Thread safety documentation (LOW) is nice-to-have
  • Debug assertion enhancement (LOW) is optional

This PR is well-implemented and ready to merge. My apologies for the confusion in the first review.


Correction posted by Claude Code

@jonahgabriel
jonahgabriel merged commit 1f19ddf into main Dec 20, 2025
29 checks passed
@jonahgabriel
jonahgabriel deleted the jonah/omn-984-wire-context-models-into-handler-dispatch-path-follow-up branch December 26, 2025 17:44
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant