Skip to content

Conversation

@NiveditJain
Copy link
Member

  • 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.

…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.
@coderabbitai
Copy link
Contributor

coderabbitai bot commented Aug 21, 2025

Note

Other AI code review bot(s) detected

CodeRabbit 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.

📝 Walkthrough

Summary by CodeRabbit

  • Refactor

    • Node templates now accept a single “unites” entry instead of a list, updating verification and state creation behavior.
    • Clearer two-phase next-state processing and stricter errors when required dependencies or identifiers are missing.
  • Tests

    • Extensive new unit tests covering state creation workflows, graph verification, models, logging, routes, and app setup.
    • Improved validation coverage for secrets, request/response models, and router surfaces.

Walkthrough

NodeTemplate.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

Cohort / File(s) Summary
Model
state-manager/app/models/node_template_model.py
NodeTemplate.unites type changed from Optional[List[Unites]] to Optional[Unites].
Next-state generation logic
state-manager/app/tasks/create_next_states.py
Added generate_next_state(...); refactored create_next_states into two phases (create non-unit states then unit states); replaced list-based unites iteration with single-unite handling; centralized dependency resolution; adjusted error handling (ValueError/AttributeError).
Graph verification
state-manager/app/tasks/verify_graph.py
verify_unites now treats node.unites as a single object and validates presence of node.unites.identifier in the dependency graph (removed per-unit loop).
Unit tests — tasks
state-manager/tests/unit/tasks/test_create_next_states.py, state-manager/tests/unit/tasks/test_verify_graph.py
Tests rewritten/expanded to exercise new helpers (Dependent, DependentString, mark_success_states, check_unites_satisfied, get_dependents, validate_dependencies) and single-unite behavior; replaced local stubs with mocked State/query interactions.
Unit tests — models & singletons & routes
state-manager/tests/unit/models/test_base.py, state-manager/tests/unit/models/test_graph_template_model.py, state-manager/tests/unit/singletons/test_logs_manager.py, state-manager/tests/unit/test_main.py, state-manager/tests/unit/test_routes.py
New and updated tests added for BaseDatabaseModel, GraphTemplate secret validation, LogsManager singleton/behavior, app route/health checks, and route/model validation.

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)
Loading

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Possibly related PRs

  • Fixes #206 #209 — Alters NodeTemplate.unites shape and updates create_next_states/verify_unites logic; strong overlap on unites handling.
  • added bulk_create_next #236 — Touches create_next_states and dependency-resolution helpers; likely related to the same refactor areas.

Suggested reviewers

  • nk-ag

Poem

"I hopped through code with careful paws,
One Unites now guides the cause.
I stitched the flow in two small hops,
Checked parents first, then finished props.
— 🐰"

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 Docstrings
🧪 Generate unit tests
  • Create PR with unit tests
  • Post copyable unit tests in a comment

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

❤️ Share
🪧 Tips

Chat

There are 3 ways to chat with CodeRabbit:

  • Review comments: Directly reply to a review comment made by CodeRabbit. Example:
    • I pushed a fix in commit <commit_id>, please review it.
    • Open a follow-up GitHub issue for this discussion.
  • Files and specific lines of code (under the "Files changed" tab): Tag @coderabbitai in a new review comment at the desired location with your query.
  • PR comments: Tag @coderabbitai in a new PR comment to ask questions about the PR branch. For the best results, please provide a very specific query, as very limited context is provided in this mode. Examples:
    • @coderabbitai gather interesting stats about this repository and render them as a table. Additionally, render a pie chart showing the language distribution in the codebase.
    • @coderabbitai read the files in the src/scheduler package and generate a class diagram using mermaid and a README in the markdown format.

Support

Need help? Create a ticket on our support page for assistance with any issues or questions.

CodeRabbit Commands (Invoked using PR/Issue comments)

Type @coderabbitai help to get the list of available commands.

Other keywords and placeholders

  • Add @coderabbitai ignore anywhere in the PR description to prevent this PR from being reviewed.
  • Add @coderabbitai summary to generate the high-level summary at a specific location in the PR description.
  • Add @coderabbitai anywhere in the PR title to generate the title automatically.

CodeRabbit Configuration File (.coderabbit.yaml)

  • You can programmatically configure CodeRabbit by adding a .coderabbit.yaml file to the root of your repository.
  • Please see the configuration documentation for more information.
  • If your editor has YAML language server enabled, you can add the path at the top of this file to enable auto-completion and validation: # yaml-language-server: $schema=https://coderabbit.ai/integrations/schema.v2.json

Status, Documentation and Community

  • Visit our Status Page to check the current availability of CodeRabbit.
  • Visit our Documentation for detailed information on how to use CodeRabbit.
  • Join our Discord Community to get help, request features, and share feedback.
  • Follow us on X/Twitter for updates and announcements.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a 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

  1. 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.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a 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
Copy link

codecov bot commented Aug 21, 2025

Codecov Report

❌ Patch coverage is 23.80952% with 32 lines in your changes missing coverage. Please review.

