diff --git a/.github/scripts/check_crap_threshold.py b/.github/scripts/check_crap_threshold.py index 0fac2e794..fbc19de37 100644 --- a/.github/scripts/check_crap_threshold.py +++ b/.github/scripts/check_crap_threshold.py @@ -1,26 +1,48 @@ #!/usr/bin/env python3 -"""Fail when pytest-crap reports a CRAP score at or above a threshold.""" +"""Enforce CRAP scores using executable-statement coverage.""" from __future__ import annotations import argparse +from dataclasses import replace import sys from pathlib import Path -from coverage import CoverageData -from pytest_crap.calculator import calculate_crap +from coverage import Coverage +from pytest_crap.calculator import FunctionScore, calculate_crap -def covered_lines_for_file(coverage_file: Path, source_file: Path) -> set[int]: - data = CoverageData(basename=str(coverage_file)) - data.read() +def scores_for_file(coverage: Coverage, source_file: Path) -> list[FunctionScore]: + measured_files = { + Path(filename).resolve() for filename in coverage.get_data().measured_files() + } + if source_file.resolve() not in measured_files: + raise SystemExit(f"No coverage data found for {source_file.resolve()}") - source_file = source_file.resolve() - for measured_file in data.measured_files(): - if Path(measured_file).resolve() == source_file: - return set(data.lines(measured_file) or []) - - raise SystemExit(f"No coverage data found for {source_file}") + _, statements, excluded, missing, _ = coverage.analysis2(str(source_file)) + executable_lines = set(statements) - set(excluded) + missing_lines = set(missing) + scores = [] + for score in calculate_crap(str(source_file), executable_lines - missing_lines): + function_lines = { + line + for line in executable_lines + if score.start_line <= line <= score.end_line + } + fraction = ( + len(function_lines - missing_lines) / len(function_lines) + if function_lines + else 1.0 + ) + # pytest-crap's physical-line denominator penalizes comments and docstrings. + scores.append( + replace( + score, + coverage_percent=fraction * 100, + crap=score.cc**2 * (1 - fraction) ** 3 + score.cc, + ) + ) + return scores def main() -> int: @@ -30,9 +52,31 @@ def main() -> int: parser.add_argument("--max-crap", type=float, required=True) args = parser.parse_args() - covered_lines = covered_lines_for_file(args.coverage_file, args.source_file) - scores = calculate_crap(str(args.source_file), covered_lines) + if args.source_file.is_dir(): + source_files = sorted( + path + for path in args.source_file.rglob("*.py") + if not {"test", "tests"}.intersection( + path.relative_to(args.source_file).parts + ) + ) + else: + source_files = [args.source_file] + if not source_files: + raise SystemExit(f"No production Python files found in {args.source_file}") + + coverage = Coverage(data_file=str(args.coverage_file)) + coverage.load() + scores = [ + score + for source_file in source_files + for score in scores_for_file(coverage, source_file) + ] offenders = [score for score in scores if score.crap >= args.max_crap] + print( + f"Analyzed {len(scores)} functions in {len(source_files)} production files; " + f"{len(offenders)} functions scored at or above {args.max_crap:g}" + ) if not offenders: print(f"All CRAP scores are below {args.max_crap:g} for {args.source_file}") @@ -41,7 +85,7 @@ def main() -> int: print(f"CRAP scores must be below {args.max_crap:g} for {args.source_file}") for score in sorted(offenders, key=lambda item: item.crap, reverse=True): print( - f"{score.name}: CRAP={score.crap:.2f}, " + f"{score.file_path}:{score.start_line} {score.name}: CRAP={score.crap:.2f}, " f"CC={score.cc}, coverage={score.coverage_percent:.1f}%" ) return 1 diff --git a/.github/scripts/check_public_api.py b/.github/scripts/check_public_api.py index bc8a3a8d6..899ddbb8c 100644 --- a/.github/scripts/check_public_api.py +++ b/.github/scripts/check_public_api.py @@ -233,11 +233,14 @@ def _load_griffe(): def generate_snapshot() -> str: griffe = _load_griffe() + from constructor_helper_attributes import ConstructorHelperAttributes + package = griffe.load( "posthog", allow_inspection=False, search_paths=[ROOT], submodules=True, + extensions=griffe.load_extensions(ConstructorHelperAttributes()), ) records = sorted(_record(obj) for obj in _iter_module_members(package)) return HEADER + "\n".join(records) + "\n" diff --git a/.github/scripts/constructor_helper_attributes.py b/.github/scripts/constructor_helper_attributes.py new file mode 100644 index 000000000..c1ef96bfb --- /dev/null +++ b/.github/scripts/constructor_helper_attributes.py @@ -0,0 +1,110 @@ +"""Recover public instance attributes initialized by private constructor helpers.""" + +from __future__ import annotations + +import ast +from typing import Iterator + +from griffe import Class, Extension, Function, Visitor, ast_children + + +_DEFERRED_SCOPES = ( + ast.FunctionDef, + ast.AsyncFunctionDef, + ast.ClassDef, + ast.Lambda, + ast.GeneratorExp, +) + + +def _eager_nodes(node: ast.AST) -> Iterator[ast.AST]: + if isinstance(node, _DEFERRED_SCOPES): + return + yield node + for child in ast_children(node): + yield from _eager_nodes(child) + + +def _public_instance_target(node: ast.AST) -> ast.Attribute | None: + if isinstance(node, ast.Assign): + targets = node.targets + elif isinstance(node, ast.AnnAssign): + targets = [node.target] + else: + return None + if ( + len(targets) == 1 + and isinstance(target := targets[0], ast.Attribute) + and isinstance(target.value, ast.Name) + and target.value.id == "self" + and not target.attr.startswith("_") + ): + return target + return None + + +class ConstructorHelperAttributes(Extension): + """Add missing attributes from direct, same-class private helper calls. + + Griffe ignores method-body assignments outside ``__init__``. Replay only + simple public instance assignments from reachable synchronous, undecorated helpers, + keeping existing constructor and class records unchanged. + """ + + def on_class_members(self, *, node, cls: Class, agent, **kwargs) -> None: + if not isinstance(node, ast.ClassDef) or not isinstance(agent, Visitor): + return + methods = { + method.name: method + for method in node.body + if isinstance(method, ast.FunctionDef) + and not method.decorator_list + and method.args.args + and method.args.args[0].arg == "self" + } + constructor = cls.members.get("__init__") + if not isinstance(constructor, Function) or "__init__" not in methods: + return + existing_names = set(cls.members) + previous = agent.current + try: + # Griffe's constructor visitor also stores Function despite this annotation. + agent.current = constructor # type: ignore[assignment] + self._follow( + methods["__init__"], methods, agent, existing_names, set(), False + ) + finally: + agent.current = previous + + def _follow(self, method, methods, agent, existing_names, active, in_helper): + if method.name in active: + return + active.add(method.name) + try: + # ast_children supplies parent links required by Griffe's assignment visitor. + for statement in ast_children(method): + if statement not in method.body: + continue + for node in _eager_nodes(statement): + target = _public_instance_target(node) if in_helper else None + if target is not None and target.attr not in existing_names: + agent.visit(node) + if ( + isinstance(node, ast.Call) + and isinstance(node.func, ast.Attribute) + and isinstance(node.func.value, ast.Name) + and node.func.value.id == "self" + and node.func.attr.startswith("_") + and not node.func.attr.startswith("__") + and node.func.attr in methods + ): + self._follow( + methods[node.func.attr], + methods, + agent, + existing_names, + active, + True, + ) + finally: + active.remove(method.name) diff --git a/.github/scripts/test_check_crap_threshold.py b/.github/scripts/test_check_crap_threshold.py new file mode 100644 index 000000000..5d697750a --- /dev/null +++ b/.github/scripts/test_check_crap_threshold.py @@ -0,0 +1,147 @@ +"""Regression tests for executable-line CRAP scoring.""" + +from pathlib import Path +import subprocess +import sys + +from coverage import CoverageData +import pytest + + +SCRIPT_PATH = Path(__file__).with_name("check_crap_threshold.py") +FUNCTION = """def choose(value): + if value == 1: + return 'one' + if value == 2: + return 'two' + if value == 3: + return 'three' + if value == 4: + return 'four' + return 'other' + +for value in range(5): + choose(value) +""" + + +def collect_coverage(tmp_path, source): + source_file = (tmp_path / "example.py").resolve() + source_file.write_text(source) + data_file = tmp_path / ".coverage" + config_file = tmp_path / ".coveragerc" + config_file.write_text("[run]\n") + subprocess.run( + [ + sys.executable, + "-m", + "coverage", + "run", + f"--rcfile={config_file}", + f"--data-file={data_file}", + str(source_file), + ], + check=True, + capture_output=True, + text=True, + ) + return source_file, data_file + + +def check(source_file, data_file, threshold): + return subprocess.run( + [ + sys.executable, + str(SCRIPT_PATH), + str(source_file), + "--coverage-file", + str(data_file), + "--max-crap", + str(threshold), + ], + capture_output=True, + text=True, + ) + + +@pytest.mark.parametrize( + "source", + [ + FUNCTION, + FUNCTION.replace( + " if value == 1:", " # Explanation\n\n" * 30 + " if value == 1:" + ), + FUNCTION.replace( + " if value == 1:", + ' """Description.\n' + + " Details.\n" * 30 + + ' """\n if value == 1:', + ), + FUNCTION.replace("def choose(value):", "def choose(\n value,\n):"), + ], + ids=["plain", "comments-and-blanks", "docstring", "multiline-signature"], +) +def test_fully_covered_function_is_unaffected_by_non_executable_lines(tmp_path, source): + source_file, data_file = collect_coverage(tmp_path, source) + result = check(source_file, data_file, 10) + assert result.returncode == 0, result.stdout + result.stderr + at_complexity_floor = check(source_file, data_file, 5) + assert at_complexity_floor.returncode == 1 + assert "CRAP=5.00, CC=5, coverage=100.0%" in at_complexity_floor.stdout + + +def test_uncovered_executable_branch_still_increases_score(tmp_path): + source_file, data_file = collect_coverage( + tmp_path, + "def choose(value):\n if value:\n return 1\n return 0\n\nchoose(True)\n", + ) + result = check(source_file, data_file, 2.0625) + assert result.returncode == 1 + assert "CRAP=2.06, CC=2, coverage=75.0%" in result.stdout + assert check(source_file, data_file, 2.0626).returncode == 0 + + +def test_coverage_exclusions_are_not_missing_statements(tmp_path): + source_file, data_file = collect_coverage( + tmp_path, + "def choose(value):\n if value: # pragma: no cover\n return 1\n return 0\n\nchoose(False)\n", + ) + result = check(source_file, data_file, 2) + assert result.returncode == 1 + assert "CRAP=2.00, CC=2, coverage=100.0%" in result.stdout + + +def test_directory_includes_unexecuted_production_and_excludes_test_helpers(tmp_path): + package = tmp_path / "posthog" + source_file = package / "ai" / "converter.py" + source_file.parent.mkdir(parents=True) + source_file.write_text( + "def choose(value):\n if value:\n return 1\n return 0\n" + ) + for directory in ("test", "tests"): + helper = package / directory / "ai" / "utils.py" + helper.parent.mkdir(parents=True) + helper.write_text("not valid Python") + data_file = tmp_path / ".coverage" + data = CoverageData(basename=str(data_file)) + data.add_lines({str(source_file.resolve()): set()}) + data.write() + result = check(package, data_file, 6) + assert result.returncode == 1 + assert "1 functions in 1 production files" in result.stdout + assert "CRAP=6.00, CC=2, coverage=0.0%" in result.stdout + assert str(source_file) in result.stdout + + +def test_missing_coverage_does_not_pass(tmp_path): + source_file = tmp_path / "example.py" + source_file.write_text("def choose():\n return 1\n") + result = check(source_file, tmp_path / ".coverage", 10) + assert result.returncode != 0 + assert "No coverage data found" in result.stderr + + +def test_empty_source_directory_does_not_pass(tmp_path): + result = check(tmp_path, tmp_path / ".coverage", 10) + assert result.returncode != 0 + assert "No production Python files found" in result.stderr diff --git a/.github/scripts/test_check_public_api.py b/.github/scripts/test_check_public_api.py index 5aa6da473..923e2a371 100644 --- a/.github/scripts/test_check_public_api.py +++ b/.github/scripts/test_check_public_api.py @@ -7,6 +7,9 @@ from pathlib import Path from types import SimpleNamespace from tempfile import TemporaryDirectory +from textwrap import dedent + +import pytest SCRIPT_PATH = Path(__file__).with_name("check_public_api.py") @@ -75,11 +78,193 @@ def test_explicit_exports_from_private_modules_include_signatures() -> None: assert not any("Hidden" in record or "_internal" in record for record in records) +def _snapshot_for_source(tmp_path, source): + package = tmp_path / "posthog" + package.mkdir(exist_ok=True) + (package / "__init__.py").write_text(dedent(source)) + checker = load_check_public_api() + checker.ROOT = tmp_path + with pytest.MonkeyPatch.context() as patch: + patch.chdir(tmp_path) + return checker.generate_snapshot() + + +def test_constructor_helpers_preserve_public_attributes(tmp_path): + snapshot = _snapshot_for_source( + tmp_path, + """ + class Client: + def __init__(self, key, before_send): + self._configure(key) + self._set_before_send(before_send) + def _configure(self, key): + self.secret_key: str = key + self._configure_more() + def _configure_more(self): + self.options = dict(enabled=True) + def _set_before_send(self, before_send): + self.before_send = before_send + """, + ) + assert "attribute posthog.Client.secret_key: str = key\n" in snapshot + assert "attribute posthog.Client.options = dict(enabled=True)\n" in snapshot + assert "attribute posthog.Client.before_send = before_send\n" in snapshot + assert "_configure" not in snapshot + + +def test_constructor_helper_cycles_do_not_leak_class_context(tmp_path): + snapshot = _snapshot_for_source( + tmp_path, + """ + class First: + def __init__(self): self._one() + def _one(self): + self.one = 1 + self._two() + def _two(self): + self.two = 2 + self._one() + class Second: + def __init__(self): self.three = 3 + module_value = 4 + """, + ) + assert "attribute posthog.First.one = 1\n" in snapshot + assert "attribute posthog.First.two = 2\n" in snapshot + assert "attribute posthog.Second.three = 3\n" in snapshot + assert "attribute posthog.module_value = 4\n" in snapshot + assert "posthog.First.three" not in snapshot + + +def test_uncalled_and_deferred_helpers_do_not_add_attributes(tmp_path): + snapshot = _snapshot_for_source( + tmp_path, + """ + class Client: + def __init__(self): + self._configure() + saved = self._unrelated + later = lambda: self._lambda_only() + deferred = (self._generator_only() for _ in range(1)) + def nested(): self._nested_only() + class Inner: + def run(inner): self._inner_only() + self.public() + other._foreign() + self._static() + self._async() + def _configure(self): + self.visible = True + self._private = False + def _unrelated(self): self.unrelated = True + def _lambda_only(self): self.lambda_only = True + def _generator_only(self): self.generator_only = True + def _nested_only(self): self.nested_only = True + def _inner_only(self): self.inner_only = True + def public(self): self.public_only = True + def _foreign(self): self.foreign_only = True + @staticmethod + def _static(self): self.static_only = True + async def _async(self): self.async_only = True + """, + ) + attributes = { + record + for record in snapshot.splitlines() + if record.startswith("attribute posthog.Client.") + } + assert attributes == {"attribute posthog.Client.visible = True"} + + +def test_direct_constructor_records_and_griffe_branch_precedence_are_preserved( + tmp_path, +): + snapshot = _snapshot_for_source( + tmp_path, + """ + class Client: + def __init__(self): + self.direct: int = 10 + self._configure() + def _configure(self): + self.direct = 20 + if condition: + self.branch = "first" + else: + self.branch = "second" + self.root = None + if self.root is None: + try: + self.root = getcwd() + except Exception: + self.root = None + """, + ) + assert "attribute posthog.Client.direct: int = 10\n" in snapshot + assert "attribute posthog.Client.branch = 'first'\n" in snapshot + assert "attribute posthog.Client.root = getcwd()\n" in snapshot + + +def test_helper_attribute_changes_remain_visible_in_snapshot(tmp_path): + source = """ + class Client: + def __init__(self): self._configure() + def _configure(self): self.visible = True + """ + before = _snapshot_for_source(tmp_path, source) + removed = _snapshot_for_source( + tmp_path, source.replace("self.visible = True", "pass") + ) + changed = _snapshot_for_source( + tmp_path, source.replace("self.visible = True", "self.visible = False") + ) + assert "attribute posthog.Client.visible = True\n" in before + assert "attribute posthog.Client.visible" not in removed + assert "attribute posthog.Client.visible = False\n" in changed + + +def test_constructor_discovery_retains_builtin_dataclass_processing(tmp_path): + snapshot = _snapshot_for_source( + tmp_path, + """ + from dataclasses import dataclass + @dataclass + class Record: + value: str + """, + ) + assert "class posthog.Record(value: str)\n" in snapshot + + +def test_constructor_helper_discovery_restores_visitor_after_failure(monkeypatch): + import ast + import griffe + from constructor_helper_attributes import ConstructorHelperAttributes + + cls = griffe.Class("Client") + constructor = griffe.Function("__init__", parent=cls) + cls.set_member("__init__", constructor) + source = ( + "class Client:\n" + " def __init__(self): self._helper()\n" + " def _helper(self): self.value = 1\n" + ) + agent = griffe.Visitor("sample", Path("sample.py"), source, griffe.Extensions()) + agent.current = cls + + def fail_to_visit(node): + raise RuntimeError("sentinel") + + monkeypatch.setattr(agent, "visit", fail_to_visit) + with pytest.raises(RuntimeError, match="sentinel"): + ConstructorHelperAttributes().on_class_members( + node=ast.parse(source).body[0], cls=cls, agent=agent + ) + assert agent.current is cls + + def main() -> int: - test_attribute_details_uses_placeholder_values() - test_explicit_exports_from_private_modules_include_signatures() - print("check_public_api tests passed.") - return 0 + return pytest.main([str(Path(__file__)), "--timeout=30", "-q"]) if __name__ == "__main__": diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d93383859..4963753a5 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -74,7 +74,7 @@ jobs: - name: Check public API snapshot run: | - python .github/scripts/test_check_public_api.py + UV_PROJECT_ENVIRONMENT=$pythonLocation uv run --extra dev --with pytest --with pytest-timeout python .github/scripts/test_check_public_api.py make public_api_check strict-type-smoke: @@ -231,7 +231,12 @@ jobs: changed_files="$(git diff --name-only "$base" "$head" -- \ posthog/utils.py \ posthog/test/test_utils.py \ - posthog/test/test_size_limited_dict.py)" + posthog/test/test_size_limited_dict.py \ + .github/scripts/check_crap_threshold.py \ + .github/scripts/test_check_crap_threshold.py \ + .github/workflows/ci.yml \ + pyproject.toml \ + uv.lock)" if [[ -n "$changed_files" ]]; then echo "changed=true" >> "$GITHUB_OUTPUT" @@ -260,7 +265,8 @@ jobs: shell: bash run: | set -euo pipefail - UV_PROJECT_ENVIRONMENT=$pythonLocation uv run --extra test --with pytest-crap --with pytest-cov pytest posthog/test/test_utils.py posthog/test/test_size_limited_dict.py --timeout=30 --cov=posthog.utils --crap --crap-threshold=10 --crap-top-n=40 -q + UV_PROJECT_ENVIRONMENT=$pythonLocation uv run --extra test --with pytest-crap --with pytest-cov pytest .github/scripts/test_check_crap_threshold.py --timeout=30 -q + UV_PROJECT_ENVIRONMENT=$pythonLocation uv run --extra test --with pytest-crap --with pytest-cov pytest posthog/test/test_utils.py posthog/test/test_size_limited_dict.py --timeout=30 --cov=posthog.utils -q UV_PROJECT_ENVIRONMENT=$pythonLocation uv run --extra test --with pytest-crap --with pytest-cov python .github/scripts/check_crap_threshold.py posthog/utils.py --max-crap 10 import-check: diff --git a/examples/example-ai-crewai/uv.lock b/examples/example-ai-crewai/uv.lock index 2078e56c2..6ec99c2ae 100644 --- a/examples/example-ai-crewai/uv.lock +++ b/examples/example-ai-crewai/uv.lock @@ -2372,14 +2372,14 @@ wheels = [ [[package]] name = "opentelemetry-api" -version = "1.45.1" +version = "1.45.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/2e/02/6e0ae9cc61bd3169d401077b507b3ebc344745171e1051ab430be012dcd9/opentelemetry_api-1.45.1.tar.gz", hash = "sha256:aa38ed19bcc084ba42782a73255b3582283eced7ad6dddbd6695189e69adfb75", size = 72804, upload-time = "2026-10-06T17:32:58.133Z" } +sdist = { url = "https://files.pythonhosted.org/packages/19/14/cf3a6b6448849656e1e5f4e2101f3988dc6789724050096aba1bc5ccf592/opentelemetry_api-1.45.2.tar.gz", hash = "sha256:d7b0008e21e5fd61cbe7a7f34ca3b6905a78ab677d9bc395e299359c75cda6b4", size = 73277, upload-time = "2026-10-29T17:29:37.888Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/1e/41/f7dcf80b81ee8e71c1a2b59f14208bc723edbd89ed027a73b175abf6348e/opentelemetry_api-1.45.1-py3-none-any.whl", hash = "sha256:b31553efa588ae44bc306f863c785c5333a9ecc091248c6ee68b4b6c87fdedfb", size = 60256, upload-time = "2026-10-06T17:32:33.506Z" }, + { url = "https://files.pythonhosted.org/packages/a9/2d/610db864d7513158cebb8f36cc88e67e73bd6c7aa1905ef18d89e1281c81/opentelemetry_api-1.45.2-py3-none-any.whl", hash = "sha256:2cbd737377e710ee740d261c94e93fa0f08e6d037c18f17513cc0a2d94479a34", size = 60335, upload-time = "2026-10-29T17:29:06.955Z" }, ] [[package]] @@ -2476,29 +2476,29 @@ wheels = [ [[package]] name = "opentelemetry-sdk" -version = "1.45.1" +version = "1.45.2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "opentelemetry-semantic-conventions" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a1/79/7392e21a1c8f0c61d90b223e31c7e48cb9d452e91a6b820ad24cca5f23c4/opentelemetry_sdk-1.45.1.tar.gz", hash = "sha256:63d24a6ca645019a631e6a51999c73e93adcac1196ca640b8ae78a7cc4762bf3", size = 218324, upload-time = "2026-10-06T17:33:13.26Z" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/1a/74a21d4d6969fb8b432ac397e66112a8167a6deb267ea4559dbaf2fe91ef/opentelemetry_sdk-1.45.2.tar.gz", hash = "sha256:7eec4de13dd21ff3a1670dcb0e02079e692736f94d733dfdce5af461d884248d", size = 221705, upload-time = "2026-10-29T17:29:54.685Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/95/3c/87c42b4bd6dd297536f04cd9383d212ac557ecd49f2cbdcd46da1c9ef5c8/opentelemetry_sdk-1.45.1-py3-none-any.whl", hash = "sha256:c604c11dc429810812348989115fa44bd558772a3d7442afc43d024f2c250ca4", size = 140063, upload-time = "2026-10-06T17:32:55.04Z" }, + { url = "https://files.pythonhosted.org/packages/9f/a5/9a3b793ef48245f2336a16beb2324c2397dbb3f987cd68942f77cca06156/opentelemetry_sdk-1.45.2-py3-none-any.whl", hash = "sha256:ff974ea70ee4f1d51e2ec152db37c2b626c7bf82bb3f7d86114fd8704e7199c6", size = 141756, upload-time = "2026-10-29T17:29:33.722Z" }, ] [[package]] name = "opentelemetry-semantic-conventions" -version = "0.66b1" +version = "0.66b2" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "opentelemetry-api" }, { name = "typing-extensions" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/46/e4/dbbfb2a010c4db2224a5114638acede6fe563d33cc20fb1752cebcbe6298/opentelemetry_semantic_conventions-0.66b1.tar.gz", hash = "sha256:497ca63bf383723411e8eaf60c8779e9877633c936bb641080adab59d0eb6ec8", size = 150250, upload-time = "2026-10-06T17:33:14.073Z" } +sdist = { url = "https://files.pythonhosted.org/packages/37/3f/50f95972b50e8e353f5b9342034c9cd17730f1956f009ce6a3fbbd89f9ce/opentelemetry_semantic_conventions-0.66b2.tar.gz", hash = "sha256:7da5f1a5dc4383eb66711d4d3d3e71824d9403ac1559fea021cc103c1a4c4976", size = 150246, upload-time = "2026-10-29T17:29:55.682Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bc/14/67f8aa798857f8cf686f515bf93d9bb877ce952ddc8efae0fa25b45ce0d6/opentelemetry_semantic_conventions-0.66b1-py3-none-any.whl", hash = "sha256:d4cddeb4315490b35213f55e2bdc9ac54bb1e4d318927475bed62b35545e581b", size = 206279, upload-time = "2026-10-06T17:32:56.103Z" }, + { url = "https://files.pythonhosted.org/packages/67/bd/6bf33ed31521c565931c037e9e81bac47c1acb000e850c6e24220e2a0c00/opentelemetry_semantic_conventions-0.66b2-py3-none-any.whl", hash = "sha256:4cbe61e5f9275d57ab31708446e7938ef4291ebbe6a7fca6f7a629c5acae87d8", size = 206280, upload-time = "2026-10-29T17:29:35.062Z" }, ] [[package]] diff --git a/examples/example-ai-langgraph/pyproject.toml b/examples/example-ai-langgraph/pyproject.toml index 20ec5ef04..b43e11d96 100644 --- a/examples/example-ai-langgraph/pyproject.toml +++ b/examples/example-ai-langgraph/pyproject.toml @@ -3,7 +3,7 @@ name = "example-ai-langgraph" version = "0.1.0" requires-python = ">=3.10" dependencies = [ - "langgraph>=0.2.0", + "langgraph>=1.2.14", "langchain-openai>=1.1.14", "opentelemetry-instrumentation-langchain>=0.24.0", "opentelemetry-sdk>=1.30.0", diff --git a/examples/example-ai-langgraph/uv.lock b/examples/example-ai-langgraph/uv.lock index 11180b80f..a717e0111 100644 --- a/examples/example-ai-langgraph/uv.lock +++ b/examples/example-ai-langgraph/uv.lock @@ -192,7 +192,7 @@ dependencies = [ [package.metadata] requires-dist = [ { name = "langchain-openai", specifier = ">=1.1.14" }, - { name = "langgraph", specifier = ">=0.2.0" }, + { name = "langgraph", specifier = ">=1.2.14" }, { name = "opentelemetry-instrumentation-langchain", specifier = ">=0.24.0" }, { name = "opentelemetry-sdk", specifier = ">=1.30.0" }, { name = "posthog", extras = ["otel"], directory = "../../" }, @@ -434,7 +434,7 @@ wheels = [ [[package]] name = "langgraph" -version = "1.2.12" +version = "1.2.14" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, @@ -444,9 +444,9 @@ dependencies = [ { name = "pydantic" }, { name = "xxhash" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/99/7e/c9e4431adb5eb78d15f0035aae162a113c80d037ceb433ddd78835217848/langgraph-1.2.12.tar.gz", hash = "sha256:5fae9fa4ce771e90f10ff19ce0d5aaef4b41d7632c713d769e89cbce1241a4b9", size = 734541, upload-time = "2026-09-21T14:43:29.684Z" } +sdist = { url = "https://files.pythonhosted.org/packages/49/3f/6e60789eab2c5ceaf34ba2c626ac345df93b026bbe7385018aa4b9c30397/langgraph-1.2.14.tar.gz", hash = "sha256:5b1a72b7a7f994ac6947b7751fbbe00e949f8bb44d9ab78881604b703247ae55", size = 752793, upload-time = "2026-10-06T14:41:19.789Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/52/27/5a02054c535ed6fcfb2ab1e5c33c51c810cf1e773194ad28c928f0473509/langgraph-1.2.12-py3-none-any.whl", hash = "sha256:95403af7b510de8d79164742f71daaf866c69ca804a8cb72bef2cc1fa1ab9813", size = 250206, upload-time = "2026-09-21T14:43:28.196Z" }, + { url = "https://files.pythonhosted.org/packages/ec/4a/f976df0d3cf7fb7986dc39f656c6b7b5655834430e9bdfaef1e1e4fffaba/langgraph-1.2.14-py3-none-any.whl", hash = "sha256:e3f5f8ff4259f9281b73ff50fc9ef5cc3b025b8936b8283ecdca7cc0c4128690", size = 255980, upload-time = "2026-10-06T14:41:18.129Z" }, ] [[package]] @@ -834,7 +834,7 @@ wheels = [ [[package]] name = "posthog" -version = "7.86.6" +version = "7.87.2" source = { directory = "../../" } dependencies = [ { name = "backoff" }, @@ -867,6 +867,9 @@ requires-dist = [ { name = "gevent", marker = "implementation_name == 'cpython' and extra == 'test'", specifier = ">=25.4.1" }, { name = "google-genai", marker = "extra == 'test'" }, { name = "griffe", marker = "extra == 'dev'" }, + { name = "griffe", marker = "extra == 'test'" }, + { name = "griffe2md", marker = "extra == 'dev'" }, + { name = "griffe2md", marker = "extra == 'test'" }, { name = "httpx", marker = "extra == 'async'", specifier = ">=0.27.0,<1.0" }, { name = "jsonschema", marker = "extra == 'test'", specifier = ">=4.0" }, { name = "langchain", marker = "extra == 'langchain'", specifier = ">=1.3.9" }, diff --git a/posthog/__init__.py b/posthog/__init__.py index 705c4e5b7..3b5f1f876 100644 --- a/posthog/__init__.py +++ b/posthog/__init__.py @@ -1476,11 +1476,7 @@ def setup() -> Client: default_client.debug = debug default_client.privacy_mode = bool(privacy_mode) default_client._set_before_send(before_send) - default_client.enable_full_ai_capture = ( - bool(enable_full_ai_capture) - or bool(_use_ai_lane) - or bool(_enable_multimodal_capture) - ) + default_client.enable_full_ai_capture = _full_ai_capture_enabled() # Metrics config is consumed lazily on first `.metrics` access, so late # module-attr assignment (e.g. a Django ready() hook running after something # already forced setup()) still applies until the metrics API is first used. @@ -1496,6 +1492,14 @@ def setup() -> Client: return default_client +def _full_ai_capture_enabled() -> bool: + return ( + bool(enable_full_ai_capture) + or bool(_use_ai_lane) + or bool(_enable_multimodal_capture) + ) + + def _proxy(method, *args, **kwargs): """Create an analytics client if one doesn't exist and send to it.""" setup() diff --git a/posthog/_async_consumer.py b/posthog/_async_consumer.py index 5bdfd33f1..87f86f652 100644 --- a/posthog/_async_consumer.py +++ b/posthog/_async_consumer.py @@ -165,17 +165,58 @@ async def upload(self, batch: list[dict[str, Any]]) -> None: for _ in batch: self.queue.task_done() - async def next(self) -> tuple[list[dict[str, Any]], bool]: - items: list[dict[str, Any]] = [] - total_size = 0 - stop = False - started = asyncio.get_running_loop().time() + async def _prepare_queued_event( + self, queued: _QueuedEvent + ) -> Optional[tuple[dict[str, Any], int]]: + try: + process_task = queued.context.run( + asyncio.create_task, self._process_queued_event(queued.event) + ) + item = await process_task + except Exception as error: + self.log.error( + "unable to process queued event, dropping (%s)", + type(error).__name__, + ) + self.queue.task_done() + return None + if item is None: + self.queue.task_done() + return None + + try: + item_size = await _serialized_event_size(item) + except Exception: + self.log.error("unable to serialize queued event for sizing, dropping") + self.queue.task_done() + return None + + if item_size > MAX_MSG_SIZE: + self.log.error( + "Event %s (%d bytes) exceeds the %dKiB limit, dropping.", + item.get("event"), + item_size, + MAX_MSG_SIZE // 1024, + ) + self.queue.task_done() + return None + return item, item_size + + def _take_carryover(self, items: list[dict[str, Any]]) -> int: if self._carryover is not None: carried_item, carried_size = self._carryover self._carryover = None items.append(carried_item) - total_size = carried_size + return carried_size + return 0 + + async def next(self) -> tuple[list[dict[str, Any]], bool]: + items: list[dict[str, Any]] = [] + stop = False + started = asyncio.get_running_loop().time() + + total_size = self._take_carryover(items) while len(items) < self.flush_at: remaining = self.flush_interval - ( @@ -193,39 +234,10 @@ async def next(self) -> tuple[list[dict[str, Any]], bool]: stop = True break - try: - process_task = queued.context.run( - asyncio.create_task, self._process_queued_event(queued.event) - ) - item = await process_task - except Exception as error: - self.log.error( - "unable to process queued event, dropping (%s)", - type(error).__name__, - ) - self.queue.task_done() - continue - - if item is None: - self.queue.task_done() - continue - - try: - item_size = await _serialized_event_size(item) - except Exception: - self.log.error("unable to serialize queued event for sizing, dropping") - self.queue.task_done() - continue - - if item_size > MAX_MSG_SIZE: - self.log.error( - "Event %s (%d bytes) exceeds the %dKiB limit, dropping.", - item.get("event"), - item_size, - MAX_MSG_SIZE // 1024, - ) - self.queue.task_done() + prepared = await self._prepare_queued_event(queued) + if prepared is None: continue + item, item_size = prepared if items and total_size + item_size > BATCH_SIZE_LIMIT: self._carryover = (item, item_size) diff --git a/posthog/_async_request.py b/posthog/_async_request.py index 50dfe829c..ff407687f 100644 --- a/posthog/_async_request.py +++ b/posthog/_async_request.py @@ -236,18 +236,9 @@ async def async_batch_post( _process_response(response) return - location = response.headers.get("Location") or response.headers.get( - "location" + redirect_url = _capture_redirect_url( + base_url, request_url, response, redirect_count ) - redirect_url = ( - _same_origin_redirect_url(base_url, request_url, location) - if location - else None - ) - if redirect_url is None: - raise APIError(400, "Cross-origin or invalid redirect blocked") - if redirect_count >= 5: - raise APIError(400, "Too many capture redirects") request_url = redirect_url finally: if owns_client: @@ -275,3 +266,17 @@ async def async_send_v1_batch( max_retries=max_retries, historical_migration=historical_migration, ) + + +def _capture_redirect_url( + base_url: str, request_url: str, response: Any, redirect_count: int +) -> str: + location = response.headers.get("Location") or response.headers.get("location") + redirect_url = ( + _same_origin_redirect_url(base_url, request_url, location) if location else None + ) + if redirect_url is None: + raise APIError(400, "Cross-origin or invalid redirect blocked") + if redirect_count >= 5: + raise APIError(400, "Too many capture redirects") + return redirect_url diff --git a/posthog/_async_utils.py b/posthog/_async_utils.py index f81d6964b..b62cf19ba 100644 --- a/posthog/_async_utils.py +++ b/posthog/_async_utils.py @@ -139,13 +139,9 @@ async def _await_result(awaitable: Awaitable[Any]) -> Any: def _ensure_loop(self) -> asyncio.AbstractEventLoop: with self._lock: - if ( - self._loop is not None - and self._thread is not None - and self._thread.is_alive() - and not self._loop.is_closed() - ): - return self._loop + ready_loop = self._ready_loop() + if ready_loop is not None: + return ready_loop startup: _LoopStartup if self._thread is None or not self._thread.is_alive(): @@ -176,30 +172,10 @@ def _run_loop(self, startup: _LoopStartup) -> None: loop = None try: loop = asyncio.new_event_loop() - original_run_in_executor = loop.run_in_executor - - def run_in_executor(executor, func, *args): - call = _ContextExecutorCall(copy_context(), func, args) - return original_run_in_executor(executor, call) - - try: - setattr(loop, "run_in_executor", run_in_executor) - except (AttributeError, TypeError): - # A loop that cannot carry callback context is unsafe for - # lifecycle re-entry from executor threads. Use the equivalent - # context-aware platform loop rather than silently losing it. - loop.close() - loop = _ContextEventLoop() + loop = self._install_context_executor(loop) asyncio.set_event_loop(loop) except BaseException as error: - if loop is not None and not loop.is_closed(): - loop.close() - with self._lock: - if startup.error is None: - startup.error = error - if self._thread is threading.current_thread(): - self._thread = None - startup.done.set() + self._fail_startup(startup, loop, error) return with self._lock: @@ -216,17 +192,7 @@ def run_in_executor(executor, func, *args): try: loop.run_forever() finally: - pending = asyncio.all_tasks(loop) - for task in pending: - task.cancel() - if pending: - loop.run_until_complete( - asyncio.gather(*pending, return_exceptions=True) - ) - loop.run_until_complete(loop.shutdown_asyncgens()) - loop.run_until_complete(loop.shutdown_default_executor()) - asyncio.set_event_loop(None) - loop.close() + self._close_loop(loop) current = threading.current_thread() with self._lock: if self._thread is current: @@ -236,3 +202,60 @@ def run_in_executor(executor, func, *args): if self._startup is startup: self._startup = None self._closing_threads.discard(current) + + def _ready_loop(self) -> asyncio.AbstractEventLoop | None: + if ( + self._loop is not None + and self._thread is not None + and self._thread.is_alive() + and not self._loop.is_closed() + ): + return self._loop + return None + + @staticmethod + def _close_loop(loop: asyncio.AbstractEventLoop) -> None: + pending = asyncio.all_tasks(loop) + for task in pending: + task.cancel() + if pending: + loop.run_until_complete(asyncio.gather(*pending, return_exceptions=True)) + loop.run_until_complete(loop.shutdown_asyncgens()) + loop.run_until_complete(loop.shutdown_default_executor()) + asyncio.set_event_loop(None) + loop.close() + + @staticmethod + def _install_context_executor( + loop: asyncio.AbstractEventLoop, + ) -> asyncio.AbstractEventLoop: + original_run_in_executor = loop.run_in_executor + + def run_in_executor(executor, func, *args): + call = _ContextExecutorCall(copy_context(), func, args) + return original_run_in_executor(executor, call) + + try: + setattr(loop, "run_in_executor", run_in_executor) + except (AttributeError, TypeError): + # A loop that cannot carry callback context is unsafe for + # lifecycle re-entry from executor threads. Use the equivalent + # context-aware platform loop rather than silently losing it. + loop.close() + loop = _ContextEventLoop() + return loop + + def _fail_startup( + self, + startup: _LoopStartup, + loop: asyncio.AbstractEventLoop | None, + error: BaseException, + ) -> None: + if loop is not None and not loop.is_closed(): + loop.close() + with self._lock: + if startup.error is None: + startup.error = error + if self._thread is threading.current_thread(): + self._thread = None + startup.done.set() diff --git a/posthog/_metrics_autocapture.py b/posthog/_metrics_autocapture.py index 4771f2d5a..73fde54bf 100644 --- a/posthog/_metrics_autocapture.py +++ b/posthog/_metrics_autocapture.py @@ -230,11 +230,7 @@ def _clean_attributes(attributes: Any) -> Any: return None cleaned: Dict[str, Any] = {} for key, value in items: - if ( - not isinstance(key, str) - or key in _UNSAFE_ATTRIBUTES - or key.startswith(_UNSAFE_ATTRIBUTE_PREFIXES) - ): + if not _safe_attribute_key(key): continue if isinstance(value, _ATTRIBUTE_VALUE_TYPES): cleaned[key] = value @@ -730,7 +726,7 @@ def on_end(self, span): ) def _record(self, span): - from opentelemetry.trace import SpanKind, StatusCode + from opentelemetry.trace import SpanKind scope = getattr(span, "instrumentation_scope", None) if ( @@ -742,42 +738,7 @@ def _record(self, span): system = _first_string(attributes, ("db.system.name", "db.system")) if not system or span.start_time is None or span.end_time is None: return - metric_attributes: Dict[str, Any] = {"db.system.name": system} - operation = _first_string(attributes, ("db.operation.name", "db.operation")) - if operation is None: - statement = _first_string(attributes, ("db.query.text", "db.statement")) - match = _STATEMENT_KEYWORD.match(statement or "") - operation = match.group(1).upper() if match else None - else: - # Only the leading keyword is bounded, as for a statement. - words = operation.split() - operation = words[0] if words else None - collection = _first_string( - attributes, - ( - "db.collection.name", - "db.mongodb.collection", - "db.sql.table", - "db.cassandra.table", - ), - ) - namespace = _first_string(attributes, ("db.namespace", "db.name")) - address = _first_string(attributes, ("server.address", "net.peer.name")) - port = attributes.get("server.port", attributes.get("net.peer.port")) - if operation: - metric_attributes["db.operation.name"] = operation - if collection: - metric_attributes["db.collection.name"] = collection - if namespace: - metric_attributes["db.namespace"] = namespace - if address: - metric_attributes["server.address"] = address - if isinstance(port, int) and not isinstance(port, bool): - metric_attributes["server.port"] = port - if span.status.status_code == StatusCode.ERROR: - metric_attributes["error.type"] = ( - _first_string(attributes, ("error.type",)) or "_OTHER" - ) + metric_attributes = _db_metric_attributes(span, attributes, system) self._histogram.record( (span.end_time - span.start_time) / 1e9, metric_attributes ) @@ -1103,73 +1064,23 @@ def start_metrics_autocapture( instrumentors: List[Any] = [] stopped = threading.Event() try: - if not _is_default_provider( - metrics_api.get_meter_provider() - ) or not _is_default_provider(trace_api.get_tracer_provider()): - log.warning( - "OpenTelemetry is already set up in this process, so metrics autocapture does not start. " - "Point your OTLP metrics exporter at PostHog instead: https://posthog.com/docs/metrics" - ) + if not _autocapture_providers_available(metrics_api, trace_api): return None - # The same resource as `client.metrics`: the OTEL_SERVICE_NAME and - # OTEL_RESOURCE_ATTRIBUTES of the app are not read. - resource = sdk_resources.Resource( - { - **resource_attributes, - "service.name": service_name or "unknown_service", - "telemetry.sdk.language": "python", - "telemetry.sdk.name": "opentelemetry", - "telemetry.sdk.version": _import_module( - "opentelemetry.sdk.version" - ).__version__, - "telemetry.distro.name": "posthog-python", - "telemetry.distro.version": VERSION, - } + resource = _autocapture_resource( + sdk_resources, service_name, resource_attributes ) reader = _metric_reader if reader is None: - exporter_base: Any = otlp.OTLPMetricExporter - delta = sdk_metrics_export.AggregationTemporality.DELTA - cumulative = sdk_metrics_export.AggregationTemporality.CUMULATIVE - - class _GatedExporter(exporter_base): - # The exporter still sends while the client is disabled unless it is gated here. - def export(self, metrics_data, timeout_millis=10_000, **kwargs): - if not is_enabled(): - return sdk_metrics_export.MetricExportResult.SUCCESS - return super().export( - metrics_data, timeout_millis=timeout_millis, **kwargs - ) - - exporter = _GatedExporter( - endpoint=f"{host}/i/v1/metrics", - headers={"Authorization": f"Bearer {api_key}"}, - timeout=_EXPORT_TIMEOUT_SECONDS, - compression=_import_module( - "opentelemetry.exporter.otlp.proto.http" - ).Compression.NoCompression, - session=_metrics_export_session(api_key), - # Delta, like the `posthog.metrics` wire shape; values that go - # up and down stay cumulative. - preferred_temporality={ - sdk_metrics.Counter: delta, - sdk_metrics.Histogram: delta, - sdk_metrics.ObservableCounter: delta, - sdk_metrics.UpDownCounter: cumulative, - sdk_metrics.ObservableUpDownCounter: cumulative, - sdk_metrics.ObservableGauge: cumulative, - }, - preferred_aggregation={ - sdk_metrics.Histogram: _import_module( - "opentelemetry.sdk.metrics.view" - ).ExplicitBucketHistogramAggregation(), - }, - ) - reader = sdk_metrics_export.PeriodicExportingMetricReader( - exporter, - export_interval_millis=(export_interval_seconds or 10.0) * 1000, + reader = _autocapture_metric_reader( + otlp, + sdk_metrics, + sdk_metrics_export, + host, + api_key, + is_enabled, + export_interval_seconds, ) # The client exports the last window at exit, within its exit budget. # The provider's own exit hook waits up to the export timeout, also @@ -1225,79 +1136,15 @@ def get_description(self): ) ) - posthog_url_pattern = _posthog_url_pattern(host) - wanted = set(_instrumentor_names(areas)) - with _stable_http_semconv(): - for entry_point in entry_points(group="opentelemetry_instrumentor"): - if entry_point.name not in wanted: - continue - if ( - entry_point.name == "django" - and not _django_settings_configured() - ): - # Without settings, the Django instrumentor configures empty - # ones, and the app's own settings.configure() then fails. - # In AppConfig.ready() the settings are configured. - log.debug( - "Metrics autocapture skipped django: its settings are not configured yet" - ) - continue - try: - kwargs: Dict[str, Any] = { - "tracer_provider": tracer_provider - if entry_point.name in _DB_INSTRUMENTORS - else no_op_tracer_provider, - "meter_provider": bounded_meter_provider, - } - environ: Dict[str, Optional[str]] = {} - if entry_point.name in _HTTP_CLIENT_INSTRUMENTORS: - # PostHog's own uploads, including this exporter's, are not app traffic. - variable, excluded = _excluded_urls( - entry_point.name, posthog_url_pattern - ) - kwargs["excluded_urls"] = excluded - environ[variable] = excluded - with _patched_environ(environ): - instrumentor_class = entry_point.load() - # Instrumentors are singletons, and creating one runs - # `__init__` again on the existing instance. That would - # reset the config of an instance that the app uses. - existing = getattr(instrumentor_class, "_instance", None) - if ( - existing is not None - and existing.is_instrumented_by_opentelemetry - ): - continue - instrumentor = ( - instrumentor_class(config=_RUNTIME_CONFIG) - if entry_point.name == "system_metrics" - else instrumentor_class() - ) - if instrumentor.is_instrumented_by_opentelemetry: - # The app instrumented this library itself. - continue - instrumentor.instrument(**kwargs) - if instrumentor.is_instrumented_by_opentelemetry: - instrumentors.append(instrumentor) - _disable_propagation(instrumentor_class) - if entry_point.name == "aiohttp-client": - _gate_aiohttp_trace_configs( - bounded_meter_provider, stopped, posthog_url_pattern - ) - if entry_point.name == "flask": - _track_flask_apps( - instrumentor_class, - { - "tracer_provider": no_op_tracer_provider, - "meter_provider": bounded_meter_provider, - }, - ) - except Exception: - log.debug( - "Metrics autocapture skipped %s", - entry_point.name, - exc_info=True, - ) + _instrument_metrics_libraries( + instrumentors, + areas, + host, + tracer_provider, + no_op_tracer_provider, + bounded_meter_provider, + stopped, + ) _active = MetricsAutocapture( meter_provider, tracer_provider, instrumentors, reader, stopped @@ -1310,3 +1157,247 @@ def get_description(self): _uninstrument(instrumentors) _shut_down(tracer_provider, meter_provider or reader) return None + + +def _safe_attribute_key(key: Any) -> bool: + return ( + isinstance(key, str) + and key not in _UNSAFE_ATTRIBUTES + and not key.startswith(_UNSAFE_ATTRIBUTE_PREFIXES) + ) + + +def _db_operation(attributes: Any) -> Optional[str]: + operation = _first_string(attributes, ("db.operation.name", "db.operation")) + if operation is None: + statement = _first_string(attributes, ("db.query.text", "db.statement")) + match = _STATEMENT_KEYWORD.match(statement or "") + operation = match.group(1).upper() if match else None + else: + # Only the leading keyword is bounded, as for a statement. + words = operation.split() + operation = words[0] if words else None + return operation + + +def _db_metric_attributes(span: Any, attributes: Any, system: str) -> Dict[str, Any]: + from opentelemetry.trace import StatusCode + + metric_attributes: Dict[str, Any] = {"db.system.name": system} + operation = _db_operation(attributes) + collection = _first_string( + attributes, + ( + "db.collection.name", + "db.mongodb.collection", + "db.sql.table", + "db.cassandra.table", + ), + ) + namespace = _first_string(attributes, ("db.namespace", "db.name")) + address = _first_string(attributes, ("server.address", "net.peer.name")) + port = attributes.get("server.port", attributes.get("net.peer.port")) + if operation: + metric_attributes["db.operation.name"] = operation + if collection: + metric_attributes["db.collection.name"] = collection + if namespace: + metric_attributes["db.namespace"] = namespace + if address: + metric_attributes["server.address"] = address + if isinstance(port, int) and not isinstance(port, bool): + metric_attributes["server.port"] = port + if span.status.status_code == StatusCode.ERROR: + metric_attributes["error.type"] = ( + _first_string(attributes, ("error.type",)) or "_OTHER" + ) + return metric_attributes + + +def _instrument_metrics_libraries( + instrumentors: List[Any], + areas: Dict[str, bool], + host: str, + tracer_provider: Any, + no_op_tracer_provider: Any, + bounded_meter_provider: Any, + stopped: threading.Event, +) -> None: + posthog_url_pattern = _posthog_url_pattern(host) + wanted = set(_instrumentor_names(areas)) + with _stable_http_semconv(): + for entry_point in entry_points(group="opentelemetry_instrumentor"): + if entry_point.name not in wanted: + continue + if entry_point.name == "django" and not _django_settings_configured(): + # Without settings, the Django instrumentor configures empty + # ones, and the app's own settings.configure() then fails. + # In AppConfig.ready() the settings are configured. + log.debug( + "Metrics autocapture skipped django: its settings are not configured yet" + ) + continue + try: + _instrument_metric_library( + entry_point, + instrumentors, + tracer_provider, + no_op_tracer_provider, + bounded_meter_provider, + stopped, + posthog_url_pattern, + ) + except Exception: + log.debug( + "Metrics autocapture skipped %s", + entry_point.name, + exc_info=True, + ) + + +def _instrument_metric_library( + entry_point: Any, + instrumentors: List[Any], + tracer_provider: Any, + no_op_tracer_provider: Any, + bounded_meter_provider: Any, + stopped: threading.Event, + posthog_url_pattern: str, +) -> None: + kwargs: Dict[str, Any] = { + "tracer_provider": tracer_provider + if entry_point.name in _DB_INSTRUMENTORS + else no_op_tracer_provider, + "meter_provider": bounded_meter_provider, + } + environ: Dict[str, Optional[str]] = {} + if entry_point.name in _HTTP_CLIENT_INSTRUMENTORS: + # PostHog's own uploads, including this exporter's, are not app traffic. + variable, excluded = _excluded_urls(entry_point.name, posthog_url_pattern) + kwargs["excluded_urls"] = excluded + environ[variable] = excluded + with _patched_environ(environ): + instrumentor_class = entry_point.load() + instrumentor = _new_metric_instrumentor(instrumentor_class, entry_point.name) + if instrumentor is None: + return + instrumentor.instrument(**kwargs) + if instrumentor.is_instrumented_by_opentelemetry: + instrumentors.append(instrumentor) + _disable_propagation(instrumentor_class) + if entry_point.name == "aiohttp-client": + _gate_aiohttp_trace_configs( + bounded_meter_provider, stopped, posthog_url_pattern + ) + if entry_point.name == "flask": + _track_flask_apps( + instrumentor_class, + { + "tracer_provider": no_op_tracer_provider, + "meter_provider": bounded_meter_provider, + }, + ) + + +def _new_metric_instrumentor(instrumentor_class: Any, name: str) -> Any: + # Instrumentors are singletons, and creating one runs + # `__init__` again on the existing instance. That would + # reset the config of an instance that the app uses. + existing = getattr(instrumentor_class, "_instance", None) + if existing is not None and existing.is_instrumented_by_opentelemetry: + return None + instrumentor = ( + instrumentor_class(config=_RUNTIME_CONFIG) + if name == "system_metrics" + else instrumentor_class() + ) + if instrumentor.is_instrumented_by_opentelemetry: + # The app instrumented this library itself. + return None + return instrumentor + + +def _autocapture_metric_reader( + otlp: Any, + sdk_metrics: Any, + sdk_metrics_export: Any, + host: str, + api_key: str, + is_enabled: Callable[[], bool], + export_interval_seconds: Optional[float], +) -> Any: + exporter_base: Any = otlp.OTLPMetricExporter + delta = sdk_metrics_export.AggregationTemporality.DELTA + cumulative = sdk_metrics_export.AggregationTemporality.CUMULATIVE + + class _GatedExporter(exporter_base): + # The exporter still sends while the client is disabled unless it is gated here. + def export(self, metrics_data, timeout_millis=10_000, **kwargs): + if not is_enabled(): + return sdk_metrics_export.MetricExportResult.SUCCESS + return super().export(metrics_data, timeout_millis=timeout_millis, **kwargs) + + exporter = _GatedExporter( + endpoint=f"{host}/i/v1/metrics", + headers={"Authorization": f"Bearer {api_key}"}, + timeout=_EXPORT_TIMEOUT_SECONDS, + compression=_import_module( + "opentelemetry.exporter.otlp.proto.http" + ).Compression.NoCompression, + session=_metrics_export_session(api_key), + # Delta, like the `posthog.metrics` wire shape; values that go + # up and down stay cumulative. + preferred_temporality={ + sdk_metrics.Counter: delta, + sdk_metrics.Histogram: delta, + sdk_metrics.ObservableCounter: delta, + sdk_metrics.UpDownCounter: cumulative, + sdk_metrics.ObservableUpDownCounter: cumulative, + sdk_metrics.ObservableGauge: cumulative, + }, + preferred_aggregation={ + sdk_metrics.Histogram: _import_module( + "opentelemetry.sdk.metrics.view" + ).ExplicitBucketHistogramAggregation(), + }, + ) + reader = sdk_metrics_export.PeriodicExportingMetricReader( + exporter, + export_interval_millis=(export_interval_seconds or 10.0) * 1000, + ) + return reader + + +def _autocapture_resource( + sdk_resources: Any, service_name: Optional[str], resource_attributes: Dict[str, Any] +) -> Any: + # The same resource as `client.metrics`: the OTEL_SERVICE_NAME and + # OTEL_RESOURCE_ATTRIBUTES of the app are not read. + resource = sdk_resources.Resource( + { + **resource_attributes, + "service.name": service_name or "unknown_service", + "telemetry.sdk.language": "python", + "telemetry.sdk.name": "opentelemetry", + "telemetry.sdk.version": _import_module( + "opentelemetry.sdk.version" + ).__version__, + "telemetry.distro.name": "posthog-python", + "telemetry.distro.version": VERSION, + } + ) + + return resource + + +def _autocapture_providers_available(metrics_api: Any, trace_api: Any) -> bool: + if not _is_default_provider( + metrics_api.get_meter_provider() + ) or not _is_default_provider(trace_api.get_tracer_provider()): + log.warning( + "OpenTelemetry is already set up in this process, so metrics autocapture does not start. " + "Point your OTLP metrics exporter at PostHog instead: https://posthog.com/docs/metrics" + ) + return False + + return True diff --git a/posthog/ai/anthropic/_anthropic_stream.py b/posthog/ai/anthropic/_anthropic_stream.py index 7c76cc369..f02ee9bf9 100644 --- a/posthog/ai/anthropic/_anthropic_stream.py +++ b/posthog/ai/anthropic/_anthropic_stream.py @@ -53,38 +53,10 @@ def consume(self, event: Any) -> None: event_usage = extract_anthropic_usage_from_event(event) merge_usage_stats(self.usage_stats, event_usage) if getattr(event, "type", None) == "message_delta": - usage = getattr(event, "usage", None) - self._apply_cumulative_input_usage(usage) - output_tokens = getattr(usage, "output_tokens", None) - if isinstance(output_tokens, int) and not isinstance(output_tokens, bool): - # output_tokens on a message_delta is cumulative too, so the - # incremental merge above must not stack successive deltas - # (server tools can emit several in one stream). Keep the - # latest total; a field the event omits leaves the prior one. - self.usage_stats["output_tokens"] = output_tokens - # thinking tokens on a message_delta are cumulative as well - thinking_tokens = extract_anthropic_usage_from_event(event).get( - "reasoning_tokens" - ) - if thinking_tokens is not None: - self.usage_stats["reasoning_tokens"] = thinking_tokens + self._apply_delta_usage(event) if getattr(event, "type", None) == "content_block_start": - block, tool = handle_anthropic_content_block_start(event) - - self._blocks_by_index.append(block or {}) - - if block: - self.content_blocks.append(block) - if block.get("type") in ("text", "thinking"): - self.current_text_block = block - else: - self.current_text_block = None - - if tool: - tool_id = tool["block"].get("id") - if tool_id: - self.tools_in_progress[tool_id] = tool + self._start_block(event) delta_text = handle_anthropic_text_delta(event, self.current_text_block) if delta_text: @@ -105,3 +77,37 @@ def consume(self, event: Any) -> None: delta_stop_reason = getattr(delta, "stop_reason", None) if delta_stop_reason is not None: self.stop_reason = delta_stop_reason + + def _apply_delta_usage(self, event: Any) -> None: + usage = getattr(event, "usage", None) + self._apply_cumulative_input_usage(usage) + output_tokens = getattr(usage, "output_tokens", None) + if isinstance(output_tokens, int) and not isinstance(output_tokens, bool): + # output_tokens on a message_delta is cumulative too, so the + # incremental merge above must not stack successive deltas + # (server tools can emit several in one stream). Keep the + # latest total; a field the event omits leaves the prior one. + self.usage_stats["output_tokens"] = output_tokens + # thinking tokens on a message_delta are cumulative as well + thinking_tokens = extract_anthropic_usage_from_event(event).get( + "reasoning_tokens" + ) + if thinking_tokens is not None: + self.usage_stats["reasoning_tokens"] = thinking_tokens + + def _start_block(self, event: Any) -> None: + block, tool = handle_anthropic_content_block_start(event) + + self._blocks_by_index.append(block or {}) + + if block: + self.content_blocks.append(block) + if block.get("type") in ("text", "thinking"): + self.current_text_block = block + else: + self.current_text_block = None + + if tool: + tool_id = tool["block"].get("id") + if tool_id: + self.tools_in_progress[tool_id] = tool diff --git a/posthog/ai/anthropic/anthropic_converter.py b/posthog/ai/anthropic/anthropic_converter.py index 1ee525e72..b4734d206 100644 --- a/posthog/ai/anthropic/anthropic_converter.py +++ b/posthog/ai/anthropic/anthropic_converter.py @@ -6,7 +6,7 @@ """ import json -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Dict, List, Literal, Optional, Tuple from posthog.ai.media import to_plain from posthog.ai.types import ( @@ -42,57 +42,9 @@ def format_anthropic_response(response: Any) -> List[FormattedMessage]: # Process content blocks from the response if hasattr(response, "content"): for choice in response.content: - if ( - hasattr(choice, "type") - and choice.type == "text" - and hasattr(choice, "text") - and choice.text - ): - text_content: FormattedTextContent = { - "type": "text", - "text": choice.text, - } - content.append(text_content) - - elif ( - hasattr(choice, "type") - and choice.type == "tool_use" - and hasattr(choice, "name") - and hasattr(choice, "id") - ): - function_call: FormattedFunctionCall = { - "type": "function", - "id": choice.id, - "function": { - "name": choice.name, - "arguments": getattr(choice, "input", {}), - }, - } - content.append(function_call) - - elif getattr(choice, "type", None) == "thinking": - content.append( - { - "type": "thinking", - "thinking": getattr(choice, "thinking", None), - "signature": getattr(choice, "signature", None), - } - ) - - elif getattr(choice, "type", None) == "redacted_thinking": - content.append( - { - "type": "redacted_thinking", - "data": getattr(choice, "data", None), - } - ) - - else: - # Catches blocks the branches above skip on falsy values (e.g. TextBlock(text="")) - # so empty-but-valid content survives instead of being silently dropped. - plain = to_plain(choice) - if isinstance(plain, dict) and plain.get("type"): - content.append(plain) + block = _format_anthropic_response_block(choice) + if block is not None: + content.append(block) if content: message: FormattedMessage = { @@ -272,15 +224,7 @@ def extract_anthropic_usage_from_response(response: Any) -> TokenUsage: if output_tokens is not None: result["output_tokens"] = output_tokens - if hasattr(response.usage, "cache_read_input_tokens"): - cache_read = response.usage.cache_read_input_tokens - if cache_read and cache_read > 0: - result["cache_read_input_tokens"] = cache_read - - if hasattr(response.usage, "cache_creation_input_tokens"): - cache_creation = response.usage.cache_creation_input_tokens - if cache_creation and cache_creation > 0: - result["cache_creation_input_tokens"] = cache_creation + _add_anthropic_cache_usage(result, response.usage) thinking_tokens = _extract_thinking_tokens(response.usage) if thinking_tokens > 0: @@ -312,21 +256,7 @@ def extract_anthropic_usage_from_event(event: Any) -> TokenUsage: usage: TokenUsage = TokenUsage() - # Handle usage stats from message_start event - if hasattr(event, "type") and event.type == "message_start": - if hasattr(event, "message") and hasattr(event.message, "usage"): - usage["input_tokens"] = getattr(event.message.usage, "input_tokens", 0) - usage["cache_creation_input_tokens"] = getattr( - event.message.usage, "cache_creation_input_tokens", 0 - ) - usage["cache_read_input_tokens"] = getattr( - event.message.usage, "cache_read_input_tokens", 0 - ) - # Capture raw usage metadata for backend processing - # Serialize to dict here in the converter (not in utils) - serialized = serialize_raw_usage(event.message.usage) - if serialized: - usage["raw_usage"] = serialized + usage.update(_extract_anthropic_start_usage(event)) # Handle usage stats from message_delta event if hasattr(event, "usage") and event.usage: @@ -374,42 +304,7 @@ def handle_anthropic_content_block_start( if not hasattr(event, "content_block"): return None, None - block = event.content_block - - if not hasattr(block, "type"): - return None, None - - if block.type == "text": - content_block: StreamingContentBlock = {"type": "text", "text": ""} - return content_block, None - - elif block.type == "tool_use": - tool_block: StreamingContentBlock = { - "type": "function", - "id": getattr(block, "id", ""), - "function": {"name": getattr(block, "name", ""), "arguments": {}}, - } - tool_in_progress: ToolInProgress = {"block": tool_block, "input_string": ""} - return tool_block, tool_in_progress - - elif block.type == "thinking": - thinking_block: StreamingContentBlock = { - "type": "thinking", - "thinking": getattr(block, "thinking", "") or "", - } - signature = getattr(block, "signature", None) - if signature: - thinking_block["signature"] = signature - return thinking_block, None - - elif block.type == "redacted_thinking": - redacted_block: StreamingContentBlock = { - "type": "redacted_thinking", - "data": getattr(block, "data", None), - } - return redacted_block, None - - return None, None + return _start_anthropic_block(event.content_block) def handle_anthropic_text_delta( @@ -442,24 +337,7 @@ def handle_anthropic_text_delta( return delta_text - if hasattr(event, "delta") and hasattr(event.delta, "thinking"): - delta_thinking = event.delta.thinking or "" - - if current_block is not None and current_block.get("type") == "thinking": - thinking_val = current_block.get("thinking") - current_block["thinking"] = (thinking_val or "") + delta_thinking - - return None - - if hasattr(event, "delta") and hasattr(event.delta, "signature"): - delta_signature = event.delta.signature or "" - - if current_block is not None and current_block.get("type") == "thinking": - signature_val = current_block.get("signature") - current_block["signature"] = (signature_val or "") + delta_signature - - return None - + _accumulate_anthropic_thinking_delta(event, current_block) return None @@ -487,13 +365,7 @@ def handle_anthropic_tool_delta( ): return - if hasattr(event, "index") and event.index < len(content_blocks): - block = content_blocks[event.index] - - if block.get("type") == "function" and block.get("id") in tools_in_progress: - tool = tools_in_progress[block["id"]] - partial_json = getattr(event.delta, "partial_json", "") - tool["input_string"] += partial_json + _append_anthropic_tool_json(event, content_blocks, tools_in_progress) def finalize_anthropic_tool_input( @@ -570,3 +442,157 @@ def format_anthropic_streaming_output_complete( "content": [{"type": "text", "text": accumulated_content}], } ] + + +def _format_anthropic_response_block(choice: Any) -> Optional[FormattedContentItem]: + if ( + hasattr(choice, "type") + and choice.type == "text" + and hasattr(choice, "text") + and choice.text + ): + text_content: FormattedTextContent = { + "type": "text", + "text": choice.text, + } + return text_content + + elif ( + hasattr(choice, "type") + and choice.type == "tool_use" + and hasattr(choice, "name") + and hasattr(choice, "id") + ): + function_call: FormattedFunctionCall = { + "type": "function", + "id": choice.id, + "function": { + "name": choice.name, + "arguments": getattr(choice, "input", {}), + }, + } + return function_call + + return _format_anthropic_reasoning_block(choice) + + +def _format_anthropic_reasoning_block(choice: Any) -> Optional[FormattedContentItem]: + if getattr(choice, "type", None) == "thinking": + return { + "type": "thinking", + "thinking": getattr(choice, "thinking", None), + "signature": getattr(choice, "signature", None), + } + + elif getattr(choice, "type", None) == "redacted_thinking": + return { + "type": "redacted_thinking", + "data": getattr(choice, "data", None), + } + + else: + # Catches blocks the branches above skip on falsy values (e.g. TextBlock(text="")) + # so empty-but-valid content survives instead of being silently dropped. + plain = to_plain(choice) + if isinstance(plain, dict) and plain.get("type"): + return plain + + return None + + +def _add_anthropic_cache_usage(result: TokenUsage, usage: Any) -> None: + if hasattr(usage, "cache_read_input_tokens"): + cache_read = usage.cache_read_input_tokens + if cache_read and cache_read > 0: + result["cache_read_input_tokens"] = cache_read + + if hasattr(usage, "cache_creation_input_tokens"): + cache_creation = usage.cache_creation_input_tokens + if cache_creation and cache_creation > 0: + result["cache_creation_input_tokens"] = cache_creation + + +def _extract_anthropic_start_usage(event: Any) -> TokenUsage: + usage = TokenUsage() + # Handle usage stats from message_start event + if hasattr(event, "type") and event.type == "message_start": + if hasattr(event, "message") and hasattr(event.message, "usage"): + usage["input_tokens"] = getattr(event.message.usage, "input_tokens", 0) + usage["cache_creation_input_tokens"] = getattr( + event.message.usage, "cache_creation_input_tokens", 0 + ) + usage["cache_read_input_tokens"] = getattr( + event.message.usage, "cache_read_input_tokens", 0 + ) + # Capture raw usage metadata for backend processing + # Serialize to dict here in the converter (not in utils) + serialized = serialize_raw_usage(event.message.usage) + if serialized: + usage["raw_usage"] = serialized + + return usage + + +def _start_anthropic_block( + block: Any, +) -> Tuple[Optional[StreamingContentBlock], Optional[ToolInProgress]]: + if not hasattr(block, "type"): + return None, None + + if block.type == "text": + content_block: StreamingContentBlock = {"type": "text", "text": ""} + return content_block, None + + elif block.type == "tool_use": + tool_block: StreamingContentBlock = { + "type": "function", + "id": getattr(block, "id", ""), + "function": {"name": getattr(block, "name", ""), "arguments": {}}, + } + tool_in_progress: ToolInProgress = {"block": tool_block, "input_string": ""} + return tool_block, tool_in_progress + + elif block.type == "thinking": + thinking_block: StreamingContentBlock = { + "type": "thinking", + "thinking": getattr(block, "thinking", "") or "", + } + signature = getattr(block, "signature", None) + if signature: + thinking_block["signature"] = signature + return thinking_block, None + + elif block.type == "redacted_thinking": + redacted_block: StreamingContentBlock = { + "type": "redacted_thinking", + "data": getattr(block, "data", None), + } + return redacted_block, None + + return None, None + + +def _accumulate_anthropic_thinking_delta( + event: Any, current_block: Optional[StreamingContentBlock] +) -> None: + fields: Tuple[Literal["thinking", "signature"], ...] = ("thinking", "signature") + for field in fields: + if hasattr(event, "delta") and hasattr(event.delta, field): + value = getattr(event.delta, field) or "" + if current_block is not None and current_block.get("type") == "thinking": + current_block[field] = (current_block.get(field) or "") + value + return + + +def _append_anthropic_tool_json( + event: Any, + content_blocks: List[StreamingContentBlock], + tools_in_progress: Dict[str, ToolInProgress], +) -> None: + if hasattr(event, "index") and event.index < len(content_blocks): + block = content_blocks[event.index] + + if block.get("type") == "function" and block.get("id") in tools_in_progress: + tool = tools_in_progress[block["id"]] + partial_json = getattr(event.delta, "partial_json", "") + tool["input_string"] += partial_json diff --git a/posthog/ai/claude_agent_sdk/client.py b/posthog/ai/claude_agent_sdk/client.py index f9a718d80..5d3d3d458 100644 --- a/posthog/ai/claude_agent_sdk/client.py +++ b/posthog/ai/claude_agent_sdk/client.py @@ -23,9 +23,10 @@ "Please install the Claude Agent SDK to use this feature: 'pip install claude-agent-sdk'" ) -from posthog.ai.claude_agent_sdk.formatting import ( +from .formatting import ( + _format_user_input, format_assistant_blocks, - format_tool_result_content, + format_tool_result_content as format_tool_result_content, ) from posthog.ai.claude_agent_sdk.processor import ( PostHogClaudeAgentProcessor, @@ -135,95 +136,77 @@ async def receive_response(self): """Instrumented receive_response -- yields all messages, emits PostHog events.""" async for message in self._client.receive_response(): try: - if isinstance(message, StreamEvent): - self._tracker.process_stream_event(message) - - if self._tracker.has_completed_generation(): - gen = self._tracker.pop_generation() - self._generation_index += 1 - self._current_generation_span_id = gen.span_id - self._processor._emit_generation( - gen, - self._trace_id, - self._generation_index, - self._current_input, - self._pending_output or None, - self._distinct_id, - self._extra_props, - self._privacy, - self._groups, - ) - self._current_input = self._next_input - self._next_input = None - self._pending_output = [] - - elif isinstance(message, AssistantMessage): - self._tracker.set_model(message.model) - parent_id = ( - self._tracker.current_span_id - or self._current_generation_span_id - ) - for block in message.content: - if isinstance(block, ToolUseBlock): - self._processor._emit_tool_span( - block, - self._trace_id, - parent_id, - self._distinct_id, - self._extra_props, - self._privacy, - self._groups, - ) - output_content = format_assistant_blocks(message.content) - if output_content: - self._pending_output = [ - {"role": "assistant", "content": output_content} - ] - - elif isinstance(message, UserMessage): - content = message.content - if isinstance(content, str): - self._next_input = [{"role": "user", "content": content}] - elif isinstance(content, list): - formatted: List[Dict[str, Any]] = [] - for block in content: - if hasattr(block, "tool_use_id"): - formatted.append( - { - "type": "tool_result", - "tool_use_id": block.tool_use_id, - "content": format_tool_result_content( - block, self._processor._client - ), - } - ) - elif hasattr(block, "text"): - formatted.append({"type": "text", "text": block.text}) - if formatted: - self._next_input = [{"role": "user", "content": formatted}] - - elif isinstance(message, ResultMessage): - if not self._tracker.had_any_stream_events: - self._processor._emit_generation_from_result( - message, - self._trace_id, - self._tracker.last_model, - self._query_start, - self._current_input, - self._pending_output, - self._distinct_id, - self._extra_props, - self._privacy, - self._groups, - ) - # Don't emit trace here -- wait for disconnect/close - # so multi-turn sessions get one trace at the end - + self._process_message(message) except Exception as e: log.debug(f"PostHog instrumentation error (non-fatal): {e}") - yield message + def _process_message(self, message: Any) -> None: + if isinstance(message, StreamEvent): + self._tracker.process_stream_event(message) + + if self._tracker.has_completed_generation(): + gen = self._tracker.pop_generation() + self._generation_index += 1 + self._current_generation_span_id = gen.span_id + self._processor._emit_generation( + gen, + self._trace_id, + self._generation_index, + self._current_input, + self._pending_output or None, + self._distinct_id, + self._extra_props, + self._privacy, + self._groups, + ) + self._current_input = self._next_input + self._next_input = None + self._pending_output = [] + + elif isinstance(message, AssistantMessage): + self._process_assistant_message(message) + + elif isinstance(message, UserMessage): + self._next_input = _format_user_input( + message.content, self._processor._client, self._next_input + ) + + elif isinstance(message, ResultMessage): + if not self._tracker.had_any_stream_events: + self._processor._emit_generation_from_result( + message, + self._trace_id, + self._tracker.last_model, + self._query_start, + self._current_input, + self._pending_output, + self._distinct_id, + self._extra_props, + self._privacy, + self._groups, + ) + # Don't emit trace here -- wait for disconnect/close + # so multi-turn sessions get one trace at the end + + def _process_assistant_message(self, message: "AssistantMessage") -> None: + self._tracker.set_model(message.model) + parent_id = self._tracker.current_span_id or self._current_generation_span_id + for block in message.content: + if isinstance(block, ToolUseBlock): + self._processor._emit_tool_span( + block, + self._trace_id, + parent_id, + self._distinct_id, + self._extra_props, + self._privacy, + self._groups, + ) + output_content = format_assistant_blocks(message.content) + if output_content: + self._pending_output = [{"role": "assistant", "content": output_content}] + async def disconnect(self) -> None: """ Disconnect the underlying client and emit the final PostHog trace event. diff --git a/posthog/ai/claude_agent_sdk/formatting.py b/posthog/ai/claude_agent_sdk/formatting.py index eea13c269..a99c8cbd3 100644 --- a/posthog/ai/claude_agent_sdk/formatting.py +++ b/posthog/ai/claude_agent_sdk/formatting.py @@ -1,4 +1,4 @@ -from typing import Any, Dict, List +from typing import Any, Dict, List, Optional try: from claude_agent_sdk import ToolUseBlock @@ -53,3 +53,26 @@ def format_assistant_blocks(blocks: Any) -> List[Dict[str, Any]]: else: out.append(to_plain(block)) return out + + +def _format_user_input( + content: Any, client: Any, previous: Optional[List[Dict[str, Any]]] +) -> Optional[List[Dict[str, Any]]]: + if isinstance(content, str): + return [{"role": "user", "content": content}] + if isinstance(content, list): + formatted: List[Dict[str, Any]] = [] + for block in content: + if hasattr(block, "tool_use_id"): + formatted.append( + { + "type": "tool_result", + "tool_use_id": block.tool_use_id, + "content": format_tool_result_content(block, client), + } + ) + elif hasattr(block, "text"): + formatted.append({"type": "text", "text": block.text}) + if formatted: + return [{"role": "user", "content": formatted}] + return previous diff --git a/posthog/ai/claude_agent_sdk/processor.py b/posthog/ai/claude_agent_sdk/processor.py index d5405ec66..022ad8cc5 100644 --- a/posthog/ai/claude_agent_sdk/processor.py +++ b/posthog/ai/claude_agent_sdk/processor.py @@ -25,9 +25,10 @@ ) from posthog import setup -from posthog.ai.claude_agent_sdk.formatting import ( +from .formatting import ( + _format_user_input, format_assistant_blocks, - format_tool_result_content, + format_tool_result_content as format_tool_result_content, ) from posthog.ai.media import ensure_serializable as _ensure_serializable from posthog.ai.utils import ( @@ -91,31 +92,7 @@ def process_stream_event(self, event: "StreamEvent") -> None: self._current.raw_usage = dict(usage) elif event_type == "message_delta" and self._current is not None: - usage = raw.get("usage", {}) - self._current.raw_usage = {**(self._current.raw_usage or {}), **usage} - # message_delta usage is cumulative: output tokens, and with server - # tools (web search, ...) the input and cache counts grow past the - # message_start values. Keep the larger count; a field the delta - # omits keeps the message_start value. - if usage.get("output_tokens"): - self._current.output_tokens = usage["output_tokens"] - for field_name in ( - "input_tokens", - "cache_read_input_tokens", - "cache_creation_input_tokens", - ): - delta_value = usage.get(field_name) - current_value = getattr(self._current, field_name) - if isinstance(delta_value, int) and delta_value > (current_value or 0): - setattr(self._current, field_name, delta_value) - server_tool_use = usage.get("server_tool_use") or {} - web_search_requests = server_tool_use.get("web_search_requests") - if isinstance(web_search_requests, int) and web_search_requests > 0: - self._current.web_search_count = web_search_requests - # Extract stop reason from message_delta - delta_stop_reason = raw.get("delta", {}).get("stop_reason") - if delta_stop_reason is not None: - self._current.stop_reason = delta_stop_reason + self._process_usage_delta(raw) elif event_type == "message_stop" and self._current is not None: self._current.end_time = time.time() @@ -123,6 +100,38 @@ def process_stream_event(self, event: "StreamEvent") -> None: self._last_model = self._current.model self._current = None + def _process_usage_delta(self, raw: Dict[str, Any]) -> None: + assert self._current is not None + usage = raw.get("usage", {}) + self._current.raw_usage = {**(self._current.raw_usage or {}), **usage} + # message_delta usage is cumulative: output tokens, and with server + # tools (web search, ...) the input and cache counts grow past the + # message_start values. Keep the larger count; a field the delta + # omits keeps the message_start value. + if usage.get("output_tokens"): + self._current.output_tokens = usage["output_tokens"] + self._update_cumulative_input_counts(usage) + server_tool_use = usage.get("server_tool_use") or {} + web_search_requests = server_tool_use.get("web_search_requests") + if isinstance(web_search_requests, int) and web_search_requests > 0: + self._current.web_search_count = web_search_requests + # Extract stop reason from message_delta + delta_stop_reason = raw.get("delta", {}).get("stop_reason") + if delta_stop_reason is not None: + self._current.stop_reason = delta_stop_reason + + def _update_cumulative_input_counts(self, usage: Dict[str, Any]) -> None: + assert self._current is not None + for field_name in ( + "input_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + ): + delta_value = usage.get(field_name) + current_value = getattr(self._current, field_name) + if isinstance(delta_value, int) and delta_value > (current_value or 0): + setattr(self._current, field_name, delta_value) + def set_model(self, model: str) -> None: self._last_model = model @@ -265,19 +274,36 @@ async def query( elif not options.include_partial_messages: options = replace(options, include_partial_messages=True) + async for message in self._instrument_query( + prompt, + options, + transport, + distinct_id_override, + trace_id, + extra_props, + privacy, + groups, + ): + yield message + + async def _instrument_query( + self, + prompt: Any, + options: ClaudeAgentOptions, + transport: Any, + distinct_id_override: Any, + trace_id: str, + extra_props: Dict[str, Any], + privacy: bool, + groups: Dict[str, Any], + ): tracker = _GenerationTracker() query_start = time.time() generation_index = 0 current_generation_span_id: Optional[str] = None # Track input/output for generation events - initial_input: List[Dict[str, Any]] = [] - if isinstance(prompt, str): - initial_input = [{"role": "user", "content": prompt}] - if options and options.system_prompt and isinstance(options.system_prompt, str): - initial_input = [ - {"role": "system", "content": options.system_prompt} - ] + initial_input + initial_input, current_input = _query_inputs(prompt, options) # Two-slot input tracking: # - current_input: input for the generation currently in progress @@ -287,7 +313,6 @@ async def query( # message_start → content_blocks → AssistantMessage → UserMessage(tool result) → message_stop # So UserMessage arrives *before* message_stop. When message_stop fires we emit # with current_input, then promote next_input → current_input for the next turn. - current_input: Optional[List[Dict[str, Any]]] = initial_input or None next_input: Optional[List[Dict[str, Any]]] = None # Accumulate assistant output per generation @@ -324,74 +349,33 @@ async def query( pending_output = [] elif isinstance(message, AssistantMessage): - tracker.set_model(message.model) - # Use the in-progress generation's span_id as parent for tool spans. - # AssistantMessage arrives before message_stop, so current_generation_span_id - # would be stale (from the previous turn). tracker.current_span_id gives us - # the correct in-progress generation. - parent_id = tracker.current_span_id or current_generation_span_id - for block in message.content: - if isinstance(block, ToolUseBlock): - self._emit_tool_span( - block, - trace_id, - parent_id, - distinct_id_override, - extra_props, - privacy, - groups, - ) - output_content = format_assistant_blocks(message.content) - if output_content: - pending_output = [ - {"role": "assistant", "content": output_content} - ] + pending_output = self._process_assistant_message( + message, + tracker, + current_generation_span_id, + trace_id, + distinct_id_override, + extra_props, + privacy, + groups, + pending_output, + ) elif isinstance(message, UserMessage): # UserMessages carry tool results. They arrive *before* message_stop # for the current turn, so queue them as input for the *next* generation. - content = message.content - if isinstance(content, str): - next_input = [{"role": "user", "content": content}] - elif isinstance(content, list): - formatted: List[Dict[str, Any]] = [] - for block in content: - if hasattr(block, "tool_use_id"): - formatted.append( - { - "type": "tool_result", - "tool_use_id": block.tool_use_id, - "content": format_tool_result_content( - block, self._client - ), - } - ) - elif hasattr(block, "text"): - formatted.append({"type": "text", "text": block.text}) - if formatted: - next_input = [{"role": "user", "content": formatted}] + next_input = _format_user_input( + message.content, self._client, next_input + ) elif isinstance(message, ResultMessage): - # Fallback: if no StreamEvents were received, emit a single - # generation from ResultMessage aggregate data - if not tracker.had_any_stream_events: - self._emit_generation_from_result( - message, - trace_id, - tracker.last_model, - query_start, - initial_input, - pending_output, - distinct_id_override, - extra_props, - privacy, - groups, - ) - - self._emit_trace( + self._process_query_result( message, + tracker, trace_id, query_start, + initial_input, + pending_output, distinct_id_override, extra_props, privacy, @@ -403,6 +387,79 @@ async def query( yield message + def _process_query_result( + self, + message: "ResultMessage", + tracker: _GenerationTracker, + trace_id: str, + query_start: float, + initial_input: List[Dict[str, Any]], + pending_output: List[Dict[str, Any]], + distinct_id_override: Any, + extra_props: Dict[str, Any], + privacy: bool, + groups: Dict[str, Any], + ) -> None: + # Fallback: if no StreamEvents were received, emit a single + # generation from ResultMessage aggregate data + if not tracker.had_any_stream_events: + self._emit_generation_from_result( + message, + trace_id, + tracker.last_model, + query_start, + initial_input, + pending_output, + distinct_id_override, + extra_props, + privacy, + groups, + ) + + self._emit_trace( + message, + trace_id, + query_start, + distinct_id_override, + extra_props, + privacy, + groups, + ) + + def _process_assistant_message( + self, + message: "AssistantMessage", + tracker: _GenerationTracker, + current_generation_span_id: Optional[str], + trace_id: str, + distinct_id: Any, + extra_props: Dict[str, Any], + privacy: bool, + groups: Dict[str, Any], + pending_output: List[Dict[str, Any]], + ) -> List[Dict[str, Any]]: + tracker.set_model(message.model) + # Use the in-progress generation's span_id as parent for tool spans. + # AssistantMessage arrives before message_stop, so current_generation_span_id + # would be stale (from the previous turn). tracker.current_span_id gives us + # the correct in-progress generation. + parent_id = tracker.current_span_id or current_generation_span_id + for block in message.content: + if isinstance(block, ToolUseBlock): + self._emit_tool_span( + block, + trace_id, + parent_id, + distinct_id, + extra_props, + privacy, + groups, + ) + output_content = format_assistant_blocks(message.content) + if output_content: + pending_output = [{"role": "assistant", "content": output_content}] + return pending_output + def _emit_generation( self, gen: _GenerationData, @@ -441,6 +498,26 @@ def _emit_generation( **extra_props, } + self._add_generation_content( + properties, input_messages, output_choices, privacy + ) + + _add_generation_usage(properties, gen) + + if resolved_id is None: + properties["$process_person_profile"] = False + + self._capture_event( + "$ai_generation", properties, resolved_id or trace_id, groups + ) + + def _add_generation_content( + self, + properties: Dict[str, Any], + input_messages: Optional[List[Dict[str, Any]]], + output_choices: Optional[List[Dict[str, Any]]], + privacy: bool, + ) -> None: if input_messages is not None: properties["$ai_input"] = _with_privacy_mode( self._client, @@ -454,27 +531,6 @@ def _emit_generation( finalize_ai_content(output_choices, self._client), ) - if gen.cache_read_input_tokens: - properties["$ai_cache_read_input_tokens"] = gen.cache_read_input_tokens - if gen.cache_creation_input_tokens: - properties["$ai_cache_creation_input_tokens"] = ( - gen.cache_creation_input_tokens - ) - if gen.web_search_count: - properties["$ai_web_search_count"] = gen.web_search_count - if gen.raw_usage: - properties["$ai_usage"] = gen.raw_usage - - if gen.stop_reason is not None: - properties["$ai_stop_reason"] = gen.stop_reason - - if resolved_id is None: - properties["$process_person_profile"] = False - - self._capture_event( - "$ai_generation", properties, resolved_id or trace_id, groups - ) - def _emit_generation_from_result( self, result: "ResultMessage", @@ -516,30 +572,11 @@ def _emit_generation_from_result( **extra_props, } - if input_messages is not None: - properties["$ai_input"] = _with_privacy_mode( - self._client, - privacy, - finalize_ai_content(input_messages, self._client), - ) - if output_choices is not None: - properties["$ai_output_choices"] = _with_privacy_mode( - self._client, - privacy, - finalize_ai_content(output_choices, self._client), - ) - - cache_read = usage.get("cache_read_input_tokens") - cache_creation = usage.get("cache_creation_input_tokens") - if cache_read: - properties["$ai_cache_read_input_tokens"] = cache_read - if cache_creation: - properties["$ai_cache_creation_input_tokens"] = cache_creation - if usage: - properties["$ai_usage"] = dict(usage) + self._add_generation_content( + properties, input_messages, output_choices, privacy + ) - if result.total_cost_usd is not None: - properties["$ai_total_cost_usd"] = result.total_cost_usd + _add_result_usage(properties, usage, result) if resolved_id is None: properties["$process_person_profile"] = False @@ -639,3 +676,47 @@ def _resolve_distinct_id( return str(override) # Fall back to instance default return self._get_distinct_id(result) + + +def _query_inputs( + prompt: Any, options: ClaudeAgentOptions +) -> tuple[List[Dict[str, Any]], Optional[List[Dict[str, Any]]]]: + initial_input: List[Dict[str, Any]] = [] + if isinstance(prompt, str): + initial_input = [{"role": "user", "content": prompt}] + if options and options.system_prompt and isinstance(options.system_prompt, str): + initial_input = [ + {"role": "system", "content": options.system_prompt} + ] + initial_input + + return initial_input, initial_input or None + + +def _add_generation_usage(properties: Dict[str, Any], gen: _GenerationData) -> None: + if gen.cache_read_input_tokens: + properties["$ai_cache_read_input_tokens"] = gen.cache_read_input_tokens + if gen.cache_creation_input_tokens: + properties["$ai_cache_creation_input_tokens"] = gen.cache_creation_input_tokens + if gen.web_search_count: + properties["$ai_web_search_count"] = gen.web_search_count + if gen.raw_usage: + properties["$ai_usage"] = gen.raw_usage + + if gen.stop_reason is not None: + properties["$ai_stop_reason"] = gen.stop_reason + + +def _add_result_usage( + properties: Dict[str, Any], usage: Dict[str, Any], result: "ResultMessage" +) -> None: + cache_read = usage.get("cache_read_input_tokens") + cache_creation = usage.get("cache_creation_input_tokens") + if cache_read: + properties["$ai_cache_read_input_tokens"] = cache_read + if cache_creation: + properties["$ai_cache_creation_input_tokens"] = cache_creation + if usage: + properties["$ai_usage"] = dict(usage) + + if result.total_cost_usd is not None: + properties["$ai_total_cost_usd"] = result.total_cost_usd diff --git a/posthog/ai/evaluations/_async_datasets.py b/posthog/ai/evaluations/_async_datasets.py index 80cc6418c..6593a919c 100644 --- a/posthog/ai/evaluations/_async_datasets.py +++ b/posthog/ai/evaluations/_async_datasets.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field from datetime import datetime -from typing import AsyncIterator, Literal, Sequence +from typing import AsyncIterator, Literal, Sequence, cast from uuid import UUID from ._dataset_serialization import ( @@ -387,10 +387,15 @@ async def snapshot( return AsyncDatasetSnapshot( dataset, selected, dataset.current_revision_id, self ) + return await self._historical_snapshot(dataset, selected) + + async def _historical_snapshot( + self, dataset: Dataset, selected: int + ) -> AsyncDatasetSnapshot: # Revisions are consecutive and newest first. Concurrent writes can # shift their offsets, so verify the result and scan on a miss. page = await self.list_revisions( - dataset.id, limit=1, offset=dataset.current_revision - selected + dataset.id, limit=1, offset=cast(int, dataset.current_revision) - selected ) for candidate in page.results: if candidate.revision == selected: diff --git a/posthog/ai/evaluations/_client.py b/posthog/ai/evaluations/_client.py index db0a5d695..5e41a8de6 100644 --- a/posthog/ai/evaluations/_client.py +++ b/posthog/ai/evaluations/_client.py @@ -69,19 +69,23 @@ def _experiment_body(name: str, options: ExperimentOptions) -> bytes: for key, value in options.items(): if key in ("id", "started_at"): continue - if key == "run_source": - if value is not None and value not in ("ci", "local", "scheduled"): - raise ValueError("run_source must be ci, local, scheduled, or None.") - data[key] = value - elif key in ("expected_item_count", "expected_result_count"): - data[key] = _count(value, key) - elif key == "dataset_revision_id": - data[key] = None if value is None else _uuid(value) - else: - data[key] = _identifier(value, key) + data[key] = _experiment_option(key, value) return _encode(data) +def _experiment_option(key: str, value: Any) -> Any: + if key == "run_source": + if value is not None and value not in ("ci", "local", "scheduled"): + raise ValueError("run_source must be ci, local, scheduled, or None.") + return value + elif key in ("expected_item_count", "expected_result_count"): + return _count(value, key) + elif key == "dataset_revision_id": + return None if value is None else _uuid(value) + else: + return _identifier(value, key) + + @dataclass(frozen=True) class _Chunk: start: int @@ -95,9 +99,7 @@ def _upload_body(items: list[bytes], results: list[bytes]) -> bytes: ) -def _chunks(results: tuple[EvaluationResult, ...]) -> list[_Chunk]: - # Validate the finite submission before sending any chunk. Receipts and - # prepared bodies use memory proportional to this caller-sized sequence. +def _result_declarations(results: tuple[EvaluationResult, ...]) -> dict[str, bytes]: identities: set[tuple[str, str]] = set() declarations: dict[str, bytes] = {} for result in results: @@ -114,6 +116,36 @@ def _chunks(results: tuple[EvaluationResult, ...]) -> list[_Chunk]: if previous != result.item._json: raise ValueError("Conflicting declarations for the same item UUID.") + return declarations + + +def _chunk_full( + results: list[bytes], items: list[bytes], item: bytes | None, size: int +) -> bool: + return ( + len(results) >= _MAX_RESULTS + or len(items) + (item is not None) > _MAX_ITEMS + or size > _MAX_REQUEST_BYTES + ) + + +def _chunk_extra( + result: EvaluationResult, + item: bytes | None, + results: list[bytes], + items: list[bytes], +) -> int: + extra = len(result._json) + bool(results) + if item is not None: + extra += len(item) + bool(items) + return extra + + +def _chunks(results: tuple[EvaluationResult, ...]) -> list[_Chunk]: + # Validate the finite submission before sending any chunk. Receipts and + # prepared bodies use memory proportional to this caller-sized sequence. + declarations = _result_declarations(results) + chunks: list[_Chunk] = [] item_ids: set[str] = set() items: list[bytes] = [] @@ -126,14 +158,8 @@ def _chunks(results: tuple[EvaluationResult, ...]) -> list[_Chunk]: item = ( declarations.get(result.item_id) if result.item_id not in item_ids else None ) - extra = len(result._json) + bool(encoded_results) - if item is not None: - extra += len(item) + bool(items) - if encoded_results and ( - len(encoded_results) >= _MAX_RESULTS - or len(items) + (item is not None) > _MAX_ITEMS - or size + extra > _MAX_REQUEST_BYTES - ): + extra = _chunk_extra(result, item, encoded_results, items) + if encoded_results and _chunk_full(encoded_results, items, item, size + extra): chunks.append( _Chunk( start, results[start:index], _upload_body(items, encoded_results) diff --git a/posthog/ai/evaluations/_dataset_serialization.py b/posthog/ai/evaluations/_dataset_serialization.py index c95fb5aea..b96d2c6bb 100644 --- a/posthog/ai/evaluations/_dataset_serialization.py +++ b/posthog/ai/evaluations/_dataset_serialization.py @@ -97,14 +97,9 @@ def _item( identity, owner = _uuid(data["id"]), _uuid(data["dataset"]) introduced = _integer(data["dataset_revision"]) is_archived = _boolean(data["archived"]) - if (id is not None and identity != id) or ( - dataset_id is not None and owner != dataset_id - ): - raise ValueError("Mismatched dataset item identity.") - if revision is not None and introduced > revision: - raise ValueError("The item version is newer than the requested revision.") - if archived is not None and is_archived != archived: - raise ValueError("The item does not match the requested archive filter.") + _validate_item_filters( + identity, owner, introduced, is_archived, id, dataset_id, revision, archived + ) if data["input"] is None: raise ValueError("Dataset inputs cannot be null.") return DatasetItem( @@ -134,6 +129,26 @@ def _item( ) +def _validate_item_filters( + identity: str, + owner: str, + introduced: int, + is_archived: bool, + id: str | None, + dataset_id: str | None, + revision: int | None, + archived: bool | None, +) -> None: + if (id is not None and identity != id) or ( + dataset_id is not None and owner != dataset_id + ): + raise ValueError("Mismatched dataset item identity.") + if revision is not None and introduced > revision: + raise ValueError("The item version is newer than the requested revision.") + if archived is not None and is_archived != archived: + raise ValueError("The item does not match the requested archive filter.") + + def _parse(parser: Callable[[dict[str, Any]], _T], data: dict[str, Any]) -> _T: return _parse_response( parser, data, detail="The dataset API returned an invalid response." @@ -197,12 +212,7 @@ def _body(fields: dict[str, Any]) -> bytes: data = { key: value for key, value in fields.items() if not isinstance(value, _Unset) } - for key in ("name", "description"): - if key in data: - value = _text(data[key]) - maximum = 400 if key == "name" else 10_000 - if len(value) > maximum or (key == "name" and not value.strip()): - raise ValueError(f"Invalid dataset {key}.") + _validate_dataset_text(data) if "input" in data and data["input"] is None: raise ValueError("Dataset inputs cannot be null.") if "metadata" in data: @@ -213,6 +223,15 @@ def _body(fields: dict[str, Any]) -> bytes: return _encode(data) +def _validate_dataset_text(data: dict[str, Any]) -> None: + for key in ("name", "description"): + if key in data: + value = _text(data[key]) + maximum = 400 if key == "name" else 10_000 + if len(value) > maximum or (key == "name" and not value.strip()): + raise ValueError(f"Invalid dataset {key}.") + + def _create_item_body( dataset_id: str, *, diff --git a/posthog/ai/evaluations/_datasets.py b/posthog/ai/evaluations/_datasets.py index cfe359a2e..a2a364907 100644 --- a/posthog/ai/evaluations/_datasets.py +++ b/posthog/ai/evaluations/_datasets.py @@ -4,7 +4,7 @@ from dataclasses import dataclass, field from datetime import datetime -from typing import Iterator, Literal, Sequence +from typing import Iterator, Literal, Sequence, cast from uuid import UUID from ._dataset_serialization import ( @@ -381,10 +381,13 @@ def snapshot( and dataset.current_revision_id is not None ): return DatasetSnapshot(dataset, selected, dataset.current_revision_id, self) + return self._historical_snapshot(dataset, selected) + + def _historical_snapshot(self, dataset: Dataset, selected: int) -> DatasetSnapshot: # Revisions are consecutive and newest first. Concurrent writes can # shift their offsets, so verify the result and scan on a miss. page = self.list_revisions( - dataset.id, limit=1, offset=dataset.current_revision - selected + dataset.id, limit=1, offset=cast(int, dataset.current_revision) - selected ) for candidate in page.results: if candidate.revision == selected: diff --git a/posthog/ai/evaluations/_serialization.py b/posthog/ai/evaluations/_serialization.py index 83aadb280..c651bad11 100644 --- a/posthog/ai/evaluations/_serialization.py +++ b/posthog/ai/evaluations/_serialization.py @@ -109,19 +109,15 @@ def _encode(value: object, *, sort_keys: bool = True) -> bytes: ) from exc +def _invalid_json_key(key: object) -> bool: + return not isinstance(key, str) or "\x00" in key + + def _validate_json(value: object, depth: int = 1) -> None: if isinstance(value, (dict, list)): if depth > 32: raise ValueError("Payload JSON nesting must not exceed 32 levels.") - children: Iterable[object] - if isinstance(value, dict): - if any(not isinstance(key, str) or "\x00" in key for key in value): - raise ValueError( - "JSON object keys must be strings without null characters." - ) - children = value.values() - else: - children = value + children = _json_children(value) for child in children: _validate_json(child, depth + 1) elif isinstance(value, str): @@ -131,6 +127,19 @@ def _validate_json(value: object, depth: int = 1) -> None: raise ValueError("Provide JSON values, without automatic object conversion.") +def _json_children(value: dict | list) -> Iterable[object]: + children: Iterable[object] + if isinstance(value, dict): + if any(_invalid_json_key(key) for key in value): + raise ValueError( + "JSON object keys must be strings without null characters." + ) + children = value.values() + else: + children = value + return children + + def _payload( explicit: dict[str, JSONValue] | _Unset, fields: dict[str, JSONValue | _Unset], @@ -149,43 +158,32 @@ def _payload( supplied = explicit elif not supplied: return _UNSET + _validate_payload_fields(supplied) + _validate_json(supplied) + encoded = _encode(supplied) + if len(encoded) > limit: + raise ValueError(f"Payload exceeds {limit} bytes of UTF-8 JSON.") + return json.loads(encoded) + + +def _validate_payload_fields(supplied: dict[str, JSONValue]) -> None: metadata = supplied.get("metadata") if metadata is not None and not isinstance(metadata, dict): raise ValueError("metadata must be a JSON object or None.") for field in ("reasoning", "error_message"): if supplied.get(field) is not None and not isinstance(supplied[field], str): raise ValueError(f"{field} must be a string or None.") - _validate_json(supplied) - encoded = _encode(supplied) - if len(encoded) > limit: - raise ValueError(f"Payload exceeds {limit} bytes of UTF-8 JSON.") - return json.loads(encoded) def _score(value: object) -> bool | float | list[str]: if type(value) is bool: return value if type(value) in (int, float): - try: - number = float(cast(Union[int, float], value)) - except OverflowError as exc: - raise ValueError( - "Numeric scores must fit a finite binary64 number." - ) from exc - if not math.isfinite(number): - raise ValueError("Numeric scores must be finite.") - if type(value) is int and number != value: - raise ValueError( - "Integer scores must be exactly representable as a binary64 number." - ) - return number + return _numeric_score(value) if ( isinstance(value, list) and value - and all( - isinstance(key, str) and key and len(key) <= 128 and "\x00" not in key - for key in value - ) + and all(_valid_category_key(key) for key in value) ): if len(set(value)) != len(value): raise ValueError("Categorical scores must contain distinct keys.") @@ -195,6 +193,24 @@ def _score(value: object) -> bool | float | list[str]: ) +def _numeric_score(value: object) -> float: + try: + number = float(cast(Union[int, float], value)) + except OverflowError as exc: + raise ValueError("Numeric scores must fit a finite binary64 number.") from exc + if not math.isfinite(number): + raise ValueError("Numeric scores must be finite.") + if type(value) is int and number != value: + raise ValueError( + "Integer scores must be exactly representable as a binary64 number." + ) + return number + + +def _valid_category_key(key: object) -> bool: + return isinstance(key, str) and bool(key) and len(key) <= 128 and "\x00" not in key + + def _count(value: object, field: str) -> int | None: if value is not None and (type(value) is not int or not 0 <= value <= 2**63 - 1): raise ValueError( diff --git a/posthog/ai/evaluations/_transport.py b/posthog/ai/evaluations/_transport.py index f6489c502..1feab76f0 100644 --- a/posthog/ai/evaluations/_transport.py +++ b/posthog/ai/evaluations/_transport.py @@ -7,8 +7,8 @@ import time from datetime import datetime, timezone from email.utils import parsedate_to_datetime -from typing import Any -from urllib.parse import urlsplit, urlunsplit +from typing import Any, cast +from urllib.parse import SplitResult, urlsplit, urlunsplit import requests @@ -26,6 +26,15 @@ def _normalize_host(host: str | None) -> str: parsed = urlsplit((host or "https://us.posthog.com").strip()) + _validate_host(parsed) + # Accessing port also validates malformed port numbers before any I/O. + netloc = parsed.netloc + if parsed.port in {None, 443} and parsed.scheme == "https": + netloc = _CLOUD_HOSTS.get(cast(str, parsed.hostname), netloc) + return urlunsplit((parsed.scheme, netloc, parsed.path.rstrip("/"), "", "")) + + +def _validate_host(parsed: SplitResult) -> None: if ( parsed.scheme not in {"http", "https"} or not parsed.hostname @@ -35,11 +44,6 @@ def _normalize_host(host: str | None) -> str: or parsed.fragment ): raise ValueError("host must be an absolute HTTP(S) URL without credentials") - # Accessing port also validates malformed port numbers before any I/O. - netloc = parsed.netloc - if parsed.port in {None, 443} and parsed.scheme == "https": - netloc = _CLOUD_HOSTS.get(parsed.hostname, netloc) - return urlunsplit((parsed.scheme, netloc, parsed.path.rstrip("/"), "", "")) def _validate_options( @@ -57,6 +61,10 @@ def _validate_options( or any(character.isspace() for character in secret_key) ): raise ValueError("secret_key must be a nonempty credential without whitespace") + _validate_retry_options(timeout, max_retries) + + +def _validate_retry_options(timeout: float, max_retries: int) -> None: if ( isinstance(timeout, bool) or not isinstance(timeout, (int, float)) diff --git a/posthog/ai/evaluations/_types.py b/posthog/ai/evaluations/_types.py index 595f06920..8b4f9b25f 100644 --- a/posthog/ai/evaluations/_types.py +++ b/posthog/ai/evaluations/_types.py @@ -191,6 +191,17 @@ def from_dict(cls, data: Mapping[str, Any]) -> EvaluationItem: return cls(**dict(data)) +def _validate_result_status( + status: ResultStatus, value: ScoreValue | None, error_code: str | None +) -> None: + if status not in ("ok", "error", "skipped", "not_applicable"): + raise ValueError("Unsupported result status.") + if status != "ok" and value is not None: + raise ValueError("Only an ok result may contain a score.") + if status != "error" and error_code is not None: + raise ValueError("Only an error result may contain an error_code.") + + @dataclass(frozen=True, init=False) class EvaluationResult: """One scorer outcome for an item and an exact scorer-version UUID. @@ -222,12 +233,7 @@ def __init__( ) -> None: item_id = item.id if isinstance(item, EvaluationItem) else _uuid(item) version_id = _uuid(scorer_version_id) - if status not in ("ok", "error", "skipped", "not_applicable"): - raise ValueError("Unsupported result status.") - if status != "ok" and value is not None: - raise ValueError("Only an ok result may contain a score.") - if status != "error" and error_code is not None: - raise ValueError("Only an error result may contain an error_code.") + _validate_result_status(status, value, error_code) data: dict[str, Any] = { "item_id": item_id, "scorer_version_id": version_id, @@ -253,8 +259,7 @@ def __init__( _MAX_RESULT_BYTES, ) if not isinstance(prepared, _Unset): - if status != "error" and prepared.get("error_message") is not None: - raise ValueError("Only an error result may contain an error_message.") + _validate_result_payload(status, prepared) data["payload"] = prepared object.__setattr__( self, "item", item if isinstance(item, EvaluationItem) else item_id @@ -425,18 +430,39 @@ def _upload_receipt( ) for row in data["results"] ) - if len(items) != len({item.id for item in items}) or { - item.id for item in items - } != {result.item_id for result in results}: - raise ValueError("Incomplete item acknowledgments.") - if [(result.item_id, result.scorer_version_id) for result in outcomes] != [ - (result.item_id, result.scorer_version_id) for result in results - ]: - raise ValueError("Incomplete result acknowledgments.") - if len({result.id for result in outcomes}) != len(outcomes): - raise ValueError("Duplicate result identities in acknowledgment.") + _validate_upload_identities(items, outcomes, results) return UploadReceipt(items=items, results=outcomes) except (KeyError, TypeError, ValueError): raise EvaluationAPIError( code="invalid_response", persistence="unknown" ) from None + + +def _validate_upload_identities( + items: tuple[ItemReceipt, ...], + outcomes: tuple[ResultReceipt, ...], + results: tuple[EvaluationResult, ...], +) -> None: + _validate_item_acknowledgments(items, results) + if [(result.item_id, result.scorer_version_id) for result in outcomes] != [ + (result.item_id, result.scorer_version_id) for result in results + ]: + raise ValueError("Incomplete result acknowledgments.") + if len({result.id for result in outcomes}) != len(outcomes): + raise ValueError("Duplicate result identities in acknowledgment.") + + +def _validate_result_payload( + status: ResultStatus, payload: dict[str, JSONValue] +) -> None: + if status != "error" and payload.get("error_message") is not None: + raise ValueError("Only an error result may contain an error_message.") + + +def _validate_item_acknowledgments( + items: tuple[ItemReceipt, ...], results: tuple[EvaluationResult, ...] +) -> None: + if len(items) != len({item.id for item in items}) or { + item.id for item in items + } != {result.item_id for result in results}: + raise ValueError("Incomplete item acknowledgments.") diff --git a/posthog/ai/gemini/gemini_converter.py b/posthog/ai/gemini/gemini_converter.py index 81b4f630b..6d409ee40 100644 --- a/posthog/ai/gemini/gemini_converter.py +++ b/posthog/ai/gemini/gemini_converter.py @@ -71,36 +71,7 @@ def _format_part(part: Any) -> Optional[FormattedContentItem]: if "file_data" in plain: media = _format_media_payload(plain["file_data"]) return {"type": _kind_from_mime(media.get("mime_type")), "file_data": media} - # Gemini-native function_call/function_response parts are normalized to the - # provider-agnostic tool-call/tool-result blocks. Emitting the raw Gemini - # shape leaves them unrenderable and invisible to evals. - if "function_call" in plain: - call = to_plain(plain["function_call"]) or {} - block: Dict[str, Any] = { - "type": "function", - "function": { - "name": call.get("name"), - "arguments": call.get("args"), - }, - } - if isinstance(call.get("id"), str): - block["id"] = call["id"] - return block - if "function_response" in plain: - response = to_plain(plain["function_response"]) or {} - response_block: Dict[str, Any] = { - "type": "function", - "tool_name": response.get("name") or "", - "content": response.get("response"), - } - if isinstance(response.get("id"), str): - response_block["tool_use_id"] = response["id"] - return response_block - if not plain: - return None - key = next(iter(plain)) - fallback: Dict[str, Any] = {"type": key, key: to_plain(plain[key])} - return cast(FormattedContentItem, fallback) + return _format_gemini_tool_part(plain) class GeminiPart(TypedDict, total=False): @@ -157,23 +128,15 @@ def _format_dict_message(item: Dict[str, Any]) -> FormattedMessage: # Handle dict with content field if "content" in item: - content = item["content"] - - if isinstance(content, list): - # If content is a list, format it as content blocks - content_blocks = _format_parts_as_content_blocks(content) - return {"role": item.get("role", "user"), "content": content_blocks} - - elif not isinstance(content, str): - content = str(content) - - return {"role": item.get("role", "user"), "content": content} + # Lists become content blocks; other non-string content is stringified. + return { + "role": item.get("role", "user"), + "content": _format_message_content(item["content"]), + } if "role" not in item: plain = to_plain(item) - if isinstance(plain, dict) and _BARE_PART_KEYS.intersection( - normalize_part_keys(plain) - ): + if _is_bare_part(plain): return {"role": "user", "content": _format_parts_as_content_blocks([item])} # Handle dict with text field @@ -198,11 +161,7 @@ def _format_object_message(item: Any) -> FormattedMessage: # Handle object with parts attribute if hasattr(item, "parts") and hasattr(item.parts, "__iter__"): content_blocks = _format_parts_as_content_blocks(list(item.parts)) - role = getattr(item, "role", "user") if hasattr(item, "role") else "user" - - # Ensure role is a string - if not isinstance(role, str): - role = "user" + role = _object_message_role(item) return {"role": role, "content": content_blocks} @@ -210,38 +169,20 @@ def _format_object_message(item: Any) -> FormattedMessage: # be caught before the "text" branch below, which would otherwise treat an # unset `.text` as empty content and drop any other part kind it carries plain = to_plain(item) - if isinstance(plain, dict) and _BARE_PART_KEYS.intersection( - normalize_part_keys(plain) - ): + if _is_bare_part(plain): return {"role": "user", "content": _format_parts_as_content_blocks([item])} # Handle object with text attribute if hasattr(item, "text"): - role = getattr(item, "role", "user") if hasattr(item, "role") else "user" - - # Ensure role is a string - if not isinstance(role, str): - role = "user" + role = _object_message_role(item) return {"role": role, "content": item.text} # Handle object with content attribute if hasattr(item, "content"): - role = getattr(item, "role", "user") if hasattr(item, "role") else "user" - - # Ensure role is a string - if not isinstance(role, str): - role = "user" + role = _object_message_role(item) - content = item.content - - if isinstance(content, list): - content_blocks = _format_parts_as_content_blocks(content) - return {"role": role, "content": content_blocks} - - elif not isinstance(content, str): - content = str(content) - return {"role": role, "content": content} + return {"role": role, "content": _format_message_content(item.content)} # Fallback to string representation return {"role": "user", "content": str(item)} @@ -265,55 +206,9 @@ def format_gemini_response(response: Any) -> List[FormattedMessage]: if hasattr(response, "candidates") and response.candidates: for candidate in response.candidates: - if hasattr(candidate, "content") and candidate.content: - content: List[FormattedContentItem] = [] - - if hasattr(candidate.content, "parts") and candidate.content.parts: - for part in candidate.content.parts: - # Checked ahead of _format_part so loosely-specced MagicMock - # fixtures (which report a truthy .text but aren't real Part - # objects) still resolve to a text block — must stay ahead of - # the _format_part delegation below. - text = getattr(part, "text", None) - if isinstance(text, str) and text: - # `is True` so loosely-specced mocks (whose .thought is a - # truthy Mock) still land on the text branch below. - if getattr(part, "thought", None) is True: - content.append({"type": "thinking", "thinking": text}) - else: - content.append({"type": "text", "text": text}) - continue - - if hasattr(part, "function_call") and part.function_call: - function_call = part.function_call - content.append( - _function_call_block( - function_call.name, - function_call.args, - getattr(function_call, "id", None), - ) - ) - continue - - block = _format_part(part) - if block is not None: - content.append(block) - - if content: - output.append( - { - "role": "assistant", - "content": content, - } - ) - - elif hasattr(candidate, "text") and candidate.text: - output.append( - { - "role": "assistant", - "content": [{"type": "text", "text": candidate.text}], - } - ) + message = _format_response_candidate(candidate) + if message is not None: + output.append(message) elif hasattr(response, "text") and response.text: output.append( @@ -479,40 +374,10 @@ def extract_gemini_web_search_count(response: Any) -> int: # Check for grounding_metadata in candidates if hasattr(response, "candidates") and response.candidates: for candidate in response.candidates: - if ( - hasattr(candidate, "grounding_metadata") - and candidate.grounding_metadata - ): - grounding_metadata = candidate.grounding_metadata - - # Check if web_search_queries exists and is non-empty - if hasattr(grounding_metadata, "web_search_queries"): - queries = grounding_metadata.web_search_queries - - if queries is not None and len(queries) > 0: - return 1 - - # Check if grounding_chunks exists and is non-empty - if hasattr(grounding_metadata, "grounding_chunks"): - chunks = grounding_metadata.grounding_chunks - - if chunks is not None and len(chunks) > 0: - return 1 - - # Also check for google_search or grounding in function call names - if hasattr(candidate, "content") and candidate.content: - if hasattr(candidate.content, "parts") and candidate.content.parts: - for part in candidate.content.parts: - if hasattr(part, "function_call") and part.function_call: - function_name = getattr( - part.function_call, "name", "" - ).lower() - - if ( - "google_search" in function_name - or "grounding" in function_name - ): - return 1 + if _has_grounding_metadata(candidate): + return 1 + if _has_search_function(candidate): + return 1 return 0 @@ -556,23 +421,7 @@ def _extract_usage_from_metadata( output_tokens=getattr(metadata, "candidates_token_count", 0), ) - # Add cache tokens if present (don't add if 0) - if hasattr(metadata, "cached_content_token_count"): - cache_tokens = metadata.cached_content_token_count - if cache_tokens and cache_tokens > 0: - usage["cache_read_input_tokens"] = cache_tokens - # Gemini counts cached_content_token_count inside prompt_token_count, so - # declare the accounting model rather than leaving ingestion to infer it. - # Under explicit context caching the two counts come from separate - # measurements and can disagree by a few percent, which makes inference - # from the token counts alone unreliable. - usage["cache_reporting_exclusive"] = False - - # Add reasoning tokens if present (don't add if 0) - if hasattr(metadata, "thoughts_token_count"): - reasoning_tokens = metadata.thoughts_token_count - if reasoning_tokens and reasoning_tokens > 0: - usage["reasoning_tokens"] = reasoning_tokens + _add_gemini_detail_usage(usage, metadata) # Capture raw usage metadata for backend processing # Serialize to dict here in the converter (not in utils) @@ -665,23 +514,7 @@ def extract_gemini_content_from_chunk( if hasattr(chunk, "candidates") and chunk.candidates: for candidate in chunk.candidates: - if hasattr(candidate, "content") and candidate.content: - if hasattr(candidate.content, "parts") and candidate.content.parts: - for part in candidate.content.parts: - if hasattr(part, "function_call") and part.function_call: - function_call = part.function_call - blocks.append( - _function_call_block( - function_call.name, - function_call.args, - getattr(function_call, "id", None), - ) - ) - continue - - block = _format_part(part) - if block is not None: - blocks.append(block) + blocks.extend(_format_candidate_parts(candidate, streaming=True)) if not blocks and hasattr(chunk, "text") and chunk.text: blocks.append({"type": "text", "text": chunk.text}) @@ -711,60 +544,8 @@ def format_gemini_streaming_output( } ] - # Handle list input if isinstance(accumulated_content, list): - content: List[FormattedContentItem] = [] - text_parts = [] - - for item in accumulated_content: - if isinstance(item, str): - # Legacy support: accumulate strings - text_parts.append(item) - elif isinstance(item, dict): - # New format: content blocks - if item.get("type") == "text": - text_parts.append(item.get("text", "")) - elif item.get("type") == "function": - # If we have accumulated text, add it first - if text_parts: - content.append( - { - "type": "text", - "text": "".join(text_parts), - } - ) - text_parts = [] - - # Add the function call - function_block: Dict[str, Any] = { - "type": "function", - "function": item.get("function", {}), - } - if isinstance(item.get("id"), str): - function_block["id"] = item["id"] - content.append(function_block) - else: - if text_parts: - content.append( - { - "type": "text", - "text": "".join(text_parts), - } - ) - text_parts = [] - - content.append(item) - - # Add any remaining text - if text_parts: - content.append( - { - "type": "text", - "text": "".join(text_parts), - } - ) - - # If we have content, return it + content = _merge_streaming_blocks(accumulated_content) if content: return [{"role": "assistant", "content": content}] @@ -788,3 +569,235 @@ def extract_gemini_embedding_token_count(response) -> Optional[int]: total += int(token_count) reported = True return total if reported else None + + +def _format_gemini_tool_part(plain: Dict[str, Any]) -> Optional[FormattedContentItem]: + # Gemini-native function_call/function_response parts are normalized to the + # provider-agnostic tool-call/tool-result blocks. Emitting the raw Gemini + # shape leaves them unrenderable and invisible to evals. + if "function_call" in plain: + call = to_plain(plain["function_call"]) or {} + block: Dict[str, Any] = { + "type": "function", + "function": { + "name": call.get("name"), + "arguments": call.get("args"), + }, + } + if isinstance(call.get("id"), str): + block["id"] = call["id"] + return block + if "function_response" in plain: + response = to_plain(plain["function_response"]) or {} + response_block: Dict[str, Any] = { + "type": "function", + "tool_name": response.get("name") or "", + "content": response.get("response"), + } + if isinstance(response.get("id"), str): + response_block["tool_use_id"] = response["id"] + return response_block + if not plain: + return None + key = next(iter(plain)) + fallback: Dict[str, Any] = {"type": key, key: to_plain(plain[key])} + return cast(FormattedContentItem, fallback) + + +def _format_message_content(content: Any) -> Any: + if isinstance(content, list): + return _format_parts_as_content_blocks(content) + if not isinstance(content, str): + return str(content) + return content + + +def _is_bare_part(plain: Any) -> bool: + return isinstance(plain, dict) and bool( + _BARE_PART_KEYS.intersection(normalize_part_keys(plain)) + ) + + +def _object_message_role(item: Any) -> str: + role = getattr(item, "role", "user") if hasattr(item, "role") else "user" + # Ensure role is a string + return role if isinstance(role, str) else "user" + + +def _format_response_part(part: Any) -> Optional[FormattedContentItem]: + # Checked ahead of _format_part so loosely-specced MagicMock + # fixtures (which report a truthy .text but aren't real Part + # objects) still resolve to a text block — must stay ahead of + # the _format_part delegation below. + text = getattr(part, "text", None) + if isinstance(text, str) and text: + # `is True` so loosely-specced mocks (whose .thought is a + # truthy Mock) still land on the text branch below. + if getattr(part, "thought", None) is True: + return {"type": "thinking", "thinking": text} + else: + return {"type": "text", "text": text} + + return _format_streaming_part(part) + + +def _format_response_candidate(candidate: Any) -> Optional[FormattedMessage]: + if hasattr(candidate, "content") and candidate.content: + content = _format_candidate_parts(candidate, streaming=False) + + if content: + return { + "role": "assistant", + "content": content, + } + + elif hasattr(candidate, "text") and candidate.text: + return { + "role": "assistant", + "content": [{"type": "text", "text": candidate.text}], + } + + return None + + +def _has_grounding_metadata(candidate: Any) -> bool: + if hasattr(candidate, "grounding_metadata") and candidate.grounding_metadata: + grounding_metadata = candidate.grounding_metadata + + # Check if web_search_queries exists and is non-empty + if hasattr(grounding_metadata, "web_search_queries"): + queries = grounding_metadata.web_search_queries + + if queries is not None and len(queries) > 0: + return True + + # Check if grounding_chunks exists and is non-empty + if hasattr(grounding_metadata, "grounding_chunks"): + chunks = grounding_metadata.grounding_chunks + + if chunks is not None and len(chunks) > 0: + return True + + return False + + +def _has_search_function(candidate: Any) -> bool: + # Also check for google_search or grounding in function call names + if hasattr(candidate, "content") and candidate.content: + if hasattr(candidate.content, "parts") and candidate.content.parts: + for part in candidate.content.parts: + if _is_search_function_part(part): + return True + return False + + +def _add_gemini_detail_usage(usage: TokenUsage, metadata: Any) -> None: + # Add cache tokens if present (don't add if 0) + if hasattr(metadata, "cached_content_token_count"): + cache_tokens = metadata.cached_content_token_count + if cache_tokens and cache_tokens > 0: + usage["cache_read_input_tokens"] = cache_tokens + # Gemini counts cached_content_token_count inside prompt_token_count, so + # declare the accounting model rather than leaving ingestion to infer it. + # Under explicit context caching the two counts come from separate + # measurements and can disagree by a few percent, which makes inference + # from the token counts alone unreliable. + usage["cache_reporting_exclusive"] = False + + # Add reasoning tokens if present (don't add if 0) + if hasattr(metadata, "thoughts_token_count"): + reasoning_tokens = metadata.thoughts_token_count + if reasoning_tokens and reasoning_tokens > 0: + usage["reasoning_tokens"] = reasoning_tokens + + +def _format_streaming_part(part: Any) -> Optional[FormattedContentItem]: + if hasattr(part, "function_call") and part.function_call: + function_call = part.function_call + return _function_call_block( + function_call.name, + function_call.args, + getattr(function_call, "id", None), + ) + + return _format_part(part) + + +def _merge_streaming_blocks( + accumulated_content: List[Any], +) -> List[FormattedContentItem]: + content: List[FormattedContentItem] = [] + text_parts = [] + + for item in accumulated_content: + if isinstance(item, str): + # Legacy support: accumulate strings + text_parts.append(item) + elif isinstance(item, dict): + # New format: content blocks + if item.get("type") == "text": + text_parts.append(item.get("text", "")) + else: + # If we have accumulated text, add it first + if text_parts: + content.append( + { + "type": "text", + "text": "".join(text_parts), + } + ) + text_parts = [] + + if item.get("type") == "function": + content.append(_format_streaming_function(item)) + else: + content.append(item) + + # Add any remaining text + if text_parts: + content.append( + { + "type": "text", + "text": "".join(text_parts), + } + ) + + return content + + +def _format_streaming_function(item: Dict[str, Any]) -> FormattedContentItem: + # Add the function call + function_block: Dict[str, Any] = { + "type": "function", + "function": item.get("function", {}), + } + if isinstance(item.get("id"), str): + function_block["id"] = item["id"] + return function_block + + +def _format_candidate_parts( + candidate: Any, streaming: bool +) -> List[FormattedContentItem]: + content: List[FormattedContentItem] = [] + if not (hasattr(candidate, "content") and candidate.content): + return content + if hasattr(candidate.content, "parts") and candidate.content.parts: + for part in candidate.content.parts: + block = ( + _format_streaming_part(part) + if streaming + else _format_response_part(part) + ) + if block is not None: + content.append(block) + return content + + +def _is_search_function_part(part: Any) -> bool: + if hasattr(part, "function_call") and part.function_call: + function_name = getattr(part.function_call, "name", "").lower() + + if "google_search" in function_name or "grounding" in function_name: + return True + return False diff --git a/posthog/ai/langchain/callbacks.py b/posthog/ai/langchain/callbacks.py index f670d6794..0589f472c 100644 --- a/posthog/ai/langchain/callbacks.py +++ b/posthog/ai/langchain/callbacks.py @@ -39,6 +39,7 @@ ToolMessage, ) from langchain_core.outputs import ChatGeneration, LLMResult +from langchain_core.outputs import Generation as _Generation from pydantic import BaseModel from posthog import setup @@ -641,33 +642,7 @@ def _capture_generation( parent_run_id: Optional[UUID] = None, include_parent_id: bool = True, ): - # The served tier comes from the response, because a requested tier can be refused. - model_params = run.model_params - served_tier = None - if isinstance(output, LLMResult) and isinstance(output.llm_output, dict): - served_tier = output.llm_output.get("service_tier") - if served_tier is not None: - model_params = {**(model_params or {}), "service_tier": served_tier} - - event_properties = { - "$ai_trace_id": trace_id, - "$ai_span_id": run_id, - "$ai_span_name": run.name, - "$ai_provider": run.provider, - "$ai_model": run.model, - "$ai_model_parameters": model_params, - # The explicit served-tier signal cost processing prices from. - **({"$ai_service_tier": served_tier} if served_tier is not None else {}), - "$ai_input": with_privacy_mode( - self._ph_client, - self._privacy_mode, - finalize_ai_content(run.input, self._ph_client), - ), - "$ai_http_status": 200, - "$ai_latency": run.latency, - "$ai_base_url": run.base_url, - "$ai_framework": "langchain", - } + event_properties = self._generation_properties(trace_id, run_id, run, output) if include_parent_id: event_properties["$ai_parent_id"] = parent_run_id @@ -696,60 +671,11 @@ def _capture_generation( output, event_properties ) else: - # Add usage - usage = _parse_usage(output, run.provider, run.model) - # Omitted when the provider never reported a count: absent means - # unknown, 0 is a report of nothing. - if usage.input_tokens is not None: - event_properties["$ai_input_tokens"] = usage.input_tokens - if usage.output_tokens is not None: - event_properties["$ai_output_tokens"] = usage.output_tokens - if usage.cache_write_tokens is not None: - event_properties["$ai_cache_creation_input_tokens"] = ( - usage.cache_write_tokens - ) - if ( - usage.cache_write_5m_tokens is not None - and usage.cache_write_1h_tokens is not None - ): - event_properties["$ai_cache_creation_5m_input_tokens"] = ( - usage.cache_write_5m_tokens - ) - event_properties["$ai_cache_creation_1h_input_tokens"] = ( - usage.cache_write_1h_tokens - ) - if usage.cache_read_tokens is not None: - event_properties["$ai_cache_read_input_tokens"] = ( - usage.cache_read_tokens - ) - if usage.reasoning_tokens is not None: - event_properties["$ai_reasoning_tokens"] = usage.reasoning_tokens - if usage.web_search_count is not None and usage.web_search_count > 0: - event_properties["$ai_web_search_count"] = usage.web_search_count - - # Generation results - generation_result = output.generations[-1] - if isinstance(generation_result[-1], ChatGeneration): - completions = [ - _convert_message_to_dict(cast(ChatGeneration, generation).message) - for generation in generation_result - ] - else: - completions = [ - _extract_raw_response(generation) - for generation in generation_result - ] - event_properties["$ai_output_choices"] = with_privacy_mode( - self._ph_client, - self._privacy_mode, - finalize_ai_content(completions, self._ph_client), + _add_generation_usage( + event_properties, _parse_usage(output, run.provider, run.model) ) - # Extract the stop reason from the generation and its metadata - if output.generations and output.generations[-1]: - stop_reason = _extract_stop_reason(output.generations[-1][-1]) - if stop_reason is not None: - event_properties["$ai_stop_reason"] = stop_reason + self._add_generation_output(event_properties, output) _capture_ai_event( self._ph_client, @@ -759,6 +685,68 @@ def _capture_generation( groups=self._groups, ) + def _generation_properties( + self, + trace_id: Any, + run_id: UUID, + run: GenerationMetadata, + output: Union[LLMResult, BaseException], + ) -> dict[str, Any]: + # The served tier comes from the response, because a requested tier can be refused. + model_params = run.model_params + served_tier = None + if isinstance(output, LLMResult) and isinstance(output.llm_output, dict): + served_tier = output.llm_output.get("service_tier") + if served_tier is not None: + model_params = {**(model_params or {}), "service_tier": served_tier} + + event_properties = { + "$ai_trace_id": trace_id, + "$ai_span_id": run_id, + "$ai_span_name": run.name, + "$ai_provider": run.provider, + "$ai_model": run.model, + "$ai_model_parameters": model_params, + # The explicit served-tier signal cost processing prices from. + **({"$ai_service_tier": served_tier} if served_tier is not None else {}), + "$ai_input": with_privacy_mode( + self._ph_client, + self._privacy_mode, + finalize_ai_content(run.input, self._ph_client), + ), + "$ai_http_status": 200, + "$ai_latency": run.latency, + "$ai_base_url": run.base_url, + "$ai_framework": "langchain", + } + return event_properties + + def _add_generation_output( + self, event_properties: dict[str, Any], output: LLMResult + ) -> None: + # Generation results + generation_result = output.generations[-1] + if isinstance(generation_result[-1], ChatGeneration): + completions = [ + _convert_message_to_dict(cast(ChatGeneration, generation).message) + for generation in generation_result + ] + else: + completions = [ + _extract_raw_response(generation) for generation in generation_result + ] + event_properties["$ai_output_choices"] = with_privacy_mode( + self._ph_client, + self._privacy_mode, + finalize_ai_content(completions, self._ph_client), + ) + + # Extract the stop reason from the generation and its metadata + if output.generations and output.generations[-1]: + stop_reason = _extract_stop_reason(output.generations[-1][-1]) + if stop_reason is not None: + event_properties["$ai_stop_reason"] = stop_reason + def _log_debug_event( self, event_name: str, @@ -822,6 +810,18 @@ def _convert_lc_tool_calls_to_oai( def _convert_message_to_dict(message: BaseMessage) -> dict[str, Any]: + message_dict = _message_role_and_content(message) + + if message.additional_kwargs: + message_dict.update(message.additional_kwargs) + + if "content" in message_dict and not message_dict["content"]: + message_dict["content"] = "" + + return message_dict + + +def _message_role_and_content(message: BaseMessage) -> dict[str, Any]: # assistant message if isinstance(message, HumanMessage): message_dict = {"role": "user", "content": message.content} @@ -840,12 +840,6 @@ def _convert_message_to_dict(message: BaseMessage) -> dict[str, Any]: else: message_dict = {"role": message.type, "content": str(message.content)} - if message.additional_kwargs: - message_dict.update(message.additional_kwargs) - - if "content" in message_dict and not message_dict["content"]: - message_dict["content"] = "" - return message_dict @@ -885,63 +879,9 @@ def _parse_usage_model( if isinstance(usage, BaseModel): usage = usage.__dict__ - conversion_list = [ - # https://pypi.org/project/langchain-anthropic/ (works also for Bedrock-Anthropic) - ("input_tokens", "input"), - ("output_tokens", "output"), - ("cache_creation_input_tokens", "cache_write"), - ("cache_read_input_tokens", "cache_read"), - # https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/get-token-count - ("prompt_token_count", "input"), - ("candidates_token_count", "output"), - ("cached_content_token_count", "cache_read"), - ("thoughts_token_count", "reasoning"), - # Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/monitoring-cw.html#runtime-cloudwatch-metrics - ("inputTokenCount", "input"), - ("outputTokenCount", "output"), - ("cacheCreationInputTokenCount", "cache_write"), - ("cacheReadInputTokenCount", "cache_read"), - # Bedrock Anthropic - ("prompt_tokens", "input"), - ("completion_tokens", "output"), - ("cache_creation_input_tokens", "cache_write"), - ("cache_read_input_tokens", "cache_read"), - # langchain-ibm https://pypi.org/project/langchain-ibm/ - ("input_token_count", "input"), - ("generated_token_count", "output"), - ] - - parsed_usage = {} - for model_key, type_key in conversion_list: - if model_key in usage: - captured_count = usage[model_key] - final_count = ( - sum(captured_count) - if isinstance(captured_count, list) - else captured_count - ) # For Bedrock, the token count is a list when streamed - - parsed_usage[type_key] = final_count - - # Caching (OpenAI & langchain 0.3.9+) - if "input_token_details" in usage and isinstance( - usage["input_token_details"], dict - ): - input_token_details = usage["input_token_details"] - parsed_usage["cache_write"] = input_token_details.get("cache_creation") - cache_write_ttl = _extract_cache_creation_ttl_breakdown(input_token_details) - if cache_write_ttl is not None: - cache_write_5m, cache_write_1h = cache_write_ttl - parsed_usage["cache_write_5m"] = cache_write_5m - parsed_usage["cache_write_1h"] = cache_write_1h - parsed_usage["cache_write"] = cache_write_5m + cache_write_1h - parsed_usage["cache_read"] = input_token_details.get("cache_read") + parsed_usage = _flat_usage_counts(usage) - # Reasoning (OpenAI & langchain 0.3.9+) - if "output_token_details" in usage and isinstance( - usage["output_token_details"], dict - ): - parsed_usage["reasoning"] = usage["output_token_details"].get("reasoning") + _add_usage_details(parsed_usage, usage) # Anthropic reports the exact number of web searches of a response web_search_requests = _extract_web_search_requests(usage) @@ -964,6 +904,13 @@ def _parse_usage_model( cache_write_1h_tokens=parsed_usage.get("cache_write_1h"), web_search_count=parsed_usage.get("web_search"), ) + _subtract_anthropic_cache(normalized_usage, provider, model) + return normalized_usage + + +def _subtract_anthropic_cache( + normalized_usage: ModelUsage, provider: Optional[str], model: Optional[str] +) -> None: # For Anthropic providers, LangChain reports input_tokens as the sum of all input tokens. # Our cost calculation expects them to be separate for Anthropic, so we subtract cache tokens. # Both cache_read and cache_write tokens should be subtracted since Anthropic's raw API @@ -971,11 +918,7 @@ def _parse_usage_model( # For other providers (OpenAI, etc.), input_tokens already excludes cache tokens as expected. # Match logic consistent with plugin-server: exact match on provider, substring # match on model, or a Claude model under another provider (e.g. Vertex AI) - is_anthropic = False - if provider and provider.lower() == "anthropic": - is_anthropic = True - elif model and ("anthropic" in model.lower() or model.lower().startswith("claude")): - is_anthropic = True + is_anthropic = _is_anthropic_model(provider, model) if is_anthropic and normalized_usage.input_tokens: cache_tokens = (normalized_usage.cache_read_tokens or 0) + ( @@ -985,7 +928,70 @@ def _parse_usage_model( normalized_usage.input_tokens = max( normalized_usage.input_tokens - cache_tokens, 0 ) - return normalized_usage + + +def _add_usage_details(parsed_usage: dict, usage: dict) -> None: + # Caching (OpenAI & langchain 0.3.9+) + if "input_token_details" in usage and isinstance( + usage["input_token_details"], dict + ): + input_token_details = usage["input_token_details"] + parsed_usage["cache_write"] = input_token_details.get("cache_creation") + cache_write_ttl = _extract_cache_creation_ttl_breakdown(input_token_details) + if cache_write_ttl is not None: + cache_write_5m, cache_write_1h = cache_write_ttl + parsed_usage["cache_write_5m"] = cache_write_5m + parsed_usage["cache_write_1h"] = cache_write_1h + parsed_usage["cache_write"] = cache_write_5m + cache_write_1h + parsed_usage["cache_read"] = input_token_details.get("cache_read") + + # Reasoning (OpenAI & langchain 0.3.9+) + if "output_token_details" in usage and isinstance( + usage["output_token_details"], dict + ): + parsed_usage["reasoning"] = usage["output_token_details"].get("reasoning") + + +def _flat_usage_counts(usage: dict) -> dict: + conversion_list = [ + # https://pypi.org/project/langchain-anthropic/ (works also for Bedrock-Anthropic) + ("input_tokens", "input"), + ("output_tokens", "output"), + ("cache_creation_input_tokens", "cache_write"), + ("cache_read_input_tokens", "cache_read"), + # https://cloud.google.com/vertex-ai/generative-ai/docs/multimodal/get-token-count + ("prompt_token_count", "input"), + ("candidates_token_count", "output"), + ("cached_content_token_count", "cache_read"), + ("thoughts_token_count", "reasoning"), + # Bedrock: https://docs.aws.amazon.com/bedrock/latest/userguide/monitoring-cw.html#runtime-cloudwatch-metrics + ("inputTokenCount", "input"), + ("outputTokenCount", "output"), + ("cacheCreationInputTokenCount", "cache_write"), + ("cacheReadInputTokenCount", "cache_read"), + # Bedrock Anthropic + ("prompt_tokens", "input"), + ("completion_tokens", "output"), + ("cache_creation_input_tokens", "cache_write"), + ("cache_read_input_tokens", "cache_read"), + # langchain-ibm https://pypi.org/project/langchain-ibm/ + ("input_token_count", "input"), + ("generated_token_count", "output"), + ] + + parsed_usage = {} + for model_key, type_key in conversion_list: + if model_key in usage: + captured_count = usage[model_key] + final_count = ( + sum(captured_count) + if isinstance(captured_count, list) + else captured_count + ) # For Bedrock, the token count is a list when streamed + + parsed_usage[type_key] = final_count + + return parsed_usage def _parse_usage( @@ -1022,45 +1028,9 @@ def _parse_usage( llm_usage = _parse_usage_model(generation["usage"], provider, model) break - for generation_chunk in generation: - if generation_chunk.generation_info and ( - "usage_metadata" in generation_chunk.generation_info - ): - llm_usage = _parse_usage_model( - generation_chunk.generation_info["usage_metadata"], - provider, - model, - ) - break - - message_chunk = getattr(generation_chunk, "message", {}) - response_metadata = getattr(message_chunk, "response_metadata", {}) - - bedrock_anthropic_usage = ( - response_metadata.get("usage", None) # for Bedrock-Anthropic - if isinstance(response_metadata, dict) - else None - ) - bedrock_titan_usage = ( - response_metadata.get( - "amazon-bedrock-invocationMetrics", None - ) # for Bedrock-Titan - if isinstance(response_metadata, dict) - else None - ) - # LangChain's standard usage, which also covers Ollama - usage_metadata = getattr(message_chunk, "usage_metadata", None) - - # Prefer the standard usage: its input_tokens include cache reads - # and writes, as the Anthropic subtraction in _parse_usage_model - # expects. ChatAnthropic also copies the raw API usage, whose - # input_tokens already exclude the cache, into response_metadata. - chunk_usage = ( - usage_metadata or bedrock_anthropic_usage or bedrock_titan_usage - ) - if chunk_usage: - llm_usage = _parse_usage_model(chunk_usage, provider, model) - break + parsed = _generation_usage(generation, provider, model) + if parsed is not None: + llm_usage = parsed if llm_usage.web_search_count is None: llm_usage.web_search_count = web_search_count @@ -1118,3 +1088,88 @@ def _stringify_exception(exception: BaseException) -> str: if description: return f"{exception.__class__.__name__}: {description}" return exception.__class__.__name__ + + +def _add_generation_usage( + event_properties: dict[str, Any], usage: "ModelUsage" +) -> None: + # Omitted when the provider never reported a count: absent means unknown, + # 0 is a report of nothing. + for field, property_name in ( + ("input_tokens", "$ai_input_tokens"), + ("output_tokens", "$ai_output_tokens"), + ("cache_write_tokens", "$ai_cache_creation_input_tokens"), + ): + value = getattr(usage, field) + if value is not None: + event_properties[property_name] = value + if ( + usage.cache_write_5m_tokens is not None + and usage.cache_write_1h_tokens is not None + ): + event_properties["$ai_cache_creation_5m_input_tokens"] = ( + usage.cache_write_5m_tokens + ) + event_properties["$ai_cache_creation_1h_input_tokens"] = ( + usage.cache_write_1h_tokens + ) + for field, property_name in ( + ("cache_read_tokens", "$ai_cache_read_input_tokens"), + ("reasoning_tokens", "$ai_reasoning_tokens"), + ): + value = getattr(usage, field) + if value is not None: + event_properties[property_name] = value + if usage.web_search_count is not None and usage.web_search_count > 0: + event_properties["$ai_web_search_count"] = usage.web_search_count + + +def _generation_usage( + generation: list[_Generation], provider: Optional[str], model: Optional[str] +) -> Optional[ModelUsage]: + for generation_chunk in generation: + if generation_chunk.generation_info and ( + "usage_metadata" in generation_chunk.generation_info + ): + return _parse_usage_model( + generation_chunk.generation_info["usage_metadata"], + provider, + model, + ) + + message_chunk = getattr(generation_chunk, "message", {}) + response_metadata = getattr(message_chunk, "response_metadata", {}) + + bedrock_anthropic_usage = ( + response_metadata.get("usage", None) # for Bedrock-Anthropic + if isinstance(response_metadata, dict) + else None + ) + bedrock_titan_usage = ( + response_metadata.get( + "amazon-bedrock-invocationMetrics", None + ) # for Bedrock-Titan + if isinstance(response_metadata, dict) + else None + ) + # LangChain's standard usage, which also covers Ollama + usage_metadata = getattr(message_chunk, "usage_metadata", None) + + # Prefer the standard usage: its input_tokens include cache reads + # and writes, as the Anthropic subtraction in _parse_usage_model + # expects. ChatAnthropic also copies the raw API usage, whose + # input_tokens already exclude the cache, into response_metadata. + chunk_usage = usage_metadata or bedrock_anthropic_usage or bedrock_titan_usage + if chunk_usage: + return _parse_usage_model(chunk_usage, provider, model) + return None + + +def _is_anthropic_model(provider: Optional[str], model: Optional[str]) -> bool: + is_anthropic = False + if provider and provider.lower() == "anthropic": + is_anthropic = True + elif model and ("anthropic" in model.lower() or model.lower().startswith("claude")): + is_anthropic = True + + return is_anthropic diff --git a/posthog/ai/langchain/middleware.py b/posthog/ai/langchain/middleware.py index 8cf282a87..e434c8473 100644 --- a/posthog/ai/langchain/middleware.py +++ b/posthog/ai/langchain/middleware.py @@ -252,22 +252,9 @@ def _finish_model( if not isinstance(run, GenerationMetadata): return - output: LLMResult | BaseException - if isinstance(result, BaseException): - output = result - else: - generations = [ - ChatGeneration(message=message) - for message in result.result - if isinstance(message, AIMessage) - ] - if not generations: - return - metadata = generations[-1].message.response_metadata - output = LLMResult( - generations=[generations], - llm_output=metadata if isinstance(metadata, dict) else None, - ) + output = _model_output(result) + if output is None: + return root_id = _root_id(state) self._safely_call( @@ -368,4 +355,24 @@ def _public_state(state: Mapping[str, Any]) -> dict[str, Any]: } +def _model_output( + result: ModelResponse[Any] | BaseException, +) -> LLMResult | BaseException | None: + if isinstance(result, BaseException): + return result + else: + generations = [ + ChatGeneration(message=message) + for message in result.result + if isinstance(message, AIMessage) + ] + if not generations: + return None + metadata = generations[-1].message.response_metadata + return LLMResult( + generations=[generations], + llm_output=metadata if isinstance(metadata, dict) else None, + ) + + __all__ = ["PostHogMiddleware"] diff --git a/posthog/ai/media.py b/posthog/ai/media.py index 6c4b4f009..ef227061e 100644 --- a/posthog/ai/media.py +++ b/posthog/ai/media.py @@ -16,18 +16,22 @@ def to_plain(obj: Any) -> Any: if isinstance(obj, dict): - return {k: v for k, v in obj.items() if v is not None} + return _without_none(obj) model_dump = getattr(obj, "model_dump", None) if callable(model_dump): try: return model_dump(exclude_none=True) except TypeError: - return {k: v for k, v in model_dump().items() if v is not None} + return _without_none(model_dump()) if dataclasses.is_dataclass(obj) and not isinstance(obj, type): - return {k: v for k, v in dataclasses.asdict(obj).items() if v is not None} + return _without_none(dataclasses.asdict(obj)) return obj +def _without_none(value: dict) -> dict: + return {k: v for k, v in value.items() if v is not None} + + def ensure_serializable(obj: Any) -> Any: """Ensure an object is JSON-serializable, converting to str as fallback. diff --git a/posthog/ai/openai/_streaming.py b/posthog/ai/openai/_streaming.py index 3a7d059b4..49234bee0 100644 --- a/posthog/ai/openai/_streaming.py +++ b/posthog/ai/openai/_streaming.py @@ -26,10 +26,7 @@ class _ResponsesStreamState: def process_chunk(self, chunk: Any) -> None: response = getattr(chunk, "response", None) - if response and self.model is None and hasattr(response, "model"): - self.model = response.model - if response and self.service_tier is None: - self.service_tier = getattr(response, "service_tier", None) + self._record_response_metadata(response) chunk_usage = extract_openai_usage_from_chunk(chunk, "responses") if chunk_usage: @@ -46,6 +43,12 @@ def process_chunk(self, chunk: Any) -> None: if stop_reason is not None: self.stop_reason = stop_reason + def _record_response_metadata(self, response: Any) -> None: + if response and self.model is None and hasattr(response, "model"): + self.model = response.model + if response and self.service_tier is None: + self.service_tier = getattr(response, "service_tier", None) + @dataclass class _ChoiceStream: @@ -82,25 +85,28 @@ def process_chunk(self, chunk: Any) -> None: merge_usage_stats(self.usage_stats, chunk_usage) for choice in getattr(chunk, "choices", None) or []: - index = getattr(choice, "index", None) - if not isinstance(index, int): - index = 0 - view = _ChoiceStream([choice]) - - content = extract_openai_content_from_chunk(view, "chat") - if content is not None: - self._outputs.setdefault(index, []).append(content) - - chunk_tool_calls = extract_openai_tool_calls_from_chunk(view) - if chunk_tool_calls: - accumulate_openai_tool_calls( - self._tool_calls.setdefault(index, {}), chunk_tool_calls - ) - - # Like a non-streamed completion, the stop reason is the first choice's. - finish_reason = getattr(choice, "finish_reason", None) - if finish_reason is not None and index == 0: - self.stop_reason = finish_reason + self._process_choice(choice) + + def _process_choice(self, choice: Any) -> None: + index = getattr(choice, "index", None) + if not isinstance(index, int): + index = 0 + view = _ChoiceStream([choice]) + + content = extract_openai_content_from_chunk(view, "chat") + if content is not None: + self._outputs.setdefault(index, []).append(content) + + chunk_tool_calls = extract_openai_tool_calls_from_chunk(view) + if chunk_tool_calls: + accumulate_openai_tool_calls( + self._tool_calls.setdefault(index, {}), chunk_tool_calls + ) + + # Like a non-streamed completion, the stop reason is the first choice's. + finish_reason = getattr(choice, "finish_reason", None) + if finish_reason is not None and index == 0: + self.stop_reason = finish_reason @property def tool_calls(self) -> Optional[List[Dict[str, Any]]]: diff --git a/posthog/ai/openai/openai_converter.py b/posthog/ai/openai/openai_converter.py index 6faf1a6e9..72f6dc457 100644 --- a/posthog/ai/openai/openai_converter.py +++ b/posthog/ai/openai/openai_converter.py @@ -37,39 +37,7 @@ def _format_responses_output_items(items: Any) -> List[FormattedContentItem]: item_type = _item_attr(item, "type") if item_type == "message": - message_content = _item_attr(item, "content") - if isinstance(message_content, list): - for content_item in message_content: - content_item_type = _item_attr(content_item, "type") - content_item_text = _item_attr(content_item, "text") - content_item_refusal = _item_attr(content_item, "refusal") - - if content_item_type == "output_text" and ( - content_item_text is not None - ): - content.append({"type": "text", "text": content_item_text}) - - elif content_item_type == "refusal" and ( - content_item_refusal is not None - ): - content.append( - {"type": "refusal", "refusal": content_item_refusal} - ) - - elif content_item_text is not None: - content.append({"type": "text", "text": content_item_text}) - - elif content_item_type == "input_image" and ( - _item_attr(content_item, "image_url") is not None - ): - image_content: FormattedImageContent = { - "type": "image", - "image": _item_attr(content_item, "image_url"), - } - content.append(image_content) - - elif message_content is not None: - content.append({"type": "text", "text": str(message_content)}) + content.extend(_format_response_message_content(item)) elif item_type == "function_call": call_id = _item_attr(item, "call_id") @@ -139,63 +107,9 @@ def format_openai_response(response: Any) -> List[FormattedMessage]: # Handle Chat Completions response format if hasattr(response, "choices"): - content: List[FormattedContentItem] = [] - role = "assistant" - - for choice in response.choices: - if hasattr(choice, "message") and choice.message: - if choice.message.role: - role = choice.message.role - - if choice.message.content: - content.append( - { - "type": "text", - "text": choice.message.content, - } - ) - - # A refused request carries its text in `refusal` and leaves `content` empty - refusal = getattr(choice.message, "refusal", None) - if refusal: - content.append({"type": "refusal", "refusal": refusal}) - - if hasattr(choice.message, "tool_calls") and choice.message.tool_calls: - for tool_call in choice.message.tool_calls: - # Custom tools carry their free-form text in `custom.input`, - # with no `function` object. - is_custom = getattr(tool_call, "type", None) == "custom" - call = getattr( - tool_call, "custom" if is_custom else "function", None - ) - if call is None: - continue - content.append( - { - "type": "function", - "id": tool_call.id, - "function": { - "name": call.name, - "arguments": call.input - if is_custom - else call.arguments, - }, - } - ) - - # Handle audio output (gpt-4o-audio-preview) - if hasattr(choice.message, "audio") and choice.message.audio: - # Convert Pydantic model to dict to capture all fields from OpenAI - audio_dict = choice.message.audio.model_dump() - content.append({"type": "audio", **audio_dict}) - - if content: - output.append( - { - "role": role, - "content": content, - } - ) + message = _format_chat_completion_choices(response.choices) + if message is not None: + output.append(message) # Handle Responses API format if hasattr(response, "output"): @@ -233,53 +147,13 @@ def format_openai_input( if messages is not None: for msg in messages: - plain = to_plain(msg) - if not isinstance(plain, dict): - plain = {"role": "user", "content": str(plain)} - - formatted: Dict[str, Any] = { - "role": plain.get("role", "user"), - "content": plain.get("content"), - } - - for key in ("tool_calls", "tool_call_id", "name", "audio", "refusal"): - if plain.get(key) is not None: - formatted[key] = ( - to_plain(plain[key]) if key == "audio" else plain[key] - ) - - formatted_messages.append(cast(FormattedMessage, formatted)) + formatted_messages.append(_format_chat_input_message(msg)) # Handle Responses API format if input_data is not None: if isinstance(input_data, list): for item in input_data: - if not isinstance(item, (dict, str)): - item = to_plain(item) - - if ( - isinstance(item, dict) - and "type" in item - and "role" not in item - and "content" not in item - ): - formatted_messages.append(cast(FormattedMessage, to_plain(item))) - continue - - role = "user" - content = "" - - if isinstance(item, dict): - role = item.get("role", "user") - content = item.get("content", "") - - elif isinstance(item, str): - content = item - - else: - content = str(item) - - formatted_messages.append({"role": role, "content": content}) + formatted_messages.append(_format_response_input_item(item)) elif isinstance(input_data, str): formatted_messages.append({"role": "user", "content": input_data}) @@ -367,93 +241,15 @@ def extract_openai_web_search_count(response: Any) -> int: Number of web search requests (exact count or binary 1/0) """ - # Priority 1: Check for exact count in Responses API output - if hasattr(response, "output"): - web_search_count = 0 - - for item in response.output: - if hasattr(item, "type") and item.type == "web_search_call": - web_search_count += 1 - - web_search_count = max(0, web_search_count) - - if web_search_count > 0: - return web_search_count - - # Priority 2: Binary detection (returns 1 or 0) - - # Check root-level indicators (Perplexity) - if hasattr(response, "citations"): - citations = getattr(response, "citations") - - if citations and len(citations) > 0: - return 1 - - if hasattr(response, "search_results"): - search_results = getattr(response, "search_results") - - if search_results and len(search_results) > 0: - return 1 - - if hasattr(response, "usage") and hasattr(response.usage, "search_context_size"): - if response.usage.search_context_size: - return 1 - - # Check for url_citation annotations in choices (Chat Completions) - if hasattr(response, "choices"): - for choice in response.choices: - # Check message.annotations (non-streaming or final chunk) - if hasattr(choice, "message") and hasattr(choice.message, "annotations"): - annotations = choice.message.annotations - - if annotations: - for annotation in annotations: - # Support both dict and object formats - annotation_type = ( - annotation.get("type") - if isinstance(annotation, dict) - else getattr(annotation, "type", None) - ) - - if annotation_type == "url_citation": - return 1 - - # Check delta.annotations (streaming chunks) - if hasattr(choice, "delta") and hasattr(choice.delta, "annotations"): - annotations = choice.delta.annotations - - if annotations: - for annotation in annotations: - # Support both dict and object formats - annotation_type = ( - annotation.get("type") - if isinstance(annotation, dict) - else getattr(annotation, "type", None) - ) - - if annotation_type == "url_citation": - return 1 - - # Check for url_citation annotations in output (Responses API) - if hasattr(response, "output"): - for item in response.output: - if hasattr(item, "content") and isinstance(item.content, list): - for content_item in item.content: - if hasattr(content_item, "annotations"): - annotations = content_item.annotations - - if annotations: - for annotation in annotations: - # Support both dict and object formats - annotation_type = ( - annotation.get("type") - if isinstance(annotation, dict) - else getattr(annotation, "type", None) - ) - - if annotation_type == "url_citation": - return 1 - + count = _count_response_web_search_calls(response) + if count > 0: + return count + if _has_root_web_search(response): + return 1 + if _has_choice_web_citations(response): + return 1 + if _has_response_web_citations(response): + return 1 return 0 @@ -482,49 +278,7 @@ def extract_openai_usage_from_response(response: Any) -> TokenUsage: if not hasattr(response, "usage") or not response.usage: return TokenUsage() - cached_tokens = None - input_tokens = None - output_tokens = None - reasoning_tokens = None - - # Responses API format - if hasattr(response.usage, "input_tokens"): - input_tokens = response.usage.input_tokens - if hasattr(response.usage, "output_tokens"): - output_tokens = response.usage.output_tokens - if hasattr(response.usage, "input_tokens_details") and hasattr( - response.usage.input_tokens_details, "cached_tokens" - ): - cached_tokens = response.usage.input_tokens_details.cached_tokens - if hasattr(response.usage, "output_tokens_details") and hasattr( - response.usage.output_tokens_details, "reasoning_tokens" - ): - reasoning_tokens = response.usage.output_tokens_details.reasoning_tokens - - # Chat Completions format - if hasattr(response.usage, "prompt_tokens"): - input_tokens = response.usage.prompt_tokens - if hasattr(response.usage, "completion_tokens"): - output_tokens = response.usage.completion_tokens - if hasattr(response.usage, "prompt_tokens_details") and hasattr( - response.usage.prompt_tokens_details, "cached_tokens" - ): - cached_tokens = response.usage.prompt_tokens_details.cached_tokens - if hasattr(response.usage, "completion_tokens_details") and hasattr( - response.usage.completion_tokens_details, "reasoning_tokens" - ): - reasoning_tokens = response.usage.completion_tokens_details.reasoning_tokens - - result = TokenUsage() - if input_tokens is not None: - result["input_tokens"] = input_tokens - if output_tokens is not None: - result["output_tokens"] = output_tokens - - if cached_tokens is not None and cached_tokens > 0: - result["cache_read_input_tokens"] = cached_tokens - if reasoning_tokens is not None and reasoning_tokens > 0: - result["reasoning_tokens"] = reasoning_tokens + result = _extract_openai_response_token_usage(response.usage) web_search_count = extract_openai_web_search_count(response) if web_search_count > 0: @@ -577,74 +331,11 @@ def extract_openai_usage_from_chunk( if not hasattr(chunk, "usage") or not chunk.usage: return usage - # Chat Completions API uses prompt_tokens and completion_tokens - # Standardize to input_tokens and output_tokens - usage["input_tokens"] = getattr(chunk.usage, "prompt_tokens", 0) - usage["output_tokens"] = getattr(chunk.usage, "completion_tokens", 0) - - # Handle cached tokens - if hasattr(chunk.usage, "prompt_tokens_details") and hasattr( - chunk.usage.prompt_tokens_details, "cached_tokens" - ): - cached = chunk.usage.prompt_tokens_details.cached_tokens - if cached is not None: - usage["cache_read_input_tokens"] = cached - - # Handle reasoning tokens - if hasattr(chunk.usage, "completion_tokens_details") and hasattr( - chunk.usage.completion_tokens_details, "reasoning_tokens" - ): - reasoning = chunk.usage.completion_tokens_details.reasoning_tokens - if reasoning is not None: - usage["reasoning_tokens"] = reasoning - - # Capture raw usage metadata for backend processing - # Serialize to dict here in the converter (not in utils) - serialized = serialize_raw_usage(chunk.usage) - if serialized: - usage["raw_usage"] = serialized + usage.update(_extract_openai_stream_token_usage(chunk.usage, "chat")) + _add_openai_raw_usage(usage, chunk.usage) elif provider_type == "responses": - # For Responses API, usage is only in chunk.response.usage on the terminal - # event. A run cut short (e.g. by max_output_tokens) ends on - # response.incomplete instead of response.completed, and is still billed. - if getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: - if ( - hasattr(chunk, "response") - and hasattr(chunk.response, "usage") - and chunk.response.usage - ): - response_usage = chunk.response.usage - usage["input_tokens"] = getattr(response_usage, "input_tokens", 0) - usage["output_tokens"] = getattr(response_usage, "output_tokens", 0) - - # Handle cached tokens - if hasattr(response_usage, "input_tokens_details") and hasattr( - response_usage.input_tokens_details, "cached_tokens" - ): - cached = response_usage.input_tokens_details.cached_tokens - if cached is not None: - usage["cache_read_input_tokens"] = cached - - # Handle reasoning tokens - if hasattr(response_usage, "output_tokens_details") and hasattr( - response_usage.output_tokens_details, "reasoning_tokens" - ): - reasoning = response_usage.output_tokens_details.reasoning_tokens - if reasoning is not None: - usage["reasoning_tokens"] = reasoning - - # Extract web search count from the complete response - if hasattr(chunk, "response"): - web_search_count = extract_openai_web_search_count(chunk.response) - if web_search_count > 0: - usage["web_search_count"] = web_search_count - - # Capture raw usage metadata for backend processing - # Serialize to dict here in the converter (not in utils) - serialized = serialize_raw_usage(response_usage) - if serialized: - usage["raw_usage"] = serialized + usage.update(_extract_openai_response_chunk_usage(chunk)) return usage @@ -669,28 +360,7 @@ def extract_openai_content_from_chunk( """ if provider_type == "chat": - # Chat Completions API format - if ( - hasattr(chunk, "choices") - and chunk.choices - and len(chunk.choices) > 0 - and chunk.choices[0].delta - ): - delta = chunk.choices[0].delta - - if delta.content: - return delta.content - - audio_delta = getattr(delta, "audio", None) - if audio_delta is not None: - plain_audio = to_plain(audio_delta) - if isinstance(plain_audio, dict): - return {"type": "audio", **plain_audio} - return {"type": "audio"} - - refusal_delta = getattr(delta, "refusal", None) - if refusal_delta: - return {"type": "refusal", "refusal": refusal_delta} + return _extract_openai_chat_content(chunk) elif provider_type == "responses": # Responses API format @@ -714,37 +384,14 @@ def extract_openai_tool_calls_from_chunk(chunk: Any) -> Optional[List[Dict[str, List of tool call deltas if present, None otherwise """ if ( - hasattr(chunk, "choices") - and chunk.choices - and len(chunk.choices) > 0 - and chunk.choices[0].delta + _has_openai_chat_delta(chunk) and hasattr(chunk.choices[0].delta, "tool_calls") and chunk.choices[0].delta.tool_calls ): tool_calls = [] for tool_call in chunk.choices[0].delta.tool_calls: - tc_dict = { - "index": getattr(tool_call, "index", None), - } - - if hasattr(tool_call, "id") and tool_call.id: - tc_dict["id"] = tool_call.id - - if hasattr(tool_call, "type") and tool_call.type: - tc_dict["type"] = tool_call.type - - if hasattr(tool_call, "function") and tool_call.function: - function_dict = {} - if hasattr(tool_call.function, "name") and tool_call.function.name: - function_dict["name"] = tool_call.function.name - if ( - hasattr(tool_call.function, "arguments") - and tool_call.function.arguments - ): - function_dict["arguments"] = tool_call.function.arguments - tc_dict["function"] = function_dict + tool_calls.append(_format_openai_tool_delta(tool_call)) - tool_calls.append(tc_dict) return tool_calls return None @@ -768,27 +415,7 @@ def accumulate_openai_tool_calls( chunk_tool_calls: List of tool call deltas from current chunk """ for tool_call_delta in chunk_tool_calls: - index = tool_call_delta.get("index") - if index is None: - # Some OpenAI-compatible providers omit the index. A delta with a new id - # or a function name starts a call, and one with neither continues the - # latest call. Those calls get keys of their own, so they can never - # collide with an index the provider sends later. - call_id = tool_call_delta.get("id") - index = next( - ( - key - for key, tc in accumulated_tool_calls.items() - if call_id and tc["id"] == call_id - ), - None, - ) - if index is None: - has_name = bool((tool_call_delta.get("function") or {}).get("name")) - if accumulated_tool_calls and not call_id and not has_name: - index = next(reversed(accumulated_tool_calls)) - else: - index = f"no-index-{len(accumulated_tool_calls)}" + index = _openai_tool_delta_index(accumulated_tool_calls, tool_call_delta) # Initialize tool call if first time seeing this index if index not in accumulated_tool_calls: @@ -801,22 +428,7 @@ def accumulate_openai_tool_calls( }, } - # Update with new data from delta - tc = accumulated_tool_calls[index] - - if "id" in tool_call_delta and tool_call_delta["id"]: - tc["id"] = tool_call_delta["id"] - - if "type" in tool_call_delta and tool_call_delta["type"]: - tc["type"] = tool_call_delta["type"] - - if "function" in tool_call_delta: - func_delta = tool_call_delta["function"] - if "name" in func_delta and func_delta["name"]: - tc["function"]["name"] = func_delta["name"] - if "arguments" in func_delta and func_delta["arguments"]: - # Arguments are sent incrementally, concatenate them - tc["function"]["arguments"] += func_delta["arguments"] + _apply_openai_tool_delta(accumulated_tool_calls[index], tool_call_delta) def format_openai_streaming_output( @@ -837,93 +449,12 @@ def format_openai_streaming_output( """ if provider_type == "chat": - content_items: List[FormattedContentItem] = [] - - # Add text content if present - if isinstance(accumulated_content, str) and accumulated_content: - content_items.append({"type": "text", "text": accumulated_content}) - elif isinstance(accumulated_content, list): - text_parts: List[str] = [] - audio_id: Optional[str] = None - audio_data_parts: List[str] = [] - audio_transcript_parts: List[str] = [] - refusal_parts: List[str] = [] - - for item in accumulated_content: - if isinstance(item, str): - if item: - text_parts.append(item) - elif isinstance(item, dict) and item.get("type") == "audio": - if audio_id is None and item.get("id"): - audio_id = item["id"] - if item.get("data"): - audio_data_parts.append(item["data"]) - if item.get("transcript"): - audio_transcript_parts.append(item["transcript"]) - elif isinstance(item, dict) and item.get("type") == "refusal": - if item.get("refusal"): - refusal_parts.append(item["refusal"]) - - if text_parts: - content_items.append({"type": "text", "text": "".join(text_parts)}) - - if audio_data_parts or audio_transcript_parts: - audio_block: Dict[str, Any] = {"type": "audio"} - if audio_id is not None: - audio_block["id"] = audio_id - if audio_data_parts: - audio_block["data"] = "".join(audio_data_parts) - if audio_transcript_parts: - audio_block["transcript"] = "".join(audio_transcript_parts) - content_items.append(audio_block) - - if refusal_parts: - content_items.append( - {"type": "refusal", "refusal": "".join(refusal_parts)} - ) - - # Add tool calls if present - if tool_calls: - for tool_call in tool_calls: - if "function" in tool_call: - function_call: FormattedFunctionCall = { - "type": "function", - "id": tool_call.get("id", ""), - "function": tool_call["function"], - } - content_items.append(function_call) - - # Return formatted message with content - if content_items: - return [{"role": "assistant", "content": content_items}] - else: - # Empty response - return [{"role": "assistant", "content": []}] + return _format_openai_chat_stream_output(accumulated_content, tool_calls) elif provider_type == "responses": - if isinstance(accumulated_content, list) and not accumulated_content: - return [] - if isinstance(accumulated_content, list) and accumulated_content: - items: List[Any] = [] - for entry in accumulated_content: - if isinstance(entry, list): - items.extend(entry) - else: - items.append(entry) - - content_items = _format_responses_output_items(items) - role = _responses_output_role(items) - - if content_items: - return [{"role": role, "content": content_items}] - return [] - elif isinstance(accumulated_content, str): - return [ - { - "role": "assistant", - "content": [{"type": "text", "text": accumulated_content}], - } - ] + output = _format_openai_responses_stream_output(accumulated_content) + if output is not None: + return output # Fallback for any other format return [ @@ -950,3 +481,558 @@ def format_openai_streaming_input( from posthog.ai.utils import merge_system_prompt return merge_system_prompt(kwargs, "openai") + + +def _format_response_message_content(item: Any) -> List[FormattedContentItem]: + content: List[FormattedContentItem] = [] + message_content = _item_attr(item, "content") + if isinstance(message_content, list): + for content_item in message_content: + block = _format_response_content_item(content_item) + if block is not None: + content.append(block) + + elif message_content is not None: + content.append({"type": "text", "text": str(message_content)}) + + return content + + +def _format_response_content_item(content_item: Any) -> Optional[FormattedContentItem]: + content_item_type = _item_attr(content_item, "type") + content_item_text = _item_attr(content_item, "text") + content_item_refusal = _item_attr(content_item, "refusal") + + if content_item_type == "output_text" and (content_item_text is not None): + return {"type": "text", "text": content_item_text} + + elif content_item_type == "refusal" and (content_item_refusal is not None): + return {"type": "refusal", "refusal": content_item_refusal} + + elif content_item_text is not None: + return {"type": "text", "text": content_item_text} + + elif content_item_type == "input_image" and ( + _item_attr(content_item, "image_url") is not None + ): + image_content: FormattedImageContent = { + "type": "image", + "image": _item_attr(content_item, "image_url"), + } + return image_content + + return None + + +def _format_chat_message(message: Any) -> List[FormattedContentItem]: + content: List[FormattedContentItem] = [] + if message.content: + content.append( + { + "type": "text", + "text": message.content, + } + ) + + # A refused request carries its text in `refusal` and leaves `content` empty + refusal = getattr(message, "refusal", None) + if refusal: + content.append({"type": "refusal", "refusal": refusal}) + + if hasattr(message, "tool_calls") and message.tool_calls: + for tool_call in message.tool_calls: + block = _format_chat_tool_call(tool_call) + if block is not None: + content.append(block) + + # Handle audio output (gpt-4o-audio-preview) + if hasattr(message, "audio") and message.audio: + # Convert Pydantic model to dict to capture all fields from OpenAI + audio_dict = message.audio.model_dump() + content.append({"type": "audio", **audio_dict}) + + return content + + +def _format_chat_tool_call(tool_call: Any) -> Optional[FormattedContentItem]: + # Custom tools carry their free-form text in `custom.input`, + # with no `function` object. + is_custom = getattr(tool_call, "type", None) == "custom" + call = getattr(tool_call, "custom" if is_custom else "function", None) + if call is None: + return None + return { + "type": "function", + "id": tool_call.id, + "function": { + "name": call.name, + "arguments": call.input if is_custom else call.arguments, + }, + } + + +def _format_chat_input_message(msg: Any) -> FormattedMessage: + plain = to_plain(msg) + if not isinstance(plain, dict): + plain = {"role": "user", "content": str(plain)} + + formatted: Dict[str, Any] = { + "role": plain.get("role", "user"), + "content": plain.get("content"), + } + + for key in ("tool_calls", "tool_call_id", "name", "audio", "refusal"): + if plain.get(key) is not None: + formatted[key] = to_plain(plain[key]) if key == "audio" else plain[key] + + return cast(FormattedMessage, formatted) + + +def _format_response_input_item(item: Any) -> FormattedMessage: + if not isinstance(item, (dict, str)): + item = to_plain(item) + + if ( + isinstance(item, dict) + and "type" in item + and "role" not in item + and "content" not in item + ): + return cast(FormattedMessage, to_plain(item)) + + role = "user" + content = "" + + if isinstance(item, dict): + role = item.get("role", "user") + content = item.get("content", "") + + elif isinstance(item, str): + content = item + + else: + content = str(item) + + return {"role": role, "content": content} + + +def _count_response_web_search_calls(response: Any) -> int: + # Priority 1: exact count in Responses API output. + count = 0 + if hasattr(response, "output"): + for item in response.output: + if hasattr(item, "type") and item.type == "web_search_call": + count += 1 + return max(0, count) + + +def _has_root_web_search(response: Any) -> bool: + # Priority 2: binary detection, including Perplexity root-level indicators. + for field in ("citations", "search_results"): + if hasattr(response, field): + values = getattr(response, field) + if values and len(values) > 0: + return True + if hasattr(response, "usage") and hasattr(response.usage, "search_context_size"): + if response.usage.search_context_size: + return True + return False + + +def _has_url_citation_annotations(item: Any) -> bool: + if hasattr(item, "annotations"): + annotations = item.annotations + if annotations: + for annotation in annotations: + if _item_attr(annotation, "type") == "url_citation": + return True + return False + + +def _has_choice_web_citations(response: Any) -> bool: + if hasattr(response, "choices"): + for choice in response.choices: + # Check message before delta, including streaming chunks. + for field in ("message", "delta"): + if hasattr(choice, field): + if _has_url_citation_annotations(getattr(choice, field)): + return True + return False + + +def _has_response_web_citations(response: Any) -> bool: + if hasattr(response, "output"): + for item in response.output: + if hasattr(item, "content") and isinstance(item.content, list): + for content_item in item.content: + if _has_url_citation_annotations(content_item): + return True + return False + + +def _extract_openai_response_token_usage(usage: Any) -> TokenUsage: + # Read Responses fields first, then Chat Completions fields. The latter + # override only fields they report, including explicitly null values. + counts = _read_openai_response_usage_fields(usage, "responses") + counts.update(_read_openai_response_usage_fields(usage, "chat")) + result: Dict[str, Any] = {} + for field in ("input_tokens", "output_tokens"): + value = counts.get(field) + if value is not None: + result[field] = value + _add_openai_response_details(result, counts) + return cast(TokenUsage, result) + + +def _read_openai_response_usage_fields(usage: Any, api_type: str) -> Dict[str, Any]: + if api_type == "responses": + token_fields = ( + ("input_tokens", "input_tokens"), + ("output_tokens", "output_tokens"), + ) + detail_fields = ( + ("cache_read_input_tokens", "input_tokens_details", "cached_tokens"), + ("reasoning_tokens", "output_tokens_details", "reasoning_tokens"), + ) + else: + token_fields = ( + ("input_tokens", "prompt_tokens"), + ("output_tokens", "completion_tokens"), + ) + detail_fields = ( + ("cache_read_input_tokens", "prompt_tokens_details", "cached_tokens"), + ("reasoning_tokens", "completion_tokens_details", "reasoning_tokens"), + ) + counts = {} + for output, field in token_fields: + if hasattr(usage, field): + counts[output] = getattr(usage, field) + for output, field, detail in detail_fields: + if hasattr(usage, field) and hasattr(getattr(usage, field), detail): + counts[output] = getattr(getattr(usage, field), detail) + return counts + + +def _add_openai_response_details( + result: Dict[str, Any], counts: Dict[str, Any] +) -> None: + for field in ("cache_read_input_tokens", "reasoning_tokens"): + value = counts.get(field) + if value is not None and value > 0: + result[field] = value + + +def _extract_openai_stream_token_usage(raw: Any, provider_type: str) -> TokenUsage: + if provider_type == "chat": + input_field, output_field = "prompt_tokens", "completion_tokens" + input_details, output_details = ( + "prompt_tokens_details", + "completion_tokens_details", + ) + else: + input_field, output_field = "input_tokens", "output_tokens" + input_details, output_details = "input_tokens_details", "output_tokens_details" + usage: Dict[str, Any] = { + "input_tokens": getattr(raw, input_field, 0), + "output_tokens": getattr(raw, output_field, 0), + } + # Unlike non-streaming usage, explicitly reported zero details are retained. + for output, field, detail in ( + ("cache_read_input_tokens", input_details, "cached_tokens"), + ("reasoning_tokens", output_details, "reasoning_tokens"), + ): + if hasattr(raw, field) and hasattr(getattr(raw, field), detail): + value = getattr(getattr(raw, field), detail) + if value is not None: + usage[output] = value + return cast(TokenUsage, usage) + + +def _add_openai_raw_usage(usage: TokenUsage, raw: Any) -> None: + # Capture raw usage metadata for backend processing. Serialize in the converter. + serialized = serialize_raw_usage(raw) + if serialized: + usage["raw_usage"] = serialized + + +def _has_openai_chat_delta(chunk: Any) -> bool: + return bool( + hasattr(chunk, "choices") + and chunk.choices + and len(chunk.choices) > 0 + and chunk.choices[0].delta + ) + + +def _format_openai_tool_delta(tool_call: Any) -> Dict[str, Any]: + tc_dict = { + "index": getattr(tool_call, "index", None), + } + + if hasattr(tool_call, "id") and tool_call.id: + tc_dict["id"] = tool_call.id + + if hasattr(tool_call, "type") and tool_call.type: + tc_dict["type"] = tool_call.type + + if hasattr(tool_call, "function") and tool_call.function: + tc_dict["function"] = _format_openai_function_delta(tool_call.function) + + return tc_dict + + +def _format_openai_function_delta(function: Any) -> Dict[str, Any]: + function_dict = {} + if hasattr(function, "name") and function.name: + function_dict["name"] = function.name + if hasattr(function, "arguments") and function.arguments: + function_dict["arguments"] = function.arguments + return function_dict + + +def _openai_tool_delta_index( + accumulated_tool_calls: Dict[Union[int, str], Dict[str, Any]], + tool_call_delta: Dict[str, Any], +) -> Union[int, str]: + index = tool_call_delta.get("index") + if index is None: + # Some OpenAI-compatible providers omit the index. A delta with a new id + # or a function name starts a call, and one with neither continues the + # latest call. Those calls get keys of their own, so they can never + # collide with an index the provider sends later. + call_id = tool_call_delta.get("id") + index = _openai_tool_index_by_id(accumulated_tool_calls, call_id) + + if index is None: + has_name = bool((tool_call_delta.get("function") or {}).get("name")) + if accumulated_tool_calls and not call_id and not has_name: + index = next(reversed(accumulated_tool_calls)) + else: + index = f"no-index-{len(accumulated_tool_calls)}" + + return index + + +def _apply_openai_tool_delta( + tc: Dict[str, Any], tool_call_delta: Dict[str, Any] +) -> None: + # Update with new data from delta + + for field in ("id", "type"): + if field in tool_call_delta and tool_call_delta[field]: + tc[field] = tool_call_delta[field] + + if "function" in tool_call_delta: + func_delta = tool_call_delta["function"] + if "name" in func_delta and func_delta["name"]: + tc["function"]["name"] = func_delta["name"] + if "arguments" in func_delta and func_delta["arguments"]: + # Arguments are sent incrementally, concatenate them + tc["function"]["arguments"] += func_delta["arguments"] + + +def _format_openai_chat_stream_output( + accumulated_content: Any, tool_calls: Optional[List[Dict[str, Any]]] +) -> List[FormattedMessage]: + content_items: List[FormattedContentItem] = [] + + # Add text content if present + if isinstance(accumulated_content, str) and accumulated_content: + content_items.append({"type": "text", "text": accumulated_content}) + elif isinstance(accumulated_content, list): + content_items.extend(_merge_openai_chat_stream_parts(accumulated_content)) + + # Add tool calls if present + if tool_calls: + for tool_call in tool_calls: + if "function" in tool_call: + function_call: FormattedFunctionCall = { + "type": "function", + "id": tool_call.get("id", ""), + "function": tool_call["function"], + } + content_items.append(function_call) + + # Return formatted message with content + if content_items: + return [{"role": "assistant", "content": content_items}] + else: + # Empty response + return [{"role": "assistant", "content": []}] + + +def _merge_openai_chat_stream_parts( + accumulated_content: List[Any], +) -> List[FormattedContentItem]: + text_parts: List[str] = [] + audio_id: Optional[str] = None + audio_data_parts: List[str] = [] + audio_transcript_parts: List[str] = [] + refusal_parts: List[str] = [] + + for item in accumulated_content: + if isinstance(item, str): + if item: + text_parts.append(item) + elif isinstance(item, dict) and item.get("type") == "audio": + audio_id = _append_openai_audio_delta( + item, audio_id, audio_data_parts, audio_transcript_parts + ) + elif isinstance(item, dict) and item.get("type") == "refusal": + if item.get("refusal"): + refusal_parts.append(item["refusal"]) + + return _format_openai_chat_stream_parts( + text_parts, audio_id, audio_data_parts, audio_transcript_parts, refusal_parts + ) + + +def _append_openai_audio_delta( + item: Dict[str, Any], + audio_id: Optional[str], + data_parts: List[str], + transcript_parts: List[str], +) -> Optional[str]: + if audio_id is None and item.get("id"): + audio_id = item["id"] + if item.get("data"): + data_parts.append(item["data"]) + if item.get("transcript"): + transcript_parts.append(item["transcript"]) + return audio_id + + +def _format_openai_chat_stream_parts( + text_parts: List[str], + audio_id: Optional[str], + audio_data_parts: List[str], + audio_transcript_parts: List[str], + refusal_parts: List[str], +) -> List[FormattedContentItem]: + content_items: List[FormattedContentItem] = [] + if text_parts: + content_items.append({"type": "text", "text": "".join(text_parts)}) + + if audio_data_parts or audio_transcript_parts: + audio_block: Dict[str, Any] = {"type": "audio"} + if audio_id is not None: + audio_block["id"] = audio_id + if audio_data_parts: + audio_block["data"] = "".join(audio_data_parts) + if audio_transcript_parts: + audio_block["transcript"] = "".join(audio_transcript_parts) + content_items.append(audio_block) + + if refusal_parts: + content_items.append({"type": "refusal", "refusal": "".join(refusal_parts)}) + + return content_items + + +def _format_openai_responses_stream_output( + accumulated_content: Any, +) -> Optional[List[FormattedMessage]]: + if isinstance(accumulated_content, list) and not accumulated_content: + return [] + if isinstance(accumulated_content, list) and accumulated_content: + items: List[Any] = [] + for entry in accumulated_content: + if isinstance(entry, list): + items.extend(entry) + else: + items.append(entry) + + content_items = _format_responses_output_items(items) + role = _responses_output_role(items) + + if content_items: + return [{"role": role, "content": content_items}] + return [] + elif isinstance(accumulated_content, str): + return [ + { + "role": "assistant", + "content": [{"type": "text", "text": accumulated_content}], + } + ] + + return None + + +def _format_chat_completion_choices(choices: Any) -> Optional[FormattedMessage]: + content: List[FormattedContentItem] = [] + role = "assistant" + + for choice in choices: + if hasattr(choice, "message") and choice.message: + if choice.message.role: + role = choice.message.role + content.extend(_format_chat_message(choice.message)) + + if content: + return { + "role": role, + "content": content, + } + + return None + + +def _extract_openai_response_chunk_usage(chunk: Any) -> TokenUsage: + usage = TokenUsage() + # For Responses API, usage is only in chunk.response.usage on the terminal + # event. A run cut short (e.g. by max_output_tokens) ends on + # response.incomplete instead of response.completed, and is still billed. + if getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: + if ( + hasattr(chunk, "response") + and hasattr(chunk.response, "usage") + and chunk.response.usage + ): + response_usage = chunk.response.usage + usage.update( + _extract_openai_stream_token_usage(response_usage, "responses") + ) + # Extract web search count from the complete response. + if hasattr(chunk, "response"): + web_search_count = extract_openai_web_search_count(chunk.response) + if web_search_count > 0: + usage["web_search_count"] = web_search_count + _add_openai_raw_usage(usage, response_usage) + + return usage + + +def _extract_openai_chat_content(chunk: Any) -> Optional[Any]: + # Chat Completions API format + if _has_openai_chat_delta(chunk): + delta = chunk.choices[0].delta + + if delta.content: + return delta.content + + audio_delta = getattr(delta, "audio", None) + if audio_delta is not None: + plain_audio = to_plain(audio_delta) + if isinstance(plain_audio, dict): + return {"type": "audio", **plain_audio} + return {"type": "audio"} + + refusal_delta = getattr(delta, "refusal", None) + if refusal_delta: + return {"type": "refusal", "refusal": refusal_delta} + + return None + + +def _openai_tool_index_by_id( + accumulated_tool_calls: Dict[Union[int, str], Dict[str, Any]], call_id: Any +) -> Optional[Union[int, str]]: + return next( + ( + key + for key, tc in accumulated_tool_calls.items() + if call_id and tc["id"] == call_id + ), + None, + ) diff --git a/posthog/ai/openai_agents/processor.py b/posthog/ai/openai_agents/processor.py index 951dd95e2..5dd8ed40a 100644 --- a/posthog/ai/openai_agents/processor.py +++ b/posthog/ai/openai_agents/processor.py @@ -248,21 +248,7 @@ def on_trace_end(self, trace: Trace) -> None: "$ai_framework": "openai-agents", } - if latency is not None: - properties["$ai_latency"] = latency - - # The Agents SDK group_id links traces from one conversation, which is - # exactly what PostHog calls a session. $ai_group_id is still emitted - # for anyone already querying it. - if group_id: - properties["$ai_session_id"] = group_id - properties["$ai_group_id"] = group_id - - # Include trace metadata if present - if metadata: - properties["$ai_trace_metadata"] = self._with_privacy_mode( - _ensure_serializable(metadata) - ) + self._add_trace_details(properties, latency, group_id, metadata) if distinct_id is None: properties["$process_person_profile"] = False @@ -275,6 +261,29 @@ def on_trace_end(self, trace: Trace) -> None: except Exception as e: log.debug(f"Error in on_trace_end: {e}") + def _add_trace_details( + self, + properties: Dict[str, Any], + latency: Optional[float], + group_id: Any, + metadata: Any, + ) -> None: + if latency is not None: + properties["$ai_latency"] = latency + + # The Agents SDK group_id links traces from one conversation, which is + # exactly what PostHog calls a session. $ai_group_id is still emitted + # for anyone already querying it. + if group_id: + properties["$ai_session_id"] = group_id + properties["$ai_group_id"] = group_id + + # Include trace metadata if present + if metadata: + properties["$ai_trace_metadata"] = self._with_privacy_mode( + _ensure_serializable(metadata) + ) + def on_span_start(self, span: Span[Any]) -> None: """Called when a new span begins.""" try: @@ -292,15 +301,7 @@ def on_span_end(self, span: Span[Any]) -> None: parent_id = span.parent_id span_data = span.span_data - # Calculate latency - start_time = self._span_start_times.pop(span_id, None) - if start_time: - latency = time.time() - start_time - else: - # Fall back to parsing timestamps - started = _parse_iso_timestamp(span.started_at) - ended = _parse_iso_timestamp(span.ended_at) - latency = (ended - started) if (started and ended) else 0 + latency = self._span_latency(span, span_id) # Get user-provided distinct ID from trace metadata (resolved at trace start). # None means no user-provided ID — use trace_id as fallback in personless mode, @@ -311,175 +312,84 @@ def on_span_end(self, span: Span[Any]) -> None: # Get group_id from trace metadata for linking group_id = self._get_group_id(trace_id) - # Get error info if present - error_info = span.error - error_properties = {} - if error_info: - if isinstance(error_info, dict): - error_message = error_info.get("message", str(error_info)) - error_type_raw = error_info.get("type", "") - else: - error_message = str(error_info) - error_type_raw = "" - - # Categorize error type for cross-provider filtering/alerting - error_type = "unknown" - if ( - "ModelBehaviorError" in error_type_raw - or "ModelBehaviorError" in error_message - ): - error_type = "model_behavior_error" - elif "UserError" in error_type_raw or "UserError" in error_message: - error_type = "user_error" - elif ( - "InputGuardrailTripwireTriggered" in error_type_raw - or "InputGuardrailTripwireTriggered" in error_message - ): - error_type = "input_guardrail_triggered" - elif ( - "OutputGuardrailTripwireTriggered" in error_type_raw - or "OutputGuardrailTripwireTriggered" in error_message - ): - error_type = "output_guardrail_triggered" - elif ( - "MaxTurnsExceeded" in error_type_raw - or "MaxTurnsExceeded" in error_message - ): - error_type = "max_turns_exceeded" - - error_properties = { - "$ai_is_error": True, - "$ai_error": ( - self._summarize_span_error(error_info) - if self._privacy_enabled() - else error_message - ), - "$ai_error_type": error_type, - } + error_properties = self._span_error_properties(span.error) # Personless mode: no user-provided distinct_id, fallback to trace_id if distinct_id is None: error_properties["$process_person_profile"] = False distinct_id = trace_id - # Dispatch based on span data type - if isinstance(span_data, GenerationSpanData): - self._handle_generation_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) - elif isinstance(span_data, FunctionSpanData): - self._handle_function_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) - elif isinstance(span_data, AgentSpanData): - self._handle_agent_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) - elif isinstance(span_data, HandoffSpanData): - self._handle_handoff_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) - elif isinstance(span_data, GuardrailSpanData): - self._handle_guardrail_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) - elif isinstance(span_data, ResponseSpanData): - self._handle_response_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) - elif isinstance(span_data, CustomSpanData): - self._handle_custom_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) - elif isinstance( - span_data, (TranscriptionSpanData, SpeechSpanData, SpeechGroupSpanData) - ): - self._handle_audio_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) - elif isinstance(span_data, MCPListToolsSpanData): - self._handle_mcp_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) - else: - # Unknown span type - capture as generic span - self._handle_generic_span( - span_data, - trace_id, - span_id, - parent_id, - latency, - distinct_id, - group_id, - error_properties, - ) + handler = self._span_handler(span_data) + handler( + span_data, + trace_id, + span_id, + parent_id, + latency, + distinct_id, + group_id, + error_properties, + ) except Exception as e: log.debug(f"Error in on_span_end: {e}") + def _span_latency(self, span: Span[Any], span_id: str) -> float: + # Calculate latency + start_time = self._span_start_times.pop(span_id, None) + if start_time: + return time.time() - start_time + else: + # Fall back to parsing timestamps + started = _parse_iso_timestamp(span.started_at) + ended = _parse_iso_timestamp(span.ended_at) + return (ended - started) if (started and ended) else 0 + + def _span_error_properties(self, error_info: Any) -> Dict[str, Any]: + # Get error info if present + error_properties = {} + if error_info: + if isinstance(error_info, dict): + error_message = error_info.get("message", str(error_info)) + error_type_raw = error_info.get("type", "") + else: + error_message = str(error_info) + error_type_raw = "" + + error_type = _categorize_span_error(error_type_raw, error_message) + + error_properties = { + "$ai_is_error": True, + "$ai_error": ( + self._summarize_span_error(error_info) + if self._privacy_enabled() + else error_message + ), + "$ai_error_type": error_type, + } + + return error_properties + + def _span_handler(self, span_data: Any) -> Callable[..., None]: + for span_types, handler_name in ( + (GenerationSpanData, "_handle_generation_span"), + (FunctionSpanData, "_handle_function_span"), + (AgentSpanData, "_handle_agent_span"), + (HandoffSpanData, "_handle_handoff_span"), + (GuardrailSpanData, "_handle_guardrail_span"), + (ResponseSpanData, "_handle_response_span"), + (CustomSpanData, "_handle_custom_span"), + ( + (TranscriptionSpanData, SpeechSpanData, SpeechGroupSpanData), + "_handle_audio_span", + ), + (MCPListToolsSpanData, "_handle_mcp_span"), + ): + if isinstance(span_data, span_types): + return getattr(self, handler_name) + # Unknown span type - capture as generic span. + return self._handle_generic_span + def _base_properties( self, trace_id: str, @@ -554,20 +464,19 @@ def _handle_generation_span( _ensure_serializable(span_data.output), self._client ) ), - **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), - **( - {"$ai_output_tokens": output_tokens} - if output_tokens is not None - else {} - ), - # Sum of the reported sides; omitted when neither side was reported. - **( - {"$ai_total_tokens": (input_tokens or 0) + (output_tokens or 0)} - if input_tokens is not None or output_tokens is not None - else {} - ), + **_token_properties(input_tokens, output_tokens), } + self._add_generation_usage_details(properties, usage, span_data) + + self._capture_event("$ai_generation", properties, distinct_id) + + def _add_generation_usage_details( + self, + properties: Dict[str, Any], + usage: Dict[str, Any], + span_data: GenerationSpanData, + ) -> None: # Add optional token fields if present. The Agents SDK nests cached # tokens under the usage details fields; flat keys are kept for custom # models. Nested reasoning is not forwarded: compatible providers count @@ -601,8 +510,6 @@ def _handle_generation_span( if finish_reason is not None: properties["$ai_stop_reason"] = finish_reason - self._capture_event("$ai_generation", properties, distinct_id) - def _handle_function_span( self, span_data: FunctionSpanData, @@ -727,14 +634,9 @@ def _handle_response_span( response = span_data.response response_id = response.id if response else None - # Try to extract usage from response - usage = getattr(response, "usage", None) if response else None - total_cost_usd = getattr(usage, "cost", None) if usage else None - # None when the response never reported a count: absent means unknown, - # 0 is a report of nothing. - input_tokens = getattr(usage, "input_tokens", None) if usage else None - output_tokens = getattr(usage, "output_tokens", None) if usage else None - + usage, total_cost_usd, input_tokens, output_tokens = _response_usage_counts( + response + ) # Try to extract model from response model = getattr(response, "model", None) if response else None @@ -747,18 +649,7 @@ def _handle_response_span( "$ai_input": self._with_privacy_mode( finalize_ai_content(_ensure_serializable(span_data.input), self._client) ), - **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), - **( - {"$ai_output_tokens": output_tokens} - if output_tokens is not None - else {} - ), - # Sum of the reported sides; omitted when neither side was reported. - **( - {"$ai_total_tokens": (input_tokens or 0) + (output_tokens or 0)} - if input_tokens is not None or output_tokens is not None - else {} - ), + **_token_properties(input_tokens, output_tokens), } if total_cost_usd is not None: @@ -837,6 +728,20 @@ def _handle_audio_span( "$ai_span_type": span_type, } + self._add_audio_metadata(properties, span_data) + + self._add_audio_input(properties, span_data, span_type) + + if hasattr(span_data, "output") and isinstance(span_data.output, str): + properties["$ai_output_state"] = self._with_privacy_mode( + finalize_ai_content( + _ensure_serializable(span_data.output), self._client + ) + ) + + self._capture_event("$ai_span", properties, distinct_id) + + def _add_audio_metadata(self, properties: Dict[str, Any], span_data: Any) -> None: # Add model info if available if hasattr(span_data, "model") and span_data.model: properties["$ai_model"] = span_data.model @@ -857,6 +762,9 @@ def _handle_audio_span( if hasattr(span_data, "output_format"): properties["audio_output_format"] = span_data.output_format + def _add_audio_input( + self, properties: Dict[str, Any], span_data: Any, span_type: str + ) -> None: # Capture the input. For speech (TTS) spans the input is the text to # synthesize. For transcription spans the input is the base64-encoded # audio — a bare string with no key/parent context the structural @@ -883,15 +791,6 @@ def _handle_audio_span( ) ) - if hasattr(span_data, "output") and isinstance(span_data.output, str): - properties["$ai_output_state"] = self._with_privacy_mode( - finalize_ai_content( - _ensure_serializable(span_data.output), self._client - ) - ) - - self._capture_event("$ai_span", properties, distinct_id) - def _redact_audio_model_config(self, model_config: Any) -> Any: # The STT prompt and TTS instructions are user-authored text. if not isinstance(model_config, dict) or not self._privacy_enabled(): @@ -976,3 +875,43 @@ def force_flush(self) -> None: self._client.flush() except Exception as e: log.debug(f"Error in force_flush: {e}") + + +def _categorize_span_error(error_type_raw: str, error_message: str) -> str: + # Ordered categories preserve precedence when a message names several errors. + for name, category in ( + ("ModelBehaviorError", "model_behavior_error"), + ("UserError", "user_error"), + ("InputGuardrailTripwireTriggered", "input_guardrail_triggered"), + ("OutputGuardrailTripwireTriggered", "output_guardrail_triggered"), + ("MaxTurnsExceeded", "max_turns_exceeded"), + ): + if name in error_type_raw or name in error_message: + return category + return "unknown" + + +def _token_properties(input_tokens: Any, output_tokens: Any) -> Dict[str, Any]: + properties = {} + if input_tokens is not None: + properties["$ai_input_tokens"] = input_tokens + if output_tokens is not None: + properties["$ai_output_tokens"] = output_tokens + # Sum of the reported sides; omitted when neither side was reported. + if input_tokens is not None or output_tokens is not None: + properties["$ai_total_tokens"] = (input_tokens or 0) + (output_tokens or 0) + return properties + + +def _response_usage_counts( + response: Any, +) -> tuple[Any, Optional[float], Optional[int], Optional[int]]: + # Try to extract usage from response + usage = getattr(response, "usage", None) if response else None + total_cost_usd = getattr(usage, "cost", None) if usage else None + # None when the response never reported a count: absent means unknown, + # 0 is a report of nothing. + input_tokens = getattr(usage, "input_tokens", None) if usage else None + output_tokens = getattr(usage, "output_tokens", None) if usage else None + + return usage, total_cost_usd, input_tokens, output_tokens diff --git a/posthog/ai/otel/_privacy.py b/posthog/ai/otel/_privacy.py index 5a2114ee5..8b87b4a3a 100644 --- a/posthog/ai/otel/_privacy.py +++ b/posthog/ai/otel/_privacy.py @@ -213,13 +213,7 @@ def _redact_tool_definitions( return None if drop_unparseable else value redacted: List[Any] = [] for item in value: - definition = ( - _redact_tool_definition_json(item, drop_unparseable) - if isinstance(item, str) - else None - if drop_unparseable - else item - ) + definition = _redact_tool_definition_item(item, drop_unparseable) if definition is not None: redacted.append(definition) if len(redacted) == len(value) and all(a is b for a, b in zip(redacted, value)): @@ -227,6 +221,16 @@ def _redact_tool_definitions( return redacted +def _redact_tool_definition_item(item: Any, drop_unparseable: bool) -> Any: + return ( + _redact_tool_definition_json(item, drop_unparseable) + if isinstance(item, str) + else None + if drop_unparseable + else item + ) + + def _is_tool_definition_attribute(key: str) -> bool: return ( key in _TOOL_DEFINITION_ATTRIBUTES diff --git a/posthog/ai/otel/_request_context.py b/posthog/ai/otel/_request_context.py index b34bf4c61..dc0018c5a 100644 --- a/posthog/ai/otel/_request_context.py +++ b/posthog/ai/otel/_request_context.py @@ -102,9 +102,7 @@ def with_request_context(span: ReadableSpan, context: RequestContext) -> Readabl def to_attribute_value(value: Any) -> Optional[AttributeValue]: """Return ``value`` as an OpenTelemetry attribute value, JSON-encoding anything OTel can't hold.""" - if value is None or (isinstance(value, str) and not value): - return None - if isinstance(value, float) and not math.isfinite(value): + if _invalid_attribute_value(value): return None if _is_out_of_int64_range(value): return str(value) @@ -112,6 +110,18 @@ def to_attribute_value(value: Any) -> Optional[AttributeValue]: return value if isinstance(value, (list, tuple)) and _is_homogeneous(value): return list(value) + return _json_attribute_value(value) + + +def _invalid_attribute_value(value: Any) -> bool: + return ( + value is None + or (isinstance(value, str) and not value) + or (isinstance(value, float) and not math.isfinite(value)) + ) + + +def _json_attribute_value(value: Any) -> Optional[AttributeValue]: try: encoded = json.dumps(value, allow_nan=False, default=_encode_string_like) except (TypeError, ValueError): diff --git a/posthog/ai/prompts.py b/posthog/ai/prompts.py index 3a785f0ce..6b40c4a62 100644 --- a/posthog/ai/prompts.py +++ b/posthog/ai/prompts.py @@ -17,6 +17,8 @@ from dataclasses import dataclass from typing import Any, Dict, List, Literal, Optional, Union, overload +from requests import Response as _Response + from posthog.capture_v1 import _parse_retry_after from posthog.request import USER_AGENT, _get_session from posthog.utils import remove_trailing_slash @@ -462,6 +464,42 @@ def get_all(self, *, label: Optional[str] = None) -> Dict[str, PromptResult]: self._maybe_capture_error(error, name="*", version=None, label=label) raise + resolved_rows, skipped = self._resolve_prompt_rows(rows, label) + + now = time.time() + results: Dict[str, PromptResult] = {} + for row in resolved_rows: + config = _extract_config(row) + self._cache[_cache_key(row["name"], None, label)] = CachedPrompt( + prompt=row["prompt"], + fetched_at=now, + name=row["name"], + version=row["version"], + label=label, + config=config, + ) + results[row["name"]] = PromptResult( + source="api", + prompt=row["prompt"], + name=row["name"], + version=row["version"], + label=label, + config=copy.deepcopy(config), + ) + + if skipped: + log.warning( + "[PostHog Prompts] Skipped %d prompt(s) that did not resolve label %r: %s", + len(skipped), + label, + ", ".join(skipped), + ) + + return results + + def _resolve_prompt_rows( + self, rows: List[Dict[str, Any]], label: Optional[str] + ) -> tuple[List[Dict[str, Any]], List[str]]: reference = _prompt_list_reference(label) # Validate every row before caching any, so a rejected batch leaves @@ -514,36 +552,7 @@ def get_all(self, *, label: Optional[str] = None) -> Dict[str, PromptResult]: self._maybe_capture_error(compat_error, name="*", version=None, label=label) raise compat_error - now = time.time() - results: Dict[str, PromptResult] = {} - for row in resolved_rows: - config = _extract_config(row) - self._cache[_cache_key(row["name"], None, label)] = CachedPrompt( - prompt=row["prompt"], - fetched_at=now, - name=row["name"], - version=row["version"], - label=label, - config=config, - ) - results[row["name"]] = PromptResult( - source="api", - prompt=row["prompt"], - name=row["name"], - version=row["version"], - label=label, - config=copy.deepcopy(config), - ) - - if skipped: - log.warning( - "[PostHog Prompts] Skipped %d prompt(s) that did not resolve label %r: %s", - len(skipped), - label, - ", ".join(skipped), - ) - - return results + return resolved_rows, skipped def _get_internal( self, @@ -569,38 +578,9 @@ def _get_internal( cached = self._cache.get(cache_key) now = time.time() - if cached is not None: - is_fresh = (now - cached.fetched_at) < ttl - - if is_fresh: - return PromptResult( - source="cache", - prompt=cached.prompt, - name=cached.name, - version=cached.version, - label=cached.label, - # Copied so a caller mutating result.config can't pollute the - # cache entry that later cache hits are served from. - config=copy.deepcopy(cached.config), - ) - - # A failed refetch left this entry in cooldown. Serving it keeps one - # throttled client from turning every later get() into another - # request, which is what holds it against the limit. A zero TTL is - # an explicit request to refetch on every read, so it opts out. - if ( - ttl > 0 - and cached.retry_not_before is not None - and now < cached.retry_not_before - ): - return PromptResult( - source="stale_cache", - prompt=cached.prompt, - name=cached.name, - version=cached.version, - label=cached.label, - config=copy.deepcopy(cached.config), - ) + cached_result = _cached_prompt_result(cached, now, ttl) + if cached_result is not None: + return cached_result # Try to fetch from API try: @@ -669,11 +649,7 @@ def _fetch_prompt_once( self._in_flight[cache_key] = pending if not is_leader: - pending.done.wait() - if pending.error is not None: - raise pending.error - assert pending.result is not None - return pending.result + return _wait_for_prompt(pending) try: data = self._fetch_prompt_from_api(name, version, label) @@ -851,17 +827,7 @@ def _fetch_prompt_list_from_api(self, label: Optional[str]) -> List[Dict[str, An response = _get_session().get(url, headers=headers, timeout=10) - if not response.ok: - if response.status_code == 401: - raise _authentication_error(reference) - if response.status_code == 403: - raise _access_denied_error(reference) - raise PromptFetchError( - f"[PostHog Prompts] Failed to fetch {reference}: HTTP {response.status_code}", - _parse_retry_after_seconds(response.headers.get("Retry-After")) - if response.status_code == 429 - else None, - ) + _raise_prompt_response_error(response, reference) try: data = response.json() @@ -877,18 +843,7 @@ def _fetch_prompt_list_from_api(self, label: Optional[str]) -> List[Dict[str, An rows.extend(data["results"]) - # The Authorization header goes to every followed link, so a link - # off the configured host must never be requested. - next_url = data.get("next") - if next_url is not None and ( - not isinstance(next_url, str) - or not _is_same_origin(next_url, self._host) - ): - raise Exception( - f"[PostHog Prompts] Refusing to follow a pagination link off the " - f"configured host while fetching {reference}." - ) - url = next_url + url = _prompt_next_url(data, self._host, reference) pages += 1 return rows @@ -954,16 +909,91 @@ def _fetch_prompt_from_api( else None, ) - try: - data = response.json() - except Exception: - raise Exception( - f"[PostHog Prompts] Invalid response format for {prompt_title}" + return _prompt_response_data(response, prompt_title) + + +def _cached_prompt_result( + cached: Optional[CachedPrompt], now: float, ttl: float +) -> Optional[PromptResult]: + if cached is not None: + is_fresh = (now - cached.fetched_at) < ttl + + if is_fresh: + return PromptResult( + source="cache", + prompt=cached.prompt, + name=cached.name, + version=cached.version, + label=cached.label, + # Copied so a caller mutating result.config can't pollute the + # cache entry that later cache hits are served from. + config=copy.deepcopy(cached.config), ) - if not _is_prompt_api_response(data): - raise Exception( - f"[PostHog Prompts] Invalid response format for {prompt_title}" + # A failed refetch left this entry in cooldown. Serving it keeps one + # throttled client from turning every later get() into another + # request, which is what holds it against the limit. A zero TTL is + # an explicit request to refetch on every read, so it opts out. + if ( + ttl > 0 + and cached.retry_not_before is not None + and now < cached.retry_not_before + ): + return PromptResult( + source="stale_cache", + prompt=cached.prompt, + name=cached.name, + version=cached.version, + label=cached.label, + config=copy.deepcopy(cached.config), ) - return data + return None + + +def _wait_for_prompt(pending: _InFlightFetch) -> CachedPrompt: + pending.done.wait() + if pending.error is not None: + raise pending.error + assert pending.result is not None + return pending.result + + +def _raise_prompt_response_error(response: _Response, reference: str) -> None: + if not response.ok: + if response.status_code == 401: + raise _authentication_error(reference) + if response.status_code == 403: + raise _access_denied_error(reference) + raise PromptFetchError( + f"[PostHog Prompts] Failed to fetch {reference}: HTTP {response.status_code}", + _parse_retry_after_seconds(response.headers.get("Retry-After")) + if response.status_code == 429 + else None, + ) + + +def _prompt_response_data(response: _Response, prompt_title: str) -> Dict[str, Any]: + try: + data = response.json() + except Exception: + raise Exception(f"[PostHog Prompts] Invalid response format for {prompt_title}") + + if not _is_prompt_api_response(data): + raise Exception(f"[PostHog Prompts] Invalid response format for {prompt_title}") + + return data + + +def _prompt_next_url(data: Dict[str, Any], host: str, reference: str) -> Optional[str]: + # The Authorization header goes to every followed link, so a link + # off the configured host must never be requested. + next_url = data.get("next") + if next_url is not None and ( + not isinstance(next_url, str) or not _is_same_origin(next_url, host) + ): + raise Exception( + f"[PostHog Prompts] Refusing to follow a pagination link off the " + f"configured host while fetching {reference}." + ) + return next_url diff --git a/posthog/ai/utils.py b/posthog/ai/utils.py index 874ff3fcb..dfc4c8c90 100644 --- a/posthog/ai/utils.py +++ b/posthog/ai/utils.py @@ -35,6 +35,10 @@ ) +def _is_token_count(value: Any) -> bool: + return isinstance(value, int) and not isinstance(value, bool) and value >= 0 + + def _extract_cache_creation_ttl_breakdown( token_details: Any, ) -> Optional[Tuple[int, int]]: @@ -46,11 +50,7 @@ def _extract_cache_creation_ttl_breakdown( token_details.get("ephemeral_5m_input_tokens"), token_details.get("ephemeral_1h_input_tokens"), ) - if not any(value is not None for value in values) or not all( - isinstance(value, int) and not isinstance(value, bool) and value >= 0 - for value in values - if value is not None - ): + if not all(_is_token_count(value) for value in values if value is not None): return None breakdown = (values[0] or 0, values[1] or 0) @@ -131,19 +131,15 @@ def serialize_raw_usage(raw_usage: Any) -> Optional[Dict[str, Any]]: if isinstance(raw_usage, dict): return raw_usage - # Try Pydantic model_dump() (OpenAI/Anthropic) - if hasattr(raw_usage, "model_dump") and callable(raw_usage.model_dump): - try: - return raw_usage.model_dump() - except Exception: - pass - - # Try to_dict() (some protobuf objects) - if hasattr(raw_usage, "to_dict") and callable(raw_usage.to_dict): - try: - return raw_usage.to_dict() - except Exception: - pass + # Try Pydantic model_dump() first, then protobuf-style to_dict(). + for method_name in ("model_dump", "to_dict"): + if hasattr(raw_usage, method_name) and callable( + getattr(raw_usage, method_name) + ): + try: + return getattr(raw_usage, method_name)() + except Exception: + pass # Try __dict__ / vars() for simple objects try: @@ -151,6 +147,10 @@ def serialize_raw_usage(raw_usage: Any) -> Optional[Dict[str, Any]]: except Exception: pass + return _raw_usage_string(raw_usage) + + +def _raw_usage_string(raw_usage: Any) -> Optional[Dict[str, Any]]: # Last resort: convert to string representation # This ensures we always return something rather than failing try: @@ -174,81 +174,55 @@ def merge_usage_stats( source: TokenUsage that may contain None values mode: Either "incremental" or "cumulative" """ + fields = ( + "input_tokens", + "output_tokens", + "cache_read_input_tokens", + "cache_creation_input_tokens", + "reasoning_tokens", + ) + target_fields = cast(Dict[str, Any], target) + source_fields = cast(Dict[str, Any], source) if mode == "incremental": - # Add new values to existing totals - source_input = source.get("input_tokens") - if source_input is not None: - current = target.get("input_tokens") or 0 - target["input_tokens"] = current + source_input - - source_output = source.get("output_tokens") - if source_output is not None: - current = target.get("output_tokens") or 0 - target["output_tokens"] = current + source_output - - source_cache_read = source.get("cache_read_input_tokens") - if source_cache_read is not None: - current = target.get("cache_read_input_tokens") or 0 - target["cache_read_input_tokens"] = current + source_cache_read - - source_cache_creation = source.get("cache_creation_input_tokens") - if source_cache_creation is not None: - current = target.get("cache_creation_input_tokens") or 0 - target["cache_creation_input_tokens"] = current + source_cache_creation - - source_reasoning = source.get("reasoning_tokens") - if source_reasoning is not None: - current = target.get("reasoning_tokens") or 0 - target["reasoning_tokens"] = current + source_reasoning - - source_web_search = source.get("web_search_count") - if source_web_search is not None: - current = target.get("web_search_count") or 0 - target["web_search_count"] = max(current, source_web_search) - - # Carried across, never accumulated: this describes the provider's accounting - # model, so it is the same on every chunk. - source_exclusive = source.get("cache_reporting_exclusive") - if source_exclusive is not None: - target["cache_reporting_exclusive"] = source_exclusive - - # Merge raw_usage to avoid losing data from earlier events - # For Anthropic streaming: message_start has input tokens, message_delta has output - # Note: raw_usage is already serialized by converters, so it's a dict - source_raw_usage = source.get("raw_usage") - if source_raw_usage is not None and isinstance(source_raw_usage, dict): - current_raw_value = target.get("raw_usage") - current_raw: Dict[str, Any] = ( - current_raw_value if isinstance(current_raw_value, dict) else {} - ) - target["raw_usage"] = {**current_raw, **source_raw_usage} - + for field in fields: + value = source_fields.get(field) + if value is not None: + target_fields[field] = (target_fields.get(field) or 0) + value + _merge_incremental_usage_metadata(target, source) elif mode == "cumulative": - # Replace with latest values (already cumulative) - if source.get("input_tokens") is not None: - target["input_tokens"] = source["input_tokens"] - if source.get("output_tokens") is not None: - target["output_tokens"] = source["output_tokens"] - if source.get("cache_read_input_tokens") is not None: - target["cache_read_input_tokens"] = source["cache_read_input_tokens"] - if source.get("cache_creation_input_tokens") is not None: - target["cache_creation_input_tokens"] = source[ - "cache_creation_input_tokens" - ] - if source.get("reasoning_tokens") is not None: - target["reasoning_tokens"] = source["reasoning_tokens"] - if source.get("web_search_count") is not None: - target["web_search_count"] = source["web_search_count"] - if source.get("cache_reporting_exclusive") is not None: - target["cache_reporting_exclusive"] = source["cache_reporting_exclusive"] - # Note: raw_usage is already serialized by converters, so it's a dict - if source.get("raw_usage") is not None: - target["raw_usage"] = source["raw_usage"] - + for field in ( + *fields, + "web_search_count", + "cache_reporting_exclusive", + "raw_usage", + ): + if source_fields.get(field) is not None: + target_fields[field] = source_fields[field] else: raise ValueError(f"Invalid mode: {mode}. Must be 'incremental' or 'cumulative'") +def _merge_incremental_usage_metadata(target: TokenUsage, source: TokenUsage) -> None: + source_web_search = source.get("web_search_count") + if source_web_search is not None: + target["web_search_count"] = max( + target.get("web_search_count") or 0, source_web_search + ) + # Carried across, never accumulated: this describes the provider's accounting + # model, so it is the same on every chunk. + source_exclusive = source.get("cache_reporting_exclusive") + if source_exclusive is not None: + target["cache_reporting_exclusive"] = source_exclusive + # Anthropic message_start and message_delta report different parts of usage. + # Converters have already serialized raw_usage, so merge dicts without losing + # fields from earlier chunks. + source_raw_usage = source.get("raw_usage") + if source_raw_usage is not None and isinstance(source_raw_usage, dict): + current_raw_value = target.get("raw_usage") + current_raw = current_raw_value if isinstance(current_raw_value, dict) else {} + target["raw_usage"] = {**current_raw, **source_raw_usage} + + def get_model_params( kwargs: Dict[str, Any], served_service_tier: Optional[str] = None ) -> Dict[str, Any]: @@ -410,49 +384,56 @@ def merge_system_prompt( elif provider == "openai": from posthog.ai.openai.openai_converter import format_openai_input - # For OpenAI, handle both Chat Completions and Responses API - messages_param = kwargs.get("messages") - input_param = kwargs.get("input") + return _merge_openai_system_prompt(kwargs, format_openai_input) - # Get base formatted messages - messages = format_openai_input(messages_param, input_param) + # Default case - return empty list + return [] - # Check if system prompt is provided as a separate parameter - if kwargs.get("system") is not None: - has_system = any(msg.get("role") == "system" for msg in messages) - if not has_system: - system_msg = cast( - FormattedMessage, - {"role": "system", "content": kwargs.get("system")}, - ) - messages = [system_msg] + messages - - # For Responses API, add instructions to the system prompt if provided - if kwargs.get("instructions") is not None: - # Find the system message if it exists - system_idx = next( - (i for i, msg in enumerate(messages) if msg.get("role") == "system"), - None, - ) - if system_idx is not None: - # Append instructions to existing system message - system_content = messages[system_idx].get("content", "") - messages[system_idx]["content"] = ( - f"{system_content}\n\n{kwargs.get('instructions')}" - ) - else: - # Create a new system message with instructions - instruction_msg = cast( - FormattedMessage, - {"role": "system", "content": kwargs.get("instructions")}, - ) - messages = [instruction_msg] + messages +def _merge_openai_system_prompt( + kwargs: Dict[str, Any], + format_openai_input: Callable[..., List[FormattedMessage]], +) -> List[FormattedMessage]: + # For OpenAI, handle both Chat Completions and Responses API + messages_param = kwargs.get("messages") + input_param = kwargs.get("input") + + # Get base formatted messages + messages = format_openai_input(messages_param, input_param) + + # Check if system prompt is provided as a separate parameter + if kwargs.get("system") is not None: + has_system = any(msg.get("role") == "system" for msg in messages) + if not has_system: + system_msg = cast( + FormattedMessage, + {"role": "system", "content": kwargs.get("system")}, + ) + messages = [system_msg] + messages + + # For Responses API, add instructions to the system prompt if provided + if kwargs.get("instructions") is not None: + # Find the system message if it exists + system_idx = next( + (i for i, msg in enumerate(messages) if msg.get("role") == "system"), + None, + ) - return messages + if system_idx is not None: + # Append instructions to existing system message + system_content = messages[system_idx].get("content", "") + messages[system_idx]["content"] = ( + f"{system_content}\n\n{kwargs.get('instructions')}" + ) + else: + # Create a new system message with instructions + instruction_msg = cast( + FormattedMessage, + {"role": "system", "content": kwargs.get("instructions")}, + ) + messages = [instruction_msg] + messages - # Default case - return empty list - return [] + return messages def _get_http_status(exc: BaseException) -> int: @@ -490,7 +471,6 @@ def call_llm_and_track_usage( response = None error = None http_status = 200 - usage: TokenUsage = TokenUsage() error_params: Dict[str, Any] = {} with new_context(client=ph_client, capture_exceptions=False): @@ -510,128 +490,21 @@ def call_llm_and_track_usage( # enable_exception_autocapture is True, similar to LangChain callbacks. # See _capture_exception_and_update_properties in langchain/callbacks.py finally: - end_time = time.time() - latency = end_time - start_time - - if posthog_trace_id is None: - posthog_trace_id = str(uuid.uuid4()) - - # Check if we have a real user distinct_id (from param or outer context) - has_person_distinct_id = ( - posthog_distinct_id is not None - or contexts.get_context_distinct_id() is not None - ) - - if not has_person_distinct_id: - # Fall back to trace_id as distinct_id when no real user id is available. - identify_context(posthog_trace_id) - - if response and ( - hasattr(response, "usage") - or (provider == "gemini" and hasattr(response, "usage_metadata")) - ): - usage = get_usage(response, provider) - - messages = merge_system_prompt(kwargs, provider) - sanitized_messages = finalize_ai_content(messages, ph_client) - - tag("$ai_provider", provider) - tag("$ai_model", kwargs.get("model") or getattr(response, "model", None)) - served_service_tier = getattr(response, "service_tier", None) - tag("$ai_model_parameters", get_model_params(kwargs, served_service_tier)) - if served_service_tier is not None: - # The explicit served-tier signal cost processing prices from. - tag("$ai_service_tier", served_service_tier) - tag( - "$ai_input", - with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages), + _capture_llm_call( + start_time, + response, + posthog_trace_id, + posthog_distinct_id, + ph_client, + provider, + kwargs, + posthog_privacy_mode, + http_status, + base_url, + posthog_properties, + error_params, + posthog_groups, ) - tag( - "$ai_output_choices", - with_privacy_mode( - ph_client, - posthog_privacy_mode, - finalize_ai_content(format_response(response, provider), ph_client), - ), - ) - tag("$ai_http_status", http_status) - # Omitted when the provider never reported a count: absent means - # unknown, 0 is a report of nothing. - input_tokens = usage.get("input_tokens") - if input_tokens is not None: - tag("$ai_input_tokens", input_tokens) - output_tokens = usage.get("output_tokens") - if output_tokens is not None: - tag("$ai_output_tokens", output_tokens) - tag("$ai_latency", latency) - tag("$ai_trace_id", posthog_trace_id) - tag("$ai_base_url", str(base_url)) - warn_if_posthog_ai_gateway(base_url) - - available_tool_calls = extract_available_tool_calls(provider, kwargs) - - if available_tool_calls: - tag("$ai_tools", available_tool_calls) - - cache_read = usage.get("cache_read_input_tokens") - if cache_read is not None and cache_read > 0: - tag("$ai_cache_read_input_tokens", cache_read) - - cache_creation = usage.get("cache_creation_input_tokens") - if cache_creation is not None and cache_creation > 0: - tag("$ai_cache_creation_input_tokens", cache_creation) - - cache_reporting_exclusive = usage.get("cache_reporting_exclusive") - if cache_reporting_exclusive is not None: - tag("$ai_cache_reporting_exclusive", cache_reporting_exclusive) - - reasoning = usage.get("reasoning_tokens") - if reasoning is not None and reasoning > 0: - tag("$ai_reasoning_tokens", reasoning) - - web_search_count = usage.get("web_search_count") - if web_search_count is not None and web_search_count > 0: - tag("$ai_web_search_count", web_search_count) - - raw_usage = usage.get("raw_usage") - if raw_usage is not None: - # Already serialized by converters - tag("$ai_usage", raw_usage) - - stop_reason = extract_stop_reason(response, provider) - if stop_reason is not None: - tag("$ai_stop_reason", stop_reason) - - if not has_person_distinct_id: - tag("$process_person_profile", False) - - # Process instructions for Responses API - if provider == "openai" and kwargs.get("instructions") is not None: - tag( - "$ai_instructions", - with_privacy_mode( - ph_client, posthog_privacy_mode, kwargs.get("instructions") - ), - ) - - # send the event to posthog - if hasattr(ph_client, "capture") and callable(ph_client.capture): - sdk_tags = get_tags() - merged_properties = { - **sdk_tags, - **(posthog_properties or {}), - **(error_params or {}), - } - merged_properties["$ai_tokens_source"] = _get_tokens_source( - sdk_tags, posthog_properties - ) - _capture_ai_event( - ph_client, - "$ai_generation", - distinct_id=contexts.get_context_distinct_id(), - properties=merged_properties, - groups=posthog_groups, - ) if error: raise error @@ -655,7 +528,6 @@ async def call_llm_and_track_usage_async( response = None error = None http_status = 200 - usage: TokenUsage = TokenUsage() error_params: Dict[str, Any] = {} with new_context(client=ph_client, capture_exceptions=False): @@ -682,128 +554,21 @@ async def call_llm_and_track_usage_async( # enable_exception_autocapture is True, similar to LangChain callbacks. # See _capture_exception_and_update_properties in langchain/callbacks.py finally: - end_time = time.time() - latency = end_time - start_time - - if posthog_trace_id is None: - posthog_trace_id = str(uuid.uuid4()) - - # Check if we have a real user distinct_id (from param or outer context) - has_person_distinct_id = ( - posthog_distinct_id is not None - or contexts.get_context_distinct_id() is not None - ) - - if not has_person_distinct_id: - # Fall back to trace_id as distinct_id when no real user id is available. - identify_context(posthog_trace_id) - - if response and ( - hasattr(response, "usage") - or (provider == "gemini" and hasattr(response, "usage_metadata")) - ): - usage = get_usage(response, provider) - - messages = merge_system_prompt(kwargs, provider) - sanitized_messages = finalize_ai_content(messages, ph_client) - - tag("$ai_provider", provider) - tag("$ai_model", kwargs.get("model") or getattr(response, "model", None)) - served_service_tier = getattr(response, "service_tier", None) - tag("$ai_model_parameters", get_model_params(kwargs, served_service_tier)) - if served_service_tier is not None: - # The explicit served-tier signal cost processing prices from. - tag("$ai_service_tier", served_service_tier) - tag( - "$ai_input", - with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages), - ) - tag( - "$ai_output_choices", - with_privacy_mode( - ph_client, - posthog_privacy_mode, - finalize_ai_content(format_response(response, provider), ph_client), - ), + _capture_llm_call( + start_time, + response, + posthog_trace_id, + posthog_distinct_id, + ph_client, + provider, + kwargs, + posthog_privacy_mode, + http_status, + base_url, + posthog_properties, + error_params, + posthog_groups, ) - tag("$ai_http_status", http_status) - # Omitted when the provider never reported a count: absent means - # unknown, 0 is a report of nothing. - input_tokens = usage.get("input_tokens") - if input_tokens is not None: - tag("$ai_input_tokens", input_tokens) - output_tokens = usage.get("output_tokens") - if output_tokens is not None: - tag("$ai_output_tokens", output_tokens) - tag("$ai_latency", latency) - tag("$ai_trace_id", posthog_trace_id) - tag("$ai_base_url", str(base_url)) - warn_if_posthog_ai_gateway(base_url) - - available_tool_calls = extract_available_tool_calls(provider, kwargs) - - if available_tool_calls: - tag("$ai_tools", available_tool_calls) - - cache_read = usage.get("cache_read_input_tokens") - if cache_read is not None and cache_read > 0: - tag("$ai_cache_read_input_tokens", cache_read) - - cache_creation = usage.get("cache_creation_input_tokens") - if cache_creation is not None and cache_creation > 0: - tag("$ai_cache_creation_input_tokens", cache_creation) - - cache_reporting_exclusive = usage.get("cache_reporting_exclusive") - if cache_reporting_exclusive is not None: - tag("$ai_cache_reporting_exclusive", cache_reporting_exclusive) - - reasoning = usage.get("reasoning_tokens") - if reasoning is not None and reasoning > 0: - tag("$ai_reasoning_tokens", reasoning) - - web_search_count = usage.get("web_search_count") - if web_search_count is not None and web_search_count > 0: - tag("$ai_web_search_count", web_search_count) - - raw_usage = usage.get("raw_usage") - if raw_usage is not None: - # Already serialized by converters - tag("$ai_usage", raw_usage) - - stop_reason = extract_stop_reason(response, provider) - if stop_reason is not None: - tag("$ai_stop_reason", stop_reason) - - if not has_person_distinct_id: - tag("$process_person_profile", False) - - # Process instructions for Responses API - if provider == "openai" and kwargs.get("instructions") is not None: - tag( - "$ai_instructions", - with_privacy_mode( - ph_client, posthog_privacy_mode, kwargs.get("instructions") - ), - ) - - # send the event to posthog - if hasattr(ph_client, "capture") and callable(ph_client.capture): - sdk_tags = get_tags() - merged_properties = { - **sdk_tags, - **(posthog_properties or {}), - **(error_params or {}), - } - merged_properties["$ai_tokens_source"] = _get_tokens_source( - sdk_tags, posthog_properties - ) - _capture_ai_event( - ph_client, - "$ai_generation", - distinct_id=contexts.get_context_distinct_id(), - properties=merged_properties, - groups=posthog_groups, - ) if error: raise error @@ -811,6 +576,178 @@ async def call_llm_and_track_usage_async( return response +def _capture_llm_call( + start_time: float, + response: Any, + posthog_trace_id: Optional[str], + posthog_distinct_id: Optional[str], + ph_client: PostHogClient, + provider: str, + kwargs: Dict[str, Any], + posthog_privacy_mode: bool, + http_status: int, + base_url: str, + posthog_properties: Optional[Dict[str, Any]], + error_params: Dict[str, Any], + posthog_groups: Optional[Dict[str, Any]], +) -> None: + usage: TokenUsage = TokenUsage() + end_time = time.time() + latency = end_time - start_time + + if posthog_trace_id is None: + posthog_trace_id = str(uuid.uuid4()) + + # Check if we have a real user distinct_id (from param or outer context) + has_person_distinct_id = ( + posthog_distinct_id is not None + or contexts.get_context_distinct_id() is not None + ) + + if not has_person_distinct_id: + # Fall back to trace_id as distinct_id when no real user id is available. + identify_context(posthog_trace_id) + + if response and ( + hasattr(response, "usage") + or (provider == "gemini" and hasattr(response, "usage_metadata")) + ): + usage = get_usage(response, provider) + + _tag_llm_content(response, provider, kwargs, ph_client, posthog_privacy_mode) + tag("$ai_http_status", http_status) + _tag_reported_token_counts(usage) + tag("$ai_latency", latency) + tag("$ai_trace_id", posthog_trace_id) + tag("$ai_base_url", str(base_url)) + warn_if_posthog_ai_gateway(base_url) + + available_tool_calls = extract_available_tool_calls(provider, kwargs) + + if available_tool_calls: + tag("$ai_tools", available_tool_calls) + + _tag_optional_usage(usage) + _capture_tagged_llm_event( + response, + provider, + kwargs, + ph_client, + posthog_privacy_mode, + has_person_distinct_id, + posthog_properties, + error_params, + posthog_groups, + ) + + +def _tag_llm_content( + response: Any, + provider: str, + kwargs: Dict[str, Any], + ph_client: PostHogClient, + posthog_privacy_mode: bool, +) -> None: + messages = merge_system_prompt(kwargs, provider) + sanitized_messages = finalize_ai_content(messages, ph_client) + + tag("$ai_provider", provider) + tag("$ai_model", kwargs.get("model") or getattr(response, "model", None)) + served_service_tier = getattr(response, "service_tier", None) + tag("$ai_model_parameters", get_model_params(kwargs, served_service_tier)) + if served_service_tier is not None: + # The explicit served-tier signal cost processing prices from. + tag("$ai_service_tier", served_service_tier) + tag( + "$ai_input", + with_privacy_mode(ph_client, posthog_privacy_mode, sanitized_messages), + ) + tag( + "$ai_output_choices", + with_privacy_mode( + ph_client, + posthog_privacy_mode, + finalize_ai_content(format_response(response, provider), ph_client), + ), + ) + + +def _tag_reported_token_counts(usage: TokenUsage) -> None: + # Omitted when the provider never reported a count: absent means + # unknown, 0 is a report of nothing. + input_tokens = usage.get("input_tokens") + if input_tokens is not None: + tag("$ai_input_tokens", input_tokens) + output_tokens = usage.get("output_tokens") + if output_tokens is not None: + tag("$ai_output_tokens", output_tokens) + + +def _tag_optional_usage(usage: TokenUsage) -> None: + for field in ("cache_read_input_tokens", "cache_creation_input_tokens"): + value = cast(Optional[int], usage.get(field)) + if value is not None and value > 0: + tag(f"$ai_{field}", value) + exclusive = usage.get("cache_reporting_exclusive") + if exclusive is not None: + tag("$ai_cache_reporting_exclusive", exclusive) + for field in ("reasoning_tokens", "web_search_count"): + value = cast(Optional[int], usage.get(field)) + if value is not None and value > 0: + tag(f"$ai_{field}", value) + raw_usage = usage.get("raw_usage") + if raw_usage is not None: + # Already serialized by converters + tag("$ai_usage", raw_usage) + + +def _capture_tagged_llm_event( + response: Any, + provider: str, + kwargs: Dict[str, Any], + ph_client: PostHogClient, + posthog_privacy_mode: bool, + has_person_distinct_id: bool, + posthog_properties: Optional[Dict[str, Any]], + error_params: Dict[str, Any], + posthog_groups: Optional[Dict[str, Any]], +) -> None: + stop_reason = extract_stop_reason(response, provider) + if stop_reason is not None: + tag("$ai_stop_reason", stop_reason) + + if not has_person_distinct_id: + tag("$process_person_profile", False) + + # Process instructions for Responses API + if provider == "openai" and kwargs.get("instructions") is not None: + tag( + "$ai_instructions", + with_privacy_mode( + ph_client, posthog_privacy_mode, kwargs.get("instructions") + ), + ) + + # send the event to posthog + if hasattr(ph_client, "capture") and callable(ph_client.capture): + sdk_tags = get_tags() + merged_properties = { + **sdk_tags, + **(posthog_properties or {}), + **(error_params or {}), + } + merged_properties["$ai_tokens_source"] = _get_tokens_source( + sdk_tags, posthog_properties + ) + _capture_ai_event( + ph_client, + "$ai_generation", + distinct_id=contexts.get_context_distinct_id(), + properties=merged_properties, + groups=posthog_groups, + ) + + def finalize_ai_content(value: Any, ph_client: Any = None) -> Any: """Single choke point for AI content properties: structural media redaction (or bytes->base64 passthrough when the client opted into multimodal capture). @@ -1074,46 +1011,9 @@ def capture_streaming_event( input_tokens = event_data["usage_stats"].get("input_tokens") output_tokens = event_data["usage_stats"].get("output_tokens") - # Build base event properties - event_properties = { - "$ai_provider": event_data["provider"], - "$ai_model": event_data["model"], - "$ai_model_parameters": get_model_params( - event_data["kwargs"], event_data.get("service_tier") - ), - **( - {"$ai_service_tier": event_data["service_tier"]} - if event_data.get("service_tier") is not None - else {} - ), - "$ai_input": with_privacy_mode( - ph_client, - event_data["privacy_mode"], - finalize_ai_content(event_data["formatted_input"], ph_client), - ), - "$ai_output_choices": with_privacy_mode( - ph_client, - event_data["privacy_mode"], - finalize_ai_content(event_data["formatted_output"], ph_client), - ), - "$ai_http_status": (_get_http_status(error) if error is not None else 200), - **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), - **({"$ai_output_tokens": output_tokens} if output_tokens is not None else {}), - "$ai_latency": event_data["latency"], - "$ai_trace_id": trace_id, - "$ai_base_url": str(event_data["base_url"]), - **(event_data.get("properties") or {}), - **( - { - "$ai_is_error": True, - "$ai_error": _format_ai_error( - ph_client, event_data["privacy_mode"], error - ), - } - if error is not None - else {} - ), - } + event_properties = _streaming_base_properties( + ph_client, event_data, trace_id, error, input_tokens, output_tokens + ) warn_if_posthog_ai_gateway(event_data["base_url"]) @@ -1134,6 +1034,44 @@ def capture_streaming_event( if available_tools: event_properties["$ai_tools"] = available_tools + _add_streaming_usage(event_properties, event_data, input_tokens, output_tokens) + + # Add stop reason if present + stop_reason = event_data.get("stop_reason") + if stop_reason is not None: + event_properties["$ai_stop_reason"] = stop_reason + + # Handle provider-specific fields + if ( + event_data["provider"] == "openai" + and event_data["kwargs"].get("instructions") is not None + ): + event_properties["$ai_instructions"] = with_privacy_mode( + ph_client, + event_data["privacy_mode"], + event_data["kwargs"]["instructions"], + ) + + if event_data.get("distinct_id") is None: + event_properties["$process_person_profile"] = False + + # Send event to PostHog + if hasattr(ph_client, "capture"): + _capture_ai_event( + ph_client, + "$ai_generation", + distinct_id=event_data.get("distinct_id") or trace_id, + properties=event_properties, + groups=event_data.get("groups"), + ) + + +def _add_streaming_usage( + event_properties: Dict[str, Any], + event_data: StreamingEventData, + input_tokens: Optional[int], + output_tokens: Optional[int], +) -> None: # Add optional token fields. The zero-defaults below only apply when the # provider reported usage at all: a stream interrupted before any report # has nothing to default, and a fabricated 0 would read as a report of @@ -1151,42 +1089,54 @@ def capture_streaming_event( if isinstance(reasoning, int) and reasoning > 0: event_properties["$ai_reasoning_tokens"] = reasoning else: - optional_token_fields = [ - "cache_read_input_tokens", - "cache_creation_input_tokens", - "reasoning_tokens", - ] + _add_optional_streaming_tokens(event_properties, event_data, usage_was_reported) - for field in optional_token_fields: - value = event_data["usage_stats"].get(field) - property_name = f"$ai_{field}" + _add_streaming_usage_metadata(event_properties, event_data["usage_stats"]) - # OpenAI async streams historically included these fields even when 0. - # Keep those defaults in the shared path so they are not mistaken for - # caller-supplied token passthrough properties. - if ( - event_data["provider"] == "openai" - and usage_was_reported - and field - in { - "cache_read_input_tokens", - "reasoning_tokens", - } - ): - event_properties.setdefault( - property_name, event_data["usage_stats"].get(field, 0) - ) - elif value is not None and isinstance(value, int) and value > 0: - event_properties[property_name] = value - cache_reporting_exclusive = event_data["usage_stats"].get( - "cache_reporting_exclusive" - ) +def _add_optional_streaming_tokens( + event_properties: Dict[str, Any], + event_data: StreamingEventData, + usage_was_reported: bool, +) -> None: + optional_token_fields = [ + "cache_read_input_tokens", + "cache_creation_input_tokens", + "reasoning_tokens", + ] + + for field in optional_token_fields: + value = event_data["usage_stats"].get(field) + property_name = f"$ai_{field}" + + # OpenAI async streams historically included these fields even when 0. + # Keep those defaults in the shared path so they are not mistaken for + # caller-supplied token passthrough properties. + if ( + event_data["provider"] == "openai" + and usage_was_reported + and field + in { + "cache_read_input_tokens", + "reasoning_tokens", + } + ): + event_properties.setdefault( + property_name, event_data["usage_stats"].get(field, 0) + ) + elif value is not None and isinstance(value, int) and value > 0: + event_properties[property_name] = value + + +def _add_streaming_usage_metadata( + event_properties: Dict[str, Any], usage: TokenUsage +) -> None: + cache_reporting_exclusive = usage.get("cache_reporting_exclusive") if cache_reporting_exclusive is not None: event_properties["$ai_cache_reporting_exclusive"] = cache_reporting_exclusive # Add web search count if present (all providers) - web_search_count = event_data["usage_stats"].get("web_search_count") + web_search_count = usage.get("web_search_count") if ( web_search_count is not None and isinstance(web_search_count, int) @@ -1195,36 +1145,57 @@ def capture_streaming_event( event_properties["$ai_web_search_count"] = web_search_count # Add raw usage metadata if present (all providers) - raw_usage = event_data["usage_stats"].get("raw_usage") + raw_usage = usage.get("raw_usage") if raw_usage is not None: # Already serialized by converters event_properties["$ai_usage"] = raw_usage - # Add stop reason if present - stop_reason = event_data.get("stop_reason") - if stop_reason is not None: - event_properties["$ai_stop_reason"] = stop_reason - # Handle provider-specific fields - if ( - event_data["provider"] == "openai" - and event_data["kwargs"].get("instructions") is not None - ): - event_properties["$ai_instructions"] = with_privacy_mode( +def _streaming_base_properties( + ph_client: PostHogClient, + event_data: StreamingEventData, + trace_id: str, + error: Optional[BaseException], + input_tokens: Optional[int], + output_tokens: Optional[int], +) -> Dict[str, Any]: + # Build base event properties + return { + "$ai_provider": event_data["provider"], + "$ai_model": event_data["model"], + "$ai_model_parameters": get_model_params( + event_data["kwargs"], event_data.get("service_tier") + ), + **( + {"$ai_service_tier": event_data["service_tier"]} + if event_data.get("service_tier") is not None + else {} + ), + "$ai_input": with_privacy_mode( ph_client, event_data["privacy_mode"], - event_data["kwargs"]["instructions"], - ) - - if event_data.get("distinct_id") is None: - event_properties["$process_person_profile"] = False - - # Send event to PostHog - if hasattr(ph_client, "capture"): - _capture_ai_event( + finalize_ai_content(event_data["formatted_input"], ph_client), + ), + "$ai_output_choices": with_privacy_mode( ph_client, - "$ai_generation", - distinct_id=event_data.get("distinct_id") or trace_id, - properties=event_properties, - groups=event_data.get("groups"), - ) + event_data["privacy_mode"], + finalize_ai_content(event_data["formatted_output"], ph_client), + ), + "$ai_http_status": (_get_http_status(error) if error is not None else 200), + **({"$ai_input_tokens": input_tokens} if input_tokens is not None else {}), + **({"$ai_output_tokens": output_tokens} if output_tokens is not None else {}), + "$ai_latency": event_data["latency"], + "$ai_trace_id": trace_id, + "$ai_base_url": str(event_data["base_url"]), + **(event_data.get("properties") or {}), + **( + { + "$ai_is_error": True, + "$ai_error": _format_ai_error( + ph_client, event_data["privacy_mode"], error + ), + } + if error is not None + else {} + ), + } diff --git a/posthog/async_client.py b/posthog/async_client.py index 7e5c3b3ff..ae63d0853 100644 --- a/posthog/async_client.py +++ b/posthog/async_client.py @@ -10,7 +10,8 @@ import warnings import weakref from datetime import datetime, timezone -from typing import Any, Callable, Dict, Mapping, Optional, Union +from functools import partial as _partial +from typing import AbstractSet, Any, Callable, Dict, Mapping, Optional, Union from uuid import UUID, uuid4 from typing_extensions import Unpack @@ -165,24 +166,7 @@ def __init__( capture_compression, gzip_fallback=gzip ) self.capture_trace_context = capture_trace_context - if personal_api_key is not None and secret_key is None: - warnings.warn( - "`personal_api_key` is deprecated; use `secret_key` instead.", - DeprecationWarning, - stacklevel=2, - ) - elif personal_api_key is not None and secret_key is not None: - self.log.warning( - "[FEATURE FLAGS] Both secret_key and personal_api_key were provided; " - "using secret_key." - ) - resolved_secret_key = secret_key if secret_key is not None else personal_api_key - self.secret_key = ( - resolved_secret_key.strip() - if isinstance(resolved_secret_key, str) - else resolved_secret_key - ) or None - self.personal_api_key = self.secret_key + self._configure_secret_key(secret_key, personal_api_key) _validate_options( email_request_timeout_seconds, email_request_max_retries, migrated_from ) @@ -203,35 +187,16 @@ def __init__( self._feature_flag_evaluations_host_cache: Optional[ _FeatureFlagEvaluationsHost ] = None - self.log_captured_exceptions = log_captured_exceptions - self.capture_exception_code_variables = capture_exception_code_variables - self.code_variables_mask_patterns = ( - code_variables_mask_patterns - if code_variables_mask_patterns is not None - else DEFAULT_CODE_VARIABLES_MASK_PATTERNS - ) - self.code_variables_ignore_patterns = ( - code_variables_ignore_patterns - if code_variables_ignore_patterns is not None - else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS - ) - self.code_variables_mask_url_credentials = ( - code_variables_mask_url_credentials - if code_variables_mask_url_credentials is not None - else DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS + self._configure_exception_capture( + log_captured_exceptions, + capture_exception_code_variables, + code_variables_mask_patterns, + code_variables_ignore_patterns, + code_variables_mask_url_credentials, + code_variables_detect_secrets, + in_app_modules, + project_root, ) - self.code_variables_detect_secrets = ( - code_variables_detect_secrets - if code_variables_detect_secrets is not None - else DEFAULT_CODE_VARIABLES_DETECT_SECRETS - ) - self.in_app_modules = in_app_modules - self.project_root = project_root - if self.project_root is None: - try: - self.project_root = os.getcwd() - except Exception: - self.project_root = None if before_send is not None and not callable(before_send): self.log.warning("before_send is not callable, it will be ignored") @@ -269,6 +234,69 @@ def __init__( self._duplicate_client_registry_key: Optional[tuple[str, str]] = None self._register_duplicate_client() + def _configure_secret_key( + self, secret_key: Optional[str], personal_api_key: Optional[str] + ) -> None: + if personal_api_key is not None and secret_key is None: + warnings.warn( + "`personal_api_key` is deprecated; use `secret_key` instead.", + DeprecationWarning, + stacklevel=3, + ) + elif personal_api_key is not None and secret_key is not None: + self.log.warning( + "[FEATURE FLAGS] Both secret_key and personal_api_key were provided; " + "using secret_key." + ) + resolved_secret_key = secret_key if secret_key is not None else personal_api_key + self.secret_key = ( + resolved_secret_key.strip() + if isinstance(resolved_secret_key, str) + else resolved_secret_key + ) or None + self.personal_api_key = self.secret_key + + def _configure_exception_capture( + self, + log_captured_exceptions: bool, + capture_exception_code_variables: bool, + code_variables_mask_patterns, + code_variables_ignore_patterns, + code_variables_mask_url_credentials, + code_variables_detect_secrets, + in_app_modules: Optional[list[str]], + project_root: Optional[str], + ) -> None: + self.log_captured_exceptions = log_captured_exceptions + self.capture_exception_code_variables = capture_exception_code_variables + self.code_variables_mask_patterns = ( + code_variables_mask_patterns + if code_variables_mask_patterns is not None + else DEFAULT_CODE_VARIABLES_MASK_PATTERNS + ) + self.code_variables_ignore_patterns = ( + code_variables_ignore_patterns + if code_variables_ignore_patterns is not None + else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS + ) + self.code_variables_mask_url_credentials = ( + code_variables_mask_url_credentials + if code_variables_mask_url_credentials is not None + else DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS + ) + self.code_variables_detect_secrets = ( + code_variables_detect_secrets + if code_variables_detect_secrets is not None + else DEFAULT_CODE_VARIABLES_DETECT_SECRETS + ) + self.in_app_modules = in_app_modules + self.project_root = project_root + if self.project_root is None: + try: + self.project_root = os.getcwd() + except Exception: + self.project_root = None + _emails: Optional[AsyncEmails] = None @property @@ -428,28 +456,40 @@ def _enqueue_prepared_event(self, prepared: dict[str, Any]) -> bool: self._ensure_workers_started() self._queue.put_nowait(queued_event) return True + return self._enqueue_from_thread(self._loop, queued_event, running_loop) + + def _enqueue_from_thread( + self, + loop: asyncio.AbstractEventLoop, + queued_event: _QueuedEvent, + running_loop: Optional[asyncio.AbstractEventLoop], + ) -> bool: if running_loop is not None: raise RuntimeError("AsyncClient cannot be shared across event loops") - if self._loop_thread_id == threading.get_ident() or not self._loop.is_running(): + if self._loop_thread_id == threading.get_ident() or not loop.is_running(): raise RuntimeError("AsyncClient event loop is not running") admitted: concurrent.futures.Future[bool] = concurrent.futures.Future() - def enqueue_on_bound_loop() -> None: - try: - if not self._accepting or self._closed: - admitted.set_result(False) - return - self._ensure_workers_started() - self._queue.put_nowait(queued_event) - except BaseException as error: - admitted.set_exception(error) - else: - admitted.set_result(True) - - self._loop.call_soon_threadsafe(enqueue_on_bound_loop) + loop.call_soon_threadsafe( + _partial(self._enqueue_on_bound_loop, queued_event, admitted) + ) return admitted.result() + def _enqueue_on_bound_loop( + self, queued_event: _QueuedEvent, admitted: concurrent.futures.Future[bool] + ) -> None: + try: + if not self._accepting or self._closed: + admitted.set_result(False) + return + self._ensure_workers_started() + self._queue.put_nowait(queued_event) + except BaseException as error: + admitted.set_exception(error) + else: + admitted.set_result(True) + def _normalize_uuid(self, msg: dict[str, Any]) -> str: raw_uuid = msg.pop("uuid", None) if raw_uuid is not None: @@ -477,19 +517,28 @@ def _prepare_event( if self.disabled or not self._accepting: return None, None - timestamp = msg.get("timestamp") - if timestamp is None: - timestamp = datetime.now(tz=timezone.utc) - try: - msg["timestamp"] = _normalize_timestamp(timestamp) - except ValueError: - self.log.warning( - "Invalid timestamp %r. Falling back to the current UTC time.", - timestamp, - ) - msg["timestamp"] = datetime.now(tz=timezone.utc).isoformat() + msg["timestamp"] = self._event_timestamp(msg.get("timestamp")) sent_uuid = self._normalize_uuid(msg) + self._add_capture_metadata(msg, disable_geoip) + # After every enrichment step, so a super property never changes or adds + # config version 2 flag attribution. + if pinned_properties is not None: + msg["properties"] = pinned_properties.apply(msg["properties"]) + if property_allowlist is not None: + msg["properties"] = { + key: value + for key, value in msg["properties"].items() + if key in property_allowlist + } + + msg["distinct_id"] = _stringify_id(msg.get("distinct_id")) + cleaned = clean(msg) + return cleaned, sent_uuid + + def _add_capture_metadata( + self, msg: dict[str, Any], disable_geoip: Optional[bool] + ) -> None: properties = msg.setdefault("properties", {}) properties["$lib"] = "posthog-python" properties["$lib_version"] = VERSION @@ -504,20 +553,19 @@ def _prepare_event( msg["properties"].setdefault("$release_id", self._release_id) if self.is_server: msg["properties"]["$is_server"] = True - # After every enrichment step, so a super property never changes or adds - # config version 2 flag attribution. - if pinned_properties is not None: - msg["properties"] = pinned_properties.apply(msg["properties"]) - if property_allowlist is not None: - msg["properties"] = { - key: value - for key, value in msg["properties"].items() - if key in property_allowlist - } - msg["distinct_id"] = _stringify_id(msg.get("distinct_id")) - cleaned = clean(msg) - return cleaned, sent_uuid + def _event_timestamp(self, value: Any) -> str: + timestamp = value + if timestamp is None: + timestamp = datetime.now(tz=timezone.utc) + try: + return _normalize_timestamp(timestamp) + except ValueError: + self.log.warning( + "Invalid timestamp %r. Falling back to the current UTC time.", + timestamp, + ) + return datetime.now(tz=timezone.utc).isoformat() async def _process_event(self, msg: dict[str, Any]) -> Optional[dict[str, Any]]: if self.before_send is None: @@ -623,13 +671,10 @@ async def capture_immediate( self, event: str, **kwargs: Unpack[OptionalCaptureArgs] ) -> Optional[str]: """Capture one event and wait until its delivery attempt completes.""" - current = asyncio.current_task() - if current is None: # pragma: no cover - async functions always have a task + admitted = self._start_immediate_call() + if admitted is None: return None - if not self._accepting: - return None - completion = self._start_inflight_operation() - self._immediate_callers[current] = self._immediate_callers.get(current, 0) + 1 + current, completion = admitted error_batch: list[dict[str, Any]] = [] try: msg, disable_geoip, property_allowlist, pinned_properties = ( @@ -646,33 +691,14 @@ async def capture_immediate( error_batch = [processed] if not self.send: return sent_uuid - try: - event_size = await _serialized_event_size(processed) - except Exception: - self.log.error( - "unable to serialize immediate event for sizing, dropping" - ) - return None - if event_size > MAX_MSG_SIZE: - self.log.error( - "Event %s (%d bytes) exceeds the %dKiB limit, dropping.", - processed.get("event"), - event_size, - MAX_MSG_SIZE // 1024, - ) + if not await self._immediate_event_is_sendable(processed): return None consumer = self._new_consumer() await consumer.request(error_batch) return sent_uuid except Exception as error: - if self.on_error: - try: - await _invoke_callback(self.on_error, error, error_batch) - except Exception as callback_error: - self.log.error( - "on_error handler failed (%s)", type(callback_error).__name__ - ) + await self._notify_immediate_error(error, error_batch) if self.debug: raise self.log.error( @@ -682,12 +708,56 @@ async def capture_immediate( ) return None finally: - remaining_calls = self._immediate_callers[current] - 1 - if remaining_calls: - self._immediate_callers[current] = remaining_calls - else: - del self._immediate_callers[current] - self._finish_inflight_operation(completion) + self._finish_immediate_call(current, completion) + + def _start_immediate_call( + self, + ) -> Optional[tuple[asyncio.Task[Any], asyncio.Future[None]]]: + current = asyncio.current_task() + if current is None: # pragma: no cover - async functions always have a task + return None + if not self._accepting: + return None + completion = self._start_inflight_operation() + self._immediate_callers[current] = self._immediate_callers.get(current, 0) + 1 + return current, completion + + def _finish_immediate_call( + self, current: asyncio.Task[Any], completion: asyncio.Future[None] + ) -> None: + remaining_calls = self._immediate_callers[current] - 1 + if remaining_calls: + self._immediate_callers[current] = remaining_calls + else: + del self._immediate_callers[current] + self._finish_inflight_operation(completion) + + async def _notify_immediate_error( + self, error: Exception, error_batch: list[dict[str, Any]] + ) -> None: + if self.on_error: + try: + await _invoke_callback(self.on_error, error, error_batch) + except Exception as callback_error: + self.log.error( + "on_error handler failed (%s)", type(callback_error).__name__ + ) + + async def _immediate_event_is_sendable(self, processed: dict[str, Any]) -> bool: + try: + event_size = await _serialized_event_size(processed) + except Exception: + self.log.error("unable to serialize immediate event for sizing, dropping") + return False + if event_size > MAX_MSG_SIZE: + self.log.error( + "Event %s (%d bytes) exceeds the %dKiB limit, dropping.", + processed.get("event"), + event_size, + MAX_MSG_SIZE // 1024, + ) + return False + return True def _build_person_properties_event( self, event: str, property_key: str, kwargs: OptionalSetArgs @@ -829,16 +899,8 @@ def capture_exception( ) -> Optional[str]: """Capture an exception. This method never raises, including in debug mode.""" try: - if exception is not None and exception_is_already_captured(exception): - self.log.debug("Exception already captured, skipping") - return None - exc_info = ( - exc_info_from_error(exception) - if exception is not None - else sys.exc_info() - ) - if exc_info is None or exc_info == (None, None, None): - self.log.warning("No exception information available") + exc_info = self._exception_info(exception) + if exc_info is None: return None exceptions = exceptions_from_error_tuple(exc_info) @@ -852,43 +914,7 @@ def capture_exception( "$exception_list": exceptions, **(kwargs.get("properties") or {}), } - context_enabled = get_capture_exception_code_variables_context() - context_mask = get_code_variables_mask_patterns_context() - context_ignore = get_code_variables_ignore_patterns_context() - context_mask_url_credentials = ( - get_code_variables_mask_url_credentials_context() - ) - context_detect_secrets = get_code_variables_detect_secrets_context() - enabled = ( - context_enabled - if context_enabled is not None - else self.capture_exception_code_variables - ) - if enabled: - try_attach_code_variables_to_frames( - exceptions, - exc_info, - mask_patterns=( - context_mask - if context_mask is not None - else self.code_variables_mask_patterns - ), - ignore_patterns=( - context_ignore - if context_ignore is not None - else self.code_variables_ignore_patterns - ), - mask_url_credentials=( - context_mask_url_credentials - if context_mask_url_credentials is not None - else self.code_variables_mask_url_credentials - ), - detect_secrets=( - context_detect_secrets - if context_detect_secrets is not None - else self.code_variables_detect_secrets - ), - ) + self._attach_exception_code_variables(exceptions, exc_info) if self.log_captured_exceptions: self.log.exception(exception, extra=kwargs) @@ -909,6 +935,56 @@ def capture_exception( self.log.exception("Failed to capture exception: %s", error) return None + def _attach_exception_code_variables(self, exceptions, exc_info) -> None: + context_enabled = get_capture_exception_code_variables_context() + context_mask = get_code_variables_mask_patterns_context() + context_ignore = get_code_variables_ignore_patterns_context() + context_mask_url_credentials = get_code_variables_mask_url_credentials_context() + context_detect_secrets = get_code_variables_detect_secrets_context() + enabled = ( + context_enabled + if context_enabled is not None + else self.capture_exception_code_variables + ) + if enabled: + try_attach_code_variables_to_frames( + exceptions, + exc_info, + mask_patterns=( + context_mask + if context_mask is not None + else self.code_variables_mask_patterns + ), + ignore_patterns=( + context_ignore + if context_ignore is not None + else self.code_variables_ignore_patterns + ), + mask_url_credentials=( + context_mask_url_credentials + if context_mask_url_credentials is not None + else self.code_variables_mask_url_credentials + ), + detect_secrets=( + context_detect_secrets + if context_detect_secrets is not None + else self.code_variables_detect_secrets + ), + ) + + def _exception_info(self, exception: Optional[ExceptionArg]): + if exception is not None and exception_is_already_captured(exception): + self.log.debug("Exception already captured, skipping") + return None + exc_info = ( + exc_info_from_error(exception) if exception is not None else sys.exc_info() + ) + if exc_info is None or exc_info == (None, None, None): + self.log.warning("No exception information available") + return None + + return exc_info + def _get_feature_flag_evaluations_host(self) -> _FeatureFlagEvaluationsHost: if self._feature_flag_evaluations_host_cache is None: self._feature_flag_evaluations_host_cache = _FeatureFlagEvaluationsHost( @@ -974,10 +1050,7 @@ async def evaluate_flags( ``unresolved_flags`` is always empty. """ host = self._get_feature_flag_evaluations_host() - if distinct_id is None: - distinct_id = _get_context_distinct_id() - if device_id is None: - device_id = _get_context_device_id() + distinct_id, device_id = _evaluation_context_ids(distinct_id, device_id) if distinct_id is None or distinct_id == "" or self.disabled: return FeatureFlagEvaluations(host=host, distinct_id="", flags={}) @@ -991,12 +1064,6 @@ async def evaluate_flags( disable_geoip=disable_geoip, ) - records: dict[str, _EvaluatedFlagRecord] = {} - request_id: Optional[str] = None - evaluated_at: Optional[int] = None - errors_while_computing = False - quota_limited = False - minimal_flag_called_events = False requested_keys = set(flag_keys) if flag_keys else None if not self._accepting: return FeatureFlagEvaluations( @@ -1006,12 +1073,43 @@ async def evaluate_flags( groups=resolved_groups, disable_geoip=disable_geoip, ) + return await self._request_flag_snapshot( + host, + distinct_id, + requested_keys, + groups=resolved_groups, + person_properties=person_properties, + group_properties=group_properties, + disable_geoip=disable_geoip, + flag_keys=flag_keys, + device_id=device_id, + ) + + async def _request_flag_snapshot( + self, + host: _FeatureFlagEvaluationsHost, + distinct_id: ID_TYPES, + requested_keys: Optional[AbstractSet[str]], + *, + groups: Mapping[str, Union[str, int]], + person_properties: Optional[Dict[str, Any]], + group_properties: Optional[Dict[str, Dict[str, Any]]], + disable_geoip: Optional[bool], + flag_keys: Optional[list[str]], + device_id: Optional[str], + ) -> FeatureFlagEvaluations: + records: dict[str, _EvaluatedFlagRecord] = {} + request_id: Optional[str] = None + evaluated_at: Optional[int] = None + errors_while_computing = False + quota_limited = False + minimal_flag_called_events = False completion = self._start_inflight_operation() try: response = await self._get_flags_decision( distinct_id, - groups=resolved_groups, + groups=groups, person_properties=person_properties, group_properties=group_properties, disable_geoip=disable_geoip, @@ -1049,7 +1147,7 @@ async def evaluate_flags( host=host, distinct_id=str(distinct_id), flags=records, - groups=resolved_groups, + groups=groups, disable_geoip=disable_geoip, request_id=request_id, evaluated_at=evaluated_at, @@ -1237,13 +1335,7 @@ async def shutdown(self) -> None: self._accepting = False errors: list[Exception] = [] - pending_operations = [ - completion - for completion in self._inflight_operations - if not completion.done() - ] - if pending_operations: - await asyncio.gather(*pending_operations, return_exceptions=True) + await self._wait_for_inflight_operations() try: await self.flush(timeout_seconds=None) @@ -1251,37 +1343,9 @@ async def shutdown(self) -> None: self.log.exception("Failed to flush async capture queue") errors.append(error) - try: - live_workers = [task for task in self._worker_tasks if not task.done()] - for _ in live_workers: - await self._queue.put(_STOP) - if self._worker_tasks: - await asyncio.gather(*self._worker_tasks, return_exceptions=True) - except Exception as error: - self.log.exception("Failed to stop async capture workers") - errors.append(error) - finally: - self._worker_tasks.clear() - self._consumers.clear() - - try: - await self._close_transport() - except Exception as error: - self.log.exception("Failed to close async capture transport") - errors.append(error) - - if self._emails is not None: - try: - await self._emails.aclose() - except Exception as error: - self.log.exception("Failed to close the async email transport") - errors.append(error) + await self._stop_workers(errors) - try: - self._unregister_duplicate_client() - except Exception as error: - self.log.exception("Failed to unregister async client") - errors.append(error) + await self._close_resources(errors) self.distinct_ids_feature_flags_reported.clear() self._experiment_exposures_reported.clear() @@ -1289,6 +1353,49 @@ async def shutdown(self) -> None: if errors and self.debug: raise errors[0] + async def _close_resources(self, errors: list[Exception]) -> None: + try: + await self._close_transport() + except Exception as error: + self.log.exception("Failed to close async capture transport") + errors.append(error) + + if self._emails is not None: + try: + await self._emails.aclose() + except Exception as error: + self.log.exception("Failed to close the async email transport") + errors.append(error) + + try: + self._unregister_duplicate_client() + except Exception as error: + self.log.exception("Failed to unregister async client") + errors.append(error) + + async def _stop_workers(self, errors: list[Exception]) -> None: + try: + live_workers = [task for task in self._worker_tasks if not task.done()] + for _ in live_workers: + await self._queue.put(_STOP) + if self._worker_tasks: + await asyncio.gather(*self._worker_tasks, return_exceptions=True) + except Exception as error: + self.log.exception("Failed to stop async capture workers") + errors.append(error) + finally: + self._worker_tasks.clear() + self._consumers.clear() + + async def _wait_for_inflight_operations(self) -> None: + pending_operations = [ + completion + for completion in self._inflight_operations + if not completion.done() + ] + if pending_operations: + await asyncio.gather(*pending_operations, return_exceptions=True) + async def join(self) -> None: await self.shutdown() @@ -1297,3 +1404,13 @@ class AsyncPosthog(AsyncClient): """Customer-facing name for :class:`AsyncClient`.""" pass + + +def _evaluation_context_ids( + distinct_id: Optional[ID_TYPES], device_id: Optional[str] +) -> tuple[Optional[ID_TYPES], Optional[str]]: + if distinct_id is None: + distinct_id = _get_context_distinct_id() + if device_id is None: + device_id = _get_context_device_id() + return distinct_id, device_id diff --git a/posthog/capture_compression.py b/posthog/capture_compression.py index 7b6760d06..41e3d8e1d 100644 --- a/posthog/capture_compression.py +++ b/posthog/capture_compression.py @@ -108,6 +108,10 @@ def _resolve_capture_compression( fallback = CaptureCompression.GZIP if gzip_fallback else CaptureCompression.NONE + return _resolve_env_compression(fallback) + + +def _resolve_env_compression(fallback: CaptureCompression) -> CaptureCompression: raw = os.environ.get(CAPTURE_COMPRESSION_ENV_VAR) if raw is None or raw.strip() == "": return fallback diff --git a/posthog/capture_v1.py b/posthog/capture_v1.py index 8f29cd706..fb8afa7aa 100644 --- a/posthog/capture_v1.py +++ b/posthog/capture_v1.py @@ -166,18 +166,7 @@ def _to_v1_event(msg: dict) -> dict: """ properties = dict(msg.get("properties") or {}) - # Relocate top-level $set/$set_once into properties; v1 has no top-level - # form. On the unusual collision where properties already carries the key, - # the properties value wins. - for key in _RELOCATE_TO_PROPERTIES: - top_val = msg.get(key) - if top_val is None: - continue - existing = properties.get(key) - if isinstance(top_val, dict) and isinstance(existing, dict): - properties[key] = {**top_val, **existing} - elif key not in properties: - properties[key] = top_val + _relocate_v1_properties(msg, properties) for key in _STRIP_FROM_PROPERTIES: properties.pop(key, None) @@ -403,21 +392,9 @@ def _parse_v1_response(res: "requests.Response") -> _V1ParsedResponse: # don't loop forever re-sending against a broken success. return _V1ParsedResponse(status, True, retry_after, malformed=True) - message = "" - try: - payload = res.json() - if isinstance(payload, dict): - message = ( - payload.get("error_description") - or payload.get("error") - or payload.get("detail") - or "" - ) - except (ValueError, AttributeError): - pass - if not message: - message = res.text or f"capture v1 request failed with status {status}" - return _V1ParsedResponse(status, False, retry_after, error_message=message) + return _V1ParsedResponse( + status, False, retry_after, error_message=_v1_error_message(res, status) + ) def _backoff(attempt_index: int, retry_after: Optional[float]) -> None: @@ -536,78 +513,155 @@ def _send_v1_batch( parsed = _parse_v1_response(res) if parsed.is_success: - if parsed.malformed: - raise CaptureV1Error( - parsed.status_code, - "capture v1 returned a success status with an unparseable body", - request_id=request_id, - attempts=attempt, - drops=all_drops, - ) - results = parsed.results or {} - _log_result_summary(request_id, attempt, results) - - retry_events: list[dict] = [] - retry_uuids: list[str] = [] - for event, uid in zip(pending_events, pending_uuids): - directive = results.get(uid) - if directive is None: - # Absent from the map: treated as accepted (matches posthog-rs). - continue - if directive.result == _RESULT_RETRY: - retry_events.append(event) - retry_uuids.append(uid) - elif directive.result == _RESULT_DROP: - # Terminal per-event rejection; keep it so it is surfaced - # even when the rest of the batch succeeds (see below). - all_drops.append((uid, directive.details)) - # ok / warning / unrecognized -> terminal success. - - if not retry_uuids: - # Nothing left to resend. If the server dropped any events, - # surface them via on_error even though the request was a 2xx — - # a success status does not mean every event was delivered. - if all_drops: - raise CaptureV1Error( - parsed.status_code, - f"{len(all_drops)} event(s) dropped by the server", - request_id=request_id, - attempts=attempt, - drops=all_drops, - ) + pending_events, pending_uuids = _apply_v1_results( + parsed, + pending_events, + pending_uuids, + all_drops, + request_id, + attempt, + last_attempt, + ) + if not pending_uuids: return - if last_attempt: - raise CaptureV1Error( - parsed.status_code, - f"{len(retry_uuids)} event(s) still pending retry after {attempt} attempt(s)", - request_id=request_id, - attempts=attempt, - retry_exhausted=retry_uuids, - drops=all_drops, - ) - pending_events, pending_uuids = retry_events, retry_uuids _backoff(attempt_index, parsed.retry_after) continue # Non-2xx. Retryable transient statuses back off; everything else # (400/401/402/413/415/...) is terminal. Any drops collected from a # prior 2xx attempt ride along so on_error still sees them. - v1_error = CaptureV1Error( + last_exc = _v1_batch_failure( + parsed, request_id, attempt, all_drops, last_attempt + ) + _backoff(attempt_index, parsed.retry_after) + + # Unreachable in practice (every branch returns or continues). + if last_exc: + raise last_exc + + +def _relocate_v1_properties(msg: dict, properties: dict) -> None: + # Relocate top-level $set/$set_once into properties; v1 has no top-level + # form. On the unusual collision where properties already carries the key, + # the properties value wins. + for key in _RELOCATE_TO_PROPERTIES: + top_val = msg.get(key) + if top_val is None: + continue + existing = properties.get(key) + if isinstance(top_val, dict) and isinstance(existing, dict): + properties[key] = {**top_val, **existing} + elif key not in properties: + properties[key] = top_val + + +def _v1_error_message(res: "requests.Response", status: int) -> str: + message = "" + try: + payload = res.json() + if isinstance(payload, dict): + message = ( + payload.get("error_description") + or payload.get("error") + or payload.get("detail") + or "" + ) + except (ValueError, AttributeError): + pass + if not message: + message = res.text or f"capture v1 request failed with status {status}" + return message + + +def _apply_v1_results( + parsed: _V1ParsedResponse, + pending_events: list[dict], + pending_uuids: list[str], + all_drops: list[tuple[str, Optional[str]]], + request_id: str, + attempt: int, + last_attempt: bool, +) -> tuple[list[dict], list[str]]: + if parsed.malformed: + raise CaptureV1Error( parsed.status_code, - parsed.error_message, - retry_after=parsed.retry_after, + "capture v1 returned a success status with an unparseable body", request_id=request_id, attempts=attempt, drops=all_drops, ) - if parsed.status_code in _RETRYABLE_STATUSES: - last_exc = v1_error - if last_attempt: - raise v1_error - _backoff(attempt_index, parsed.retry_after) + results = parsed.results or {} + _log_result_summary(request_id, attempt, results) + + retry_events, retry_uuids = _partition_v1_results( + pending_events, pending_uuids, results, all_drops + ) + + if not retry_uuids: + # Nothing left to resend. If the server dropped any events, + # surface them via on_error even though the request was a 2xx — + # a success status does not mean every event was delivered. + if all_drops: + raise CaptureV1Error( + parsed.status_code, + f"{len(all_drops)} event(s) dropped by the server", + request_id=request_id, + attempts=attempt, + drops=all_drops, + ) + return [], [] + if last_attempt: + raise CaptureV1Error( + parsed.status_code, + f"{len(retry_uuids)} event(s) still pending retry after {attempt} attempt(s)", + request_id=request_id, + attempts=attempt, + retry_exhausted=retry_uuids, + drops=all_drops, + ) + return retry_events, retry_uuids + + +def _partition_v1_results( + pending_events: list[dict], + pending_uuids: list[str], + results: dict[str, _V1EventResult], + all_drops: list[tuple[str, Optional[str]]], +) -> tuple[list[dict], list[str]]: + retry_events: list[dict] = [] + retry_uuids: list[str] = [] + for event, uid in zip(pending_events, pending_uuids): + directive = results.get(uid) + if directive is None: + # Absent from the map: treated as accepted (matches posthog-rs). continue - raise v1_error + if directive.result == _RESULT_RETRY: + retry_events.append(event) + retry_uuids.append(uid) + elif directive.result == _RESULT_DROP: + # Terminal per-event rejection; keep it so it is surfaced + # even when the rest of the batch succeeds (see below). + all_drops.append((uid, directive.details)) + # ok / warning / unrecognized -> terminal success. - # Unreachable in practice (every branch returns or continues). - if last_exc: - raise last_exc + return retry_events, retry_uuids + + +def _v1_batch_failure( + parsed: _V1ParsedResponse, + request_id: str, + attempt: int, + all_drops: list[tuple[str, Optional[str]]], + last_attempt: bool, +) -> CaptureV1Error: + error = CaptureV1Error( + parsed.status_code, + parsed.error_message, + retry_after=parsed.retry_after, + request_id=request_id, + attempts=attempt, + drops=all_drops, + ) + if parsed.status_code not in _RETRYABLE_STATUSES or last_attempt: + raise error + return error diff --git a/posthog/client.py b/posthog/client.py index 35380d979..73181a82c 100644 --- a/posthog/client.py +++ b/posthog/client.py @@ -11,6 +11,7 @@ import warnings import weakref from contextvars import ContextVar +from dataclasses import dataclass as _dataclass from datetime import datetime, timedelta, timezone from typing import ( Any, @@ -178,6 +179,32 @@ _atexit_deadline_lock = threading.Lock() +@_dataclass +class _FlagReadResult: + """Per-read result and attribution, including metadata retained on remote failure.""" + + result: Optional[FeatureFlagResult] = None + details: Optional[FeatureFlag] = None + request_id: Optional[str] = None + evaluated_at: Optional[int] = None + error: Optional[str] = None + minimal_flag_called_events: bool = False + # Only a read whose value came from details may attribute a v2 experiment split. + value_from_details: bool = False + + +@_dataclass +class _FlagEvaluationResult: + """Mutable only during one evaluation pass, then handed to its snapshot.""" + + records: Dict[str, _EvaluatedFlagRecord] + request_id: Optional[str] = None + evaluated_at: Optional[int] = None + errors_while_computing: bool = False + quota_limited: bool = False + minimal_flag_called_events: bool = False + + class _LocalEvaluationSnapshot(FlagDefinitionCacheData): flags_by_key: Dict[str, Any] flag_definition_version: int @@ -526,6 +553,12 @@ def _rule_attribution(details: Optional[FeatureFlag]) -> Optional[_RuleAttributi if _v4_field(metadata, "config_version") != 2: return None reason_code = (details.reason.code if details.reason else None) or None + return _rule_attribution_from_metadata(details, metadata, reason_code) + + +def _rule_attribution_from_metadata( + details: FeatureFlag, metadata: FlagMetadata, reason_code: Optional[str] +) -> _RuleAttribution: allowed_rule_types = _RULE_TYPES_BY_REASON_CODE.get(reason_code or "") if _v4_field(details, "failed") is True or allowed_rule_types is None: return _RuleAttribution(reason_code, None, None, None) @@ -708,56 +741,112 @@ def _capture_flag_events( ) if reported_flags is not None: - if attribution is not None: - properties.pop("$feature_flag_reason", None) - properties.update(attribution.flag_called_properties()) - # Record the server's experiment signal when known, so minimization's impact - # can be measured by segmenting on it. - if has_experiment is not None: - properties["$feature_flag_has_experiment"] = has_experiment - - # Minimize iff the server-controlled gate is on AND the flag is known to have - # no linked experiment. Any missing signal (gate absent, has_experiment - # missing) fails safe to the full legacy shape. - should_minimize = minimal_flag_called_events and has_experiment is False - # Only thread the internal allowlist through when minimizing, so the - # full-property path's capture() call signature stays unchanged. - extra_capture_kwargs: Dict[str, Any] = {} - if should_minimize: - extra_capture_kwargs["_property_allowlist"] = ( - _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES - if attribution is None - else _MINIMAL_V2_FLAG_CALLED_EVENT_PROPERTIES - ) + _capture_flag_called_event( + capture, + reported_flags, + feature_flag_reported_key, + distinct_id=distinct_id, + properties=properties, + capture_groups=capture_groups, + disable_geoip=disable_geoip, + has_experiment=has_experiment, + minimal_flag_called_events=minimal_flag_called_events, + attribution=attribution, + flag_identity=flag_identity, + ) + + _capture_experiment_exposure( + capture, + exposures_reported, + distinct_id=distinct_id, + key=key, + groups_key=groups_key, + properties=properties, + capture_groups=capture_groups, + disable_geoip=disable_geoip, + attribution=attribution, + flag_identity=flag_identity, + ) + + +def _capture_flag_called_event( + capture, + reported_flags, + feature_flag_reported_key, + *, + distinct_id, + properties, + capture_groups, + disable_geoip, + has_experiment, + minimal_flag_called_events, + attribution, + flag_identity, +) -> None: + if attribution is not None: + properties.pop("$feature_flag_reason", None) + properties.update(attribution.flag_called_properties()) + # Record the server's experiment signal when known, so minimization's impact + # can be measured by segmenting on it. + if has_experiment is not None: + properties["$feature_flag_has_experiment"] = has_experiment + + # Minimize iff the server-controlled gate is on AND the flag is known to have + # no linked experiment. Any missing signal (gate absent, has_experiment + # missing) fails safe to the full legacy shape. + should_minimize = minimal_flag_called_events and has_experiment is False + # Only thread the internal allowlist through when minimizing, so the + # full-property path's capture() call signature stays unchanged. + extra_capture_kwargs: Dict[str, Any] = {} + if should_minimize: + extra_capture_kwargs["_property_allowlist"] = ( + _MINIMAL_FLAG_CALLED_EVENT_PROPERTIES + if attribution is None + else _MINIMAL_V2_FLAG_CALLED_EVENT_PROPERTIES + ) + + if attribution is None: + # Config version 1 events are sent and deduped as before. + capture( + "$feature_flag_called", + distinct_id=distinct_id, + properties=properties, + groups=capture_groups, + disable_geoip=disable_geoip, + **extra_capture_kwargs, + ) + reported_flags.add(feature_flag_reported_key) + else: + _capture_once( + capture, + reported_flags, + feature_flag_reported_key, + "$feature_flag_called", + distinct_id=distinct_id, + properties=properties, + groups=capture_groups, + disable_geoip=disable_geoip, + _pinned_properties=_PinnedProperties( + _V2_FLAG_CALLED_PINNED, + {**attribution.flag_called_properties(), **flag_identity}, + ), + **extra_capture_kwargs, + ) - if attribution is None: - # Config version 1 events are sent and deduped as before. - capture( - "$feature_flag_called", - distinct_id=distinct_id, - properties=properties, - groups=capture_groups, - disable_geoip=disable_geoip, - **extra_capture_kwargs, - ) - reported_flags.add(feature_flag_reported_key) - else: - _capture_once( - capture, - reported_flags, - feature_flag_reported_key, - "$feature_flag_called", - distinct_id=distinct_id, - properties=properties, - groups=capture_groups, - disable_geoip=disable_geoip, - _pinned_properties=_PinnedProperties( - _V2_FLAG_CALLED_PINNED, - {**attribution.flag_called_properties(), **flag_identity}, - ), - **extra_capture_kwargs, - ) +def _capture_experiment_exposure( + capture, + exposures_reported, + *, + distinct_id, + key, + groups_key, + properties, + capture_groups, + disable_geoip, + attribution, + flag_identity, +) -> None: if attribution is None or attribution.variant_key is None: return @@ -1013,6 +1102,29 @@ def join(self) -> None: ) errors.append(error) + self._stop_consumers(errors) + try: + self.discard_undrainable_queued_work() + except Exception as error: + self.log.exception( + "Failed to discard queued %s lane work during lifecycle cleanup", + self.name, + ) + errors.append(error) + if drain_requested: + try: + self._drain_signal.complete() + except Exception as error: + self.log.exception( + "Failed to complete %s lane drain during lifecycle cleanup", + self.name, + ) + errors.append(error) + + if errors: + raise errors[0] + + def _stop_consumers(self, errors: list[Exception]) -> None: for consumer in self.consumers: try: consumer._pause(drain=True) @@ -1034,26 +1146,6 @@ def join(self) -> None: self.name, ) errors.append(error) - try: - self.discard_undrainable_queued_work() - except Exception as error: - self.log.exception( - "Failed to discard queued %s lane work during lifecycle cleanup", - self.name, - ) - errors.append(error) - if drain_requested: - try: - self._drain_signal.complete() - except Exception as error: - self.log.exception( - "Failed to complete %s lane drain during lifecycle cleanup", - self.name, - ) - errors.append(error) - - if errors: - raise errors[0] def reset_sync_send_state_after_fork(self) -> None: """Replace sync-send state inherited from threads that did not survive fork.""" @@ -1458,75 +1550,25 @@ def __init__( self._logged_untyped_flags_response = False self.capture_trace_context = capture_trace_context - self.capture_exception_code_variables = capture_exception_code_variables - self.code_variables_mask_patterns = ( - code_variables_mask_patterns - if code_variables_mask_patterns is not None - else DEFAULT_CODE_VARIABLES_MASK_PATTERNS + self._configure_exception_properties( + capture_exception_code_variables, + code_variables_mask_patterns, + code_variables_ignore_patterns, + code_variables_mask_url_credentials, + code_variables_detect_secrets, + in_app_modules, + project_root, ) - self.code_variables_ignore_patterns = ( - code_variables_ignore_patterns - if code_variables_ignore_patterns is not None - else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS - ) - self.code_variables_mask_url_credentials = ( - code_variables_mask_url_credentials - if code_variables_mask_url_credentials is not None - else DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS - ) - self.code_variables_detect_secrets = ( - code_variables_detect_secrets - if code_variables_detect_secrets is not None - else DEFAULT_CODE_VARIABLES_DETECT_SECRETS - ) - self.in_app_modules = in_app_modules - if project_root is None: - try: - project_root = os.getcwd() - except Exception: - project_root = None - - self.project_root = project_root - - if personal_api_key is not None and secret_key is None: - warnings.warn( - "`personal_api_key` is deprecated; use `secret_key` instead. " - "`secret_key` accepts a Personal API Key or a Project Secret API Key.", - DeprecationWarning, - stacklevel=2, - ) - elif secret_key is not None and personal_api_key is not None: - self.log.warning( - "[FEATURE FLAGS] Both `secret_key` and `personal_api_key` were " - "provided; using `secret_key` and ignoring `personal_api_key`." - ) - resolved_secret_key = secret_key if secret_key is not None else personal_api_key - self.secret_key = ( - resolved_secret_key.strip() - if isinstance(resolved_secret_key, str) - else resolved_secret_key - ) or None - self.personal_api_key = self.secret_key - _validate_options( - email_request_timeout_seconds, email_request_max_retries, migrated_from + self._configure_credentials( + personal_api_key, + secret_key, + app_host, + email_request_timeout_seconds, + email_request_max_retries, + migrated_from, ) - self.app_host = app_host - self.email_request_timeout_seconds = email_request_timeout_seconds - self.email_request_max_retries = email_request_max_retries - self.migrated_from = migrated_from - if debug: - # Ensures that debug level messages are logged when debug mode is on. - # Otherwise, defaults to WARNING level. See https://docs.python.org/3/howto/logging.html#what-happens-if-no-configuration-is-provided - logging.basicConfig() - self.log.setLevel(logging.DEBUG) - else: - self.log.setLevel(logging.WARNING) - - if not self.api_key: - self.log.error( - "api_key is empty after trimming whitespace; check your project API key" - ) + self._configure_logging(debug) self._set_before_send(before_send) @@ -1592,13 +1634,107 @@ def __init__( self._metrics_autocapture_tried = False self._start_metrics_autocapture_once(metrics) + self._register_fork_handler() + + self._warn_if_duplicate_async_client() + + def _register_fork_handler(self) -> None: if hasattr(os, "register_at_fork"): weak_self = weakref.ref(self) os.register_at_fork( after_in_child=lambda: Client._reinit_after_fork_weak(weak_self) ) - self._warn_if_duplicate_async_client() + def _configure_exception_properties( + self, + capture_exception_code_variables, + code_variables_mask_patterns, + code_variables_ignore_patterns, + code_variables_mask_url_credentials, + code_variables_detect_secrets, + in_app_modules, + project_root, + ) -> None: + self.capture_exception_code_variables = capture_exception_code_variables + self.code_variables_mask_patterns = ( + code_variables_mask_patterns + if code_variables_mask_patterns is not None + else DEFAULT_CODE_VARIABLES_MASK_PATTERNS + ) + self.code_variables_ignore_patterns = ( + code_variables_ignore_patterns + if code_variables_ignore_patterns is not None + else DEFAULT_CODE_VARIABLES_IGNORE_PATTERNS + ) + self.code_variables_mask_url_credentials = ( + code_variables_mask_url_credentials + if code_variables_mask_url_credentials is not None + else DEFAULT_CODE_VARIABLES_MASK_URL_CREDENTIALS + ) + self.code_variables_detect_secrets = ( + code_variables_detect_secrets + if code_variables_detect_secrets is not None + else DEFAULT_CODE_VARIABLES_DETECT_SECRETS + ) + self.in_app_modules = in_app_modules + + if project_root is None: + try: + project_root = os.getcwd() + except Exception: + project_root = None + + self.project_root = project_root + + def _configure_credentials( + self, + personal_api_key, + secret_key, + app_host, + email_request_timeout_seconds, + email_request_max_retries, + migrated_from, + ) -> None: + if personal_api_key is not None and secret_key is None: + warnings.warn( + "`personal_api_key` is deprecated; use `secret_key` instead. " + "`secret_key` accepts a Personal API Key or a Project Secret API Key.", + DeprecationWarning, + stacklevel=3, + ) + elif secret_key is not None and personal_api_key is not None: + self.log.warning( + "[FEATURE FLAGS] Both `secret_key` and `personal_api_key` were " + "provided; using `secret_key` and ignoring `personal_api_key`." + ) + resolved_secret_key = secret_key if secret_key is not None else personal_api_key + self.secret_key = ( + resolved_secret_key.strip() + if isinstance(resolved_secret_key, str) + else resolved_secret_key + ) or None + self.personal_api_key = self.secret_key + _validate_options( + email_request_timeout_seconds, email_request_max_retries, migrated_from + ) + self.app_host = app_host + self.email_request_timeout_seconds = email_request_timeout_seconds + self.email_request_max_retries = email_request_max_retries + self.migrated_from = migrated_from + + def _configure_logging(self, debug) -> None: + if debug: + # Ensures that debug level messages are logged when debug mode is on. + # Otherwise, defaults to WARNING level. See https://docs.python.org/3/howto/logging.html#what-happens-if-no-configuration-is-provided + logging.basicConfig() + self.log.setLevel(logging.DEBUG) + else: + self.log.setLevel(logging.WARNING) + + if not self.api_key: + self.log.error( + "api_key is empty after trimming whitespace; check your project API key" + ) def _start_metrics_autocapture_once(self, metrics_config: Any) -> None: """Starts metrics autocapture when the ``metrics`` option asks for it. A @@ -2122,35 +2258,19 @@ def _get_flags_decision( if self.disabled: return normalize_flags_response({}) - groups = groups or {} - person_properties = person_properties or {} - group_properties = group_properties or {} - - if distinct_id is None: - distinct_id = get_context_distinct_id() - - if device_id is None: - device_id = get_context_device_id() + request_data = self._flags_request_data( + distinct_id, + groups, + person_properties, + group_properties, + disable_geoip, + device_id, + ) - if disable_geoip is None: - disable_geoip = self.disable_geoip - - if not groups: - groups = {} - - request_data: Dict[str, Any] = { - "distinct_id": distinct_id, - "groups": groups, - "person_properties": person_properties, - "group_properties": group_properties, - "geoip_disable": disable_geoip, - "device_id": device_id, - } - - if flag_keys_to_evaluate: - request_data["flag_keys_to_evaluate"] = flag_keys_to_evaluate - if self._sdk_info != _USER_AGENT: - request_data["_user_agent"] = self._sdk_info + if flag_keys_to_evaluate: + request_data["flag_keys_to_evaluate"] = flag_keys_to_evaluate + if self._sdk_info != _USER_AGENT: + request_data["_user_agent"] = self._sdk_info resp_data = flags( self.api_key, @@ -2174,6 +2294,42 @@ def _get_flags_decision( ) return response + def _flags_request_data( + self, + distinct_id, + groups, + person_properties, + group_properties, + disable_geoip, + device_id, + ): + groups = groups or {} + person_properties = person_properties or {} + group_properties = group_properties or {} + + if distinct_id is None: + distinct_id = get_context_distinct_id() + + if device_id is None: + device_id = get_context_device_id() + + if disable_geoip is None: + disable_geoip = self.disable_geoip + + if not groups: + groups = {} + + request_data: Dict[str, Any] = { + "distinct_id": distinct_id, + "groups": groups, + "person_properties": person_properties, + "group_properties": group_properties, + "geoip_disable": disable_geoip, + "device_id": device_id, + } + + return request_data + @no_throw() def capture( self, event: str, **kwargs: Unpack[OptionalCaptureArgs] @@ -2278,13 +2434,7 @@ def _capture( pinned_properties = kwargs.get("_pinned_properties", None) before_send_drops = kwargs.get("_before_send_drops", None) - properties = {**(properties or {}), **system_context()} - - if self.capture_trace_context: - properties = {**_get_current_otel_span_properties(), **properties} - - properties = add_context_tags(properties) - assert properties is not None # Type hint for mypy + properties = self._capture_properties(properties) (distinct_id, personless) = get_identity_state(distinct_id) @@ -2316,78 +2466,9 @@ def _capture( ) extra_properties = flags_snapshot._get_event_properties() else: - feature_variants: Optional[dict[str, Union[bool, str]]] = {} - - # Parse and normalize send_feature_flags parameter - flag_options = self._parse_send_feature_flags(send_feature_flags) - - if flag_options["should_send"]: - warnings.warn( - "`send_feature_flags` is deprecated and will be removed in a future major " - "version. Pass a `flags` snapshot from `posthog.evaluate_flags(...)` instead " - "— it avoids a second `/flags` request per capture and guarantees the event " - "carries the exact flag values your code branched on.", - DeprecationWarning, - stacklevel=2, - ) - try: - if flag_options["only_evaluate_locally"] is True: - # Local evaluation explicitly requested - feature_variants = self.get_all_flags( - distinct_id, - groups=(groups or {}), - person_properties=flag_options["person_properties"], - group_properties=flag_options["group_properties"], - disable_geoip=disable_geoip, - only_evaluate_locally=True, - flag_keys_to_evaluate=flag_options["flag_keys_filter"], - ) - elif flag_options["only_evaluate_locally"] is False: - # Remote evaluation explicitly requested - feature_variants = self.get_feature_variants( - distinct_id, - groups, - person_properties=flag_options["person_properties"], - group_properties=flag_options["group_properties"], - disable_geoip=disable_geoip, - flag_keys_to_evaluate=flag_options["flag_keys_filter"], - ) - elif self.feature_flags: - # Local flags available, prefer local evaluation - feature_variants = self.get_all_flags( - distinct_id, - groups=(groups or {}), - person_properties=flag_options["person_properties"], - group_properties=flag_options["group_properties"], - disable_geoip=disable_geoip, - only_evaluate_locally=True, - flag_keys_to_evaluate=flag_options["flag_keys_filter"], - ) - else: - # Fall back to remote evaluation - feature_variants = self.get_feature_variants( - distinct_id, - groups, - person_properties=flag_options["person_properties"], - group_properties=flag_options["group_properties"], - disable_geoip=disable_geoip, - flag_keys_to_evaluate=flag_options["flag_keys_filter"], - ) - except Exception as e: - self.log.exception( - f"[FEATURE FLAGS] Unable to get feature variants: {e}" - ) - - for feature, variant in (feature_variants or {}).items(): - extra_properties[f"$feature/{feature}"] = variant - - active_feature_flags = [ - key - for (key, value) in (feature_variants or {}).items() - if value is not False - ] - if active_feature_flags: - extra_properties["$active_feature_flags"] = active_feature_flags + extra_properties = self._capture_legacy_flag_properties( + distinct_id, groups, disable_geoip, send_feature_flags + ) if extra_properties: properties = {**extra_properties, **properties} @@ -2402,6 +2483,100 @@ def _capture( before_send_drops=before_send_drops, ) + def _capture_properties(self, properties): + properties = {**(properties or {}), **system_context()} + + if self.capture_trace_context: + properties = {**_get_current_otel_span_properties(), **properties} + + properties = add_context_tags(properties) + assert properties is not None # Type hint for mypy + return properties + + def _capture_legacy_flag_properties( + self, distinct_id, groups, disable_geoip, send_feature_flags + ): + extra_properties: dict[str, Any] = {} + feature_variants: Optional[dict[str, Union[bool, str]]] = {} + + # Parse and normalize send_feature_flags parameter + flag_options = self._parse_send_feature_flags(send_feature_flags) + + if flag_options["should_send"]: + warnings.warn( + "`send_feature_flags` is deprecated and will be removed in a future major " + "version. Pass a `flags` snapshot from `posthog.evaluate_flags(...)` instead " + "— it avoids a second `/flags` request per capture and guarantees the event " + "carries the exact flag values your code branched on.", + DeprecationWarning, + stacklevel=3, + ) + try: + feature_variants = self._evaluate_capture_flags( + distinct_id, groups, disable_geoip, flag_options + ) + except Exception as e: + self.log.exception( + f"[FEATURE FLAGS] Unable to get feature variants: {e}" + ) + + for feature, variant in (feature_variants or {}).items(): + extra_properties[f"$feature/{feature}"] = variant + + active_feature_flags = [ + key + for (key, value) in (feature_variants or {}).items() + if value is not False + ] + if active_feature_flags: + extra_properties["$active_feature_flags"] = active_feature_flags + + return extra_properties + + def _evaluate_capture_flags(self, distinct_id, groups, disable_geoip, flag_options): + if flag_options["only_evaluate_locally"] is True: + # Local evaluation explicitly requested + return self.get_all_flags( + distinct_id, + groups=(groups or {}), + person_properties=flag_options["person_properties"], + group_properties=flag_options["group_properties"], + disable_geoip=disable_geoip, + only_evaluate_locally=True, + flag_keys_to_evaluate=flag_options["flag_keys_filter"], + ) + elif flag_options["only_evaluate_locally"] is False: + # Remote evaluation explicitly requested + return self.get_feature_variants( + distinct_id, + groups, + person_properties=flag_options["person_properties"], + group_properties=flag_options["group_properties"], + disable_geoip=disable_geoip, + flag_keys_to_evaluate=flag_options["flag_keys_filter"], + ) + elif self.feature_flags: + # Local flags available, prefer local evaluation + return self.get_all_flags( + distinct_id, + groups=(groups or {}), + person_properties=flag_options["person_properties"], + group_properties=flag_options["group_properties"], + disable_geoip=disable_geoip, + only_evaluate_locally=True, + flag_keys_to_evaluate=flag_options["flag_keys_filter"], + ) + else: + # Fall back to remote evaluation + return self.get_feature_variants( + distinct_id, + groups, + person_properties=flag_options["person_properties"], + group_properties=flag_options["group_properties"], + disable_geoip=disable_geoip, + flag_keys_to_evaluate=flag_options["flag_keys_filter"], + ) + def _parse_send_feature_flags(self, send_feature_flags) -> SendFeatureFlagsOptions: """ Parse and normalize send_feature_flags parameter into a standard format. @@ -2727,134 +2902,17 @@ def capture_exception( try: properties = properties or {} - # Check if this exception has already been captured - if exception is not None and exception_is_already_captured(exception): - self.log.debug("Exception already captured, skipping") - return None - - if exception is not None: - exc_info = exc_info_from_error(exception) - else: - exc_info = sys.exc_info() - - if exc_info is None or exc_info == (None, None, None): - self.log.warning("No exception information available") + exc_info = self._exception_info(exception) + if exc_info is None: return None - # Format stack trace for cymbal - capture_metadata_input = dict(kwargs).get("_capture_metadata") - capture_metadata = ( - capture_metadata_input - if isinstance(capture_metadata_input, dict) - else {} - ) - mechanism = capture_metadata.get("mechanism") - all_exceptions_with_trace = exceptions_from_error_tuple( - exc_info, mechanism=mechanism if isinstance(mechanism, dict) else None - ) - - # Add in-app property to frames in the exceptions - event = handle_in_app( - { - "exception": { - "values": all_exceptions_with_trace, - }, - }, - in_app_include=self.in_app_modules, - project_root=self.project_root, - ) - all_exceptions_with_trace_and_in_app = event["exception"]["values"] - - reserved_properties = { - "$exception_list", - "$exception_level", - "$exception_source", - "$debug_images", - "$exception_handled", - "$exception_types", - "$exception_values", - "$exception_sources", - "$exception_functions", - "$exception_fingerprint_version", - "$exception_fingerprint_record", - "$exception_issue_id", - "$exception_release", - "$cymbal_errors", - } - reserved_property_overrides = reserved_properties.intersection(properties) - if reserved_property_overrides: - try: - warnings.warn( - "Reserved exception properties passed through " - "`capture_exception(properties=...)` currently override " - "SDK-owned metadata, but this behavior is deprecated and will " - "be removed in the next major version: " - + ", ".join(sorted(reserved_property_overrides)), - DeprecationWarning, - stacklevel=2, - ) - except DeprecationWarning: - # capture_exception must not drop an event when applications - # promote deprecation warnings to errors. - pass - - caller_properties = properties - properties = { - **_get_current_otel_span_properties(), - "$exception_list": all_exceptions_with_trace_and_in_app, - "$exception_level": _normalize_exception_level( - capture_metadata.get("level") - ) - or "error", - } - source = capture_metadata.get("source") - if isinstance(source, str) and source: - properties["$exception_source"] = source - properties.update(caller_properties) - - context_enabled = get_capture_exception_code_variables_context() - context_mask = get_code_variables_mask_patterns_context() - context_ignore = get_code_variables_ignore_patterns_context() - context_mask_url_credentials = ( - get_code_variables_mask_url_credentials_context() + properties, all_exceptions_with_trace_and_in_app = ( + self._exception_properties(exc_info, properties, kwargs) ) - context_detect_secrets = get_code_variables_detect_secrets_context() - enabled = ( - context_enabled - if context_enabled is not None - else self.capture_exception_code_variables - ) - mask_patterns = ( - context_mask - if context_mask is not None - else self.code_variables_mask_patterns - ) - ignore_patterns = ( - context_ignore - if context_ignore is not None - else self.code_variables_ignore_patterns - ) - mask_url_credentials = ( - context_mask_url_credentials - if context_mask_url_credentials is not None - else self.code_variables_mask_url_credentials + self._attach_exception_code_variables( + all_exceptions_with_trace_and_in_app, exc_info ) - detect_secrets = ( - context_detect_secrets - if context_detect_secrets is not None - else self.code_variables_detect_secrets - ) - - if enabled: - try_attach_code_variables_to_frames( - all_exceptions_with_trace_and_in_app, - exc_info, - mask_patterns=mask_patterns, - ignore_patterns=ignore_patterns, - mask_url_credentials=mask_url_credentials, - detect_secrets=detect_secrets, - ) if self.log_captured_exceptions: self.log.exception(exception, extra=kwargs) @@ -2883,6 +2941,140 @@ def capture_exception( self.log.exception(f"Failed to capture exception: {e}") return None + def _exception_info(self, exception): + # Check if this exception has already been captured + if exception is not None and exception_is_already_captured(exception): + self.log.debug("Exception already captured, skipping") + return None + + if exception is not None: + exc_info = exc_info_from_error(exception) + else: + exc_info = sys.exc_info() + + if exc_info is None or exc_info == (None, None, None): + self.log.warning("No exception information available") + return None + + return exc_info + + def _exception_properties(self, exc_info, properties, kwargs): + # Format stack trace for cymbal + capture_metadata_input = dict(kwargs).get("_capture_metadata") + capture_metadata = ( + capture_metadata_input if isinstance(capture_metadata_input, dict) else {} + ) + mechanism = capture_metadata.get("mechanism") + all_exceptions_with_trace = exceptions_from_error_tuple( + exc_info, mechanism=mechanism if isinstance(mechanism, dict) else None + ) + + # Add in-app property to frames in the exceptions + event = handle_in_app( + { + "exception": { + "values": all_exceptions_with_trace, + }, + }, + in_app_include=self.in_app_modules, + project_root=self.project_root, + ) + all_exceptions_with_trace_and_in_app = event["exception"]["values"] + + reserved_properties = { + "$exception_list", + "$exception_level", + "$exception_source", + "$debug_images", + "$exception_handled", + "$exception_types", + "$exception_values", + "$exception_sources", + "$exception_functions", + "$exception_fingerprint_version", + "$exception_fingerprint_record", + "$exception_issue_id", + "$exception_release", + "$cymbal_errors", + } + reserved_property_overrides = reserved_properties.intersection(properties) + if reserved_property_overrides: + try: + warnings.warn( + "Reserved exception properties passed through " + "`capture_exception(properties=...)` currently override " + "SDK-owned metadata, but this behavior is deprecated and will " + "be removed in the next major version: " + + ", ".join(sorted(reserved_property_overrides)), + DeprecationWarning, + stacklevel=3, + ) + except DeprecationWarning: + # capture_exception must not drop an event when applications + # promote deprecation warnings to errors. + pass + + caller_properties = properties + properties = { + **_get_current_otel_span_properties(), + "$exception_list": all_exceptions_with_trace_and_in_app, + "$exception_level": _normalize_exception_level( + capture_metadata.get("level") + ) + or "error", + } + source = capture_metadata.get("source") + if isinstance(source, str) and source: + properties["$exception_source"] = source + properties.update(caller_properties) + + return properties, all_exceptions_with_trace_and_in_app + + def _attach_exception_code_variables( + self, all_exceptions_with_trace_and_in_app, exc_info + ): + context_enabled = get_capture_exception_code_variables_context() + context_mask = get_code_variables_mask_patterns_context() + context_ignore = get_code_variables_ignore_patterns_context() + context_mask_url_credentials = get_code_variables_mask_url_credentials_context() + context_detect_secrets = get_code_variables_detect_secrets_context() + + enabled = ( + context_enabled + if context_enabled is not None + else self.capture_exception_code_variables + ) + mask_patterns = ( + context_mask + if context_mask is not None + else self.code_variables_mask_patterns + ) + ignore_patterns = ( + context_ignore + if context_ignore is not None + else self.code_variables_ignore_patterns + ) + mask_url_credentials = ( + context_mask_url_credentials + if context_mask_url_credentials is not None + else self.code_variables_mask_url_credentials + ) + detect_secrets = ( + context_detect_secrets + if context_detect_secrets is not None + else self.code_variables_detect_secrets + ) + + if enabled: + try_attach_code_variables_to_frames( + all_exceptions_with_trace_and_in_app, + exc_info, + mask_patterns=mask_patterns, + ignore_patterns=ignore_patterns, + mask_url_credentials=mask_url_credentials, + detect_secrets=detect_secrets, + ) + @classmethod def _reinit_client_registry_after_fork(cls): """Replace the inherited registry lock once in each forked child.""" @@ -2982,6 +3174,9 @@ def _reinit_after_fork(self): reset_sessions() # Start child threads only after replacing every lock they can touch. + self._restart_flag_poller_after_fork(terminal_requested) + + def _restart_flag_poller_after_fork(self, terminal_requested): if terminal_requested: self.poller = None elif self.enable_local_evaluation: @@ -3026,6 +3221,30 @@ def _enqueue( if self.disabled: return None + self._normalize_event_timestamp(msg) + + self._normalize_event_uuid(msg) + + self._enrich_event_properties( + msg, disable_geoip, property_allowlist, pinned_properties + ) + + msg = clean(msg) + + msg = self._run_before_send(msg, before_send_drops) + if msg is None: + return None + + # Re-normalized after before_send, which may have replaced or removed + # msg["uuid"], so the returned uuid always matches the wire event. + self._normalize_event_uuid(msg) + sent_uuid = msg["uuid"] + + self.log.debug("queueing: %s", msg) + + return self._submit_event(msg, lane, sent_uuid) + + def _normalize_event_timestamp(self, msg): timestamp = msg["timestamp"] if timestamp is None: timestamp = datetime.now(tz=timezone.utc) @@ -3039,8 +3258,9 @@ def _enqueue( ) msg["timestamp"] = datetime.now(tz=timezone.utc).isoformat() - self._normalize_event_uuid(msg) - + def _enrich_event_properties( + self, msg, disable_geoip, property_allowlist, pinned_properties + ): if not msg.get("properties"): msg["properties"] = {} msg["properties"]["$lib"] = self._library_id @@ -3065,6 +3285,11 @@ def _enqueue( if self.is_server: msg["properties"]["$is_server"] = True + self._filter_event_properties(msg, property_allowlist, pinned_properties) + + msg["distinct_id"] = stringify_id(msg.get("distinct_id", None)) + + def _filter_event_properties(self, msg, property_allowlist, pinned_properties): # Applied after every enrichment step (system context, context tags, super # properties, $lib/$lib_version) so the final event shape is exactly the # allowlist regardless of where a property came from. @@ -3078,10 +3303,7 @@ def _enqueue( k: v for k, v in msg["properties"].items() if k in property_allowlist } - msg["distinct_id"] = stringify_id(msg.get("distinct_id", None)) - - msg = clean(msg) - + def _run_before_send(self, msg, before_send_drops): if self.before_send: try: modified_msg = self.before_send(msg) @@ -3097,13 +3319,9 @@ def _enqueue( self.log.exception(f"Error in before_send callback: {e}") return None - # Re-normalized after before_send, which may have replaced or removed - # msg["uuid"], so the returned uuid always matches the wire event. - self._normalize_event_uuid(msg) - sent_uuid = msg["uuid"] - - self.log.debug("queueing: %s", msg) + return msg + def _submit_event(self, msg, lane, sent_uuid): # if send is False, return msg as if it was successfully queued, unless # shutdown has already closed this lane's admission. if not self.send: @@ -3119,35 +3337,7 @@ def _enqueue( if self.sync_mode: self.log.debug("enqueued with blocking %s.", msg["event"]) - def send_sync() -> None: - # Sync mode bypasses the lane's queue but keeps its wire config: - # the AI lane is pinned to v0, so its events post to the AI - # endpoint regardless of `capture_mode`. - if lane.capture_mode == CaptureMode.V1: - _send_v1_batch( - self.api_key, - self.host, - [msg], - compression=self.capture_compression, - timeout=self.timeout, - max_retries=self.max_retries, - historical_migration=self.historical_migration, - sdk_info=self._sdk_info, - ) - return - - batch_post( - self.api_key, - self.host, - gzip=self.gzip, - timeout=self.timeout, - batch=[msg], - historical_migration=self.historical_migration, - path=lane.endpoint, - **self._request_identity_kwargs(), - ) - - if lane.run_sync_if_open(send_sync): + if lane.run_sync_if_open(lambda: self._send_sync(msg, lane)): return sent_uuid self.log.warning( "%s lane received event %s after shutdown, dropping it", @@ -3182,6 +3372,34 @@ def send_sync() -> None: ) return None + def _send_sync(self, msg, lane) -> None: + # Sync mode bypasses the lane's queue but keeps its wire config: + # the AI lane is pinned to v0, so its events post to the AI + # endpoint regardless of `capture_mode`. + if lane.capture_mode == CaptureMode.V1: + _send_v1_batch( + self.api_key, + self.host, + [msg], + compression=self.capture_compression, + timeout=self.timeout, + max_retries=self.max_retries, + historical_migration=self.historical_migration, + sdk_info=self._sdk_info, + ) + return + + batch_post( + self.api_key, + self.host, + gzip=self.gzip, + timeout=self.timeout, + batch=[msg], + historical_migration=self.historical_migration, + path=lane.endpoint, + **self._request_identity_kwargs(), + ) + @property def metrics(self) -> PostHogMetrics: """ @@ -3230,30 +3448,33 @@ def _traces_pipeline(self) -> Optional[PostHogTraces]: # Re-checked under the lock, which shutdown takes before it # reads the pipeline, so neither can miss the other. if self._traces is None and not self._shutdown_requested: - try: - config = resolve_traces_config( - self._traces_config, host_resource_attributes() - ) - drops = DropLog(config.flush_interval) - self._traces = PostHogTraces( - self, - config, - self._tracing_context, - self._active_span_var, - SpanExporter(self, config, drops), - drops, - ) - # Sync mode has no exit hook for events; spans still - # queue for the timer, so they need their own. - if self.sync_mode and self.send: - atexit.register(self._atexit_spans) - except Exception: - # Off rather than defaults: defaults would drop a - # before_span_send hook and export unscrubbed spans. - self.log.exception("Error initializing traces; tracing is off") - self._traces_config = False + self._initialize_traces_pipeline() return self._traces + def _initialize_traces_pipeline(self) -> None: + try: + config = resolve_traces_config( + self._traces_config, host_resource_attributes() + ) + drops = DropLog(config.flush_interval) + self._traces = PostHogTraces( + self, + config, + self._tracing_context, + self._active_span_var, + SpanExporter(self, config, drops), + drops, + ) + # Sync mode has no exit hook for events; spans still + # queue for the timer, so they need their own. + if self.sync_mode and self.send: + atexit.register(self._atexit_spans) + except Exception: + # Off rather than defaults: defaults would drop a + # before_span_send hook and export unscrubbed spans. + self.log.exception("Error initializing traces; tracing is off") + self._traces_config = False + def _tracing_context(self) -> Dict[str, Optional[str]]: # The request context's identity, which the Django middleware fills # from the X-POSTHOG-* headers: the span's person/session join keys. @@ -3393,14 +3614,7 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: if include_metrics_autocapture else None ) - deadline: Optional[float] = None - if timeout_seconds is None: - for lane in self._lanes: - lane.flush(None) - else: - deadline = time.monotonic() + timeout_seconds - for lane in self._lanes: - lane.flush(max(0.0, deadline - time.monotonic())) + deadline = self._flush_lanes(timeout_seconds) if metrics_flush is not None: metrics_flush( None if deadline is None else max(0.0, deadline - time.monotonic()) @@ -3413,6 +3627,17 @@ def flush(self, timeout_seconds: Optional[float] = 10) -> None: self.log.exception("error flushing queue: %s", e) return + def _flush_lanes(self, timeout_seconds): + deadline: Optional[float] = None + if timeout_seconds is None: + for lane in self._lanes: + lane.flush(None) + else: + deadline = time.monotonic() + timeout_seconds + for lane in self._lanes: + lane.flush(max(0.0, deadline - time.monotonic())) + return deadline + def _start_metrics_autocapture_flush( self, ) -> Optional[Callable[[Optional[float]], None]]: @@ -3607,6 +3832,9 @@ def _join_once( # means every worker cleanup step was attempted once. self._workers_joined = True + self._close_lifecycle_resources(errors) + + def _close_lifecycle_resources(self, errors: list[Exception]) -> None: if not self._join_cleanup_complete: if self.poller: self._run_lifecycle_cleanup( @@ -3705,15 +3933,7 @@ def _shutdown_once(self, errors: list[Exception]) -> None: def _run_lifecycle(self, require_shutdown: bool = False) -> None: while True: with self._lifecycle_condition: - if require_shutdown and self._shutdown_complete: - if self.debug and self._lifecycle_cleanup_failed: - raise RuntimeError("client lifecycle cleanup failed") - return - if not require_shutdown and ( - self._join_cleanup_complete or self._shutdown_complete - ): - if self.debug and self._lifecycle_cleanup_failed: - raise RuntimeError("client lifecycle cleanup failed") + if self._lifecycle_has_completed(require_shutdown): return if self._lifecycle_owner is not None: if self._is_lifecycle_callback_thread() or ( @@ -3725,47 +3945,66 @@ def _run_lifecycle(self, require_shutdown: bool = False) -> None: self._lifecycle_owner = threading.current_thread() try: - errors: list[Exception] = [] - while True: - with self._lifecycle_lock: - run_shutdown = self._shutdown_requested - - if run_shutdown: - self._shutdown_once(errors) - else: - self._join_once(errors) - - with self._lifecycle_condition: - if ( - self._shutdown_requested - and not self._shutdown_complete - and not run_shutdown - ): - continue - if errors: - self._lifecycle_cleanup_failed = True - raise errors[0] - self._lifecycle_owner = None - self._lifecycle_condition.notify_all() - return + self._finish_lifecycle() + return except BaseException: with self._lifecycle_condition: self._lifecycle_owner = None self._lifecycle_condition.notify_all() raise - @no_throw() - def _atexit(self) -> None: - """Make a bounded delivery attempt, then stop daemon workers.""" - with self._lifecycle_condition: - # A daemon lifecycle worker already owns cleanup. Do not wait for it - # at interpreter exit; the process must remain free to terminate. - if self._lifecycle_owner is not None: - return - self._lifecycle_owner = threading.current_thread() + def _lifecycle_has_completed(self, require_shutdown: bool) -> bool: + if require_shutdown and self._shutdown_complete: + self._raise_lifecycle_failure() + return True + if not require_shutdown and ( + self._join_cleanup_complete or self._shutdown_complete + ): + self._raise_lifecycle_failure() + return True + return False - try: - try: + def _raise_lifecycle_failure(self) -> None: + if self.debug and self._lifecycle_cleanup_failed: + raise RuntimeError("client lifecycle cleanup failed") + + def _finish_lifecycle(self) -> None: + errors: list[Exception] = [] + while True: + with self._lifecycle_lock: + run_shutdown = self._shutdown_requested + + if run_shutdown: + self._shutdown_once(errors) + else: + self._join_once(errors) + + with self._lifecycle_condition: + if ( + self._shutdown_requested + and not self._shutdown_complete + and not run_shutdown + ): + continue + if errors: + self._lifecycle_cleanup_failed = True + raise errors[0] + self._lifecycle_owner = None + self._lifecycle_condition.notify_all() + return + + @no_throw() + def _atexit(self) -> None: + """Make a bounded delivery attempt, then stop daemon workers.""" + with self._lifecycle_condition: + # A daemon lifecycle worker already owns cleanup. Do not wait for it + # at interpreter exit; the process must remain free to terminate. + if self._lifecycle_owner is not None: + return + self._lifecycle_owner = threading.current_thread() + + try: + try: for lane in self._lanes: lane.close() @@ -4065,18 +4304,7 @@ def _local_evaluation_snapshot(self) -> _LocalEvaluationSnapshot: } def _load_feature_flags(self): - should_fetch = True - if self._flag_definition_cache_provider: - try: - should_fetch = self._resolve_flag_definition_cache_provider_result( - self._flag_definition_cache_provider.should_fetch_flag_definitions() - ) - except Exception as e: - self.log.error( - f"[FEATURE FLAGS] Cache provider should_fetch error: {e}" - ) - # Fail-safe: fetch from API if cache provider errors - should_fetch = True + should_fetch = self._should_fetch_flag_definitions() # If not fetching, try to get from cache if not should_fetch and self._flag_definition_cache_provider: @@ -4110,6 +4338,22 @@ def _load_feature_flags(self): if should_fetch: self._fetch_feature_flags_from_api() + def _should_fetch_flag_definitions(self): + should_fetch = True + if self._flag_definition_cache_provider: + try: + should_fetch = self._resolve_flag_definition_cache_provider_result( + self._flag_definition_cache_provider.should_fetch_flag_definitions() + ) + except Exception as e: + self.log.error( + f"[FEATURE FLAGS] Cache provider should_fetch error: {e}" + ) + # Fail-safe: fetch from API if cache provider errors + should_fetch = True + + return should_fetch + def _fetch_feature_flags_from_api(self): """Fetch feature flags from the PostHog API.""" personal_api_key = self.personal_api_key @@ -4127,171 +4371,194 @@ def _fetch_feature_flags_from_api(self): cache_data_to_store: Optional[FlagDefinitionCacheData] = None old_fingerprint: Optional[str] = None try: - request_get = _get_with_identity if self._request_identity_kwargs() else get - response = request_get( - personal_api_key, - f"/flags/definitions?token={self.api_key}&send_cohorts", - self.host, - timeout=10, - etag=request_etag, - **self._request_identity_kwargs(), + response, fingerprint = self._request_flag_definitions( + personal_api_key, request_etag ) - - # Canonical serialization can be expensive; do it before publication. - fingerprint = ( - self._hash_flag_definitions(response.data) - if response.data is not None and not response.not_modified - else None + published, old_fingerprint, cache_data_to_store = ( + self._publish_flag_definitions_response( + response, fingerprint, fetch_generation, request_etag + ) + ) + if not published: + return + self._warn_unresolvable_flags() + self._store_flag_definitions_if_current( + cache_data_to_store, fetch_generation ) + except APIError as e: with self._flag_definition_publication_lock: - if fetch_generation <= self._flag_definition_published_generation: - self.log.debug( - "[FEATURE FLAGS] Ignoring stale flag definition response" - ) - self._last_feature_flag_poll = datetime.now(tz=timezone.utc) - return + if ( + fetch_generation > self._flag_definition_published_generation + and e.status in (401, 402) + ): + old_fingerprint = self._flag_definition_fingerprint + debug_error = self._handle_flag_definition_api_error( + e, fetch_generation + ) + if debug_error is not None: + raise debug_error + except Exception as e: + self.log.warning( + "[FEATURE FLAGS] Fetching feature flags failed with following error. We will retry in %s seconds." + % self.poll_interval + ) + self.log.warning(e) + finally: + # Legacy readers ignore snapshot metadata. Cleanup is best effort and + # must not hold the publication lock, including when debug raises. + if isinstance(self.flag_cache, RedisFlagCache): + self.flag_cache._invalidate_snapshot(old_fingerprint) - # A 304 is valid only for the ETag used by this request. Another - # overlapping response may already have installed newer definitions. - if response.not_modified: - if self._flags_etag != request_etag: - self.log.debug( - "[FEATURE FLAGS] Ignoring stale 304 flag definition response" - ) - self._last_feature_flag_poll = datetime.now(tz=timezone.utc) - return + self._last_feature_flag_poll = datetime.now(tz=timezone.utc) - self._flags_etag = response.etag - self._flag_definition_published_generation = fetch_generation - self._flag_definition_cache_generation = fetch_generation - self.log.debug( - "[FEATURE FLAGS] Flags not modified (304), using cached data" - ) - self._last_feature_flag_poll = datetime.now(tz=timezone.utc) - return + def _request_flag_definitions(self, personal_api_key, request_etag): + request_get = _get_with_identity if self._request_identity_kwargs() else get + response = request_get( + personal_api_key, + f"/flags/definitions?token={self.api_key}&send_cohorts", + self.host, + timeout=10, + etag=request_etag, + **self._request_identity_kwargs(), + ) - if response.data is None: - self.log.error( - "[FEATURE FLAGS] Unexpected empty response data in non-304 response" - ) - return + # Canonical serialization can be expensive; do it before publication. + fingerprint = ( + self._hash_flag_definitions(response.data) + if response.data is not None and not response.not_modified + else None + ) + + return response, fingerprint - old_fingerprint = self._update_flag_state( - response.data, _fingerprint=fingerprint + def _publish_flag_definitions_response( + self, response, fingerprint, fetch_generation, request_etag + ): + cache_data_to_store = None + with self._flag_definition_publication_lock: + if fetch_generation <= self._flag_definition_published_generation: + self.log.debug( + "[FEATURE FLAGS] Ignoring stale flag definition response" ) + self._last_feature_flag_poll = datetime.now(tz=timezone.utc) + return False, None, None - if self._flag_definition_cache_provider: - cache_data_to_store = { - "flags": self.feature_flags or [], - "group_type_mapping": self.group_type_mapping or {}, - "cohorts": self.cohorts or {}, - "minimal_flag_called_events": self._minimal_flag_called_events, - "property_matching_version": self._property_matching_version, - } + # A 304 is valid only for the ETag used by this request. Another + # overlapping response may already have installed newer definitions. + if response.not_modified: + if self._flags_etag != request_etag: + self.log.debug( + "[FEATURE FLAGS] Ignoring stale 304 flag definition response" + ) + self._last_feature_flag_poll = datetime.now(tz=timezone.utc) + return False, None, None - # Publish the ETag only after its matching flag state is installed. self._flags_etag = response.etag self._flag_definition_published_generation = fetch_generation self._flag_definition_cache_generation = fetch_generation + self.log.debug( + "[FEATURE FLAGS] Flags not modified (304), using cached data" + ) + self._last_feature_flag_poll = datetime.now(tz=timezone.utc) + return False, None, None - self._warn_unresolvable_flags() + if response.data is None: + self.log.error( + "[FEATURE FLAGS] Unexpected empty response data in non-304 response" + ) + return False, None, None - if cache_data_to_store and self._flag_definition_cache_provider: - # Keep provider I/O out of the publication lock. The separate lock - # preserves cache write order without delaying newer API fetches or - # in-memory publication. - with self._flag_definition_cache_write_lock: - with self._flag_definition_publication_lock: - should_store = ( - fetch_generation == self._flag_definition_cache_generation - ) - if should_store: - try: - self._resolve_flag_definition_cache_provider_result( - self._flag_definition_cache_provider.on_flag_definitions_received( - cache_data_to_store - ) - ) - except Exception as e: - self.log.error( - f"[FEATURE FLAGS] Cache provider store error: {e}" - ) - # Flags are already in memory, so continue normally + old_fingerprint = self._update_flag_state( + response.data, _fingerprint=fingerprint + ) - except APIError as e: - with self._flag_definition_publication_lock: - if fetch_generation <= self._flag_definition_published_generation: - self.log.debug("[FEATURE FLAGS] Ignoring stale API error response") - elif e.status == 401: - detail = ( - f"Error loading feature flags: {e.message}. " - "Please verify both your project_api_key and secret_key. " - "More information: https://posthog.com/docs/api/overview" - ) - self.log.error("[FEATURE FLAGS] %s", detail) - self.feature_flags = [] - self.group_type_mapping = {} - self.cohorts = {} - self._property_matching_version = 1 - old_fingerprint = self._flag_definition_fingerprint - self._flag_definition_fingerprint = "" - self._flags_etag = None - self._flag_definition_published_generation = fetch_generation - self._flag_definition_cache_generation = fetch_generation - - if self.flag_cache: - self.flag_definition_version += 1 - self.flag_cache._advance_generation( - self.flag_definition_version, - self._flag_definition_fingerprint, - ) + if self._flag_definition_cache_provider: + cache_data_to_store = { + "flags": self.feature_flags or [], + "group_type_mapping": self.group_type_mapping or {}, + "cohorts": self.cohorts or {}, + "minimal_flag_called_events": self._minimal_flag_called_events, + "property_matching_version": self._property_matching_version, + } - if self.debug: - raise APIError(status=401, message=detail) - elif e.status == 402: - self.log.warning( - "[FEATURE FLAGS] PostHog feature flags quota limited, resetting feature flag data. Learn more about billing limits at https://posthog.com/docs/billing/limits-alerts" + # Publish the ETag only after its matching flag state is installed. + self._flags_etag = response.etag + self._flag_definition_published_generation = fetch_generation + self._flag_definition_cache_generation = fetch_generation + + return True, old_fingerprint, cache_data_to_store + + def _store_flag_definitions_if_current(self, cache_data_to_store, fetch_generation): + if cache_data_to_store and self._flag_definition_cache_provider: + # Keep provider I/O out of the publication lock. The separate lock + # preserves cache write order without delaying newer API fetches or + # in-memory publication. + with self._flag_definition_cache_write_lock: + with self._flag_definition_publication_lock: + should_store = ( + fetch_generation == self._flag_definition_cache_generation ) - # Reset all feature flag data when quota limited - self.feature_flags = [] - self.group_type_mapping = {} - self.cohorts = {} - self._property_matching_version = 1 - old_fingerprint = self._flag_definition_fingerprint - self._flag_definition_fingerprint = "" - self._flags_etag = None - self._flag_definition_published_generation = fetch_generation - self._flag_definition_cache_generation = fetch_generation - - # Invalidate results without waiting for external cache I/O. - if self.flag_cache: - self.flag_definition_version += 1 - self.flag_cache._advance_generation( - self.flag_definition_version, - self._flag_definition_fingerprint, + if should_store: + try: + self._resolve_flag_definition_cache_provider_result( + self._flag_definition_cache_provider.on_flag_definitions_received( + cache_data_to_store + ) ) - - if self.debug: - raise APIError( - status=402, - message="PostHog feature flags quota limited", + except Exception as e: + self.log.error( + f"[FEATURE FLAGS] Cache provider store error: {e}" ) - else: - self.log.error(f"[FEATURE FLAGS] Error loading feature flags: {e}") - except Exception as e: + # Flags are already in memory, so continue normally + + def _handle_flag_definition_api_error(self, e, fetch_generation): + debug_error = None + if fetch_generation <= self._flag_definition_published_generation: + self.log.debug("[FEATURE FLAGS] Ignoring stale API error response") + elif e.status == 401: + detail = ( + f"Error loading feature flags: {e.message}. " + "Please verify both your project_api_key and secret_key. " + "More information: https://posthog.com/docs/api/overview" + ) + self.log.error("[FEATURE FLAGS] %s", detail) + self._clear_flag_definitions_after_api_error(fetch_generation) + + if self.debug: + debug_error = APIError(status=401, message=detail) + elif e.status == 402: self.log.warning( - "[FEATURE FLAGS] Fetching feature flags failed with following error. We will retry in %s seconds." - % self.poll_interval + "[FEATURE FLAGS] PostHog feature flags quota limited, resetting feature flag data. Learn more about billing limits at https://posthog.com/docs/billing/limits-alerts" ) - self.log.warning(e) - finally: - # Legacy readers ignore snapshot metadata. Cleanup is best effort and - # must not hold the publication lock, including when debug raises. - if isinstance(self.flag_cache, RedisFlagCache): - self.flag_cache._invalidate_snapshot(old_fingerprint) + # Reset all feature flag data when quota limited + self._clear_flag_definitions_after_api_error(fetch_generation) - self._last_feature_flag_poll = datetime.now(tz=timezone.utc) + if self.debug: + debug_error = APIError( + status=402, + message="PostHog feature flags quota limited", + ) + else: + self.log.error(f"[FEATURE FLAGS] Error loading feature flags: {e}") + return debug_error + + def _clear_flag_definitions_after_api_error(self, fetch_generation): + self.feature_flags = [] + self.group_type_mapping = {} + self.cohorts = {} + self._property_matching_version = 1 + self._flag_definition_fingerprint = "" + self._flags_etag = None + self._flag_definition_published_generation = fetch_generation + self._flag_definition_cache_generation = fetch_generation + + if self.flag_cache: + self.flag_definition_version += 1 + self.flag_cache._advance_generation( + self.flag_definition_version, + self._flag_definition_fingerprint, + ) def load_feature_flags(self): """ @@ -4349,8 +4616,6 @@ def _compute_flag_locally( if _definition_snapshot is not None else self._local_evaluation_snapshot() ) - flags_by_key = snapshot["flags_by_key"] - property_matching_version = snapshot.get("property_matching_version", 1) groups = groups or {} person_properties = person_properties or {} group_properties = group_properties or {} @@ -4369,87 +4634,136 @@ def _compute_flag_locally( flag_filters = feature_flag.get("filters") or {} aggregation_group_type_index = flag_filters.get("aggregation_group_type_index") - group_type_mapping = snapshot["group_type_mapping"] if aggregation_group_type_index is not None: - group_name = group_type_mapping.get(str(aggregation_group_type_index)) + return self._compute_group_flag_locally( + feature_flag, + groups, + group_properties, + warn_on_unknown_groups, + device_id, + snapshot, + evaluation_cache, + aggregation_group_type_index, + ) + else: + return self._compute_person_flag_locally( + feature_flag, + distinct_id, + person_properties, + device_id, + snapshot, + evaluation_cache, + groups, + group_properties, + ) + + def _compute_group_flag_locally( + self, + feature_flag, + groups, + group_properties, + warn_on_unknown_groups, + device_id, + snapshot, + evaluation_cache, + aggregation_group_type_index, + ) -> FlagValue: + group_type_mapping = snapshot["group_type_mapping"] + flags_by_key = snapshot["flags_by_key"] + property_matching_version = snapshot.get("property_matching_version", 1) + group_name = group_type_mapping.get(str(aggregation_group_type_index)) - if not group_name: + if not group_name: + self.log.warning( + f"[FEATURE FLAGS] Unknown group type index {aggregation_group_type_index} for feature flag {feature_flag['key']}" + ) + # failover to `/flags` + raise InconclusiveMatchError( + "Flag has unknown group type index", + reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, + ) + + if group_name not in groups: + # Group flags are never enabled in `groups` aren't passed in + # don't failover to `/flags`, since response will be the same + if warn_on_unknown_groups: self.log.warning( - f"[FEATURE FLAGS] Unknown group type index {aggregation_group_type_index} for feature flag {feature_flag['key']}" + f"[FEATURE FLAGS] Can't compute group feature flag: {feature_flag['key']} without group names passed in" ) - # failover to `/flags` - raise InconclusiveMatchError( - "Flag has unknown group type index", - reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, + else: + self.log.debug( + f"[FEATURE FLAGS] Can't compute group feature flag: {feature_flag['key']} without group names passed in" ) + return False - if group_name not in groups: - # Group flags are never enabled in `groups` aren't passed in - # don't failover to `/flags`, since response will be the same - if warn_on_unknown_groups: - self.log.warning( - f"[FEATURE FLAGS] Can't compute group feature flag: {feature_flag['key']} without group names passed in" - ) - else: - self.log.debug( - f"[FEATURE FLAGS] Can't compute group feature flag: {feature_flag['key']} without group names passed in" - ) - return False + if group_name not in group_properties: + raise InconclusiveMatchError( + f"Flag has no group properties for group '{group_name}'", + reason=UnresolvedFlagReason.MISSING_CONTEXT, + ) + focused_group_properties = group_properties[group_name] + group_key = groups[group_name] + return match_feature_flag_properties( + feature_flag, + group_key, + focused_group_properties, + cohort_properties=snapshot["cohorts"], + flags_by_key=flags_by_key, + property_matching_version=property_matching_version, + evaluation_cache=evaluation_cache, + device_id=device_id, + bucketing_value=group_key, + group_type_mapping=group_type_mapping, + groups=groups, + group_properties=group_properties, + ) - if group_name not in group_properties: - raise InconclusiveMatchError( - f"Flag has no group properties for group '{group_name}'", - reason=UnresolvedFlagReason.MISSING_CONTEXT, - ) - focused_group_properties = group_properties[group_name] - group_key = groups[group_name] - return match_feature_flag_properties( - feature_flag, - group_key, - focused_group_properties, - cohort_properties=snapshot["cohorts"], - flags_by_key=flags_by_key, - property_matching_version=property_matching_version, - evaluation_cache=evaluation_cache, - device_id=device_id, - bucketing_value=group_key, - group_type_mapping=group_type_mapping, - groups=groups, - group_properties=group_properties, + def _compute_person_flag_locally( + self, + feature_flag, + distinct_id, + person_properties, + device_id, + snapshot, + evaluation_cache, + groups, + group_properties, + ) -> FlagValue: + group_type_mapping = snapshot["group_type_mapping"] + flags_by_key = snapshot["flags_by_key"] + property_matching_version = snapshot.get("property_matching_version", 1) + try: + bucketing_value = resolve_bucketing_value( + feature_flag, distinct_id, device_id ) - else: - try: - bucketing_value = resolve_bucketing_value( - feature_flag, distinct_id, device_id - ) - except InconclusiveMatchError as e: - # No condition is checked without a bucketing value, so a cause in - # the definition also applies and outranks missing_context. - e.reason = ( - _get_definition_unresolved_reason( - feature_flag, - flags_by_key, - snapshot["cohorts"], - group_type_mapping, - ) - or e.reason + except InconclusiveMatchError as e: + # No condition is checked without a bucketing value, so a cause in + # the definition also applies and outranks missing_context. + e.reason = ( + _get_definition_unresolved_reason( + feature_flag, + flags_by_key, + snapshot["cohorts"], + group_type_mapping, ) - raise - return match_feature_flag_properties( - feature_flag, - distinct_id, - person_properties, - cohort_properties=snapshot["cohorts"], - flags_by_key=flags_by_key, - property_matching_version=property_matching_version, - evaluation_cache=evaluation_cache, - device_id=device_id, - bucketing_value=bucketing_value, - group_type_mapping=group_type_mapping, - groups=groups, - group_properties=group_properties, + or e.reason ) + raise + return match_feature_flag_properties( + feature_flag, + distinct_id, + person_properties, + cohort_properties=snapshot["cohorts"], + flags_by_key=flags_by_key, + property_matching_version=property_matching_version, + evaluation_cache=evaluation_cache, + device_id=device_id, + bucketing_value=bucketing_value, + group_type_mapping=group_type_mapping, + groups=groups, + group_properties=group_properties, + ) def feature_enabled( self, @@ -4555,20 +4869,33 @@ def _get_feature_flag_result( ) ) # Ensure non-None values for type checking - groups = groups or {} - person_properties = person_properties or {} - group_properties = group_properties or {} - - flag_result = None - flag_details = None - # Whether flag_details supplied the value this read reports. Only then may the - # read report config version 2 attribution or send $experiment_exposure. - value_from_details = False - request_id = None - evaluated_at = None - feature_flag_error: Optional[str] = None - remote_minimal_flag_called_events = False + return self._get_feature_flag_result_with_context( + key, + distinct_id, + override_match_value, + groups or {}, + person_properties or {}, + group_properties or {}, + only_evaluate_locally, + send_feature_flag_events, + disable_geoip, + device_id, + ) + def _get_feature_flag_result_with_context( + self, + key, + distinct_id, + override_match_value, + groups, + person_properties, + group_properties, + only_evaluate_locally, + send_feature_flag_events, + disable_geoip, + device_id, + ): + read = _FlagReadResult() # Resolve device_id from context if not provided if device_id is None: device_id = get_context_device_id() @@ -4584,40 +4911,16 @@ def _get_feature_flag_result( group_properties, device_id, ) - local_definition_version = snapshot["flag_definition_version"] flag_was_locally_evaluated = flag_value is not None if flag_value is not None: - lookup_match_value = ( - override_match_value if override_match_value is not None else flag_value - ) - payload = ( - self._compute_payload_locally( - key, lookup_match_value, _definition_snapshot=snapshot - ) - if lookup_match_value is not None - else None - ) - flag_result = FeatureFlagResult.from_value_and_payload( - key, lookup_match_value, payload + read.result = self._resolve_local_flag_result( + key, + distinct_id, + flag_value, + override_match_value, + snapshot, ) - - # Cache the local evaluation, not a payload lookup override. - cached_flag_result = flag_result - if override_match_value is not None: - cached_flag_result = FeatureFlagResult.from_value_and_payload( - key, - flag_value, - self._compute_payload_locally( - key, flag_value, _definition_snapshot=snapshot - ), - ) - if self.flag_cache and cached_flag_result: - # The cache rejects invalidated generations, including writes - # already in flight when new definitions are published. - self.flag_cache.set_cached_flag( - distinct_id, key, cached_flag_result, local_definition_version - ) elif only_evaluate_locally: if self.feature_flags is None: self.log.warning( @@ -4625,128 +4928,225 @@ def _get_feature_flag_result( "Returning None. You can call load_feature_flags() to load flags explicitly." ) else: - try: - ( - flag_details, - request_id, - evaluated_at, - errors_while_computing, - remote_minimal_flag_called_events, - ) = self._get_feature_flag_details_from_server( - key, - distinct_id, - groups, - person_properties, - group_properties, - disable_geoip, - device_id=device_id, - ) - errors = [] - if errors_while_computing: - errors.append(FeatureFlagError.ERRORS_WHILE_COMPUTING) - if flag_details is None: - errors.append(FeatureFlagError.FLAG_MISSING) - if errors: - feature_flag_error = ",".join(errors) - - flag_result = FeatureFlagResult.from_flag_details( - flag_details, override_match_value - ) + self._get_remote_flag_result( + read, + key, + distinct_id, + groups, + person_properties, + group_properties, + disable_geoip, + device_id, + override_match_value, + snapshot["flag_definition_version"], + ) - if ( - flag_details is not None - and flag_details._malformed_field is not None - ): - # The record could not be read, so it reads as disabled. It must - # not replace the last good cached value, which it falls back to - # like a failed request, and it reports what a failed request - # reports: an unknown error and no request details. - feature_flag_error = FeatureFlagError.UNKNOWN_ERROR - request_id = None - evaluated_at = None - flag_result = ( - self._get_stale_flag_fallback(distinct_id, key) or flag_result - ) - # The request-start generation is an invalidation boundary, not - # a claim about the server's definitions. Refresh rejects late writes. - elif self.flag_cache and flag_result: - self.flag_cache._set_cached_remote_flag( - distinct_id, key, flag_result, local_definition_version - ) + if send_feature_flag_events: + self._capture_flag_read_result( + read, + key, + distinct_id, + snapshot, + flag_was_locally_evaluated, + groups, + disable_geoip, + ) + return read.result - self.log.debug( - f"Successfully computed flag remotely: #{key} -> #{flag_result}" - ) - # A payload lookup for a caller's match_value reports that value, not - # the record's. A failure above falls back to a stale cached value. - value_from_details = ( - flag_details is not None and override_match_value is None - ) - except QuotaLimitError as e: - self.log.warning(f"[FEATURE FLAGS] Quota limit exceeded: {e}") - feature_flag_error = FeatureFlagError.QUOTA_LIMITED - flag_result = self._get_stale_flag_fallback(distinct_id, key) - except RequestsTimeout as e: - self.log.warning(f"[FEATURE FLAGS] Request timed out: {e}") - feature_flag_error = FeatureFlagError.TIMEOUT - flag_result = self._get_stale_flag_fallback(distinct_id, key) - except RequestsConnectionError as e: - self.log.warning(f"[FEATURE FLAGS] Connection error: {e}") - feature_flag_error = FeatureFlagError.CONNECTION_ERROR - flag_result = self._get_stale_flag_fallback(distinct_id, key) - except APIError as e: - self.log.warning(f"[FEATURE FLAGS] API error: {e}") - feature_flag_error = FeatureFlagError.api_error(e.status) - flag_result = self._get_stale_flag_fallback(distinct_id, key) - except Exception as e: - self.log.exception(f"[FEATURE FLAGS] Unable to get flag remotely: {e}") - feature_flag_error = FeatureFlagError.UNKNOWN_ERROR - flag_result = self._get_stale_flag_fallback(distinct_id, key) + def _resolve_local_flag_result( + self, key, distinct_id, flag_value, override_match_value, snapshot + ): + local_definition_version = snapshot["flag_definition_version"] + lookup_match_value = ( + override_match_value if override_match_value is not None else flag_value + ) + payload = ( + self._compute_payload_locally( + key, lookup_match_value, _definition_snapshot=snapshot + ) + if lookup_match_value is not None + else None + ) + flag_result = FeatureFlagResult.from_value_and_payload( + key, lookup_match_value, payload + ) - if send_feature_flag_events: - # Locally-evaluated flags carry has_experiment in the stored definition; - # remotely-evaluated flags carry it in the response metadata. None when - # the server (older deployment) does not report it. - has_experiment: Optional[bool] = None - # Source the gate the same way as has_experiment above; see - # _capture_feature_flag_called_if_needed for why. - minimal_flag_called_events = self._minimal_flag_called_events - # Remotely-evaluated flags carry id and version in the response metadata; - # locally-evaluated flags read them from the definition that produced the value. - local_flag_id: Optional[int] = None - local_flag_version: Optional[int] = None - if flag_was_locally_evaluated: - local_def = snapshot["flags_by_key"].get(key) - if isinstance(local_def, dict): - has_experiment = _parse_has_experiment( - local_def.get("has_experiment") - ) - local_flag_id = local_def.get("id") - local_flag_version = local_def.get("version") - elif isinstance(flag_details, FeatureFlag): - has_experiment = _metadata_has_experiment(flag_details.metadata) - minimal_flag_called_events = remote_minimal_flag_called_events + # Cache the local evaluation, not a payload lookup override. + cached_flag_result = flag_result + if override_match_value is not None: + cached_flag_result = FeatureFlagResult.from_value_and_payload( + key, + flag_value, + self._compute_payload_locally( + key, flag_value, _definition_snapshot=snapshot + ), + ) + if self.flag_cache and cached_flag_result: + # The cache rejects invalidated generations, including writes + # already in flight when new definitions are published. + self.flag_cache.set_cached_flag( + distinct_id, key, cached_flag_result, local_definition_version + ) + return flag_result - self._capture_feature_flag_called( - distinct_id, + def _get_remote_flag_result( + self, + read, + key, + distinct_id, + groups, + person_properties, + group_properties, + disable_geoip, + device_id, + override_match_value, + local_definition_version, + ): + try: + ( + read.details, + read.request_id, + read.evaluated_at, + errors_while_computing, + read.minimal_flag_called_events, + ) = self._get_feature_flag_details_from_server( key, - flag_result.get_value() if flag_result else None, - flag_result.payload if flag_result else None, - flag_was_locally_evaluated, + distinct_id, groups, + person_properties, + group_properties, disable_geoip, - request_id, - evaluated_at, - flag_details, - feature_flag_error, - has_experiment, - minimal_flag_called_events, - local_flag_id=local_flag_id, - local_flag_version=local_flag_version, - value_from_details=value_from_details, + device_id=device_id, + ) + self._resolve_remote_flag_result( + read, + key, + distinct_id, + override_match_value, + local_definition_version, + errors_while_computing, ) - return flag_result + self.log.debug( + f"Successfully computed flag remotely: #{key} -> #{read.result}" + ) + # A payload lookup for a caller's match_value reports that value, not + # the record's. A failure above falls back to a stale cached value. + read.value_from_details = ( + read.details is not None and override_match_value is None + ) + except QuotaLimitError as e: + self.log.warning(f"[FEATURE FLAGS] Quota limit exceeded: {e}") + read.error = FeatureFlagError.QUOTA_LIMITED + read.result = self._get_stale_flag_fallback(distinct_id, key) + except RequestsTimeout as e: + self.log.warning(f"[FEATURE FLAGS] Request timed out: {e}") + read.error = FeatureFlagError.TIMEOUT + read.result = self._get_stale_flag_fallback(distinct_id, key) + except RequestsConnectionError as e: + self.log.warning(f"[FEATURE FLAGS] Connection error: {e}") + read.error = FeatureFlagError.CONNECTION_ERROR + read.result = self._get_stale_flag_fallback(distinct_id, key) + except APIError as e: + self.log.warning(f"[FEATURE FLAGS] API error: {e}") + read.error = FeatureFlagError.api_error(e.status) + read.result = self._get_stale_flag_fallback(distinct_id, key) + except Exception as e: + self.log.exception(f"[FEATURE FLAGS] Unable to get flag remotely: {e}") + read.error = FeatureFlagError.UNKNOWN_ERROR + read.result = self._get_stale_flag_fallback(distinct_id, key) + + def _resolve_remote_flag_result( + self, + read, + key, + distinct_id, + override_match_value, + local_definition_version, + errors_while_computing, + ): + read.error = self._remote_flag_result_error( + errors_while_computing, read.details + ) + + read.result = FeatureFlagResult.from_flag_details( + read.details, override_match_value + ) + + if read.details is not None and read.details._malformed_field is not None: + # The record could not be read, so it reads as disabled. It must + # not replace the last good cached value, which it falls back to + # like a failed request, and it reports what a failed request + # reports: an unknown error and no request details. + read.error = FeatureFlagError.UNKNOWN_ERROR + read.request_id = None + read.evaluated_at = None + read.result = self._get_stale_flag_fallback(distinct_id, key) or read.result + # The request-start generation is an invalidation boundary, not + # a claim about the server's definitions. Refresh rejects late writes. + elif self.flag_cache and read.result: + self.flag_cache._set_cached_remote_flag( + distinct_id, key, read.result, local_definition_version + ) + + @staticmethod + def _remote_flag_result_error(errors_while_computing, flag_details): + errors = [] + if errors_while_computing: + errors.append(FeatureFlagError.ERRORS_WHILE_COMPUTING) + if flag_details is None: + errors.append(FeatureFlagError.FLAG_MISSING) + return ",".join(errors) if errors else None + + def _capture_flag_read_result( + self, + read, + key, + distinct_id, + snapshot, + flag_was_locally_evaluated, + groups, + disable_geoip, + ): + # Locally-evaluated flags carry has_experiment in the stored definition; + # remotely-evaluated flags carry it in the response metadata. None when + # the server (older deployment) does not report it. + has_experiment: Optional[bool] = None + # Source the gate the same way as has_experiment above; see + # _capture_feature_flag_called_if_needed for why. + minimal_flag_called_events = self._minimal_flag_called_events + # Remotely-evaluated flags carry id and version in the response metadata; + # locally-evaluated flags read them from the definition that produced the value. + local_flag_id: Optional[int] = None + local_flag_version: Optional[int] = None + if flag_was_locally_evaluated: + local_def = snapshot["flags_by_key"].get(key) + if isinstance(local_def, dict): + has_experiment = _parse_has_experiment(local_def.get("has_experiment")) + local_flag_id = local_def.get("id") + local_flag_version = local_def.get("version") + elif isinstance(read.details, FeatureFlag): + has_experiment = _metadata_has_experiment(read.details.metadata) + minimal_flag_called_events = read.minimal_flag_called_events + + self._capture_feature_flag_called( + distinct_id, + key, + read.result.get_value() if read.result else None, + read.result.payload if read.result else None, + flag_was_locally_evaluated, + groups, + disable_geoip, + read.request_id, + read.evaluated_at, + read.details, + read.error, + has_experiment, + minimal_flag_called_events, + local_flag_id=local_flag_id, + local_flag_version=local_flag_version, + value_from_details=read.value_from_details, + ) def get_feature_flag_result( self, @@ -5059,14 +5459,7 @@ def _capture_feature_flag_called( properties["$feature_flag_request_id"] = request_id if evaluated_at: properties["$feature_flag_evaluated_at"] = evaluated_at - if isinstance(flag_details, FeatureFlag): - if flag_details.reason and flag_details.reason.description: - properties["$feature_flag_reason"] = flag_details.reason.description - if isinstance(flag_details.metadata, FlagMetadata): - if flag_details.metadata.version: - properties["$feature_flag_version"] = flag_details.metadata.version - if flag_details.metadata.id: - properties["$feature_flag_id"] = flag_details.metadata.id + self._add_remote_flag_called_metadata(properties, flag_details) if local_flag_version: properties["$feature_flag_version"] = local_flag_version if local_flag_id: @@ -5087,6 +5480,17 @@ def _capture_feature_flag_called( value_from_details=value_from_details, ) + @staticmethod + def _add_remote_flag_called_metadata(properties, flag_details): + if isinstance(flag_details, FeatureFlag): + if flag_details.reason and flag_details.reason.description: + properties["$feature_flag_reason"] = flag_details.reason.description + if isinstance(flag_details.metadata, FlagMetadata): + if flag_details.metadata.version: + properties["$feature_flag_version"] = flag_details.metadata.version + if flag_details.metadata.id: + properties["$feature_flag_id"] = flag_details.metadata.id + def _capture_feature_flag_called_if_needed( self, *, @@ -5399,15 +5803,9 @@ def evaluate_flags( """ host = self._get_feature_flag_evaluations_host() - if distinct_id is None: - distinct_id = get_context_distinct_id() - - # Resolve device_id from context when not explicitly provided. The context value - # may be set via tracing headers; the explicit parameter is an override for callers - # who want to bypass it. Used by the remote /flags request for experience-continuity - # flag matching. - if device_id is None: - device_id = get_context_device_id() + distinct_id, device_id = self._resolve_feature_flag_identity( + distinct_id, device_id + ) if not distinct_id or self.disabled: # Empty snapshot. The class short-circuits on empty distinct_id so calling @@ -5431,14 +5829,45 @@ def evaluate_flags( ) ) groups = groups or {} - requested_keys = set(flag_keys) if flag_keys else None + return self._evaluate_flags_with_context( + host, + distinct_id, + groups, + person_properties, + group_properties, + only_evaluate_locally, + disable_geoip, + flag_keys, + device_id, + ) - records: Dict[str, _EvaluatedFlagRecord] = {} - request_id: Optional[str] = None - evaluated_at: Optional[int] = None - errors_while_computing = False - quota_limited = False - locally_evaluated_keys: set[str] = set() + @staticmethod + def _resolve_feature_flag_identity(distinct_id, device_id): + if distinct_id is None: + distinct_id = get_context_distinct_id() + + # Resolve device_id from context when not explicitly provided. The context value + # may be set via tracing headers; the explicit parameter is an override for callers + # who want to bypass it. Used by the remote /flags request for experience-continuity + # flag matching. + if device_id is None: + device_id = get_context_device_id() + + return distinct_id, device_id + + def _evaluate_flags_with_context( + self, + host, + distinct_id, + groups, + person_properties, + group_properties, + only_evaluate_locally, + disable_geoip, + flag_keys, + device_id, + ) -> FeatureFlagEvaluations: + requested_keys = set(flag_keys) if flag_keys else None unresolved: Dict[str, UnresolvedFlagReason] = {} # Try local evaluation first when the poller has loaded definitions. local_person_properties = self._person_properties_for_local_evaluation( @@ -5461,14 +5890,55 @@ def evaluate_flags( minimal_flag_called_events = definition_snapshot.get( "minimal_flag_called_events", False ) - feature_flags_by_key = definition_snapshot["flags_by_key"] - local_flags = local_result.get("featureFlags") or {} - local_payloads = local_result.get("featureFlagPayloads") or {} - if requested_keys and not requested_keys.issubset(local_flags): + result = _FlagEvaluationResult( + records=self._local_evaluation_records(local_result, definition_snapshot), + minimal_flag_called_events=minimal_flag_called_events, + ) + if requested_keys and not requested_keys.issubset(result.records): # A requested flag may have been created since the last definitions poll. # Ask /flags for the caller's original scope unless this is a local-only call. fallback_to_server = True + # Fall back to remote evaluation for any flags the poller couldn't resolve locally. + # Use the flags decision path directly so the resulting records carry id/version/reason + # and fired ``$feature_flag_called`` events match what ``get_feature_flag()`` emits. + if fallback_to_server and not only_evaluate_locally: + self._evaluate_flags_remotely( + result, + distinct_id, + groups, + person_properties, + group_properties, + disable_geoip, + flag_keys, + device_id, + requested_keys, + ) + + return FeatureFlagEvaluations( + host=host, + distinct_id=str(distinct_id), + flags=result.records, + groups=groups, + disable_geoip=disable_geoip, + request_id=result.request_id, + evaluated_at=result.evaluated_at, + errors_while_computing=result.errors_while_computing, + quota_limited=result.quota_limited, + minimal_flag_called_events=result.minimal_flag_called_events, + unresolved_flags={ + key: reason + for key, reason in unresolved.items() + if key not in result.records + }, + ) + + @staticmethod + def _local_evaluation_records(local_result, definition_snapshot): + records: Dict[str, _EvaluatedFlagRecord] = {} + feature_flags_by_key = definition_snapshot["flags_by_key"] + local_flags = local_result.get("featureFlags") or {} + local_payloads = local_result.get("featureFlagPayloads") or {} for key, value in local_flags.items(): flag_def = feature_flags_by_key.get(key) or {} records[key] = _LocalFlagRecord( @@ -5482,62 +5952,62 @@ def evaluate_flags( locally_evaluated=True, has_experiment=_parse_has_experiment(flag_def.get("has_experiment")), ) - locally_evaluated_keys.add(key) - # Fall back to remote evaluation for any flags the poller couldn't resolve locally. - # Use the flags decision path directly so the resulting records carry id/version/reason - # and fired ``$feature_flag_called`` events match what ``get_feature_flag()`` emits. - if fallback_to_server and not only_evaluate_locally: - try: - response = self._get_flags_decision( - distinct_id, - groups=groups, - person_properties=person_properties, - group_properties=group_properties, - disable_geoip=disable_geoip, - flag_keys_to_evaluate=flag_keys, - device_id=device_id, - ) - request_id = response.get("requestId") - raw_evaluated_at = response.get("evaluatedAt") - evaluated_at = ( - raw_evaluated_at if isinstance(raw_evaluated_at, int) else None - ) - errors_while_computing = bool( - response.get("errorsWhileComputingFlags", False) - ) - minimal_flag_called_events = ( - response.get("minimalFlagCalledEvents") is True - ) - for key, detail in response.get("flags", {}).items(): - if requested_keys is not None and key not in requested_keys: - continue - if key in locally_evaluated_keys: - continue - records[key] = _RemoteFlagRecord(detail) - except QuotaLimitError as e: - self.log.warning(f"[FEATURE FLAGS] Quota limit exceeded: {e}") - quota_limited = True - except Exception as e: - self.log.exception( - f"[FEATURE FLAGS] Unable to evaluate flags remotely: {e}" - ) + return records - return FeatureFlagEvaluations( - host=host, - distinct_id=str(distinct_id), - flags=records, - groups=groups, - disable_geoip=disable_geoip, - request_id=request_id, - evaluated_at=evaluated_at, - errors_while_computing=errors_while_computing, - quota_limited=quota_limited, - minimal_flag_called_events=minimal_flag_called_events, - unresolved_flags={ - key: reason for key, reason in unresolved.items() if key not in records - }, - ) + def _evaluate_flags_remotely( + self, + result, + distinct_id, + groups, + person_properties, + group_properties, + disable_geoip, + flag_keys, + device_id, + requested_keys, + ): + local_keys = set(result.records) + try: + response = self._get_flags_decision( + distinct_id, + groups=groups, + person_properties=person_properties, + group_properties=group_properties, + disable_geoip=disable_geoip, + flag_keys_to_evaluate=flag_keys, + device_id=device_id, + ) + result.request_id = response.get("requestId") + raw_evaluated_at = response.get("evaluatedAt") + result.evaluated_at = ( + raw_evaluated_at if isinstance(raw_evaluated_at, int) else None + ) + result.errors_while_computing = bool( + response.get("errorsWhileComputingFlags", False) + ) + result.minimal_flag_called_events = ( + response.get("minimalFlagCalledEvents") is True + ) + self._merge_remote_evaluation_records( + result.records, response, requested_keys, local_keys + ) + except QuotaLimitError as e: + self.log.warning(f"[FEATURE FLAGS] Quota limit exceeded: {e}") + result.quota_limited = True + except Exception as e: + self.log.exception( + f"[FEATURE FLAGS] Unable to evaluate flags remotely: {e}" + ) + + @staticmethod + def _merge_remote_evaluation_records(records, response, requested_keys, local_keys): + for key, detail in response.get("flags", {}).items(): + if requested_keys is not None and key not in requested_keys: + continue + if key in local_keys: + continue + records[key] = _RemoteFlagRecord(detail) _feature_flag_evaluations_host_cache: Optional[_FeatureFlagEvaluationsHost] = None @@ -5573,43 +6043,19 @@ def _get_all_flags_and_payloads_locally( snapshot = self._local_evaluation_snapshot() # If loading in previous line failed if snapshot["flags"]: - # Filter flags based on flag_keys_to_evaluate if provided - flags_to_process = snapshot["flags"] - if flag_keys_to_evaluate: - flag_keys_set = set(flag_keys_to_evaluate) - flags_to_process = [ - flag for flag in snapshot["flags"] if flag["key"] in flag_keys_set - ] - - for flag in flags_to_process: - try: - flags[flag["key"]] = self._compute_flag_locally( - flag, - distinct_id, - groups=groups, - person_properties=person_properties, - group_properties=group_properties, - warn_on_unknown_groups=warn_on_unknown_groups, - device_id=device_id, - _definition_snapshot=snapshot, - ) - matched_payload = self._compute_payload_locally( - flag["key"], - flags[flag["key"]], - _definition_snapshot=snapshot, - ) - if matched_payload is not None: - payloads[flag["key"]] = matched_payload - except (InconclusiveMatchError, RequiresServerEvaluation) as e: - # No need to log this, since it's just telling us to fall back to `/flags` - fallback_to_flags = True - if unresolved is not None: - unresolved[flag["key"]] = e.reason - except Exception as e: - self.log.exception( - f"[FEATURE FLAGS] Error while computing variant and payload: {e}" - ) - fallback_to_flags = True + fallback_to_flags = self._evaluate_local_flag_set( + snapshot, + distinct_id, + groups, + person_properties, + group_properties, + warn_on_unknown_groups, + flag_keys_to_evaluate, + device_id, + unresolved, + flags, + payloads, + ) else: fallback_to_flags = True @@ -5622,6 +6068,67 @@ def _get_all_flags_and_payloads_locally( snapshot, ) + @staticmethod + def _select_flags_for_local_evaluation(snapshot, flag_keys_to_evaluate): + # Filter flags based on flag_keys_to_evaluate if provided + flags_to_process = snapshot["flags"] + if flag_keys_to_evaluate: + flag_keys_set = set(flag_keys_to_evaluate) + flags_to_process = [ + flag for flag in snapshot["flags"] if flag["key"] in flag_keys_set + ] + + return flags_to_process + + def _evaluate_local_flag_set( + self, + snapshot, + distinct_id, + groups, + person_properties, + group_properties, + warn_on_unknown_groups, + flag_keys_to_evaluate, + device_id, + unresolved, + flags, + payloads, + ): + fallback_to_flags = False + flags_to_process = self._select_flags_for_local_evaluation( + snapshot, flag_keys_to_evaluate + ) + for flag in flags_to_process: + try: + flags[flag["key"]] = self._compute_flag_locally( + flag, + distinct_id, + groups=groups, + person_properties=person_properties, + group_properties=group_properties, + warn_on_unknown_groups=warn_on_unknown_groups, + device_id=device_id, + _definition_snapshot=snapshot, + ) + matched_payload = self._compute_payload_locally( + flag["key"], + flags[flag["key"]], + _definition_snapshot=snapshot, + ) + if matched_payload is not None: + payloads[flag["key"]] = matched_payload + except (InconclusiveMatchError, RequiresServerEvaluation) as e: + # No need to log this, since it's just telling us to fall back to `/flags` + fallback_to_flags = True + if unresolved is not None: + unresolved[flag["key"]] = e.reason + except Exception as e: + self.log.exception( + f"[FEATURE FLAGS] Error while computing variant and payload: {e}" + ) + fallback_to_flags = True + return fallback_to_flags + def _initialize_flag_cache(self, cache_url): """Initialize feature flag cache for graceful degradation during service outages. @@ -5672,37 +6179,7 @@ def _initialize_flag_cache(self, cache_url): return FlagCache(size, ttl) elif scheme in ("redis", "rediss"): - try: - # Not worth importing redis if we're not using it - import redis - - redis_url = f"{parsed.scheme}://" - if parsed.username or parsed.password: - redis_url += f"{parsed.username or ''}:{parsed.password or ''}@" - host = parsed.hostname or "localhost" - if ":" in host: - host = f"[{host}]" - redis_url += f"{host}:{parsed.port or 6379}" - if parsed.path: - redis_url += parsed.path - - client = redis.from_url(redis_url) - - # Test connection before using it - client.ping() - - return RedisFlagCache(client, default_ttl=ttl) - - except ImportError: - self.log.warning( - "[FEATURE FLAGS] Redis not available, flag caching disabled" - ) - return None - except Exception as e: - self.log.warning( - f"[FEATURE FLAGS] Redis connection failed: {e}, flag caching disabled" - ) - return None + return self._initialize_redis_flag_cache(parsed, ttl) else: raise ValueError( f"Unknown cache URL scheme: {scheme}. Supported schemes: memory, redis, rediss" @@ -5714,6 +6191,45 @@ def _initialize_flag_cache(self, cache_url): ) return None + def _initialize_redis_flag_cache(self, parsed, ttl): + try: + # Not worth importing redis if we're not using it + import redis + + redis_url = self._redis_flag_cache_url(parsed) + + client = redis.from_url(redis_url) + + # Test connection before using it + client.ping() + + return RedisFlagCache(client, default_ttl=ttl) + + except ImportError: + self.log.warning( + "[FEATURE FLAGS] Redis not available, flag caching disabled" + ) + return None + except Exception as e: + self.log.warning( + f"[FEATURE FLAGS] Redis connection failed: {e}, flag caching disabled" + ) + return None + + @staticmethod + def _redis_flag_cache_url(parsed): + redis_url = f"{parsed.scheme}://" + if parsed.username or parsed.password: + redis_url += f"{parsed.username or ''}:{parsed.password or ''}@" + host = parsed.hostname or "localhost" + if ":" in host: + host = f"[{host}]" + redis_url += f"{host}:{parsed.port or 6379}" + if parsed.path: + redis_url += parsed.path + + return redis_url + def feature_flag_definitions(self): """ Return feature flag definitions loaded for local evaluation. diff --git a/posthog/consumer.py b/posthog/consumer.py index 0190dcf01..d9de6d441 100644 --- a/posthog/consumer.py +++ b/posthog/consumer.py @@ -73,14 +73,17 @@ def draining(self, consumer) -> bool: with self._queue.mutex: return self._requests > 0 and (consumer.running or consumer._drain_on_stop) + def _is_draining_locked(self, consumer) -> bool: + return self._requests > 0 and ( + consumer is None or consumer.running or consumer._drain_on_stop + ) + def get(self, timeout: float, consumer=None): """Get an item, or wake with ``Empty`` when draining or stopping.""" with self._queue.not_empty: deadline = time.monotonic() + timeout while True: - draining = self._requests > 0 and ( - consumer is None or consumer.running or consumer._drain_on_stop - ) + draining = self._is_draining_locked(consumer) if consumer is not None and not consumer.running and not draining: raise Empty if self._queue._qsize(): @@ -213,6 +216,44 @@ def _can_upload(self) -> bool: with self.queue.mutex: return self.running or self._drain_on_stop + def _event_size(self, item) -> Optional[int]: + try: + item_size = len(json.dumps(item, cls=DatetimeSerializer).encode()) + except Exception: + # Callback-modified events can still contain invalid mapping + # keys or circular references. Never log the payload here. + self.log.error("Unable to serialize queued event for sizing, dropping.") + return None + if item_size > self.max_msg_size: + # Log only name and size: AI events may carry unredacted + # multimodal payloads that must not leak into logs. + self.log.error( + "Event %s (%d bytes) exceeds the %dKiB limit for %s, dropping.", + item.get("event") if isinstance(item, dict) else type(item), + item_size, + self.max_msg_size // 1024, + self.endpoint, + ) + return None + return item_size + + def _next_queue_timeout(self, start_time: float) -> Optional[float]: + # While draining we take only what is already queued, never waiting + # for `flush_interval` to elapse or for `flush_at` to be reached. + draining = self._draining() + if not self.running and not draining: + return None + remaining = self.flush_interval - (time.monotonic() - start_time) + if not draining and remaining <= 0: + return None + + return 0 if draining else remaining + + @staticmethod + def _acknowledge(queue, count: int) -> None: + for _ in range(count): + queue.task_done() + def next(self): """Return the next batch of items to upload.""" queue = self.queue @@ -224,47 +265,21 @@ def next(self): try: while len(items) < self.flush_at: - # While draining we take only what is already queued, never waiting - # for `flush_interval` to elapse or for `flush_at` to be reached. - draining = self._draining() - if not self.running and not draining: - break - remaining = self.flush_interval - (time.monotonic() - start_time) - if not draining and remaining <= 0: + timeout = self._next_queue_timeout(start_time) + if timeout is None: break try: if self._drain_signal is not None: item = self._drain_signal.get( - timeout=0 if draining else remaining, + timeout=timeout, consumer=self, ) else: - item = queue.get(block=True, timeout=remaining) + item = queue.get(block=True, timeout=timeout) pending_items += 1 - try: - item_size = len( - json.dumps(item, cls=DatetimeSerializer).encode() - ) - except Exception: - # Callback-modified events can still contain invalid mapping - # keys or circular references. Never log the payload here. - self.log.error( - "Unable to serialize queued event for sizing, dropping." - ) - queue.task_done() - pending_items -= 1 - continue - if item_size > self.max_msg_size: - # Log only name and size: AI events may carry unredacted - # multimodal payloads that must not leak into logs. - self.log.error( - "Event %s (%d bytes) exceeds the %dKiB limit for %s, dropping.", - item.get("event") if isinstance(item, dict) else type(item), - item_size, - self.max_msg_size // 1024, - self.endpoint, - ) + item_size = self._event_size(item) + if item_size is None: queue.task_done() pending_items -= 1 continue @@ -276,13 +291,11 @@ def next(self): except Empty: break except BaseException: - for _ in range(pending_items): - queue.task_done() + self._acknowledge(queue, pending_items) raise if not self._can_upload(): - for _ in range(pending_items): - queue.task_done() + self._acknowledge(queue, pending_items) return [] return items diff --git a/posthog/emails.py b/posthog/emails.py index 42e70b971..c8a89b7f9 100644 --- a/posthog/emails.py +++ b/posthog/emails.py @@ -378,14 +378,7 @@ def _prepare_code_send(args: Mapping[str, Any]) -> _PreparedSend: raise ValueError( "template_key can only have letters, digits, hyphens (-) and underscores (_), at most 400." ) - content = {"subject": _nonempty_text(args, "subject")} - content.update( - {field: _string(args, field) for field in ("html", "text") if field in args} - ) - if "html" not in content and "text" not in content: - raise TypeError("send() needs html, text or both.") - if not (content.get("html") or content.get("text")): - raise ValueError("html or text must be a nonempty string.") + content = _code_content(args) payload: dict[str, Any] = {"to": to, "template_key": template_key, **content} if distinct_id is not None: @@ -399,6 +392,19 @@ def _prepare_code_send(args: Mapping[str, Any]) -> _PreparedSend: ) +def _code_content(args: Mapping[str, Any]) -> dict[str, str]: + content = {"subject": _nonempty_text(args, "subject")} + content.update( + {field: _string(args, field) for field in ("html", "text") if field in args} + ) + if "html" not in content and "text" not in content: + raise TypeError("send() needs html, text or both.") + if not (content.get("html") or content.get("text")): + raise ValueError("html or text must be a nonempty string.") + + return content + + def _sender(args: Mapping[str, Any]) -> dict[str, Any]: sender: dict[str, Any] = {} if "from_" in args: @@ -523,13 +529,7 @@ def _status_path(message_id: Union[str, UUID]) -> str: def _validate_options( timeout: float, max_retries: int, migrated_from: Optional[str] ) -> None: - if ( - isinstance(timeout, bool) - or not isinstance(timeout, (int, float)) - or not math.isfinite(timeout) - or timeout <= 0 - ): - raise ValueError("timeout must be a finite positive number of seconds.") + _validate_timeout(timeout) if ( isinstance(max_retries, bool) or not isinstance(max_retries, int) @@ -544,6 +544,16 @@ def _validate_options( ) +def _validate_timeout(timeout: float) -> None: + if ( + isinstance(timeout, bool) + or not isinstance(timeout, (int, float)) + or not math.isfinite(timeout) + or timeout <= 0 + ): + raise ValueError("timeout must be a finite positive number of seconds.") + + def _text(value: Any) -> Optional[str]: return value if isinstance(value, str) and value else None @@ -553,32 +563,33 @@ def _error_name(status: Optional[int], code: str) -> str: return "missing_api_key" if status is None: return "service_unavailable" if code == "network_error" else "application_error" - if status == 400: - return "validation_error" - if status == 403: - return "invalid_permission" - if status == 404: - return "not_found" + if status in (409, 422): + return _idempotency_error_name(status, code) + if status == 429: + return ( + "monthly_quota_exceeded" if code in _QUOTA_CODES else "rate_limit_exceeded" + ) + return { + 400: "validation_error", + 403: "invalid_permission", + 404: "not_found", + 503: "service_unavailable", + }.get(status, "application_error") + + +def _idempotency_error_name(status: int, code: str) -> str: if status == 409: return ( "validation_error" if code == "template_outdated" else "concurrent_idempotent_requests" ) - if status == 422: - # The template send route also answers 422 for send checks such as recipient_suppressed. - return ( - "invalid_idempotent_request" - if code == "idempotency_key_reused" - else "validation_error" - ) - if status == 429: - return ( - "monthly_quota_exceeded" if code in _QUOTA_CODES else "rate_limit_exceeded" - ) - if status == 503: - return "service_unavailable" - return "application_error" + # The template send route also answers 422 for send checks such as recipient_suppressed. + return ( + "invalid_idempotent_request" + if code == "idempotency_key_reused" + else "validation_error" + ) def _sdk_error(code: str, message: str, status: Optional[int] = None) -> EmailError: @@ -608,21 +619,23 @@ def _parse_response( return SendEmailResult( id=message_id, replayed=fields.get("replayed") is True, - warnings=tuple( - EmailWarning( - code=_text(warning.get("code")) or "warning", - message=_text(warning.get("detail")) - or _text(warning.get("message")) - or "", - fix=_text(warning.get("fix")), - ) - for warning in (warnings if isinstance(warnings, list) else []) - if isinstance(warning, dict) - ), + warnings=_parse_warnings(warnings), ) raise _response_error(fields, status, retry_after_header) +def _parse_warnings(warnings: Any) -> tuple[EmailWarning, ...]: + return tuple( + EmailWarning( + code=_text(warning.get("code")) or "warning", + message=_text(warning.get("detail")) or _text(warning.get("message")) or "", + fix=_text(warning.get("fix")), + ) + for warning in (warnings if isinstance(warnings, list) else []) + if isinstance(warning, dict) + ) + + def _timestamp(value: Any) -> Optional[datetime]: text = _text(value) if text is None: diff --git a/posthog/exception_utils.py b/posthog/exception_utils.py index a6f5883b2..d462d23cc 100644 --- a/posthog/exception_utils.py +++ b/posthog/exception_utils.py @@ -415,20 +415,7 @@ def get_lines_from_file( ): # type: (...) -> Tuple[List[Annotated[str]], Optional[Annotated[str]], List[Annotated[str]]] context_lines = 5 - source = None - if loader is not None and hasattr(loader, "get_source"): - try: - source_str = loader.get_source(module) # type: Optional[str] - except (ImportError, IOError): - source_str = None - if source_str is not None: - source = source_str.splitlines() - - if source is None: - try: - source = linecache.getlines(filename) - except (OSError, IOError): - return [], None, [] + source = _get_source_lines(filename, loader, module) if not source: return [], None, [] @@ -452,6 +439,26 @@ def get_lines_from_file( return [], None, [] +def _get_source_lines(filename, loader, module): + # type: (str, Optional[Any], Optional[str]) -> Optional[List[str]] + source = None + if loader is not None and hasattr(loader, "get_source"): + try: + source_str = loader.get_source(module) # type: Optional[str] + except (ImportError, IOError): + source_str = None + if source_str is not None: + source = source_str.splitlines() + + if source is None: + try: + source = linecache.getlines(filename) + except (OSError, IOError): + return None + + return source + + def get_source_context( frame, # type: FrameType tb_lineno, # type: int @@ -583,14 +590,16 @@ def _valid_mechanism(mechanism): if key not in {"type", "handled", "source", "synthetic", "exception_id", "parent_id"} } - if isinstance(mechanism.get("type"), str) and mechanism["type"]: - result["type"] = mechanism["type"] - if isinstance(mechanism.get("handled"), bool): - result["handled"] = mechanism["handled"] - if isinstance(mechanism.get("source"), str) and mechanism["source"]: - result["source"] = mechanism["source"] - if isinstance(mechanism.get("synthetic"), bool): - result["synthetic"] = mechanism["synthetic"] + for field, field_type in ( + ("type", str), + ("handled", bool), + ("source", str), + ("synthetic", bool), + ): + if isinstance(mechanism.get(field), field_type) and ( + field_type is bool or mechanism[field] + ): + result[field] = mechanism[field] return result @@ -653,9 +662,36 @@ def single_exception_from_error_tuple( Creates a dict that goes into the events `exception.values` list """ exception_value = {} # type: Dict[str, Any] - exception_value["mechanism"] = _valid_mechanism(mechanism) + exception_value["mechanism"] = _exception_mechanism( + exc_value, tb, mechanism, exception_id, parent_id, source + ) + + exception_value["module"] = get_type_module(exc_type) + exception_value["type"] = get_type_name(exc_type) + exception_value["value"] = get_error_message(exc_value) + + max_value_length = DEFAULT_MAX_VALUE_LENGTH # fallback + + frames = [ + serialize_frame( + tb.tb_frame, + tb_lineno=tb.tb_lineno, + max_value_length=max_value_length, + ) + for tb in iter_stacks(tb) + ] + + if frames: + exception_value["stacktrace"] = {"frames": frames, "type": "raw"} + + return exception_value + + +def _exception_mechanism(exc_value, tb, mechanism, exception_id, parent_id, source): + # type: (Optional[BaseException], Optional[TracebackType], Optional[Dict[str, Any]], Optional[int], Optional[int], Optional[str]) -> Dict[str, Any] + result = _valid_mechanism(mechanism) if exception_id is not None: - exception_value["mechanism"]["exception_id"] = exception_id + result["exception_id"] = exception_id if exc_value is not None: errno = get_errno(exc_value) @@ -663,27 +699,15 @@ def single_exception_from_error_tuple( errno = None if errno is not None: - exception_value["mechanism"].setdefault("meta", {}).setdefault( - "errno", {} - ).setdefault("number", errno) - - if isinstance(source, str) and source: - exception_value["mechanism"]["source"] = source - - is_root_exception = exception_id == 0 - if not is_root_exception and parent_id is not None: - exception_value["mechanism"]["parent_id"] = parent_id - exception_value["mechanism"]["type"] = "chained" - exception_value["mechanism"].pop("handled", None) + result.setdefault("meta", {}).setdefault("errno", {}).setdefault( + "number", errno + ) - if is_root_exception: - exception_value["mechanism"].setdefault("type", "generic") - exception_value["mechanism"].setdefault("handled", True) - exception_value["mechanism"].pop("source", None) + _set_exception_linkage(result, exception_id, parent_id, source) # The stack is synthetic only when the SDK built it, which happens for an # exception that was passed to capture without ever being raised. - exception_value["mechanism"].setdefault( + result.setdefault( "synthetic", tb is not None and _is_synthesized_traceback(exc_value, tb), ) @@ -692,27 +716,26 @@ def single_exception_from_error_tuple( exc_value, BaseExceptionGroup ) if is_exception_group: - exception_value["mechanism"]["is_exception_group"] = True + result["is_exception_group"] = True - exception_value["module"] = get_type_module(exc_type) - exception_value["type"] = get_type_name(exc_type) - exception_value["value"] = get_error_message(exc_value) + return result - max_value_length = DEFAULT_MAX_VALUE_LENGTH # fallback - frames = [ - serialize_frame( - tb.tb_frame, - tb_lineno=tb.tb_lineno, - max_value_length=max_value_length, - ) - for tb in iter_stacks(tb) - ] +def _set_exception_linkage(result, exception_id, parent_id, source): + # type: (Dict[str, Any], Optional[int], Optional[int], Optional[str]) -> None + if isinstance(source, str) and source: + result["source"] = source - if frames: - exception_value["stacktrace"] = {"frames": frames, "type": "raw"} + is_root_exception = exception_id == 0 + if not is_root_exception and parent_id is not None: + result["parent_id"] = parent_id + result["type"] = "chained" + result.pop("handled", None) - return exception_value + if is_root_exception: + result.setdefault("type", "generic") + result.setdefault("handled", True) + result.pop("source", None) HAS_CHAINED_EXCEPTIONS = hasattr(Exception, "__suppress_context__") @@ -803,17 +826,7 @@ def _exceptions_from_error( parent_id = exception_id exception_id += 1 - causing_exception = None # type: Optional[BaseException] - relationship = None # type: Optional[str] - should_suppress_context = ( - hasattr(exc_value, "__suppress_context__") and exc_value.__suppress_context__ # type: ignore - ) - if should_suppress_context and exc_value is not None: - causing_exception = getattr(exc_value, "__cause__", None) - relationship = "cause" - else: - causing_exception = getattr(exc_value, "__context__", None) - relationship = "context" + causing_exception, relationship = _exception_cause(exc_value) if causing_exception is not None and exception_id < _MAX_EXCEPTION_ENTRIES: (exception_id, child_exceptions) = _exceptions_from_error( @@ -828,6 +841,32 @@ def _exceptions_from_error( ) exceptions.extend(child_exceptions) + return _append_exception_group_members( + exc_value, exception_id, parent_id, traversal_state, exceptions + ) + + +def _exception_cause(exc_value): + # type: (Optional[BaseException]) -> Tuple[Optional[BaseException], str] + causing_exception = None # type: Optional[BaseException] + relationship = None # type: Optional[str] + should_suppress_context = ( + hasattr(exc_value, "__suppress_context__") and exc_value.__suppress_context__ # type: ignore + ) + if should_suppress_context and exc_value is not None: + causing_exception = getattr(exc_value, "__cause__", None) + relationship = "cause" + else: + causing_exception = getattr(exc_value, "__context__", None) + relationship = "context" + + return causing_exception, relationship + + +def _append_exception_group_members( + exc_value, exception_id, parent_id, traversal_state, exceptions +): + # type: (Optional[BaseException], int, int, _ExceptionTraversalState, List[Dict[str, Any]]) -> Tuple[int, List[Dict[str, Any]]] # Aggregate-member inspection has its own shared budget. Duplicate and cyclic # references still consume it even though they do not produce output entries. # Cause/context traversal does not consume this budget, so the final inspected @@ -923,13 +962,16 @@ def iter_event_stacktraces(event): if "stacktrace" in event: yield event["stacktrace"] if "threads" in event: - for thread in event["threads"].get("values") or (): - if "stacktrace" in thread: - yield thread["stacktrace"] + yield from _iter_nested_stacktraces(event["threads"].get("values") or ()) if "exception" in event: - for exception in event["exception"].get("values") or (): - if "stacktrace" in exception: - yield exception["stacktrace"] + yield from _iter_nested_stacktraces(event["exception"].get("values") or ()) + + +def _iter_nested_stacktraces(entries): + # type: (Any) -> Iterator[Dict[str, Any]] + for entry in entries: + if "stacktrace" in entry: + yield entry["stacktrace"] def iter_event_frames(event): @@ -1429,11 +1471,24 @@ def _is_high_entropy_secret(value): if not _SECRET_REJECT_CHARS.isdisjoint(distinct): return False + if _secret_character_class_count(distinct) < _SECRET_MIN_CHAR_CLASSES: + return False # identifiers, enums, slugs, numbers + + n = len(value) + entropy = 0.0 + for occurrences in counts.values(): + p = occurrences / n + entropy -= p * math.log2(p) + return entropy >= _SECRET_MIN_ENTROPY_BITS + + +def _secret_character_class_count(distinct): + """Count token character classes, rejecting whitespace and pure-hex ids.""" has_lower = has_upper = has_digit = has_symbol = False hex_only = True for ch in distinct: if ch.isspace(): - return False + return 0 if ch.islower(): has_lower = True if ch not in _HEX_DIGITS: @@ -1449,16 +1504,8 @@ def _is_high_entropy_secret(value): hex_only = False if hex_only: # ObjectId / SHA / md5 digest - return False - if (has_lower + has_upper + has_digit + has_symbol) < _SECRET_MIN_CHAR_CLASSES: - return False # identifiers, enums, slugs, numbers - - n = len(value) - entropy = 0.0 - for occurrences in counts.values(): - p = occurrences / n - entropy -= p * math.log2(p) - return entropy >= _SECRET_MIN_ENTROPY_BITS + return 0 + return has_lower + has_upper + has_digit + has_symbol def _looks_like_secret(value): @@ -1656,41 +1703,30 @@ def _mask_mapping(items, config, seen, depth): placeholder, because the key text reaches the output as-is.""" result: Dict[Any, Any] = {} for key, value in items: - if type(key) is str: - out_key = key_str = key - key_is_json_safe = True - else: - try: - key_str = key if isinstance(key, str) else str(key) - except Exception: - # Without the key text, nothing shows whether the key names a secret, so - # the value is redacted too. - result[_redacted_key(result)] = CODE_VARIABLES_REDACTED_VALUE - continue - # json.dumps only accepts str/int/float/bool/None keys; coerce anything else to - # its string form so one exotic key can't make json.dumps fail. That string form - # is a repr, which can embed field values, not only a name. - key_is_json_safe = ( - key is None - or isinstance(key, (str, int)) # bool is an int subclass - or (isinstance(key, float) and math.isfinite(key)) - ) - out_key = key if key_is_json_safe else key_str + try: + key_str = key if isinstance(key, str) else str(key) + except Exception: + # Without the key text, nothing shows whether the key names a secret, so + # the value is redacted too. + result[_redacted_key(result)] = CODE_VARIABLES_REDACTED_VALUE + continue + out_key, key_is_json_safe = _mapping_key(key, key_str) if len(key_str) > _MAX_VALUE_LENGTH_FOR_PATTERN_MATCH: # Too long to scan, so the key text can't be vouched for either. result[_redacted_key(result)] = CODE_VARIABLES_TOO_LONG_VALUE continue key_matches_mask = _matcher_matches(key_str, config.mask) - if key_matches_mask and not _is_field_name(key, key_is_json_safe): - # A name such as `password` is safe to keep, but other text that matches can - # hold the value itself, e.g. `BasicAuth(login='u', password='...')`. - out_key = _redacted_key(result) - elif config.detect_secrets and _looks_like_secret(key_str): - out_key = _redacted_key(result) - elif not key_is_json_safe and _key_parts_fail_masking(key, config, seen, depth): - out_key = _redacted_key(result) - elif isinstance(out_key, str): - out_key = _redact_embedded_credentials(out_key, config) + out_key = _mask_mapping_key( + key, + key_str, + key_is_json_safe, + out_key, + key_matches_mask, + result, + config, + seen, + depth, + ) if out_key in result: # Two keys can end up with the same text, for example URLs that differ only in # their credentials. A placeholder keeps the later entry from overwriting. @@ -1702,6 +1738,45 @@ def _mask_mapping(items, config, seen, depth): return result +def _mapping_key(key, key_str): + if type(key) is str: + return key, True + # json.dumps only accepts str/int/float/bool/None keys; coerce anything else to + # its string form so one exotic key can't make json.dumps fail. That string form + # is a repr, which can embed field values, not only a name. + key_is_json_safe = ( + key is None + or isinstance(key, (str, int)) # bool is an int subclass + or (isinstance(key, float) and math.isfinite(key)) + ) + out_key = key if key_is_json_safe else key_str + return out_key, key_is_json_safe + + +def _mask_mapping_key( + key, + key_str, + key_is_json_safe, + out_key, + key_matches_mask, + result, + config, + seen, + depth, +): + if key_matches_mask and not _is_field_name(key, key_is_json_safe): + # A name such as `password` is safe to keep, but other text that matches can + # hold the value itself, e.g. `BasicAuth(login='u', password='...')`. + out_key = _redacted_key(result) + elif config.detect_secrets and _looks_like_secret(key_str): + out_key = _redacted_key(result) + elif not key_is_json_safe and _key_parts_fail_masking(key, config, seen, depth): + out_key = _redacted_key(result) + elif isinstance(out_key, str): + out_key = _redact_embedded_credentials(out_key, config) + return out_key + + def _mask_value(value, config, seen=None, depth=0): """Turn any Python value into a JSON-safe, masked value. Single source of truth for what gets redacted; the result contains only JSON-native types so it can be handed @@ -1723,10 +1798,18 @@ def _mask_value(value, config, seen=None, depth=0): return value if t is float: # Non-finite floats (NaN/Infinity) are invalid JSON, so render them as strings. - return value if math.isfinite(value) else str(value) + return _mask_float(value) if value is None: return value + return _mask_non_plain_value(value, config, seen, depth) + + +def _mask_float(value): + return value if math.isfinite(value) else str(value) + +def _mask_non_plain_value(value, config, seen, depth): + t = type(value) if t is not dict and t is not list and t is not tuple: # A scalar subclass (e.g. IntEnum, str subclass) - handle before traversing. if isinstance(value, str): @@ -1735,7 +1818,10 @@ def _mask_value(value, config, seen=None, depth=0): return str(value) if isinstance(value, (bool, int, float)): return value + return _mask_traversable_value(value, config, seen, depth) + +def _mask_traversable_value(value, config, seen, depth): if depth >= _MAX_MASK_DEPTH: # Too deep to keep traversing; fail closed rather than repr (which could leak). return CODE_VARIABLES_TOO_LONG_VALUE @@ -1750,7 +1836,11 @@ def _mask_value(value, config, seen=None, depth=0): if len(seen) > _MAX_TOTAL_NODES_TO_MASK: return CODE_VARIABLES_TOO_LONG_VALUE - if t is dict or isinstance(value, dict): + return _mask_collection_or_object(value, config, seen, depth) + + +def _mask_collection_or_object(value, config, seen, depth): + if type(value) is dict or isinstance(value, dict): if len(value) > _MAX_COLLECTION_ITEMS_TO_SCAN: return CODE_VARIABLES_TOO_LONG_VALUE return _mask_mapping(value.items(), config, seen, depth) @@ -1758,24 +1848,35 @@ def _mask_value(value, config, seen=None, depth=0): # namedtuples are tuples but their fields have names: traverse like an object (so a # field named `password` is caught) and emit a dict the encoder can serialize directly. if isinstance(value, tuple) and hasattr(value, "_fields"): - fields = value._fields - if len(fields) > _MAX_COLLECTION_ITEMS_TO_SCAN: - return CODE_VARIABLES_TOO_LONG_VALUE - masked = _mask_mapping(zip(fields, value), config, seen, depth) - masked["__class__"] = _safe_type_name(value) - return masked + return _mask_namedtuple(value, config, seen, depth) if isinstance(value, (list, tuple)): - if len(value) > _MAX_COLLECTION_ITEMS_TO_SCAN: - return CODE_VARIABLES_TOO_LONG_VALUE - masked_items = [_mask_value(item, config, seen, depth + 1) for item in value] - try: - return type(value)(masked_items) - except Exception: - # A list/tuple subclass whose constructor rejects a single iterable; the items - # are already masked, so fall back to a plain list. - return masked_items + return _mask_sequence(value, config, seen, depth) + return _mask_object(value, config, seen, depth) + + +def _mask_sequence(value, config, seen, depth): + if len(value) > _MAX_COLLECTION_ITEMS_TO_SCAN: + return CODE_VARIABLES_TOO_LONG_VALUE + masked_items = [_mask_value(item, config, seen, depth + 1) for item in value] + try: + return type(value)(masked_items) + except Exception: + # A list/tuple subclass whose constructor rejects a single iterable; the items + # are already masked, so fall back to a plain list. + return masked_items + + +def _mask_namedtuple(value, config, seen, depth): + fields = value._fields + if len(fields) > _MAX_COLLECTION_ITEMS_TO_SCAN: + return CODE_VARIABLES_TOO_LONG_VALUE + masked = _mask_mapping(zip(fields, value), config, seen, depth) + masked["__class__"] = _safe_type_name(value) + return masked + +def _mask_object(value, config, seen, depth): # Custom objects: traverse their real attributes so a field named e.g. `password` is # caught by name, rather than repr-scanning (which a custom __repr__ could relabel). attrs = _extract_object_attrs(value) @@ -1813,31 +1914,7 @@ def _encode_variable(value, config, limiter): try: safe = _mask_value(value, config) - if safe is None: - result = "None" - elif isinstance(safe, bool): - result = str(safe) - elif isinstance(safe, float) and not math.isfinite(safe): - # Only reachable when masking is disabled; keep NaN/Infinity out of the JSON. - result = str(safe) - elif isinstance(safe, (int, float)): - # Numbers are emitted as raw JSON numbers, so they skip string truncation. - result_size = len(str(safe)) - if not limiter.can_add(result_size): - return None - limiter.add(result_size) - return safe - elif isinstance(safe, str): - result = safe - else: - # `default` is a safety net for anything _mask_value left non-serializable. - result = json.dumps( - safe, - default=lambda o: _safe_repr(o, config), - allow_nan=False, - ) - - return _finalize(result, limiter, config.max_length) + return _encode_masked_variable(safe, config, limiter) except Exception: # Fail closed: even if json.dumps chokes, re-render through the masking-aware repr. try: @@ -1847,6 +1924,34 @@ def _encode_variable(value, config, limiter): return _finalize(rendered, limiter, config.max_length) +def _encode_masked_variable(safe, config, limiter): + if safe is None: + result = "None" + elif isinstance(safe, bool): + result = str(safe) + elif isinstance(safe, float) and not math.isfinite(safe): + # Only reachable when masking is disabled; keep NaN/Infinity out of the JSON. + result = str(safe) + elif isinstance(safe, (int, float)): + # Numbers are emitted as raw JSON numbers, so they skip string truncation. + result_size = len(str(safe)) + if not limiter.can_add(result_size): + return None + limiter.add(result_size) + return safe + elif isinstance(safe, str): + result = safe + else: + # `default` is a safety net for anything _mask_value left non-serializable. + result = json.dumps( + safe, + default=lambda o: _safe_repr(o, config), + allow_nan=False, + ) + + return _finalize(result, limiter, config.max_length) + + def _is_simple_type(value): return isinstance(value, (type(None), bool, int, float, str)) diff --git a/posthog/feature_flag_evaluations.py b/posthog/feature_flag_evaluations.py index 8ddd10f49..eb4600174 100644 --- a/posthog/feature_flag_evaluations.py +++ b/posthog/feature_flag_evaluations.py @@ -338,13 +338,25 @@ def _record_access(self, key: str, send_feature_flag_events: bool) -> None: return flag = self._flags.get(key) - if flag is None: - response: Optional[FlagValue] = None - elif not flag.enabled: - response = False - else: - response = flag.variant if flag.variant is not None else True + response = _flag_response(flag) + properties = self._get_flag_called_properties(key, flag, response) + self._host.capture_flag_called_event_if_needed( + distinct_id=self._distinct_id, + key=key, + response=response, + groups=self._groups, + disable_geoip=self._disable_geoip, + properties=properties, + has_experiment=flag.has_experiment if flag else None, + minimal_flag_called_events=self._minimal_flag_called_events, + # The record that supplied the value. A remotely evaluated experiment + # split also sends $experiment_exposure, whatever this read returned. + details=flag.details if isinstance(flag, _RemoteFlagRecord) else None, + value_from_details=True, + ) + + def _get_flag_called_properties(self, key, flag, response) -> Dict[str, Any]: properties: Dict[str, Any] = { "$feature_flag": key, "$feature_flag_response": response, @@ -353,20 +365,20 @@ def _record_access(self, key: str, send_feature_flag_events: bool) -> None: } if flag is not None: - if flag.payload is not None: - properties["$feature_flag_payload"] = flag.payload - if flag.id: - properties["$feature_flag_id"] = flag.id - if flag.version: - properties["$feature_flag_version"] = flag.version - if flag.reason: - properties["$feature_flag_reason"] = flag.reason + _add_flag_metadata_properties(properties, flag) if self._request_id: properties["$feature_flag_request_id"] = self._request_id if self._evaluated_at and not (flag and flag.locally_evaluated): properties["$feature_flag_evaluated_at"] = self._evaluated_at + errors = self._get_access_errors(key, flag) + if errors: + properties["$feature_flag_error"] = ",".join(errors) + + return properties + + def _get_access_errors(self, key, flag) -> List[str]: # Build the comma-joined `$feature_flag_error` matching the single-flag path's # granularity: response-level errors (errors-while-computing, quota-limited) are # combined with per-flag errors (flag-missing or local-evaluation-inconclusive) @@ -382,20 +394,26 @@ def _record_access(self, key: str, send_feature_flag_events: bool) -> None: if key in self._unresolved_flags else "flag_missing" ) - if errors: - properties["$feature_flag_error"] = ",".join(errors) + return errors - self._host.capture_flag_called_event_if_needed( - distinct_id=self._distinct_id, - key=key, - response=response, - groups=self._groups, - disable_geoip=self._disable_geoip, - properties=properties, - has_experiment=flag.has_experiment if flag else None, - minimal_flag_called_events=self._minimal_flag_called_events, - # The record that supplied the value. A remotely evaluated experiment - # split also sends $experiment_exposure, whatever this read returned. - details=flag.details if isinstance(flag, _RemoteFlagRecord) else None, - value_from_details=True, - ) + +def _flag_response(flag: Optional[_EvaluatedFlagRecord]) -> Optional[FlagValue]: + if flag is None: + return None + elif not flag.enabled: + return False + else: + return flag.variant if flag.variant is not None else True + + +def _add_flag_metadata_properties( + properties: Dict[str, Any], flag: _EvaluatedFlagRecord +) -> None: + if flag.payload is not None: + properties["$feature_flag_payload"] = flag.payload + if flag.id: + properties["$feature_flag_id"] = flag.id + if flag.version: + properties["$feature_flag_version"] = flag.version + if flag.reason: + properties["$feature_flag_reason"] = flag.reason diff --git a/posthog/feature_flags.py b/posthog/feature_flags.py index 6634ae367..2a5b955cb 100644 --- a/posthog/feature_flags.py +++ b/posthog/feature_flags.py @@ -292,6 +292,41 @@ def evaluate_flag_dependency( _warn_malformed_flag_dependency_once(flag_key, shape_problem) return False + dependency_chain = _get_dependency_chain(property) + _assert_dependency_definitions_loaded( + dependency_chain, flags_by_key, evaluation_cache + ) + _evaluate_dependency_chain( + dependency_chain, + flags_by_key, + evaluation_cache, + distinct_id, + properties, + cohort_properties, + device_id, + property_matching_version, + ) + + # The condition matches iff the referenced flag's value matches the expected + # value. None means inconclusive or not evaluated — distinct from a + # definitive False, which must be allowed to match `expected_value=False`. + actual_value = evaluation_cache.get(flag_key) + if actual_value is None: + # A referenced flag that has no loaded definition or is missing from the + # chain is a definition problem, not an unresolved dependency. + reason = ( + UnresolvedFlagReason.UNRESOLVED_DEPENDENCY + if flag_key in flags_by_key and flag_key in evaluation_cache + else UnresolvedFlagReason.UNSUPPORTED_DEFINITION + ) + raise InconclusiveMatchError( + f"Flag dependency '{flag_key}' was inconclusive or not evaluated", + reason=reason, + ) + return matches_dependency_value(expected_value, actual_value) + + +def _get_dependency_chain(property): # Check if dependency_chain is present - it should always be provided for flag dependencies if "dependency_chain" not in property: # Missing dependency_chain indicates malformed server data @@ -300,6 +335,7 @@ def evaluate_flag_dependency( reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, ) + flag_key = property.get("key") dependency_chain = property["dependency_chain"] if not isinstance(dependency_chain, list): @@ -316,6 +352,12 @@ def evaluate_flag_dependency( reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, ) + return dependency_chain + + +def _assert_dependency_definitions_loaded( + dependency_chain, flags_by_key, evaluation_cache +): # A dependency with no definition outranks one that is only inconclusive, so # check the whole chain before evaluating any of it. missing_dependency = next( @@ -333,86 +375,100 @@ def evaluate_flag_dependency( reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, ) + +def _evaluate_dependency_chain( + dependency_chain, + flags_by_key, + evaluation_cache, + distinct_id, + properties, + cohort_properties, + device_id, + property_matching_version, +): # Evaluate and cache each flag in the chain; members already cached are # skipped. This does not decide the outcome — it only populates the cache. for dep_flag_key in dependency_chain: if dep_flag_key in evaluation_cache: continue + _evaluate_dependency( + dep_flag_key, + flags_by_key, + evaluation_cache, + distinct_id, + properties, + cohort_properties, + device_id, + property_matching_version, + ) - dep_flag = flags_by_key[dep_flag_key] - # Same order as client._compute_flag_locally: an inactive flag is false - # whatever its experience continuity setting. - if not dep_flag.get("active"): - evaluation_cache[dep_flag_key] = False - continue +def _evaluate_dependency( + dep_flag_key, + flags_by_key, + evaluation_cache, + distinct_id, + properties, + cohort_properties, + device_id, + property_matching_version, +): + dep_flag = flags_by_key[dep_flag_key] - if dep_flag.get("ensure_experience_continuity"): - evaluation_cache[dep_flag_key] = None - raise InconclusiveMatchError( - f"Cannot evaluate flag dependency '{dep_flag_key}' - flag has experience continuity enabled", - reason=UnresolvedFlagReason.UNRESOLVED_DEPENDENCY, - ) + # Same order as client._compute_flag_locally: an inactive flag is false + # whatever its experience continuity setting. + if not dep_flag.get("active"): + evaluation_cache[dep_flag_key] = False + return - # Recursively evaluate the dependency - try: - dep_flag_filters = dep_flag.get("filters") or {} - dep_aggregation_group_type_index = dep_flag_filters.get( - "aggregation_group_type_index" - ) - if dep_aggregation_group_type_index is not None: - # Group flags should continue bucketing by the group key - # from the current evaluation context. - dep_bucketing_value = distinct_id - else: - dep_bucketing_value = resolve_bucketing_value( - dep_flag, distinct_id, device_id - ) - dep_result = match_feature_flag_properties( - dep_flag, - distinct_id, - properties, - cohort_properties=cohort_properties, - flags_by_key=flags_by_key, - evaluation_cache=evaluation_cache, - device_id=device_id, - property_matching_version=property_matching_version, - bucketing_value=dep_bucketing_value, - ) - evaluation_cache[dep_flag_key] = dep_result - except RequiresServerEvaluation as e: - # The dependent flag needs the server too, so its later conditions must - # not resolve it locally. Nothing is cached, so another filter on this - # dependency also raises this error. - raise RequiresServerEvaluation( - f"Cannot evaluate flag dependency '{dep_flag_key}': {e}", - reason=UnresolvedFlagReason.UNRESOLVED_DEPENDENCY, - ) from e - except InconclusiveMatchError as e: - # If we can't evaluate a dependency, store None and propagate the error - evaluation_cache[dep_flag_key] = None - raise InconclusiveMatchError( - f"Cannot evaluate flag dependency '{dep_flag_key}': {e}", - reason=UnresolvedFlagReason.UNRESOLVED_DEPENDENCY, - ) from e + if dep_flag.get("ensure_experience_continuity"): + evaluation_cache[dep_flag_key] = None + raise InconclusiveMatchError( + f"Cannot evaluate flag dependency '{dep_flag_key}' - flag has experience continuity enabled", + reason=UnresolvedFlagReason.UNRESOLVED_DEPENDENCY, + ) - # The condition matches iff the referenced flag's value matches the expected - # value. None means inconclusive or not evaluated — distinct from a - # definitive False, which must be allowed to match `expected_value=False`. - actual_value = evaluation_cache.get(flag_key) - if actual_value is None: - # A referenced flag that has no loaded definition or is missing from the - # chain is a definition problem, not an unresolved dependency. - reason = ( - UnresolvedFlagReason.UNRESOLVED_DEPENDENCY - if flag_key in flags_by_key and flag_key in evaluation_cache - else UnresolvedFlagReason.UNSUPPORTED_DEFINITION + # Recursively evaluate the dependency + try: + dep_flag_filters = dep_flag.get("filters") or {} + dep_aggregation_group_type_index = dep_flag_filters.get( + "aggregation_group_type_index" ) - raise InconclusiveMatchError( - f"Flag dependency '{flag_key}' was inconclusive or not evaluated", - reason=reason, + if dep_aggregation_group_type_index is not None: + # Group flags should continue bucketing by the group key + # from the current evaluation context. + dep_bucketing_value = distinct_id + else: + dep_bucketing_value = resolve_bucketing_value( + dep_flag, distinct_id, device_id + ) + dep_result = match_feature_flag_properties( + dep_flag, + distinct_id, + properties, + cohort_properties=cohort_properties, + flags_by_key=flags_by_key, + evaluation_cache=evaluation_cache, + device_id=device_id, + property_matching_version=property_matching_version, + bucketing_value=dep_bucketing_value, ) - return matches_dependency_value(expected_value, actual_value) + evaluation_cache[dep_flag_key] = dep_result + except RequiresServerEvaluation as e: + # The dependent flag needs the server too, so its later conditions must + # not resolve it locally. Nothing is cached, so another filter on this + # dependency also raises this error. + raise RequiresServerEvaluation( + f"Cannot evaluate flag dependency '{dep_flag_key}': {e}", + reason=UnresolvedFlagReason.UNRESOLVED_DEPENDENCY, + ) from e + except InconclusiveMatchError as e: + # If we can't evaluate a dependency, store None and propagate the error + evaluation_cache[dep_flag_key] = None + raise InconclusiveMatchError( + f"Cannot evaluate flag dependency '{dep_flag_key}': {e}", + reason=UnresolvedFlagReason.UNRESOLVED_DEPENDENCY, + ) from e def matches_dependency_value(expected_value, actual_value): @@ -499,19 +555,7 @@ def match_feature_flag_properties( flag_filters = flag.get("filters") or {} if flag_filters.get("feature_enrollment") is True: - enrollment_key = f"$feature_enrollment/{flag.get('key')}" - if flag_filters.get("aggregation_group_type_index") is not None: - raise InconclusiveMatchError( - "Group flag has feature enrollment, which reads person properties", - reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, - ) - if enrollment_key not in properties: - raise InconclusiveMatchError( - f"Flag has feature enrollment and '{enrollment_key}' was not passed", - reason=UnresolvedFlagReason.MISSING_CONTEXT, - ) - enrollment = properties[enrollment_key] - return enrollment is True or enrollment == "true" + return _match_feature_enrollment(flag, flag_filters, properties) # Holdouts are evaluated before release conditions, so a held-out value is excluded # from the flag's targeting entirely rather than being bucketed into a variant. @@ -519,101 +563,89 @@ def match_feature_flag_properties( if holdout_variant is not None: return holdout_variant + return _match_flag_conditions( + flag, + flag_filters, + distinct_id, + properties, + cohort_properties or {}, + flags_by_key, + evaluation_cache, + device_id, + bucketing_value, + group_type_mapping or {}, + groups or {}, + group_properties or {}, + property_matching_version, + ) + + +def _match_flag_conditions( + flag, + flag_filters, + distinct_id, + properties, + cohort_properties, + flags_by_key, + evaluation_cache, + device_id, + bucketing_value, + group_type_mapping, + groups, + group_properties, + property_matching_version, +) -> FlagValue: flag_conditions = flag_filters.get("groups") or [] flag_aggregation = flag_filters.get("aggregation_group_type_index") early_exit_enabled = flag_filters.get("early_exit") inconclusive_reason: Optional[UnresolvedFlagReason] = None - cohort_properties = cohort_properties or {} - groups = groups or {} - group_properties = group_properties or {} - group_type_mapping = group_type_mapping or {} - # Some filters can be explicitly set to null, which require accessing variants like so - flag_variants = (flag_filters.get("multivariate") or {}).get("variants") or [] - valid_variant_keys = [variant["key"] for variant in flag_variants] + valid_variant_keys = _get_valid_variant_keys(flag_filters) for condition in flag_conditions: try: - # Per-condition aggregation overrides only when the condition explicitly - # sets its own aggregation_group_type_index (mixed targeting). - # When absent, use the properties/bucketing already resolved by the caller. - condition_aggregation = condition.get( - "aggregation_group_type_index", flag_aggregation - ) - - # Mixed-override path: condition-level aggregation differs from flag-level. - # This assumes flag-level aggregation is None for mixed flags. - if condition_aggregation != flag_aggregation: - if condition_aggregation is not None: - group_name = group_type_mapping.get(str(condition_aggregation)) - if not group_name or group_name not in groups: - log.debug( - "Skipping group condition for flag '%s': group type index %s not available", - flag.get("key", ""), - condition_aggregation, - ) - continue - if group_name not in group_properties: - inconclusive_reason = _pick_unresolved_reason( - inconclusive_reason, UnresolvedFlagReason.MISSING_CONTEXT - ) - continue - effective_properties = group_properties[group_name] - effective_bucketing = groups[group_name] - else: - effective_properties = properties - effective_bucketing = bucketing_value - else: - effective_properties = properties - effective_bucketing = bucketing_value - - match_result = is_condition_match( + match_result, value = _match_flag_condition( flag, distinct_id, condition, - effective_properties, + properties, cohort_properties, flags_by_key, evaluation_cache, - bucketing_value=effective_bucketing, - device_id=device_id, - property_matching_version=property_matching_version, + device_id, + bucketing_value, + group_type_mapping, + groups, + group_properties, + property_matching_version, + flag_aggregation, + valid_variant_keys, ) - if match_result == ConditionMatch.MATCH: - variant_override = condition.get("variant") - if variant_override and variant_override in valid_variant_keys: - variant = variant_override - else: - variant = get_matching_variant(flag, effective_bucketing) - return variant or True - elif ( - early_exit_enabled - and match_result == ConditionMatch.OUT_OF_ROLLOUT_BOUND - ): - # The condition's property filters (if any) matched and only the rollout check - # failed, so re-evaluating later groups can't change the outcome. If an earlier - # condition was inconclusive, stop here but preserve that result for fallback. - if inconclusive_reason is not None: - break - return False except RequiresServerEvaluation as e: - # Static cohort or other missing server-side data - must fallback to API. - # The loop stops here, so a filter it did not reach can still carry a - # higher-precedence reason in the flag's definition. - try: - definition_reason = _get_definition_unresolved_reason( - flag, flags_by_key or {}, cohort_properties, group_type_mapping - ) - except Exception: - # A malformed definition must not stop the fallback to the server. - definition_reason = None - if definition_reason is not None: - e.reason = _pick_unresolved_reason(e.reason, definition_reason) + _enrich_server_evaluation_reason( + e, + flag, + flags_by_key, + cohort_properties, + group_type_mapping, + ) raise except InconclusiveMatchError as e: # Evaluation error (bad regex, invalid date, missing property, etc.) # Track that we had an inconclusive match, but try other conditions inconclusive_reason = _pick_unresolved_reason(inconclusive_reason, e.reason) + continue + if match_result == ConditionMatch.MATCH: + return value + elif early_exit_enabled and match_result == ConditionMatch.OUT_OF_ROLLOUT_BOUND: + # The condition's property filters (if any) matched and only the rollout check + # failed, so re-evaluating later groups can't change the outcome. If an earlier + # condition was inconclusive, stop here but preserve that result for fallback. + return _unmatched_flag_result(inconclusive_reason) + + return _unmatched_flag_result(inconclusive_reason) + +def _unmatched_flag_result(inconclusive_reason) -> bool: if inconclusive_reason is not None: raise InconclusiveMatchError( "Can't determine if feature flag is enabled or not with given properties", @@ -625,6 +657,143 @@ def match_feature_flag_properties( return False +def _get_condition_context( + flag, + condition, + flag_aggregation, + properties, + bucketing_value, + group_type_mapping, + groups, + group_properties, +): + # Per-condition aggregation overrides only when the condition explicitly + # sets its own aggregation_group_type_index (mixed targeting). + # When absent, use the properties/bucketing already resolved by the caller. + condition_aggregation = condition.get( + "aggregation_group_type_index", flag_aggregation + ) + + # Mixed-override path: condition-level aggregation differs from flag-level. + # This assumes flag-level aggregation is None for mixed flags. + if condition_aggregation != flag_aggregation: + if condition_aggregation is not None: + group_name = group_type_mapping.get(str(condition_aggregation)) + if not group_name or group_name not in groups: + log.debug( + "Skipping group condition for flag '%s': group type index %s not available", + flag.get("key", ""), + condition_aggregation, + ) + return None + if group_name not in group_properties: + raise InconclusiveMatchError( + "Group properties were not passed", + reason=UnresolvedFlagReason.MISSING_CONTEXT, + ) + effective_properties = group_properties[group_name] + effective_bucketing = groups[group_name] + else: + effective_properties = properties + effective_bucketing = bucketing_value + else: + effective_properties = properties + effective_bucketing = bucketing_value + + return effective_properties, effective_bucketing + + +def _match_flag_condition( + flag, + distinct_id, + condition, + properties, + cohort_properties, + flags_by_key, + evaluation_cache, + device_id, + bucketing_value, + group_type_mapping, + groups, + group_properties, + property_matching_version, + flag_aggregation, + valid_variant_keys, +): + context = _get_condition_context( + flag, + condition, + flag_aggregation, + properties, + bucketing_value, + group_type_mapping, + groups, + group_properties, + ) + if context is None: + return ConditionMatch.NO_MATCH, None + effective_properties, effective_bucketing = context + match_result = is_condition_match( + flag, + distinct_id, + condition, + effective_properties, + cohort_properties, + flags_by_key, + evaluation_cache, + bucketing_value=effective_bucketing, + device_id=device_id, + property_matching_version=property_matching_version, + ) + if match_result == ConditionMatch.MATCH: + variant_override = condition.get("variant") + if variant_override and variant_override in valid_variant_keys: + variant = variant_override + else: + variant = get_matching_variant(flag, effective_bucketing) + return match_result, variant or True + return match_result, None + + +def _enrich_server_evaluation_reason( + e, flag, flags_by_key, cohort_properties, group_type_mapping +): + # Static cohort or other missing server-side data - must fallback to API. + # The loop stops here, so a filter it did not reach can still carry a + # higher-precedence reason in the flag's definition. + try: + definition_reason = _get_definition_unresolved_reason( + flag, flags_by_key or {}, cohort_properties, group_type_mapping + ) + except Exception: + # A malformed definition must not stop the fallback to the server. + definition_reason = None + if definition_reason is not None: + e.reason = _pick_unresolved_reason(e.reason, definition_reason) + + +def _get_valid_variant_keys(flag_filters): + # Some filters can be explicitly set to null, which require accessing variants like so + flag_variants = (flag_filters.get("multivariate") or {}).get("variants") or [] + return [variant["key"] for variant in flag_variants] + + +def _match_feature_enrollment(flag, flag_filters, properties) -> bool: + enrollment_key = f"$feature_enrollment/{flag.get('key')}" + if flag_filters.get("aggregation_group_type_index") is not None: + raise InconclusiveMatchError( + "Group flag has feature enrollment, which reads person properties", + reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, + ) + if enrollment_key not in properties: + raise InconclusiveMatchError( + f"Flag has feature enrollment and '{enrollment_key}' was not passed", + reason=UnresolvedFlagReason.MISSING_CONTEXT, + ) + enrollment = properties[enrollment_key] + return enrollment is True or enrollment == "true" + + def _match_condition_filter( prop, properties, @@ -677,57 +846,18 @@ def is_condition_match( ) -> ConditionMatch: rollout_percentage = condition.get("rollout_percentage") if len(condition.get("properties") or []) > 0: - # Filters in a condition have no defined order, so an inconclusive filter does - # not stop the loop: a later filter can still rule the condition out or carry a - # higher-precedence reason. - inconclusive_reason: Optional[UnresolvedFlagReason] = None - server_evaluation_error: Optional[RequiresServerEvaluation] = None - for prop in condition.get("properties"): - try: - matches = _match_condition_filter( - prop, - properties, - cohort_properties, - flags_by_key, - evaluation_cache, - distinct_id, - device_id=device_id, - property_matching_version=property_matching_version, - ) - except RequiresServerEvaluation as e: - # Fall back to the server only when this is the first filter that could - # not be checked, as when the loop stopped there. After another - # inconclusive filter, later conditions must still be able to resolve - # the flag locally. - if inconclusive_reason is None: - server_evaluation_error = e - inconclusive_reason = _pick_unresolved_reason( - inconclusive_reason, e.reason - ) - continue - except InconclusiveMatchError as e: - inconclusive_reason = _pick_unresolved_reason( - inconclusive_reason, e.reason - ) - continue - except Exception: - # After an inconclusive filter, an unexpected error ends the condition - # as inconclusive, so later conditions can still resolve the flag. - if inconclusive_reason is None: - raise - break - if not matches: - return ConditionMatch.NO_MATCH - - # Re-raise the original error type: match_feature_flag_properties stops at it - # and falls back to the server instead of trying later conditions. - if server_evaluation_error is not None: - raise server_evaluation_error - if inconclusive_reason is not None: - raise InconclusiveMatchError( - f"Can't determine if a condition of flag {feature_flag.get('key')} matches", - reason=inconclusive_reason, - ) + if not _match_condition_properties( + feature_flag, + distinct_id, + condition, + properties, + cohort_properties, + flags_by_key, + evaluation_cache, + device_id, + property_matching_version, + ): + return ConditionMatch.NO_MATCH if rollout_percentage is None: return ConditionMatch.MATCH @@ -742,6 +872,71 @@ def is_condition_match( return ConditionMatch.MATCH +def _match_condition_properties( + feature_flag, + distinct_id, + condition, + properties, + cohort_properties, + flags_by_key, + evaluation_cache, + device_id, + property_matching_version, +) -> bool: + # Filters in a condition have no defined order, so an inconclusive filter does + # not stop the loop: a later filter can still rule the condition out or carry a + # higher-precedence reason. + inconclusive_reason: Optional[UnresolvedFlagReason] = None + server_evaluation_error: Optional[RequiresServerEvaluation] = None + for prop in condition.get("properties"): + try: + matches = _match_condition_filter( + prop, + properties, + cohort_properties, + flags_by_key, + evaluation_cache, + distinct_id, + device_id=device_id, + property_matching_version=property_matching_version, + ) + except RequiresServerEvaluation as e: + # Fall back to the server only when this is the first filter that could + # not be checked, as when the loop stopped there. After another + # inconclusive filter, later conditions must still be able to resolve + # the flag locally. + if inconclusive_reason is None: + server_evaluation_error = e + inconclusive_reason = _pick_unresolved_reason(inconclusive_reason, e.reason) + continue + except InconclusiveMatchError as e: + inconclusive_reason = _pick_unresolved_reason(inconclusive_reason, e.reason) + continue + except Exception: + # After an inconclusive filter, an unexpected error ends the condition + # as inconclusive, so later conditions can still resolve the flag. + if inconclusive_reason is None: + raise + break + if not matches: + return False + + _raise_condition_errors(feature_flag, server_evaluation_error, inconclusive_reason) + return True + + +def _raise_condition_errors(feature_flag, server_evaluation_error, inconclusive_reason): + # Re-raise the original error type: match_feature_flag_properties stops at it + # and falls back to the server instead of trying later conditions. + if server_evaluation_error is not None: + raise server_evaluation_error + if inconclusive_reason is not None: + raise InconclusiveMatchError( + f"Can't determine if a condition of flag {feature_flag.get('key')} matches", + reason=inconclusive_reason, + ) + + # Raised when an operator passes the PROPERTY_OPERATORS gate but has no dispatch # branch in match_property. Distinct from the unknown-operator rejection at the top # of the function so the dispatch-completeness test can tell the two apart. @@ -775,6 +970,17 @@ def _format_json_float(value: float, reason: UnresolvedFlagReason) -> str: def _json_value_to_string(value, reason: UnresolvedFlagReason) -> str: + if isinstance(value, list): + return ( + "[" + ",".join(_json_value_to_string(item, reason) for item in value) + "]" + ) + if isinstance(value, dict): + return _json_object_to_string(value, reason) + + return _json_scalar_to_string(value, reason) + + +def _json_scalar_to_string(value, reason: UnresolvedFlagReason) -> str: if value is None: return "null" if isinstance(value, bool): @@ -785,28 +991,25 @@ def _json_value_to_string(value, reason: UnresolvedFlagReason) -> str: return str(value) if isinstance(value, float): return _format_json_float(value, reason) - if isinstance(value, list): - return ( - "[" + ",".join(_json_value_to_string(item, reason) for item in value) + "]" - ) - if isinstance(value, dict): - if not all(isinstance(key, str) for key in value): - raise InconclusiveMatchError( - "Property object keys must be strings for local evaluation", - reason=reason, - ) - items = ( - f"{json.dumps(key, ensure_ascii=False)}:{_json_value_to_string(value[key], reason)}" - for key in sorted(value) - ) - return "{" + ",".join(items) + "}" - raise InconclusiveMatchError( f"Property value of type {type(value).__name__} is not JSON-compatible", reason=reason, ) +def _json_object_to_string(value: dict, reason: UnresolvedFlagReason) -> str: + if not all(isinstance(key, str) for key in value): + raise InconclusiveMatchError( + "Property object keys must be strings for local evaluation", + reason=reason, + ) + items = ( + f"{json.dumps(key, ensure_ascii=False)}:{_json_value_to_string(value[key], reason)}" + for key in sorted(value) + ) + return "{" + ",".join(items) + "}" + + def _value_to_string(value, reason: UnresolvedFlagReason) -> str: if isinstance(value, str): return value @@ -862,159 +1065,202 @@ class _RegexMetadata(NamedTuple): uses_quantifier_or_alternation: bool -@_lru_cache(maxsize=256) -def _compile_regex(pattern: str) -> Optional[_RegexMetadata]: - """Cache only pattern metadata; exceptions are not retained by lru_cache.""" - # Successful re compilation alone does not establish Rust/fancy-regex compatibility. - translated = [] - class_start = None - previous = "" - group_has_quantifier = [] - closed_group_has_quantifier = False - uses_word_class = False - uses_space_class = False - uses_end_anchor = False - uses_quantifier_or_alternation = False - index = 0 - - def defer() -> NoReturn: - raise InconclusiveMatchError( - "Regex requires server evaluation", - reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, - ) +def _defer_regex() -> NoReturn: + raise InconclusiveMatchError( + "Regex requires server evaluation", + reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, + ) - while index < len(pattern): - character = pattern[index] - if 0xD800 <= ord(character) <= 0xDFFF: - defer() - if class_start is None: - uses_end_anchor |= character == "$" - uses_quantifier_or_alternation |= character in "*+?{|" - if character == "(": - group_has_quantifier.append(False) - elif character == ")": - closed_group_has_quantifier = ( - group_has_quantifier.pop() if group_has_quantifier else False - ) - if group_has_quantifier: - group_has_quantifier[-1] |= closed_group_has_quantifier - elif character in "*+?{": - # Strict anchors can turn nested repetition into expensive backtracking. - if previous == ")" and closed_group_has_quantifier: - defer() - if group_has_quantifier: - group_has_quantifier[-1] = True + +class _RegexTranslator: + """Scan Python/Rust shared syntax, retaining the metadata needed at match time.""" + + def __init__(self, pattern: str) -> None: + self.pattern = pattern + self.translated: list[str] = [] + self.class_start: Optional[int] = None + self.previous = "" + self.group_has_quantifier: list[bool] = [] + self.closed_group_has_quantifier = False + self.uses_word_class = False + self.uses_space_class = False + self.uses_end_anchor = False + self.uses_quantifier_or_alternation = False + self.index = 0 + + def translate(self) -> str: + while self.index < len(self.pattern): + character = self.pattern[self.index] + if 0xD800 <= ord(character) <= 0xDFFF: + _defer_regex() + if self.class_start is None: + self._scan_group(character) + self._translate_character(character) + self.index += 1 + return "".join(self.translated) + + def _scan_group(self, character: str) -> None: + self.uses_end_anchor |= character == "$" + self.uses_quantifier_or_alternation |= character in "*+?{|" + if character == "(": + self.group_has_quantifier.append(False) + elif character == ")": + self.closed_group_has_quantifier = ( + self.group_has_quantifier.pop() if self.group_has_quantifier else False + ) + if self.group_has_quantifier: + self.group_has_quantifier[-1] |= self.closed_group_has_quantifier + elif character in "*+?{": + # Strict anchors can turn nested repetition into expensive backtracking. + if self.previous == ")" and self.closed_group_has_quantifier: + _defer_regex() + if self.group_has_quantifier: + self.group_has_quantifier[-1] = True + + def _translate_character(self, character: str) -> None: if character == "\\": - index += 1 - if index == len(pattern): - translated.append("\\") - break - escaped = pattern[index] - if escaped == "x": - digits = pattern[index + 1 : index + 3] - if len(digits) != 2 or any( - c not in "0123456789abcdefABCDEF" for c in digits - ): - defer() - translated.append("\\x" + digits) - index += 3 - previous = "literal" - continue - if escaped in "bB<>" and class_start is None: - # Even short ASCII word-boundary searches can exceed the service's step limit. - defer() - if escaped not in "dDsSwWbfnrt": - if not (0x20 <= ord(escaped) <= 0x7E and not escaped.isalnum()): - defer() - if class_start is not None and escaped not in r"\-[]^.$*+?(){}|": - defer() - uses_word_class |= escaped in "wW" - uses_space_class |= escaped in "sS" - translated.append("\\" + escaped) - previous = "literal" + self._translate_escape() elif character == "[": - if class_start is not None: - defer() - class_start = index - translated.append(character) - elif character == "]" and class_start is not None: - content_start = class_start + ( - 2 if pattern[class_start + 1 : class_start + 2] == "^" else 1 - ) - if index == content_start: - defer() - class_start = None - previous = "literal" - translated.append(character) - elif class_start is not None: - if pattern[index : index + 2] in ("&&", "--", "~~"): - defer() - # Rust and Python disagree on shorthand endpoints in ranges. - if ( - character == "-" - and index + 1 < len(pattern) - and pattern[index + 1] != "]" - ): - if pattern[index + 1 : index + 3] in ( - r"\d", - r"\D", - r"\s", - r"\S", - r"\w", - r"\W", - ) or ( - index >= 2 - and pattern[index - 2] == "\\" - and pattern[index - 1] in "dDsSwW" - ): - defer() - translated.append(character) - elif character == "(" and pattern[index + 1 : index + 2] == "?": - if pattern[index + 2 : index + 3] != ":": - defer() - translated.append("(?:") - index += 2 - previous = "(" + if self.class_start is not None: + _defer_regex() + self.class_start = self.index + self.translated.append(character) + elif character == "]" and self.class_start is not None: + self._close_class() + elif self.class_start is not None: + self._translate_class_character(character) + else: + self._translate_outside_class(character) + + def _translate_escape(self) -> None: + self.index += 1 + if self.index == len(self.pattern): + self.translated.append("\\") + return + escaped = self.pattern[self.index] + if escaped == "x": + self._translate_hex_escape() + return + if escaped in "bB<>" and self.class_start is None: + # Even short ASCII word-boundary searches can exceed the service's step limit. + _defer_regex() + self._assert_escape_supported(escaped) + self.uses_word_class |= escaped in "wW" + self.uses_space_class |= escaped in "sS" + self.translated.append("\\" + escaped) + self.previous = "literal" + + def _translate_hex_escape(self) -> None: + digits = self.pattern[self.index + 1 : self.index + 3] + if len(digits) != 2 or any(c not in "0123456789abcdefABCDEF" for c in digits): + _defer_regex() + self.translated.append("\\x" + digits) + # The outer scan advances the last character of the escape. + self.index += 2 + self.previous = "literal" + + def _assert_escape_supported(self, escaped: str) -> None: + if escaped not in "dDsSwWbfnrt": + if not (0x20 <= ord(escaped) <= 0x7E and not escaped.isalnum()): + _defer_regex() + if self.class_start is not None and escaped not in r"\-[]^.$*+?(){}|": + _defer_regex() + + def _close_class(self) -> None: + assert self.class_start is not None + content_start = self.class_start + ( + 2 if self.pattern[self.class_start + 1 : self.class_start + 2] == "^" else 1 + ) + if self.index == content_start: + _defer_regex() + self.class_start = None + self.previous = "literal" + self.translated.append("]") + + def _translate_class_character(self, character: str) -> None: + if self.pattern[self.index : self.index + 2] in ("&&", "--", "~~"): + _defer_regex() + # Rust and Python disagree on shorthand endpoints in ranges. + if ( + character == "-" + and self.index + 1 < len(self.pattern) + and self.pattern[self.index + 1] != "]" + ): + self._assert_range_supported() + self.translated.append(character) + + def _assert_range_supported(self) -> None: + if self.pattern[self.index + 1 : self.index + 3] in ( + r"\d", + r"\D", + r"\s", + r"\S", + r"\w", + r"\W", + ) or ( + self.index >= 2 + and self.pattern[self.index - 2] == "\\" + and self.pattern[self.index - 1] in "dDsSwW" + ): + _defer_regex() + + def _translate_outside_class(self, character: str) -> None: + if character == "(" and self.pattern[self.index + 1 : self.index + 2] == "?": + if self.pattern[self.index + 2 : self.index + 3] != ":": + _defer_regex() + self.translated.append("(?:") + self.index += 2 + self.previous = "(" elif character == "{": - repetition = re.match(r"\{([0-9]+)(?:,([0-9]*))?\}", pattern[index:]) - if repetition is None or previous in ( - "", - "|", - "(", - "^", - "$", - "empty group", - "quantifier", - ): - defer() - if any( - len(n) > 10 or int(n) > 0xFFFFFFFF for n in repetition.groups() if n - ): - defer() - translated.append(repetition.group()) - index += len(repetition.group()) - 1 - if pattern[index + 1 : index + 2] in ("+", "{"): - defer() - previous = "quantifier" + self._translate_repetition() else: - if character in "*+?": - if previous in ("^", "$", "empty group") or pattern[ - index + 1 : index + 2 - ] in ("+", "{"): - defer() - previous = "quantifier" - else: - previous = ( - "empty group" - if character == ")" and previous in ("(", "empty group") - else character - ) - # Python's $ also matches before a final newline; Rust's default $ is strict. - translated.append(r"\Z" if character == "$" else character) - index += 1 + self._translate_literal_or_quantifier(character) + + def _translate_repetition(self) -> None: + repetition = re.match(r"\{([0-9]+)(?:,([0-9]*))?\}", self.pattern[self.index :]) + if repetition is None or self.previous in ( + "", + "|", + "(", + "^", + "$", + "empty group", + "quantifier", + ): + _defer_regex() + if any(len(n) > 10 or int(n) > 0xFFFFFFFF for n in repetition.groups() if n): + _defer_regex() + self.translated.append(repetition.group()) + self.index += len(repetition.group()) - 1 + if self.pattern[self.index + 1 : self.index + 2] in ("+", "{"): + _defer_regex() + self.previous = "quantifier" + + def _translate_literal_or_quantifier(self, character: str) -> None: + if character in "*+?": + if self.previous in ("^", "$", "empty group") or self.pattern[ + self.index + 1 : self.index + 2 + ] in ("+", "{"): + _defer_regex() + self.previous = "quantifier" + else: + self.previous = ( + "empty group" + if character == ")" and self.previous in ("(", "empty group") + else character + ) + # Python's $ also matches before a final newline; Rust's default $ is strict. + self.translated.append(r"\Z" if character == "$" else character) + +@_lru_cache(maxsize=256) +def _compile_regex(pattern: str) -> Optional[_RegexMetadata]: + """Cache only pattern metadata; exceptions are not retained by lru_cache.""" + # Successful re compilation alone does not establish Rust/fancy-regex compatibility. + translator = _RegexTranslator(pattern) + translated = translator.translate() try: - regex = re.compile("".join(translated)) + regex = re.compile(translated) except re.error: return None except Exception as error: @@ -1025,10 +1271,10 @@ def defer() -> NoReturn: return _RegexMetadata( regex, - uses_word_class, - uses_space_class, - uses_end_anchor, - uses_quantifier_or_alternation, + translator.uses_word_class, + translator.uses_space_class, + translator.uses_end_anchor, + translator.uses_quantifier_or_alternation, ) @@ -1039,42 +1285,42 @@ def _match_regex(pattern: str, value: str) -> Optional[bool]: metadata = _compile_regex(pattern) if metadata is None: return None - ( - regex, - uses_word_class, - uses_space_class, - uses_end_anchor, - uses_quantifier_or_alternation, - ) = metadata + regex = metadata.regex - def defer() -> NoReturn: + _assert_regex_value_supported(metadata, value) + try: + return regex.search(value) is not None + except Exception as error: raise InconclusiveMatchError( - "Regex requires server evaluation", + "Regex execution requires server evaluation", reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, - ) + ) from error + +def _assert_regex_value_supported(metadata: _RegexMetadata, value: str) -> None: # Tightening $ can expose backtracking on a final LF, including repeated alternatives. - if uses_end_anchor and uses_quantifier_or_alternation and value.endswith("\n"): - defer() + if ( + metadata.uses_end_anchor + and metadata.uses_quantifier_or_alternation + and value.endswith("\n") + ): + _defer_regex() # Unicode letters/decimal digits agree, but Rust's word class also includes marks and # joiners, while Python includes other numbers. Python's whitespace adds U+001C..U+001F. if not value.isascii(): - if any(0xD800 <= ord(c) <= 0xDFFF for c in value): - defer() - if uses_word_class and any( - not c.isascii() and not (c.isalpha() or c.isdecimal()) for c in value - ): - defer() - if uses_space_class and any(c in value for c in "\x1c\x1d\x1e\x1f"): - defer() - try: - return regex.search(value) is not None - except Exception as error: - raise InconclusiveMatchError( - "Regex execution requires server evaluation", - reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, - ) from error + _assert_regex_unicode_supported(metadata, value) + if metadata.uses_space_class and any(c in value for c in "\x1c\x1d\x1e\x1f"): + _defer_regex() + + +def _assert_regex_unicode_supported(metadata: _RegexMetadata, value: str) -> None: + if any(0xD800 <= ord(c) <= 0xDFFF for c in value): + _defer_regex() + if metadata.uses_word_class and any( + not c.isascii() and not (c.isalpha() or c.isdecimal()) for c in value + ): + _defer_regex() _TEXT_OPERATORS = ( @@ -1159,7 +1405,12 @@ def _assert_filter_supported(property) -> None: _value_to_string(item, UnresolvedFlagReason.UNSUPPORTED_DEFINITION) elif operator in _TEXT_OPERATORS: _value_to_string(value, UnresolvedFlagReason.UNSUPPORTED_DEFINITION) - elif operator in ("regex", "not_regex"): + else: + _assert_parsed_filter_supported(operator, value) + + +def _assert_parsed_filter_supported(operator, value) -> None: + if operator in ("regex", "not_regex"): pattern = str(value) if pattern: _compile_regex(pattern) @@ -1212,148 +1463,166 @@ def _match_property_value( override_value = property_values[key] if operator in NUMERIC_OPERATORS: - parsed_value = _parse_numeric_filter(value) - parsed_override = _parse_numeric_value(override_value) - if parsed_override is None: - return False - if operator == "gt": - return parsed_override > parsed_value - if operator == "gte": - return parsed_override >= parsed_value - if operator == "lt": - return parsed_override < parsed_value - return parsed_override <= parsed_value + return _match_numeric_property(operator, value, override_value) - if ( + if _is_unmatchable_null(operator, override_value): + return False + + if operator in ("exact", "is_not"): + matched = _match_exact_property( + value, override_value, property_matching_version + ) + return matched if operator == "exact" else not matched + + return _match_scalar_property(operator, value, override_value) + + +def _is_unmatchable_null(operator, override_value) -> bool: + return ( operator not in NONE_VALUES_ALLOWED_OPERATORS and operator not in STRING_OPERATORS and operator != "exact" and override_value is None - ): - return False + ) - if operator in ("exact", "is_not"): - def compute_exact_match(value, override_value): - override_string = _value_to_string( - override_value, UnresolvedFlagReason.MISSING_CONTEXT +def _match_numeric_property(operator, value, override_value) -> bool: + parsed_value = _parse_numeric_filter(value) + parsed_override = _parse_numeric_value(override_value) + if parsed_override is None: + return False + if operator == "gt": + return parsed_override > parsed_value + if operator == "gte": + return parsed_override >= parsed_value + if operator == "lt": + return parsed_override < parsed_value + return parsed_override <= parsed_value + + +def _match_exact_property(value, override_value, property_matching_version) -> bool: + override_string = _value_to_string( + override_value, UnresolvedFlagReason.MISSING_CONTEXT + ).lower() + # Empty filters retain recursive truthiness in both matching modes. + if value == []: + return _is_truthy_property_value(override_value) + if property_matching_version != 2 and _is_truthy_or_falsy_property_value(value): + return _is_truthy_property_value(value) == _is_truthy_property_value( + override_value + ) + + if isinstance(value, list): + return override_string in [ + _value_to_string( + candidate, UnresolvedFlagReason.UNSUPPORTED_DEFINITION ).lower() - # Empty filters retain recursive truthiness in both matching modes. - if value == []: - return _is_truthy_property_value(override_value) - if property_matching_version != 2 and _is_truthy_or_falsy_property_value( - value - ): - return _is_truthy_property_value(value) == _is_truthy_property_value( - override_value - ) + for candidate in value + ] + return ( + _value_to_string(value, UnresolvedFlagReason.UNSUPPORTED_DEFINITION).lower() + == override_string + ) - if isinstance(value, list): - return override_string in [ - _value_to_string( - candidate, UnresolvedFlagReason.UNSUPPORTED_DEFINITION - ).lower() - for candidate in value - ] - return ( - _value_to_string( - value, UnresolvedFlagReason.UNSUPPORTED_DEFINITION - ).lower() - == override_string - ) - if operator == "exact": - return compute_exact_match(value, override_value) - else: - return not compute_exact_match(value, override_value) +def _match_text_property(operator, value, override_value) -> bool: + property_string = _ascii_lower(override_value, UnresolvedFlagReason.MISSING_CONTEXT) + filter_string = _ascii_lower(value, UnresolvedFlagReason.UNSUPPORTED_DEFINITION) - if operator in _TEXT_OPERATORS: - property_string = _ascii_lower( - override_value, UnresolvedFlagReason.MISSING_CONTEXT - ) - filter_string = _ascii_lower(value, UnresolvedFlagReason.UNSUPPORTED_DEFINITION) + if operator in ("icontains", "not_icontains"): + matched = filter_string in property_string + elif operator in ("starts_with", "not_starts_with"): + matched = property_string.startswith(filter_string) + else: + matched = property_string.endswith(filter_string) - if operator in ("icontains", "not_icontains"): - matched = filter_string in property_string - elif operator in ("starts_with", "not_starts_with"): - matched = property_string.startswith(filter_string) - else: - matched = property_string.endswith(filter_string) + return not matched if operator.startswith("not_") else matched - if operator.startswith("not_"): - return not matched - return matched - if operator in ("regex", "not_regex"): - property_string = _value_to_string( - override_value, UnresolvedFlagReason.MISSING_CONTEXT - ) - regex_match = _match_regex(str(value), property_string) - if regex_match is None: - return False - return not regex_match if operator == "not_regex" else regex_match +def _match_regex_property(operator, value, override_value) -> bool: + property_string = _value_to_string( + override_value, UnresolvedFlagReason.MISSING_CONTEXT + ) + regex_match = _match_regex(str(value), property_string) + if regex_match is None: + return False + return not regex_match if operator == "not_regex" else regex_match - if operator in DATE_OPERATORS: - parsed_date = _parse_date_filter(value) - - if isinstance(override_value, datetime.datetime): - override_date = convert_to_datetime_aware(override_value) - if operator == "is_date_before": - return override_date < parsed_date - else: - return override_date > parsed_date - elif isinstance(override_value, datetime.date): - if operator == "is_date_before": - return override_value < parsed_date.date() - else: - return override_value > parsed_date.date() - elif isinstance(override_value, str): - try: - override_date = parse_datetime(override_value) - override_date = convert_to_datetime_aware(override_date) - if operator == "is_date_before": - return override_date < parsed_date - else: - return override_date > parsed_date - except Exception: - raise InconclusiveMatchError( - "The date provided is not a valid format", - reason=UnresolvedFlagReason.MISSING_CONTEXT, - ) - else: - raise InconclusiveMatchError( - "The date provided must be a string or date object", - reason=UnresolvedFlagReason.MISSING_CONTEXT, - ) - if operator in SEMVER_OPERATORS: - flag_semver = _parse_semver_filter(operator, value) +def _match_date_property(operator, value, override_value) -> bool: + parsed_date = _parse_date_filter(value) + + if isinstance(override_value, datetime.datetime): + override_date = convert_to_datetime_aware(override_value) + return _compare_dates(operator, override_date, parsed_date) + elif isinstance(override_value, datetime.date): + return _compare_dates(operator, override_value, parsed_date.date()) + elif isinstance(override_value, str): try: - override_parsed = parse_semver(override_value) - except (ValueError, TypeError): + override_date = parse_datetime(override_value) + override_date = convert_to_datetime_aware(override_date) + return _compare_dates(operator, override_date, parsed_date) + except Exception: raise InconclusiveMatchError( - f"Person property value '{override_value}' is not a valid semver", + "The date provided is not a valid format", reason=UnresolvedFlagReason.MISSING_CONTEXT, ) + else: + raise InconclusiveMatchError( + "The date provided must be a string or date object", + reason=UnresolvedFlagReason.MISSING_CONTEXT, + ) + + +def _compare_dates(operator, override_date, parsed_date) -> bool: + if operator == "is_date_before": + return override_date < parsed_date + return override_date > parsed_date + - if operator in SEMVER_RANGE_OPERATORS: - lower, upper = flag_semver - return lower <= override_parsed < upper - if operator == "semver_eq": - return override_parsed == flag_semver - if operator == "semver_neq": - return override_parsed != flag_semver - if operator == "semver_gt": - return override_parsed > flag_semver - if operator == "semver_gte": - return override_parsed >= flag_semver - if operator == "semver_lt": - return override_parsed < flag_semver - if operator == "semver_lte": - return override_parsed <= flag_semver +def _match_semver_property(operator, value, override_value) -> bool: + flag_semver = _parse_semver_filter(operator, value) + try: + override_parsed = parse_semver(override_value) + except (ValueError, TypeError): + raise InconclusiveMatchError( + f"Person property value '{override_value}' is not a valid semver", + reason=UnresolvedFlagReason.MISSING_CONTEXT, + ) + if operator in SEMVER_RANGE_OPERATORS: + lower, upper = flag_semver + return lower <= override_parsed < upper + if operator == "semver_eq": + return override_parsed == flag_semver + if operator == "semver_neq": + return override_parsed != flag_semver + if operator == "semver_gt": + return override_parsed > flag_semver + if operator == "semver_gte": + return override_parsed >= flag_semver + if operator == "semver_lt": + return override_parsed < flag_semver + if operator == "semver_lte": + return override_parsed <= flag_semver + return _raise_unhandled_operator(operator) + + +def _match_scalar_property(operator, value, override_value) -> bool: + if operator in _TEXT_OPERATORS: + return _match_text_property(operator, value, override_value) + if operator in ("regex", "not_regex"): + return _match_regex_property(operator, value, override_value) + if operator in DATE_OPERATORS: + return _match_date_property(operator, value, override_value) + if operator in SEMVER_OPERATORS: + return _match_semver_property(operator, value, override_value) + return _raise_unhandled_operator(operator) + + +def _raise_unhandled_operator(operator) -> NoReturn: # Unreachable: all operators in PROPERTY_OPERATORS are handled above, - # and unknown operators are rejected at the top of this function. + # and unknown operators are rejected at the top of _match_property_value. raise InconclusiveMatchError( f"Operator {operator} {_UNHANDLED_OPERATOR_MESSAGE}", reason=UnresolvedFlagReason.UNSUPPORTED_DEFINITION, @@ -1426,19 +1695,7 @@ def match_property_group( device_id=None, property_matching_version: int = 1, ) -> bool: - # The backend serializes its canonical empty PropertyGroup as {}. - if property_group == {}: - return True - if not isinstance(property_group, dict): - raise RequiresServerEvaluation("Cohort property group must be an object") - - property_group_type = property_group.get("type") - if property_group_type not in ("AND", "OR"): - raise RequiresServerEvaluation("Cohort property group type must be AND or OR") - is_and = property_group_type == "AND" - properties = property_group.get("values") - if not isinstance(properties, list): - raise RequiresServerEvaluation("Cohort property group values must be a list") + is_and, properties = _get_property_group_values(property_group) if not properties: return True @@ -1447,58 +1704,18 @@ def match_property_group( for prop in properties: try: - if not isinstance(prop, dict): - raise RequiresServerEvaluation( - "Cohort property group entry must be an object" - ) - - if _is_nested_property_group(prop): - matches = match_property_group( - prop, - property_values, - cohort_properties, - flags_by_key, - evaluation_cache, - distinct_id, - device_id=device_id, - property_matching_version=property_matching_version, - ) - negation = False - elif prop.get("type") == "cohort": - matches = match_cohort( - prop, - property_values, - cohort_properties, - flags_by_key, - evaluation_cache, - distinct_id, - device_id=device_id, - property_matching_version=property_matching_version, - ) - negation = prop.get("negation", False) - elif prop.get("type") == "flag": - matches = evaluate_flag_dependency( - prop, - flags_by_key, - evaluation_cache, - distinct_id, - property_values, - cohort_properties, - device_id=device_id, - property_matching_version=property_matching_version, - ) - negation = prop.get("negation", False) - else: - matches = match_property( - prop, property_values, property_matching_version - ) - negation = prop.get("negation", False) - - effective_match = matches != bool(negation) - if is_and and not effective_match: - decisive_result = False - elif not is_and and effective_match: - decisive_result = True + effective_match = _match_property_group_entry( + prop, + property_values, + cohort_properties, + flags_by_key, + evaluation_cache, + distinct_id, + device_id, + property_matching_version, + ) + if effective_match != is_and: + decisive_result = effective_match except RequiresServerEvaluation: # Static/missing cohorts and malformed definitions always require server evaluation, # even when another branch would otherwise resolve the group locally. @@ -1519,6 +1736,62 @@ def match_property_group( return is_and +def _get_property_group_values(property_group): + # The backend serializes its canonical empty PropertyGroup as {}. + if property_group == {}: + return True, [] + if not isinstance(property_group, dict): + raise RequiresServerEvaluation("Cohort property group must be an object") + + property_group_type = property_group.get("type") + if property_group_type not in ("AND", "OR"): + raise RequiresServerEvaluation("Cohort property group type must be AND or OR") + is_and = property_group_type == "AND" + properties = property_group.get("values") + if not isinstance(properties, list): + raise RequiresServerEvaluation("Cohort property group values must be a list") + return is_and, properties + + +def _match_property_group_entry( + prop, + property_values, + cohort_properties, + flags_by_key, + evaluation_cache, + distinct_id, + device_id, + property_matching_version, +): + if not isinstance(prop, dict): + raise RequiresServerEvaluation("Cohort property group entry must be an object") + if _is_nested_property_group(prop): + matches = match_property_group( + prop, + property_values, + cohort_properties, + flags_by_key, + evaluation_cache, + distinct_id, + device_id=device_id, + property_matching_version=property_matching_version, + ) + negation = False + else: + matches = _match_condition_filter( + prop, + property_values, + cohort_properties, + flags_by_key, + evaluation_cache, + distinct_id, + device_id=device_id, + property_matching_version=property_matching_version, + ) + negation = prop.get("negation", False) + return matches != bool(negation) + + def _get_definition_unresolved_reason( flag, flags_by_key, cohort_properties, group_type_mapping ) -> Optional[UnresolvedFlagReason]: @@ -1533,11 +1806,17 @@ def _get_definition_unresolved_reason( aggregation = flag_filters.get("aggregation_group_type_index") if aggregation is not None and not group_type_mapping.get(str(aggregation)): return UnresolvedFlagReason.UNSUPPORTED_DEFINITION + if _has_unsupported_condition(flag_filters, flags_by_key, cohort_properties): + return UnresolvedFlagReason.UNSUPPORTED_DEFINITION + return None + + +def _has_unsupported_condition(flag_filters, flags_by_key, cohort_properties) -> bool: for condition in flag_filters.get("groups") or []: for prop in condition.get("properties") or []: if _is_unsupported_property(prop, flags_by_key, cohort_properties, set()): - return UnresolvedFlagReason.UNSUPPORTED_DEFINITION - return None + return True + return False def _is_unsupported_property( @@ -1545,13 +1824,7 @@ def _is_unsupported_property( ) -> bool: property_type = prop.get("type") if property_type == "flag": - # A malformed condition shape is a definitive local no-match, not a gap. - if _flag_dependency_shape_problem(prop) is not None: - return False - dependency_chain = prop.get("dependency_chain") - if not isinstance(dependency_chain, list) or not dependency_chain: - return True - return any(not flags_by_key.get(key) for key in dependency_chain) + return _is_unsupported_dependency(prop, flags_by_key) if property_type != "cohort": try: _assert_filter_supported(prop) @@ -1574,6 +1847,16 @@ def _is_unsupported_property( ) +def _is_unsupported_dependency(prop, flags_by_key) -> bool: + # A malformed condition shape is a definitive local no-match, not a gap. + if _flag_dependency_shape_problem(prop) is not None: + return False + dependency_chain = prop.get("dependency_chain") + if not isinstance(dependency_chain, list) or not dependency_chain: + return True + return any(not flags_by_key.get(key) for key in dependency_chain) + + def _is_unsupported_property_group( property_group, flags_by_key, cohort_properties, visited_cohort_ids ) -> bool: diff --git a/posthog/integrations/asgi.py b/posthog/integrations/asgi.py index d58f818db..f20b4779b 100644 --- a/posthog/integrations/asgi.py +++ b/posthog/integrations/asgi.py @@ -120,10 +120,13 @@ def _server_host(scope: Mapping[str, Any]) -> Optional[str]: hostname = _sanitize_text(hostname) if not hostname: return None + return _format_server_host(hostname, port, scope.get("scheme")) + + +def _format_server_host(hostname: str, port: int, scheme: Any) -> str: if ":" in hostname and not hostname.startswith("["): hostname = f"[{hostname}]" - scheme = scope.get("scheme") default_port = (scheme == "http" and port == 80) or ( scheme == "https" and port == 443 ) @@ -148,6 +151,23 @@ def _extract_properties( properties["$user_agent"] = user_agent properties["$raw_user_agent"] = user_agent + ip_address = _client_ip(scope, headers) + if ip_address: + properties["$ip"] = ip_address + + scheme = _sanitize_text(scope.get("scheme"), max_length=16) + host = _decode_header(headers.get(b"host")) or _server_host(scope) + if scheme and host and path: + # Deliberately omit query strings: they commonly contain secrets and + # high-cardinality values. The path remains available separately. + properties["$current_url"] = f"{scheme}://{host}{path}" + + return properties + + +def _client_ip( + scope: Mapping[str, Any], headers: Mapping[bytes, bytes] +) -> Optional[str]: forwarded_for = _decode_header(headers.get(b"x-forwarded-for")) if forwarded_for: ip_address = _sanitize_text(forwarded_for.split(",", 1)[0]) @@ -160,17 +180,7 @@ def _extract_properties( and isinstance(client[0], str) else None ) - if ip_address: - properties["$ip"] = ip_address - - scheme = _sanitize_text(scope.get("scheme"), max_length=16) - host = _decode_header(headers.get(b"host")) or _server_host(scope) - if scheme and host and path: - # Deliberately omit query strings: they commonly contain secrets and - # high-cardinality values. The path remains available separately. - properties["$current_url"] = f"{scheme}://{host}{path}" - - return properties + return ip_address async def _resolve_callback_result(value): @@ -252,13 +262,7 @@ async def __call__(self, scope, receive, send) -> None: _enrich_current_span() - properties = _extract_properties(scope, headers) - if self.extra_properties: - extra_properties = await _resolve_callback_result( - self.extra_properties(scope) - ) - if extra_properties: - properties.update(extra_properties) + properties = await self._request_properties(scope, headers) for key, value in properties.items(): contexts.tag(key, value) @@ -282,3 +286,13 @@ async def __call__(self, scope, receive, send) -> None: exception, _capture_metadata=_CAPTURE_METADATA ) raise + + async def _request_properties(self, scope, headers) -> dict[str, Any]: + properties = _extract_properties(scope, headers) + if self.extra_properties: + extra_properties = await _resolve_callback_result( + self.extra_properties(scope) + ) + if extra_properties: + properties.update(extra_properties) + return properties diff --git a/posthog/integrations/celery.py b/posthog/integrations/celery.py index 0e37379f9..0d59c2a52 100644 --- a/posthog/integrations/celery.py +++ b/posthog/integrations/celery.py @@ -392,18 +392,7 @@ def _handle_task_end( task_properties["error_type"] = type(exception).__name__ task_properties["error_message"] = str(exception) - task_name = task_properties.get("celery_task_name") - should_track = False - if self.capture_task_lifecycle_events or ( - exception and self.capture_exceptions - ): - should_track = self._should_track(task_name, task_properties) - - if exception and self.capture_exceptions and should_track: - self._capture_exception(exception) - - if self.capture_task_lifecycle_events and should_track: - self._capture_event(f"celery task {state}", properties=task_properties) + self._capture_task_end(state, task_properties, exception) except Exception: logger.exception("Failed to process Celery %s state", state) finally: @@ -411,6 +400,22 @@ def _handle_task_end( if ctx is not None: ctx.__exit__(None, None, None) + def _capture_task_end( + self, state: str, task_properties: dict, exception: Any + ) -> None: + task_name = task_properties.get("celery_task_name") + should_track = False + if self.capture_task_lifecycle_events or ( + exception and self.capture_exceptions + ): + should_track = self._should_track(task_name, task_properties) + + if exception and self.capture_exceptions and should_track: + self._capture_exception(exception) + + if self.capture_task_lifecycle_events and should_track: + self._capture_event(f"celery task {state}", properties=task_properties) + def _apply_propagated_identity(self, request: Any) -> None: headers = self._extract_headers(request) distinct_id = headers.get(CONTEXT_DISTINCT_ID_HEADER) diff --git a/posthog/integrations/django.py b/posthog/integrations/django.py index 9c305a75c..a49527181 100644 --- a/posthog/integrations/django.py +++ b/posthog/integrations/django.py @@ -96,6 +96,35 @@ def __init__(self, get_response): from django.conf import settings + self._configure_callbacks(settings) + + capture_exceptions = getattr(settings, "POSTHOG_MW_CAPTURE_EXCEPTIONS", True) + self.capture_exceptions: Optional[bool] = ( + capture_exceptions + if capture_exceptions is None or isinstance(capture_exceptions, bool) + else True + ) + capture_tracing_headers = getattr( + settings, "POSTHOG_MW_CAPTURE_TRACING_HEADERS", None + ) + trust_tracing_headers = getattr( + settings, "POSTHOG_MW_TRUST_TRACING_HEADERS", None + ) + self.capture_tracing_headers = _resolve_capture_tracing_headers( + capture_tracing_headers, + trust_tracing_headers, + canonical_name="POSTHOG_MW_CAPTURE_TRACING_HEADERS", + deprecated_name="POSTHOG_MW_TRUST_TRACING_HEADERS", + ) + + if hasattr(settings, "POSTHOG_MW_CLIENT") and isinstance( + settings.POSTHOG_MW_CLIENT, Client + ): + self.client = cast("Optional[Client]", settings.POSTHOG_MW_CLIENT) + else: + self.client = None + + def _configure_callbacks(self, settings) -> None: extra_properties = ( settings.POSTHOG_MW_EXTRA_PROPERTIES if hasattr(settings, "POSTHOG_MW_EXTRA_PROPERTIES") @@ -128,32 +157,6 @@ def __init__(self, get_response): else None ) - capture_exceptions = getattr(settings, "POSTHOG_MW_CAPTURE_EXCEPTIONS", True) - self.capture_exceptions: Optional[bool] = ( - capture_exceptions - if capture_exceptions is None or isinstance(capture_exceptions, bool) - else True - ) - capture_tracing_headers = getattr( - settings, "POSTHOG_MW_CAPTURE_TRACING_HEADERS", None - ) - trust_tracing_headers = getattr( - settings, "POSTHOG_MW_TRUST_TRACING_HEADERS", None - ) - self.capture_tracing_headers = _resolve_capture_tracing_headers( - capture_tracing_headers, - trust_tracing_headers, - canonical_name="POSTHOG_MW_CAPTURE_TRACING_HEADERS", - deprecated_name="POSTHOG_MW_TRUST_TRACING_HEADERS", - ) - - if hasattr(settings, "POSTHOG_MW_CLIENT") and isinstance( - settings.POSTHOG_MW_CLIENT, Client - ): - self.client = cast("Optional[Client]", settings.POSTHOG_MW_CLIENT) - else: - self.client = None - @property def trust_tracing_headers(self) -> bool: """Deprecated alias for ``capture_tracing_headers``.""" @@ -236,6 +239,9 @@ def _build_tags(self, request, user_id, user_email): tags["$user_agent"] = user_agent tags["$raw_user_agent"] = user_agent + return self._transform_properties(request, tags) + + def _transform_properties(self, request, tags): # Apply application-specific properties, then the optional transform. if self.extra_properties: extra = self.extra_properties(request) diff --git a/posthog/mcp/__init__.py b/posthog/mcp/__init__.py index 98a73bede..28b34e1b3 100644 --- a/posthog/mcp/__init__.py +++ b/posthog/mcp/__init__.py @@ -248,32 +248,7 @@ def instrument( :param options: Optional :class:`MCPAnalyticsOptions`. """ opts = options or MCPAnalyticsOptions() - server_build = validate_server_build(opts.server_build) - - # Install the logger first so the version advisory below (and any warning) is - # actually visible rather than going to the default no-op sink. - if opts.logger: - set_logger(opts.logger) - - # The wrapping path hooks the official MCP SDK's server internals, so it needs the - # `mcp` package. It's a peer dependency (you already have it — you built the server - # with it), imported lazily here rather than bundled. PostHogMCP (custom dispatchers) - # doesn't need it at all. Raise a clear error rather than a silent no-op below. - try: - import mcp # noqa: F401 - except ImportError: - raise ModuleNotFoundError( - "instrument() needs the MCP SDK. Install it with: pip install 'mcp>=1.26'. " - "(PostHogMCP for custom dispatchers works without it.)" - ) - _warn_if_unsupported_mcp_version() - - # Fail fast on a `collect_feedback` config error (reserved extra key, - # undeclared extra_required) — before the try below, so it raises instead of - # degrading to the no-op handle and first surfacing at tools/list time. - feedback_options = resolve_collect_feedback_options(opts.collect_feedback) - if feedback_options is not None: - get_feedback_tool_descriptor(feedback_options) + server_build = _prepare_instrumentation(opts) key = tracking_key(server) @@ -281,15 +256,9 @@ def instrument( # MCP is an optional peer: load adapters only when instrumentation is # requested, inside the no-crash boundary. Class probes stay major-specific. from ._compatibility import ( - is_fastmcp, is_fastmcp_v2, - is_low_level_server, - is_mcpserver, uses_v2_handler_registry, ) - from ._instrument_fastmcp import instrument_fastmcp - from ._instrument_lowlevel import instrument_fastmcp_v2, instrument_low_level - from ._instrument_v2 import instrument_lowlevel_v2, instrument_mcpserver_v2 client = _resolve_client(posthog_client) if client is None: @@ -298,13 +267,7 @@ def instrument( existing_data = get_server_tracking_data(key) data = existing_data if data is None: - sink = McpEventSink(client) if client is not None else None - data = MCPAnalyticsData( - options=opts, - sink=sink, - session_id=new_session_id(), - server_build=server_build, - ) + data = _new_tracking_data(opts, client, server_build) if is_fastmcp_v2(server) and uses_v2_handler_registry(key): data.standalone_fastmcp = weakref.ref(server) @@ -319,26 +282,7 @@ def instrument( set_server_tracking_data(key, data) - if is_fastmcp(server): - instrument_fastmcp(server, data) - elif is_mcpserver(server): - instrument_mcpserver_v2(server, data) - elif is_fastmcp_v2(server): - if uses_v2_handler_registry(server._mcp_server): - instrument_lowlevel_v2(server._mcp_server, data) - else: - instrument_fastmcp_v2(server, data) - elif is_low_level_server(server): - if uses_v2_handler_registry(server): - instrument_lowlevel_v2(server, data) - else: - instrument_low_level(server, data) - else: - raise TypeError( - f"Unsupported server type: {type(server)!r}. Pass a high-level server " - "(mcp.server.fastmcp.FastMCP on SDK 1.x, mcp.server.mcpserver.MCPServer " - "on SDK 2.x, or jlowin's fastmcp.FastMCP) or a low-level mcp.server.Server." - ) + _instrument_adapter(server, data) if client is not None: from ._lib_identity import apply_mcp_lib_identity @@ -354,3 +298,77 @@ def instrument( except Exception as error: # noqa: BLE001 log(f"Warning: failed to instrument server - {error}") return _NoopAnalytics() + + +def _instrument_adapter(server: Any, data: MCPAnalyticsData) -> None: + from ._compatibility import ( + is_fastmcp, + is_fastmcp_v2, + is_low_level_server, + is_mcpserver, + uses_v2_handler_registry, + ) + from ._instrument_fastmcp import instrument_fastmcp + from ._instrument_lowlevel import instrument_fastmcp_v2, instrument_low_level + from ._instrument_v2 import instrument_lowlevel_v2, instrument_mcpserver_v2 + + if is_fastmcp(server): + instrument_fastmcp(server, data) + elif is_mcpserver(server): + instrument_mcpserver_v2(server, data) + elif is_fastmcp_v2(server): + if uses_v2_handler_registry(server._mcp_server): + instrument_lowlevel_v2(server._mcp_server, data) + else: + instrument_fastmcp_v2(server, data) + elif is_low_level_server(server): + if uses_v2_handler_registry(server): + instrument_lowlevel_v2(server, data) + else: + instrument_low_level(server, data) + else: + raise TypeError( + f"Unsupported server type: {type(server)!r}. Pass a high-level server " + "(mcp.server.fastmcp.FastMCP on SDK 1.x, mcp.server.mcpserver.MCPServer " + "on SDK 2.x, or jlowin's fastmcp.FastMCP) or a low-level mcp.server.Server." + ) + + +def _new_tracking_data( + opts: MCPAnalyticsOptions, client: Optional[Client], server_build: Optional[str] +) -> MCPAnalyticsData: + sink = McpEventSink(client) if client is not None else None + return MCPAnalyticsData( + options=opts, sink=sink, session_id=new_session_id(), server_build=server_build + ) + + +def _prepare_instrumentation(opts: MCPAnalyticsOptions) -> Optional[str]: + server_build = validate_server_build(opts.server_build) + + # Install the logger first so the version advisory below (and any warning) is + # actually visible rather than going to the default no-op sink. + if opts.logger: + set_logger(opts.logger) + + # The wrapping path hooks the official MCP SDK's server internals, so it needs the + # `mcp` package. It's a peer dependency (you already have it — you built the server + # with it), imported lazily here rather than bundled. PostHogMCP (custom dispatchers) + # doesn't need it at all. Raise a clear error rather than a silent no-op below. + try: + import mcp # noqa: F401 + except ImportError: + raise ModuleNotFoundError( + "instrument() needs the MCP SDK. Install it with: pip install 'mcp>=1.26'. " + "(PostHogMCP for custom dispatchers works without it.)" + ) + _warn_if_unsupported_mcp_version() + + # Fail fast on a `collect_feedback` config error (reserved extra key, + # undeclared extra_required) — before the try below, so it raises instead of + # degrading to the no-op handle and first surfacing at tools/list time. + feedback_options = resolve_collect_feedback_options(opts.collect_feedback) + if feedback_options is not None: + get_feedback_tool_descriptor(feedback_options) + + return server_build diff --git a/posthog/mcp/_capture.py b/posthog/mcp/_capture.py index ef48ef4a4..ec7349d11 100644 --- a/posthog/mcp/_capture.py +++ b/posthog/mcp/_capture.py @@ -15,6 +15,7 @@ from ._internal import MCPAnalyticsData from .logger import log from ._sink import McpCaptureOptions +from .types import UserIdentity from .version import __version__ @@ -53,9 +54,7 @@ def capture_event( "protocol_version": event_input.get("protocol_version"), "client_user_agent": event_input.get("client_user_agent"), "vendor_client": event_input.get("vendor_client"), - "identify_actor_given_id": actor.distinct_id if actor else None, - "identify_actor_data": (actor.properties or {}) if actor else {}, - "groups": actor.groups if actor else None, + **_actor_fields(actor), "resource_name": event_input.get("resource_name"), "tool_category": event_input.get("tool_category"), "tool_description": event_input.get("tool_description"), @@ -82,3 +81,11 @@ def capture_event( def log_capture_skipped() -> None: log("Warning: Server tracking data not found. Event will not be published.") + + +def _actor_fields(actor: Optional[UserIdentity]) -> Dict[str, Any]: + return { + "identify_actor_given_id": actor.distinct_id if actor else None, + "identify_actor_data": (actor.properties or {}) if actor else {}, + "groups": actor.groups if actor else None, + } diff --git a/posthog/mcp/_context_parameters.py b/posthog/mcp/_context_parameters.py index 943d13764..f0a742c31 100644 --- a/posthog/mcp/_context_parameters.py +++ b/posthog/mcp/_context_parameters.py @@ -55,11 +55,7 @@ def add_context_parameter_to_schema( """ schema = input_schema - if ( - schema - and isinstance(schema.get("properties"), dict) - and "context" in schema["properties"] - ): + if schema_has_param(schema, "context"): log( f"WARN: Tool \"{tool_name}\" already has 'context' parameter. Skipping context injection." ) @@ -89,11 +85,15 @@ def add_context_parameter_to_schema( } if required: - required_list = schema.get("required") - if isinstance(required_list, list): - if "context" not in required_list: - required_list.append("context") - else: - schema["required"] = ["context"] + _require_context(schema) return schema + + +def _require_context(schema: Dict[str, Any]) -> None: + required_list = schema.get("required") + if isinstance(required_list, list): + if "context" not in required_list: + required_list.append("context") + else: + schema["required"] = ["context"] diff --git a/posthog/mcp/_conversation_id.py b/posthog/mcp/_conversation_id.py index d3c774516..9f6a91fe9 100644 --- a/posthog/mcp/_conversation_id.py +++ b/posthog/mcp/_conversation_id.py @@ -14,6 +14,7 @@ import re from typing import Any, Dict, Optional, Tuple +from ._context_parameters import schema_has_param from .constants import DEFAULT_CONVERSATION_ID_DESCRIPTION from ._ids import _uuid7 from .logger import log @@ -39,11 +40,7 @@ def add_conversation_id_to_schema( ``$ref``/``oneOf``/``allOf``/``anyOf``. """ schema = input_schema - if ( - schema - and isinstance(schema.get("properties"), dict) - and CONVERSATION_ID_PARAM_NAME in schema["properties"] - ): + if schema and schema_has_param(schema, CONVERSATION_ID_PARAM_NAME): if _is_our_declaration(schema["properties"][CONVERSATION_ID_PARAM_NAME]): return schema log( @@ -151,17 +148,26 @@ def inject_prompt_back(result: Any, conversation_id: str) -> Any: if isinstance(result, list): return [*result, block] - target = getattr(result, "root", result) - content = getattr(target, "content", None) - if not isinstance(content, list): - return result + return _inject_model_prompt_back(result, block) + + +def _copy_with_attr(value: Any, attr: str, updated: Any) -> Optional[Any]: try: - import mcp.types as mcp_types # noqa: PLC0415 + copied = copy.copy(value) + except Exception: # noqa: BLE001 - delivery must not break a tool call + return None + if copied is value: + return None + try: + setattr(copied, attr, updated) + except Exception: # noqa: BLE001 - read-only objects fail closed + return None + return copied - block = mcp_types.TextContent(type="text", text=block["text"]) - except ImportError: - pass +def _copy_prompt_back_result( + result: Any, target: Any, content: list, block: Any +) -> Any: copy_model = getattr(target, "model_copy", None) if callable(copy_model): try: @@ -185,15 +191,16 @@ def inject_prompt_back(result: Any, conversation_id: str) -> Any: return result if wrapped is None else wrapped -def _copy_with_attr(value: Any, attr: str, updated: Any) -> Optional[Any]: - try: - copied = copy.copy(value) - except Exception: # noqa: BLE001 - delivery must not break a tool call - return None - if copied is value: - return None +def _inject_model_prompt_back(result: Any, block: Any) -> Any: + target = getattr(result, "root", result) + content = getattr(target, "content", None) + if not isinstance(content, list): + return result try: - setattr(copied, attr, updated) - except Exception: # noqa: BLE001 - read-only objects fail closed - return None - return copied + import mcp.types as mcp_types # noqa: PLC0415 + + block = mcp_types.TextContent(type="text", text=block["text"]) + except ImportError: + pass + + return _copy_prompt_back_result(result, target, content, block) diff --git a/posthog/mcp/_instrument_lowlevel.py b/posthog/mcp/_instrument_lowlevel.py index a8ab58671..01e972332 100644 --- a/posthog/mcp/_instrument_lowlevel.py +++ b/posthog/mcp/_instrument_lowlevel.py @@ -724,14 +724,7 @@ async def _listed_input_schema( # spend rate limits, and its registry can disagree with what that # middleware advertises, so the advertised schema is unknown. return None - candidates = [tool for tool in await list_tools() if tool.name == name] - if version is not None: - candidates = [ - tool - for tool in candidates - if str(getattr(tool, "version", "")) == version - ] - tool = candidates[-1] if candidates else None + tool = _last_listed_tool(await list_tools(), name, version) schema = getattr(tool, "parameters", None) except Exception as error: # noqa: BLE001 - analytics must not break dispatch log(f"PostHog MCP: could not resolve schema for tool {name!r} - {error}") @@ -765,13 +758,7 @@ def _schema_view(schema: Dict[str, Any], *, dereferenced: bool) -> Tuple[set, bo node.get(key) for node in nodes for key in ("oneOf", "allOf", "anyOf") ): return set(_INJECTED_KEYS), False - declared = { - key - for node in nodes - if isinstance(node.get("properties"), dict) - for key in node["properties"] - if key in _INJECTED_KEYS - } + declared = _declared_injected_keys(nodes) injectable = can_inject_model_parameter(nodes[-1]) and "llm_model" not in declared return declared, injectable @@ -897,10 +884,7 @@ def _reference_chain(schema: Dict[str, Any]) -> Optional[List[Dict[str, Any]]]: return nodes if not isinstance(ref, str) or not ref.startswith("#/"): return None - target: Any = schema - for part in ref[2:].split("/"): - key = part.replace("~1", "/").replace("~0", "~") # JSON Pointer escapes - target = target.get(key) if isinstance(target, dict) else None + target = _resolve_local_pointer(schema, ref) if not isinstance(target, dict) or any(target is node for node in nodes): return None nodes.append(target) @@ -946,3 +930,31 @@ def _mcp_session_id(server: Any) -> Optional[str]: except Exception: # noqa: BLE001 pass return None + + +def _resolve_local_pointer(schema: Dict[str, Any], ref: str) -> Any: + target: Any = schema + for part in ref[2:].split("/"): + key = part.replace("~1", "/").replace("~0", "~") # JSON Pointer escapes + target = target.get(key) if isinstance(target, dict) else None + return target + + +def _declared_injected_keys(nodes: List[Dict[str, Any]]) -> set: + declared = { + key + for node in nodes + if isinstance(node.get("properties"), dict) + for key in node["properties"] + if key in _INJECTED_KEYS + } + return declared + + +def _last_listed_tool(tools: list, name: str, version: Optional[str]) -> Any: + candidates = [tool for tool in tools if tool.name == name] + if version is not None: + candidates = [ + tool for tool in candidates if str(getattr(tool, "version", "")) == version + ] + return candidates[-1] if candidates else None diff --git a/posthog/mcp/_instrument_v2.py b/posthog/mcp/_instrument_v2.py index b7d597789..9f63f1c11 100644 --- a/posthog/mcp/_instrument_v2.py +++ b/posthog/mcp/_instrument_v2.py @@ -554,24 +554,7 @@ async def _standalone_injected_parameters( Listings from other requests can have different application-owned parameters. Without a schema, stripping could delete application arguments. """ - try: - from fastmcp.utilities.versions import VersionSpec, version_sort_key - - version_spec = VersionSpec(eq=version) if version else None - # Middleware can shadow registered tools, so resolve the effective listing. - candidates = [ - tool - for tool in await server.list_tools() - if tool.name == name - and (version_spec is None or version_spec.matches(tool.version)) - ] - tool = max(candidates, key=version_sort_key, default=None) - if tool is None: - tool = await server.get_tool(name, version=version_spec) - schema = getattr(tool, "parameters", None) - except Exception as error: # noqa: BLE001 - schema lookup must not prevent dispatch - log(f"PostHog MCP: could not resolve schema for tool {name!r} - {error}") - return None, None + schema = await _standalone_input_schema(server, name, version) if not isinstance(schema, dict): return None, None injected = set() @@ -874,3 +857,27 @@ def _name_owned_by_real_tool_v2(high_level: Any, name: str) -> Optional[bool]: except Exception as err: # noqa: BLE001 - analytics must not break the call warn_ownership_lookup_failed(name, err) return None + + +async def _standalone_input_schema( + server: Any, name: str, version: Optional[str] +) -> Any: + try: + from fastmcp.utilities.versions import VersionSpec, version_sort_key + + version_spec = VersionSpec(eq=version) if version else None + # Middleware can shadow registered tools, so resolve the effective listing. + candidates = [ + tool + for tool in await server.list_tools() + if tool.name == name + and (version_spec is None or version_spec.matches(tool.version)) + ] + tool = max(candidates, key=version_sort_key, default=None) + if tool is None: + tool = await server.get_tool(name, version=version_spec) + schema = getattr(tool, "parameters", None) + except Exception as error: # noqa: BLE001 - schema lookup must not prevent dispatch + log(f"PostHog MCP: could not resolve schema for tool {name!r} - {error}") + return None + return schema diff --git a/posthog/mcp/_instrumentation.py b/posthog/mcp/_instrumentation.py index 9b1db44c7..f480a98ef 100644 --- a/posthog/mcp/_instrumentation.py +++ b/posthog/mcp/_instrumentation.py @@ -57,7 +57,12 @@ from .session import resolve_session_id, resolve_session_id_with_source from .session_token import SessionTokenPayload, decode_session_id from .tools import resolve_missing_capability_tool_name -from .types import CollectFeedbackOptions, FeedbackReport, ToolInputOptions +from .types import ( + CollectFeedbackOptions, + FeedbackReport, + InputAliasMap, + ToolInputOptions, +) # The virtual tools this SDK advertises into tools/list. Every piece of per-tool # policy -- enable switch, configured name, warning text, warn-once @@ -813,22 +818,10 @@ async def record_tool_call( event["llm_model"] = model event["llm_model_source"] = source - if error is not None: - event["is_error"] = True - event["error"] = capture_exception(error) - elif result is not None: - event["response"] = _wrap_response(result) - if is_tool_result_error(result): - event["is_error"] = True - event["error"] = capture_exception(result) + _set_tool_outcome(event, result, error) props = await resolve_event_properties(data, request, extra) - input_aliases = None - if data.options.resolve_input_aliases is not None: - try: - input_aliases = data.options.resolve_input_aliases(name) - except Exception as err: # noqa: BLE001 - analytics callbacks are isolated - log(f"Warning: resolve_input_aliases failed for tool {name}: {err}") + input_aliases = _resolve_input_aliases(data, name) event["properties"] = { **(props or {}), **get_tool_input_properties( @@ -1033,18 +1026,7 @@ def resolve_virtual_tool_injection( listed = advertised_tool_names(tools) if not is_first_page: - for kind, name in enabled.items(): - # Only warn about a tool we actually shadowed. If this kind never - # made it onto the first page -- a real tool already held the name - # ("blocked"), or the other virtual tool won it ("duplicate") -- - # then nothing of ours is advertised under it and the host's tool - # runs untouched. Telling them otherwise sends them chasing a bug - # that isn't there. - if name in listed and not any( - (kind, name, variant) in data.warned_virtual_tool_collisions - for variant in ("blocked", "duplicate") - ): - _warn_virtual_tool_collision(data, kind, name, "shadowed") + _warn_shadowed_virtual_tools(data, enabled, listed) return {} injectable: Dict[str, str] = {} @@ -1215,11 +1197,7 @@ def mutate_tool_schema( schema, "conversation_id" ): schema = add_conversation_id_to_schema(schema, tool.name) - if schema is not original_schema: - try: - setattr(tool, schema_attribute, schema) - except Exception: # noqa: BLE001 - some schema attrs may be read-only - log(f"WARN: could not set {schema_attribute} on tool {tool.name}") + _write_tool_schema(tool, schema_attribute, original_schema, schema) if data.options.enable_conversation_id: data.tool_output_instructions[tool.name] = add_instructions_to_output_schema( tool @@ -1578,3 +1556,53 @@ async def record_resource_request( fire_and_forget(capture_event(data, event), data) except Exception as err: # noqa: BLE001 - isolate analytics from the request path log(f"record_resource_request failed (event dropped): {err}") + + +def _set_tool_outcome(event: Dict[str, Any], result: Any, error: Any) -> None: + if error is not None: + event["is_error"] = True + event["error"] = capture_exception(error) + elif result is not None: + event["response"] = _wrap_response(result) + if is_tool_result_error(result): + event["is_error"] = True + event["error"] = capture_exception(result) + + +def _resolve_input_aliases( + data: MCPAnalyticsData, name: str +) -> Optional[InputAliasMap]: + input_aliases = None + if data.options.resolve_input_aliases is not None: + try: + input_aliases = data.options.resolve_input_aliases(name) + except Exception as err: # noqa: BLE001 - analytics callbacks are isolated + log(f"Warning: resolve_input_aliases failed for tool {name}: {err}") + return input_aliases + + +def _warn_shadowed_virtual_tools( + data: MCPAnalyticsData, enabled: Dict[str, str], listed: Set[str] +) -> None: + for kind, name in enabled.items(): + # Only warn about a tool we actually shadowed. If this kind never + # made it onto the first page -- a real tool already held the name + # ("blocked"), or the other virtual tool won it ("duplicate") -- + # then nothing of ours is advertised under it and the host's tool + # runs untouched. Telling them otherwise sends them chasing a bug + # that isn't there. + if name in listed and not any( + (kind, name, variant) in data.warned_virtual_tool_collisions + for variant in ("blocked", "duplicate") + ): + _warn_virtual_tool_collision(data, kind, name, "shadowed") + + +def _write_tool_schema( + tool: Any, schema_attribute: str, original_schema: Any, schema: Any +) -> None: + if schema is not original_schema: + try: + setattr(tool, schema_attribute, schema) + except Exception: # noqa: BLE001 - some schema attrs may be read-only + log(f"WARN: could not set {schema_attribute} on tool {tool.name}") diff --git a/posthog/mcp/_output_instructions.py b/posthog/mcp/_output_instructions.py index ee391f1e7..6841ef817 100644 --- a/posthog/mcp/_output_instructions.py +++ b/posthog/mcp/_output_instructions.py @@ -198,6 +198,28 @@ def mirror_instructions_into_structured_content( }, True return result, False + return _mirror_model_instructions(result, payload) + + +def _copy_with_attr(value: Any, attr: str, updated: Any) -> Optional[Any]: + """Return a shallow copy with one changed attribute, or ``None`` when the + object cannot be copied safely.""" + try: + copied = copy.copy(value) + except Exception: # noqa: BLE001 - analytics must not break tool results + return None + if copied is value: + return None + try: + setattr(copied, attr, updated) + except Exception: # noqa: BLE001 - read-only result objects fail closed + return None + return copied + + +def _mirror_model_instructions( + result: Any, payload: Dict[str, Any] +) -> Tuple[Any, bool]: # CallToolResult model (or the ServerResult wrapper around one). target = getattr(result, "root", result) attr, structured = _read_attr(target, _STRUCTURED_CONTENT_ATTRS) @@ -212,6 +234,12 @@ def mirror_instructions_into_structured_content( # result object, and pinning one conversation's handle onto it would serve # that handle to every later caller (and, through the handle, collapse # unrelated clients into one session). + return _copy_structured_result(result, target, attr, updated) + + +def _copy_structured_result( + result: Any, target: Any, attr: str, updated: Any +) -> Tuple[Any, bool]: copy_model = getattr(target, "model_copy", None) if callable(copy_model): try: @@ -234,19 +262,3 @@ def mirror_instructions_into_structured_content( if new_result is None: return result, False return new_result, True - - -def _copy_with_attr(value: Any, attr: str, updated: Any) -> Optional[Any]: - """Return a shallow copy with one changed attribute, or ``None`` when the - object cannot be copied safely.""" - try: - copied = copy.copy(value) - except Exception: # noqa: BLE001 - analytics must not break tool results - return None - if copied is value: - return None - try: - setattr(copied, attr, updated) - except Exception: # noqa: BLE001 - read-only result objects fail closed - return None - return copied diff --git a/posthog/mcp/_posthog_events.py b/posthog/mcp/_posthog_events.py index f7b58f0e7..e446e5c78 100644 --- a/posthog/mcp/_posthog_events.py +++ b/posthog/mcp/_posthog_events.py @@ -112,18 +112,7 @@ def _is_tool_call(event: Event) -> bool: def _add_common_properties(event: Event, properties: Dict[str, Any]) -> None: - if event.get("resource_name"): - is_unknown_tool = ( - event.get("event_type") == MCPAnalyticsEventType.MCP_UNKNOWN_TOOL - ) - if not is_unknown_tool: - properties[_P.RESOURCE_NAME] = event["resource_name"] - if is_unknown_tool or _is_tool_call(event): - properties[_P.TOOL_NAME] = event["resource_name"] - if event.get("tool_description") and _is_tool_call(event): - properties[_P.TOOL_DESCRIPTION] = event["tool_description"] - if event.get("tool_category") and _is_tool_call(event): - properties[_P.TOOL_CATEGORY] = event["tool_category"] + _add_resource_properties(event, properties) if ( event.get("listed_tool_names") and len(event["listed_tool_names"]) > 0 @@ -132,42 +121,29 @@ def _add_common_properties(event: Event, properties: Dict[str, Any]) -> None: properties[_P.LISTED_TOOL_NAMES] = event["listed_tool_names"] if event.get("duration") is not None: properties[_P.DURATION_MS] = event["duration"] - if event.get("server_name"): - properties[_P.SERVER_NAME] = event["server_name"] - if event.get("server_version"): - properties[_P.SERVER_VERSION] = event["server_version"] - if event.get("client_name"): - properties[_P.CLIENT_NAME] = event["client_name"] - if event.get("client_version"): - properties[_P.CLIENT_VERSION] = event["client_version"] - # HTTP transports only, and only for the request that carried the header — - # stdio and in-memory servers simply never set these. - if event.get("client_user_agent"): - properties[_P.CLIENT_USER_AGENT] = event["client_user_agent"] - if event.get("vendor_client"): - properties[_P.VENDOR_CLIENT] = event["vendor_client"] - if event.get("protocol_version"): - properties[_P.PROTOCOL_VERSION] = event["protocol_version"] - if event.get("user_intent"): - properties[_P.INTENT] = event["user_intent"] - if event.get("user_intent_source"): - properties[_P.INTENT_SOURCE] = event["user_intent_source"] - if event.get("llm_model"): - properties[_P.LLM_MODEL] = event["llm_model"] - if event.get("llm_model_source"): - properties[_P.LLM_MODEL_SOURCE] = event["llm_model_source"] + # HTTP header fields are set only by the request that carried them. + _add_truthy_properties( + event, + properties, + ( + ("server_name", _P.SERVER_NAME), + ("server_version", _P.SERVER_VERSION), + ("client_name", _P.CLIENT_NAME), + ("client_version", _P.CLIENT_VERSION), + ("client_user_agent", _P.CLIENT_USER_AGENT), + ("vendor_client", _P.VENDOR_CLIENT), + ("protocol_version", _P.PROTOCOL_VERSION), + ("user_intent", _P.INTENT), + ("user_intent_source", _P.INTENT_SOURCE), + ("llm_model", _P.LLM_MODEL), + ("llm_model_source", _P.LLM_MODEL_SOURCE), + ), + ) if event.get("is_error") is not None: properties[_P.IS_ERROR] = event["is_error"] if event.get("is_error"): _add_error_details(event, properties) - if event.get("parameters") is not None: - properties[_P.PARAMETERS] = event["parameters"] - if event.get("response") is not None: - properties[_P.RESPONSE] = event["response"] - identify_actor_data = event.get("identify_actor_data") - if identify_actor_data and len(identify_actor_data) > 0: - # Person properties from identify().properties go straight to $set. - properties["$set"] = {**identify_actor_data} + _add_payload_properties(event, properties) # The dispatch wrappers that say nothing the event does not already say: the @@ -287,16 +263,17 @@ def _build_exception_event(event: Event) -> PostHogCaptureEvent: properties[_P.TOOL_DESCRIPTION] = event["tool_description"] if event.get("tool_category") and _is_tool_call(event): properties[_P.TOOL_CATEGORY] = event["tool_category"] - if event.get("server_name"): - properties[_P.SERVER_NAME] = event["server_name"] - if event.get("server_version"): - properties[_P.SERVER_VERSION] = event["server_version"] - if event.get("client_name"): - properties[_P.CLIENT_NAME] = event["client_name"] - if event.get("client_version"): - properties[_P.CLIENT_VERSION] = event["client_version"] - if event.get("protocol_version"): - properties[_P.PROTOCOL_VERSION] = event["protocol_version"] + _add_truthy_properties( + event, + properties, + ( + ("server_name", _P.SERVER_NAME), + ("server_version", _P.SERVER_VERSION), + ("client_name", _P.CLIENT_NAME), + ("client_version", _P.CLIENT_VERSION), + ("protocol_version", _P.PROTOCOL_VERSION), + ), + ) _add_custom_properties(event, properties) _add_server_build(event, properties) @@ -307,3 +284,37 @@ def _build_exception_event(event: Event) -> PostHogCaptureEvent: "properties": properties, "timestamp": _get_timestamp(event), } + + +def _add_truthy_properties( + event: Event, properties: Dict[str, Any], fields: tuple[tuple[str, str], ...] +) -> None: + for field, property_name in fields: + if event.get(field): + properties[property_name] = event[field] + + +def _add_resource_properties(event: Event, properties: Dict[str, Any]) -> None: + if event.get("resource_name"): + is_unknown_tool = ( + event.get("event_type") == MCPAnalyticsEventType.MCP_UNKNOWN_TOOL + ) + if not is_unknown_tool: + properties[_P.RESOURCE_NAME] = event["resource_name"] + if is_unknown_tool or _is_tool_call(event): + properties[_P.TOOL_NAME] = event["resource_name"] + if event.get("tool_description") and _is_tool_call(event): + properties[_P.TOOL_DESCRIPTION] = event["tool_description"] + if event.get("tool_category") and _is_tool_call(event): + properties[_P.TOOL_CATEGORY] = event["tool_category"] + + +def _add_payload_properties(event: Event, properties: Dict[str, Any]) -> None: + if event.get("parameters") is not None: + properties[_P.PARAMETERS] = event["parameters"] + if event.get("response") is not None: + properties[_P.RESPONSE] = event["response"] + identify_actor_data = event.get("identify_actor_data") + if identify_actor_data and len(identify_actor_data) > 0: + # Person properties from identify().properties go straight to $set. + properties["$set"] = {**identify_actor_data} diff --git a/posthog/mcp/_sanitization.py b/posthog/mcp/_sanitization.py index 426d2d29f..63bbe87eb 100644 --- a/posthog/mcp/_sanitization.py +++ b/posthog/mcp/_sanitization.py @@ -193,21 +193,9 @@ def _sanitize_single_url(value: str, *, nested: bool, in_prose: bool) -> str: try: url = urlsplit(url_text) query, sanitized_query = _sanitize_url_fields(url.query, nested=nested) - # A `?` inside a fragment splits it in two, and each half stands on its - # own: `#/callback?token=x` is text then fields, `#a=1?b=2` is fields - # twice. Shape cannot tell them apart — `#/token=x` looks like a route and - # is a field list — so each half is judged only on whether it holds a `=`. - head, separator, fragment_tail = url.fragment.partition("?") - sanitized_head, head_rewrote_last = _sanitize_fragment_part(head, nested=nested) - # A `?` can fall inside a credential (`#password=pre?fix`), and nothing - # here can tell that from a real boundary. When the head ends in a value - # just redacted, the tail may be the rest of it, so it goes too. - sanitized_tail, tail_rewrote_last = ( - (_REDACTED_VALUE, True) - if head_rewrote_last and fragment_tail - else _sanitize_fragment_part(fragment_tail, nested=nested) + fragment, head_rewrote_last, tail_rewrote_last = _sanitize_url_fragment( + url.fragment, nested=nested ) - fragment = sanitized_head + separator + sanitized_tail netloc = _redact_userinfo(url.netloc) if (netloc, sanitized_query, fragment) == ( url.netloc, @@ -806,3 +794,24 @@ def _build_captured_mcp_arguments( continue captured[key] = sanitize_captured_value(value) return captured + + +def _sanitize_url_fragment( + fragment_value: str, *, nested: bool +) -> tuple[str, bool, bool]: + # A `?` inside a fragment splits it in two, and each half stands on its + # own: `#/callback?token=x` is text then fields, `#a=1?b=2` is fields + # twice. Shape cannot tell them apart — `#/token=x` looks like a route and + # is a field list — so each half is judged only on whether it holds a `=`. + head, separator, fragment_tail = fragment_value.partition("?") + sanitized_head, head_rewrote_last = _sanitize_fragment_part(head, nested=nested) + # A `?` can fall inside a credential (`#password=pre?fix`), and nothing + # here can tell that from a real boundary. When the head ends in a value + # just redacted, the tail may be the rest of it, so it goes too. + sanitized_tail, tail_rewrote_last = ( + (_REDACTED_VALUE, True) + if head_rewrote_last and fragment_tail + else _sanitize_fragment_part(fragment_tail, nested=nested) + ) + fragment = sanitized_head + separator + sanitized_tail + return fragment, head_rewrote_last, tail_rewrote_last diff --git a/posthog/mcp/_tool_input.py b/posthog/mcp/_tool_input.py index 60135a6f0..006b876f7 100644 --- a/posthog/mcp/_tool_input.py +++ b/posthog/mcp/_tool_input.py @@ -34,10 +34,7 @@ def _root_reference_chain(schema: Dict[str, Any]) -> Optional[List[Dict[str, Any return nodes if not isinstance(ref, str) or not ref.startswith("#/"): return None - target: Any = schema - for part in ref[2:].split("/"): - key = part.replace("~1", "/").replace("~0", "~") - target = target.get(key) if isinstance(target, dict) else None + target = _resolve_local_pointer(schema, ref) if not isinstance(target, dict) or any(target is node for node in nodes): return None nodes.append(target) @@ -75,24 +72,12 @@ def _aliases_used( return [] used = [] for canonical, names in aliases.items(): - if ( - not isinstance(canonical, str) - or canonical in input_value - or len(canonical) > _MAX_KEY_LENGTH - or not can_record(canonical) - ): + if not _can_record_canonical(canonical, input_value, can_record): continue names_value: Any = names if not isinstance(names_value, (list, tuple)): continue - alias = next( - ( - name - for name in names_value - if isinstance(name, str) and name in input_value - ), - None, - ) + alias = _first_present_alias(names_value, input_value) if alias and len(alias) <= _MAX_KEY_LENGTH and can_record(alias): used.append(f"{alias}:{canonical}") return sorted(used)[:_MAX_INPUT_KEYS] @@ -114,44 +99,13 @@ def get_tool_input_properties( resolved = options or ToolInputOptions() properties = _declared_properties(input_schema) known = set(properties) | set(_alias_names(resolved.input_aliases)) - keys = [ - key - for key in input_value - if isinstance(key, str) and (key in known or key not in _ANALYTICS_KEYS) - ] - declared: List[str] = [] - undeclared: List[str] = [] - has_redacted = False - recording_decisions: Dict[tuple[str, bool], bool] = {} - - def can_record(key: str, is_declared: bool) -> bool: - cache_key = (key, is_declared) - if cache_key in recording_decisions: - return recording_decisions[cache_key] - record = is_declared - if resolved.should_record_input_key is not None: - try: - record = ( - resolved.should_record_input_key(key, {"declared": is_declared}) - is True - ) - except Exception: # noqa: BLE001 - analytics callbacks fail closed - record = False - recording_decisions[cache_key] = record - return record - - for key in keys: - is_declared = key in known - if len(key) <= _MAX_KEY_LENGTH and can_record(key, is_declared): - (declared if is_declared else undeclared).append(key) - else: - has_redacted = True - - visible = [*sorted(declared), *sorted(undeclared)][:_MAX_INPUT_KEYS] - if has_redacted and len(visible) < _MAX_INPUT_KEYS: - visible.append("[redacted]") + keys = _input_keys(input_value, known) + recorder = _InputKeyRecorder(resolved) + visible = _visible_input_keys(keys, known, recorder) aliases_used = _aliases_used( - resolved.input_aliases, input_value, lambda key: can_record(key, True) + resolved.input_aliases, + input_value, + lambda key: recorder.can_record(key, True), ) result: JsonRecord = { PostHogMCPAnalyticsProperty.INPUT_KEYS: visible, @@ -161,3 +115,77 @@ def can_record(key: str, is_declared: bool) -> bool: return result except Exception: # noqa: BLE001 - analytics must not change tool dispatch return {} + + +def _resolve_local_pointer(schema: Dict[str, Any], ref: str) -> Any: + target: Any = schema + for part in ref[2:].split("/"): + key = part.replace("~1", "/").replace("~0", "~") # JSON Pointer escapes + target = target.get(key) if isinstance(target, dict) else None + return target + + +def _first_present_alias(names: Any, input_value: Dict[str, Any]) -> Optional[str]: + return next( + (name for name in names if isinstance(name, str) and name in input_value), None + ) + + +def _can_record_canonical( + canonical: Any, input_value: Dict[str, Any], can_record: Callable[[str], bool] +) -> bool: + return ( + isinstance(canonical, str) + and canonical not in input_value + and len(canonical) <= _MAX_KEY_LENGTH + and can_record(canonical) + ) + + +def _visible_input_keys( + keys: List[str], known: set, recorder: _InputKeyRecorder +) -> List[str]: + declared: List[str] = [] + undeclared: List[str] = [] + has_redacted = False + for key in keys: + is_declared = key in known + if len(key) <= _MAX_KEY_LENGTH and recorder.can_record(key, is_declared): + (declared if is_declared else undeclared).append(key) + else: + has_redacted = True + + visible = [*sorted(declared), *sorted(undeclared)][:_MAX_INPUT_KEYS] + if has_redacted and len(visible) < _MAX_INPUT_KEYS: + visible.append("[redacted]") + return visible + + +def _input_keys(input_value: Dict[str, Any], known: set) -> List[str]: + return [ + key + for key in input_value + if isinstance(key, str) and (key in known or key not in _ANALYTICS_KEYS) + ] + + +class _InputKeyRecorder: + def __init__(self, options: ToolInputOptions) -> None: + self.options = options + self.decisions: Dict[tuple[str, bool], bool] = {} + + def can_record(self, key: str, is_declared: bool) -> bool: + cache_key = (key, is_declared) + if cache_key in self.decisions: + return self.decisions[cache_key] + record = is_declared + if self.options.should_record_input_key is not None: + try: + record = ( + self.options.should_record_input_key(key, {"declared": is_declared}) + is True + ) + except Exception: # noqa: BLE001 - analytics callbacks fail closed + record = False + self.decisions[cache_key] = record + return record diff --git a/posthog/mcp/_truncation.py b/posthog/mcp/_truncation.py index 382c64639..67fa0b939 100644 --- a/posthog/mcp/_truncation.py +++ b/posthog/mcp/_truncation.py @@ -78,12 +78,7 @@ def _visit( if isinstance(value, bool): # before int — bool is an int subclass return value if isinstance(value, (int, float)): - if isinstance(value, float): - if math.isnan(value): - return "[NaN]" - if math.isinf(value): - return "[Infinity]" if value > 0 else "[-Infinity]" - return value + return _normalize_number(value) if isinstance(value, str): return truncate_text(value, max_string_length) if isinstance(value, datetime): @@ -91,33 +86,9 @@ def _visit( if callable(value): return f"[Function: {getattr(value, '__name__', '') or ''}]" - if isinstance(value, (list, tuple)): - oid = id(value) - if oid in memo: - return "[Circular ~]" - if remaining_depth <= 0: - return "[Array]" - memo.add(oid) - result: Any = _visit_array( - list(value), remaining_depth - 1, max_breadth, max_string_length, memo - ) - memo.discard(oid) - return result - - if isinstance(value, dict): - oid = id(value) - if oid in memo: - return "[Circular ~]" - if remaining_depth <= 0: - return "[Object]" - memo.add(oid) - result = _visit_object( - value, remaining_depth - 1, max_breadth, max_string_length, memo - ) - memo.discard(oid) - return result - - return str(value) + return _visit_container( + value, remaining_depth, max_breadth, max_string_length, memo + ) def _visit_array( @@ -255,12 +226,7 @@ def _collect_string_paths( return if isinstance(obj, dict): for key, value in obj.items(): - if ( - not current_path - and key == "server_build" - and isinstance(value, str) - and len(value) <= MAX_SERVER_BUILD_LENGTH - ): + if _is_preserved_build(current_path, key, value): continue _collect_string_paths(value, current_path + [str(key)], results) @@ -373,3 +339,57 @@ def truncate_event(event: Dict[str, Any]) -> Dict[str, Any]: # Layer 3: size-targeted normalization return _truncate_to_size(result) + + +def _visit_container( + value: Any, + remaining_depth: int, + max_breadth: int, + max_string_length: int, + memo: set, +) -> Any: + if isinstance(value, (list, tuple)): + oid = id(value) + if oid in memo: + return "[Circular ~]" + if remaining_depth <= 0: + return "[Array]" + memo.add(oid) + result: Any = _visit_array( + list(value), remaining_depth - 1, max_breadth, max_string_length, memo + ) + memo.discard(oid) + return result + + if isinstance(value, dict): + oid = id(value) + if oid in memo: + return "[Circular ~]" + if remaining_depth <= 0: + return "[Object]" + memo.add(oid) + result = _visit_object( + value, remaining_depth - 1, max_breadth, max_string_length, memo + ) + memo.discard(oid) + return result + + return str(value) + + +def _normalize_number(value: Any) -> Any: + if isinstance(value, float): + if math.isnan(value): + return "[NaN]" + if math.isinf(value): + return "[Infinity]" if value > 0 else "[-Infinity]" + return value + + +def _is_preserved_build(current_path: List[str], key: Any, value: Any) -> bool: + return ( + not current_path + and key == "server_build" + and isinstance(value, str) + and len(value) <= MAX_SERVER_BUILD_LENGTH + ) diff --git a/posthog/mcp/asgi.py b/posthog/mcp/asgi.py index a66b41113..c450f026a 100644 --- a/posthog/mcp/asgi.py +++ b/posthog/mcp/asgi.py @@ -116,15 +116,7 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: if scope.get("type") != "http": await self.app(scope, receive, send) return - # Tell the in-process fallback that this request's handshake is captured - # here. Two stacked middlewares see the same request; the outer one owns it. - # The SSE transport stays unmarked: its session is in a query parameter, so - # its first request captures the handshake, in the session its calls use. - observer = None - if HANDSHAKE_OBSERVER_KEY not in scope and not _is_sse(scope): - observer = self._analytics() - if observer is not None: - scope[HANDSHAKE_OBSERVER_KEY] = observer + observer = self._handshake_observer(scope) try: incoming_header = read_mcp_session_header(_headers_dict(scope)) @@ -148,15 +140,7 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: # Guard the whole sniff+mint path: on any failure (a raising receive() or # adversarial body) the request passes through untouched, per the contract. - token = None - initialize = None - try: - body, receive = await _buffer_body(receive) - initialize = _parse_initialize(body) - if initialize is not None: - token = _mint_token(initialize) - except Exception as error: # noqa: BLE001 - log(f"PostHog MCP session middleware: mint failed - {error}") + initialize, token, receive = await _sniff_initialize(receive) if initialize is None: await self.app(scope, receive, send) @@ -176,6 +160,19 @@ async def __call__(self, scope: Any, receive: Any, send: Any) -> None: await self.app(scope, receive, send) + def _handshake_observer(self, scope: Any) -> Optional[MCPAnalyticsData]: + # Tell the in-process fallback that this request's handshake is captured + # here. Two stacked middlewares see the same request; the outer one owns it. + # The SSE transport stays unmarked: its session is in a query parameter, so + # its first request captures the handshake, in the session its calls use. + observer = None + if HANDSHAKE_OBSERVER_KEY not in scope and not _is_sse(scope): + observer = self._analytics() + if observer is not None: + scope[HANDSHAKE_OBSERVER_KEY] = observer + + return observer + def _is_sse(scope: Any) -> bool: """The deprecated SSE transport, which keys its session on ``?session_id=``.""" @@ -482,3 +479,17 @@ def factory(*args: Any, **kwargs: Any) -> Any: setattr(factory, _AUTOWIRED, True) return factory + + +async def _sniff_initialize(receive: Any) -> tuple[Any, Optional[str], Any]: + token = None + initialize = None + try: + body, receive = await _buffer_body(receive) + initialize = _parse_initialize(body) + if initialize is not None: + token = _mint_token(initialize) + except Exception as error: # noqa: BLE001 + log(f"PostHog MCP session middleware: mint failed - {error}") + + return initialize, token, receive diff --git a/posthog/mcp/feedback.py b/posthog/mcp/feedback.py index 5cb732b6c..080359db4 100644 --- a/posthog/mcp/feedback.py +++ b/posthog/mcp/feedback.py @@ -143,17 +143,7 @@ def get_feedback_tool_descriptor( extra_properties = options.extra_properties or {} extra_required = options.extra_required or [] - for key in extra_properties: - if key in _RESERVED_EXTRA_PROPERTY_KEYS: - raise ValueError( - f'collect_feedback extra_properties key "{key}" collides with a core or ' - "SDK-reserved send_feedback field. Rename it." - ) - for key in extra_required: - if key not in extra_properties: - raise ValueError( - f'collect_feedback extra_required key "{key}" is not declared in extra_properties.' - ) + _validate_feedback_extras(extra_properties, extra_required) # Deep copies: the host may reuse or mutate the fragments it handed us, and # schema injection mutates the advertised descriptor. @@ -182,18 +172,7 @@ def _matches_extra_schema(value: Any, schema: Dict[str, Any]) -> bool: ``enum`` — the same advisory-schema enforcement the core fields get, so ``extras`` only ever holds schema-conforming values and a misbehaving agent shows up as absence rather than as an unexpected shape in the host's handler.""" - if isinstance(value, list): - actual = "array" - elif value is None: - actual = "null" - elif isinstance(value, bool): - actual = "boolean" - elif isinstance(value, (int, float)): - actual = "number" - elif isinstance(value, str): - actual = "string" - else: - actual = "object" + actual = _json_type(value) declared = schema.get("type") if declared == "integer" and actual == "number": # A JSON Schema `integer` also accepts a whole-valued float (`3.0`), but @@ -219,12 +198,7 @@ def parse_feedback_report( ``type``/``enum`` are lifted into ``extras`` — mismatches and anything the agent invented reach the handler via ``raw`` only and are never captured.""" raw = args or {} - declared = (options.extra_properties if options is not None else None) or {} - extras = { - key: raw[key] - for key, schema in declared.items() - if raw.get(key) is not None and _matches_extra_schema(raw[key], schema) - } + extras = _parse_feedback_extras(raw, options) feedback_type = raw.get("feedback_type") sentiment = raw.get("sentiment") task_completed = raw.get("task_completed") @@ -309,10 +283,7 @@ def build_feedback_event_properties(report: FeedbackReport) -> JsonRecord: properties[PostHogMCPAnalyticsProperty.FEEDBACK_TASK_COMPLETED] = ( report.task_completed ) - for key, value in report.extras.items(): - captured = _capture_extra_value(value) - if captured is not None: - properties[f"$mcp_feedback_{key}"] = captured + _add_feedback_extras(report, properties) return properties @@ -354,3 +325,54 @@ async def handle_feedback( f"({type(error).__name__}); returning the default acknowledgement" ) return _SEND_FEEDBACK_RESULT_TEXT + + +def _validate_feedback_extras( + extra_properties: Dict[str, Any], extra_required: list[str] +) -> None: + for key in extra_properties: + if key in _RESERVED_EXTRA_PROPERTY_KEYS: + raise ValueError( + f'collect_feedback extra_properties key "{key}" collides with a core or ' + "SDK-reserved send_feedback field. Rename it." + ) + for key in extra_required: + if key not in extra_properties: + raise ValueError( + f'collect_feedback extra_required key "{key}" is not declared in extra_properties.' + ) + + +def _json_type(value: Any) -> str: + if isinstance(value, list): + actual = "array" + elif value is None: + actual = "null" + elif isinstance(value, bool): + actual = "boolean" + elif isinstance(value, (int, float)): + actual = "number" + elif isinstance(value, str): + actual = "string" + else: + actual = "object" + return actual + + +def _parse_feedback_extras( + raw: JsonRecord, options: Optional[CollectFeedbackOptions] +) -> JsonRecord: + declared = (options.extra_properties if options is not None else None) or {} + extras = { + key: raw[key] + for key, schema in declared.items() + if raw.get(key) is not None and _matches_extra_schema(raw[key], schema) + } + return extras + + +def _add_feedback_extras(report: FeedbackReport, properties: JsonRecord) -> None: + for key, value in report.extras.items(): + captured = _capture_extra_value(value) + if captured is not None: + properties[f"$mcp_feedback_{key}"] = captured diff --git a/posthog/mcp/posthog_mcp.py b/posthog/mcp/posthog_mcp.py index 2e1c4d3e5..09dbe7f40 100644 --- a/posthog/mcp/posthog_mcp.py +++ b/posthog/mcp/posthog_mcp.py @@ -469,40 +469,46 @@ def prepare_tool_list( # descriptor, which a host re-preparing an already-prepared list hands # straight back. Warning about that would be warning about ourselves. if report_missing: - name = self._missing_capability_tool_name - existing = _find_tool(prepared, name) - if existing is not None and not self._is_sdk_virtual_tool(existing): - self._warn_virtual_tool_collision( - VIRTUAL_TOOL_MISSING_CAPABILITY, - name, - 'PostHogMCP(missing_capability_tool_name="...")', - ) - elif existing is None: - prepared.append(build_report_missing_descriptor(name)) + self._append_missing_capability(prepared) if collect_feedback and self._collect_feedback is not None: - name = self._feedback_tool_name - existing = _find_tool(prepared, name) - # Both virtual tools under one name: missing-capability wins in - # `prepare_tool_call`, so advertising this one too would dead-letter - # the feedback path. Same precedence as `instrument()`. - duplicate = name == self._missing_capability_tool_name - if duplicate or ( - existing is not None and not self._is_sdk_virtual_tool(existing) - ): - self._warn_virtual_tool_collision( - VIRTUAL_TOOL_FEEDBACK, - name, - 'PostHogMCP(collect_feedback=CollectFeedbackOptions(tool_name="..."))', - variant="duplicate" if duplicate else "blocked", - ) - elif existing is None: - prepared.append(get_feedback_tool_descriptor(self._collect_feedback)) + self._append_feedback(prepared) # Read ownership before any analytics field lands on the schemas. conversation_ownership = _collect_conversation_ownership(prepared) prepared = self._inject_models(prepared) prepared = self._inject_conversation(prepared, conversation_ownership) return prepared + def _append_missing_capability(self, prepared: List[Any]) -> None: + name = self._missing_capability_tool_name + existing = _find_tool(prepared, name) + if existing is not None and not self._is_sdk_virtual_tool(existing): + self._warn_virtual_tool_collision( + VIRTUAL_TOOL_MISSING_CAPABILITY, + name, + 'PostHogMCP(missing_capability_tool_name="...")', + ) + elif existing is None: + prepared.append(build_report_missing_descriptor(name)) + + def _append_feedback(self, prepared: List[Any]) -> None: + name = self._feedback_tool_name + existing = _find_tool(prepared, name) + # Both virtual tools under one name: missing-capability wins in + # `prepare_tool_call`, so advertising this one too would dead-letter + # the feedback path. Same precedence as `instrument()`. + duplicate = name == self._missing_capability_tool_name + if duplicate or ( + existing is not None and not self._is_sdk_virtual_tool(existing) + ): + self._warn_virtual_tool_collision( + VIRTUAL_TOOL_FEEDBACK, + name, + 'PostHogMCP(collect_feedback=CollectFeedbackOptions(tool_name="..."))', + variant="duplicate" if duplicate else "blocked", + ) + elif existing is None: + prepared.append(get_feedback_tool_descriptor(self._collect_feedback)) + def _warn_virtual_tool_collision( self, kind: str, @@ -571,6 +577,63 @@ def prepare_tool_call( if isinstance(raw_context, str) and raw_context.strip() else None ) + prepared_args, llm_model, llm_model_source = self._prepare_model( + name, args, request_meta, original_tool + ) + + prepared_args, session_id, conversation_id, conversation_state = ( + self._prepare_conversation( + name, args, prepared_args, original_tool, session_id + ) + ) + # A supplied `original_tool` is a real application tool by this name (it + # comes from the host's own list, which never holds a virtual tool), so + # the real tool wins — the stateless twin of the ownership check + # instrument() runs. Without it the name match stands, and the + # documented remedy for a collision is renaming PostHog's tool. + # The name check against missing-capability keeps one precedence when + # both tools share a name: without it both flags are true, and a + # dispatcher testing `is_feedback` first misroutes every call. + is_feedback = self._is_feedback_call(name, original_tool) + # Same guard for the missing-capability tool. Unlike feedback it has no + # constructor enable switch on this path (the name is always populated), + # so `original_tool` is the only ownership signal available here. + is_missing_capability = ( + name == self._missing_capability_tool_name and original_tool is None + ) + return PreparedToolCall( + args=prepared_args, + intent=intent, + intent_source="context_parameter" if intent else None, + llm_model=llm_model, + llm_model_source=llm_model_source, + is_missing_capability=is_missing_capability, + is_feedback=is_feedback, + feedback_report=( + parse_feedback_report(args, self._collect_feedback) + if is_feedback + else None + ), + session_id=session_id, + conversation_id=conversation_id, + _conversation_state=conversation_state, + ) + + def _is_feedback_call(self, name: str, original_tool: Any) -> bool: + return ( + self._collect_feedback is not None + and name == self._feedback_tool_name + and name != self._missing_capability_tool_name + and original_tool is None + ) + + def _prepare_model( + self, + name: str, + args: Optional[JsonRecord], + request_meta: Optional[JsonRecord], + original_tool: Any, + ) -> Tuple[Optional[JsonRecord], Optional[str], Optional[MCPAnalyticsModelSource]]: analytics_owns_model = False llm_model: Optional[str] = None llm_model_source: Optional[MCPAnalyticsModelSource] = None @@ -588,6 +651,27 @@ def prepare_tool_call( if analytics_owns_model: prepared_args = _strip_args(prepared_args, "llm_model") + return prepared_args, llm_model, llm_model_source + + def _owns_conversation_parameter( + self, ownership: Optional[_ConversationOwnership] + ) -> bool: + return ( + self._enable_conversation_id + and ownership is not None + and ownership.conversation_id + ) + + def _prepare_conversation( + self, + name: str, + args: Optional[JsonRecord], + prepared_args: Optional[JsonRecord], + original_tool: Any, + session_id: Optional[str], + ) -> Tuple[ + Optional[JsonRecord], Optional[str], Optional[str], PreparedConversationState + ]: ownership = ( _conversation_ownership_of(original_tool) if original_tool is not None @@ -606,53 +690,21 @@ def prepare_tool_call( conversation_id = None if minted else resolved_conversation_id if conversation_id: session_id = derive_session_id_from_conversation(conversation_id) - if ( - self._enable_conversation_id - and ownership is not None - and ownership.conversation_id - ): + if self._owns_conversation_parameter(ownership): prepared_args = _strip_args(prepared_args, "conversation_id") - # A supplied `original_tool` is a real application tool by this name (it - # comes from the host's own list, which never holds a virtual tool), so - # the real tool wins — the stateless twin of the ownership check - # instrument() runs. Without it the name match stands, and the - # documented remedy for a collision is renaming PostHog's tool. - # The name check against missing-capability keeps one precedence when - # both tools share a name: without it both flags are true, and a - # dispatcher testing `is_feedback` first misroutes every call. - is_feedback = ( - self._collect_feedback is not None - and name == self._feedback_tool_name - and name != self._missing_capability_tool_name - and original_tool is None - ) - # Same guard for the missing-capability tool. Unlike feedback it has no - # constructor enable switch on this path (the name is always populated), - # so `original_tool` is the only ownership signal available here. - is_missing_capability = ( - name == self._missing_capability_tool_name and original_tool is None - ) - return PreparedToolCall( - args=prepared_args, - intent=intent, - intent_source="context_parameter" if intent else None, - llm_model=llm_model, - llm_model_source=llm_model_source, - is_missing_capability=is_missing_capability, - is_feedback=is_feedback, - feedback_report=( - parse_feedback_report(args, self._collect_feedback) - if is_feedback - else None - ), - session_id=session_id, - conversation_id=conversation_id, - _conversation_state=PreparedConversationState( - minted_conversation_id=minted_conversation_id, - output_instructions=self._enable_conversation_id - and ownership is not None - and ownership.output_instructions, - ), + state = self._prepared_conversation_state(ownership, minted_conversation_id) + return prepared_args, session_id, conversation_id, state + + def _prepared_conversation_state( + self, + ownership: Optional[_ConversationOwnership], + minted_conversation_id: Optional[str], + ) -> PreparedConversationState: + return PreparedConversationState( + minted_conversation_id=minted_conversation_id, + output_instructions=self._enable_conversation_id + and ownership is not None + and ownership.output_instructions, ) def prepare_tool_result( @@ -682,10 +734,7 @@ def prepare_tool_result( prepared, handle ) if minted_conversation_id: - injected = inject_prompt_back(prepared, handle) - if injected is not prepared: - delivered = True - prepared = injected + prepared, delivered = _deliver_minted_handle(prepared, handle, delivered) if delivered: conversation_id = handle session_id = derive_session_id_from_conversation(handle) @@ -981,3 +1030,12 @@ def _tool_schema(tool: Any) -> Optional[Dict[str, Any]]: if schema is None: schema = getattr(tool, "inputSchema", None) return schema if isinstance(schema, dict) else None + + +def _deliver_minted_handle( + prepared: Any, handle: str, delivered: bool +) -> Tuple[Any, bool]: + injected = inject_prompt_back(prepared, handle) + if injected is not prepared: + delivered = True + return injected, delivered diff --git a/posthog/mcp/session_token.py b/posthog/mcp/session_token.py index ac2b38593..a562b9b27 100644 --- a/posthog/mcp/session_token.py +++ b/posthog/mcp/session_token.py @@ -86,10 +86,31 @@ def decode_session_id(value: Any) -> Optional[SessionTokenPayload]: return None if not isinstance(parsed, dict): return None - sid = parsed.get("sid") - if not isinstance(sid, str) or not sid or len(sid) > _MAX_SESSION_ID_LENGTH: + return _decode_token_payload(parsed) + + +def read_mcp_session_header(headers: Any) -> Optional[str]: + """Read the ``mcp-session-id`` value off a headers mapping. + + Handles case-insensitive keys (transports lowercase them, but hand-built + mappings may not) and list-valued headers, and trims whitespace. + """ + if headers is None: return None - payload = SessionTokenPayload(session_id=sid) + value = _session_header_value(headers) + first = value[0] if isinstance(value, (list, tuple)) and value else value + if not isinstance(first, str): + return None + trimmed = first.strip() + return trimmed or None + + +def _base64url_to_bytes(value: str) -> bytes: + padding = "=" * (-len(value) % 4) + return base64.urlsafe_b64decode(value + padding) + + +def _decode_client_fields(parsed: dict, payload: SessionTokenPayload) -> None: # A bad cn/cv/pv just means no client info -- it does not reject the token. cn = parsed.get("cn") if isinstance(cn, str) and cn: @@ -100,17 +121,9 @@ def decode_session_id(value: Any) -> Optional[SessionTokenPayload]: pv = parsed.get("pv") if isinstance(pv, str) and pv: payload.protocol_version = pv[:_MAX_CLIENT_FIELD_LENGTH] - return payload - -def read_mcp_session_header(headers: Any) -> Optional[str]: - """Read the ``mcp-session-id`` value off a headers mapping. - Handles case-insensitive keys (transports lowercase them, but hand-built - mappings may not) and list-valued headers, and trims whitespace. - """ - if headers is None: - return None +def _session_header_value(headers: Any) -> Any: value: Any = None # Mapping-like: prefer a direct get, else scan case-insensitively. get = getattr(headers, "get", None) @@ -125,13 +138,13 @@ def read_mcp_session_header(headers: Any) -> Optional[str]: if isinstance(key, str) and key.lower() == MCP_SESSION_HEADER: value = candidate break - first = value[0] if isinstance(value, (list, tuple)) and value else value - if not isinstance(first, str): - return None - trimmed = first.strip() - return trimmed or None + return value -def _base64url_to_bytes(value: str) -> bytes: - padding = "=" * (-len(value) % 4) - return base64.urlsafe_b64decode(value + padding) +def _decode_token_payload(parsed: dict[str, Any]) -> Optional[SessionTokenPayload]: + sid = parsed.get("sid") + if not isinstance(sid, str) or not sid or len(sid) > _MAX_SESSION_ID_LENGTH: + return None + payload = SessionTokenPayload(session_id=sid) + _decode_client_fields(parsed, payload) + return payload diff --git a/posthog/metrics_capture.py b/posthog/metrics_capture.py index 3d8bc15d4..86dd5c5db 100644 --- a/posthog/metrics_capture.py +++ b/posthog/metrics_capture.py @@ -102,16 +102,7 @@ def _to_otlp_any_value(value: Any) -> dict: if isinstance(value, int): return {"intValue": value} if isinstance(value, float): - # proto3 JSON has no representation for non-finite floats; encode the proto3 - # literal strings (not Python's "inf"/"nan") so both SDKs emit identical bytes. - if not math.isfinite(value): - if math.isnan(value): - return {"stringValue": "NaN"} - return {"stringValue": "Infinity" if value > 0 else "-Infinity"} - # Integral floats encode as intValue, matching the JS Number.isInteger branch. - if value.is_integer(): - return {"intValue": int(value)} - return {"doubleValue": value} + return _float_otlp_value(value) if isinstance(value, str): return {"stringValue": value} if isinstance(value, (list, tuple)): @@ -343,18 +334,9 @@ def __init__(self, client, config: Optional[dict] = None): self._flush_interval: float = _validated_flush_interval( config.get("flush_interval", _DEFAULT_FLUSH_INTERVAL_SECONDS) ) - max_series = config.get("max_series_per_flush", _DEFAULT_MAX_SERIES_PER_FLUSH) - if ( - not isinstance(max_series, int) - or isinstance(max_series, bool) - or max_series <= 0 - ): - log.warning( - "Ignoring metrics max_series_per_flush %r: expected a positive integer", - max_series, - ) - max_series = _DEFAULT_MAX_SERIES_PER_FLUSH - self._max_series_per_flush: int = max_series + self._max_series_per_flush: int = _validated_max_series( + config.get("max_series_per_flush", _DEFAULT_MAX_SERIES_PER_FLUSH) + ) self._before_send: Optional[Callable] = _validated_before_send( config.get("before_send") ) @@ -450,17 +432,8 @@ def _capture( "attributes": attributes, } if self._before_send is not None: - try: - filtered = self._before_send(sample) - except Exception as e: - log.error("Error in metrics before_send: %s", e) - return - if not filtered: - return - if not isinstance(filtered, dict): - log.warning( - "Dropping metric: before_send must return the sample dict or a falsy value" - ) + filtered = _filter_metric_sample(self._before_send, sample) + if filtered is None: return sample_dict: dict[str, Any] = filtered # Defaults keep the static types closed; a hook that removed the field @@ -471,26 +444,7 @@ def _capture( unit = sample_dict.get("unit") attributes = sample_dict.get("attributes") - if metric_type not in _VALID_METRIC_TYPES: - log.warning( - "Dropping metric '%s': unknown metric type '%s'", name, metric_type - ) - return - - if not name or not isinstance(name, str): - log.warning("Dropping metric with empty name") - return - if ( - not isinstance(value, (int, float)) - or isinstance(value, bool) - or not math.isfinite(value) - ): - log.warning("Dropping metric '%s': value must be a finite number", name) - return - if metric_type == "count" and value < 0: - log.warning( - "Dropping count '%s': counters are monotonic, value must be >= 0", name - ) + if not _valid_metric_sample(metric_type, name, value): return if attributes: @@ -503,45 +457,53 @@ def _capture( with self._lock: self._reset_after_fork_locked() - state = self._series.get(key) + state = self._admit_series_locked(key, name, metric_type, unit, attributes) if state is None: - if len(self._series) >= self._max_series_per_flush: - if not self._series_cap_warned: - self._series_cap_warned = True - log.warning( - "Metric series cap reached (%s per flush window); dropping new series " - "until the next flush. Reduce attribute cardinality.", - self._max_series_per_flush, - ) - return - state = _SeriesState(name, metric_type, unit, attributes) - self._series[key] = state - - # Bookkeeping only for admitted samples, so name-cardinality misuse (IDs in - # metric names) can't grow this map past the series cap. - seen_type = self._type_by_name.get(name) - if seen_type is None: - self._type_by_name[name] = metric_type - elif seen_type != metric_type and name not in self._type_collision_warned: - self._type_collision_warned.add(name) - log.warning( - "Metric name '%s' is already used as a %s; recording it as a %s too will blend " - "both series in charts. Use a distinct name.", - name, - seen_type, - metric_type, - ) + return self._fold(state, float(value)) - if trace_context is not None: - state.exemplar = { - "trace_id": trace_context[0], - "span_id": trace_context[1], - "value": float(value), - "time_ms": int(time.time() * 1000), - } + _record_exemplar(state, trace_context, value) self._arm_flush_timer() + def _admit_series_locked( + self, + key: str, + name: str, + metric_type: str, + unit: Optional[str], + attributes: Optional[dict], + ) -> Optional[_SeriesState]: + state = self._series.get(key) + if state is None: + if len(self._series) >= self._max_series_per_flush: + if not self._series_cap_warned: + self._series_cap_warned = True + log.warning( + "Metric series cap reached (%s per flush window); dropping new series " + "until the next flush. Reduce attribute cardinality.", + self._max_series_per_flush, + ) + return None + state = _SeriesState(name, metric_type, unit, attributes) + self._series[key] = state + + # Bookkeeping only for admitted samples, so name-cardinality misuse (IDs in + # metric names) can't grow this map past the series cap. + seen_type = self._type_by_name.get(name) + if seen_type is None: + self._type_by_name[name] = metric_type + elif seen_type != metric_type and name not in self._type_collision_warned: + self._type_collision_warned.add(name) + log.warning( + "Metric name '%s' is already used as a %s; recording it as a %s too will blend " + "both series in charts. Use a distinct name.", + name, + seen_type, + metric_type, + ) + + return state + def _reinit_after_fork(self) -> None: # Runs in a forked child (via the client's os.register_at_fork hook) before # user code. The inherited locks may be held by parent threads that do not @@ -734,23 +696,7 @@ def _merge_window_back(self, window: dict) -> None: continue self._series[key] = old continue - current.window_start_ms = min(current.window_start_ms, old.window_start_ms) - # The live window's exemplar is newer; the old one only fills a gap. - if current.exemplar is None: - current.exemplar = old.exemplar - if current.type == "count": - current.total = (current.total or 0.0) + (old.total or 0.0) - elif current.type == "histogram" and old.hist: - if current.hist is None: - current.hist = old.hist - else: - current.hist["count"] += old.hist["count"] - current.hist["sum"] += old.hist["sum"] - current.hist["min"] = min(current.hist["min"], old.hist["min"]) - current.hist["max"] = max(current.hist["max"], old.hist["max"]) - for i, count in enumerate(old.hist["bucket_counts"]): - current.hist["bucket_counts"][i] += count - # Gauge: the live window's value is newer — keep it. + _merge_series_state(current, old) if dropped: log.warning( "Dropped %s unsent metric series while merging a failed flush back (series cap %s)", @@ -795,64 +741,172 @@ def _build_metrics(self, window: dict) -> list: metric_key = (state.type, state.name, state.unit or "") metric = by_metric.get(metric_key) if metric is None: - metric = {"name": state.name} - if state.unit: - metric["unit"] = state.unit - if state.type == "count": - metric["sum"] = { - "aggregationTemporality": _OTLP_TEMPORALITY_DELTA, - "isMonotonic": True, - "dataPoints": [], - } - elif state.type == "gauge": - metric["gauge"] = {"dataPoints": []} - else: - metric["histogram"] = { - "aggregationTemporality": _OTLP_TEMPORALITY_DELTA, - "dataPoints": [], - } + metric = _metric_descriptor(state) by_metric[metric_key] = metric - attributes = _to_otlp_key_value_list(state.attributes or {}) - start_nano = _ms_to_unix_nano(state.window_start_ms) - exemplars = _build_exemplars(state.exemplar) - - if state.type == "count": - metric["sum"]["dataPoints"].append( - { - "attributes": attributes, - "startTimeUnixNano": start_nano, - "timeUnixNano": now_nano, - "asDouble": state.total or 0.0, - **exemplars, - } - ) - elif state.type == "gauge": - metric["gauge"]["dataPoints"].append( - { - "attributes": attributes, - "timeUnixNano": now_nano, - "asDouble": state.last or 0.0, - **exemplars, - } - ) - elif state.hist is not None: - # Encoding pinned by the ingest's JSON deserializer: nano timestamps are decimal - # strings, but count/bucketCounts are plain JSON numbers — string-encoded u64s in - # those fields are silently dropped upstream (opentelemetry-rust#3328). - metric["histogram"]["dataPoints"].append( - { - "attributes": attributes, - "startTimeUnixNano": start_nano, - "timeUnixNano": now_nano, - "count": state.hist["count"], - "sum": state.hist["sum"], - "min": state.hist["min"], - "max": state.hist["max"], - "bucketCounts": state.hist["bucket_counts"], - "explicitBounds": DEFAULT_HISTOGRAM_BOUNDS, - **exemplars, - } - ) + _append_metric_point(metric, state, now_nano) return list(by_metric.values()) + + +def _float_otlp_value(value: float) -> dict: + # proto3 JSON has no representation for non-finite floats; encode the proto3 + # literal strings (not Python's "inf"/"nan") so both SDKs emit identical bytes. + if not math.isfinite(value): + if math.isnan(value): + return {"stringValue": "NaN"} + return {"stringValue": "Infinity" if value > 0 else "-Infinity"} + # Integral floats encode as intValue, matching the JS Number.isInteger branch. + if value.is_integer(): + return {"intValue": int(value)} + return {"doubleValue": value} + + +def _validated_max_series(max_series: Any) -> int: + if ( + not isinstance(max_series, int) + or isinstance(max_series, bool) + or max_series <= 0 + ): + log.warning( + "Ignoring metrics max_series_per_flush %r: expected a positive integer", + max_series, + ) + max_series = _DEFAULT_MAX_SERIES_PER_FLUSH + return max_series + + +def _valid_metric_sample(metric_type: str, name: str, value: float) -> bool: + if metric_type not in _VALID_METRIC_TYPES: + log.warning("Dropping metric '%s': unknown metric type '%s'", name, metric_type) + return False + + if not name or not isinstance(name, str): + log.warning("Dropping metric with empty name") + return False + if ( + not isinstance(value, (int, float)) + or isinstance(value, bool) + or not math.isfinite(value) + ): + log.warning("Dropping metric '%s': value must be a finite number", name) + return False + if metric_type == "count" and value < 0: + log.warning( + "Dropping count '%s': counters are monotonic, value must be >= 0", name + ) + return False + + return True + + +def _filter_metric_sample(before_send: Callable, sample: dict) -> Optional[dict]: + try: + filtered = before_send(sample) + except Exception as e: + log.error("Error in metrics before_send: %s", e) + return None + if not filtered: + return None + if not isinstance(filtered, dict): + log.warning( + "Dropping metric: before_send must return the sample dict or a falsy value" + ) + return None + return filtered + + +def _merge_series_state(current: _SeriesState, old: _SeriesState) -> None: + current.window_start_ms = min(current.window_start_ms, old.window_start_ms) + # The live window's exemplar is newer; the old one only fills a gap. + if current.exemplar is None: + current.exemplar = old.exemplar + if current.type == "count": + current.total = (current.total or 0.0) + (old.total or 0.0) + elif current.type == "histogram" and old.hist: + if current.hist is None: + current.hist = old.hist + else: + current.hist["count"] += old.hist["count"] + current.hist["sum"] += old.hist["sum"] + current.hist["min"] = min(current.hist["min"], old.hist["min"]) + current.hist["max"] = max(current.hist["max"], old.hist["max"]) + for i, count in enumerate(old.hist["bucket_counts"]): + current.hist["bucket_counts"][i] += count + # Gauge: the live window's value is newer — keep it. + + +def _metric_descriptor(state: _SeriesState) -> dict: + metric: dict = {"name": state.name} + if state.unit: + metric["unit"] = state.unit + if state.type == "count": + metric["sum"] = { + "aggregationTemporality": _OTLP_TEMPORALITY_DELTA, + "isMonotonic": True, + "dataPoints": [], + } + elif state.type == "gauge": + metric["gauge"] = {"dataPoints": []} + else: + metric["histogram"] = { + "aggregationTemporality": _OTLP_TEMPORALITY_DELTA, + "dataPoints": [], + } + return metric + + +def _record_exemplar( + state: _SeriesState, trace_context: Optional[tuple[str, str]], value: float +) -> None: + if trace_context is not None: + state.exemplar = { + "trace_id": trace_context[0], + "span_id": trace_context[1], + "value": float(value), + "time_ms": int(time.time() * 1000), + } + + +def _append_metric_point(metric: dict, state: _SeriesState, now_nano: str) -> None: + attributes = _to_otlp_key_value_list(state.attributes or {}) + start_nano = _ms_to_unix_nano(state.window_start_ms) + exemplars = _build_exemplars(state.exemplar) + + if state.type == "count": + metric["sum"]["dataPoints"].append( + { + "attributes": attributes, + "startTimeUnixNano": start_nano, + "timeUnixNano": now_nano, + "asDouble": state.total or 0.0, + **exemplars, + } + ) + elif state.type == "gauge": + metric["gauge"]["dataPoints"].append( + { + "attributes": attributes, + "timeUnixNano": now_nano, + "asDouble": state.last or 0.0, + **exemplars, + } + ) + elif state.hist is not None: + # Encoding pinned by the ingest's JSON deserializer: nano timestamps are decimal + # strings, but count/bucketCounts are plain JSON numbers — string-encoded u64s in + # those fields are silently dropped upstream (opentelemetry-rust#3328). + metric["histogram"]["dataPoints"].append( + { + "attributes": attributes, + "startTimeUnixNano": start_nano, + "timeUnixNano": now_nano, + "count": state.hist["count"], + "sum": state.hist["sum"], + "min": state.hist["min"], + "max": state.hist["max"], + "bucketCounts": state.hist["bucket_counts"], + "explicitBounds": DEFAULT_HISTOGRAM_BOUNDS, + **exemplars, + } + ) diff --git a/posthog/request.py b/posthog/request.py index 52854f922..e7f179374 100644 --- a/posthog/request.py +++ b/posthog/request.py @@ -296,24 +296,7 @@ def _process_response( ) raise QuotaLimitError(res.status_code, "Feature flags quota limited") return response - retry_after = None - retry_after_header = res.headers.get("Retry-After") - if retry_after_header: - try: - retry_after = float(retry_after_header) - except (ValueError, TypeError): - try: - from email.utils import parsedate_to_datetime - - retry_after = max( - 0.0, - ( - parsedate_to_datetime(retry_after_header) - - datetime.now(timezone.utc) - ).total_seconds(), - ) - except (ValueError, TypeError): - pass + retry_after = _response_retry_after(res) try: payload = res.json() @@ -550,3 +533,26 @@ def default(self, obj: Any): def encode(self, o: Any) -> str: return super().encode(_replace_non_finite_floats(o)) + + +def _response_retry_after(res: requests.Response) -> Optional[float]: + retry_after = None + retry_after_header = res.headers.get("Retry-After") + if retry_after_header: + try: + retry_after = float(retry_after_header) + except (ValueError, TypeError): + try: + from email.utils import parsedate_to_datetime + + retry_after = max( + 0.0, + ( + parsedate_to_datetime(retry_after_header) + - datetime.now(timezone.utc) + ).total_seconds(), + ) + except (ValueError, TypeError): + pass + + return retry_after diff --git a/posthog/test/ai/anthropic/test_anthropic_converter.py b/posthog/test/ai/anthropic/test_anthropic_converter.py index aef31a03a..9294c465d 100644 --- a/posthog/test/ai/anthropic/test_anthropic_converter.py +++ b/posthog/test/ai/anthropic/test_anthropic_converter.py @@ -213,3 +213,95 @@ def test_stream_message_delta_thinking_tokens_not_double_counted(): ) assert acc.usage_stats["output_tokens"] == 50 assert acc.usage_stats["reasoning_tokens"] == 30 + + +def test_stream_keeps_block_indexes_thinking_private_from_text_and_cumulative_usage(): + from types import SimpleNamespace as NS + + from posthog.ai.anthropic._anthropic_stream import _AnthropicStreamAccumulator + + accumulator = _AnthropicStreamAccumulator() + events = [ + NS( + type="message_start", + message=NS(usage=NS(input_tokens=10, cache_read_input_tokens=1)), + ), + NS( + type="content_block_start", + index=0, + content_block=NS(type="thinking", thinking="", signature=""), + ), + NS( + type="content_block_delta", + index=0, + delta=NS(type="thinking_delta", thinking="reason"), + ), + NS( + type="content_block_delta", + index=0, + delta=NS(type="signature_delta", signature="sig"), + ), + NS(type="content_block_stop", index=0), + NS( + type="content_block_start", + index=1, + content_block=NS(type="server_tool_use", id="server-1"), + ), + NS(type="content_block_stop", index=1), + NS( + type="content_block_start", + index=2, + content_block=NS(type="tool_use", id="call-2", name="lookup"), + ), + NS( + type="content_block_delta", + index=2, + delta=NS(type="input_json_delta", partial_json='{"q":'), + ), + NS( + type="content_block_delta", + index=2, + delta=NS(type="input_json_delta", partial_json='"term"}'), + ), + NS(type="content_block_stop", index=2), + NS(type="content_block_start", index=3, content_block=NS(type="text")), + NS( + type="content_block_delta", + index=3, + delta=NS(type="text_delta", text="answer"), + ), + NS(type="content_block_stop", index=3), + NS( + type="message_delta", + usage=NS(input_tokens=30, cache_read_input_tokens=4, output_tokens=5), + ), + NS( + type="message_delta", + usage=NS(input_tokens=28, output_tokens=8), + delta=NS(stop_reason="end_turn"), + ), + ] + for event in events: + accumulator.consume(event) + + assert accumulator.accumulated_content == "answer" + assert accumulator.stop_reason == "end_turn" + assert accumulator.usage_stats["input_tokens"] == 30 + assert accumulator.usage_stats["cache_read_input_tokens"] == 4 + assert accumulator.usage_stats["output_tokens"] == 8 + assert format_anthropic_streaming_output_complete( + accumulator.content_blocks, accumulator.accumulated_content + ) == [ + { + "role": "assistant", + "content": [ + {"type": "thinking", "thinking": "reason", "signature": "sig"}, + { + "type": "function", + "id": "call-2", + "function": {"name": "lookup", "arguments": {"q": "term"}}, + }, + {"type": "text", "text": "answer"}, + ], + } + ] diff --git a/posthog/test/ai/claude_agent_sdk/test_client.py b/posthog/test/ai/claude_agent_sdk/test_client.py new file mode 100644 index 000000000..acdd38ac3 --- /dev/null +++ b/posthog/test/ai/claude_agent_sdk/test_client.py @@ -0,0 +1,372 @@ +"""Stateful Claude client delivery, turn boundaries, and teardown contracts.""" + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +pytest.importorskip("claude_agent_sdk") +from claude_agent_sdk.types import ( # noqa: E402 + AssistantMessage, + ClaudeAgentOptions, + ResultMessage, + StreamEvent, + TextBlock, + ToolResultBlock, + ToolUseBlock, + UserMessage, +) + +from posthog.ai.claude_agent_sdk import PostHogClaudeSDKClient # noqa: E402 +from posthog.ai.claude_agent_sdk.processor import ( # noqa: E402 + PostHogClaudeAgentProcessor, +) + + +def stream(kind, **data): + return StreamEvent(uuid="evt", session_id="session", event={"type": kind, **data}) + + +def result(**overrides): + return ResultMessage( + **{ + "subtype": "success", + "duration_ms": 2000, + "duration_api_ms": 1500, + "is_error": False, + "num_turns": 1, + "session_id": "session", + "usage": {"input_tokens": 10, "output_tokens": 0}, + "total_cost_usd": 0.01, + **overrides, + } + ) + + +@pytest.fixture +def sdk(monkeypatch): + client = SimpleNamespace( + connect=AsyncMock(), + query=AsyncMock(), + disconnect=AsyncMock(), + interrupt=AsyncMock(), + set_permission_mode=AsyncMock(), + set_model=AsyncMock(), + ) + constructor = MagicMock(return_value=client) + monkeypatch.setattr( + "posthog.ai.claude_agent_sdk.client.ClaudeSDKClient", constructor + ) + client.constructor = constructor + return client + + +def respond(sdk, messages): + async def receive(): + for message in messages: + yield message + + sdk.receive_response = receive + + +def analytics(): + return MagicMock(privacy_mode=False, enable_full_ai_capture=False) + + +def events(ph, name): + return [ + call.kwargs + for call in ph.capture.call_args_list + if call.kwargs["event"] == name + ] + + +@pytest.mark.asyncio +async def test_multiturn_stream_tool_parentage_input_promotion_and_final_trace(sdk): + ph = analytics() + tool = ToolUseBlock(id="tool-1", name="lookup", input={"query": "stock"}) + first = [ + stream( + "message_start", message={"model": "claude", "usage": {"input_tokens": 12}} + ), + AssistantMessage(model="claude", content=[TextBlock(text="checking"), tool]), + UserMessage( + content=[ + ToolResultBlock( + tool_use_id="tool-1", + content=[ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "A" * 400, + }, + } + ], + ), + TextBlock(text="tool done"), + ] + ), + # Empty/unrecognized content must not overwrite already queued tool results. + UserMessage(content=[]), + stream( + "message_delta", + usage={"output_tokens": 4}, + delta={"stop_reason": "tool_use"}, + ), + stream("message_stop"), + stream( + "message_start", message={"model": "claude", "usage": {"input_tokens": 20}} + ), + AssistantMessage(model="claude", content=[TextBlock(text="answer")]), + UserMessage(content="next input"), + stream("message_delta", usage={"output_tokens": 2}), + stream("message_stop"), + result(), + ] + async with PostHogClaudeSDKClient( + posthog_client=ph, + posthog_distinct_id="user", + posthog_trace_id="trace", + posthog_properties={"environment": "test"}, + posthog_groups={"team": "sdk"}, + ) as client: + await client.query("question", session_id="session") + respond(sdk, first) + assert [msg async for msg in client.receive_response()] == first + assert not events(ph, "$ai_trace") + generations = events(ph, "$ai_generation") + assert len(generations) == 2 + one, two = [event["properties"] for event in generations] + assert one["$ai_input"] == [{"role": "user", "content": "question"}] + assert one["$ai_output_choices"][0]["content"][0] == { + "type": "text", + "text": "checking", + } + assert one["$ai_input_tokens"] == 12 + assert one["$ai_output_tokens"] == 4 + assert one["$ai_stop_reason"] == "tool_use" + assert two["$ai_input"] == [ + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "tool-1", + "content": [ + { + "type": "image", + "source": { + "type": "base64", + "media_type": "image/png", + "data": "[base64 image redacted]", + }, + }, + ], + }, + {"type": "text", "text": "tool done"}, + ], + } + ] + span = events(ph, "$ai_span")[0]["properties"] + assert span["$ai_parent_id"] == one["$ai_span_id"] + assert span["$ai_input_state"] == {"query": "stock"} + assert span["$ai_span_name"] == "lookup" + + await client.query("follow up") + second = [ + stream("message_start", message={"model": "claude", "usage": {}}), + AssistantMessage(model="claude", content=[]), + stream("message_stop"), + result(), + ] + respond(sdk, second) + assert [msg async for msg in client.receive_response()] == second + third = events(ph, "$ai_generation")[2]["properties"] + assert third["$ai_span_name"] == "generation_3" + assert third["$ai_input"] == [{"role": "user", "content": "follow up"}] + assert "$ai_output_choices" not in third + assert "$ai_input_tokens" not in third + assert not events(ph, "$ai_trace") + traces = events(ph, "$ai_trace") + assert len(traces) == 1 + assert traces[0]["properties"]["$ai_trace_name"] == "claude_agent_sdk_session" + assert { + event["properties"]["$ai_trace_id"] for event in events(ph, "$ai_generation") + } == {"trace"} + for event in ph.capture.call_args_list: + assert event.kwargs["distinct_id"] == "user" + assert event.kwargs["groups"] == {"team": "sdk"} + assert event.kwargs["properties"]["environment"] == "test" + ph.flush.assert_called_once_with() + sdk.connect.assert_awaited_once_with(None) + assert sdk.query.await_args_list[0].args == ("question", "session") + sdk.disconnect.assert_awaited_once_with() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("private", [False, True]) +async def test_result_fallback_and_privacy_without_person_identity(sdk, private): + ph = analytics() + client = PostHogClaudeSDKClient( + posthog_client=ph, posthog_trace_id="trace", posthog_privacy_mode=private + ) + await client.query("secret") + messages = [ + AssistantMessage(model="claude", content=[TextBlock(text="response")]), + result(), + ] + respond(sdk, messages) + assert [msg async for msg in client.receive_response()] == messages + generation = events(ph, "$ai_generation")[0] + props = generation["properties"] + assert generation["distinct_id"] == "trace" + assert props["$process_person_profile"] is False + assert props["$ai_input"] == ( + None if private else [{"role": "user", "content": "secret"}] + ) + assert props["$ai_output_choices"] == ( + None + if private + else [{"role": "assistant", "content": [{"type": "text", "text": "response"}]}] + ) + assert props["$ai_input_tokens"] == 10 + assert props["$ai_output_tokens"] == 0 + assert props["$ai_latency"] == 1.5 + assert props["$ai_total_cost_usd"] == 0.01 + assert not events(ph, "$ai_trace") + await client.disconnect() + assert events(ph, "$ai_trace")[0]["properties"]["$process_person_profile"] is False + + +@pytest.mark.asyncio +@pytest.mark.parametrize("partial", [None, False, True]) +async def test_options_are_copied_and_delegation_preserves_arguments(sdk, partial): + options = ( + None + if partial is None + else ClaudeAgentOptions( + include_partial_messages=partial, system_prompt="system" + ) + ) + transport = object() + client = PostHogClaudeSDKClient(options, transport, posthog_client=analytics()) + actual, actual_transport = sdk.constructor.call_args.args + assert actual.include_partial_messages is True + assert actual_transport is transport + if options is not None: + assert options.include_partial_messages is partial + assert actual.system_prompt == "system" + assert (actual is options) is partial + await client.connect("initial") + await client.interrupt() + await client.set_permission_mode("acceptEdits") + await client.set_model("claude") + await client.set_model() + sdk.connect.assert_awaited_once_with("initial") + sdk.interrupt.assert_awaited_once_with() + sdk.set_permission_mode.assert_awaited_once_with("acceptEdits") + assert [call.args for call in sdk.set_model.await_args_list] == [ + ("claude",), + (None,), + ] + + +@pytest.mark.asyncio +async def test_instrumentation_failures_never_interrupt_delivery_or_disconnect(sdk): + ph = analytics() + ph.capture.side_effect = RuntimeError("capture failure") + ph.flush.side_effect = RuntimeError("flush failure") + client = PostHogClaudeSDKClient(posthog_client=ph) + await client.query("prompt") + messages = [stream("message_start", message=None), result()] + respond(sdk, messages) + assert [msg async for msg in client.receive_response()] == messages + await client.disconnect() + sdk.disconnect.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_context_manager_propagates_application_exception_and_disconnects(sdk): + ph = analytics() + with pytest.raises(ValueError, match="application failure"): + async with PostHogClaudeSDKClient(posthog_client=ph): + raise ValueError("application failure") + sdk.disconnect.assert_awaited_once_with() + assert len(events(ph, "$ai_trace")) == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("resolved", [None, "person", 123]) +async def test_processor_callable_identity_receives_result_only( + sdk, monkeypatch, resolved +): + ph = analytics() + resolve = MagicMock(return_value=resolved) + processor = PostHogClaudeAgentProcessor(client=ph, distinct_id=resolve) + message = result() + + async def query(**kwargs): + yield message + + monkeypatch.setattr("posthog.ai.claude_agent_sdk.processor.original_query", query) + assert [msg async for msg in processor.query(prompt="question")] == [message] + resolve.assert_called_once_with(message) + generation = events(ph, "$ai_generation")[0] + trace = events(ph, "$ai_trace")[0] + assert generation["distinct_id"] == generation["properties"]["$ai_trace_id"] + assert generation["properties"]["$process_person_profile"] is False + assert trace["distinct_id"] == ( + str(resolved) if resolved else trace["properties"]["$ai_trace_id"] + ) + assert ("$process_person_profile" in trace["properties"]) is (not bool(resolved)) + + +@pytest.mark.parametrize("resolved", [None, "person", 123, ""]) +def test_processor_identity_fallback_defers_callable_until_result(resolved): + """Generation events have no result; only the final trace may resolve a callable ID.""" + resolve = MagicMock(return_value=resolved) + processor = PostHogClaudeAgentProcessor(client=analytics(), distinct_id=resolve) + assert processor._resolve_distinct_id(None) is None + resolve.assert_not_called() + message = result() + assert processor._resolve_distinct_id(None, message) == ( + str(resolved) if resolved else None + ) + resolve.assert_called_once_with(message) + + +@pytest.mark.asyncio +async def test_processor_cancellation_propagates_without_synthetic_generation_or_trace( + monkeypatch, +): + import asyncio + + entered = asyncio.Event() + released = asyncio.Event() + finalized = asyncio.Event() + + async def query(**kwargs): + try: + yield stream("message_start", message={"model": "claude", "usage": {}}) + entered.set() + await released.wait() + finally: + finalized.set() + + monkeypatch.setattr("posthog.ai.claude_agent_sdk.processor.original_query", query) + ph = analytics() + processor = PostHogClaudeAgentProcessor(client=ph) + + async def consume(): + return [message async for message in processor.query(prompt="question")] + + task = asyncio.create_task(consume()) + await entered.wait() + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert finalized.is_set() + ph.capture.assert_not_called() + ph.flush.assert_not_called() diff --git a/posthog/test/ai/gemini/test_gemini_converter.py b/posthog/test/ai/gemini/test_gemini_converter.py index 4227b6d08..7614f7037 100644 --- a/posthog/test/ai/gemini/test_gemini_converter.py +++ b/posthog/test/ai/gemini/test_gemini_converter.py @@ -361,3 +361,64 @@ def test_input_tokens_without_tool_use(self): ) usage = extract_gemini_usage_from_response(resp) assert usage["input_tokens"] == 20 + + +@pytest.mark.parametrize( + "message, expected", + [ + ({"role": "model", "content": 42}, {"role": "model", "content": "42"}), + ({"content": None}, {"role": "user", "content": "None"}), + ({"content": ""}, {"role": "user", "content": ""}), + ( + {"role": "model", "content": [{"text": "answer"}]}, + {"role": "model", "content": [{"type": "text", "text": "answer"}]}, + ), + ( + {"role": "model", "parts": [], "content": "ignored"}, + {"role": "model", "content": []}, + ), + ( + {"text": "bare part"}, + {"role": "user", "content": [{"type": "text", "text": "bare part"}]}, + ), + ({"role": "model", "text": "answer"}, {"role": "model", "content": "answer"}), + ({"unknown": 1}, {"role": "user", "content": "{'unknown': 1}"}), + ( + {"parts": "not a list"}, + {"role": "user", "content": "{'parts': 'not a list'}"}, + ), + ], +) +def test_input_dictionary_field_precedence_and_content_coercion(message, expected): + assert format_gemini_input(message) == [expected] + + +@pytest.mark.parametrize( + "attributes, expected", + [ + ( + {"role": 7, "parts": ("first", {"text": "second"})}, + { + "role": "user", + "content": [ + {"type": "text", "text": "first"}, + {"type": "text", "text": "second"}, + ], + }, + ), + ( + {"role": None, "text": 0, "content": "ignored"}, + {"role": "user", "content": 0}, + ), + ( + {"role": "model", "content": [{"text": "answer"}]}, + {"role": "model", "content": [{"type": "text", "text": "answer"}]}, + ), + ({"role": "model", "content": 42}, {"role": "model", "content": "42"}), + ({"role": "model", "content": ""}, {"role": "model", "content": ""}), + ], +) +def test_input_object_field_precedence_and_role_validation(attributes, expected): + from types import SimpleNamespace + + assert format_gemini_input(SimpleNamespace(**attributes)) == [expected] diff --git a/posthog/test/ai/openai/test_openai_converter.py b/posthog/test/ai/openai/test_openai_converter.py index 2034f533f..155a03d73 100644 --- a/posthog/test/ai/openai/test_openai_converter.py +++ b/posthog/test/ai/openai/test_openai_converter.py @@ -420,3 +420,129 @@ def test_indexless_calls_keep_their_order_around_indexed_ones(): ("call_b", "b", "{}"), ("call_c", "c", "{}"), ] + + +@pytest.mark.parametrize( + "text, tool_calls, expected", + [ + ("", None, []), + ("hello", None, [{"type": "text", "text": "hello"}]), + ("", [{}], [{"type": "function", "id": None, "function": {}}]), + ( + "hello", + [{"id": "call-1", "function": {"name": "f", "arguments": "{}"}}], + [ + {"type": "text", "text": "hello"}, + { + "type": "function", + "id": "call-1", + "function": {"name": "f", "arguments": "{}"}, + }, + ], + ), + ], +) +def test_legacy_streaming_content_preserves_text_then_tools(text, tool_calls, expected): + from copy import deepcopy + + from posthog.ai.openai.openai_converter import format_openai_streaming_content + + original = deepcopy(tool_calls) + assert format_openai_streaming_content(text, tool_calls) == expected + assert tool_calls == original + + +def test_responses_content_preserves_empty_values_and_field_precedence(): + from posthog.ai.openai.openai_converter import format_openai_response + + response = types.SimpleNamespace( + output=[ + { + "type": "message", + "role": "developer", + "content": [ + {"type": "output_text", "text": ""}, + {"type": "refusal", "refusal": "", "text": "ignored"}, + types.SimpleNamespace(type="unknown", text="fallback"), + {"type": "input_image", "image_url": "https://image"}, + { + "type": "input_image", + "image_url": "ignored", + "text": "text wins", + }, + {"type": "unknown"}, + ], + }, + {"type": "message", "content": 0}, + {"type": "message", "content": None}, + ] + ) + assert format_openai_response(response) == [ + { + "role": "developer", + "content": [ + {"type": "text", "text": ""}, + {"type": "refusal", "refusal": ""}, + {"type": "text", "text": "fallback"}, + {"type": "image", "image": "https://image"}, + {"type": "text", "text": "text wins"}, + {"type": "text", "text": "0"}, + ], + } + ] + + +@pytest.mark.parametrize( + "item, expected", + [ + (42, {"role": "user", "content": "42"}), + (None, {"role": "user", "content": "None"}), + ("prompt", {"role": "user", "content": "prompt"}), + ( + {"type": "function_call_output", "output": "ok"}, + {"type": "function_call_output", "output": "ok"}, + ), + ({"type": "message", "role": "system"}, {"role": "system", "content": ""}), + ({"type": "message", "content": None}, {"role": "user", "content": None}), + ( + types.SimpleNamespace(role="assistant", content="typed"), + {"role": "user", "content": "namespace(role='assistant', content='typed')"}, + ), + ], +) +def test_responses_input_distinguishes_messages_from_typed_items(item, expected): + assert format_openai_input(input_data=[item]) == [expected] + + +def test_usage_chat_fields_override_responses_even_when_null_or_zero(): + from posthog.ai.openai.openai_converter import ( + extract_openai_usage_from_chunk, + extract_openai_usage_from_response, + ) + + usage = types.SimpleNamespace( + input_tokens=100, + output_tokens=50, + input_tokens_details=types.SimpleNamespace(cached_tokens=9), + output_tokens_details=types.SimpleNamespace(reasoning_tokens=7), + prompt_tokens=None, + completion_tokens=0, + prompt_tokens_details=types.SimpleNamespace(cached_tokens=None), + completion_tokens_details=types.SimpleNamespace(reasoning_tokens=0), + ) + response = types.SimpleNamespace(usage=usage) + non_streamed = extract_openai_usage_from_response(response) + assert set(non_streamed) == {"output_tokens", "raw_usage"} + assert non_streamed["output_tokens"] == 0 + assert non_streamed["raw_usage"]["prompt_tokens"] is None + + streamed = extract_openai_usage_from_chunk(response, "chat") + assert set(streamed) == { + "input_tokens", + "output_tokens", + "reasoning_tokens", + "raw_usage", + } + assert streamed["input_tokens"] is None + assert streamed["output_tokens"] == streamed["reasoning_tokens"] == 0 + assert streamed["raw_usage"] == non_streamed["raw_usage"] diff --git a/posthog/test/ai/openai_agents/test_processor.py b/posthog/test/ai/openai_agents/test_processor.py index 36e9a5803..13f0c6c7a 100644 --- a/posthog/test/ai/openai_agents/test_processor.py +++ b/posthog/test/ai/openai_agents/test_processor.py @@ -1222,3 +1222,70 @@ def test_ai_lane_client_routes_through_capture_ai(mock_client, mock_trace): mock_client.capture.assert_not_called() mock_client.capture_ai.assert_called_once() assert mock_client.capture_ai.call_args[1]["event"] == "$ai_trace" + + +@pytest.mark.parametrize("full_capture", [False, True]) +@pytest.mark.parametrize("privacy", [False, True]) +def test_transcription_privacy_wins_over_full_capture_without_mutating_config( + mock_client, mock_span, full_capture, privacy +): + mock_client.enable_full_ai_capture = full_capture + processor = PostHogTracingProcessor( + client=mock_client, distinct_id="audio-user", privacy_mode=privacy + ) + config = { + "prompt": "private prompt", + "instructions": "private instructions", + "language": "en", + } + data = TranscriptionSpanData( + input="AAAA", + input_format="pcm", + output="private transcript", + model_config=config, + ) + mock_span.span_data = data + processor.on_span_end(mock_span) + + capture = mock_client.capture_ai if full_capture else mock_client.capture + capture.assert_called_once() + props = capture.call_args.kwargs["properties"] + assert capture.call_args.kwargs["distinct_id"] == "audio-user" + assert props["$ai_input"] == ( + None if privacy else "AAAA" if full_capture else "[base64 audio redacted]" + ) + assert props["$ai_output_state"] == (None if privacy else "private transcript") + assert props["model_config"] == { + "prompt": None if privacy else "private prompt", + "instructions": None if privacy else "private instructions", + "language": "en", + } + assert ( + data.model_config + == config + == { + "prompt": "private prompt", + "instructions": "private instructions", + "language": "en", + } + ) + + +def test_span_error_category_keeps_precedence_and_privacy_summary( + mock_client, mock_span +): + processor = PostHogTracingProcessor(client=mock_client, privacy_mode=True) + mock_client.enable_full_ai_capture = False + mock_span.span_data = GenerationSpanData(model="gpt-4o", usage={"input_tokens": 0}) + mock_span.error = { + "type": "UserError", + "message": "private ModelBehaviorError and MaxTurnsExceeded", + } + processor.on_span_end(mock_span) + + props = mock_client.capture.call_args.kwargs["properties"] + assert props["$ai_error_type"] == "model_behavior_error" + assert props["$ai_error"] == "UserError" + assert props["$ai_input_tokens"] == props["$ai_total_tokens"] == 0 + assert "$ai_output_tokens" not in props + assert props["$process_person_profile"] is False diff --git a/posthog/test/ai/test_sanitization.py b/posthog/test/ai/test_sanitization.py index 34be067b4..9f40ed885 100644 --- a/posthog/test/ai/test_sanitization.py +++ b/posthog/test/ai/test_sanitization.py @@ -672,3 +672,32 @@ def test_sanitize_messages_importable_from_utils(): if __name__ == "__main__": unittest.main() + + +class TestMediaPlaceholderMimeHints(unittest.TestCase): + def test_media_kind_in_placeholders_preserves_mime_and_format_classification(self): + cases = { + "audio/wav": "audio", + "video/mp4": "video", + "application/pdf": "file", + "AUDIO": "audio", + "video": "video", + "image": "image", + "mp3": "audio", + "webm": "video", + "png": "image", + "pdf": "file", + } + for mime, kind in cases.items(): + with self.subTest(mime=mime): + value = {"inline_data": {"mime_type": mime, "data": LONG_RAW_BASE64}} + self.assertEqual( + redact_media(value), + { + "inline_data": { + "mime_type": mime, + "data": f"[base64 {kind} redacted]", + } + }, + ) + self.assertEqual(value["inline_data"]["data"], LONG_RAW_BASE64) diff --git a/posthog/test/ai/test_utils.py b/posthog/test/ai/test_utils.py new file mode 100644 index 000000000..49add2881 --- /dev/null +++ b/posthog/test/ai/test_utils.py @@ -0,0 +1,166 @@ +"""Usage conversion fallback order and streaming merge ownership contracts.""" + +from copy import deepcopy +from types import SimpleNamespace + +import pytest + +from posthog.ai.types import TokenUsage +from posthog.ai.utils import merge_usage_stats, serialize_raw_usage + + +class PydanticUsage: + def model_dump(self): + return {"input_tokens": 4} + + def to_dict(self): + raise AssertionError("model_dump must take precedence") + + +class ProtobufUsage: + def model_dump(self): + raise ValueError("unsupported model serialization") + + def to_dict(self): + return {"output_tokens": 7} + + +class AttributeUsage: + def __init__(self): + self.input_tokens = 9 + + def model_dump(self): + raise ValueError("failed") + + def to_dict(self): + raise ValueError("failed") + + +class StringUsage: + __slots__ = () + model_dump = None + to_dict = None + + def __str__(self): + return "provider usage" + + +class UnreadableUsage(StringUsage): + __slots__ = () + + def __str__(self): + raise ValueError("cannot stringify") + + +@pytest.mark.parametrize( + ("raw", "expected"), + [ + (None, None), + ({"input_tokens": 1}, {"input_tokens": 1}), + (PydanticUsage(), {"input_tokens": 4}), + (ProtobufUsage(), {"output_tokens": 7}), + (AttributeUsage(), {"input_tokens": 9}), + (SimpleNamespace(output_tokens=3), {"output_tokens": 3}), + (StringUsage(), {"_raw": "provider usage"}), + (UnreadableUsage(), None), + ], +) +def test_serialize_raw_usage_follows_fallback_order_without_raising(raw, expected): + assert serialize_raw_usage(raw) == expected + + +def test_serialize_raw_usage_preserves_already_serialized_mapping(): + raw = {"input_tokens": 0, "nested": {"details": 1}} + assert serialize_raw_usage(raw) is raw + + +def test_merge_usage_stats_incremental_keeps_max_search_and_merges_raw_shallowly(): + target = TokenUsage( + input_tokens=None, + output_tokens=3, + cache_read_input_tokens=2, + cache_creation_input_tokens=4, + reasoning_tokens=1, + web_search_count=5, + cache_reporting_exclusive=False, + raw_usage={"input_tokens": 12, "details": {"old": 1}}, + ) + source = TokenUsage( + input_tokens=10, + output_tokens=0, + cache_read_input_tokens=3, + cache_creation_input_tokens=6, + reasoning_tokens=2, + web_search_count=2, + cache_reporting_exclusive=True, + raw_usage={"output_tokens": 7, "details": {"new": 2}}, + ) + original_source = deepcopy(source) + previous_raw = target["raw_usage"] + merge_usage_stats(target, source) + assert target == { + "input_tokens": 10, + "output_tokens": 3, + "cache_read_input_tokens": 5, + "cache_creation_input_tokens": 10, + "reasoning_tokens": 3, + "web_search_count": 5, + "cache_reporting_exclusive": True, + "raw_usage": {"input_tokens": 12, "output_tokens": 7, "details": {"new": 2}}, + } + assert previous_raw == {"input_tokens": 12, "details": {"old": 1}} + assert source == original_source + assert target["raw_usage"] is not source["raw_usage"] + merge_usage_stats( + target, TokenUsage(input_tokens=None, output_tokens=None, raw_usage=None) + ) + assert target["input_tokens"] == 10 + assert target["output_tokens"] == 3 + assert target["raw_usage"] == { + "input_tokens": 12, + "output_tokens": 7, + "details": {"new": 2}, + } + + +def test_merge_usage_stats_cumulative_replaces_zero_and_false_but_not_none(): + target = TokenUsage( + input_tokens=10, + output_tokens=8, + cache_read_input_tokens=3, + cache_creation_input_tokens=6, + reasoning_tokens=4, + web_search_count=7, + cache_reporting_exclusive=True, + raw_usage={"previous": 1}, + ) + source = TokenUsage( + input_tokens=None, + output_tokens=0, + cache_read_input_tokens=0, + cache_creation_input_tokens=0, + reasoning_tokens=0, + web_search_count=0, + cache_reporting_exclusive=False, + raw_usage={"latest": 2}, + ) + merge_usage_stats(target, source, mode="cumulative") + assert target == { + "input_tokens": 10, + "output_tokens": 0, + "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 0, + "reasoning_tokens": 0, + "web_search_count": 0, + "cache_reporting_exclusive": False, + "raw_usage": {"latest": 2}, + } + assert source["input_tokens"] is None + assert target["raw_usage"] is source["raw_usage"] + + +def test_merge_usage_stats_invalid_mode_does_not_mutate_target(): + target = TokenUsage(input_tokens=10) + with pytest.raises(ValueError, match="Invalid mode: unknown"): + merge_usage_stats(target, TokenUsage(input_tokens=4), mode="unknown") + assert target == {"input_tokens": 10} diff --git a/posthog/test/integrations/test_asgi_integration.py b/posthog/test/integrations/test_asgi_integration.py index 922f6fcde..480b098b6 100644 --- a/posthog/test/integrations/test_asgi_integration.py +++ b/posthog/test/integrations/test_asgi_integration.py @@ -439,15 +439,32 @@ async def app(received_scope, receive, send): @pytest.mark.asyncio -async def test_builds_url_from_server_when_host_header_is_absent(): +@pytest.mark.parametrize( + "scheme,server,expected_url", + [ + ("http", ("2001:db8::1", 8080), "http://[2001:db8::1]:8080/health"), + ("https", ("2001:db8::1", 443), "https://[2001:db8::1]/health"), + ("http", ("[2001:db8::1]", 80), "http://[2001:db8::1]/health"), + ("http", ("example.com", 80), "http://example.com/health"), + ("https", ("example.com", 443), "https://example.com/health"), + ("wss", ("example.com", 443), "wss://example.com:443/health"), + ("https", ("example.com", "443"), None), + ("https", (None, 443), None), + ("https", (" \r\n ", 443), None), + ("https", ("example.com",), None), + ("https", None, None), + ], +) +async def test_builds_url_from_server_when_host_header_is_absent( + scheme, server, expected_url +): observed = {} async def app(scope, receive, send): observed.update(contexts.get_tags()) - scope = http_scope( - headers=[], scheme="http", server=("2001:db8::1", 8080), path="/health" - ) + scope = http_scope(headers=[], scheme=scheme, server=server, path="/health") await PosthogASGIMiddleware(app)(scope, noop_receive, noop_send) - assert observed["$current_url"] == "http://[2001:db8::1]:8080/health" + assert observed.get("$current_url") == expected_url + assert observed["$request_path"] == "/health" diff --git a/posthog/test/mcp/test_asgi.py b/posthog/test/mcp/test_asgi.py new file mode 100644 index 000000000..395d826e8 --- /dev/null +++ b/posthog/test/mcp/test_asgi.py @@ -0,0 +1,28 @@ +"""Cancellation at the request-body boundary must not start the downstream app.""" + +import asyncio + +import pytest + +from posthog.mcp.asgi import PostHogMcpStatelessSessionMiddleware, get_mcp_session + + +async def test_cancellation_while_sniffing_initialize_propagates(): + calls = [] + scope = {"type": "http", "method": "POST", "headers": []} + + async def receive(): + raise asyncio.CancelledError() + + async def app(*args): + calls.append("app") + + async def send(message): + calls.append(message) + + middleware = PostHogMcpStatelessSessionMiddleware(app) + with pytest.raises(asyncio.CancelledError): + await middleware(scope, receive, send) + + assert calls == [] + assert get_mcp_session(scope) is None diff --git a/posthog/test/mcp/test_capture.py b/posthog/test/mcp/test_capture.py new file mode 100644 index 000000000..1190f9941 --- /dev/null +++ b/posthog/test/mcp/test_capture.py @@ -0,0 +1,69 @@ +"""Custom event handles must use the same privacy and identity pipeline as auto capture.""" + +import pytest +from mcp.server.lowlevel import Server + +from posthog.mcp import McpAnalytics, instrument +from posthog.mcp._internal import get_server_tracking_data +from posthog.mcp.types import UserIdentity +from posthog.test.mcp._helpers import FakeClient + + +async def test_custom_capture_is_awaited_and_preserves_session_identity(): + server = Server("custom-capture") + client = FakeClient() + handle = instrument(server, client) + data = get_server_tracking_data(server) + data.identified_sessions.set( + data.session_id, + UserIdentity( + distinct_id="user-7", properties={"plan": "pro"}, groups={"org": "acme"} + ), + ) + + await handle.capture("report_exported", {"format": "csv"}) + + assert len(client.events) == 1 + event = client.events[0] + assert event["event"] == "report_exported" + assert event["distinct_id"] == "user-7" + assert event["properties"]["$session_id"] == data.session_id + assert event["properties"]["$set"] == {"plan": "pro"} + assert event["properties"]["$groups"] == {"org": "acme"} + assert event["properties"]["format"] == "csv" + assert "$process_person_profile" not in event["properties"] + + +@pytest.mark.parametrize("name", [None, "", 3]) +async def test_custom_capture_rejects_invalid_event_names(name): + server = Server("custom-capture") + client = FakeClient() + handle = instrument(server, client) + + with pytest.raises(ValueError, match="requires an event name"): + await handle.capture(name) + + assert client.events == [] + + +async def test_custom_capture_after_tracking_removed_is_a_noop(): + server = Server("custom-capture") + client = FakeClient() + instrument(server, client) + # A handle without tracking state is inert, not a fallback to another server. + handle = McpAnalytics(Server("untracked")) + + await handle.capture("report_exported") + + assert client.events == [] + + +async def test_custom_capture_with_no_sink_is_a_noop(): + server = Server("custom-capture") + client = FakeClient() + handle = instrument(server, client) + get_server_tracking_data(server).sink = None + + await handle.capture("report_exported") + + assert client.events == [] diff --git a/posthog/test/mcp/test_instrument_lowlevel_ownership.py b/posthog/test/mcp/test_instrument_lowlevel_ownership.py new file mode 100644 index 000000000..e10eb00b1 --- /dev/null +++ b/posthog/test/mcp/test_instrument_lowlevel_ownership.py @@ -0,0 +1,70 @@ +"""Ownership inference stays safe when a host exposes only a function registry.""" + +from types import SimpleNamespace + +import pytest + +from posthog.mcp._instrument_lowlevel import ( + _standalone_ownership, + registered_tool_exists, +) +from posthog.mcp._internal import MCPAnalyticsData +from posthog.mcp.types import MCPAnalyticsOptions + + +def _application_arguments(context, conversation_id, llm_model): + pass + + +class _UninspectableFunction: + @property + def __signature__(self): + raise ValueError("signature unavailable") + + def __call__(self): + pass + + +@pytest.mark.parametrize( + ("tool", "expected_strip", "expected_model"), + [ + (None, {"context", "conversation_id"}, None), + (SimpleNamespace(), {"context", "conversation_id", "llm_model"}, True), + (SimpleNamespace(fn=_application_arguments), set(), False), + ( + SimpleNamespace(fn=_UninspectableFunction()), + {"context", "conversation_id", "llm_model"}, + True, + ), + ], +) +async def test_signature_only_registry_preserves_application_arguments( + tool, expected_strip, expected_model +): + async def get_tool(name): + assert name == "search" + return tool + + server = SimpleNamespace(get_tool=get_tool, middleware=[]) + data = MCPAnalyticsData(options=MCPAnalyticsOptions()) + + strip, model_owned, schema = await _standalone_ownership( + data, server, "search", None + ) + + assert strip == expected_strip + assert model_owned is expected_model + assert schema is None + + +async def test_registry_outage_is_not_evidence_of_an_unknown_tool(): + fastmcp = pytest.importorskip("fastmcp") + server = fastmcp.FastMCP("registry-outage") + + async def unavailable(name): + raise RuntimeError("registry temporarily unavailable") + + server.get_tool = unavailable + + assert await registered_tool_exists(server, "search") is None + assert await registered_tool_exists(None, "search") is None diff --git a/posthog/test/mcp/test_instrument_v2.py b/posthog/test/mcp/test_instrument_v2.py new file mode 100644 index 000000000..cb9c97418 --- /dev/null +++ b/posthog/test/mcp/test_instrument_v2.py @@ -0,0 +1,30 @@ +"""The v2 ownership lookup must not consume host request cancellation.""" + +import asyncio +from types import SimpleNamespace + +import pytest + +from posthog.mcp._instrument_v2 import _standalone_injected_parameters +from posthog.mcp._internal import MCPAnalyticsData +from posthog.mcp.types import MCPAnalyticsOptions + + +async def test_cancellation_during_standalone_schema_lookup_propagates(): + pytest.importorskip("fastmcp") + calls = [] + + async def list_tools(): + calls.append("list") + raise asyncio.CancelledError() + + async def get_tool(*args, **kwargs): + calls.append("get") + + server = SimpleNamespace(list_tools=list_tools, get_tool=get_tool) + data = MCPAnalyticsData(options=MCPAnalyticsOptions()) + with pytest.raises(asyncio.CancelledError): + await _standalone_injected_parameters(server, data, "search", None) + + assert calls == ["list"] + assert data.tool_model_parameter_injected == {} diff --git a/posthog/test/mcp/test_model_parameters.py b/posthog/test/mcp/test_model_parameters.py index 1c2620f8c..d13e1b151 100644 --- a/posthog/test/mcp/test_model_parameters.py +++ b/posthog/test/mcp/test_model_parameters.py @@ -102,3 +102,27 @@ def test_request_meta_from_context_supports_pydantic_style_metadata(): assert request_meta_from_context(SimpleNamespace(meta=meta)) == { "x-codex-turn-metadata": {"model": "gpt-5.6-sol"} } + + +@pytest.mark.parametrize("method", ["model_dump", "dict"]) +@pytest.mark.parametrize("outcome", ["metadata", "non_object", "raises"]) +def test_request_metadata_serialization_is_best_effort(method, outcome): + metadata = {"x-codex-turn-metadata": {"model": "reported-model"}} + calls = [] + + def serialize(**kwargs): + calls.append(kwargs) + if outcome == "raises": + raise RuntimeError("metadata unavailable") + return metadata if outcome == "metadata" else ["not metadata"] + + context = SimpleNamespace(meta=SimpleNamespace(**{method: serialize})) + + assert request_meta_from_context(context) == ( + metadata if outcome == "metadata" else None + ) + assert calls == [ + {"mode": "python", "by_alias": True} + if method == "model_dump" + else {"by_alias": True} + ] diff --git a/posthog/test/mcp/test_posthog_mcp.py b/posthog/test/mcp/test_posthog_mcp.py index fb2e69961..76c724bb6 100644 --- a/posthog/test/mcp/test_posthog_mcp.py +++ b/posthog/test/mcp/test_posthog_mcp.py @@ -797,3 +797,69 @@ def test_capture_without_a_session_gets_a_fresh_one_per_call(): assert session.startswith("ses_") assert first["distinct_id"] == session assert second["properties"]["$session_id"] != session + + +class _UncopyableToolResult: + def __init__(self): + self.content = [] + self.structuredContent = {"rows": []} + + def __copy__(self): + return self + + +class _FailingModelCopyResult(_UncopyableToolResult): + def model_copy(self, **kwargs): + raise ValueError("result is read-only") + + +class _FailingResultWrapper: + def __init__(self, root): + self.root = root + + def model_copy(self, **kwargs): + raise ValueError("wrapper is read-only") + + +@pytest.mark.parametrize( + "make_result", + [ + _UncopyableToolResult, + _FailingModelCopyResult, + lambda: _FailingResultWrapper( + SimpleNamespace(content=[], structuredContent={"rows": []}) + ), + lambda: SimpleNamespace( + root=SimpleNamespace(content=[], structuredContent={"rows": []}) + ), + ], + ids=["uncopyable", "model-copy-fails", "rewrap-fails", "plain-wrapper"], +) +def test_prepare_tool_result_copy_failures_never_pin_a_shared_handle(make_result): + client, _ = make_client() + client.prepare_tool_list([_sql_tool()]) + call = client.prepare_tool_call("execute-sql", {}) + result = make_result() + target = getattr(result, "root", result) + + prepared = client.prepare_tool_result(result, call) + + assert target.content == [] + assert target.structuredContent == {"rows": []} + if isinstance(result, SimpleNamespace): + assert prepared.result is not result + assert prepared.result.root is not target + assert prepared.result.root.structuredContent[MCP_INSTRUCTIONS_KEY] == { + "conversation_id": prepared.conversation_id + } + assert ( + prepared.result.root.content[-1].text + == _handle_block(prepared.conversation_id)["text"] + ) + assert prepared.session_id == derive_session_id_from_conversation( + prepared.conversation_id + ) + else: + assert prepared.result is result + assert prepared.conversation_id is None + assert prepared.session_id is None diff --git a/posthog/test/test_async_client.py b/posthog/test/test_async_client.py index fcb6e2298..97ea959fe 100644 --- a/posthog/test/test_async_client.py +++ b/posthog/test/test_async_client.py @@ -3,6 +3,9 @@ import asyncio import contextvars import logging +from datetime import datetime, timezone +from pathlib import Path +from uuid import UUID import subprocess import sys import threading @@ -13,7 +16,9 @@ from posthog import AsyncClient, AsyncPosthog, CaptureCompression, CaptureMode from posthog.consumer import MAX_MSG_SIZE from posthog.contexts import ( + identify_context, new_context, + set_context_session, set_capture_exception_code_variables_context, set_code_variables_detect_secrets_context, set_code_variables_ignore_patterns_context, @@ -771,3 +776,245 @@ async def batch_post(*args, **kwargs): assert "items pending" in caplog.text allow_upload.set() await client.shutdown() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method_name,method_kwargs,expected_id,expected_properties", + [ + ( + "group_identify", + { + "group_type": "company", + "group_key": "company-1", + "properties": {"plan": "pro", "nested": {"drop": None, "count": 0}}, + "distinct_id": "group-user", + }, + "group-user", + { + "$group_type": "company", + "$group_key": "company-1", + "$group_set": {"plan": "pro", "nested": {"drop": None, "count": 0}}, + }, + ), + ( + "alias", + {"previous_id": 123, "distinct_id": "alias-user"}, + "123", + {"distinct_id": "123", "alias": "alias-user"}, + ), + ], +) +async def test_identity_events_preserve_payload_session_and_normalized_envelope( + method_name, method_kwargs, expected_id, expected_properties +): + delivered = [] + + async def batch_post(*args, **kwargs): + delivered.extend(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + async with AsyncPosthog("test-key", flush_interval=30) as client: + with new_context(fresh=True): + set_context_session("request-session") + event_uuid = getattr(client, method_name)( + **method_kwargs, + uuid="{4BF92F35-77B3-4DA6-A3CE-929D0E0E4736}", + timestamp="2025-01-02T03:04:05.123456+02:30", + disable_geoip=False, + ) + await client.flush(timeout_seconds=1) + + assert len(delivered) == 1 + event = delivered[0] + assert event_uuid == event["uuid"] == "4bf92f35-77b3-4da6-a3ce-929d0e0e4736" + assert event["timestamp"] == "2025-01-02T00:34:05.123456+00:00" + assert event["distinct_id"] == expected_id + assert { + key: event["properties"][key] for key in expected_properties + } == expected_properties + assert event["properties"]["$session_id"] == "request-session" + assert "$geoip_disable" not in event["properties"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "method_name,method_kwargs", + [ + ("group_identify", {"group_type": "", "group_key": "company-1"}), + ("group_identify", {"group_type": "company", "group_key": ""}), + ("alias", {"previous_id": "", "distinct_id": "user-1"}), + ("alias", {"previous_id": "anon-1", "distinct_id": None}), + ], +) +async def test_identity_events_reject_missing_required_ids(method_name, method_kwargs): + with ( + new_context(fresh=True), + mock.patch( + "posthog._async_consumer.async_batch_post", new=mock.AsyncMock() + ) as send, + ): + async with AsyncPosthog("test-key") as client: + assert getattr(client, method_name)(**method_kwargs) is None + await client.flush() + send.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_alias_uses_request_identity_when_no_explicit_distinct_id_is_supplied(): + delivered = [] + + async def batch_post(*args, **kwargs): + delivered.extend(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + async with AsyncPosthog("test-key") as client: + with new_context(fresh=True): + identify_context("request-user") + assert client.alias("anon-1", None) is not None + await client.flush() + + assert delivered[0]["distinct_id"] == "anon-1" + assert delivered[0]["properties"]["alias"] == "request-user" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("cancel_immediate", [False, True]) +async def test_concurrent_shutdown_closes_once_after_immediate_operation_finishes( + cancel_immediate, +): + started, release = asyncio.Event(), asyncio.Event() + http_client = mock.Mock(aclose=mock.AsyncMock()) + on_error = mock.AsyncMock() + + async def batch_post(*args, **kwargs): + started.set() + await release.wait() + + with ( + mock.patch("posthog.async_client._build_client", return_value=http_client), + mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post), + ): + client = AsyncPosthog("test-key", on_error=on_error) + capture = asyncio.create_task(client.capture_immediate("in flight")) + await started.wait() + shutdowns = [asyncio.create_task(client.shutdown()) for _ in range(2)] + await asyncio.sleep(0) + assert client.capture("too late") is None + assert not any(task.done() for task in shutdowns) + http_client.aclose.assert_not_awaited() + if cancel_immediate: + capture.cancel() + with pytest.raises(asyncio.CancelledError): + await capture + else: + release.set() + assert await capture is not None + await asyncio.wait_for(asyncio.gather(*shutdowns), timeout=1) + + http_client.aclose.assert_awaited_once_with() + on_error.assert_not_awaited() + assert client._immediate_callers == {} + assert client._inflight_operations == set() + assert await client.capture_immediate("after shutdown") is None + + +@pytest.mark.asyncio +async def test_cancelled_shutdown_can_be_retried_without_cancelling_buffered_delivery(): + started, release = asyncio.Event(), asyncio.Event() + delivered = [] + + async def batch_post(*args, **kwargs): + started.set() + await release.wait() + delivered.extend(kwargs["batch"]) + + with mock.patch("posthog._async_consumer.async_batch_post", side_effect=batch_post): + client = AsyncPosthog("test-key", flush_at=1) + event_uuid = client.capture("buffered") + await started.wait() + shutdown = asyncio.create_task(client.shutdown()) + await asyncio.sleep(0) + shutdown.cancel() + with pytest.raises(asyncio.CancelledError): + await shutdown + assert client.capture("after cancelled shutdown") is None + release.set() + await asyncio.wait_for(client.shutdown(), timeout=1) + + assert [event["uuid"] for event in delivered] == [event_uuid] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("debug", [False, True]) +async def test_shutdown_attempts_remaining_cleanup_before_reporting_first_error(debug): + first_error = RuntimeError("flush failed") + http_client = mock.Mock( + aclose=mock.AsyncMock(side_effect=ValueError("close failed")) + ) + with ( + mock.patch("posthog.async_client._build_client", return_value=http_client), + mock.patch( + "posthog.async_client._async_flags", new=mock.AsyncMock(return_value={}) + ), + ): + client = AsyncPosthog("test-key", debug=debug, secret_key="secret-key") + await client.evaluate_flags("user-1") + emails = client.emails + with ( + mock.patch.object( + client, "flush", new=mock.AsyncMock(side_effect=first_error) + ), + mock.patch.object( + emails, + "aclose", + new=mock.AsyncMock(side_effect=OSError("email close failed")), + ) as close_email, + ): + if debug: + with pytest.raises(RuntimeError) as raised: + await client.shutdown() + assert raised.value is first_error + else: + await client.shutdown() + await client.shutdown() + + http_client.aclose.assert_awaited_once_with() + close_email.assert_awaited_once_with() + assert client._duplicate_client_registry_key is None + assert client._closed + assert client.capture("after shutdown") is None + + +@pytest.mark.asyncio +async def test_immediate_capture_replaces_invalid_uuid_and_timestamp_before_hook(): + observed = [] + + async def before_send(event): + observed.append(event.copy()) + return {**event, "uuid": "hook-cannot-replace-event-identity"} + + with mock.patch( + "posthog._async_consumer.async_batch_post", new=mock.AsyncMock() + ) as send: + client = AsyncPosthog("test-key", before_send=before_send) + before = datetime.now(timezone.utc) + event_uuid = await client.capture_immediate( + "event", uuid="invalid", timestamp="not-a-date" + ) + after = datetime.now(timezone.utc) + await client.shutdown() + + assert str(UUID(event_uuid)) == event_uuid + assert observed[0]["uuid"] == event_uuid + assert before <= datetime.fromisoformat(observed[0]["timestamp"]) <= after + assert send.await_args.kwargs["batch"][0]["uuid"] == event_uuid + + +@pytest.mark.asyncio +async def test_deprecated_secret_key_warning_points_to_constructor_caller(): + with pytest.warns(DeprecationWarning, match="personal_api_key") as warnings_seen: + client = AsyncPosthog("test-key", personal_api_key=" legacy-key ", send=False) + assert Path(warnings_seen[0].filename) == Path(__file__) + assert client.secret_key == client.personal_api_key == "legacy-key" + await client.shutdown() diff --git a/posthog/test/test_async_consumer.py b/posthog/test/test_async_consumer.py index 9bb3f682f..7ccae1efd 100644 --- a/posthog/test/test_async_consumer.py +++ b/posthog/test/test_async_consumer.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextvars import json from unittest import mock @@ -8,7 +9,7 @@ import pytest from freezegun import freeze_time -from posthog._async_consumer import _AsyncConsumer +from posthog._async_consumer import _AsyncConsumer, _QueuedEvent, _STOP from posthog.capture_compression import CaptureCompression from posthog.capture_mode import CaptureMode from posthog.request import APIError @@ -207,3 +208,113 @@ async def test_get_or_flush_preserves_results_and_cleans_up_waiters(queued, flus assert not (asyncio.all_tasks() - tasks_before) if queued: consumer.queue.task_done() + + +@pytest.mark.asyncio +async def test_run_drops_invalid_events_and_preserves_queued_context(monkeypatch): + marker = contextvars.ContextVar("queued_test_marker", default="outside") + consumer = make_consumer(retries=0) + monkeypatch.setattr("posthog._async_consumer.MAX_MSG_SIZE", 256) + + async def process(event): + if event["event"] == "broken": + raise ValueError("processing failed") + if event["event"] == "filtered": + return None + return {**event, "context": marker.get()} + + consumer.process_event = process + token = marker.set("queued") + try: + for event in [ + {"event": "broken"}, + {"event": "filtered"}, + {"event": "invalid-json", "properties": {("tuple",): True}}, + {"event": "too-large", "properties": {"large": "x" * 256}}, + {"event": "accepted"}, + ]: + consumer.queue.put_nowait(_QueuedEvent(event, contextvars.copy_context())) + finally: + marker.reset(token) + consumer.queue.put_nowait(_STOP) + + with mock.patch( + "posthog._async_consumer.async_batch_post", new=mock.AsyncMock() + ) as batch_post: + await consumer.run() + await asyncio.wait_for(consumer.queue.join(), 1) + + assert batch_post.await_count == 1 + assert batch_post.await_args.kwargs["batch"] == [ + {"event": "accepted", "context": "queued"} + ] + assert marker.get() == "outside" + + +@pytest.mark.asyncio +async def test_next_carries_oversized_batch_tail_without_reprocessing(monkeypatch): + consumer = make_consumer(retries=0) + events = [{"event": "first"}, {"event": "second"}] + monkeypatch.setattr( + "posthog._async_consumer.BATCH_SIZE_LIMIT", + sum(len(json.dumps(event).encode()) for event in events) - 1, + ) + for event in events: + consumer.queue.put_nowait(_QueuedEvent(event, contextvars.copy_context())) + consumer.queue.put_nowait(_STOP) + + first, stopped_first = await consumer.next() + second, stopped_second = await consumer.next() + + assert first == [events[0]] + assert second == [events[1]] + assert (stopped_first, stopped_second) == (False, True) + assert consumer.process_event.await_count == 2 + with mock.patch.object(consumer, "request", new=mock.AsyncMock()): + await consumer.upload(first) + await consumer.upload(second) + await asyncio.wait_for(consumer.queue.join(), 1) + + +@pytest.mark.asyncio +async def test_upload_acknowledges_failed_batch_even_when_on_error_raises(): + consumer = make_consumer(retries=0) + consumer.on_error = mock.AsyncMock(side_effect=RuntimeError("callback failed")) + event = {"event": "failed"} + consumer.queue.put_nowait(event) + consumer.queue.get_nowait() + error = APIError(400, "terminal") + + with mock.patch( + "posthog._async_consumer.async_batch_post", + new=mock.AsyncMock(side_effect=error), + ): + await consumer.upload([event]) + await asyncio.wait_for(consumer.queue.join(), 1) + + consumer.on_error.assert_awaited_once_with(error, [event]) + + +@pytest.mark.asyncio +async def test_next_propagates_cancellation_during_event_processing(): + consumer = make_consumer(retries=0) + started = asyncio.Event() + cancelled = asyncio.Event() + + async def process(event): + started.set() + try: + await asyncio.Event().wait() + finally: + cancelled.set() + + consumer.process_event = process + consumer.queue.put_nowait( + _QueuedEvent({"event": "blocked"}, contextvars.copy_context()) + ) + task = asyncio.create_task(consumer.next()) + await asyncio.wait_for(started.wait(), 1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert cancelled.is_set() diff --git a/posthog/test/test_client_crap_flags.py b/posthog/test/test_client_crap_flags.py new file mode 100644 index 000000000..c6356c84c --- /dev/null +++ b/posthog/test/test_client_crap_flags.py @@ -0,0 +1,170 @@ +"""Client flag initialization and failure contracts protected during decomposition.""" + +import sys +from unittest import mock + +import pytest + +from posthog.client import Client +from posthog.request import APIError +from posthog.test.test_property_matching_version import definitions +from posthog.test.test_utils import FAKE_TEST_API_KEY, FakeRedis +from posthog.utils import FlagCache, RedisFlagCache + + +@pytest.fixture +def client(): + instance = Client( + FAKE_TEST_API_KEY, + secret_key="test-secret", + send=False, + enable_local_evaluation=False, + ) + instance.feature_flags = [] + yield instance + instance.shutdown() + + +@pytest.mark.parametrize( + "configured,expected,ttl", + [ + ("redis://", "redis://localhost:6379", 300), + ("redis://cache/2?ttl=17", "redis://cache:6379/2", 17), + ("redis://user@cache:6380/0", "redis://user:@cache:6380/0", 300), + ("redis://:secret@cache/0", "redis://:secret@cache:6379/0", 300), + ( + "rediss://user:secret@[::1]:6380/3?ttl=11", + "rediss://user:secret@[::1]:6380/3", + 11, + ), + ], +) +def test_redis_cache_initialization_preserves_url_and_ttl( + monkeypatch, configured, expected, ttl +): + redis_module = mock.Mock() + monkeypatch.setitem(sys.modules, "redis", redis_module) + instance = Client( + FAKE_TEST_API_KEY, + send=False, + enable_local_evaluation=False, + flag_fallback_cache_url=configured, + ) + try: + redis_module.from_url.assert_called_once_with(expected) + redis_module.from_url.return_value.ping.assert_called_once_with() + assert isinstance(instance.flag_cache, RedisFlagCache) + assert instance.flag_cache.default_ttl == ttl + finally: + instance.shutdown() + + +@pytest.mark.parametrize("failure", ["missing_dependency", "connection_failure"]) +def test_redis_cache_initialization_failure_disables_only_cache( + monkeypatch, caplog, failure +): + if failure == "missing_dependency": + monkeypatch.setitem(sys.modules, "redis", None) + expected_warning = "Redis not available" + else: + redis_module = mock.Mock() + redis_module.from_url.return_value.ping.side_effect = RuntimeError("offline") + monkeypatch.setitem(sys.modules, "redis", redis_module) + expected_warning = "Redis connection failed: offline" + instance = Client( + FAKE_TEST_API_KEY, + send=False, + enable_local_evaluation=False, + flag_fallback_cache_url="redis://cache/0", + ) + try: + assert instance.flag_cache is None + assert not instance.disabled + assert expected_warning in caplog.text + finally: + instance.shutdown() + + +@pytest.mark.parametrize("status", [401, 402]) +@pytest.mark.parametrize("debug", [False, True]) +def test_definition_error_invalidates_redis_snapshot_even_when_debug_raises( + client, status, debug +): + client.flag_cache = RedisFlagCache(FakeRedis()) + client._update_flag_state(definitions(2)) + fingerprint = client._flag_definition_fingerprint + previous_version = client.flag_definition_version + client.debug = debug + with ( + mock.patch("posthog.client.get", side_effect=APIError(status, "rejected")), + mock.patch.object(client.flag_cache, "_invalidate_snapshot") as invalidate, + ): + if debug: + with pytest.raises(APIError) as error: + client._load_feature_flags() + assert error.value.status == status + else: + client._load_feature_flags() + invalidate.assert_called_once_with(fingerprint) + assert client.feature_flags == [] + assert client._property_matching_version == 1 + assert client._flags_etag is None + assert client.flag_definition_version == previous_version + 1 + assert ( + client._flag_definition_published_generation + == client._flag_definition_fetch_generation + == client._flag_definition_cache_generation + ) + + +def test_remote_cache_write_failure_keeps_stale_value_and_received_metadata(client): + client.flag_cache = FlagCache() + initial = {"featureFlags": {"checkout": "control"}} + with mock.patch("posthog.client.flags", return_value=initial): + assert ( + client.get_feature_flag_result( + "checkout", "user", send_feature_flag_events=False + ).get_value() + == "control" + ) + + response = { + "requestId": "new-request", + "evaluatedAt": 123, + "flags": { + "checkout": { + "value": "test", + "reason": {"description": "New assignment"}, + "metadata": { + "id": 7, + "version": 2, + "config_version": 2, + "rule_type": "experiment", + "rule_id": "rule-id", + "variant_key": "test", + }, + } + }, + } + with ( + mock.patch("posthog.client.flags", return_value=response), + mock.patch.object( + client.flag_cache, + "_set_cached_remote_flag", + side_effect=RuntimeError("offline"), + ), + mock.patch.object(client, "capture") as capture, + ): + result = client.get_feature_flag_result("checkout", "user") + assert result.get_value() == "control" + capture.assert_called_once() + event = capture.call_args.kwargs + assert capture.call_args.args == ("$feature_flag_called",) + properties = event["properties"] + assert properties["$feature_flag_response"] == "control" + assert properties["$feature_flag_error"] == "unknown_error" + assert properties["$feature_flag_request_id"] == "new-request" + assert properties["$feature_flag_evaluated_at"] == 123 + assert properties["$feature_flag_id"] == 7 + assert properties["$feature_flag_reason"] == "New assignment" + assert "$feature_flag_rule_id" not in properties diff --git a/posthog/test/test_code_variables.py b/posthog/test/test_code_variables.py index 6263b4b8d..2efa0bbec 100644 --- a/posthog/test/test_code_variables.py +++ b/posthog/test/test_code_variables.py @@ -1400,3 +1400,74 @@ def trigger_error(): ) assert "code_variables" in output assert "n8fK2pQ9vX7mL4wR8tY3uZ6bC1dE5gH" in output # detection off -> retained + + +@pytest.mark.parametrize( + "base, value, expected", + [ + (str, "contains_password_here", REDACTED), + (int, 42, 42), + (float, 1.5, 1.5), + (float, float("inf"), "inf"), + ], +) +def test_scalar_subclasses_are_masked_before_object_attribute_traversal( + base, value, expected +): + class Scalar(base): + pass + + scalar = Scalar(value) + scalar.extra = {"password": "not part of the scalar wire representation"} + assert mask(scalar) == expected + assert encode(scalar) == expected + + +def test_oversized_namedtuple_is_replaced_without_reading_items(): + class Wide(collections.namedtuple("Wide", [f"field{n}" for n in range(51)])): + def __iter__(self): + raise AssertionError("oversized values must not be read") + + assert mask(Wide(*range(51))) == TOO_LONG + + +@pytest.mark.parametrize( + "value, expected", [(float("nan"), "nan"), (float("inf"), "inf")] +) +def test_disabled_masking_still_encodes_non_finite_numbers_as_strings(value, expected): + config = _MaskingConfig.build( + [], [], mask_url_credentials=False, detect_secrets=False + ) + assert _encode_variable(value, config, VariableSizeLimiter()) == expected + + +def test_numeric_variables_charge_the_shared_budget_without_string_truncation(): + limiter = VariableSizeLimiter(max_size=3) + assert encode(123, limiter=limiter) == 123 + assert limiter.get_remaining_space() == 0 + assert encode(1, limiter=limiter) is None + + +def test_encoding_failure_uses_masked_repr_and_charges_the_budget(): + # Reading this dict fails before it can be serialized structurally. Its repr holds + # a named secret, so the fallback must redact rather than expose that repr. + broken = _BrokenDict(password="hunter2") + limiter = VariableSizeLimiter(max_size=len(REDACTED)) + assert encode(broken, limiter=limiter) == REDACTED + assert limiter.get_remaining_space() == 0 + assert encode(broken, limiter=limiter) is None + + +def test_disabling_name_patterns_does_not_read_property_getters(): + class Opaque: + __slots__ = () + + @property + def password(self): + raise AssertionError("masking must not read getters") + + def __repr__(self): + return "" + + assert mask(Opaque(), patterns=[]) == "" + assert mask(Opaque())["password"] == REDACTED diff --git a/posthog/test/test_emails.py b/posthog/test/test_emails.py index eaf1c0949..dc076030d 100644 --- a/posthog/test/test_emails.py +++ b/posthog/test/test_emails.py @@ -799,3 +799,100 @@ def test_rejects_a_code_send_that_mixes_or_misses_fields(_name, overrides, error ) assert adapter.requests == [] + + +@pytest.mark.parametrize("factory", [Emails, AsyncEmails]) +@pytest.mark.parametrize( + "options, message", + [ + ({"timeout": True}, "timeout must be a finite positive number of seconds."), + ({"timeout": "3"}, "timeout must be a finite positive number of seconds."), + ( + {"timeout": float("inf")}, + "timeout must be a finite positive number of seconds.", + ), + ( + {"timeout": float("nan")}, + "timeout must be a finite positive number of seconds.", + ), + ({"timeout": 0}, "timeout must be a finite positive number of seconds."), + ({"max_retries": True}, "max_retries must be a nonnegative integer."), + ({"max_retries": None}, "max_retries must be a nonnegative integer."), + ({"max_retries": -1}, "max_retries must be a nonnegative integer."), + ( + {"migrated_from": "re send"}, + "migrated_from must be visible ASCII characters, such as 'resend'.", + ), + ], +) +def test_invalid_transport_options_are_rejected_at_construction( + factory, options, message +): + with pytest.raises(ValueError) as raised: + factory(SECRET_KEY, **options) + assert str(raised.value) == message + + +@pytest.mark.parametrize("warnings", [None, {}, "not a warning list"]) +def test_non_list_warnings_are_ignored(warnings): + emails, _ = emails_with([accepted(warnings=warnings, replayed=1)]) + result = send(emails) + assert result.warnings == () + assert result.replayed is False + + +def test_warning_parsing_keeps_order_fallback_fields_and_safe_defaults(caplog): + emails, _ = emails_with( + [ + accepted( + warnings=[ + None, + { + "code": "first", + "detail": "preferred", + "message": "ignored", + "fix": "update", + }, + {"code": "", "detail": 42, "message": "fallback", "fix": False}, + {}, + ] + ) + ] + ) + with caplog.at_level(logging.WARNING, logger="posthog"): + result = send(emails) + assert [(w.code, w.message, w.fix) for w in result.warnings] == [ + ("first", "preferred", "update"), + ("warning", "fallback", None), + ("warning", "", None), + ] + assert len(caplog.records) == 3 + assert all( + record.getMessage().endswith(message) + for record, message in zip( + caplog.records, + [ + "PostHog email warning first: preferred", + "PostHog email warning warning: fallback", + "PostHog email warning warning: ", + ], + ) + ) + + +def test_code_send_preserves_utf8_and_compact_json_bytes(): + emails, adapter = emails_with([accepted()]) + send_code(emails, subject="Bonjour é", html="

你好

") + assert adapter.requests[0].body == ( + '{"to":"user@example.com","template_key":"password-reset",' + '"subject":"Bonjour é","html":"

你好

","capture_engagement":true}' + ).encode("utf-8") + + +def test_code_send_validation_order_still_checks_recipient_before_template_and_content(): + emails, adapter = emails_with([]) + with pytest.raises(TypeError, match="to must be one email address as a string"): + send_code(emails, to=[], template_key="bad key", subject=123, html=123) + with pytest.raises(ValueError, match="template_key can only have letters"): + send_code(emails, template_key="bad key", subject=123, html=123) + assert adapter.requests == [] diff --git a/posthog/test/test_evaluate_flags.py b/posthog/test/test_evaluate_flags.py index dbe33583c..bb76ed85a 100644 --- a/posthog/test/test_evaluate_flags.py +++ b/posthog/test/test_evaluate_flags.py @@ -1207,6 +1207,96 @@ def test_unsupported_definition_logs_warning(self, _name, checkout): len(self._warnings(logs, "checkout", "unsupported_definition")), 1 ) + @parameterized.expand( + [ + ("non_object", []), + ("unknown_group_type", {"type": "XOR", "values": []}), + ("non_list_values", {"type": "AND", "values": {}}), + ("non_object_entry", {"type": "AND", "values": [42]}), + ( + "nested_unsupported_leaf", + { + "type": "OR", + "values": [ + { + "type": "AND", + "values": [ + {"key": "plan", "operator": "bogus", "value": "pro"} + ], + } + ], + }, + ), + ] + ) + @mock.patch("posthog.client.flags") + def test_malformed_dynamic_cohort_warns_and_remains_unresolved( + self, _name, cohort, patch_flags + ): + client = self._client() + definitions = self._definitions( + _flag_with("checkout", {"type": "cohort", "value": 42}) + ) + definitions["cohorts"] = {"42": cohort} + + with self.assertLogs("posthog", level=logging.WARNING) as logs: + self._refresh( + client, GetResponse(data=definitions, etag=None, not_modified=False) + ) + self.assertEqual( + len(self._warnings(logs, "checkout", "unsupported_definition")), 1 + ) + snapshot = client.evaluate_flags( + "user", person_properties={"plan": "pro"}, only_evaluate_locally=True + ) + self.assertEqual(snapshot.keys, []) + self.assertEqual( + dict(snapshot.unresolved_flags), + {"checkout": UnresolvedFlagReason.UNSUPPORTED_DEFINITION}, + ) + patch_flags.assert_not_called() + + @parameterized.expand( + [ + ("canonical_empty", {}), + ( + "nested_supported_group", + { + "type": "AND", + "values": [ + {}, + { + "type": "OR", + "values": [{"key": "plan", "value": "pro"}], + }, + ], + }, + ), + ] + ) + @mock.patch("posthog.client.flags") + def test_supported_dynamic_cohort_does_not_warn_or_fall_back( + self, _name, cohort, patch_flags + ): + client = self._client() + definitions = self._definitions( + _flag_with("checkout", {"type": "cohort", "value": 42}) + ) + definitions["cohorts"] = {"42": cohort} + + with self.assertNoLogs("posthog", level=logging.WARNING): + self._refresh( + client, GetResponse(data=definitions, etag=None, not_modified=False) + ) + snapshot = client.evaluate_flags( + "user", person_properties={"plan": "pro"}, only_evaluate_locally=True + ) + self.assertIs( + snapshot.get_flag("checkout", send_feature_flag_events=False), True + ) + self.assertEqual(dict(snapshot.unresolved_flags), {}) + patch_flags.assert_not_called() + def test_does_not_warn_for_flags_that_depend_on_the_call(self): client = self._client() diff --git a/posthog/test/test_exception_utils.py b/posthog/test/test_exception_utils.py index 007d575b4..e719c8c31 100644 --- a/posthog/test/test_exception_utils.py +++ b/posthog/test/test_exception_utils.py @@ -1,5 +1,10 @@ +import sys from datetime import datetime, timedelta, timezone +from types import SimpleNamespace +import pytest + +from posthog import exception_utils from posthog.exception_utils import format_timestamp @@ -16,3 +21,262 @@ def test_format_timestamp_converts_aware_value_to_utc(): ) assert format_timestamp(value) == "2026-06-27T12:00:00.123456Z" + + +# Source context, traversal and frame classification are observable payload contracts. +def test_source_context_keeps_file_order_and_five_lines_on_each_side(tmp_path): + source = tmp_path / "application.py" + source.write_text("\n".join(f"line {n}" for n in range(13))) + + assert exception_utils.get_lines_from_file(str(source), 6) == ( + ["line 1", "line 2", "line 3", "line 4", "line 5"], + "line 6", + ["line 7", "line 8", "line 9", "line 10", "line 11"], + ) + + +def test_source_context_prefers_loader_over_file(tmp_path): + source = tmp_path / "application.py" + source.write_text("stale contents") + loader = SimpleNamespace(get_source=lambda module: "before\r\ncurrent\r\nafter") + + assert exception_utils.get_lines_from_file( + str(source), 1, loader=loader, module="app" + ) == ( + ["before"], + "current", + ["after"], + ) + + +@pytest.mark.parametrize("failure", [ImportError, OSError]) +def test_source_context_falls_back_when_loader_cannot_read_source(tmp_path, failure): + def unavailable(module): + raise failure("source unavailable") + + source = tmp_path / "application.py" + source.write_text("current\nafter\n") + assert exception_utils.get_lines_from_file( + str(source), 0, loader=SimpleNamespace(get_source=unavailable), module="app" + ) == ([], "current", ["after"]) + + +def test_source_context_handles_changed_files(tmp_path): + source = tmp_path / "application.py" + source.write_text("current\n") + assert exception_utils.get_lines_from_file(str(source), 99) == ([], None, []) + + +def test_source_context_does_not_replace_empty_loader_source_with_stale_file(tmp_path): + source = tmp_path / "application.py" + source.write_text("stale contents") + loader = SimpleNamespace(get_source=lambda module: "") + assert exception_utils.get_lines_from_file(str(source), 0, loader=loader) == ( + [], + None, + [], + ) + + +def test_source_context_handles_unreadable_file(monkeypatch): + def unavailable(filename): + raise OSError("source unavailable") + + monkeypatch.setattr(exception_utils.linecache, "getlines", unavailable) + assert exception_utils.get_lines_from_file("application.py", 0) == ([], None, []) + + +@pytest.mark.parametrize("relationship", ["cause", "context"]) +def test_walk_exception_chain_is_outermost_first_and_stops_at_cycles(relationship): + outer, inner = RuntimeError("outer"), ValueError("inner") + if relationship == "cause": + outer.__cause__ = inner + inner.__cause__ = outer + else: + outer.__context__ = inner + inner.__context__ = outer + + assert list(exception_utils.walk_exception_chain((type(outer), outer, None))) == [ + (RuntimeError, outer, None), + (ValueError, inner, None), + ] + + +def test_walk_exception_chain_stops_at_suppressed_context(): + outer = RuntimeError("outer") + outer.__context__ = ValueError("ignored") + outer.__suppress_context__ = True + assert list(exception_utils.walk_exception_chain((RuntimeError, outer, None))) == [ + (RuntimeError, outer, None) + ] + assert list(exception_utils.walk_exception_chain((None, None, None))) == [] + + +def test_event_frame_iteration_keeps_top_level_thread_and_exception_order(): + first, second, third = ( + {"function": "main"}, + {"function": "thread"}, + {"function": "fail"}, + ) + event = { + "stacktrace": {"frames": [first]}, + "threads": {"values": [{}, {"stacktrace": {"frames": [second]}}]}, + "exception": { + "values": [{"stacktrace": {}}, {}, {"stacktrace": {"frames": [third]}}] + }, + } + assert list(exception_utils.iter_event_frames(event)) == [first, second, third] + assert ( + list(exception_utils.iter_event_frames({"threads": {}, "exception": {}})) == [] + ) + + +@pytest.mark.parametrize( + "frame, expected", + [ + ({"in_app": False, "module": "app"}, False), + ({"module": "app.handlers"}, True), + ({"module": "app2.handlers"}, None), + ({"module": "vendor.plugin"}, False), + ({"abs_path": "/venv/site-packages/pkg.py"}, False), + ({"abs_path": "/venv/dist-packages/pkg.py"}, False), + ({"abs_path": "/project/app.py"}, True), + ({"abs_path": "/elsewhere/app.py"}, None), + ({}, None), + ], +) +def test_frame_classification_preserves_overrides_and_module_boundaries( + frame, expected +): + frames = [frame] + assert ( + exception_utils.set_in_app_in_frames( + frames, + in_app_exclude=["vendor", "app"], + in_app_include=["app"], + project_root="/project/", + ) + is frames + ) + assert frame.get("in_app") is expected + + +def test_frame_classification_accepts_missing_options_and_empty_frames(): + frame = {"module": "app", "abs_path": "/project/app.py"} + exception_utils.set_in_app_in_frames([frame], None, None) + assert "in_app" not in frame + assert exception_utils.set_in_app_in_frames([], None, None) is None + + +def test_exc_info_preserves_explicit_tuple(): + error = ValueError("reported") + assert exception_utils.exc_info_from_error((ValueError, error, None)) == ( + ValueError, + error, + None, + ) + + +def test_exc_info_preserves_active_traceback_when_synthesis_fails(monkeypatch): + def unavailable(error): + raise RuntimeError("cannot synthesize traceback") + + monkeypatch.setattr(exception_utils, "construct_artificial_traceback", unavailable) + error = ValueError("active") + try: + raise error + except ValueError: + expected = sys.exc_info() + assert exception_utils.exc_info_from_error(error) == expected + + +def test_exc_info_does_not_borrow_an_unrelated_active_exception(monkeypatch): + monkeypatch.setattr( + exception_utils, "construct_artificial_traceback", lambda error: None + ) + error = ValueError("reported") + try: + raise RuntimeError("unrelated") + except RuntimeError: + assert exception_utils.exc_info_from_error(error) == (ValueError, error, None) + assert exception_utils.exc_info_from_error(error) == (ValueError, error, None) + + +@pytest.mark.parametrize("invalid", [None, "error", (ValueError, ValueError("bad"))]) +def test_exc_info_rejects_non_exception_inputs(invalid): + with pytest.raises(ValueError, match="Expected Exception object"): + exception_utils.exc_info_from_error(invalid) + + +@pytest.mark.parametrize( + "mechanism, expected", + [ + ( + { + "type": "custom", + "handled": False, + "source": "cause", + "synthetic": False, + "extra": "kept", + }, + { + "type": "custom", + "handled": False, + "source": "cause", + "synthetic": False, + "extra": "kept", + }, + ), + ( + { + "type": "", + "handled": 1, + "source": False, + "synthetic": "yes", + "extra": "kept", + }, + {"synthetic": False, "extra": "kept"}, + ), + (None, {"synthetic": False}), + ], +) +def test_exception_serialization_validates_mechanism_fields_independently( + mechanism, expected +): + error = ValueError("reported") + result = exception_utils.single_exception_from_error_tuple( + ValueError, error, None, mechanism=mechanism + ) + assert result == { + "mechanism": expected, + "module": None, + "type": "ValueError", + "value": "reported", + } + + +def test_exception_tree_truncates_long_cyclic_cause_chain_in_outermost_first_order(): + errors = [ValueError(str(n)) for n in range(60)] + for outer, inner in zip(errors, errors[1:]): + outer.__cause__ = inner + errors[-1].__cause__ = errors[0] + result = exception_utils.exceptions_from_error_tuple((ValueError, errors[0], None)) + assert [entry["value"] for entry in result] == [str(n) for n in range(50)] + assert [entry["mechanism"]["exception_id"] for entry in result] == list(range(50)) + assert [entry["mechanism"]["parent_id"] for entry in result[1:]] == list(range(49)) + assert all(entry["mechanism"]["source"] == "cause" for entry in result[1:]) + + +def test_exception_tree_stops_at_a_cycle_before_the_entry_limit(): + outer, inner = ValueError("outer"), RuntimeError("inner") + outer.__context__ = inner + inner.__context__ = outer + result = exception_utils.exceptions_from_error_tuple((ValueError, outer, None)) + assert [entry["value"] for entry in result] == ["outer", "inner"] + assert result[1]["mechanism"] == { + "exception_id": 1, + "parent_id": 0, + "type": "chained", + "source": "context", + "synthetic": False, + } diff --git a/posthog/test/test_feature_flags.py b/posthog/test/test_feature_flags.py index a18592db4..a994adeb8 100644 --- a/posthog/test/test_feature_flags.py +++ b/posthog/test/test_feature_flags.py @@ -506,6 +506,9 @@ def test_early_exit_preserves_fallback_after_inconclusive_presence_condition( "rollout_percentage": 100, }, {"properties": [], "rollout_percentage": 0}, + # Early exit must stop before this conclusive match, + # preserving the earlier missing-context fallback. + {"properties": [], "rollout_percentage": 100}, ], }, } diff --git a/posthog/test/test_metrics.py b/posthog/test/test_metrics.py index e25bb34bf..7185e631b 100644 --- a/posthog/test/test_metrics.py +++ b/posthog/test/test_metrics.py @@ -868,3 +868,47 @@ def test_failed_flushes_back_off_exponentially_capped(self): c.metrics.count("m", 1) assert c.metrics._flush_timer.interval == 1.0 c.metrics.reset() + + +@pytest.mark.parametrize("metric_type", ["count", "gauge", "histogram"]) +@pytest.mark.parametrize("current_sampled", [False, True]) +def test_failed_flush_merges_values_recorded_during_send( + client, metric_type, current_sampled +): + record = getattr(client.metrics, metric_type) + with trace.use_span(otel_span(sampled=True)): + record("in-flight", 3) + session = mock_session(503) + + def fail_with_concurrent_sample(*args, **kwargs): + with trace.use_span(otel_span(sampled=current_sampled)): + record("in-flight", 4) + return mock.Mock(status_code=503) + + session.post.side_effect = fail_with_concurrent_sample + with mock.patch("posthog.metrics_capture._get_session", return_value=session): + client.metrics.flush() + + payload, _, _ = flush_and_capture(client) + (metric,) = metrics_from(payload) + if metric_type == "count": + assert metric["sum"]["dataPoints"][0]["asDouble"] == 7.0 + elif metric_type == "gauge": + assert metric["gauge"]["dataPoints"][0]["asDouble"] == 4.0 + else: + (point,) = metric["histogram"]["dataPoints"] + assert (point["count"], point["sum"], point["min"], point["max"]) == ( + 2, + 7.0, + 3.0, + 4.0, + ) + assert point["bucketCounts"] == [0, 2] + [0] * 14 + assert len(metrics_from(payload)) == 1 + (point,) = data_points_by_type(payload).values() + (exemplar,) = point["exemplars"] + assert (exemplar["traceId"], exemplar["spanId"], exemplar["asDouble"]) == ( + OTEL_TRACE_ID, + OTEL_SPAN_ID, + 4.0 if current_sampled else 3.0, + ) diff --git a/posthog/test/test_metrics_autocapture.py b/posthog/test/test_metrics_autocapture.py index 57221289d..ab77e2c66 100644 --- a/posthog/test/test_metrics_autocapture.py +++ b/posthog/test/test_metrics_autocapture.py @@ -443,6 +443,60 @@ def test_limits_http_attribute_sets(self, posthog_host, app_server, started): assert len(attributes) == 2 assert {"otel.metric.overflow": True} in attributes + def test_fork_reset_replaces_held_collection_locks_without_acquiring_them( + self, posthog_host, started + ): + # Exercise the reset callback in-process: child coverage is not recorded, + # and a held inherited lock must never be acquired by the reset itself. + handle = _start( + posthog_host, + started, + export_interval_seconds=3600, + areas={"runtime": True}, + ) + meter = handle._meter_provider.get_meter("fork-reset-test") + histogram = meter.create_histogram("fork.reset.duration") + histogram.record(1) + consumer = handle._meter_provider._measurement_consumer + owners = [(meter, "_instrument_registration_lock"), (consumer, "_lock")] + for storage in consumer._reader_storages.values(): + owners.append((storage, "_lock")) + for matches in storage._instrument_view_instrument_matches.values(): + for match in matches: + owners.append((match, "_lock")) + for aggregation in match._attributes_aggregation.values(): + owners.append((aggregation, "_lock")) + reservoir = getattr(aggregation, "_reservoir", None) + if getattr(reservoir, "_lock", None) is not None: + owners.append((reservoir, "_lock")) + owners.extend( + [(handle._reader, "_export_lock"), (metrics_autocapture, "_lock")] + ) + original_locks = [getattr(owner, name) for owner, name in owners] + runtime = next(i for i in handle._instrumentors if hasattr(i, "_proc")) + original_process = runtime._proc + handle._flush_requested = 8 + handle._flush_exported = 5 + + with ExitStack() as stack: + for lock in original_locks: + stack.enter_context(_held_by_another_thread(lock)) + metrics_autocapture._reinit_after_fork() + for (owner, name), original in zip(owners, original_locks): + replacement = getattr(owner, name) + assert replacement is not original + assert type(replacement) is type(original) + assert replacement.acquire(blocking=False) + replacement.release() + histogram.record(2) + handle.force_flush() + + assert runtime._proc is not original_process + assert runtime._proc.pid == os.getpid() + assert handle._flush_thread is None + assert handle._flush_requested == handle._flush_exported == 0 + assert _metrics_exports(posthog_host) + @pytest.mark.skipif(not hasattr(os, "fork"), reason="needs os.fork") @pytest.mark.filterwarnings("ignore:This process .* is multi-threaded") def test_keeps_working_in_a_forked_child(self, posthog_host, app_server, started): diff --git a/posthog/test/tracing/test_export.py b/posthog/test/tracing/test_export.py index 52cff1515..0b6851035 100644 --- a/posthog/test/tracing/test_export.py +++ b/posthog/test/tracing/test_export.py @@ -394,6 +394,55 @@ def test_never_returns_a_span_it_failed_to_encode(self): assert sender.payloads == [] assert queued(pipeline) == [] + def test_encoding_failure_keeps_the_retry_batch_separate_from_later_spans( + self, clock, caplog + ): + caplog.set_level("WARNING", logger="posthog") + sender = FakeSender(SendOutcome("retry-later"), SendOutcome("ok")) + pipeline, _, _ = make_traces(sender=sender, max_export_batch_size=4) + pipeline.start_span("bad").end() + pipeline.start_span("retry-me").end() + encode = export_module.build_otlp_span + + def encode_unless_bad(record): + if record.name == "bad": + raise RuntimeError("bad") + return encode(record) + + with mock.patch.object(export_module, "build_otlp_span", encode_unless_bad): + pipeline.flush() + pipeline.start_span("later").end() + clock["now"] += 100 + pipeline.flush() + + assert [[span["name"] for span in batch] for batch in sender.batches()] == [ + ["retry-me"], + ["retry-me"], + ["later"], + ] + assert sender.batches()[0] == sender.batches()[1] + assert queued(pipeline) == [] + warnings = [ + r.getMessage() for r in caplog.records if "Dropping" in r.getMessage() + ] + assert len(warnings) == 1 + assert "Dropping 1 span(s)" in warnings[0] + + def test_close_during_encoding_abandons_the_batch_without_sending(self): + pipeline, sender, _ = make_traces() + pipeline.start_span("closing").end() + pipeline.start_span("queued").end() + + def close_and_fail(record): + pipeline.close() + raise RuntimeError("encoder failed after shutdown") + + with mock.patch.object(export_module, "build_otlp_span", close_and_fail): + pipeline.flush() + + assert sender.payloads == [] + assert queued(pipeline) == [] + def test_counts_a_span_it_failed_to_encode_once(self, caplog): caplog.set_level("WARNING", logger="posthog") sender = FakeSender(SendOutcome("fatal")) diff --git a/posthog/tracing/_config.py b/posthog/tracing/_config.py index 579d9ff48..0951080d5 100644 --- a/posthog/tracing/_config.py +++ b/posthog/tracing/_config.py @@ -143,27 +143,31 @@ def _resolve_before_span_send(value: Any) -> Tuple[Callable, ...]: return tuple(supplied) -def resolve_traces_config( - config: Any, host_resource_attributes: Optional[Mapping[str, str]] = None -) -> ResolvedTracesConfig: - """Validate the ``traces`` option, falling back to documented defaults per key. - - Host resource attributes (``os.name``, ``os.version``) merge first so user - ``resource_attributes`` win; a string ``service.name`` / ``service.version`` - / ``deployment.environment`` there wins over the named fields. - """ +def _config_mapping(config: Any) -> Mapping: if not isinstance(config, Mapping): if config is not None: log.warning( "Ignoring traces config: expected a dict, got %s", type(config).__name__ ) config = {} - unknown = [key for key in config if key not in _KNOWN_KEYS] if unknown: log.warning( "Ignoring unknown traces option(s): %s", ", ".join(map(str, unknown)) ) + return config + + +def resolve_traces_config( + config: Any, host_resource_attributes: Optional[Mapping[str, str]] = None +) -> ResolvedTracesConfig: + """Validate the ``traces`` option, falling back to documented defaults per key. + + Host resource attributes (``os.name``, ``os.version``) merge first so user + ``resource_attributes`` win; a string ``service.name`` / ``service.version`` + / ``deployment.environment`` there wins over the named fields. + """ + config = _config_mapping(config) resource_attributes: Dict[str, Any] = dict(host_resource_attributes or {}) resource_attributes.update( diff --git a/posthog/tracing/_context.py b/posthog/tracing/_context.py index 25c5299da..24ea98a42 100644 --- a/posthog/tracing/_context.py +++ b/posthog/tracing/_context.py @@ -28,6 +28,18 @@ def _enrichment_is_enabled() -> bool: return bool((project_api_key or "").strip() or (api_key or "").strip()) +def _missing_context_attributes(attributes: Mapping) -> dict: + context = contexts._get_active_span_context() + extra = {} + session_id = context["session_id"] + distinct_id = context["distinct_id"] + if session_id and "sessionId" not in attributes: + extra["sessionId"] = session_id + if distinct_id and not any(key in attributes for key in _IDENTITY_ATTRIBUTES): + extra["posthogDistinctId"] = distinct_id + return extra + + def _enrich_span(span: object) -> None: try: if not _enrichment_is_enabled() or not getattr(span, "is_recording")(): @@ -35,14 +47,7 @@ def _enrich_span(span: object) -> None: attributes = getattr(span, "attributes", None) if not isinstance(attributes, Mapping): return - context = contexts._get_active_span_context() - extra = {} - session_id = context["session_id"] - distinct_id = context["distinct_id"] - if session_id and "sessionId" not in attributes: - extra["sessionId"] = session_id - if distinct_id and not any(key in attributes for key in _IDENTITY_ATTRIBUTES): - extra["posthogDistinctId"] = distinct_id + extra = _missing_context_attributes(attributes) if not extra: return # The SDK exposes no public capacity getter and evicts old attributes diff --git a/posthog/tracing/_export.py b/posthog/tracing/_export.py index e68385f78..a68c5716e 100644 --- a/posthog/tracing/_export.py +++ b/posthog/tracing/_export.py @@ -202,11 +202,7 @@ def _drain_within_budget(self, deadline: Optional[float], retry: bool) -> None: while True: removed, stop = self._drain(deadline) if not stop: - if ( - removed - and self._queue - and (deadline is None or time.monotonic() < deadline) - ): + if self._followup_due(removed, deadline): self._drain(deadline) return if deadline is None or not retry: @@ -218,6 +214,13 @@ def _drain_within_budget(self, deadline: Optional[float], retry: bool) -> None: if self._wait_for_retry(min(wait, remaining)): return + def _followup_due(self, removed: int, deadline: Optional[float]) -> bool: + return bool( + removed + and self._queue + and (deadline is None or time.monotonic() < deadline) + ) + def _retry_wait_locked(self) -> Optional[float]: """The backoff to wait out before retrying, or ``None`` when there is nothing to retry.""" with self._lock: @@ -281,6 +284,11 @@ def _drain(self, deadline: Optional[float]) -> Tuple[int, bool]: if self._is_closed() or self._discard_if_disabled_locked(): return 0, True remaining = len(self._queue) + return self._drain_batches(remaining, deadline) + + def _drain_batches( + self, remaining: int, deadline: Optional[float] + ) -> Tuple[int, bool]: removed = 0 sent_any = False # A batch the SDK measured as too large says nothing about the batches @@ -292,70 +300,88 @@ def _drain(self, deadline: Optional[float]) -> Tuple[int, bool]: return removed, True if not self._queue: break - # A retried batch cannot grow to take in spans behind it. - cap = ( - min(self._max_export_batch_size, self._head_batch_size) - if self._head_batch_failures - else self._max_export_batch_size - ) - if local_cap is not None: - cap = min(cap, local_cap) - size = max(1, min(cap, remaining, len(self._queue))) - batch = self._queue[:size] - # Read before the send. A send inside an open Retry-After - # window is caller-driven and exempt from the wait, and so from - # the charge. - chargeable = ( - time.monotonic() >= self._head_batch_chargeable_at - and not self._retry_after.is_open() - ) + batch, chargeable = self._batch_locked(remaining, local_cap) + size = len(batch) - spans, failed = self._encode(batch) - if failed: - # Out of the queue before the send is settled, so a failed - # send does not count them a second time. - with self._lock: - if self._is_closed(): - return removed, True - for index in reversed(failed): - del self._queue[index] - if not spans: - self._reset_head_batch_budget_locked() - size -= len(failed) - remaining -= len(failed) - removed += len(failed) - if not spans: - continue + spans, failed = self._encode_queued_batch(batch) + if spans is None: + return removed, True + size -= failed + remaining -= failed + removed += failed + if not spans: + continue # The first request is exempt, so a flush called with no budget left # (a serverless handler, say) still ships a batch. - if sent_any and deadline is not None and time.monotonic() >= deadline: + if _request_budget_spent(sent_any, deadline): return removed, True sent_any = True - try: - outcome = self._send( - self._client, build_traces_payload(spans, self._resource) - ) - except Exception: - log.debug("Span batch send failed", exc_info=True) - outcome = SendOutcome("retry-later") - - with self._lock: - if self._is_closed(): - return removed, True - taken, stop, note = self._apply_outcome_locked( - outcome, size, chargeable - ) - if note: - log.debug(note) - if outcome.measured_locally and not taken: - local_cap = max(1, size // 2) + outcome, taken, stop = self._send_and_settle(spans, size, chargeable) + local_cap = _local_batch_cap(outcome, taken, size, local_cap) remaining -= taken removed += taken if stop: return removed, True return removed, False + def _send_and_settle( + self, spans: List[dict], size: int, chargeable: bool + ) -> Tuple[SendOutcome, int, bool]: + try: + outcome = self._send( + self._client, build_traces_payload(spans, self._resource) + ) + except Exception: + log.debug("Span batch send failed", exc_info=True) + outcome = SendOutcome("retry-later") + + with self._lock: + if self._is_closed(): + return outcome, 0, True + taken, stop, note = self._apply_outcome_locked(outcome, size, chargeable) + if note: + log.debug(note) + return outcome, taken, stop + + def _encode_queued_batch( + self, batch: List[SpanRecord] + ) -> Tuple[Optional[List[dict]], int]: + spans, failed = self._encode(batch) + if failed: + # Out of the queue before the send is settled, so a failed + # send does not count them a second time. + with self._lock: + if self._is_closed(): + return None, 0 + for index in reversed(failed): + del self._queue[index] + if not spans: + self._reset_head_batch_budget_locked() + return spans, len(failed) + + def _batch_locked( + self, remaining: int, local_cap: Optional[int] + ) -> Tuple[List[SpanRecord], bool]: + # A retried batch cannot grow to take in spans behind it. + cap = ( + min(self._max_export_batch_size, self._head_batch_size) + if self._head_batch_failures + else self._max_export_batch_size + ) + if local_cap is not None: + cap = min(cap, local_cap) + size = max(1, min(cap, remaining, len(self._queue))) + batch = self._queue[:size] + # Read before the send. A send inside an open Retry-After + # window is caller-driven and exempt from the wait, and so from + # the charge. + chargeable = ( + time.monotonic() >= self._head_batch_chargeable_at + and not self._retry_after.is_open() + ) + return batch, chargeable + def _apply_outcome_locked( self, outcome: SendOutcome, size: int, chargeable: bool ) -> Tuple[int, bool, Optional[str]]: @@ -373,30 +399,7 @@ def _apply_outcome_locked( return size, False, None if outcome.kind == "too-large": - if size == 1: - del self._queue[:1] - self._end_failure_sequence_locked() - self._drops.record(1, "it is too large for the ingestion endpoint") - # The oversized span is gone; the batches after it are not suspect. - if self._batch_size_before_halving is not None: - self._max_export_batch_size = self._batch_size_before_halving - self._batch_size_before_halving = None - return 1, False, None - # Halve the refused batch, not the configured size: a shallow queue - # would otherwise resend the same body. - halved = max(1, size // 2) - if not outcome.measured_locally: - if self._batch_size_before_halving is None: - self._batch_size_before_halving = self._max_export_batch_size - self._max_export_batch_size = halved - self._reset_head_batch_budget_locked() - return ( - 0, - False, - "Span batch too large; retrying the same spans in batches of {}".format( - halved - ), - ) + return self._apply_too_large_locked(outcome, size) if outcome.kind == "retry-later": self._consecutive_failures += 1 @@ -428,6 +431,34 @@ def _apply_outcome_locked( self._drops.record(size, "the ingestion endpoint rejected the batch") return size, False, None + def _apply_too_large_locked( + self, outcome: SendOutcome, size: int + ) -> Tuple[int, bool, Optional[str]]: + if size == 1: + del self._queue[:1] + self._end_failure_sequence_locked() + self._drops.record(1, "it is too large for the ingestion endpoint") + # The oversized span is gone; the batches after it are not suspect. + if self._batch_size_before_halving is not None: + self._max_export_batch_size = self._batch_size_before_halving + self._batch_size_before_halving = None + return 1, False, None + # Halve the refused batch, not the configured size: a shallow queue + # would otherwise resend the same body. + halved = max(1, size // 2) + if not outcome.measured_locally: + if self._batch_size_before_halving is None: + self._batch_size_before_halving = self._max_export_batch_size + self._max_export_batch_size = halved + self._reset_head_batch_budget_locked() + return ( + 0, + False, + "Span batch too large; retrying the same spans in batches of {}".format( + halved + ), + ) + def _is_closed(self) -> bool: # A method, so the check after each send is not narrowed away. return self._closed @@ -534,3 +565,15 @@ def _clear_timer_locked(self) -> None: def _draw_jitter() -> float: return 1 - FLUSH_BACKOFF_JITTER + random.random() * FLUSH_BACKOFF_JITTER * 2 + + +def _request_budget_spent(sent_any: bool, deadline: Optional[float]) -> bool: + return sent_any and deadline is not None and time.monotonic() >= deadline + + +def _local_batch_cap( + outcome: SendOutcome, taken: int, size: int, current: Optional[int] +) -> Optional[int]: + if outcome.measured_locally and not taken: + return max(1, size // 2) + return current diff --git a/posthog/tracing/_limits.py b/posthog/tracing/_limits.py index faeb55428..ae6f06d59 100644 --- a/posthog/tracing/_limits.py +++ b/posthog/tracing/_limits.py @@ -55,19 +55,7 @@ def _truncate(value: Any, max_length: int, state: _WalkState, depth: int) -> Any return None is_container = isinstance(value, (Mapping, list, tuple, set, frozenset)) if not is_container: - # Leaves are charged too: a shared subtree is walked once per path. - if state.remaining_nodes <= 0: - return value - state.remaining_nodes -= 1 - if isinstance(value, str): - return truncate_string(value, max_length) - if isinstance(value, (bool, int, float, date)) or callable(value): - return value - # The encoder stringifies anything else, so bound that text. - try: - return truncate_string(str(value), max_length) - except Exception: - return UNSERIALIZABLE_VALUE + return _truncate_leaf(value, max_length, state) marker = id(value) if marker in state.ancestors: @@ -86,6 +74,22 @@ def _truncate(value: Any, max_length: int, state: _WalkState, depth: int) -> Any state.ancestors.discard(marker) +def _truncate_leaf(value: Any, max_length: int, state: _WalkState) -> Any: + # Leaves are charged too: a shared subtree is walked once per path. + if state.remaining_nodes <= 0: + return value + state.remaining_nodes -= 1 + if isinstance(value, str): + return truncate_string(value, max_length) + if isinstance(value, (bool, int, float, date)) or callable(value): + return value + # The encoder stringifies anything else, so bound that text. + try: + return truncate_string(str(value), max_length) + except Exception: + return UNSERIALIZABLE_VALUE + + def _truncate_mapping( value: Mapping, max_length: int, state: _WalkState, depth: int ) -> dict: @@ -188,6 +192,17 @@ def apply_span_limits( non_negative_count(record.dropped_attributes_count) + dropped_attributes ) + _apply_event_limits(record, max_events, max_attributes_per_event, max_length) + + if record.status is not None and record.status.message: + record.status = SpanStatus( + record.status.code, truncate_string(record.status.message, max_length) + ) + + +def _apply_event_limits( + record: SpanRecord, max_events: int, max_attributes_per_event: int, max_length: int +) -> None: kept_events: list = [] dropped_events = 0 for event in record.events: @@ -208,8 +223,3 @@ def apply_span_limits( record.dropped_events_count = ( non_negative_count(record.dropped_events_count) + dropped_events ) - - if record.status is not None and record.status.message: - record.status = SpanStatus( - record.status.code, truncate_string(record.status.message, max_length) - ) diff --git a/posthog/tracing/_otlp.py b/posthog/tracing/_otlp.py index d5e0745b3..aba6e3922 100644 --- a/posthog/tracing/_otlp.py +++ b/posthog/tracing/_otlp.py @@ -153,47 +153,59 @@ def _encode(value: Any, state: _EncodeState, depth: int) -> dict: if isinstance(value, bool): return {"boolValue": value} if isinstance(value, int): - if value > INT64_MAX or value < INT64_MIN: - log.debug( - "Attribute %s is outside the int64 range; encoding it as a string", - value, - ) - return {"stringValue": str(value)} - # int(): an IntEnum stringifies as its name. - return {"intValue": str(int(value))} + return _encode_int(value) if isinstance(value, float): - if not math.isfinite(value): - if math.isnan(value): - return {"stringValue": "NaN"} - return {"stringValue": "Infinity" if value > 0 else "-Infinity"} - return {"doubleValue": value} + return _encode_float(value) if isinstance(value, str): return {"stringValue": sanitize_string(value)} if isinstance(value, (datetime, date)): return {"stringValue": value.isoformat()} - if isinstance(value, Mapping) or isinstance(value, (list, tuple, set, frozenset)): - marker = id(value) - if marker in state.ancestors: - return {"stringValue": CIRCULAR_VALUE} - if depth >= MAX_VALUE_DEPTH: - return {"stringValue": TRUNCATED_VALUE} - state.ancestors.add(marker) - try: - if isinstance(value, Mapping): - return { - "kvlistValue": { - "values": _encode_key_value_list(value, state, depth + 1)[0] - } - } - return {"arrayValue": {"values": _encode_array(value, state, depth + 1)}} - finally: - # Siblings sharing one object are not a cycle. - state.ancestors.discard(marker) + if isinstance(value, (Mapping, list, tuple, set, frozenset)): + return _encode_container(value, state, depth) if callable(value): return {"stringValue": FUNCTION_VALUE} return {"stringValue": wire_string(value)} +def _encode_int(value: int) -> dict: + if value > INT64_MAX or value < INT64_MIN: + log.debug( + "Attribute %s is outside the int64 range; encoding it as a string", + value, + ) + return {"stringValue": str(value)} + # int(): an IntEnum stringifies as its name. + return {"intValue": str(int(value))} + + +def _encode_float(value: float) -> dict: + if not math.isfinite(value): + if math.isnan(value): + return {"stringValue": "NaN"} + return {"stringValue": "Infinity" if value > 0 else "-Infinity"} + return {"doubleValue": value} + + +def _encode_container(value: Any, state: _EncodeState, depth: int) -> dict: + marker = id(value) + if marker in state.ancestors: + return {"stringValue": CIRCULAR_VALUE} + if depth >= MAX_VALUE_DEPTH: + return {"stringValue": TRUNCATED_VALUE} + state.ancestors.add(marker) + try: + if isinstance(value, Mapping): + return { + "kvlistValue": { + "values": _encode_key_value_list(value, state, depth + 1)[0] + } + } + return {"arrayValue": {"values": _encode_array(value, state, depth + 1)}} + finally: + # Siblings sharing one object are not a cycle. + state.ancestors.discard(marker) + + def _encode_array(values: Any, state: _EncodeState, depth: int) -> list: result: list = [] count = 0 @@ -328,14 +340,21 @@ def build_otlp_span(record: SpanRecord) -> dict: dropped_events = non_negative_count(record.dropped_events_count) if dropped_events: span["droppedEventsCount"] = dropped_events - if record.status is not None and record.status.code in SPAN_STATUS_TO_OTLP: - status: dict = {"code": SPAN_STATUS_TO_OTLP[record.status.code]} - if record.status.message: - status["message"] = wire_string(record.status.message) + status = _to_otlp_status(record.status) + if status is not None: span["status"] = status return span +def _to_otlp_status(value: Optional[SpanStatus]) -> Optional[dict]: + if value is not None and value.code in SPAN_STATUS_TO_OTLP: + status: dict = {"code": SPAN_STATUS_TO_OTLP[value.code]} + if value.message: + status["message"] = wire_string(value.message) + return status + return None + + def build_traces_payload(spans: List[dict], resource: list) -> dict: """Wrap spans in the OTLP envelope: one resource, one scope, N spans per batch. diff --git a/posthog/tracing/_pipeline.py b/posthog/tracing/_pipeline.py index b09d25dff..1562773f7 100644 --- a/posthog/tracing/_pipeline.py +++ b/posthog/tracing/_pipeline.py @@ -140,10 +140,7 @@ def _start_span( if self._closed or getattr(self._client, "disabled", False): return inert_span(parent, tracestate, self._active_var) - parent = traceparent_header(parent) - if isinstance(parent, str) and not parent.strip(): - # `headers.get("traceparent", "")` with no header: the default applies. - parent = None + parent = _normalize_parent(parent) if parent is not None and not isinstance(parent, (str, RecordingSpan)): if isinstance(parent, Span): # A child of an inert handle is inert too, still forwarding any @@ -156,10 +153,26 @@ def _start_span( else: parent_context = self._resolve_parent(parent, tracestate) + span_id = self._reserve_span() + if span_id is None: + return inert_span(parent, tracestate, self._active_var) + + try: + return self._build_span( + span_id, name, kind, attributes, parent_context, start_time + ) + except Exception: + # The reservation would otherwise hold its slot until age eviction. + with self._lock: + self._live_spans.pop(span_id, None) + raise + + def _reserve_span(self) -> Optional[str]: + """Reserve a live-span slot, reporting evictions outside the lock.""" with self._lock: if self._is_closed(): # close() may have run since the check above. - return inert_span(parent, tracestate, self._active_var) + return None # Swept only at the bound, so a long span that does end is still # exported, while a process that leaked its way there recovers. aged = 0 @@ -181,17 +194,9 @@ def _start_span( "the live-span limit ({}) was reached; spans are being started " "and never ended".format(self._config.max_live_spans), ) - return inert_span(parent, tracestate, self._active_var) + return None - try: - return self._build_span( - span_id, name, kind, attributes, parent_context, start_time - ) - except Exception: - # The reservation would otherwise hold its slot until age eviction. - with self._lock: - self._live_spans.pop(span_id, None) - raise + return span_id def _build_span( self, @@ -215,7 +220,7 @@ def _build_span( start_ns=start_ns, backdated=start_ns != now_ns, on_end=self._on_span_end, - kind=kind if isinstance(kind, str) and kind else "internal", + kind=_recording_kind(kind), attributes=span_attributes, parent_span_id=parent_context.parent_span_id if parent_context else None, trace_state=parent_context.trace_state if parent_context else None, @@ -331,3 +336,15 @@ def _report_limit_drops(record: SpanRecord) -> None: record.dropped_events_count, event_attributes, ) + + +def _recording_kind(kind: Any) -> str: + return kind if isinstance(kind, str) and kind else "internal" + + +def _normalize_parent(parent: Any) -> Any: + parent = traceparent_header(parent) + if isinstance(parent, str) and not parent.strip(): + # `headers.get("traceparent", "")` with no header: the default applies. + return None + return parent diff --git a/posthog/tracing/_sanitize.py b/posthog/tracing/_sanitize.py index d87299da8..8f57770ff 100644 --- a/posthog/tracing/_sanitize.py +++ b/posthog/tracing/_sanitize.py @@ -62,16 +62,10 @@ def to_epoch_ns(value: Any) -> Optional[int]: """A ``datetime`` or epoch seconds as unix nanoseconds; ``None`` when unusable.""" if value is None or isinstance(value, bool): return None - ns: int + ns: Optional[int] if isinstance(value, datetime): - try: - if value.tzinfo is None: - value = value.astimezone() - delta = value - _EPOCH - ns = ( - delta.days * 86400 + delta.seconds - ) * 10**9 + delta.microseconds * 1000 - except Exception: + ns = _datetime_epoch_ns(value) + if ns is None: return None elif isinstance(value, int): ns = value * 10**9 @@ -82,11 +76,21 @@ def to_epoch_ns(value: Any) -> Optional[int]: ns = int(round(scaled)) else: return None - if ns < MIN_TIMESTAMP_NS or ns > MAX_TIMESTAMP_NS: + if not MIN_TIMESTAMP_NS <= ns <= MAX_TIMESTAMP_NS: return None return ns +def _datetime_epoch_ns(value: datetime) -> Optional[int]: + try: + if value.tzinfo is None: + value = value.astimezone() + delta = value - _EPOCH + return (delta.days * 86400 + delta.seconds) * 10**9 + delta.microseconds * 1000 + except Exception: + return None + + def resolve_start_ns(value: Any, now_ns: int) -> int: """A caller-supplied start time, or now; warns when the server will clamp it.""" supplied = to_epoch_ns(value) diff --git a/posthog/tracing/_span.py b/posthog/tracing/_span.py index f78ae464f..4dd0b3a1e 100644 --- a/posthog/tracing/_span.py +++ b/posthog/tracing/_span.py @@ -288,9 +288,7 @@ def _write_attribute(self, key: str, value: Any) -> None: return if value is None: # None removes the key, freeing its slot. - if key in self._attributes and key not in self._auto_keys: - self._user_attribute_count -= 1 - self._attributes.pop(key, None) + self._remove_attribute_locked(key) return # Checked before the value is walked, which is the costly part. if key not in self._auto_keys and key not in self._attributes: @@ -305,6 +303,11 @@ def _write_attribute(self, key: str, value: Any) -> None: if not self._ended: self._attributes[key] = bounded + def _remove_attribute_locked(self, key: str) -> None: + if key in self._attributes and key not in self._auto_keys: + self._user_attribute_count -= 1 + self._attributes.pop(key, None) + def set_attribute(self, key: str, value: Any) -> "Span": if self._mutable("set_attribute"): key_str = attribute_key(key) diff --git a/posthog/tracing/_traceparent.py b/posthog/tracing/_traceparent.py index 974527b74..e8ee9072e 100644 --- a/posthog/tracing/_traceparent.py +++ b/posthog/tracing/_traceparent.py @@ -50,6 +50,12 @@ def _header_text(value: object) -> Optional[str]: return value if isinstance(value, str) else None +def _valid_version_fields(version: str, trailing: Optional[str]) -> bool: + if version == "ff" or (version == "00" and trailing): + return False + return not (trailing and _NON_PRINTABLE_ASCII_RE.search(trailing)) + + def _match_traceparent(value: object) -> Optional[_TraceparentFields]: value = _header_text(value) if value is None: @@ -62,9 +68,7 @@ def _match_traceparent(value: object) -> Optional[_TraceparentFields]: if match is None: return None version, trace_id, span_id, flags, trailing = match.group(1, 2, 3, 4, 5) - if version == "ff" or (version == "00" and trailing): - return None - if trailing and _NON_PRINTABLE_ASCII_RE.search(trailing): + if not _valid_version_fields(version, trailing): return None if trace_id == INVALID_TRACE_ID or span_id == INVALID_SPAN_ID: return None diff --git a/posthog/tracing/_transport.py b/posthog/tracing/_transport.py index a261a0e41..9474ce693 100644 --- a/posthog/tracing/_transport.py +++ b/posthog/tracing/_transport.py @@ -66,23 +66,30 @@ def parse_retry_after(value: Any, now: Optional[datetime] = None) -> Optional[fl raw = value.strip() if re.match(r"^\d+\s*,", raw): raw = raw.split(",", 1)[0].strip() + seconds: Optional[float] if _DELTA_SECONDS_RE.match(raw): seconds = float(raw) elif _NUMERIC_RE.match(raw): return None else: - try: - when = parsedate_to_datetime(raw) - except (TypeError, ValueError, IndexError): + seconds = _retry_after_date_seconds(raw, now) + if seconds is None: return None - if when.tzinfo is None: - when = when.replace(tzinfo=timezone.utc) - seconds = (when - (now or datetime.now(timezone.utc))).total_seconds() if not math.isfinite(seconds) or seconds <= 0: return None return seconds +def _retry_after_date_seconds(raw: str, now: Optional[datetime]) -> Optional[float]: + try: + when = parsedate_to_datetime(raw) + except (TypeError, ValueError, IndexError): + return None + if when.tzinfo is None: + when = when.replace(tzinfo=timezone.utc) + return (when - (now or datetime.now(timezone.utc))).total_seconds() + + def send_traces_batch(client: Any, payload: dict) -> SendOutcome: """POST one OTLP batch with bearer auth and gzip, classifying the response. diff --git a/posthog/types.py b/posthog/types.py index 229b25a16..625b5dc5a 100644 --- a/posthog/types.py +++ b/posthog/types.py @@ -358,6 +358,24 @@ def _legacy_rendering(value: Any) -> tuple[bool, Optional[str], Optional[str]]: return True, None, json.dumps(value, separators=(",", ":"), ensure_ascii=False) +def _parse_typed_reason(value: Any) -> Optional[FlagReason]: + reason = None + raw_reason = _read_typed(value, "reason", dict) + if raw_reason is not None: + reason = FlagReason( + code=_read_typed(raw_reason.get("code"), "reason.code", str) or "", + condition_index=_read_int( + raw_reason.get("condition_index"), "reason.condition_index" + ), + description=_read_typed( + raw_reason.get("description"), "reason.description", str + ) + or "", + ) + + return reason + + def _parse_typed_record(cls: Type[_F], key: str, resp: Any) -> _F: """Parse a ``/flags?v=4`` record, which carries a typed ``value``. @@ -372,19 +390,7 @@ def _parse_typed_record(cls: Type[_F], key: str, resp: Any) -> _F: value = _read_value(resp.get("value")) failed = _read_typed(resp.get("failed"), "failed", bool) is True - reason = None - raw_reason = _read_typed(resp.get("reason"), "reason", dict) - if raw_reason is not None: - reason = FlagReason( - code=_read_typed(raw_reason.get("code"), "reason.code", str) or "", - condition_index=_read_int( - raw_reason.get("condition_index"), "reason.condition_index" - ), - description=_read_typed( - raw_reason.get("description"), "reason.description", str - ) - or "", - ) + reason = _parse_typed_reason(resp.get("reason")) if failed: # A failed record has no value, whatever it carries. diff --git a/references/public_api_snapshot.txt b/references/public_api_snapshot.txt index bfffd5d64..b19c54964 100644 --- a/references/public_api_snapshot.txt +++ b/references/public_api_snapshot.txt @@ -789,6 +789,7 @@ attribute posthog.capture_v1.CaptureV1Error.request_id = request_id attribute posthog.capture_v1.CaptureV1Error.retry_exhausted = retry_exhausted or [] attribute posthog.client.Client.api_key = (project_api_key or '').strip() attribute posthog.client.Client.app_host = app_host +attribute posthog.client.Client.before_send = before_send attribute posthog.client.Client.capture_compression = _resolve_capture_compression(capture_compression, gzip_fallback=gzip) attribute posthog.client.Client.capture_exception_code_variables = capture_exception_code_variables attribute posthog.client.Client.capture_mode = _resolve_capture_mode(capture_mode)