perf(track): bulk metadata upserts + drain pacing + single-pass scrub

fleet/track/api.py

@@ -34,6 +34,7 @@
 
 DEFAULT_TIMEOUT = 30.0
 SERVER_UPLOAD_URL_BATCH_CAP = 100  # /v1/track/upload-urls returns 400 above this.
+SERVER_BULK_UPSERT_BATCH_CAP = 100  # /v1/track/sessions/bulk; chunk to match.
 
 
 AuthInfo = Union[str, Tuple[str, str]]
@@ -56,6 +57,22 @@ class TrackTextMatch:
     negate: bool = False
 
 
+@dataclass(frozen=True)
+class BulkSessionUpsert:
+    """One item in a /v1/track/sessions/bulk request.
+
+    Mirrors the per-arg shape of `upsert_session`. `content_codec`,
+    `raw_bytes`, `stored_bytes` are only sent when this row carries
+    a fresh upload (i.e. include_content_metadata=True equivalent).
+    """
+
+    path: str
+    session: Any
+    content_codec: Optional[str] = None
+    raw_bytes: Optional[int] = None
+    stored_bytes: Optional[int] = None
+
+
 @dataclass(frozen=True)
 class TrackSessionSearchRequest:
     """Structured body for `POST /v1/track/sessions/search`."""
@@ -223,6 +240,34 @@ def upsert_session(
         )
         _raise(resp)
 
