perf: anonymize traces inline during extraction instead of a second pass

src/teich/anonymize.py

@@ -69,14 +69,14 @@ def anonymize_path(input_path: Path, output_path: Path, *, in_place: bool = Fals
     report = AnonymizeReport(input_path=input_path, output_path=output_path)
     if input_path.is_file():
         destination = output_path / input_path.name if output_path.exists() and output_path.is_dir() else output_path
-        file_report = _anonymize_file(input_path, destination)
+        file_report = anonymize_file(input_path, destination)
         report.files.append(file_report)
         return report
 
     for source_file in sorted(path for path in input_path.rglob("*") if path.is_file()):
         relative_path = source_file.relative_to(input_path)
         destination = source_file if in_place else output_path / relative_path
-        file_report = _anonymize_file(source_file, destination)
+        file_report = anonymize_file(source_file, destination)
         report.files.append(file_report)
     return report
 
@@ -89,7 +89,7 @@ def _path_contains(parent: Path, child: Path) -> bool:
         return False
 
 
-def _anonymize_file(source: Path, destination: Path) -> AnonymizeFileReport:
+def anonymize_file(source: Path, destination: Path) -> AnonymizeFileReport:
     destination.parent.mkdir(parents=True, exist_ok=True)
     if source.suffix.lower() not in TEXT_EXTENSIONS:
         if source.resolve() != destination.resolve():

src/teich/cli.py

@@ -450,6 +450,7 @@ def _run_extract_command(
         sources=sessions_dir,
         model_filter=model_filter,
         clear_destination=True,
+        anonymize=not skip_anonymize,
     )
     if not result.source_paths:
         console.print(f"[red]No local {provider} session folders found.[/red]")
@@ -464,19 +465,18 @@ def _run_extract_command(
     stale_readme_path = output / "README.md"
     if stale_readme_path.exists() and stale_readme_path.is_file():
         stale_readme_path.unlink()
-    anonymize_report = None if skip_anonymize else anonymize_path(output, output, in_place=True)
     readme_path = _write_extract_readme(provider, output, model_filter=model_filter)
     extracted_message = f"Extracted {result.count} {provider} trace{'s' if result.count != 1 else ''}"
     if model_filter:
         extracted_message += f" with {model_filter}"
     console.print(f"[green]{extracted_message}[/green]", soft_wrap=True)
-    if anonymize_report is None:
+    if result.anonymize_totals is None:
         console.print(
             "[yellow]Skipped anonymization because --no-anon was passed. Review the data before sharing or uploading.[/yellow]",
             soft_wrap=True,
         )
     else:
-        totals = anonymize_report.totals
+        totals = result.anonymize_totals
         console.print(
             "[cyan]Automatically scrambled[/cyan] "
             f"{totals.get('api_key', 0)} API keys, "

src/teich/extract.py

@@ -15,6 +15,7 @@
 from tempfile import TemporaryDirectory
 from typing import Any, Literal
 
+from .anonymize import TraceAnonymizer, anonymize_file
 from .tool_schema import CURSOR_BUILTIN_TOOLS
 
 ExtractProvider = Literal["claude", "codex", "cursor", "hermes", "pi"]
@@ -39,12 +40,34 @@ class ExtractResult:
     destination_dir: Path
     copied_files: list[Path] = field(default_factory=list)
     source_paths: list[Path] = field(default_factory=list)
+    anonymize_totals: dict[str, int] | None = None
 
     @property
     def count(self) -> int:
         return len(self.copied_files)
 
 
+class _InlineAnonymizer:
+    """Anonymize traces as they are written, one TraceAnonymizer per file."""
+
+    def __init__(self) -> None:
+        self.totals: dict[str, int] = {"email": 0, "username": 0, "api_key": 0}
+
+    def _add(self, counts: dict[str, int]) -> None:
+        for key, count in counts.items():
+            self.totals[key] = self.totals.get(key, 0) + count
+
+    def copy_file(self, source: Path, destination: Path) -> None:
+        report = anonymize_file(source, destination)
+        self._add(report.replacements)
+
+    def anonymize_events(self, events: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        anonymizer = TraceAnonymizer()
+        events = [anonymizer.anonymize_value(event) for event in events]
+        self._add(anonymizer.counts)
+        return events
+
+
 def default_session_sources(provider: ExtractProvider, home: Path | None = None) -> list[Path]:
     """Return likely local session stores for a supported provider."""
     home = home or Path.home()
@@ -94,6 +117,7 @@ def extract_local_sessions(
     model_filter: str | None = None,
     clear_destination: bool = False,
     progress: ProgressCallback | None = None,
+    anonymize: bool = False,
 ) -> ExtractResult:
     """Extract local sessions for provider into output_dir."""
     source_candidates = list(sources) if sources is not None else default_session_sources(provider, home)
@@ -102,27 +126,33 @@ def extract_local_sessions(
     destination_dir.mkdir(parents=True, exist_ok=True)
     if clear_destination:
         _clear_extract_destination(destination_dir)
+    anonymizer = _InlineAnonymizer() if anonymize else None
     if provider == "hermes":
-        copied_files = _extract_hermes_state_dbs(resolved_sources, destination_dir, model_filter=model_filter)
+        copied_files = _extract_hermes_state_dbs(
+            resolved_sources, destination_dir, model_filter=model_filter, anonymizer=anonymizer
+        )
     elif provider == "cursor":
         copied_files = _extract_cursor_databases(
             resolved_sources,
             destination_dir,
             model_filter=model_filter,
             progress=progress,
+            anonymizer=anonymizer,
         )
     else:
         copied_files = _extract_jsonl_session_files(
             provider,
             resolved_sources,
             destination_dir,
             model_filter=model_filter,
+            anonymizer=anonymizer,
         )
     return ExtractResult(
         provider=provider,
         destination_dir=destination_dir,
         copied_files=copied_files,
         source_paths=resolved_sources,
+        anonymize_totals=anonymizer.totals if anonymizer is not None else None,
     )
 
 
@@ -283,6 +313,7 @@ def _extract_cursor_databases(
     *,
     model_filter: str | None = None,
     progress: ProgressCallback | None = None,
+    anonymizer: _InlineAnonymizer | None = None,
 ) -> list[Path]:
     indexed_rows: list[tuple[int, dict[str, Any]]] = []
     row_index = 0
@@ -322,11 +353,13 @@ def _extract_cursor_databases(
     if model_filter:
         rows = [row for row in rows if trace_matches_model([row], model_filter)]
     if not rows:
-        return _extract_cursor_project_files(sources, destination_dir, model_filter=model_filter)
+        return _extract_cursor_project_files(
+            sources, destination_dir, model_filter=model_filter, anonymizer=anonymizer
+        )
     copied: list[Path] = []
     for row_index, row in enumerate(rows, start=1):
         destination = _unique_destination(destination_dir, _cursor_session_file_name(row, row_index))
-        _write_jsonl_dict_events(destination, _cursor_row_to_transcript_events(row, row_index))
+        _write_jsonl_dict_events(destination, _cursor_row_to_transcript_events(row, row_index), anonymizer=anonymizer)
         copied.append(destination)
     return copied
 
@@ -336,6 +369,7 @@ def _extract_cursor_project_files(
     destination_dir: Path,
     *,
     model_filter: str | None = None,
+    anonymizer: _InlineAnonymizer | None = None,
 ) -> list[Path]:
     copied: list[Path] = []
     for source in sources:
@@ -354,7 +388,9 @@ def _extract_cursor_project_files(
                 except ValueError:
                     relative = transcript.name
             destination = _unique_nested_destination(destination_dir / relative)
-            _write_jsonl_dict_events(destination, _cursor_native_transcript_events(events, transcript, project_dir))
+            _write_jsonl_dict_events(
+                destination, _cursor_native_transcript_events(events, transcript, project_dir), anonymizer=anonymizer
+            )
             copied.append(destination)
     return copied
 
@@ -464,12 +500,13 @@ def _extract_jsonl_session_files(
     destination_dir: Path,
     *,
     model_filter: str | None = None,
+    anonymizer: _InlineAnonymizer | None = None,
 ) -> list[Path]:
     copied: list[Path] = []
     for source in sources:
         for path in _jsonl_files(source):
             destination = _unique_destination(destination_dir, path.name)
-            if _copy_provider_jsonl(provider, path, destination, model_filter=model_filter):
+            if _copy_provider_jsonl(provider, path, destination, model_filter=model_filter, anonymizer=anonymizer):
                 copied.append(destination)
     return copied
 
@@ -480,17 +517,24 @@ def _copy_provider_jsonl(
     destination: Path,
     *,
     model_filter: str | None = None,
+    anonymizer: _InlineAnonymizer | None = None,
 ) -> bool:
     events = _read_jsonl_dict_events(source)
     if events is None:
         if model_filter:
             return False
         destination.parent.mkdir(parents=True, exist_ok=True)
-        shutil.copy2(source, destination)
+        if anonymizer is not None:
+            anonymizer.copy_file(source, destination)
+        else:
+            shutil.copy2(source, destination)
         return True
     if model_filter and not trace_matches_model(events, model_filter):
         return False
     destination.parent.mkdir(parents=True, exist_ok=True)
+    if anonymizer is not None:
+        anonymizer.copy_file(source, destination)
+        return True
     shutil.copy2(source, destination)
     try:
         shutil.copystat(source, destination)
@@ -516,7 +560,13 @@ def _read_jsonl_dict_events(path: Path) -> list[dict[str, Any]] | None:
     return events
 
 
-def _write_jsonl_dict_events(path: Path, events: list[dict[str, Any]]) -> None:
+def _write_jsonl_dict_events(
+    path: Path,
+    events: list[dict[str, Any]],
+    anonymizer: _InlineAnonymizer | None = None,
+) -> None:
+    if anonymizer is not None:
+        events = anonymizer.anonymize_events(events)
     path.parent.mkdir(parents=True, exist_ok=True)
     with path.open("w", encoding="utf-8") as handle:
         for event in events:
@@ -529,11 +579,12 @@ def _write_individual_jsonl_rows(
     destination_dir: Path,
     rows: list[dict[str, Any]],
     file_name_for_row: Callable[[dict[str, Any], int], str],
+    anonymizer: _InlineAnonymizer | None = None,
 ) -> list[Path]:
     copied: list[Path] = []
     for row_index, row in enumerate(rows, start=1):
         destination = _unique_destination(destination_dir, file_name_for_row(row, row_index))
-        _write_jsonl_dict_events(destination, [row])
+        _write_jsonl_dict_events(destination, [row], anonymizer=anonymizer)
         copied.append(destination)
     return copied
 
@@ -2106,13 +2157,16 @@ def _extract_hermes_state_dbs(
     destination_dir: Path,
     *,
     model_filter: str | None = None,
+    anonymizer: _InlineAnonymizer | None = None,
 ) -> list[Path]:
     copied: list[Path] = []
     for source in sources:
         state_dbs = [source] if source.is_file() else sorted(source.rglob("state.db"))
         for state_db in state_dbs:
             rows = _hermes_state_db_session_rows(state_db, model_filter=model_filter)
-            copied.extend(_write_individual_jsonl_rows(destination_dir, rows, _hermes_session_file_name))
+            copied.extend(
+                _write_individual_jsonl_rows(destination_dir, rows, _hermes_session_file_name, anonymizer=anonymizer)
+            )
     return copied
 
 

src/teich/studio/extraction.py

@@ -9,7 +9,6 @@
 from pathlib import Path
 from typing import Any
 
-from ..anonymize import anonymize_path
 from ..config import Config
 from ..extract import CURSOR_EXTRACTION_NOTICE, ExtractProvider, extract_local_sessions
 from ..trace_readme import extraction_readme_tags, write_traces_readme
@@ -98,13 +97,16 @@ def emit_progress(event: dict[str, Any]) -> None:
                 last_progress_at = now
                 self.events.append(event)
 
+            if not self.skip_anonymize:
+                self._emit_status("running", "Extracting and anonymizing traces...")
             result = extract_local_sessions(
                 self.provider,
                 output_dir=self.output_dir,
                 sources=self.source_paths or None,
                 model_filter=self.model_filter,
                 clear_destination=True,
                 progress=emit_progress,
+                anonymize=not self.skip_anonymize,
             )
             self.detected_sources = [str(path) for path in result.source_paths]
             self.result_files = [path.name for path in result.copied_files]
@@ -137,9 +139,7 @@ def emit_progress(event: dict[str, Any]) -> None:
                     }
                 )
             else:
-                self._emit_status("running", "Anonymizing staged traces...")
-                report = anonymize_path(self.output_dir, self.output_dir, in_place=True)
-                self.anonymize_totals = report.totals
+                self.anonymize_totals = result.anonymize_totals or {}
                 self.events.append(
                     {
                         "kind": "extract_anonymize",

tests/test_studio.py

@@ -804,7 +804,7 @@ def test_extract_endpoint_warns_cursor_may_take_a_while(client):
     source = client.project_dir / "state.vscdb"
     source.write_text("", encoding="utf-8")
 
-    def fake_extract(provider, *, output_dir, sources=None, model_filter=None, clear_destination=False, progress=None):
+    def fake_extract(provider, *, output_dir, sources=None, model_filter=None, clear_destination=False, progress=None, anonymize=False):
         assert provider == "cursor"
         assert sources == [source]
         assert model_filter is None