fix tool calls data loss

agentx/__init__.py

@@ -5,6 +5,8 @@
 from agentx.exceptions import (
     AgentXError,
     AgentXAuthError,
+    AgentXValidationError,
+    AgentXConnectionError,
     AgentXAPIError,
     DatasetNotFound,
     CINotEnabled,
@@ -21,6 +23,8 @@
     "AgentX",
     "AgentXError",
     "AgentXAuthError",
+    "AgentXValidationError",
+    "AgentXConnectionError",
     "AgentXAPIError",
     "DatasetNotFound",
     "CINotEnabled",

agentx/agentx.py

@@ -3,7 +3,7 @@
 import os
 import logging
 
-from agentx.util import get_headers, api_base
+from agentx.util import get_headers, api_base, normalize_base
 from agentx.resources.agent import Agent
 from agentx.resources.workforce import Workforce
 
@@ -16,16 +16,23 @@ def __init__(
         base_url: Optional[str] = None,
         workspace_id: Optional[str] = None,
     ):
+        # The api_key is NOT written back into os.environ (it used to be): every sub-client
+        # below receives it explicitly, and mutating process-global state from a constructor
+        # re-pointed unrelated code - the same leak the base_url write below had (deep-dive
+        # round 3, bug #1). Static flows that still read the env (AgentX.list_workforces,
+        # bare get_headers()) now require the caller to set AGENTX_API_KEY themselves.
         self.api_key = api_key or os.getenv("AGENTX_API_KEY")
-        if self.api_key and not os.getenv("AGENTX_API_KEY"):
-            os.environ["AGENTX_API_KEY"] = self.api_key
 
         # base_url overrides AGENTX_API_BASE_URL env var (and the SDK default). It is
         # deliberately NOT written back into os.environ: the constructor used to do that, which
         # made the last-constructed client silently re-point every other client in the process
         # (deep-dive round 3, bug #1). Each sub-client below receives this value explicitly and
-        # captures it at construction instead.
+        # captures it at construction instead. Normalized (trailing slash and the
+        # /custom-agent-evaluations suffix stripped) so an evaluations-shaped URL works for
+        # every sub-client, not just evaluations.
         self.base_url = base_url or os.getenv("AGENTX_API_BASE_URL")
+        if self.base_url:
+            self.base_url = normalize_base(self.base_url)
 
         self.workspace_id = workspace_id or os.getenv("AGENTX_WORKSPACE_ID")
 
@@ -91,6 +98,27 @@ def __init__(
             workspace_id=self.workspace_id,
         )
         self.tracer = Tracer(_ingest_client)
+        self._ingest_client = _ingest_client
+
+    # ------------------------------------------------------------------
+    # Lifecycle
+    # ------------------------------------------------------------------
+
+    def close(self, timeout: float = 5.0) -> bool:
+        """Flush queued traces and stop the tracer's background ingest worker.
+
+        Returns ``True`` when everything drained before ``timeout`` seconds elapsed. Optional -
+        an ``atexit`` hook already flushes queued traces on interpreter shutdown - but a
+        long-running service that tears clients down mid-process should call it (or use the
+        client as a context manager) so worker threads don't accumulate.
+        """
+        return self._ingest_client.close(timeout)
+
+    def __enter__(self) -> "AgentX":
+        return self
+
+    def __exit__(self, exc_type, exc_val, tb) -> None:
+        self.close()
 
     @classmethod
     def from_env(cls) -> "AgentX":
@@ -129,7 +157,9 @@ def list_agents(self) -> List[Agent]:
 
     @staticmethod
     def list_workforces() -> List["Workforce"]:
-        """List all workforces/teams."""
+        """List all workforces/teams. Static, so it reads AGENTX_API_KEY from the environment
+        directly - the constructor no longer writes ``api_key`` into os.environ, so set the
+        env var yourself before calling this."""
         url = f"{api_base()}/access/teams"
         response = requests.get(url, headers=get_headers())
         if response.status_code == 200:

agentx/evaluations/client.py

@@ -24,7 +24,12 @@
 
 logger = logging.getLogger(__name__)
 
-from agentx.util import _DEFAULT_API_BASE as _UTIL_API_BASE
+from agentx.util import _DEFAULT_API_BASE as _UTIL_API_BASE, normalize_base
+
+# The canonical error classes (agentx.exceptions) are raised - and re-exported here for
+# compat with code that imported them from this module - so `except agentx.AgentXAuthError`
+# works whichever client raised.
+from agentx.exceptions import AgentXError, AgentXAuthError, AgentXValidationError
 
 _DEFAULT_BASE_URL = f"{_UTIL_API_BASE}/custom-agent-evaluations"
 SDK_NAME = "agentx-python"
@@ -42,7 +47,7 @@
 _SELF_HOST_SCORING_TIMEOUT = 900
 
 
-class AgentXEvaluationsError(Exception):
+class AgentXEvaluationsError(AgentXError):
     """An evaluations API call failed.
 
     ``status_code`` carries the HTTP status when the failure came from a response rather
@@ -55,14 +60,6 @@ def __init__(self, message: str, status_code: Optional[int] = None) -> None:
         self.status_code = status_code
 
 
-class AgentXAuthError(AgentXEvaluationsError):
-    pass
-
-
-class AgentXValidationError(AgentXEvaluationsError):
-    pass
-
-
 class EvaluationSubmissionError(AgentXEvaluationsError):
     """A result batch could not be submitted (after one retry). The run is left unfinalized;
     re-running execute() on the same context resumes past already-submitted cases."""
@@ -96,13 +93,10 @@ def __init__(
         # whatever workspace the API key's user defaults to, not the one the caller intended.
         self._workspace_id = workspace_id
         # Priority: constructor arg > env var > SDK default
-        # Always append /custom-agent-evaluations so users only need to provide /api/v1
-        _api_base = (
-            base_url or os.getenv("AGENTX_API_BASE_URL", _UTIL_API_BASE)
-        ).rstrip("/")
-        if not _api_base.endswith("/custom-agent-evaluations"):
-            _api_base = f"{_api_base}/custom-agent-evaluations"
-        self._base_url = _api_base
+        # normalize_base strips a trailing slash and any /custom-agent-evaluations suffix,
+        # then the suffix is appended - users only need to provide /api/v1 either way.
+        _api_base = normalize_base(base_url or os.getenv("AGENTX_API_BASE_URL", _UTIL_API_BASE))
+        self._base_url = f"{_api_base}/custom-agent-evaluations"
         # None until an analysis call tells us which engine this is; see _api_root.
         self._analysis_on_dashboard_router: Optional[bool] = None
         self._session = requests.Session()
@@ -184,13 +178,16 @@ def _request(
                 continue
 
             if resp.status_code == 401:
-                raise AgentXAuthError("Invalid or missing API key")
+                raise AgentXAuthError("Invalid or missing API key", status_code=401)
             if resp.status_code == 422:
                 raise AgentXValidationError(resp.text)
+            # Gate on the schedule itself so HTTP-status retries walk the SAME full backoff
+            # schedule connection errors do - the old `attempt < _MAX_RETRIES - 1` gate left
+            # the schedule's last entry unreachable for HTTP retries (ingest_client precedent).
             if (
                 resp.status_code in _RETRYABLE_STATUS
                 and retry
-                and attempt < _MAX_RETRIES - 1
+                and attempt < len(schedule) - 1
             ):
                 logger.debug(
                     "Retryable status %d (attempt %d)", resp.status_code, attempt + 1
@@ -500,6 +497,17 @@ def get_report(self, run_id: str) -> Report:
         return self._report_from_dashboard(run_id)
 
     def get_missing_results(self, run_id: str) -> List[Dict[str, Any]]:
+        """Deprecated: on self-host the route's response body has no top-level list, so this
+        always returns ``[]``. Use :meth:`get_submitted_keys` - the same route's
+        ``submittedKeys`` - to find out what a run still needs."""
+        import warnings
+
+        warnings.warn(
+            "get_missing_results() always returns [] on self-host - use get_submitted_keys() "
+            "to resume a run instead.",
+            DeprecationWarning,
+            stacklevel=2,
+        )
         data = self._request("GET", f"/runs/{run_id}/missing-results")
         return data if isinstance(data, list) else data.get("missing", [])
 

agentx/evaluations/models.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from typing import Any, Dict, List, Literal, Optional, Union
-from pydantic import AliasChoices, BaseModel, Field, model_validator
+from pydantic import AliasChoices, BaseModel, Field, field_validator, model_validator
 
 # ---------------------------------------------------------------------------
 # Observable trace
@@ -78,6 +78,17 @@ class Dataset(BaseModel):
     # Custom code scorers attached to this dataset - [{ id, name, code, enabled }]. Retrievable,
     # so a fetched dataset round-trips them (import_dataset copies them to the new dataset).
     code_scorers: Optional[List[Dict[str, Any]]] = Field(default=None, alias="codeScorers")
+    # Grading config carried on the dataset itself - similarity metric toggles (each a
+    # {"enabled": bool, ...} object on the wire), LLM-as-judge overrides, and the raw
+    # sovereigntyIndex object. Modeled so a fetched Dataset round-trips them: extra="ignore"
+    # used to silently drop all of these on read, and import_dataset lost them on the copy.
+    vector_similarity: Optional[Any] = Field(default=None, alias="vectorSimilarity")
+    jaccard_similarity: Optional[Any] = Field(default=None, alias="jaccardSimilarity")
+    bleu_score: Optional[Any] = Field(default=None, alias="bleuScore")
+    rouge_score: Optional[Any] = Field(default=None, alias="rougeScore")
+    judge_prompt: Optional[str] = Field(default=None, alias="judgePrompt")
+    judge_model: Optional[str] = Field(default=None, alias="judgeModel")
+    sovereignty_index: Optional[Dict[str, Any]] = Field(default=None, alias="sovereigntyIndex")
     status: str = "published"
     version_id: Optional[str] = Field(default=None, alias="versionId")
     # Sovereignty & Portability - models selected to compare on this dataset.
@@ -430,7 +441,12 @@ class RunResultRow(BaseModel):
     question_index: Optional[int] = Field(default=None, alias="questionIndex")
     run_number: Optional[int] = Field(default=None, alias="runNumber")
     question_text: Optional[str] = Field(default=None, alias="questionText")
-    response: Optional[str] = None
+    # The engine sends the agent's answer as an `output` OBJECT ({"text": ...}), not a
+    # `response` string - accept both spellings and lift the dict's text (see the
+    # validator below), so row.response actually populates on self-host.
+    response: Optional[str] = Field(
+        default=None, validation_alias=AliasChoices("response", "output")
+    )
     trace_id: Optional[str] = Field(default=None, alias="traceId")
     latency_ms: Optional[float] = Field(default=None, alias="latencyMs")
     input_tokens: Optional[int] = Field(default=None, alias="inputTokens")
@@ -452,6 +468,15 @@ class Config:
         populate_by_name = True
         extra = "ignore"
 
+    @field_validator("response", mode="before")
+    @classmethod
+    def _lift_output_text(cls, value: Any) -> Any:
+        # The `output` alias delivers the wire's whole output object - keep the declared
+        # Optional[str] by lifting its text field.
+        if isinstance(value, dict):
+            return value.get("text")
+        return value
+
     @classmethod
     def from_wire(cls, wire: Dict[str, Any]) -> "RunResultRow":
         row = cls.model_validate(wire)

agentx/evaluations/runner.py

@@ -11,7 +11,11 @@
 from agentx.evaluations.adapters.raw import RawCallableAdapter
 from agentx.evaluations.adapters.precomputed import PrecomputedAdapter
 from agentx.evaluations.adapters.http_endpoint import HttpEndpointAdapter
-from agentx.evaluations.client import EvaluationsClient, EvaluationSubmissionError
+from agentx.evaluations.client import (
+    AgentXEvaluationsError,
+    EvaluationsClient,
+    EvaluationSubmissionError,
+)
 from agentx.evaluations.models import (
     AnalysisStatus,
     Dataset,
@@ -227,6 +231,7 @@ def produce(case: EvaluationCase) -> EvaluationResult:
         if concurrency > 1:
             import concurrent.futures
             import contextvars
+            from collections import deque
 
             def in_scope(case: EvaluationCase) -> EvaluationResult:
                 # ContextVars (the eval-run scope) do not cross thread boundaries on their own -
@@ -240,45 +245,68 @@ def in_scope(case: EvaluationCase) -> EvaluationResult:
                 if _idem_key(self._run.run_id, case.case_id, case.run_number) not in already_done
             ]
             executor = concurrent.futures.ThreadPoolExecutor(max_workers=concurrency)
-            # map() yields in submission order, so batching/submission below stays deterministic.
-            mapped = executor.map(in_scope, pending)
 
-            def ordered() -> "Iterator[EvaluationResult]":
+            def bounded() -> "Iterator[EvaluationResult]":
+                # Bounded submit loop instead of executor.map(): map() dispatches EVERY case
+                # up front, so a fail-fast flush failure (EvaluationSubmissionError below)
+                # still paid for the whole rest of the run in agent calls. Keep at most
+                # `concurrency` cases in flight, topping up as results are consumed; yields
+                # stay in submission order so batching below is deterministic. On teardown
+                # (an exception in the consuming loop closes this generator) whatever is
+                # queued but unstarted is cancelled.
+                import itertools
+
+                case_iter = iter(pending)
+                in_flight: "deque[concurrent.futures.Future]" = deque()
                 try:
-                    yield from mapped
+                    for case in itertools.islice(case_iter, concurrency):
+                        in_flight.append(executor.submit(in_scope, case))
+                    while in_flight:
+                        result = in_flight.popleft().result()
+                        next_case = next(case_iter, None)
+                        if next_case is not None:
+                            in_flight.append(executor.submit(in_scope, next_case))
+                        yield result
                 finally:
-                    executor.shutdown(wait=True)
+                    executor.shutdown(wait=False, cancel_futures=True)
 
-            results_iter = ordered()
+            results_iter = bounded()
         else:
             results_iter = None  # sequential path below produces inline
 
-        for idx, case in enumerate(cases, start=1):
-            idem_key = _idem_key(self._run.run_id, case.case_id, case.run_number)
-
-            if idem_key in already_done:
-                logger.debug("Skipping already-submitted case: %s", idem_key)
-                _print_progress(idx, total, case, skipped=True)
-                continue
-
-            result = next(results_iter) if results_iter is not None else produce(case)
-            result.idempotency_key = idem_key
-            # Tag the result with the case's model so the server can group it into
-            # the Sovereignty & Portability matrix (the callable may also set it).
-            if case.model:
-                meta = dict(result.metadata or {})
-                meta.setdefault("model", case.model)
-                result.metadata = meta
-            result = EvaluationResult(
-                **{**result.model_dump(), "idempotencyKey": idem_key}
-            )
-            self._results.append(result)
-            batch.append(result)
-            _print_progress(idx, total, case, result=result)
+        try:
+            for idx, case in enumerate(cases, start=1):
+                idem_key = _idem_key(self._run.run_id, case.case_id, case.run_number)
+
+                if idem_key in already_done:
+                    logger.debug("Skipping already-submitted case: %s", idem_key)
+                    _print_progress(idx, total, case, skipped=True)
+                    continue
+
+                result = next(results_iter) if results_iter is not None else produce(case)
+                result.idempotency_key = idem_key
+                # Tag the result with the case's model so the server can group it into
+                # the Sovereignty & Portability matrix (the callable may also set it).
+                if case.model:
+                    meta = dict(result.metadata or {})
+                    meta.setdefault("model", case.model)
+                    result.metadata = meta
+                result = EvaluationResult(
+                    **{**result.model_dump(), "idempotencyKey": idem_key}
+                )
+                self._results.append(result)
+                batch.append(result)
+                _print_progress(idx, total, case, result=result)
 
-            if len(batch) >= max_batch:
-                self._flush_batch(batch)
-                batch = []
+                if len(batch) >= max_batch:
+                    self._flush_batch(batch)
+                    batch = []
+        finally:
+            # Deterministic teardown: a flush failure mid-run must stop the in-flight agent
+            # dispatch NOW (bounded()'s finally cancels queued cases), not whenever the
+            # generator happens to be garbage-collected.
+            if results_iter is not None:
+                results_iter.close()
 
         if batch:
             self._flush_batch(batch)
@@ -333,9 +361,21 @@ def _fetch_submitted_keys(self) -> Set[str]:
         so a re-execute() after a crash skips (and never re-pays for) finished cases."""
         try:
             return set(self._client.get_submitted_keys(self._run.run_id))
-        except Exception:
-            # Older engines without the route: no resume, identical to the historical behavior.
-            return set()
+        except AgentXEvaluationsError as exc:
+            if exc.status_code == 404:
+                # Older engines without the route: no resume, identical to the historical
+                # behavior. ONLY the 404 qualifies - a transient 502/timeout here used to be
+                # swallowed too, and an empty resume set silently re-runs (and re-bills)
+                # every already-finished case.
+                return set()
+            _say(f"  {red('✗')}  Could not fetch already-submitted keys: {dim(str(exc))}")
+            logger.error(
+                "Resume-key fetch for run %s failed (%s) - refusing to re-run the whole run "
+                "blind; retry execute() once the engine is reachable",
+                self._run.run_id,
+                exc,
+            )
+            raise
 
     # ------------------------------------------------------------------
     # Step 2: finalize
@@ -465,6 +505,8 @@ def analyze(
 
         Args:
             mode: "auto" (default), "sync", or "batch" - how item scoring executes server-side.
+                Hosted-only: self-host runs the analysis synchronously regardless; see the
+                response's mode field for what actually ran.
             quality_mode: "quality_first" or "balanced" - how many items get a second/third judge.
             judges: 1-3 model ids from ``client.evaluations.list_models()``, e.g.
                 ``["gpt-5.6-luna", "claude-opus-4-8"]``. Omit to let the engine score with its