Files with missing lines Patch % Lines
state-manager/app/tasks/create_next_states.py 15.78% 32 Missing ⚠️

📢 Thoughts on this report? Let us know!

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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.

📥 Commits

Reviewing files that changed from the base of the PR and between e49d90a and bb16cf7.

📒 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 correct

The 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 identifiers

The function now properly raises a ValueError when 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 generation

The new generate_next_state function 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 dependencies

The refactored logic properly separates regular next states from unite-dependent states, ensuring unite conditions are satisfied before creating dependent states. The use of pending_unites to defer unite state creation is a clean approach.


232-232: Verify correct current_state passed to generate_next_state

  • In state-manager/app/tasks/create_next_states.py (line 232), the call
    new_unit_states.append(
        generate_next_state(
            next_state_input_model,
            next_state_node_template,
            identifier,
            parents,
            parent_state
        )
    )
    passes parent_state as the current_state parameter to generate_next_state.
  • Please confirm the signature of generate_next_state expects that final argument to be the “triggering” state, not the parent of the dependency.
  • Audit all other invocations of generate_next_state to 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.

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 cases

To ensure the new single-Unites model is fully exercised, add two more async tests in state-manager/tests/unit/tasks/test_create_next_states.py alongside the existing test_check_unites_satisfied_true/false:

• Test that when node_tpl.unites is set to None, check_unites_satisfied immediately returns True without inspecting the parents mapping.
• Test that when node_tpl.unites.identifier is not present in the supplied parents dict, check_unites_satisfied raises a ValueError("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.

📥 Commits

Reviewing files that changed from the base of the PR and between bb16cf7 and aa05226.

📒 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 signature

Using PydanticObjectId directly in tests matches the cns.check_unites_satisfied type contract and avoids stringly-typed IDs. Looks good.

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 calls

validate_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_unites loop), 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.

📥 Commits

Reviewing files that changed from the base of the PR and between aa05226 and d28ff50.

📒 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 the GraphTemplate class definition, check whether get_valid (and fallback get) are defined on it, and list all call sites. Once you have its output, we can confirm if get_valid truly 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.
Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 correct PydanticObjectId. 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.

📥 Commits

Reviewing files that changed from the base of the PR and between d28ff50 and fd60994.

📒 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 env

This test clearly validates the intended priority of command-line over environment variable.


168-181: Good sanity checks on logger methods

Validates that basic logging calls don’t raise—useful as a smoke test.

@NiveditJain
Copy link
Member Author

/gemini review

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a 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.

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 check

Catching Exception is 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 wording

This 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_at is 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 characters

Construct 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 bytes

Add 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 mock

Currently 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 routes

Equality 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: ignore
state-manager/tests/unit/test_main.py (4)

32-46: Strengthen health route detection: assert APIRoute, endpoint binding, and method

Use FastAPI’s APIRoute and verify the bound endpoint is app_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_model is None unless explicitly set; verify response_class instead.

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
+                break

Add 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 slash

Reduces 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 instead

Asserting instance1 is instance2 couples 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 test

Seven 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_logger

This 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 runs

Use a module reload or clear the singleton cache, then assert configure is 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 approach
state-manager/tests/unit/tasks/test_create_next_states.py (3)

203-218: Assert unites satisfaction query includes the parents. equality filter

Strengthens contract of check_unites_satisfied by 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 exist

Prevents 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 AttributeError

Locks in current behavior of validate_dependencies when 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 absent

Currently a missing unites.identifier in parents raises 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 False

Note: 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_string concatenates 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.

📥 Commits

Reviewing files that changed from the base of the PR and between fd60994 and 7219c0d.

📒 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)

Copy link
Contributor

@coderabbitai coderabbitai bot left a 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 performant

Raising 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 parents to parent_ids to disambiguate from the parents dict 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 False

Follow-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.value may receive non-strings (int/bool/etc.). generate_string() concatenates them with base and tail, 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, but generate_next_state resolves placeholders with parent_state for unites. Validate with parent_state.identifier after 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.

📥 Commits

Reviewing files that changed from the base of the PR and between 7219c0d and 7a60904.

📒 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 key

Switching from the class attribute to current_state.parents and keying by current_state.identifier avoids the FieldInfo unpack TypeError and prevents key-mismatch bugs.


203-205: Pending unites now store identifiers (not templates or dependency IDs) — correct

This 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 match check_unites_satisfied signature change

If 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 branches

The loop at state-manager/app/tasks/create_next_states.py (lines 217–235) picks just one parent_state (the unites.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 whose unites.identifier matches, not just the anchor.
  • Pass that collection (or their outputs) into generate_next_state so templates can resolve any sibling-branch placeholders.
  • Update or overload generate_next_state to accept a context: dict[str, State] (or similar) representing all merged branches.
  • Adjust the new state’s lineage (its parents map) 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.

@NiveditJain NiveditJain merged commit 1e3316c into exospherehost:main Aug 21, 2025
4 of 5 checks passed
@coderabbitai coderabbitai bot mentioned this pull request Sep 3, 2025
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