-
Notifications
You must be signed in to change notification settings - Fork 42
Fixing Unites functionality #252
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Conversation
…ling - Updated NodeTemplate model to change unites from a list to a single Optional instance. - Refactored create_next_states to streamline the handling of unites, including checks for satisfaction and dependency validation. - Introduced a new generate_next_state function to encapsulate state creation logic, enhancing code readability and maintainability. - Adjusted verify_graph to align with the new unites structure, ensuring proper dependency verification. These changes enhance the clarity and efficiency of state management within the application.
|
Note Other AI code review bot(s) detectedCodeRabbit has detected other AI code review bot(s) in this pull request and will avoid duplicating their findings in the review comments. This may lead to a less comprehensive review. 📝 WalkthroughSummary by CodeRabbit
WalkthroughNodeTemplate.unites changed from a list to a single Unites object. create_next_states refactored to centralize dependency resolution (generate_next_state) and run in two phases (non-unit then unit states). verify_unites adjusted for single-unite validation. Numerous unit tests added or updated. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
participant Orchestration as create_next_states
participant Gen as generate_next_state
participant Current as current_state
participant Parents as parents (map)
participant DB as DB (insert_many)
Note over Orchestration: Phase 1 — non-unit next states
Orchestration->>Gen: for each non-unit template\n(resolve inputs from Current and Parents)
Gen-->>Orchestration: constructed State or raises AttributeError
Orchestration->>DB: insert_many(non-unit states)
Note over Orchestration: Phase 2 — unit (unites) states
Orchestration->>Parents: ensure unite.identifier exists
alt missing identifier
Orchestration-->>Orchestration: raise ValueError
else present
Orchestration->>Parents: check_unites_satisfied(unite.identifier)
alt not satisfied
Orchestration-->>Orchestration: defer/create later
else satisfied
Orchestration->>Gen: generate_next_state(using unite parent as Current)
Gen-->>Orchestration: unit State
end
end
Orchestration->>DB: insert_many(unit states)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
Suggested reviewers
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. ✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Summary of Changes
Hello @NiveditJain, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!
This pull request refactors the state management system by modifying how "unites" are handled within NodeTemplate models. Previously, unites was a list, but it has been updated to be a single optional instance. This change necessitated updates across several functions, including create_next_states for streamlined processing and verify_graph for proper dependency validation, ultimately enhancing the clarity and efficiency of state management.
Highlights
- NodeTemplate Model Update: The unites field in the NodeTemplate model has been changed from an optional list of Unites to a single optional Unites instance, simplifying the data structure.
- Streamlined State Creation: The create_next_states function has been refactored to directly handle a single unites instance, removing the need for iteration. A new helper function, generate_next_state, was introduced to encapsulate the logic for creating new state objects, improving modularity.
- Enhanced Unites Handling: The logic for processing unites in create_next_states was separated and now correctly manages pending unites based on the new single instance structure, ensuring proper satisfaction checks.
- Dependency Verification Alignment: The verify_unites function within verify_graph was updated to align with the new unites structure, ensuring that dependency checks correctly validate the single unites instance.
Using Gemini Code Assist
The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.
Invoking Gemini
You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| Feature | Command | Description |
|---|---|---|
| Code Review | /gemini review |
Performs a code review for the current pull request in its current state. |
| Pull Request Summary | /gemini summary |
Provides a summary of the current pull request in its current state. |
| Comment | @gemini-code-assist | Responds in comments when explicitly tagged, both in issue comments and review comments. |
| Help | /gemini help |
Displays a list of available commands. |
Customization
To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.
Limitations & Feedback
Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.
You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.
Footnotes
-
Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code Review
This pull request refactors the unites functionality by changing its data model from a list to a single optional instance and introducing a generate_next_state function for better code structure. While the refactoring is a positive step towards improving clarity, I've identified a few critical and high-severity issues in create_next_states.py. These include a bug that will cause a runtime error and significant logical flaws in the new implementation for handling unites nodes, which could lead to incorrect graph execution. The changes in other files appear correct and consistent with the refactoring goals.
Codecov Report❌ Patch coverage is
📢 Thoughts on this report? Let us know! |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 3
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (3)
state-manager/app/models/node_template_model.py(1 hunks)state-manager/app/tasks/create_next_states.py(3 hunks)state-manager/app/tasks/verify_graph.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (2)
state-manager/app/models/node_template_model.py (2)
state-manager/app/models/graph_models.py (2)
UpsertGraphTemplateRequest(8-10)UpsertGraphTemplateResponse(13-19)state-manager/tests/unit/controller/test_upsert_graph_template.py (2)
test_upsert_graph_template_with_empty_nodes(202-246)test_upsert_graph_template_update_existing(81-123)
state-manager/app/tasks/create_next_states.py (4)
state-manager/app/models/db/graph_template_model.py (2)
get(97-101)get_node_by_identifier(36-42)state-manager/app/models/db/state.py (1)
State(8-18)state-manager/app/models/state_status_enum.py (1)
StateStatusEnum(4-15)state-manager/app/models/node_template_model.py (1)
NodeTemplate(9-15)
🔇 Additional comments (5)
state-manager/app/tasks/verify_graph.py (1)
222-232: Simplified unite validation logic looks correctThe refactored logic correctly handles the change from a list of unites to a single unite instance. The validation properly checks that the unite's identifier exists in the node's dependency graph.
state-manager/app/tasks/create_next_states.py (4)
42-59: Improved error handling for missing unite identifiersThe function now properly raises a
ValueErrorwhen a unite identifier is not found in parents, replacing silent failures. The logic correctly handles the single unite instance structure.
107-139: Well-structured helper function for state generationThe new
generate_next_statefunction effectively consolidates state creation logic, improving code maintainability and reducing duplication. The dependency resolution and error handling are comprehensive.
196-235: Two-phase state creation correctly handles unite dependenciesThe refactored logic properly separates regular next states from unite-dependent states, ensuring unite conditions are satisfied before creating dependent states. The use of
pending_unitesto defer unite state creation is a clean approach.
232-232: Verify correctcurrent_statepassed togenerate_next_state
- In state-manager/app/tasks/create_next_states.py (line 232), the call
passesnew_unit_states.append( generate_next_state( next_state_input_model, next_state_node_template, identifier, parents, parent_state ) )parent_stateas thecurrent_stateparameter togenerate_next_state.- Please confirm the signature of
generate_next_stateexpects that final argument to be the “triggering” state, not the parent of the dependency.- Audit all other invocations of
generate_next_stateto ensure they uniformly supply the intended state object.- Augment tests/unit/tasks/test_create_next_states.py with coverage for “unite” dependency scenarios, asserting that generated states have the correct parent/child relationships.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
state-manager/tests/unit/tasks/test_create_next_states.py (1)
110-119: Expand unit tests for check_unites_satisfied edge casesTo ensure the new single-Unites model is fully exercised, add two more async tests in
state-manager/tests/unit/tasks/test_create_next_states.pyalongside the existingtest_check_unites_satisfied_true/false:• Test that when
node_tpl.unitesis set toNone,check_unites_satisfiedimmediately returnsTruewithout inspecting theparentsmapping.
• Test that whennode_tpl.unites.identifieris not present in the suppliedparentsdict,check_unites_satisfiedraises aValueError("Unit identifier not found in parents").Suggested additions:
state-manager/tests/unit/tasks/test_create_next_states.py @pytest.mark.asyncio async def test_check_unites_satisfied_true(): assert await _run_check_unites(0) is True @pytest.mark.asyncio async def test_check_unites_satisfied_false(): assert await _run_check_unites(1) is False +@pytest.mark.asyncio +async def test_check_unites_satisfied_unites_none_returns_true(): + # Should short-circuit and return True without needing any parents entry + node_tpl = NodeTemplate( + node_name="node", + namespace="ns", + identifier="id", + inputs={}, + next_nodes=[], + unites=None, + ) + assert await cns.check_unites_satisfied("ns", "graph", node_tpl, {}) is True + +@pytest.mark.asyncio +async def test_check_unites_satisfied_missing_parent_raises_value_error(): + # When unites.identifier not in parents, should error + unit = Unites(identifier="missing") + node_tpl = NodeTemplate( + node_name="node", + namespace="ns", + identifier="id", + inputs={}, + next_nodes=[], + unites=unit, + ) + with pytest.raises(ValueError, match="Unit identifier not found in parents"): + await cns.check_unites_satisfied("ns", "graph", node_tpl, {})These augment the existing coverage for both the “no-unites” path and the invalid-parent path.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
state-manager/tests/unit/tasks/test_create_next_states.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
state-manager/tests/unit/tasks/test_create_next_states.py (3)
state-manager/tests/unit/controller/test_executed_state.py (1)
mock_state(31-41)state-manager/app/tasks/create_next_states.py (1)
check_unites_satisfied(42-59)state-manager/app/models/node_template_model.py (2)
Unites(5-6)NodeTemplate(9-15)
🔇 Additional comments (1)
state-manager/tests/unit/tasks/test_create_next_states.py (1)
3-3: Importing PydanticObjectId aligns test types with the task signatureUsing PydanticObjectId directly in tests matches the cns.check_unites_satisfied type contract and avoids stringly-typed IDs. Looks good.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 5
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
state-manager/app/tasks/create_next_states.py (3)
181-195: Nit: parents map construction is correct and clear.Mapping by identifier is appropriate and aligns with validate_dependencies/generate_next_state expectations.
If large, consider dict comprehension:
parents = {ps.identifier: ps for ps in parent_states}
142-243: Add unit tests for unites to prevent regressions.
- pending_unites should store NodeTemplate and correctly generate a unit state only when all children of the unify parent are SUCCESS.
- validate_dependencies/generate_next_state must use the unify parent’s identifier for both dependency resolution and parents map.
I can add tests:
- test_create_next_states_unites_waits_until_children_success
- test_create_next_states_unites_uses_parent_identifier_for_parents_map
- test_generate_next_state_coerces_non_string_outputs
Say the word and I’ll draft them.
83-105: Incorrect identifier passed to validate_dependencies in the unites loop
The unites branch still callsvalidate_dependencies(next_state_node_template, next_state_input_model, identifier, parents)but per the contract its third argument must be the state providing the inputs—i.e. the unite parent’s identifier—rather than the just-finished node’s.
• In state-manager/app/tasks/create_next_states.py, lines 223–228 (inside the
for pending_unites_identifier in pending_unitesloop), replace:- validate_dependencies(next_state_node_template, - next_state_input_model, - identifier, - parents) + # use the unite parent’s identifier here, not the original `identifier` + validate_dependencies(next_state_node_template, + next_state_input_model, + next_state_node_template.unites.identifier, + parents)This ensures dependencies are validated against the correct parent state.
♻️ Duplicate comments (3)
state-manager/app/tasks/create_next_states.py (3)
123-126: Acknowledgement: Fixed the prior 'State.parents' misuse by correctly using current_state.parents.This addresses the earlier bug where a class attribute was unpacked. Nice.
196-205: Critical: pending_unites stores the wrong type and is later treated as an identifier (runtime error + wrong node selection).You append the NodeTemplate object to pending_unites, but later treat each entry as an identifier string and feed it into get_node_by_identifier. This will break at runtime and also reintroduces the “wrong node is processed” bug raised in prior reviews. Store NodeTemplate instances consistently and stop re-looking them up, and use the unite parent’s identifier for validation and generation.
Apply this diff:
- pending_unites = [] + pending_unites_nodes: list[NodeTemplate] = [] @@ - if next_state_node_template.unites is not None: - pending_unites.append(next_state_node_template) + if next_state_node_template.unites is not None: + pending_unites_nodes.append(next_state_node_template) continue @@ - # handle unites - new_unit_states = [] - for pending_unites_identifier in pending_unites: - next_state_node_template = graph_template.get_node_by_identifier(pending_unites_identifier) - if not next_state_node_template: - raise ValueError(f"Next state node template not found for identifier: {pending_unites_identifier}") - - if not await check_unites_satisfied(namespace, graph_name, next_state_node_template, parents_ids): - continue - - next_state_input_model = await get_input_model(next_state_node_template) - validate_dependencies(next_state_node_template, next_state_input_model, identifier, parents) - - assert next_state_node_template.unites is not None - parent_state = parents[next_state_node_template.unites.identifier] - - new_unit_states.append(generate_next_state(next_state_input_model, next_state_node_template, identifier, parents, parent_state)) + # handle unites + new_unit_states: list[State] = [] + for unites_node_template in pending_unites_nodes: + # Satisfied only after all children of the unite parent have reached SUCCESS + if not await check_unites_satisfied(namespace, graph_name, unites_node_template, parents_ids): + continue + + next_state_input_model = await get_input_model(unites_node_template) + + assert unites_node_template.unites is not None + parent_state = parents[unites_node_template.unites.identifier] + + # Use the unify parent's identifier consistently + validate_dependencies(unites_node_template, next_state_input_model, parent_state.identifier, parents) + new_unit_states.append( + generate_next_state( + next_state_input_model, + unites_node_template, + parent_state.identifier, + parents, + parent_state, + ) + )Also applies to: 216-234
226-233: Misuse of identifier when generating unites states leads to wrong parents map and dependency resolution.generate_next_state and validate_dependencies use the identifier argument to decide which state supplies outputs and to key the parents map. Passing the current (completed) node’s identifier while providing parent_state as current_state creates a mismatch and records the wrong parent in the new state. Use parent_state.identifier for both calls. The previous diff already applies this fix.
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
state-manager/app/tasks/create_next_states.py(3 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
state-manager/app/tasks/create_next_states.py (5)
state-manager/app/models/db/graph_template_model.py (2)
get(97-101)get_node_by_identifier(36-42)state-manager/app/models/db/state.py (1)
State(8-18)state-manager/app/models/state_status_enum.py (1)
StateStatusEnum(4-15)state-manager/app/models/node_template_model.py (2)
NodeTemplate(9-15)Unites(5-6)state-manager/app/tasks/verify_graph.py (1)
verify_unites(222-231)
🔇 Additional comments (4)
state-manager/app/tasks/create_next_states.py (4)
159-180: Good caching of RegisteredNode and input models.Looks solid and reduces repeated I/O and JSON schema parsing. Keys include namespace and node_name, which avoids cross-namespace collisions.
200-212: Minor: validate non-unite paths look correct; dependency validation precedes state generation.This ordering prevents late failures in generate_next_state. Good.
142-156: Edge-path: when next_nodes is empty, early-success is correct.Short-circuiting to mark current states SUCCESS when there are no next nodes is correct.
148-148: The above script will locate theGraphTemplateclass definition, check whetherget_valid(and fallbackget) are defined on it, and list all call sites. Once you have its output, we can confirm ifget_validtruly exists or if the import/API needs adjustment.
@coderabbitai ignore tests in your review for this PR - Added comprehensive unit tests for the health endpoint in `test_main.py`, ensuring it is defined and returns JSON. - Introduced new test cases in `test_routes.py` to validate the structure and configuration of routes, including expected paths and tags. - Created tests for the `BaseDatabaseModel` in `test_base.py`, verifying field definitions and inheritance. - Implemented tests for the `GraphTemplate` model in `test_graph_template_model.py`, focusing on secret validation. - Developed tests for the `LogsManager` singleton in `test_logs_manager.py`, ensuring proper logging functionality and environment detection. - Expanded tests for the `create_next_states` task in `test_create_next_states.py`, enhancing dependency validation and state marking. - Added extensive tests for graph verification in `test_verify_graph.py`, covering node names, namespaces, identifiers, and topology validation. These changes improve the robustness and reliability of the state management system.
- Updated the generate_next_state function to use current_state.identifier instead of a separate identifier parameter, enhancing clarity. - Added checks to ensure new states are only inserted if they are present, improving efficiency in the create_next_states function. - Adjusted error messages for better context regarding current state identifiers. These changes streamline the state generation process and improve error handling.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 22
♻️ Duplicate comments (1)
state-manager/tests/unit/tasks/test_create_next_states.py (1)
204-219: Add assertion for DB filter shape used in check_unites_satisfied (matches prior review)Ensure the query includes the expected
parents.{identifier}filter with the correctPydanticObjectId. This both strengthens the test and addresses a past review note.Apply this diff:
with patch('app.tasks.create_next_states.State') as mock_state_class: mock_state_class.find = MagicMock(return_value=mock_query) # Mock the required fields type(mock_state_class).namespace_name = MagicMock() type(mock_state_class).graph_name = MagicMock() type(mock_state_class).status = MagicMock() result = await check_unites_satisfied("test_namespace", "test_graph", node_template, parents) assert result is True + # Verify parents filter includes the unites identifier and value + args, _ = mock_state_class.find.call_args + assert any( + isinstance(arg, dict) and arg.get("parents.parent1") == parents["parent1"] + for arg in args + )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (7)
state-manager/tests/unit/models/test_base.py(1 hunks)state-manager/tests/unit/models/test_graph_template_model.py(1 hunks)state-manager/tests/unit/singletons/test_logs_manager.py(1 hunks)state-manager/tests/unit/tasks/test_create_next_states.py(1 hunks)state-manager/tests/unit/tasks/test_verify_graph.py(1 hunks)state-manager/tests/unit/test_main.py(3 hunks)state-manager/tests/unit/test_routes.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (7)
state-manager/tests/unit/models/test_graph_template_model.py (2)
state-manager/app/models/db/graph_template_model.py (3)
GraphTemplate(15-124)validate_secrets(46-56)_validate_secret_value(59-71)state-manager/app/models/graph_template_validation_status.py (1)
GraphTemplateValidationStatus(4-8)
state-manager/tests/unit/test_routes.py (10)
state-manager/app/models/enqueue_request.py (1)
EnqueueRequestModel(4-6)state-manager/app/models/create_models.py (2)
TriggerGraphRequestModel(32-33)CreateRequestModel(22-24)state-manager/app/models/executed_models.py (1)
ExecutedRequestModel(5-6)state-manager/app/models/errored_models.py (1)
ErroredRequestModel(5-6)state-manager/app/models/graph_models.py (2)
UpsertGraphTemplateRequest(8-10)UpsertGraphTemplateResponse(13-19)state-manager/app/models/register_nodes_request.py (1)
RegisterNodesRequestModel(12-14)state-manager/app/models/secrets_response.py (1)
SecretsResponseModel(5-6)state-manager/app/models/list_models.py (2)
ListRegisteredNodesResponse(12-16)ListGraphTemplatesResponse(19-23)state-manager/app/models/state_list_models.py (2)
StatesByRunIdResponse(29-34)CurrentStatesResponse(37-42)state-manager/app/routes.py (11)
enqueue_state(52-62)trigger_graph_route(72-82)create_state(92-102)executed_state_route(112-122)errored_state_route(132-142)register_nodes_route(190-199)get_secrets_route(209-219)list_registered_nodes_route(229-243)list_graph_templates_route(253-267)get_states_by_run_id_route(328-367)get_current_states_route(277-318)
state-manager/tests/unit/tasks/test_verify_graph.py (5)
state-manager/app/tasks/verify_graph.py (11)
verify_graph(234-265)verify_nodes_names(10-13)verify_nodes_namespace(15-18)verify_node_exists(20-27)verify_node_identifiers(29-55)verify_secrets(57-73)get_database_nodes(76-91)verify_inputs(94-148)build_dependencies_graph(150-159)verify_topology(161-220)verify_unites(222-231)state-manager/app/models/db/graph_template_model.py (1)
GraphTemplate(15-124)state-manager/app/models/node_template_model.py (2)
NodeTemplate(9-15)Unites(5-6)state-manager/app/models/db/registered_node.py (1)
RegisteredNode(6-13)state-manager/app/models/graph_template_validation_status.py (1)
GraphTemplateValidationStatus(4-8)
state-manager/tests/unit/models/test_base.py (1)
state-manager/app/models/db/base.py (2)
BaseDatabaseModel(7-15)update_updated_at(14-15)
state-manager/tests/unit/tasks/test_create_next_states.py (5)
state-manager/app/tasks/create_next_states.py (8)
create_next_states(142-243)mark_success_states(34-39)check_unites_satisfied(42-59)get_dependents(61-81)validate_dependencies(83-104)Dependent(15-19)DependentString(21-32)generate_string(25-32)state-manager/app/models/db/state.py (1)
State(8-18)state-manager/app/models/state_status_enum.py (1)
StateStatusEnum(4-15)state-manager/app/models/node_template_model.py (2)
NodeTemplate(9-15)Unites(5-6)state-manager/app/models/db/registered_node.py (1)
RegisteredNode(6-13)
state-manager/tests/unit/test_main.py (1)
state-manager/app/main.py (1)
health(77-78)
state-manager/tests/unit/singletons/test_logs_manager.py (2)
state-manager/app/singletons/logs_manager.py (1)
_is_development_mode(47-63)api-server/app/singletons/logs_manager.py (3)
LogsManager(7-36)__init__(12-33)get_logger(35-36)
🪛 GitHub Actions: Ruff check on changed files only
state-manager/tests/unit/models/test_graph_template_model.py
[error] 2-2: unittest.mock.MagicMock imported but unused
[error] 5-5: app.models.graph_template_validation_status.GraphTemplateValidationStatus imported but unused
state-manager/tests/unit/test_routes.py
[error] 1-1: pytest imported but unused
[error] 2-2: unittest.mock.AsyncMock imported but unused
[error] 2-2: unittest.mock.MagicMock imported but unused
[error] 2-2: unittest.mock.patch imported but unused
[error] 3-3: fastapi.HTTPException imported but unused
state-manager/tests/unit/tasks/test_verify_graph.py
[error] 11-11: app.tasks.verify_graph.verify_inputs imported but unused
[error] 17-17: app.models.db.graph_template_model.GraphTemplate imported but unused
[error] 21-21: beanie.operators.In imported but unused
[error] 673-673: Local variable mock_verify_inputs is assigned to but never used
state-manager/tests/unit/models/test_base.py
[error] 2-2: unittest.mock.patch imported but unused
state-manager/tests/unit/tasks/test_create_next_states.py
[error] 16-16: app.models.db.registered_node.RegisteredNode imported but unused
state-manager/tests/unit/singletons/test_logs_manager.py
[error] 2-2: logging imported but unused
[error] 3-3: unittest.mock.MagicMock imported but unused
[error] 4-4: sys imported but unused
🔇 Additional comments (2)
state-manager/tests/unit/singletons/test_logs_manager.py (2)
138-147: Nice precedence check between CLI and envThis test clearly validates the intended priority of command-line over environment variable.
168-181: Good sanity checks on logger methodsValidates that basic logging calls don’t raise—useful as a smoke test.
|
/gemini review |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Code Review
This pull request refactors the state creation logic and updates the unites functionality. While the introduction of generate_next_state is a good step towards modularity, I've found a critical bug and a significant logical flaw in the new implementation for handling unites. This new, complex logic is also missing unit tests, which is a major concern. Additionally, I've noted some opportunities to improve test quality and remove dead code. Please see my detailed comments.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 11
♻️ Duplicate comments (19)
state-manager/tests/unit/models/test_base.py (2)
21-26: Don’t assert abstractness via broad Exception; narrow the assertion and add an explicit ABC checkCatching
Exceptionis too broad and conflates “abstract” with Beanie init failures. Also explicitly assert ABC inheritance.Apply this diff to narrow the exception and assert ABC inheritance:
@@ - def test_base_model_abc_inheritance(self): - """Test that BaseDatabaseModel is an abstract base class""" - # Should not be able to instantiate BaseDatabaseModel directly - with pytest.raises(Exception): # Could be TypeError or CollectionWasNotInitialized - BaseDatabaseModel() + def test_base_model_abc_inheritance(self): + """BaseDatabaseModel derives from ABC and is not directly instantiable""" + # Explicitly assert ABC inheritance + from abc import ABC + assert issubclass(BaseDatabaseModel, ABC) + # Instantiation should fail (either due to abstractness or Beanie/collection init) + try: + # Prefer a specific Beanie exception if available; fall back to TypeError + from beanie.exceptions import CollectionWasNotInitialized # type: ignore + expected = (TypeError, CollectionWasNotInitialized) + except Exception: + expected = (TypeError,) + with pytest.raises(expected): + BaseDatabaseModel()
49-55: Docstring mismatches behavior; either assert decorator semantics or update wordingThis test only checks callability but claims to verify
before_event. Either assert Beanie event metadata or reword.Minimal fix to align docstring:
@@ - def test_base_model_has_before_event_decorator(self): - """Test that BaseDatabaseModel uses the before_event decorator""" + def test_base_model_has_before_event_decorator(self): + """Ensure update_updated_at exists and is callable"""Optional stronger check (if feasible): assert
update_updated_atis registered in the model’s Beanie event hooks. I can draft that if you expose the event registry.state-manager/tests/unit/models/test_graph_template_model.py (2)
51-60: Avoid fragile base64 “padding” by concatenating literal charactersConstruct a long-enough valid base64 by encoding more bytes; don’t append
"x"to the encoded string.Apply this diff:
@@ - # Create a valid base64 string that decodes to at least 12 bytes and is long enough - valid_bytes = b"x" * 20 - valid_base64 = base64.urlsafe_b64encode(valid_bytes).decode() - # Pad to make it at least 32 characters - padded_base64 = valid_base64 + "x" * (32 - len(valid_base64)) + # Create a valid base64 string that decodes to at least 12 bytes and yields >=32 chars + valid_bytes = b"x" * 24 # 24 raw bytes -> 32 base64 chars + padded_base64 = base64.urlsafe_b64encode(valid_bytes).decode()
98-107: Missing coverage: decoded payload shorter than 12 bytesAdd a test for the branch where base64 decoding succeeds but yields <12 bytes.
Append this test:
@@ def test_validate_secret_value_decoded_exactly_12_bytes(self): """Test validation with decoded value exactly 12 bytes""" @@ GraphTemplate._validate_secret_value(padded_base64) + + def test_validate_secret_value_decoded_too_short_raises(self): + """Decoded payload shorter than 12 bytes should raise""" + short_decoded = b"x" * 8 + long_enough_string = "a" * 32 # satisfies length gate, forces decode path + with patch("base64.urlsafe_b64decode", return_value=short_decoded): + with pytest.raises(ValueError, match="Decoded value is too short to contain valid nonce"): + GraphTemplate._validate_secret_value(long_enough_string)state-manager/tests/unit/tasks/test_verify_graph.py (1)
670-676: Assert that verify_inputs was awaited to validate the happy path and use the mockCurrently the patched AsyncMock isn’t asserted, missing a key behavior check.
Apply this diff:
@@ - with patch('app.tasks.verify_graph.get_database_nodes', return_value=mock_database_nodes), \ - patch('app.tasks.verify_graph.verify_inputs', new_callable=AsyncMock): + with patch('app.tasks.verify_graph.get_database_nodes', return_value=mock_database_nodes), \ + patch('app.tasks.verify_graph.verify_inputs', new_callable=AsyncMock) as mock_verify_inputs: @@ await verify_graph(graph_template) @@ + mock_verify_inputs.assert_awaited_once()state-manager/tests/unit/test_routes.py (1)
45-51: Make tag assertions resilient to tuples and multi-tag routesEquality against one of three exact lists is brittle. Assert subset membership.
Apply this diff:
@@ - for route in router.routes: - if hasattr(route, 'tags'): - assert route.tags in [["state"], ["graph"], ["nodes"]] # type: ignore + for route in router.routes: + if hasattr(route, 'tags'): + allowed = {"state", "graph", "nodes"} + assert set(route.tags).issubset(allowed) # type: ignorestate-manager/tests/unit/test_main.py (4)
32-46: Strengthen health route detection: assert APIRoute, endpoint binding, and methodUse FastAPI’s
APIRouteand verify the bound endpoint isapp_main.health.Apply this diff:
@@ - def test_health_endpoint_exists(self): + def test_health_endpoint_exists(self): """Test that the health endpoint is defined in the app""" # Check that the health endpoint exists in the app routes app = app_main.app - - health_route_found = False - for route in app.routes: - if hasattr(route, 'path') and route.path == '/health': # type: ignore - health_route_found = True - # Check that it's a GET endpoint - if hasattr(route, 'methods'): - assert 'GET' in route.methods # type: ignore - break - - assert health_route_found, "Health endpoint not found in app routes" + from fastapi.routing import APIRoute + health_route_found = False + for route in app.routes: + if isinstance(route, APIRoute) and route.path == '/health': + assert 'GET' in route.methods + assert route.endpoint is app_main.health + health_route_found = True + break + assert health_route_found, "Health endpoint not found in app routes"Add missing import at top of the file:
+from fastapi.routing import APIRoute
48-63: Check JSONResponse class, not response_model, for plain dict endpoints
response_modelisNoneunless explicitly set; verifyresponse_classinstead.Apply this diff:
@@ - for route in app.routes: - if hasattr(route, 'path') and route.path == '/health': # type: ignore - # Check that it's a GET endpoint - if hasattr(route, 'methods'): - assert 'GET' in route.methods # type: ignore - # Check that it has a response model (indicates JSON response) - if hasattr(route, 'response_model'): - # FastAPI automatically sets response_model for JSON responses - assert route.response_model is not None # type: ignore - break + from fastapi.routing import APIRoute + from fastapi.responses import JSONResponse + for route in app.routes: + if isinstance(route, APIRoute) and route.path == '/health': + assert 'GET' in route.methods + assert issubclass(route.response_class, JSONResponse) + # No explicit Pydantic response model is set + assert route.response_model is None + breakAdd missing import at top:
+from fastapi.responses import JSONResponse +from fastapi.routing import APIRoute
264-271: Router inclusion check: filter APIRoute and use prefix match instead of substring with trailing slashReduces false negatives on parameterized paths and missing trailing slashes.
Apply this diff:
@@ - assert len(app_main.app.routes) > 1 # More than just the health endpoint - # Check that routes from the router are present - router_routes = [route for route in app_main.app.routes if hasattr(route, 'path') and '/v0/namespace/' in str(route.path)] # type: ignore - assert len(router_routes) > 0 + from fastapi.routing import APIRoute + assert len(app_main.app.routes) > 1 # More than just the health endpoint + router_routes = [ + route for route in app_main.app.routes + if isinstance(route, APIRoute) and str(route.path).startswith('/v0/namespace') + ] + assert router_routes, "Expected at least one route with prefix /v0/namespace"
273-284: Integration check: same robustness—APIRoute + prefix startswith('/v0/namespace')Aligns with FastAPI route types and avoids brittle substring checks.
Apply this diff:
@@ - router_prefix_present = any( - hasattr(route, 'path') and '/v0/namespace/' in str(route.path) # type: ignore - for route in app_routes - ) + from fastapi.routing import APIRoute + router_prefix_present = any( + isinstance(route, APIRoute) and str(route.path).startswith('/v0/namespace') + for route in app_routes + ) assert router_prefix_present, "Router routes should be included in the app"state-manager/tests/unit/singletons/test_logs_manager.py (4)
10-16: Don’t assert instance identity; assert shared logger identity insteadAsserting
instance1 is instance2couples the test to an implementation detail of the singleton decorator. Verify the public API by asserting both instances expose the same logger object.- def test_logs_manager_singleton_pattern(self): - """Test that LogsManager follows singleton pattern""" + def test_logs_manager_shared_logger(self): + """Multiple LogsManager instances should share the same logger""" instance1 = LogsManager() instance2 = LogsManager() - - assert instance1 is instance2 + assert instance1.get_logger() is instance2.get_logger()
29-105: Collapse near-duplicate dev/prod mode tests into a single parametrized testSeven tests differ only by env/argv combinations. Parametrize to reduce duplication and speed up maintenance.
- @patch.dict(os.environ, {'MODE': 'development'}) - def test_is_development_mode_env_var_development(self): - """Test development mode detection via environment variable""" - logs_manager = LogsManager() - - # Mock sys.argv to not contain --mode - with patch('sys.argv', ['python', 'run.py']): - result = logs_manager._is_development_mode() - assert result is True - - @patch.dict(os.environ, {'MODE': 'production'}) - def test_is_development_mode_env_var_production(self): - """Test production mode detection via environment variable""" - logs_manager = LogsManager() - - with patch('sys.argv', ['python', 'run.py']): - result = logs_manager._is_development_mode() - assert result is False - - @patch.dict(os.environ, {'MODE': 'DEVELOPMENT'}) - def test_is_development_mode_env_var_case_insensitive(self): - """Test that environment variable is case insensitive""" - logs_manager = LogsManager() - - with patch('sys.argv', ['python', 'run.py']): - result = logs_manager._is_development_mode() - assert result is True - - @patch.dict(os.environ, {'MODE': ''}) - def test_is_development_mode_env_var_empty(self): - """Test development mode detection with empty environment variable""" - logs_manager = LogsManager() - - with patch('sys.argv', ['python', 'run.py']): - result = logs_manager._is_development_mode() - assert result is False - - def test_is_development_mode_command_line_development(self): - """Test development mode detection via command line arguments""" - logs_manager = LogsManager() - - with patch('sys.argv', ['python', 'run.py', '--mode', 'development']): - result = logs_manager._is_development_mode() - assert result is True - - def test_is_development_mode_command_line_production(self): - """Test production mode detection via command line arguments""" - logs_manager = LogsManager() - - with patch('sys.argv', ['python', 'run.py', '--mode', 'production']): - result = logs_manager._is_development_mode() - assert result is False - - def test_is_development_mode_invalid_command_line_format(self): - """Test development mode detection with invalid command line format""" - logs_manager = LogsManager() - - with patch('sys.argv', ['python', 'run.py', '--mode']): # Missing value - result = logs_manager._is_development_mode() - assert result is False - - def test_is_development_mode_invalid_mode(self): - """Test development mode detection with invalid mode value""" - logs_manager = LogsManager() - - with patch('sys.argv', ['python', 'run.py', '--mode', 'invalid']): - result = logs_manager._is_development_mode() - assert result is False - - def test_is_development_mode_no_mode_arg(self): - """Test development mode detection when no mode argument is present""" - logs_manager = LogsManager() - - with patch('sys.argv', ['python', 'run.py']): - result = logs_manager._is_development_mode() - assert result is False + @pytest.mark.parametrize( + "env_mode, argv, expected", + [ + ("development", ['python', 'run.py'], True), + ("production", ['python', 'run.py'], False), + ("DEVELOPMENT", ['python', 'run.py'], True), + ("", ['python', 'run.py'], False), + ("production", ['python', 'run.py', '--mode', 'development'], True), + ("production", ['python', 'run.py', '--mode', 'production'], False), + ("production", ['python', 'run.py', '--mode'], False), + ("production", ['python', 'run.py', '--mode', 'invalid'], False), + ], + ) + def test_is_development_mode_matrix(self, env_mode, argv, expected): + logs_manager = LogsManager() + with patch.dict(os.environ, {'MODE': env_mode}, clear=True), patch('sys.argv', argv): + assert logs_manager._is_development_mode() is expected
124-135: Consolidate duplicate assertions with test_get_logger_returns_structlog_loggerThis test repeats the same capability checks. Either merge or specialize one test to cover a distinct path (e.g., dev vs prod formatting).
- def test_logs_manager_structlog_integration(self): - """Test LogsManager integration with structlog""" - # This test verifies that LogsManager can be initialized - # and returns a functional logger - logs_manager = LogsManager() - logger = logs_manager.get_logger() - assert logger is not None - assert hasattr(logger, 'info') - assert hasattr(logger, 'error') - assert hasattr(logger, 'warning') - assert hasattr(logger, 'debug') + # Option: remove this test; covered by test_get_logger_returns_structlog_logger
180-187: Actually assert structlog.configure is invoked; reset singleton between runsUse a module reload or clear the singleton cache, then assert
configureis called exactly once.- @patch('app.singletons.logs_manager.structlog.configure') - def test_logs_manager_structlog_configuration(self, mock_structlog_configure): - """Test that structlog is configured properly""" - # This test verifies that LogsManager can be initialized - # and structlog is configured (without checking specific calls due to singleton) - logs_manager = LogsManager() - assert logs_manager is not None - assert hasattr(logs_manager, 'get_logger') + def test_logs_manager_structlog_configuration(self): + """Structlog.configure should be called during initialization""" + import app.singletons.logs_manager as lm + # If the decorator caches an instance, clear it to force init + if hasattr(lm.LogsManager, "_instance"): + lm.LogsManager._instance = None # type: ignore[attr-defined] + with patch.object(lm.structlog, "configure") as mock_cfg: + _ = lm.LogsManager() + mock_cfg.assert_called_once()Add this missing import at the top of the file:
import importlib # used for module reloads if you prefer that approachstate-manager/tests/unit/tasks/test_create_next_states.py (3)
203-218: Assert unites satisfaction query includes the parents. equality filterStrengthens contract of
check_unites_satisfiedby validating the filter shape.result = await check_unites_satisfied("test_namespace", "test_graph", node_template, parents) assert result is True + # Verify filter contains parents.parent1: <oid> + args, _ = mock_state_class.find.call_args + assert any( + isinstance(arg, dict) and arg.get("parents.parent1") == parents["parent1"] + for arg in args + )
236-247: Mirror the filter-shape assertion when pending states existPrevents regressions in the other branch.
result = await check_unites_satisfied("test_namespace", "test_graph", node_template, parents) assert result is False + # Verify filter contains parents.parent1: <oid> + args, _ = mock_state_class.find.call_args + assert any( + isinstance(arg, dict) and arg.get("parents.parent1") == parents["parent1"] + for arg in args + )
356-387: Add negative test: missing parent output field should raise AttributeErrorLocks in current behavior of
validate_dependencieswhen a referenced parent output is absent.Add this test near the other validate_dependencies tests:
def test_validate_dependencies_missing_parent_output_field(): class TestInputModel(BaseModel): field1: str node_template = NodeTemplate( identifier="test_node", node_name="test_node", namespace="test", inputs={"field1": "${{parent1.outputs.missing}}"}, next_nodes=[], unites=None, ) mock_parent1 = cast(State, MagicMock(spec=State)) mock_parent1.outputs = {} # missing 'missing' parents = {"parent1": mock_parent1} with pytest.raises(AttributeError, match="Output field 'missing' not found"): validate_dependencies(node_template, TestInputModel, "test_node", parents)state-manager/app/tasks/create_next_states.py (2)
46-49: Policy choice: raise vs. defer when unite parent is absentCurrently a missing
unites.identifierinparentsraises and halts next-state creation. If legitimate flows exist where the unite parent isn’t available yet, consider soft-failing (warn + return False) to defer creation instead of erroring.Proposed alternative:
- if not unites_id: - raise ValueError(f"Unit identifier not found in parents: {node_template.unites.identifier}") + if not unites_id: + logger.warning("Unites identifier '%s' not found in parents for %s/%s; deferring.", + node_template.unites.identifier, namespace, graph_name) + return FalseNote: This change will require updating tests that currently expect a ValueError.
117-121: Potential TypeError when concatenating non-string outputs; coerce to str
DependentString.generate_stringconcatenates strings; parent/current outputs can be int/bool/etc. Coerce values to str at assignment to avoid runtime errors.- dependency_string.dependents[key].value = current_state.outputs[dependency_string.dependents[key].field] + dependency_string.dependents[key].value = str( + current_state.outputs[dependency_string.dependents[key].field] + ) @@ - dependency_string.dependents[key].value = parents[dependency_string.dependents[key].identifier].outputs[dependency_string.dependents[key].field] + dependency_string.dependents[key].value = str( + parents[dependency_string.dependents[key].identifier].outputs[dependency_string.dependents[key].field] + )
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (8)
state-manager/app/tasks/create_next_states.py(3 hunks)state-manager/tests/unit/models/test_base.py(1 hunks)state-manager/tests/unit/models/test_graph_template_model.py(1 hunks)state-manager/tests/unit/singletons/test_logs_manager.py(1 hunks)state-manager/tests/unit/tasks/test_create_next_states.py(1 hunks)state-manager/tests/unit/tasks/test_verify_graph.py(1 hunks)state-manager/tests/unit/test_main.py(3 hunks)state-manager/tests/unit/test_routes.py(1 hunks)
🧰 Additional context used
🧬 Code graph analysis (7)
state-manager/tests/unit/models/test_graph_template_model.py (1)
state-manager/app/models/db/graph_template_model.py (3)
GraphTemplate(15-124)validate_secrets(46-56)_validate_secret_value(59-71)
state-manager/tests/unit/tasks/test_verify_graph.py (4)
state-manager/app/tasks/verify_graph.py (10)
verify_graph(234-265)verify_nodes_names(10-13)verify_nodes_namespace(15-18)verify_node_exists(20-27)verify_node_identifiers(29-55)verify_secrets(57-73)get_database_nodes(76-91)build_dependencies_graph(150-159)verify_topology(161-220)verify_unites(222-231)state-manager/app/models/graph_template_validation_status.py (1)
GraphTemplateValidationStatus(4-8)state-manager/app/models/node_template_model.py (2)
NodeTemplate(9-15)Unites(5-6)state-manager/app/models/db/registered_node.py (1)
RegisteredNode(6-13)
state-manager/tests/unit/test_routes.py (10)
state-manager/app/models/enqueue_request.py (1)
EnqueueRequestModel(4-6)state-manager/app/models/create_models.py (2)
TriggerGraphRequestModel(32-33)CreateRequestModel(22-24)state-manager/app/models/executed_models.py (1)
ExecutedRequestModel(5-6)state-manager/app/models/errored_models.py (1)
ErroredRequestModel(5-6)state-manager/app/models/graph_models.py (2)
UpsertGraphTemplateRequest(8-10)UpsertGraphTemplateResponse(13-19)state-manager/app/models/register_nodes_request.py (1)
RegisterNodesRequestModel(12-14)state-manager/app/models/secrets_response.py (1)
SecretsResponseModel(5-6)state-manager/app/models/list_models.py (2)
ListRegisteredNodesResponse(12-16)ListGraphTemplatesResponse(19-23)state-manager/app/models/state_list_models.py (2)
StatesByRunIdResponse(29-34)CurrentStatesResponse(37-42)state-manager/app/routes.py (11)
enqueue_state(52-62)trigger_graph_route(72-82)create_state(92-102)executed_state_route(112-122)errored_state_route(132-142)register_nodes_route(190-199)get_secrets_route(209-219)list_registered_nodes_route(229-243)list_graph_templates_route(253-267)get_states_by_run_id_route(328-367)get_current_states_route(277-318)
state-manager/tests/unit/singletons/test_logs_manager.py (1)
state-manager/app/singletons/logs_manager.py (1)
_is_development_mode(47-63)
state-manager/tests/unit/models/test_base.py (1)
state-manager/app/models/db/base.py (2)
BaseDatabaseModel(7-15)update_updated_at(14-15)
state-manager/tests/unit/test_main.py (1)
state-manager/app/main.py (1)
health(77-78)
state-manager/app/tasks/create_next_states.py (5)
state-manager/app/models/db/graph_template_model.py (2)
get(97-101)get_node_by_identifier(36-42)state-manager/app/models/db/state.py (1)
State(8-18)state-manager/app/models/state_status_enum.py (1)
StateStatusEnum(4-15)state-manager/app/models/node_template_model.py (1)
NodeTemplate(9-15)state-manager/app/tasks/verify_graph.py (1)
verify_unites(222-231)
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Actionable comments posted: 2
♻️ Duplicate comments (3)
state-manager/app/tasks/create_next_states.py (3)
42-59: Unites satisfaction: prefer soft-fail when unite-parent not present; also rename param for clarity and keep query performantRaising on Line 48 turns a common “not all prerequisites are visible yet” scenario into a hard failure. Returning False (defer creation) with a warning is safer and aligns with two-phase creation. Also, rename
parentstoparent_idsto disambiguate from theparentsdict of State objects used elsewhere.Apply this diff within this range:
-async def check_unites_satisfied(namespace: str, graph_name: str, node_template: NodeTemplate, parents: dict[str, PydanticObjectId]) -> bool: +async def check_unites_satisfied(namespace: str, graph_name: str, node_template: NodeTemplate, parent_ids: dict[str, PydanticObjectId]) -> bool: @@ - unites_id = parents.get(node_template.unites.identifier) + unites_id = parent_ids.get(node_template.unites.identifier) if not unites_id: - raise ValueError(f"Unit identifier not found in parents: {node_template.unites.identifier}") + logger.warning( + "Unites parent '%s' not present in parent_ids for %s/%s; deferring unite creation", + node_template.unites.identifier, namespace, graph_name + ) + return FalseFollow-up (outside this hunk): update the sole caller accordingly (see Lines 224-225).
Indexing note (reiterating a prior review): consider a compound index on
- { namespace_name: 1, graph_name: 1, "parents.": 1, status: 1 }
or a wildcard on parents with partial filter (status != SUCCESS) to keep this check O(log N) at scale.
107-121: Coerce dependent values to str to prevent runtime TypeError during concatenation
Dependent.valuemay receive non-strings (int/bool/etc.).generate_string()concatenates them withbaseandtail, which will raise a TypeError. Cast when assigning.for key in sorted(dependency_string.dependents.keys()): if dependency_string.dependents[key].identifier == current_state.identifier: if dependency_string.dependents[key].field not in current_state.outputs: raise AttributeError(f"Output field '{dependency_string.dependents[key].field}' not found on current state '{current_state.identifier}' for template '{next_state_node_template.identifier}'") - dependency_string.dependents[key].value = current_state.outputs[dependency_string.dependents[key].field] + dependency_string.dependents[key].value = str( + current_state.outputs[dependency_string.dependents[key].field] + ) else: - dependency_string.dependents[key].value = parents[dependency_string.dependents[key].identifier].outputs[dependency_string.dependents[key].field] + dependency_string.dependents[key].value = str( + parents[dependency_string.dependents[key].identifier].outputs[dependency_string.dependents[key].field] + )
227-233: Validate unites inputs against the actual “current” context (the unite parent), not the triggering node
validate_dependencies(..., identifier, parents)uses the triggering node’s identifier, butgenerate_next_stateresolves placeholders withparent_statefor unites. Validate withparent_state.identifierafter you resolve it.- next_state_input_model = await get_input_model(next_state_node_template) - validate_dependencies(next_state_node_template, next_state_input_model, identifier, parents) + next_state_input_model = await get_input_model(next_state_node_template) @@ assert next_state_node_template.unites is not None parent_state = parents[next_state_node_template.unites.identifier] + # Validate with the actual current context for unites + validate_dependencies(next_state_node_template, next_state_input_model, parent_state.identifier, parents)
📜 Review details
Configuration used: CodeRabbit UI
Review profile: ASSERTIVE
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (1)
state-manager/app/tasks/create_next_states.py(3 hunks)
🧰 Additional context used
🧬 Code graph analysis (1)
state-manager/app/tasks/create_next_states.py (5)
state-manager/app/models/db/graph_template_model.py (2)
get(97-101)get_node_by_identifier(36-42)state-manager/app/models/db/state.py (1)
State(8-18)state-manager/app/models/state_status_enum.py (1)
StateStatusEnum(4-15)state-manager/app/models/node_template_model.py (2)
NodeTemplate(9-15)Unites(5-6)state-manager/app/tasks/verify_graph.py (1)
verify_unites(222-231)
🔇 Additional comments (4)
state-manager/app/tasks/create_next_states.py (4)
123-126: Good fix: inherit instance parents and use current state's identifier as the parent keySwitching from the class attribute to
current_state.parentsand keying bycurrent_state.identifieravoids the FieldInfo unpack TypeError and prevents key-mismatch bugs.
203-205: Pending unites now store identifiers (not templates or dependency IDs) — correctThis resolves earlier misalignment where pending items were either template objects or the dependency’s identifier. The subsequent lookup by identifier is now consistent.
Also applies to: 219-223
224-225: Update caller to matchcheck_unites_satisfiedsignature changeIf you accept the soft-fail + rename, adjust the call to pass
parent_ids.- if not await check_unites_satisfied(namespace, graph_name, next_state_node_template, parents_ids): + if not await check_unites_satisfied(namespace, graph_name, next_state_node_template, parents_ids): # name unchanged here; ensure the function signature also uses parent_ids # If you rename locally, keep names consistent across definition and call.Note: If you rename only the parameter in the function definition, this call site remains syntactically the same; include this reminder to keep naming consistent everywhere for readability.
Likely an incorrect or invalid review comment.
217-235: Unite nodes currently consume only a single “anchor” parent, not all completed branchesThe loop at state-manager/app/tasks/create_next_states.py (lines 217–235) picks just one
parent_state(theunites.identifier) when creating a new unite state. If your intent is full join semantics—i.e. waiting for all branches to succeed and making every sibling’s outputs available—you’ll need to:
- After
check_unites_satisfied(...), gather all parent states whoseunites.identifiermatches, not just the anchor.- Pass that collection (or their outputs) into
generate_next_stateso templates can resolve any sibling-branch placeholders.- Update or overload
generate_next_stateto accept acontext: dict[str, State](or similar) representing all merged branches.- Adjust the new state’s lineage (its
parentsmap) to include every branch tip, or at least provide a resolvable context for placeholders.Alternatively, if the design is that a unite only ever cares about its anchor parent, then:
- Document this contract clearly in the code and/or README.
- Add an assertion in
validate_dependencies(or upstream) that no placeholders in the unite node reference any other parent.Please review your templates and dependency patterns to confirm which behavior you need and update either the code or documentation accordingly.
These changes enhance the clarity and efficiency of state management within the application.