+    def upsert_sessions_bulk(
+        self,
+        *,
+        device_id: str,
+        items: list["BulkSessionUpsert"],
+    ) -> None:
+        """Bulk-register metadata for many sessions in one request.
+
+        Server reuses the single-row upsert translation logic per item, so
+        path → s3_key conversion and validation behave identically. Empty
+        list is a no-op. Chunks at SERVER_BULK_UPSERT_BATCH_CAP to match
+        the server cap.
+        """
+        if not items:
+            return
+        for chunk_start in range(0, len(items), SERVER_BULK_UPSERT_BATCH_CAP):
+            chunk = items[chunk_start : chunk_start + SERVER_BULK_UPSERT_BATCH_CAP]
+            body = {
+                "device_id": device_id,
+                "items": [_bulk_item_payload(item) for item in chunk],
+            }
+            resp = self._client.post(
+                "/v1/track/sessions/bulk",
+                json=body,
+                headers=self._headers(),
+            )
+            _raise(resp)
+
     def list_sessions(
         self,
         *,
@@ -395,6 +440,20 @@ def _session_payload(session: Any) -> dict[str, Any]:
     raise TypeError(f"Unsupported session payload type: {type(session)!r}")
 
 
+def _bulk_item_payload(item: "BulkSessionUpsert") -> dict[str, Any]:
+    out: dict[str, Any] = {
+        "path": item.path,
+        "session": _session_payload(item.session),
+    }
+    if item.content_codec is not None:
+        out["content_codec"] = item.content_codec
+    if item.raw_bytes is not None:
+        out["raw_bytes"] = item.raw_bytes
+    if item.stored_bytes is not None:
+        out["stored_bytes"] = item.stored_bytes
+    return out
+
+
 def _json_body(body: Any) -> dict[str, Any]:
     if is_dataclass(body):
         return asdict(body)

fleet/track/daemon.py

@@ -25,7 +25,7 @@
 import uuid
 from typing import TYPE_CHECKING, Optional
 
-from .api import TrackAPIClient, TrackAPIError
+from .api import BulkSessionUpsert, TrackAPIClient, TrackAPIError
 from .blocklist import TrackBlocklist
 from .drainer import QueueDrainer
 from .merkle import HashCache, MerkleTree
@@ -141,6 +141,11 @@ def __init__(
         # in the orchestrator metadata index.
         self._metadata_indexed: dict[str, str] = {}
         self._metadata_lock = threading.Lock()
+        # Workers append confirmed uploads here; the main loop flushes via
+        # the bulk endpoint each iteration. Buffering halves orchestrator
+        # round trips and avoids workers blocking on per-file POSTs.
+        self._metadata_buffer: list[tuple[str, str, BulkSessionUpsert]] = []
+        self._metadata_buffer_lock = threading.Lock()
 
     # ------------------------------------------------------------------ #
     # Public entry points                                                  #
@@ -187,6 +192,7 @@ def run_once(self, *, device_id: Optional[str] = None) -> "ReconcileResult":  #
 
         # Wait for in-flight uploads.
         self._pool.drain(timeout=60)
+        self._flush_metadata_buffer()
         if self._manifest_dirty:
             self._upload_manifest()
             self._manifest_dirty = False
@@ -248,6 +254,7 @@ def run(self) -> None:
                 last_queue_reset = now
 
             self._drain_queue()
+            self._flush_metadata_buffer()
             # Persist manifest opportunistically when the queue is idle, but
             # also at least every MANIFEST_FLUSH_INTERVAL seconds so a daemon
             # that stays continuously busy still publishes progress. The
@@ -269,6 +276,7 @@ def run(self) -> None:
         if self._pool:
             self._pool.drain(timeout=60)
             self._pool.shutdown()
+        self._flush_metadata_buffer()
         if self._manifest_dirty:
             self._upload_manifest()
         self._queue.close()
@@ -375,10 +383,18 @@ def _upload_manifest(self) -> None:
     # ------------------------------------------------------------------ #
 
     def _drain_queue(self) -> None:
-        """Delegate to QueueDrainer; one pass."""
+        """Drain pending work into the upload pool until the queue is empty.
+
+        Tight-loops drain_once so dispatch isn't capped by the main loop's
+        10s sleep. The thread pool's worker count is the real concurrency
+        ceiling; this just keeps it fed.
+        """
         if self._drainer is None:
             return
-        self._drainer.drain_once(self._device_id)
+        while not self._stop.is_set():
+            result = self._drainer.drain_once(self._device_id)
+            if result.claimed == 0:
+                break
 
     # ------------------------------------------------------------------ #
     # Upload callbacks                                                     #
@@ -431,12 +447,13 @@ def _upsert_session_metadata(
         *,
         upload_payload: UploadPayload | None = None,
     ) -> None:
-        """Best-effort metadata index update for a confirmed S3 object.
+        """Buffer a metadata upsert for the next bulk flush.
 
-        The v1 syncer's correctness still comes from S3 bytes + manifest. The
-        metadata index is a read-side accelerator for listing/resume, so a
-        transient failure here should be retried on a later reconcile rather
-        than marking the file upload failed.
+        Workers call this synchronously from upload-completion callbacks;
+        the actual HTTP roundtrip happens later in `_flush_metadata_buffer`,
+        which the daemon main loop invokes after each drain pass and on
+        graceful shutdown. Buffering halves orchestrator round trips and
+        keeps workers off the network for metadata writes.
 
         `upload_payload` is only present immediately after this process uploads
         the file. For files merely confirmed by the remote manifest, omit
@@ -453,27 +470,53 @@ def _upsert_session_metadata(
         if session is None:
             return
 
-        kwargs = {"include_content_metadata": False}
         if upload_payload is not None:
-            kwargs = {
-                "content_codec": upload_payload.content_codec,
-                "raw_bytes": upload_payload.raw_bytes,
-                "stored_bytes": upload_payload.stored_bytes,
-            }
+            item = BulkSessionUpsert(
+                path=rel_path,
+                session=session,
+                content_codec=upload_payload.content_codec,
+                raw_bytes=upload_payload.raw_bytes,
+                stored_bytes=upload_payload.stored_bytes,
+            )
+        else:
+            item = BulkSessionUpsert(path=rel_path, session=session)
+
+        with self._metadata_buffer_lock:
+            # Drop any prior buffered entry for this path; the latest sha wins.
+            self._metadata_buffer = [
+                (p, s, i) for p, s, i in self._metadata_buffer if p != rel_path
+            ]
+            self._metadata_buffer.append((rel_path, sha256, item))
+
+    def _flush_metadata_buffer(self) -> None:
+        """Send buffered metadata upserts to the orchestrator in one bulk call.
+
+        Best-effort: a transient failure leaves entries in the buffer for
+        the next flush. If the daemon dies, the next reconcile re-queues
+        anything missing, so we tolerate buffer loss.
+        """
+        with self._metadata_buffer_lock:
+            if not self._metadata_buffer:
+                return
+            pending = self._metadata_buffer
+            self._metadata_buffer = []
 
         try:
-            self._api.upsert_session(
+            self._api.upsert_sessions_bulk(
                 device_id=self._device_id,
-                path=rel_path,
-                session=session,
-                **kwargs,
+                items=[item for _, _, item in pending],
             )
         except Exception as e:
-            log.warning("metadata upsert failed %s: %s", rel_path, e)
+            log.warning("bulk metadata upsert failed (%d items): %s", len(pending), e)
+            with self._metadata_buffer_lock:
+                # Put failed items back at the front so they're retried first;
+                # newer items appended during the flush stay in order.
+                self._metadata_buffer = pending + self._metadata_buffer
             return
 
         with self._metadata_lock:
-            self._metadata_indexed[rel_path] = sha256
+            for rel_path, sha256, _ in pending:
+                self._metadata_indexed[rel_path] = sha256
 
     # ------------------------------------------------------------------ #
     # Status                                                               #

fleet/track/drainer.py

@@ -26,7 +26,7 @@
 
 # Server caps /v1/track/upload-urls at 100 paths per request. Match it
 # here to avoid a 400 if a single drain claims more.
-DEFAULT_BATCH_SIZE = 32
+DEFAULT_BATCH_SIZE = 100
 
 
 @dataclass(frozen=True)

fleet/track/scrubber.py

@@ -121,7 +121,8 @@ def scrub(text: str, rules: Iterable[Rule] = DEFAULT_RULES) -> ScrubResult:
 
     Hits are recorded against the *original* text's line numbers so that
     `flt track inspect` can point at the line the user can find in their
-    on-disk session file.
+    on-disk session file. The upload path uses `scrub_text` instead since
+    it doesn't need hits — see scrub_bytes.
     """
     rule_list: Sequence[Rule] = tuple(rules)
     hits: list[Hit] = []
@@ -133,21 +134,29 @@ def scrub(text: str, rules: Iterable[Rule] = DEFAULT_RULES) -> ScrubResult:
             line = text.count("\n", 0, m.start()) + 1
             hits.append(Hit(rule=rule.name, line=line, matched=m.group(0)))
 
-    # Second pass: actual substitution. We re-run regexes here rather than
-    # building offsets, because subs in earlier rules can change later
-    # rules' match positions (e.g. a long secret becoming "[REDACTED]"
-    # could expose a substring that looks like another secret).
-    for rule in rule_list:
-        text = rule.pattern.sub(rule.replacement, text)
+    return ScrubResult(text=scrub_text(text, rule_list), hits=tuple(hits))
+
 
-    return ScrubResult(text=text, hits=tuple(hits))
+def scrub_text(text: str, rules: Iterable[Rule] = DEFAULT_RULES) -> str:
+    """Apply substitution rules in order; return scrubbed text only.
+
+    Half the work of `scrub` — skips the hit-enumeration pass. Used by
+    the upload path, which discards hits anyway. Re-runs regexes for
+    substitution rather than building offsets, because subs in earlier
+    rules can change later rules' match positions (e.g. a long secret
+    becoming "[REDACTED]" could expose a substring that looks like
+    another secret).
+    """
+    for rule in rules:
+        text = rule.pattern.sub(rule.replacement, text)
+    return text
 
 
 def scrub_bytes(data: bytes, rules: Iterable[Rule] = DEFAULT_RULES) -> bytes:
     """Scrub raw bytes (decoded as UTF-8, re-encoded). Backward-compatible
     wrapper for the upload path that just wants the scrubbed payload."""
     try:
         text = data.decode("utf-8", errors="replace")
-        return scrub(text, rules).text.encode("utf-8")
+        return scrub_text(text, rules).encode("utf-8")
     except Exception:
         return data

tests/track/test_api.py

@@ -10,6 +10,8 @@
 import pytest
 
 from fleet.track.api import (
+    BulkSessionUpsert,
+    SERVER_BULK_UPSERT_BATCH_CAP,
     SERVER_UPLOAD_URL_BATCH_CAP,
     TrackAPIClient,
     TrackAPIError,
@@ -186,6 +188,84 @@ def handler(request: httpx.Request) -> httpx.Response:
     assert captured["body"]["stored_bytes"] == 200
 
 
+def test_upsert_sessions_bulk_posts_items_to_bulk_endpoint():
+    captured: dict = {}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        captured["url"] = str(request.url)
+        captured["body"] = json.loads(request.content)
+        return httpx.Response(204)
+
+    api = TrackAPIClient(client=_client_with_handler(handler), auth_provider=_auth)
+    api.upsert_sessions_bulk(
+        device_id="dev1",
+        items=[
+            BulkSessionUpsert(
+                path=".codex/sessions/a.jsonl",
+                session={"id": "sa", "tool": "codex"},
+                content_codec="gzip",
+                raw_bytes=1000,
+                stored_bytes=200,
+            ),
+            BulkSessionUpsert(
+                path=".cursor/projects/b.jsonl",
+                session={"id": "sb", "tool": "cursor"},
+            ),
+        ],
+    )
+
+    assert captured["url"] == "http://test/v1/track/sessions/bulk"
+    assert captured["body"]["device_id"] == "dev1"
+    items = captured["body"]["items"]
+    assert len(items) == 2
+    assert items[0]["path"] == ".codex/sessions/a.jsonl"
+    assert items[0]["session"]["tool"] == "codex"
+    assert items[0]["content_codec"] == "gzip"
+    assert items[0]["raw_bytes"] == 1000
+    assert items[0]["stored_bytes"] == 200
+    # Item 2 omitted content metadata; bulk payload must too.
+    assert items[1]["path"] == ".cursor/projects/b.jsonl"
+    assert "content_codec" not in items[1]
+    assert "raw_bytes" not in items[1]
+    assert "stored_bytes" not in items[1]
+
+
+def test_upsert_sessions_bulk_empty_list_skips_request():
+    called = {"n": 0}
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        called["n"] += 1
+        return httpx.Response(204)
+
+    api = TrackAPIClient(client=_client_with_handler(handler), auth_provider=_auth)
+    api.upsert_sessions_bulk(device_id="dev1", items=[])
+    assert called["n"] == 0
+
+
+def test_upsert_sessions_bulk_chunks_above_server_cap():
+    """Caller may pass more than the server cap; client splits into multiple POSTs."""
+    sizes_received: list[int] = []
+
+    def handler(request: httpx.Request) -> httpx.Response:
+        body = json.loads(request.content)
+        sizes_received.append(len(body["items"]))
+        return httpx.Response(204)
+
+    items = [
+        BulkSessionUpsert(path=f"{i}.jsonl", session={"id": f"s{i}"})
+        for i in range(SERVER_BULK_UPSERT_BATCH_CAP * 2 + 7)
+    ]
+
+    api = TrackAPIClient(client=_client_with_handler(handler), auth_provider=_auth)
+    api.upsert_sessions_bulk(device_id="dev1", items=items)
+
+    assert sizes_received == [
+        SERVER_BULK_UPSERT_BATCH_CAP,
+        SERVER_BULK_UPSERT_BATCH_CAP,
+        7,
+    ]
+
+
 def test_upsert_session_can_omit_content_metadata():
     captured: dict = {}