Refine async thread safety with aiologic locks and reasoning tools

[JohnRichard4096]

Summary by Sourcery

Switch thread-safety primitives to aiologic-based locks and adjust agent tool handling to support reasoning tools per-call while tightening concurrency tests and versioning.

Bug Fixes:

• Fix async context lock behavior in ContextThreadsafe to prevent deadlocks and ensure proper acquire/release flow.
• Ensure AdapterManager and MultiClientManager initialize their thread-safe context and locks correctly to avoid shared or uninitialized lock issues.
• Correct agent tool selection so reasoning tools are added only when needed for a given call instead of being always present.

Enhancements:

• Replace global and callback locks in chat manager and tools subsystem with aiologic.Lock for consistent async/thread-safe behavior.
• Refine thread-safety test to simulate concurrent acquisition and verify lock contention and re-entry behavior more accurately.

Build:

• Bump amrita_core version to 0.8.1 and add aiologic as a runtime dependency.

CI:

• Add a timeout to the CI build job to prevent runaway workflows.

[sourcery-ai]

Reviewer's Guide

Introduces a threadsafe async lock abstraction that supports Python 3.13+ free-threaded mode, refactors ContextThreadsafe to use it, adjusts agent tool selection so the reasoning tool is added per-call instead of at construction, ensures ProtocolSingleton runs its base class init, and bumps the project version to 0.8.1.

Sequence diagram for per-call agent tool selection

sequenceDiagram
    actor User
    participant Agent
    participant StrategyContext as StrategyContext
    participant Config as ConfigBuiltin
    participant ToolsCaller

    User->>Agent: single_execute(msg_list)
    Agent->>StrategyContext: get config
    StrategyContext-->>Agent: config
    Agent->>Config: read tool_calling_mode
    alt tool_calling_mode == agent
        Agent->>Agent: ensure STOP_TOOL in self.tools
    else tool_calling_mode == none
        Agent-->>User: return False
        Agent->>Agent: return
    end

    Agent->>Agent: tools = self.tools.copy()
    Agent->>Config: read agent_thought_mode
    alt agent_thought_mode startswith reasoning
        Agent->>Agent: tools.append(REASONING_TOOL)
    end

    Agent->>ToolsCaller: tools_caller(msg_list, tools, tool_choice,...)
    ToolsCaller-->>Agent: response_msg
    Agent-->>User: processed result

Loading

Class diagram for AsyncLockThreadsafe and ContextThreadsafe

classDiagram
    class AsyncLockThreadsafe {
        +asyncio.Lock _async_lock
        +threading.Lock _thread_lock
        +threading.Lock _meta_lock
        +int _owner_thread_id
        +bool _is_thread_locked
        +AsyncLockThreadsafe()
        +acquire() async
        +release() void
        +__aenter__() async AsyncLockThreadsafe
        +__aexit__(exc_type, exc_val, exc_tb) async void
    }

    class ContextThreadsafe {
        +AsyncLockThreadsafe _ctx_lock
        +ContextThreadsafe()
        +__aenter__() async void
        +__aexit__(exc_type, exc_val, exc_tb) async void
    }

    ContextThreadsafe --> AsyncLockThreadsafe : wraps

Loading

Class diagram for ProtocolSingleton initialization behavior

classDiagram
    class ModelAdapter {
        <<interface>>
    }

    class ProtocolSingleton {
        -static ProtocolSingleton __instance
        -dict~string, ModelAdapter~ _adapter_class
        +__new__() ProtocolSingleton
        +__init__() void
        +get_adapters() dict~string, ModelAdapter~
    }

    ProtocolSingleton --> ModelAdapter : stores adapter types

Loading

File-Level Changes

Change	Details	Files
Add a cross-thread/asyncio lock abstraction to support free-threaded Python builds and refactor the context helper to use it.	Introduce AsyncLockThreadsafe that composes asyncio.Lock with threading primitives in no-GIL builds and wraps asyncio.Lock in older versions Track ownership metadata (owner thread id, lock mode) to safely acquire/release the lock from both coroutines and threads Refactor ContextThreadsafe into an instance-based wrapper around AsyncLockThreadsafe and improve its documentation	src/amrita_core/threadsafe.py
Adjust agent tool selection so the reasoning tool is attached at call time based on the current configuration instead of being permanently added at initialization.	Stop appending the reasoning tool in the agent constructor when agent_thought_mode starts with 'reasoning' Clone the base tools list inside single_execute and append the reasoning tool dynamically when the reasoning mode is enabled, before invoking tools_caller	src/amrita_core/builtins/agent.py
Ensure the protocol singleton correctly initializes its superclass state.	Add an explicit init that calls the base class initializer while preserving the singleton behavior	src/amrita_core/protocol.py
Bump library version and refresh lockfile.	Update project version from 0.8.0 to 0.8.1 in the project configuration Regenerate the dependency lock file to reflect the new version or dependency changes	pyproject.toml uv.lock

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey - I've found 2 issues, and left some high level feedback:

• In AsyncLockThreadsafe.acquire, the path that calls await asyncio.to_thread(self._thread_lock.acquire) will block cancellation and doesn't track acquisition depth, so if a coroutine is cancelled while waiting or a caller tries to re-enter the lock from the same thread, state can become subtle; consider making the lock explicitly non‑reentrant (e.g., with an assertion) and/or documenting and guarding against cancellation while waiting on the thread lock.
• The addition of __init__ in ModelAdapterProtocol calls super().__init__() on every access to the singleton, which is unusual for a singleton pattern; if ModelAdapterProtocol or its superclass ever add state in __init__, this will re-run initialization repeatedly, so you may want to guard initialization with a flag or move it into __new__.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- In `AsyncLockThreadsafe.acquire`, the path that calls `await asyncio.to_thread(self._thread_lock.acquire)` will block cancellation and doesn't track acquisition depth, so if a coroutine is cancelled while waiting or a caller tries to re-enter the lock from the same thread, state can become subtle; consider making the lock explicitly non‑reentrant (e.g., with an assertion) and/or documenting and guarding against cancellation while waiting on the thread lock.
- The addition of `__init__` in `ModelAdapterProtocol` calls `super().__init__()` on every access to the singleton, which is unusual for a singleton pattern; if `ModelAdapterProtocol` or its superclass ever add state in `__init__`, this will re-run initialization repeatedly, so you may want to guard initialization with a flag or move it into `__new__`.

## Individual Comments

### Comment 1
<location path="src/amrita_core/threadsafe.py" line_range="51" />
<code_context>
+        if take_async:
+            await self._async_lock.acquire()
+        else:
+            await asyncio.to_thread(self._thread_lock.acquire)
+            with self._meta_lock:
+                self._owner_thread_id = current_tid
</code_context>
<issue_to_address>
**issue (bug_risk):** Cancellation while waiting on `to_thread(self._thread_lock.acquire)` can leave the lock acquired with no corresponding release.

If the coroutine is cancelled while `asyncio.to_thread(self._thread_lock.acquire)` is running, the worker thread may successfully acquire the lock, but the coroutine will never run the subsequent code to record ownership or ensure a matching `release()`, leaving the lock taken indefinitely. Please handle `asyncio.CancelledError` so that a successfully acquired lock is released on cancellation, or switch to a non-blocking acquire loop that cooperates safely with cancellation.
</issue_to_address>

### Comment 2
<location path="src/amrita_core/threadsafe.py" line_range="80-89" />
<code_context>
+    async def __aexit__(self, exc_type, exc_val, exc_tb) -> None:
+        self.release()

 class ContextThreadsafe:
-    """Add a optional async-context lock for thread-safe."""
+    """
+    Provides an optional async context manager for thread-safe operations.
+
+    This class wraps AsyncLockThreadsafe to provide a convenient context manager
+    interface for ensuring thread safety in async code.
+    """
+
+    _ctx_lock: AsyncLockThreadsafe

-    _lock = asyncio.Lock()
+    def __init__(self):
+        self._ctx_lock = AsyncLockThreadsafe()

     async def __aenter__(self):
</code_context>
<issue_to_address>
**question (bug_risk):** Changing from a class-level lock to a per-instance lock alters synchronization semantics and may break assumptions of global mutual exclusion.

Previously `_lock` was a class attribute, so all `ContextThreadsafe` uses shared a single `asyncio.Lock`. Now each instance has its own `AsyncLockThreadsafe`, so different `ContextThreadsafe` instances can run concurrently. If any code relies on this being a process/module-wide critical section, this change could introduce race conditions. If global mutual exclusion is still required, use a shared (class- or module-level) `AsyncLockThreadsafe` instead of an instance-level one.
</issue_to_address>

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[JohnRichard4096]

Well, I think I should consider using AioLogic instead of other locks.

[JohnRichard4096]

@sourcery-ai title

[JohnRichard4096]

@sourcery-ai summary