perf: Reuse the connection to Lago, and fix two silent-failure paths

CHANGELOG.md

@@ -6,6 +6,19 @@ All notable changes to this project will be documented here. Format follows [Kee
 
 ### Fixed
 
+- **Every batch opened a new connection to Lago, and the isolation path paid that cost per event.** `LagoClient.send_batch` called the module-level `requests.post`, which builds a `Session`, uses it and closes it — so the pool was discarded after every call and the client declined the keep-alive HTTP/1.1 offers by default. Counted against a local server answering `Connection: keep-alive`: **8 batches produced 8 TCP connections; one `Session` produces 1.** The cost of each discarded connection is a TCP handshake plus a TLS handshake, about **2 RTT** — measured from a client 137ms from `api.getlago.com`, ~276ms of setup in front of a request the server answers in ~147ms (440.6ms → 147.6ms median per batch). Stated as 2 RTT rather than as a multiplier deliberately: the ratio compresses toward 1x for a client sitting next to the API and grows for one far from it, so the tests assert on the connection count, which does not move.
+  - **The isolation path is where this actually hurt.** The normal path amortises one handshake over 100 events (~2.7ms/event). When one duplicate `transaction_id` 422s a batch, `_send_individually` re-sends all 100 events alone and paid the handshake **per event** (~274ms/event, two orders of magnitude worse) — on the recovery path, when the system is already degraded. Driving the real queue through a 422 against a counting server: **101 requests / 101 connections before, 101 requests / 1 connection after**, removing ~27s of pure handshake from a single bad batch.
+  - Discarding the `Session` also discarded the TLS session ticket, so no abbreviated handshake was possible on reconnect either.
+  - The one consequence pointing the other way: the isolation fan-out now lands on Lago ~3x faster, and there is no application-level rate limit in front of it. That strengthens the case for bounding the fan-out rather than weakening it.
+  - Safe as one shared `Session` because sends happen on exactly one thread, the queue's `lago-queue` worker. A second sender thread would need a thread-local `Session` or a pool sized past `requests`' default of 10.
+
+- **`flush()` returned `True` while a batch was still in flight.** `_take_batch` pops events OUT of the buffer before the send is attempted, and `flush()` decided it was done by checking that the buffer was empty — so for the duration of a POST the batch was in neither the buffer nor Lago. Measured against the real queue: **`flush()` returned `True` on five events that were in flight and then failed.** At process exit the exit drain is a second net and usually still delivers, so this mostly cost nothing there; the damage is to the contract, because `flush()` is documented as the delivery checkpoint and a caller using it as one — end of a batch-job phase, before logging "billing synced", before a container reports ready-to-terminate — got a false positive with nothing behind it. `_in_flight` now counts events taken but not yet accounted for, and every path that consumes a batch settles it: `_replay_failed` settles what it puts back in the same lock acquisition that re-queues it, `_send_individually` settles what it resolved rather than re-queued, and both drain loops settle on delivery or on final loss.
+  - Worth stating for anyone benchmarking the SDK: a `True` from `flush()` was never proof of delivery and should not be reported as one. Showing that a load test lost nothing means counting the events Lago accepted.
+
+- **An adapter that threw never reached `on_error`, only the log.** `LagoSDK.emit()` reports through `_report_error(exc, "emit")`, but the adapter call sits OUTSIDE it on the wrapper's side of the boundary — so when a provider changes its response shape the adapter raises before `emit()` is entered and the SDK's own reporting never runs. Across the five wrappers: **14 `except Exception` sites, 14 `logger.warning` calls, 0 `on_error` routes.** That is the worst direction for this failure to take, because provider drift does not hit an occasional call, it hits *every* call for that provider — turning a total billing outage into a log line nobody is tailing. `on_error` is the channel customers watch, and the one the queue already uses to report an overflow for exactly this reason. All 14 sites now report as well as log; `_report_error` already swallows anything the customer's callback raises, so the customer's LLM call still returns normally either way.
+
+- **A comment in `queue.py` contradicted the code it documented.** `_after_in_child` has always given the child a fresh, empty buffer, but the comment above it claimed the parent's events were "copied over".
+
 - **An explicitly-passed falsy `api_url` silently resolved to PRODUCTION Lago.** Preferring the config value over `""` is right — `requests` raises `MissingSchema`, which is not a `LagoApiError`, so the queue classified it transient, re-prepended the batch and retried at the 60s ceiling forever, stopping all billing with only a growing buffer as the symptom. But `LagoConfig`'s default is the production URL, so `api_url=os.environ.get("LAGO_API_URL", "")` with the var unset resolved to production with **no `on_error` and no log** — verified live: 0 reports, 0 log lines, and a client posting to `api.getlago.com`. For a CI job or a developer holding a real production key that writes live billing data, and ingested events cannot be un-ingested. The fallback is unchanged, so the original config-clobber bug stays fixed; it is now reported under `config.api_url` through the same log-plus-callback floor as every other drop path. An *unpassed* `api_url` stays silent — `None` means the caller never mentioned it, and reporting the common case would train customers to ignore the channel this fix depends on.
 - **`usage_metadata` from the Cloudflare gateway got no drift sweep, and had already lost two counters.** `extras` was a fixed three-key dict, so any counter the adapter does not map vanished with no error and no `on_error` — the one place violating the drift contract `test_drift.py` enforces for the native adapters. Not hypothetical: replaying the **14 captured fixtures** through the adapter drops **`neurons`** (Cloudflare's Workers AI billing unit) in 4 entries and **`input_text_tokens`** in 1, and a live Logs API pull also returns **`units`**, a cost quantity that appears in no fixture at all — the hand-maintained key enumeration in the module docstring had already drifted past reality, which is exactly the failure mode a snapshot invites. Unmapped keys are now swept into `extras["usage_metadata"]` against an explicit `_MAPPED_USAGE_KEYS` set, so a ninth spelling surfaces on its own instead of needing another 14-fixture audit. Deliberately **nested** rather than merged flat into `extras`: the poller reads `extras["cached"]` to decide whether to skip billing a request Cloudflare served for free, so a future `usage_metadata` key called `cached` or `step` must not be able to shadow it. The regression test iterates the fixture directory rather than a fixed key list, so a recapture that introduces a new counter fails it with no test edit. Closes #16.
 - **`prime()` re-downloaded the ~400-model OpenRouter catalogue regardless of the TTL.** It set `_openrouter_stale` unconditionally, and it is reached from `_auto_prime_pricing_for` on a matching `wrap()` and from `warm_pricing()` — both of which a server can run per request — so `pricing_ttl_seconds` never applied on that path at all, and the catalogue was refetched on essentially every flush tick, on the thread the queue drains events from. Measured with the shipped 1-hour TTL: **4 `prime()`+`maybe_refresh()` cycles produced 4 full downloads where 1 was correct; now 1.** Gated on the same "no table, or past the TTL" test `lookup()` already uses, so priming and looking up cannot disagree about what needs fetching — and a table that genuinely ages out is still re-primed, so prices do not freeze at the first fetch.
@@ -69,6 +82,10 @@ All notable changes to this project will be documented here. Format follows [Kee
 
 ### Changed
 
+- **Unit tests can no longer reach the network.** A test that stubs HTTP by patching a library function keeps passing as a shape long after it has stopped testing anything: when the code under test moves to a different call path, the patch stops matching and the request goes out for real. Not hypothetical — the two `verify_ssl` tests patched `requests.post`, and once the client held its own `Session` they began issuing live requests to `api.getlago.com`. They failed, but with Lago's `404` rather than with anything pointing at the mismatch. `pytest-socket`'s `--disable-socket` is now on by default: on those same tests it turns a confusing `LagoApiError: 404` after 0.97s into `SocketBlockedError` after 0.15s, removes the class of CI flake where a green suite depends on the network, and makes it impossible for a test run to reach production Lago.
+  - `127.0.0.1` stays allowed so a test can stand up a real server when the property under test is about real socket behaviour. `test_connection_reuse.py` does that — it counts accepts on a local HTTP/1.1 server rather than asserting that some particular function was called, and fails on the old code with "8 connections for 8 batches".
+  - `responses` is added alongside for stubbing at the `requests` transport boundary rather than by patching internals. Existing tests are left as they are; it is there for new tests and for migrating the brittle patch sites as they are touched. **No cassette/auto-record layer** — fixtures in this repo stay explicit, captured by a developer running `tests/unit/adapters/fixtures/capture_*.py` on purpose and reviewed in the diff.
+
 - **The Databricks usage read no longer selects 36 columns to bill off 14.** `read_usage` ran `SELECT * FROM system.ai_gateway.usage`, and the Statement Execution API's default `disposition=INLINE` **fails** a statement whose response exceeds 25 MiB rather than paginating past it — so the width of the projection is what sets the largest window this reader can handle, for a module whose own guidance is "read one wide window per run". Measured on the live table over 247 real rows: **1,411 bytes/row for `SELECT *` against 435** for the 14 columns the extraction actually reads, i.e. a ceiling of ~18k rows where it should be ~60k. The dropped columns are the wide ones nothing bills off (`routing_information`, `endpoint_metadata`, `url`, `user_agent`). Billing output is byte-identical.
   - The projection is one named list next to the interval regex, and the coupling runs the other way too: a column missing from it reaches the adapter as *absent*, which every field degrades to zero/empty on rather than raising — a silently under-billed event. So the test that pins it feeds the canned rows **through** the statement's own column list and asserts the events match an unprojected read, which fails if the projection is trimmed too far.
   - `ORDER BY event_time` dropped with it: nothing downstream reads the row order — the BYOK join is keyed, the unbilled-bucket report is sorted, and each event's `transaction_id` derives from the row's own ids — so it only bought the warehouse a sort over the widest read this module makes.

pyproject.toml

@@ -41,6 +41,8 @@ dev = [
     "pytest-asyncio>=0.23",
     "pytest-cov>=5",
     "hypothesis>=6",
+    "pytest-socket>=0.7",
+    "responses>=0.25",
     # tooling
     "ruff>=0.6",
     "mypy>=1.10",
@@ -65,6 +67,8 @@ where = ["src"]
 [tool.pytest.ini_options]
 testpaths = ["tests"]
 pythonpath = ["src"]
+# Localhost stays open so a test can stand up a real server (test_connection_reuse).
+addopts = "--disable-socket --allow-hosts=127.0.0.1,::1"
 
 [tool.ruff]
 line-length = 110

src/lago_agent_sdk/lago_client.py

@@ -16,6 +16,11 @@ def __init__(self, api_key: str, api_url: str, timeout: float = 10.0, verify_ssl
         self.api_url = api_url.rstrip("/")
         self.timeout = timeout
         self.verify_ssl = verify_ssl
+        # Not `requests.post`: that builds and closes a Session per call, so every
+        # batch pays a fresh TCP + TLS handshake. Safe to share across sender threads
+        # (urllib3's pool is thread-safe), but size `pool_maxsize` (default 10) to the
+        # sender count — past it connections are opened and discarded, undoing this.
+        self._session = requests.Session()
         if not verify_ssl:
             # The customer explicitly opted out via config — they've already
             # accepted the risk; requests/urllib3's warning on every single
@@ -47,6 +52,12 @@ def __repr__(self) -> str:
             f"timeout={self.timeout}, verify_ssl={self.verify_ssl})"
         )
 
+    def close(self) -> None:
+        try:
+            self._session.close()
+        except Exception:  # noqa: BLE001
+            pass
+
     def send_batch(self, events: list[dict[str, Any]]) -> None:
         if not events:
             return
@@ -56,7 +67,7 @@ def send_batch(self, events: list[dict[str, Any]]) -> None:
             "Content-Type": "application/json",
         }
         payload = {"events": events}
-        resp = requests.post(
+        resp = self._session.post(
             url, headers=headers, data=json.dumps(payload), timeout=self.timeout, verify=self.verify_ssl
         )
         if not (200 <= resp.status_code < 300):

src/lago_agent_sdk/queue.py

@@ -125,17 +125,19 @@ def __init__(
         self._stopping = threading.Event()
         self._backoff_seconds = 0.0
         self._http_calls = 0  # for tests
+        # `_take_batch` pops events out before the send, so mid-POST they are in
+        # neither the buffer nor Lago. `flush()` has to wait on this too.
+        self._in_flight = 0
         # Per-thread "already reporting an overflow" flag — see push().
         self._reporting = threading.local()
 
         self._thread = threading.Thread(target=self._run, name="lago-queue", daemon=True)
         self._thread.start()
         atexit.register(self._atexit_shutdown)
 
-        # After fork, the daemon thread is gone in the child. Recreate it
-        # along with fresh sync primitives — the buffer's contents are copied
-        # over (which is fine: child re-emits its own events) but the lock
-        # state from the parent is unsafe to reuse.
+        # After fork the daemon thread is gone in the child. Recreate it with fresh
+        # sync primitives; the parent's buffer is dropped rather than inherited (see
+        # `_after_in_child`) so the two can never both deliver the same events.
         if hasattr(os, "register_at_fork"):
             os.register_at_fork(after_in_child=self._after_in_child)
 
@@ -146,6 +148,7 @@ def _after_in_child(self) -> None:
         self._buffer = deque()  # don't replay parent's events from the child
         self._backoff_seconds = 0.0
         self._http_calls = 0
+        self._in_flight = 0
         self._reporting = threading.local()
         # Note: the PricingProvider self-heals on fork via a PID check inside
         # lookup()/maybe_refresh(); we deliberately do NOT call into it from this
@@ -210,8 +213,8 @@ def flush(self, timeout: float = 5.0) -> bool:
         deadline = time.monotonic() + timeout
         while time.monotonic() < deadline:
             with self._lock:
-                empty = not self._buffer
-            if empty:
+                settled = not self._buffer and self._in_flight == 0
+            if settled:
                 return True
             self._wake.set()
             time.sleep(0.01)
@@ -235,11 +238,21 @@ def _take_batch(self) -> list[dict[str, Any]]:
                 return []
             n = min(self._max_batch_size, len(self._buffer))
             batch = [self._buffer.popleft() for _ in range(n)]
+            self._in_flight += n
         return batch
 
+    def _settle(self, n: int) -> None:
+        """Delivered, dropped for good, or back on the buffer — no longer in flight."""
+        if n <= 0:
+            return
+        with self._lock:
+            self._in_flight = max(0, self._in_flight - n)
+
     def _replay_failed(self, batch: list[dict[str, Any]]) -> None:
         with self._lock:
             self._buffer.extendleft(reversed(batch))
+            # Same lock acquisition as the re-queue, or flush() sees neither.
+            self._in_flight = max(0, self._in_flight - len(batch))
 
     def _report_error(self, exc: Exception, where: str = "send_batch") -> None:
         """Best-effort `on_error` callback — a customer's own callback must
@@ -303,7 +316,8 @@ def _send_individually(
                         exc,
                     )
         if retry:
-            self._replay_failed(retry)
+            self._replay_failed(retry)  # settles those
+        self._settle(len(batch) - len(retry))
         return len(retry)
 
     def _next_backoff(self) -> float:
@@ -338,6 +352,7 @@ def _drain_buffer(self) -> None:
                 self._http_calls += 1
                 self._sender(batch)
                 self._backoff_seconds = 0.0
+                self._settle(len(batch))
             except Exception as exc:  # noqa: BLE001
                 if _is_permanent_failure(exc):
                     # Lago's batch endpoint is all-or-nothing: a single bad
@@ -415,6 +430,7 @@ def _run(self) -> None:
                 break
             try:
                 self._sender(batch)
+                self._settle(len(batch))
             except Exception as exc:  # noqa: BLE001
                 if _is_permanent_failure(exc):
                     # `requeue_transient=False`: re-queuing here would put the event
@@ -430,6 +446,7 @@ def _run(self) -> None:
                         len(batch),
                         exc,
                     )
+                    self._settle(len(batch))
         with self._lock:
             stranded = len(self._buffer)
         if stranded:

src/lago_agent_sdk/sdk.py

@@ -793,4 +793,6 @@ def flush(self, timeout: float = 5.0) -> bool:
         return self._queue.flush(timeout=timeout)
 
     def shutdown(self, timeout: float = 5.0) -> None:
+        # Drain first, then release the socket — the queue's exit drain still needs it.
         self._queue.shutdown(timeout=timeout)
+        self._lago_client.close()

src/lago_agent_sdk/wrappers/anthropic.py

@@ -150,6 +150,7 @@ def _emit_from(payload: Any, model_id: str, opts: dict[str, Any]) -> None:
             sdk.emit(usage, **opts)
         except Exception as exc:  # noqa: BLE001
             logger.warning("lago: anthropic emit failed: %s", exc)
+            sdk._report_error(exc, "emit")
 
     # ------------------------------------------------------------------
     # Sync messages.create — auto-detects streaming via response shape
@@ -313,6 +314,7 @@ def _emit_final(self) -> None:
                 self._sdk.emit(usage, **self._opts)
         except Exception as exc:  # noqa: BLE001
             logger.warning("lago: anthropic stream-manager emit failed: %s", exc)
+            self._sdk._report_error(exc, "emit")
 
     async def _emit_final_async(self) -> None:
         """Async path — `AsyncMessageStream.get_final_message()` is a coroutine.
@@ -332,3 +334,4 @@ async def _emit_final_async(self) -> None:
                 self._sdk.emit(usage, **self._opts)
         except Exception as exc:  # noqa: BLE001
             logger.warning("lago: anthropic async stream-manager emit failed: %s", exc)
+            self._sdk._report_error(exc, "emit")