agentx-python 0.8.19__tar.gz → 0.8.21__tar.gz

{agentx_python-0.8.19/agentx_python.egg-info → agentx_python-0.8.21}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: agentx-python
-Version: 0.8.19
+Version: 0.8.21
 Summary: Official Python SDK for AgentX (https://www.agentx.so/)
 Home-page: https://github.com/AgentX-ai/AgentX-python
 Author: Robin Wang and AgentX Team

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/__init__.py

@@ -12,11 +12,10 @@ from agentx.exceptions import (
     CIGateFailure,
 )
 
-logging.basicConfig(
-    level=logging.INFO,
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-    datefmt="%Y-%m-%d %H:%M:%S %Z",
-)
+# Library logging hygiene: a library must never call logging.basicConfig - it hijacks the
+# host application's root logger (format AND level) and turns the app's own later basicConfig
+# into a no-op. Consumers opt into our logs with logging.getLogger("agentx").setLevel(...).
+logging.getLogger("agentx").addHandler(logging.NullHandler())
 
 __all__ = [
     "AgentX",

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/evaluations/client.py

@@ -207,6 +207,12 @@ class EvaluationsClient:
                 return resp.json()
             except Exception:
                 return resp.text
+        # A timeout keeps its type: runner._flush_batch catches requests.Timeout specifically
+        # (the engine may still be scoring the batch - a blind retry double-bills every judge
+        # call), and wrapping it in AgentXEvaluationsError here made that guard unreachable.
+        # Applies to retry=False calls too, where the single attempt lands straight here.
+        if isinstance(last_exc, requests.Timeout):
+            raise last_exc
         raise AgentXEvaluationsError(f"Request failed after retries: {last_exc}")
 
     # ------------------------------------------------------------------
@@ -227,7 +233,9 @@ class EvaluationsClient:
     # ------------------------------------------------------------------
 
     def create_dataset(self, payload: dict) -> Dataset:
-        data = self._request("POST", "/datasets", json=self._with_workspace(payload))
+        # Server-side write: a timeout after the dataset row was created would be
+        # retried into a duplicate dataset, so no transport retry.
+        data = self._request("POST", "/datasets", json=self._with_workspace(payload), retry=False)
         return Dataset(**data)
 
     def delete_dataset(self, dataset_id: str) -> None:
@@ -255,8 +263,9 @@ class EvaluationsClient:
     # ------------------------------------------------------------------
 
     def create_evaluation_settings(self, payload: dict) -> EvaluationSettings:
+        # Server-side write - no transport retry (see init_run's comment).
         data = self._request(
-            "POST", "/evaluation-settings", json=self._with_workspace(payload)
+            "POST", "/evaluation-settings", json=self._with_workspace(payload), retry=False
         )
         return EvaluationSettings(**data)
 
@@ -281,12 +290,14 @@ class EvaluationsClient:
 
     # ------------------------------------------------------------------
     # Prompt registry endpoints - see agentx.evaluations.prompts.PromptClient for the concept
-    # (the external-agent analog to native autotune). Deliberately read-mostly: no publish here,
-    # a new version only ever comes from the dashboard's human-approved propose/publish flow.
+    # (the external-agent analog to native autotune). propose_prompt never publishes;
+    # publish_prompt_version below IS the explicit approval step - call it only after a human
+    # reviewed the proposal.
     # ------------------------------------------------------------------
 
     def create_prompt(self, payload: dict) -> Prompt:
-        data = self._request("POST", "/prompts", json=self._with_workspace(payload))
+        # Server-side write - no transport retry (see init_run's comment).
+        data = self._request("POST", "/prompts", json=self._with_workspace(payload), retry=False)
         return Prompt(**data)
 
     def list_prompts(self) -> List[Prompt]:
@@ -322,7 +333,9 @@ class EvaluationsClient:
         alias and keeps working. ``split`` records the named case subset this run covers.
         ``additional_scorer_ids`` (self-host): extra judge scorers that each pass their own
         verdict on every result from the same single agent execution - verdicts land in each
-        result row's ``judgeScorerResults`` and the run's ``scorerBreakdown``."""
+        result row's ``judgeScorerResults`` and the run's ``scorerBreakdown``. When
+        ``scorer_group_id`` is set, the engine nulls ``additionalScorerIds`` on the run too -
+        the group is the whole grading story, not a layer on top of extra scorers."""
         from agentx.version import VERSION
 
         grader_id = _resolve_scorer_id(scorer_id, evaluation_settings_id)
@@ -343,12 +356,15 @@ class EvaluationsClient:
         if additional_scorer_ids:
             payload["additionalScorerIds"] = additional_scorer_ids
         # Scorer group grading (self-host): the group's weighted 0-10 aggregate fills the rating
-        # column and member verdicts land per row. Mutually exclusive with scorer_id (group wins).
+        # column and member verdicts land per row. Mutually exclusive with scorer_id (group
+        # wins), and the engine also nulls additionalScorerIds when a group grades the run.
         if scorer_group_id:
             payload["scorerGroupId"] = scorer_group_id
         if split:
             payload["split"] = split
-        data = self._request("POST", "/runs", json=self._with_workspace(payload))
+        # Server-side write: a timeout after the run row was created would be
+        # retried into a duplicate run, so no transport retry.
+        data = self._request("POST", "/runs", json=self._with_workspace(payload), retry=False)
         return EvaluationRun(**data)
 
     def append_results(
@@ -406,7 +422,10 @@ class EvaluationsClient:
         # scorer's own per-result verdicts. Unknown names are a hard 400 from the engine.
         if scorer:
             params["scorer"] = scorer
-        return self._request("GET", f"/runs/{run_id}/gate", params=params)
+        # record=True is a server-side write despite the GET verb (it persists a gate-history
+        # row): a timeout after the row was stored would be retried into a duplicate verdict,
+        # so no transport retry - same precedent as init_run.
+        return self._request("GET", f"/runs/{run_id}/gate", params=params, retry=not record)
 
     def analyze_run(
         self,
@@ -432,8 +451,12 @@ class EvaluationsClient:
 
         if not self._analysis_on_dashboard_router:
             try:
+                # The self-host route runs the analysis SYNCHRONOUSLY (engine
+                # routes/evaluations.ts) - a short timeout with retries re-billed the whole
+                # multi-judge analysis up to 4x while the first was still running. Full
+                # analysis timeout, no transport retry.
                 return self._request(
-                    "POST", f"/runs/{run_id}/analyze", json=payload, timeout=30
+                    "POST", f"/runs/{run_id}/analyze", json=payload, timeout=1800, retry=False
                 )
             except AgentXEvaluationsError as exc:
                 if not self._note_missing_analysis_route(exc, "analyze"):
@@ -554,13 +577,17 @@ class EvaluationsClient:
         if isinstance(dataset_id, dict):  # populated reference, not a bare id
             dataset_id = dataset_id.get("_id") or dataset_id.get("id")
 
-        return Report(
-            runId=run_id,
-            datasetId=dataset_id or "",
-            status=envelope.get("status") or "completed",
-            statistics=envelope.get("statistics"),
+        # Built as one merged dict (explicit keys last, so they win) - passing the
+        # explicit keys as keyword arguments alongside **body raises "got multiple
+        # values" whenever the analysis body itself carries runId/datasetId/status/
+        # statistics.
+        return Report(**{
             **body,
-        )
+            "runId": run_id,
+            "datasetId": dataset_id or "",
+            "status": envelope.get("status") or "completed",
+            "statistics": envelope.get("statistics"),
+        })
 
     # ------------------------------------------------------------------
     # Prompt improvement loop (examples -> propose -> publish). These ride the engine's
@@ -589,8 +616,11 @@ class EvaluationsClient:
             payload["reasoning"] = reasoning
         if based_on_version is not None:
             payload["basedOnVersion"] = based_on_version
+        # Server-side write: a timeout after the version was stored would be
+        # retried into a duplicate version, so no transport retry.
         return self._request(
-            "POST", f"/evaluate/prompts/{prompt_id}/versions", base=self._api_root, json=payload
+            "POST", f"/evaluate/prompts/{prompt_id}/versions", base=self._api_root, json=payload,
+            retry=False,
         )
 
     # ------------------------------------------------------------------
@@ -628,7 +658,7 @@ class EvaluationsClient:
             payload["judgeModel"] = judge_model
         if both_orders:
             payload["bothOrders"] = True
-        response = self._request("POST", "/evaluate/runs/pairwise", json=payload, base=self._api_root)
+        response = self._request("POST", "/evaluate/runs/pairwise", json=payload, base=self._api_root, timeout=900, retry=False,)
         return PairwiseComparison(**response["comparison"])
 
     def get_pairwise(self, batch_id: str) -> PairwiseComparison:
@@ -662,7 +692,8 @@ class EvaluationsClient:
         payload: dict = {"name": name, "definition": definition}
         if description is not None:
             payload["description"] = description
-        return self._request("POST", "/evaluate/tool-schemas", base=self._api_root, json=payload)
+        # Server-side write - no transport retry (see init_run's comment).
+        return self._request("POST", "/evaluate/tool-schemas", base=self._api_root, json=payload, retry=False)
 
     def get_tool_schema_examples(self, tool_schema_id: str, window: Optional[str] = None) -> dict:
         params = {"window": window} if window else None
@@ -686,8 +717,10 @@ class EvaluationsClient:
             payload["reasoning"] = reasoning
         if based_on_version is not None:
             payload["basedOnVersion"] = based_on_version
+        # Server-side write - no transport retry (see init_run's comment).
         return self._request(
-            "POST", f"/evaluate/tool-schemas/{tool_schema_id}/versions", base=self._api_root, json=payload
+            "POST", f"/evaluate/tool-schemas/{tool_schema_id}/versions", base=self._api_root, json=payload,
+            retry=False,
         )
 
     # ------------------------------------------------------------------

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/evaluations/datasets.py

@@ -16,7 +16,13 @@ _REQUIRED_CSV_COLS = {"query"}
 
 
 class DatasetBuilder:
-    """Fluent builder for creating a Custom Agent Evaluations dataset."""
+    """Fluent builder for creating a Custom Agent Evaluations dataset.
+
+    ``judge_prompt``/``judge_model`` are LLM-as-judge overrides for the dataset's grading
+    config. NOTE (self-host): the engine's dataset-create route currently ignores both -
+    set them on a judge scorer / the evaluation settings instead. ``sovereignty_models``
+    is accepted on the wire but not acted on by the self-host engine.
+    """
 
     def __init__(
         self,
@@ -50,6 +56,8 @@ class DatasetBuilder:
         # LLM-as-judge overrides for this dataset's own grading config. Omit either to keep the
         # server default (raw prompt template / gpt-5.6-luna, see EVALUATIONS.md). judge_model
         # must be one of client.evaluations.list_models() (OpenAI or Anthropic).
+        # Self-host: the dataset-create route currently IGNORES judgePrompt/judgeModel - set
+        # them on a judge scorer / the evaluation settings instead (see class docstring).
         if judge_prompt is not None:
             self._payload["judgePrompt"] = judge_prompt
         if judge_model is not None:
@@ -73,7 +81,8 @@ class DatasetBuilder:
             self._payload["codeScorers"] = [
                 {
                     "id": scorer.get("id") or _uuid.uuid4().hex[:12],
-                    "name": scorer["name"],
+                    # Name may be omitted - the engine defaults it, so don't KeyError here.
+                    "name": scorer.get("name"),
                     "code": scorer["code"],
                     "enabled": scorer.get("enabled", True),
                 }
@@ -84,7 +93,8 @@ class DatasetBuilder:
         if rouge_score:
             self._payload["rougeScore"] = {"enabled": True}
         # Sovereignty & Portability - the models to compare on this dataset (use
-        # client.evaluations.list_models() to discover valid ids).
+        # client.evaluations.list_models() to discover valid ids). Self-host: accepted on
+        # the wire but not acted on by the engine (see class docstring).
         if sovereignty_models:
             self._payload["sovereigntyIndex"] = {
                 "enabled": True,
@@ -354,10 +364,13 @@ class DatasetClient:
             "acceptanceCriteria",
             "rejectionCriteria",
             "evaluationCriteria",
+            "judgePrompt",
+            "judgeModel",
             "vectorSimilarity",
             "jaccardSimilarity",
             "bleuScore",
             "rougeScore",
+            "sovereigntyIndex",
             "codeScorers",
         ):
             if wire.get(key) is not None:

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/evaluations/evaluation_settings.py

@@ -64,7 +64,8 @@ class EvaluationSettingsBuilder:
         if rouge_score:
             self._payload["rougeScore"] = {"enabled": True}
         # Sovereignty & Portability - the models to compare when this config runs
-        # (use client.evaluations.list_models() to discover valid ids).
+        # (use client.evaluations.list_models() to discover valid ids). Self-host: accepted
+        # on the wire but not acted on by the engine (same caveat as DatasetBuilder's).
         if sovereignty_models:
             self._payload["sovereigntyIndex"] = {
                 "enabled": True,

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/evaluations/models.py

@@ -248,6 +248,10 @@ class LiveStatistics(BaseModel):
     min_rating: Optional[float] = Field(default=None, alias="minRating")
     max_rating: Optional[float] = Field(default=None, alias="maxRating")
     rated_count: int = Field(default=0, alias="ratedCount")
+    # Rows the judge could not score / rows submitted with an error - the difference between
+    # "everything rated 8" and "half the run never got a verdict".
+    skipped_count: int = Field(default=0, alias="skippedCount")
+    failed_count: int = Field(default=0, alias="failedCount")
 
     class Config:
         populate_by_name = True

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/evaluations/prompts.py

@@ -17,9 +17,9 @@ class PromptClient:
     existing version-comparison view (``client.evaluations`` run comparisons on a dataset) can
     tell you which published version actually scored higher.
 
-    Deliberately read-mostly from here: there is no ``publish`` on this client. A prompt only
-    gets a new version through the dashboard's human-approved propose/publish flow, so a
-    rewritten prompt never reaches your running agent without someone explicitly approving it.
+    ``propose()`` never publishes anything; ``publish_version()`` IS the explicit approval
+    step - call it only after a human reviewed the proposal, since a published version is
+    what your running agent pulls as its live prompt.
 
     Example::
 

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/evaluations/reporting.py

@@ -1,5 +1,7 @@
 from __future__ import annotations
 
+from typing import Optional
+
 from agentx.evaluations.models import Report
 from agentx.evaluations._term import (
     bold,
@@ -21,14 +23,14 @@ _RATING_ICONS = {"high": "●", "medium": "◑", "low": "○"}
 _PRI_COLORS = {"high": red, "medium": yellow, "low": dim}
 
 
-def _rating_badge(rating: str | None) -> str:
+def _rating_badge(rating: Optional[str]) -> str:
     icon = _RATING_ICONS.get(rating or "", "·")
     color = _RATING_COLORS.get(rating or "", dim)
     label = (rating or "").upper()
     return color(f"{icon} {label}") if label else dim(icon)
 
 
-def _section(title: str, rating: str | None = None) -> None:
+def _section(title: str, rating: Optional[str] = None) -> None:
     badge = f"  {_rating_badge(rating)}" if rating else ""
     print(f"\n{bold(title)}{badge}")
     print(dim(_THIN))

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/evaluations/runner.py

@@ -3,6 +3,8 @@ from __future__ import annotations
 import logging
 import os
 import time
+
+import requests
 import uuid
 from typing import Any, Callable, Dict, Iterator, List, Optional, Set, Union
 
@@ -304,6 +306,12 @@ class EvaluationRunContext:
                         resp.failed_validation,
                     )
                     return
+                except requests.Timeout as exc:
+                    # A read timeout means the engine may STILL be scoring this batch - a
+                    # retry re-POSTs it and double-bills every judge call (idempotency keys
+                    # protect rows already inserted, not judge work mid-flight). Fail loud.
+                    last_exc = exc
+                    break
                 except Exception as exc:
                     last_exc = exc
                     if attempt == 1:
@@ -334,6 +342,10 @@ class EvaluationRunContext:
     # ------------------------------------------------------------------
 
     def finalize(self) -> "EvaluationRunContext":
+        """Mark the run completed server-side. A failed finalize is raised, not swallowed
+        (same fail-loud posture as _flush_batch): it leaves the run in_progress - a state a
+        CI pipeline MUST treat as a failure, since gates and baselines only consider
+        completed runs."""
         _say()
         with Spinner("Finalizing - submitting results"):
             try:
@@ -345,6 +357,7 @@ class EvaluationRunContext:
             except Exception as exc:
                 _say(f"  {red('✗')}  Finalize failed: {dim(str(exc))}")
                 logger.error("Finalize failed: %s", exc)
+                raise
         return self
 
     def gate(

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/integrations/anthropic.py

@@ -274,7 +274,7 @@ def _patch_stream(
                 return iter(ctx)
 
             def __aiter__(self_inner):
-                return aiter(ctx)
+                return ctx.__aiter__()  # aiter() builtin is 3.10+; python_requires is >=3.9
 
             def __getattr__(self_inner, item):
                 return getattr(ctx, item)

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/integrations/autogen.py

@@ -167,6 +167,9 @@ class AgentXAutoGenObserver:
                         "end_time": end_t,
                         "input": pending["input"] if pending else None,
                         "output": f"ERROR: {output}" if is_error else (str(output) if output is not None else None),
+                        # The engine's failure test is success === false; without
+                        # this a failed tool call would read as passing.
+                        "success": not is_error,
                     })
                 continue
 

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/integrations/crewai.py

@@ -127,6 +127,16 @@ class AgentXCrewObserver:
         except ImportError:
             return task_timings, lambda: None
 
+        # Double-instrumentation guard (bus-keyed latch, the same idea as the
+        # other integrations' _agentx_patched flag): the event bus is a global
+        # singleton, so a notebook re-run or an overlapping kickoff that
+        # already has AgentX listeners registered would otherwise get a second
+        # set and duplicate every task span. When already attached, this
+        # kickoff just falls back to the evenly-divided timing approximation.
+        if getattr(crewai_event_bus, "_agentx_attached", False):
+            return task_timings, lambda: None
+        crewai_event_bus._agentx_attached = True
+
         def on_task_started(source: Any, event: Any) -> None:
             task_id = getattr(event, "task_id", None)
             if task_id is None:
@@ -154,9 +164,14 @@ class AgentXCrewObserver:
         crewai_event_bus.on(TaskFailedEvent)(on_task_failed)
 
         def unregister() -> None:
-            crewai_event_bus.off(TaskStartedEvent, on_task_started)
-            crewai_event_bus.off(TaskCompletedEvent, on_task_completed)
-            crewai_event_bus.off(TaskFailedEvent, on_task_failed)
+            try:
+                crewai_event_bus.off(TaskStartedEvent, on_task_started)
+                crewai_event_bus.off(TaskCompletedEvent, on_task_completed)
+                crewai_event_bus.off(TaskFailedEvent, on_task_failed)
+            finally:
+                # Clear the latch even if .off() raises, so a later kickoff
+                # can re-attach instead of being locked out forever.
+                crewai_event_bus._agentx_attached = False
 
         return task_timings, unregister
 
@@ -197,6 +212,9 @@ class AgentXCrewObserver:
                 "end_time": end,
                 "input": description,
                 "output": output_text,
+                # A CrewAI task is an agent turn, not a model call - without this,
+                # _merge_child_run's default stamped every task span "llm".
+                "kind": "agent",
             })
             if description is not None or task_output_text is not None:
                 tool_calls.append({"name": name, "input": description, "output": task_output_text})

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/integrations/google_adk.py

@@ -231,6 +231,7 @@ class AgentXADKPlugin(BasePlugin):
                 model=call_start.get("model") if call_start else None,
                 input_tokens=call_input_tokens,
                 output_tokens=call_output_tokens,
+                span_kind="llm",
             )
 
     async def on_model_error_callback(
@@ -261,6 +262,7 @@ class AgentXADKPlugin(BasePlugin):
                 output=f"ERROR: {error}",
                 model=call_start.get("model") if call_start else None,
                 error=str(error),
+                span_kind="llm",
             )
 
     # ------------------------------------------------------------------
@@ -290,7 +292,8 @@ class AgentXADKPlugin(BasePlugin):
         tool_input = _safe_serialize(tool_args)
         tool_output = str(result) if result is not None else None
         state["root_span"].child_span(
-            tool_name, start_time=start_t, end_time=end_t, input=tool_input, output=tool_output
+            tool_name, start_time=start_t, end_time=end_t, input=tool_input, output=tool_output,
+            span_kind="tool",
         )
 
     async def on_tool_error_callback(
@@ -311,5 +314,6 @@ class AgentXADKPlugin(BasePlugin):
         tool_input = _safe_serialize(tool_args)
         tool_output = f"ERROR: {error}"
         state["root_span"].child_span(
-            tool_name, start_time=start_t, end_time=end_t, input=tool_input, output=tool_output, error=str(error)
+            tool_name, start_time=start_t, end_time=end_t, input=tool_input, output=tool_output, error=str(error),
+            span_kind="tool",
         )

{agentx_python-0.8.19 → agentx_python-0.8.21}/agentx/integrations/langchain.py

@@ -277,6 +277,10 @@ class AgentXCallbackHandler(BaseCallbackHandler):
     become child spans, and each LLM call / tool call / retrieval becomes a span parented under
     the node that ran it - so the engine's Execution Timeline shows the actual graph trajectory
     (which nodes ran, in what order, and what each did), not a flat step list.
+
+    Retriever runs are stamped ``retrieval`` - LangChain cannot distinguish memory-backed
+    retrievers, so a Mem0/Zep-style store exposed as a retriever classifies as retrieval too.
+    Use ``tracer.trace_memory`` for lookups that should classify as memory.
     """
 
     def __init__(
@@ -339,6 +343,16 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             self._retrieval_starts.pop(run_id, None)
             self._parents.pop(run_id, None)
 
+        # Pre-run retrieval steps waiting for a top-level chain that never came
+        # (e.g. retriever.invoke() called but agent.invoke() aborted before
+        # on_chain_start). Each step carries its own start_time, so drop the
+        # pre-cutoff ones just like the run_id-keyed structures above.
+        with self._state_lock:
+            if self._pending_retrieval_steps:
+                self._pending_retrieval_steps[:] = [
+                    step for step in self._pending_retrieval_steps if step.get("start_time", 0) >= cutoff
+                ]
+
     # ------------------------------------------------------------------
     # Chain lifecycle
     # ------------------------------------------------------------------
@@ -359,8 +373,9 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             self._prune_stale_entries()
             # Consume any retrieval steps that ran before this chain started
             # (pre-run RAG: retriever.invoke() called before agent.invoke())
-            pending = self._pending_retrieval_steps[:]
-            self._pending_retrieval_steps.clear()
+            with self._state_lock:
+                pending = self._pending_retrieval_steps[:]
+                self._pending_retrieval_steps.clear()
             self._runs[run_id] = {
                 "start": time.time(),
                 "input": _extract_input(inputs),
@@ -670,7 +685,7 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             "input": _extract_llm_input(prompts=prompts, messages=messages),
         }
         top = self._find_top_ancestor(parent_run_id)
-        if top and not self._runs[top].get("model") and model:
+        if top and top in self._runs and not self._runs[top].get("model") and model:
             self._runs[top]["model"] = model
 
     def on_llm_start(
@@ -882,15 +897,16 @@ class AgentXCallbackHandler(BaseCallbackHandler):
                 step["output"] = "\n\n---\n\n".join(contents)
 
         top = self._find_top_ancestor(parent_run_id)
-        if top and top in self._runs:
-            # Retriever ran inside an active chain - attach directly
-            retrievals = self._runs[top]["retrieval_steps"]
-            step["name"] = f"Retrieval {len(retrievals) + 1}"
-            retrievals.append(step)
-        else:
-            # Retriever ran before the chain started (pre-run RAG pattern)
-            step["name"] = f"Retrieval {len(self._pending_retrieval_steps) + 1}"
-            self._pending_retrieval_steps.append(step)
+        with self._state_lock:
+            if top and top in self._runs:
+                # Retriever ran inside an active chain - attach directly
+                retrievals = self._runs[top]["retrieval_steps"]
+                step["name"] = f"Retrieval {len(retrievals) + 1}"
+                retrievals.append(step)
+            else:
+                # Retriever ran before the chain started (pre-run RAG pattern)
+                step["name"] = f"Retrieval {len(self._pending_retrieval_steps) + 1}"
+                self._pending_retrieval_steps.append(step)
 
     def on_retriever_error(
         self,
@@ -918,15 +934,16 @@ class AgentXCallbackHandler(BaseCallbackHandler):
             step["query"] = query
 
         top = self._find_top_ancestor(parent_run_id)
-        if top and top in self._runs:
-            # Retriever ran inside an active chain - attach directly
-            retrievals = self._runs[top]["retrieval_steps"]
-            step["name"] = f"Retrieval {len(retrievals) + 1}"
-            retrievals.append(step)
-        else:
-            # Retriever ran before the chain started (pre-run RAG pattern)
-            step["name"] = f"Retrieval {len(self._pending_retrieval_steps) + 1}"
-            self._pending_retrieval_steps.append(step)
+        with self._state_lock:
+            if top and top in self._runs:
+                # Retriever ran inside an active chain - attach directly
+                retrievals = self._runs[top]["retrieval_steps"]
+                step["name"] = f"Retrieval {len(retrievals) + 1}"
+                retrievals.append(step)
+            else:
+                # Retriever ran before the chain started (pre-run RAG pattern)
+                step["name"] = f"Retrieval {len(self._pending_retrieval_steps) + 1}"
+                self._pending_retrieval_steps.append(step)
 
     # ------------------------------------------------------------------
     # Helpers