Fixes for bugs encountered #205
Conversation
|
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
WalkthroughState handling was changed to use in-memory mutations and .save() instead of DB set operations, parents now store PydanticObjectId values and ObjectId conversions were removed; create_next_state builds an expanded parents map (including current state.id) and marks terminal nodes SUCCESS; controllers now pass raw state IDs to State.get. Changes
Sequence Diagram(s)sequenceDiagram
participant Client
participant ExecCtrl as ExecutedStateController
participant DB as State DB
participant BG as BackgroundTaskScheduler
participant Creator as CreateNextState
Client->>ExecCtrl: POST executed_state (with outputs)
ExecCtrl->>DB: load State (ensure state and state.id)
alt outputs non-empty
ExecCtrl->>DB: mutate state (outputs/status/parents) and await state.save()
else outputs empty
ExecCtrl->>DB: persist empty outputs/status/parents using state.id
end
ExecCtrl->>BG: schedule CreateNextState(state_id)
ExecCtrl-->>Client: 202 Accepted
BG->>Creator: run(state_id)
Creator->>DB: fetch State
alt has next nodes
Creator->>Creator: parents = state.parents ∪ {state.identifier: state.id}
loop each next node
Creator->>DB: create new State with parents
end
Creator->>DB: set original state SUCCESS
else no next nodes
Creator->>DB: set original state SUCCESS and return
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Possibly related PRs
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. 📜 Recent review detailsConfiguration used: CodeRabbit UI 💡 Knowledge Base configuration:
You can enable these sources in your CodeRabbit configuration. 📒 Files selected for processing (1)
🧰 Additional context used🧬 Code Graph Analysis (1)state-manager/app/tasks/create_next_state.py (3)
🔇 Additional comments (5)
✨ 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 (
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @nk-ag, 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 addresses several bugs within the state management system. It refines how state objects are updated, correctly handles scenarios where a state has no subsequent nodes by marking it as successful, and improves the tracking of parent states to ensure accurate lineage for new state creations. Additionally, it updates type hints for Object IDs to align with the beanie ODM.
Highlights
- State Object Update Mechanism: Changed the state update logic in
executed_state.pyfrom usingState.find_one().set()to directly modifying thestateobject properties (outputs,status,parents) and then callingstate.save(). This ensures a more direct and potentially more reliable update of the state document. - Terminal State Handling: Modified
create_next_state.pyto gracefully handle cases where a node template has no next nodes. Instead of raising an exception, the current state's status is now set toSUCCESS, and the function returns, indicating a successful completion of a state branch. - Parent State Tracking: Improved the parent tracking logic in
create_next_state.pyby explicitly including the current state's ID in its ownparentsdictionary before passing it down to subsequent state creations. This ensures that theparentsdictionary accurately reflects the lineage when new states are generated. - Type Hinting for Object IDs: Updated the
Statemodel instate.pyto importPydanticObjectIdfrombeanieand use it as the type hint for theparentsdictionary values, replacing the genericObjectId. This aligns with thebeanieODM's specific type for MongoDB Object IDs.
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.
Code Review
This pull request introduces several important bug fixes and improvements. The fix in executed_state.py to update the state object before passing it to a background task is crucial for correctness. The change in create_next_state.py to correctly handle terminal graph nodes as successful states is also a great improvement. The refactoring to use PydanticObjectId in the State model aligns with beanie best practices.
However, I've found a couple of issues related to the transition to PydanticObjectId. The code in a few places still attempts to use bson.ObjectId where a PydanticObjectId is now expected. I've left specific comments on the affected lines in the diff.
Additionally, please note that the change to PydanticObjectId in state.py likely introduces a bug in create_next_state.py on line 87: dependent_state = await State.get(ObjectId(parent_id)). Since parent_id will now be a PydanticObjectId, ObjectId(parent_id) will likely fail. This should probably be changed to await State.get(parent_id). As this line is not in the current diff, I am mentioning it here for your consideration.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
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/controller/executed_state.py(1 hunks)state-manager/app/models/db/state.py(2 hunks)state-manager/app/tasks/create_next_state.py(3 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (3)
state-manager/app/models/db/state.py (2)
state-manager/app/models/enqueue_response.py (1)
StateModel(6-11)state-manager/app/models/db/base.py (1)
BaseDatabaseModel(7-15)
state-manager/app/controller/executed_state.py (4)
state-manager/app/models/state_status_enum.py (1)
StateStatusEnum(4-14)state-manager/app/routes.py (1)
executed_state_route(84-94)state-manager/app/models/executed_models.py (1)
ExecutedRequestModel(5-6)state-manager/app/controller/errored_state.py (1)
errored_state(11-34)
state-manager/app/tasks/create_next_state.py (1)
state-manager/app/models/state_status_enum.py (1)
StateStatusEnum(4-14)
🔇 Additional comments (2)
state-manager/app/tasks/create_next_state.py (2)
40-42: Good fix: short-circuit leaf nodes by marking SUCCESS.Early return on nodes without successors prevents unnecessary processing and double-finalization later. Looks correct.
113-114: LGTM: new_state inherits the expanded parents mapping.Passing the merged
parentsto the child state ensures placeholders can resolve fields from the current state too.
There was a problem hiding this comment.
Actionable comments posted: 1
🔭 Outside diff range comments (6)
state-manager/app/controller/get_secrets.py (2)
5-5: Fix pipeline failure: remove unused bson.ObjectId importRuff flagged this as unused. Removing it will fix the lint failure.
Apply this diff:
-from bson import ObjectId
28-35: Return proper HTTP status codes instead of ValueErrorRaising ValueError bubbles up as 500s unless an exception handler is installed. Use HTTPException with appropriate status codes (404 for not found) to match executed_state and improve API ergonomics.
Apply this diff to use HTTPException:
@@ -from app.models.db.graph_template_model import GraphTemplate -from bson import ObjectId +from app.models.db.graph_template_model import GraphTemplate +from fastapi import HTTPException, status @@ - if not state: - logger.error(f"State {state_id} not found", x_exosphere_request_id=x_exosphere_request_id) - raise ValueError(f"State {state_id} not found") + if not state: + logger.error(f"State {state_id} not found", x_exosphere_request_id=x_exosphere_request_id) + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"State {state_id} not found") @@ - if state.namespace_name != namespace_name: - logger.error(f"State {state_id} does not belong to namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) - raise ValueError(f"State {state_id} does not belong to namespace {namespace_name}") + if state.namespace_name != namespace_name: + logger.error(f"State {state_id} does not belong to namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"State {state_id} not found in namespace {namespace_name}") @@ - if not graph_template: - logger.error(f"Graph template {state.graph_name} not found in namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) - raise ValueError(f"Graph template {state.graph_name} not found in namespace {namespace_name}") + if not graph_template: + logger.error(f"Graph template {state.graph_name} not found in namespace {namespace_name}", x_exosphere_request_id=x_exosphere_request_id) + raise HTTPException(status_code=status.HTTP_404_NOT_FOUND, detail=f"Graph template {state.graph_name} not found in namespace {namespace_name}") @@ - except Exception as e: + except Exception as e: logger.error(f"Error getting secrets for state {state_id}: {str(e)}", x_exosphere_request_id=x_exosphere_request_id) raise eAlso applies to: 43-45, 54-56
state-manager/app/controller/executed_state.py (1)
2-2: Drop bson.ObjectId in favor of PydanticObjectId and use State.getTo complete the migration away from BSON and align with Beanie IDs end-to-end.
Apply this diff:
-from bson import ObjectId +from beanie import PydanticObjectId @@ -async def executed_state(namespace_name: str, state_id: ObjectId, body: ExecutedRequestModel, x_exosphere_request_id: str, background_tasks: BackgroundTasks) -> ExecutedResponseModel: +async def executed_state(namespace_name: str, state_id: PydanticObjectId, body: ExecutedRequestModel, x_exosphere_request_id: str, background_tasks: BackgroundTasks) -> ExecutedResponseModel: @@ - state = await State.find_one(State.id == state_id) + state = await State.get(state_id)Also applies to: 12-12, 17-17
state-manager/app/tasks/create_next_state.py (3)
4-4: Fix pipeline failure: remove unused bson.ObjectId importbson.ObjectId is no longer used; Ruff flags this (F401).
Apply this diff:
-from bson import ObjectId
63-70: Fix unreachable branch and guard missing inputs in templatelen(splits) == 0 is never true; also, indexing next_node_template.inputs[field_name] can KeyError. Simplify placeholder detection and validate presence.
Apply this diff:
- for field_name, _ in next_node_input_model.model_fields.items(): - temporary_input = next_node_template.inputs[field_name] - splits = temporary_input.split("${{") - - if len(splits) == 0: - next_node_input_data[field_name] = temporary_input - continue + for field_name, _ in next_node_input_model.model_fields.items(): + if field_name not in next_node_template.inputs: + raise Exception(f"Input '{field_name}' not provided in template for node '{next_node_template.identifier}'") + temporary_input = next_node_template.inputs[field_name] + if "${{" not in temporary_input: + next_node_input_data[field_name] = temporary_input + continue + splits = temporary_input.split("${{")
99-99: Be robust to non-string outputs in placeholder substitutionIf dependent_state.outputs[input_field] can be non-string, concatenation will fail. Cast to str.
Apply this diff:
- constructed_string += dependent_state.outputs[input_field] + split.split("}}")[1] + constructed_string += str(dependent_state.outputs[input_field]) + split.split("}}")[1]
♻️ Duplicate comments (1)
state-manager/app/controller/executed_state.py (1)
25-27: Critical: Mixing .set() with passing a stale in-memory State can clobber DB fields laterYou update via a query
.set(...)but then pass the stalestateobject to the background task; create_next_state eventually callsawait state.save(), which can overwrite the DB-updated parents/outputs with stale values. Unify by mutating the in-memorystateand saving it before scheduling.Apply this diff:
- if len(body.outputs) == 0: - await State.find_one(State.id == state_id).set( - {"status": StateStatusEnum.EXECUTED, "outputs": {}, "parents": {**state.parents, state.identifier: state.id}} - ) - - background_tasks.add_task(create_next_state, state) + if len(body.outputs) == 0: + state.outputs = {} + state.status = StateStatusEnum.EXECUTED + state.parents = {**state.parents, state.identifier: state.id} + await state.save() + + background_tasks.add_task(create_next_state, state)
📜 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 (4)
state-manager/app/controller/executed_state.py(2 hunks)state-manager/app/controller/get_secrets.py(1 hunks)state-manager/app/models/db/state.py(2 hunks)state-manager/app/tasks/create_next_state.py(4 hunks)
🧰 Additional context used
🧬 Code Graph Analysis (4)
state-manager/app/controller/get_secrets.py (4)
state-manager/app/models/db/state.py (1)
State(8-18)state-manager/app/routes.py (1)
get_secrets_route(181-191)state-manager/app/models/db/graph_template_model.py (2)
get_secrets(71-74)get_secret(76-81)python-sdk/exospherehost/runtime.py (1)
_get_secrets_endpoint(120-124)
state-manager/app/controller/executed_state.py (3)
state-manager/app/models/state_status_enum.py (1)
StateStatusEnum(4-14)state-manager/app/models/db/state.py (1)
State(8-18)state-manager/app/tasks/create_next_state.py (1)
create_next_state(14-127)
state-manager/app/models/db/state.py (1)
state-manager/app/models/db/base.py (1)
BaseDatabaseModel(7-15)
state-manager/app/tasks/create_next_state.py (2)
state-manager/app/models/state_status_enum.py (1)
StateStatusEnum(4-14)state-manager/app/models/db/state.py (1)
State(8-18)
🪛 GitHub Actions: Ruff check on changed files only
state-manager/app/controller/get_secrets.py
[error] 5-5: Ruff lint error: Unused import 'bson.ObjectId' detected (F401). Remove the import or run 'ruff --fix' to auto-fix.
state-manager/app/tasks/create_next_state.py
[error] 4-4: Ruff lint error: Unused import 'bson.ObjectId' detected (F401). Remove the import or run 'ruff --fix' to auto-fix.
🔇 Additional comments (9)
state-manager/app/models/db/state.py (2)
4-4: Using PydanticObjectId is the right moveSwitching to Beanie’s PydanticObjectId aligns the model with Beanie’s expected ID type and removes BSON coupling.
18-18: No bson.ObjectId wrappers for state.id found
I searched the entire state-manager module for any instances ofObjectId(state.id), re-wrapping inState.get(ObjectId(...)), or in parent maps, and found none. All callers appear to be passing and using the newPydanticObjectIddirectly.state-manager/app/controller/get_secrets.py (1)
27-27: Directly using State.get(state_id) is correct and consistent with PydanticObjectIdThis aligns with the migration away from bson.ObjectId conversions and leverages Beanie’s native ID handling.
state-manager/app/controller/executed_state.py (3)
18-18: Good guard: ensure state is present and has an idThis prevents downstream failures in create_next_state.
32-37: LGTM on in-memory mutation + save pathConsistent with the PR objective to avoid DB .set and pass up-to-date state to the background task.
51-53: Parents now store state.id directly — correct with PydanticObjectIdUsing state.id instead of ObjectId(state.id) keeps types consistent.
state-manager/app/tasks/create_next_state.py (3)
42-44: Terminal-node handling is correctMarking the current state SUCCESS and returning when there are no next nodes matches the intended semantics.
46-49: Correct: include the current state in parents and use state.id directlyparents = state.parents | {state.identifier: state.id} both deduplicates and ensures type correctness.
88-90: LGTM: stop re-wrapping ids with ObjectId when fetchingUsing await State.get(parent_id) avoids type mismatches with PydanticObjectId.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
2 participants
No description provided.