fix(concurrency): isolate per-session state for concurrent agent safety#400
fix(concurrency): isolate per-session state for concurrent agent safety#400oysteinkrog wants to merge 2 commits intoBeehiveInnovations:mainfrom
Conversation
Summary of ChangesHello @oysteinkrog, 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 significantly enhances the concurrency safety of workflow tools, particularly the Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Changelog
Activity
Using Gemini Code AssistThe 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 by creating a comment using either
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 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
|
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 14914b5f58
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # Workflow complete - release per-session state to free memory. | ||
| self._cleanup_session(continuation_id) |
There was a problem hiding this comment.
In the final consensus step, session state is deleted before customize_workflow_response() and _customize_consensus_metadata() run. Those methods call _get_session(request.continuation_id), which recreates a fresh empty session, so the final payload can lose accumulated_responses and report empty/zero model metadata (models_consulted, total_models) despite a completed workflow. This happens whenever a normal multi-step consensus run reaches its last step with a continuation_id.
Useful? React with 👍 / 👎.
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request effectively addresses a concurrency issue in workflow tools by introducing an asyncio.Lock for serialization and, in the ConsensusTool, implementing per-session state management to ensure thread safety. The changes are well-structured, splitting execute_workflow into a locking wrapper and an inner method, and isolating session-specific data in ConsensusTool. My review focuses on refining the implementation of session isolation. Specifically, I've suggested making updates to legacy instance variables conditional to prevent session state from leaking into the shared tool instance, which will make the concurrency fix more robust and less confusing. I also recommended clarifying a comment in _get_session to better explain its behavior for the legacy path. Overall, this is a solid improvement for the stability of concurrent agent operations.
| # Keep legacy instance vars in sync (single-caller / no-continuation_id path). | ||
| self.models_to_consult = session["models_to_consult"] | ||
| self.accumulated_responses = session["accumulated_responses"] |
There was a problem hiding this comment.
These lines unconditionally update instance variables (self.models_to_consult, self.accumulated_responses) with session-specific data. While the lock prevents race conditions, this pollutes the shared instance state with data from a specific session, which undermines the goal of session isolation.
The comment indicates this is for the "single-caller / no-continuation_id path". To align with that and ensure true session isolation, these updates should be conditional.
# Keep legacy instance vars in sync (single-caller / no-continuation_id path).
if not continuation_id:
self.models_to_consult = session["models_to_consult"]
self.accumulated_responses = session["accumulated_responses"]| def store_initial_issue(self, step_description: str, continuation_id: str | None = None): | ||
| """Store initial prompt for model consultations.""" | ||
| # Write into the per-session dict so concurrent callers stay isolated. | ||
| session = self._get_session(continuation_id) | ||
| session["original_proposal"] = step_description | ||
| # Also keep legacy instance vars for any code that reads them directly. | ||
| self.original_proposal = step_description | ||
| self.initial_prompt = step_description # Keep for backward compatibility | ||
|
|
There was a problem hiding this comment.
Similar to another comment, updating self.original_proposal and self.initial_prompt unconditionally pollutes the shared instance state with data from a specific session. This goes against the principle of session isolation you're introducing.
The comment "Also keep legacy instance vars for any code that reads them directly" suggests this is for backward compatibility. Therefore, it should only apply to the legacy path where continuation_id is not present.
def store_initial_issue(self, step_description: str, continuation_id: str | None = None):
"""Store initial prompt for model consultations."""
# Write into the per-session dict so concurrent callers stay isolated.
session = self._get_session(continuation_id)
session["original_proposal"] = step_description
# Also keep legacy instance vars for any code that reads them directly.
if not continuation_id:
self.original_proposal = step_description
self.initial_prompt = step_description # Keep for backward compatibility| """Return the per-session state dict for *continuation_id*. | ||
|
|
||
| When no continuation_id is provided (legacy / single-caller path) a | ||
| transient dict backed by the instance variables is returned so the | ||
| rest of the code works without special-casing. | ||
| """ |
There was a problem hiding this comment.
The comment "callers can read/write through the same interface" is a bit misleading for the legacy path (continuation_id is None). The returned dictionary is a new dictionary containing references to the instance's mutable attributes (like lists).
- In-place modifications to these mutable objects (e.g.,
session['accumulated_responses'].append(...)) will be reflected on the instance. - However, re-assigning a key in the returned dictionary (e.g.,
session['original_proposal'] = 'new') will not update the corresponding instance attribute.
This subtlety can be error-prone. While the current code seems to handle it correctly by manually syncing instance attributes, clarifying the comment would prevent future confusion.
"""Return the per-session state dict for *continuation_id*.
When no continuation_id is provided (legacy / single-caller path), a
transient dict is returned. Note that only in-place modifications to
mutable objects (like lists) in this dict will affect the instance state.
Direct assignments to keys will not be persisted back to the instance.
"""
|
All four review comments addressed in commit Codex P2 — cleanup before customize: Moved Gemini high (×2) — unconditional instance-var pollution: Wrapped both sync sites behind
When a Gemini medium — _get_session docstring: Updated to clarify the snapshot semantics: in-place mutations on mutable values (lists) propagate to the instance, but scalar key reassignments do not — callers must sync those manually via the Own review notes (for transparency):
|
Pre-existing files that needed reformatting to satisfy the CI black check after upgrading to black 26.1.0. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Fixes BeehiveInnovations#395. All workflow tools share a single singleton instance in server.py. When multiple agents call the same tool concurrently, they race on shared self.* mutable state at asyncio await yield points. Changes: - workflow_mixin.py: add asyncio.Lock per tool instance; split execute_workflow into _execute_workflow_locked so subclasses can call the locked body without re-acquiring the lock - consensus.py: add _sessions dict keyed by continuation_id to isolate original_proposal, models_to_consult, accumulated_responses per caller; restore work_history/consolidated_findings from conversation memory on continuation steps (base class logic that the custom consensus loop bypassed); fix off-by-one in customize_workflow_response status condition (step_number < total_steps, not < total_steps - 1) - all other workflow tools (analyze, codereview, planner, precommit, refactor, secaudit, testgen, tracer): update store_initial_issue call signature to match updated base class
oysteinkrog
force-pushed
the
fix/concurrent-agent-safety
branch
from
f770f5e to
86ea859
Compare
…ty (from PR BeehiveInnovations#400) Adds asyncio.Lock per tool instance and splits execute_workflow into a locking wrapper to prevent concurrent callers from silently overwriting each other's state.
1 participant
Summary
Fixes #395 — ConsensusTool (and all workflow tools) share mutable instance state across concurrent callers on the singleton tool instances, causing one agent's data to silently overwrite another's.
tools/workflow/workflow_mixin.py: Addasyncio.Lockper tool instance; splitexecute_workflowinto a locking wrapper +_execute_workflow_lockedinner method so concurrent calls to the same tool type are serialized. Updatestore_initial_issuebase signature to accept an optionalcontinuation_idparameter.tools/consensus.py: Add a per-session_sessions: dict[str, dict]keyed bycontinuation_idto fully isolateoriginal_proposal,models_to_consult, andaccumulated_responsesper caller. Splitexecute_workflowinto a locking wrapper +_execute_consensus_workflow_locked(callssuper()._execute_workflow_lockedto avoid re-entrant deadlock on the sameasyncio.Lock). Clean up the session dict on final step completion.tools/{analyze,codereview,planner,precommit,refactor,secaudit,testgen,tracer}.py: Addcontinuation_id: str | None = Noneto theirstore_initial_issueoverrides to match the updated base class signature.Root Cause
All tools are instantiated once in
server.pyas singletons (TOOLS = {"chat": ChatTool(), ...}). The asyncio event loop is single-threaded, but at everyawaityield point (waiting on model API responses) a second caller can enter the same method and overwriteself.original_proposal,self.work_history, etc. on the shared instance.Test plan
python -m pytest tests/ -v -m "not integration"— 866 passed, 4 skipped)./code_quality_checks.shpasses (ruff, black, isort, full unit suite)🤖 Generated with Claude Code