diff --git a/src/teich/anonymize.py b/src/teich/anonymize.py index c9c1b8f..f4141e1 100644 --- a/src/teich/anonymize.py +++ b/src/teich/anonymize.py @@ -84,7 +84,7 @@ 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 @@ -95,10 +95,10 @@ def anonymize_path(input_path: Path, output_path: Path, *, in_place: bool = Fals # Each file is anonymized independently (fresh TraceAnonymizer per # file), so files can be processed in parallel safely. with ProcessPoolExecutor(max_workers=workers) as executor: - report.files.extend(executor.map(_anonymize_file, source_files, destinations)) + report.files.extend(executor.map(anonymize_file, source_files, destinations)) else: for source_file, destination in zip(source_files, destinations): - report.files.append(_anonymize_file(source_file, destination)) + report.files.append(anonymize_file(source_file, destination)) return report @@ -110,7 +110,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(): diff --git a/src/teich/cli.py b/src/teich/cli.py index d3b04f9..c88a54a 100644 --- a/src/teich/cli.py +++ b/src/teich/cli.py @@ -452,6 +452,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]") @@ -466,19 +467,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, trace_files=result.copied_files) 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, " diff --git a/src/teich/extract.py b/src/teich/extract.py index c600eb6..aa5daab 100644 --- a/src/teich/extract.py +++ b/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,41 @@ 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 anonymize_remaining_files(self, root: Path, excluded: Iterable[Path]) -> None: + """Scrub pre-existing output files without re-reading newly written traces.""" + excluded_paths = {path.resolve() for path in excluded} + for path in sorted(candidate for candidate in root.rglob("*") if candidate.is_file()): + if path.resolve() not in excluded_paths: + self.copy_file(path, path) + + 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 +124,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,14 +133,18 @@ 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( @@ -117,12 +152,16 @@ def extract_local_sessions( resolved_sources, destination_dir, model_filter=model_filter, + anonymizer=anonymizer, ) + if anonymizer is not None and copied_files: + anonymizer.anonymize_remaining_files(destination_dir, copied_files) 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 +322,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 +362,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 +378,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 +397,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 +509,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 +526,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 +569,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 +588,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 +2166,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 diff --git a/src/teich/studio/extraction.py b/src/teich/studio/extraction.py index 3b74847..cbb36e2 100644 --- a/src/teich/studio/extraction.py +++ b/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,6 +97,8 @@ 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, @@ -105,6 +106,7 @@ def emit_progress(event: dict[str, Any]) -> 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", diff --git a/tests/test_extract_anonymize_cli.py b/tests/test_extract_anonymize_cli.py index c765458..4a603ed 100644 --- a/tests/test_extract_anonymize_cli.py +++ b/tests/test_extract_anonymize_cli.py @@ -10,7 +10,7 @@ from teich import anonymize as anonymize_module from teich.converter import convert_traces_to_training_data from teich import extract as extract_module -from teich.extract import CURSOR_EXTRACTION_NOTICE, default_session_sources +from teich.extract import CURSOR_EXTRACTION_NOTICE, default_session_sources, extract_local_sessions from teich.cli import app runner = CliRunner() @@ -1220,6 +1220,68 @@ def test_extract_no_anon_skips_automatic_anonymization(tmp_path: Path): assert original_key in text +def test_extract_inline_anonymization_scrubs_traces_and_leftover_text(tmp_path: Path): + sessions_dir = tmp_path / "sessions" + sessions_dir.mkdir() + original_key = "sk-proj-ABCDEFGHIJKLMNOPQRSTUVWXYZ123456" + for index in range(2): + (sessions_dir / f"session-{index}.jsonl").write_text( + json.dumps( + { + "type": "message", + "text": f"/home/alice/project alice@example.com {original_key}", + } + ) + + "\n", + encoding="utf-8", + ) + + output_dir = tmp_path / "output" + output_dir.mkdir() + stale_text = output_dir / "notes.txt" + stale_text.write_text(f"alice@example.com {original_key}\n", encoding="utf-8") + stale_binary = output_dir / "archive.bin" + stale_binary.write_bytes(b"alice@example.com sk-proj-ABCDEFGHIJKLMNOPQRSTUVWXYZ123456") + + result = extract_local_sessions( + "codex", + output_dir=output_dir, + sources=[sessions_dir], + clear_destination=True, + anonymize=True, + ) + + assert result.anonymize_totals == {"email": 3, "username": 2, "api_key": 3} + for trace in result.copied_files: + text = trace.read_text(encoding="utf-8") + assert "/home/user1/project" in text + assert "redacted-user1@example.com" in text + assert original_key not in text + assert "redacted-user1@example.com" in stale_text.read_text(encoding="utf-8") + assert original_key not in stale_text.read_text(encoding="utf-8") + assert stale_binary.read_bytes() == b"alice@example.com sk-proj-ABCDEFGHIJKLMNOPQRSTUVWXYZ123456" + + +def test_extract_does_not_scrub_leftovers_when_no_trace_was_copied(tmp_path: Path): + sessions_dir = tmp_path / "empty-sessions" + sessions_dir.mkdir() + output_dir = tmp_path / "output" + output_dir.mkdir() + stale_text = output_dir / "notes.txt" + stale_text.write_text("alice@example.com\n", encoding="utf-8") + + result = extract_local_sessions( + "codex", + output_dir=output_dir, + sources=[sessions_dir], + clear_destination=True, + anonymize=True, + ) + + assert result.copied_files == [] + assert stale_text.read_text(encoding="utf-8") == "alice@example.com\n" + + def test_extract_model_filter_for_codex_claude_cursor_pi_and_hermes(tmp_path: Path): codex_dir = tmp_path / "codex" claude_dir = tmp_path / "claude" diff --git a/tests/test_studio.py b/tests/test_studio.py index 658b8f5..fb596bd 100644 --- a/tests/test_studio.py +++ b/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