From 980f8fe9d95899ace53fc5345ee7724bdf76f32e Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:05:25 +0000 Subject: [PATCH 1/5] feat: Support public ingestion keys in `init_logger` `init_logger(ingestion_key=...)` and `BRAINTRUST_INGESTION_KEY` create a logger that writes traces to a single project with an ingestion key URL. The logger has its own queue and transport per endpoint and key, sends the key as a bearer token, and never logs in, discovers the project, or falls back to private credentials. An explicit `api_key` takes precedence over the env var, and passing both options is an error. Attachments and oversized batches go through the data plane's chunked `/v1/uploads` API before the rows that reference them are published. Rows only carry trace fields, and feedback is limited to scores, expected values, and tags, since comments and audit fields are writable provenance. Slug and W3C parents, plus global `update_span`, route through the current ingestion key logger so they don't trigger a private login. --- py/README.md | 28 ++ py/src/braintrust/env.py | 1 + py/src/braintrust/logger.py | 615 ++++++++++++++++++++---- py/src/braintrust/test_ingestion_key.py | 569 ++++++++++++++++++++++ 4 files changed, 1111 insertions(+), 102 deletions(-) create mode 100644 py/src/braintrust/test_ingestion_key.py diff --git a/py/README.md b/py/README.md index b0b3242dc..4007b660e 100644 --- a/py/README.md +++ b/py/README.md @@ -41,6 +41,34 @@ Then run: BRAINTRUST_API_KEY= braintrust eval tutorial_eval.py ``` +## Logging With an Ingestion Key + +Ingestion keys let apps that run on your users' machines, like a desktop app or a CLI, send traces to a single project without an API key. +Pass the ingestion key URL from your project settings to `init_logger`, or set `BRAINTRUST_INGESTION_KEY`: + +```python +import braintrust + +logger = braintrust.init_logger(ingestion_key="https://dp.example.com/ingest?ingestKey=bt-ik-...") +logger.log(input="question", output="answer") +``` + +The logger sends rows and attachments straight to the data plane in the URL, with the key in the `Authorization` header. +It never logs in, so it ignores `BRAINTRUST_API_KEY` and any earlier `braintrust.login()`, and it doesn't register or look up the project. +Passing an explicit `api_key` to `init_logger` uses that key instead of `BRAINTRUST_INGESTION_KEY`, and passing both `api_key` and `ingestion_key` is an error. +Feedback logged with an ingestion key can only include scores, expected values, and tags. + +To export OpenTelemetry spans with an ingestion key, point a standard OTLP exporter at the `/otel/v1/traces` path under the ingestion URL and drop the query string: + +```python +from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter + +exporter = OTLPSpanExporter( + endpoint="https://dp.example.com/ingest/otel/v1/traces", + headers={"Authorization": "Bearer bt-ik-..."}, +) +``` + ## Optional Extras Install extras as needed for specific workflows: diff --git a/py/src/braintrust/env.py b/py/src/braintrust/env.py index f08afb668..aff3b50f5 100644 --- a/py/src/braintrust/env.py +++ b/py/src/braintrust/env.py @@ -214,6 +214,7 @@ def __get__(self, instance: object, owner: type | None = None) -> bool: class BraintrustEnv: API_KEY = EnvVar("BRAINTRUST_API_KEY", EnvParser.STRING) + INGESTION_KEY = EnvVar("BRAINTRUST_INGESTION_KEY", EnvParser.STRING) API_URL = EnvVar("BRAINTRUST_API_URL", EnvParser.STRING) PROXY_URL = EnvVar("BRAINTRUST_PROXY_URL", EnvParser.STRING) HTTP_TIMEOUT = EnvVar("BRAINTRUST_HTTP_TIMEOUT", EnvParser.FLOAT) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 21155a2dd..e1964c845 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -14,6 +14,7 @@ import json import logging import os +import re import sys import textwrap import threading @@ -37,7 +38,7 @@ cast, overload, ) -from urllib.parse import quote, urlencode +from urllib.parse import parse_qsl, quote, urlencode, urlsplit, urlunsplit import chevron import exceptiongroup @@ -47,7 +48,7 @@ from . import context, id_gen from .api._routing import normalize_proxy_url -from .api._transport import HTTPConnection +from .api._transport import HTTPConnection, _parse_retry_after from .api._transport import RetryRequestExceptionsAdapter as RetryRequestExceptionsAdapter from .api.auth import LoginResult, OrganizationInfo from .api.client import BraintrustClient, BraintrustOpenApiClient @@ -57,12 +58,16 @@ BraintrustJSONDecodeError, BraintrustTransportError, ) -from .api.policies import DEFAULT_RETRYABLE_STATUSES, is_retryable_request_exception +from .api.policies import DEFAULT_MAX_ELAPSED_TIME, DEFAULT_RETRYABLE_STATUSES, is_retryable_request_exception from .bt_json import bt_dumps, bt_safe_deep_copy from .db_fields import ( + ARRAY_DELETE_FIELD, AUDIT_METADATA_FIELD, AUDIT_SOURCE_FIELD, + CREATED_FIELD, + ID_FIELD, IS_MERGE_FIELD, + MERGE_PATHS_FIELD, OBJECT_DELETE_FIELD, OBJECT_ID_KEYS, TRANSACTION_ID_FIELD, @@ -616,6 +621,12 @@ def default_get_api_conn(): lambda: _HTTPBackgroundLogger(LazyValue(default_get_api_conn, use_mutex=True)), use_mutex=True ) + # Ingestion key loggers publish independently of the private login, + # with one background logger per ingestion endpoint and key. + self._public_bg_loggers: dict[_IngestionEndpoint, _PublicHTTPBackgroundLogger] = {} + self._public_bg_loggers_lock = threading.Lock() + self._masking_function: Callable[[Any], Any] | None = None + self._id_generator = None # For unit-testing, tests may wish to temporarily override the global @@ -750,6 +761,9 @@ def copy_state(self, other: "BraintrustState"): "current_span", "_global_bg_logger", "_override_bg_logger", + "_public_bg_loggers", + "_public_bg_loggers_lock", + "_masking_function", "_context_manager", "_last_otel_setting", "_context_manager_lock", @@ -876,6 +890,15 @@ def loader_api_client(self, options: "_LoaderLoginOptions") -> "Iterator[Braintr def global_bg_logger(self) -> "_BackgroundLogger": return getattr(self._override_bg_logger, "logger", None) or self._global_bg_logger.get() + def public_bg_logger(self, endpoint: "_IngestionEndpoint") -> "_PublicHTTPBackgroundLogger": + with self._public_bg_loggers_lock: + bg_logger = self._public_bg_loggers.get(endpoint) + if bg_logger is None: + bg_logger = _PublicHTTPBackgroundLogger(endpoint) + bg_logger.set_masking_function(self._masking_function) + self._public_bg_loggers[endpoint] = bg_logger + return bg_logger + # Should only be called by the login function. def login_replace_api_conn(self, api_conn: "HTTPConnection"): self._global_bg_logger.get().internal_replace_api_conn(api_conn) @@ -891,8 +914,12 @@ def enforce_queue_size_limit(self, enforce: bool) -> None: bg_logger.enforce_queue_size_limit(enforce) def set_masking_function(self, masking_function: Callable[[Any], Any] | None) -> None: - """Set the masking function on the background logger.""" + """Set the masking function on the background loggers.""" self.global_bg_logger().set_masking_function(masking_function) + with self._public_bg_loggers_lock: + self._masking_function = masking_function + for bg_logger in self._public_bg_loggers.values(): + bg_logger.set_masking_function(masking_function) _state: BraintrustState = None # type: ignore @@ -922,6 +949,10 @@ def set_http_adapter(adapter: HTTPAdapter) -> None: _state._api_conn._reset() if _state._client: _state._client.transport._set_adapter(adapter) + with _state._public_bg_loggers_lock: + for bg_logger in _state._public_bg_loggers.values(): + bg_logger.api_conn.get()._set_adapter(adapter=adapter) + bg_logger.api_conn.get()._reset() # Per-credential loader resources may have been created with the previous # adapter. Eviction closes them once any active requests release their lease; @@ -1210,35 +1241,7 @@ def flush(self, batch_size: int | None = None): if len(all_items) == 0: return - # Construct batches of records to flush in parallel. - all_items_with_meta = [stringify_with_overflow_meta(item) for item in all_items] - max_request_size_result = self._get_max_request_size() - batches = batch_items( - items=all_items_with_meta, - batch_max_num_items=batch_size, - batch_max_num_bytes=max_request_size_result["max_request_size"] // 2, - get_byte_size=lambda item: len(item.str_value), - ) - - post_promises = [] - try: - post_promises = [ - HTTP_REQUEST_THREAD_POOL.submit(self._submit_logs_request, batch, max_request_size_result) - for batch in batches - ] - except RuntimeError: - # If the thread pool has shut down, e.g. because the process - # is terminating, run the requests the old fashioned way. - for batch in batches: - self._submit_logs_request(batch, max_request_size_result) - - concurrent.futures.wait(post_promises) - # Raise any exceptions from the promises as one group. - post_promise_exceptions = [e for e in (f.exception() for f in post_promises) if e is not None] - if post_promise_exceptions: - raise exceptiongroup.BaseExceptionGroup( - f"Encountered the following errors while logging:", post_promise_exceptions - ) + self._publish_items(all_items, batch_size) attachment_errors: list[Exception] = [] for attachment in attachments: @@ -1257,12 +1260,45 @@ def flush(self, batch_size: int | None = None): attachment_errors, ) + def _publish_items(self, all_items: list[dict[str, Any]], batch_size: int) -> None: + # Construct batches of records to flush in parallel. + all_items_with_meta = [stringify_with_overflow_meta(item) for item in all_items] + max_request_size_result = self._get_max_request_size() + batches = batch_items( + items=all_items_with_meta, + batch_max_num_items=batch_size, + batch_max_num_bytes=max_request_size_result["max_request_size"] // 2, + get_byte_size=lambda item: len(item.str_value), + ) + + post_promises = [] + try: + post_promises = [ + HTTP_REQUEST_THREAD_POOL.submit(self._submit_logs_request, batch, max_request_size_result) + for batch in batches + ] + except RuntimeError: + # If the thread pool has shut down, e.g. because the process + # is terminating, run the requests the old fashioned way. + for batch in batches: + self._submit_logs_request(batch, max_request_size_result) + + concurrent.futures.wait(post_promises) + # Raise any exceptions from the promises as one group. + post_promise_exceptions = [e for e in (f.exception() for f in post_promises) if e is not None] + if post_promise_exceptions: + raise exceptiongroup.BaseExceptionGroup( + f"Encountered the following errors while logging:", post_promise_exceptions + ) + def _unwrap_lazy_values( - self, wrapped_items: Sequence[_LazyLogRecord] + self, wrapped_items: Sequence[_LazyLogRecord], extract_attachments: bool = True ) -> tuple[list[dict[str, Any]], list["BaseAttachment"]]: for i in range(self.num_tries): try: - unwrapped_items = [record for item in wrapped_items if (record := item.get()) is not None] + unwrapped_items = [ + self._prepare_record(record) for item in wrapped_items if (record := item.get()) is not None + ] merged_items = merge_row_batch(unwrapped_items) # Logger-local hooks run after span hooks and merging. @@ -1274,8 +1310,9 @@ def _unwrap_lazy_values( ] attachments: list["BaseAttachment"] = [] - for item in merged_items: - _extract_attachments(item, attachments) + if extract_attachments: + for item in merged_items: + _extract_attachments(item, attachments) return merged_items, attachments except Exception as e: @@ -1300,6 +1337,9 @@ def _unwrap_lazy_values( ) return [], [] + def _prepare_record(self, record: dict[str, Any]) -> dict[str, Any]: + return record + def _request_logs3_overflow_upload( self, conn: HTTPConnection, payload_size_bytes: int, rows: list[dict[str, Any]] ) -> dict[str, Any]: @@ -1488,6 +1528,277 @@ def set_masking_function(self, masking_function: Callable[[Any], Any] | None): self._export_customizers = (_MaskingCustomizer(masking_function),) if masking_function is not None else () +INGESTION_KEY_QUERY_PARAM = "ingestKey" +_INGESTION_KEY_PATTERN = re.compile(r"bt-ik-[A-Za-z0-9]{48}") +# Ingestion keys may only write trace rows, so rows never carry other objects' ids or writable provenance. +_INGESTION_ROW_FIELDS = frozenset( + { + ID_FIELD, + "span_id", + "root_span_id", + "span_parents", + CREATED_FIELD, + "org_id", + "project_id", + "log_id", + "input", + "output", + "expected", + "error", + "tags", + "scores", + "metadata", + "metrics", + "context", + "span_attributes", + IS_MERGE_FIELD, + MERGE_PATHS_FIELD, + ARRAY_DELETE_FIELD, + OBJECT_DELETE_FIELD, + } +) +_UPLOAD_REFERENCE_TYPES = {"attachment": "braintrust_attachment", "logs3_overflow": LOGS3_OVERFLOW_REFERENCE_TYPE} + + +@dataclasses.dataclass(frozen=True) +class _IngestionEndpoint: + """An ingestion key URL split into the ingestion root and the key it carries.""" + + url: str + key: str = dataclasses.field(repr=False) + + +def _parse_ingestion_url(value: str) -> _IngestionEndpoint: + # Errors never include `value`, since it contains the key. + try: + parsed = urlsplit(value.strip()) + parsed.port # Raises for malformed ports. + query = parse_qsl(parsed.query, keep_blank_values=True, strict_parsing=True) + reason = None + except ValueError: + reason = "it is not a valid URL" + if reason is None: + path = parsed.path.rstrip("/") + if parsed.scheme not in ("http", "https") or not parsed.hostname: + reason = "expected an http or https URL" + elif parsed.username is not None or parsed.password is not None: + reason = "URL credentials are not allowed" + elif "#" in value: + reason = "URL fragments are not allowed" + elif not path.endswith("/ingest"): + reason = "expected the URL path to end with /ingest" + elif ( + len(query) != 1 + or query[0][0] != INGESTION_KEY_QUERY_PARAM + or not _INGESTION_KEY_PATTERN.fullmatch(query[0][1]) + ): + reason = f"expected a single {INGESTION_KEY_QUERY_PARAM} query parameter holding a bt-ik- key" + else: + return _IngestionEndpoint(url=urlunsplit((parsed.scheme, parsed.netloc, path, "", "")), key=query[0][1]) + raise ValueError(f"Invalid Braintrust ingestion key URL: {reason}") + + +class _UploadGrantExpired(Exception): + pass + + +class _PublicHTTPBackgroundLogger(_HTTPBackgroundLogger): + """Publishes rows to a data plane with an ingestion key. + + The key only allows trace ingestion into a single project, so this logger + never logs in, discovers projects, or falls back to private credentials. + Attachments and oversized batches go through the data plane's chunked + upload API before the rows that reference them are published. + """ + + def __init__(self, endpoint: _IngestionEndpoint): + conn = HTTPConnection(endpoint.url, adapter=_http_adapter) + conn.set_token(endpoint.key) + super().__init__(LazyValue(lambda: conn, use_mutex=False)) + self.endpoint = endpoint + self.timeout = BraintrustEnv.HTTP_TIMEOUT.get(60.0) + self._stripped_fields: set[str] = set() + self.enforce_queue_size_limit(True) + + def _get_max_request_size(self) -> dict[str, Any]: + # Any batch can overflow through the upload API, so there is no server limit to look up. + max_request_size = coalesce(self._max_request_size_override, DEFAULT_MAX_REQUEST_SIZE) + return {"max_request_size": max_request_size, "can_use_overflow": True} + + def flush(self, batch_size: int | None = None): + if batch_size is None: + batch_size = self.default_batch_size + + with self.flush_lock: + items, _ = self._unwrap_lazy_values(self.queue.drain_all(), extract_attachments=False) + if not items: + return + + # Attachment references are issued by the data plane, so attachments + # upload before the rows that point at them. + uploads: dict[int, AttachmentReference | Exception] = {} + rows: list[dict[str, Any]] = [] + attachment_errors: list[Exception] = [] + for item in items: + try: + _extract_attachments(item, [], lambda attachment: self._upload_attachment(attachment, uploads)) + except Exception as e: + attachment_errors.append(e) + print( + f"Dropping row {item.get(ID_FIELD)} because an attachment failed to upload: {e}", + file=self.outfile, + ) + continue + rows.append(item) + + if rows: + self._publish_items(rows, batch_size) + + if len(attachment_errors) == 1: + raise attachment_errors[0] + elif len(attachment_errors) > 1: + raise exceptiongroup.ExceptionGroup( + "Encountered errors while uploading attachments", + attachment_errors, + ) + + def _upload_attachment( + self, attachment: "BaseAttachment", uploads: "dict[int, AttachmentReference | Exception]" + ) -> AttachmentReference: + if isinstance(attachment, ExternalAttachment): + return attachment.reference + # The same attachment object can appear in several rows of one flush. + if id(attachment) not in uploads: + reference = attachment.reference + assert reference["type"] == "braintrust_attachment" + try: + uploads[id(attachment)] = self._upload( + "attachment", attachment.data, reference["content_type"], filename=reference["filename"] + ) + except Exception as e: + uploads[id(attachment)] = e + result = uploads[id(attachment)] + if isinstance(result, Exception): + raise result + return result + + def _prepare_record(self, record: dict[str, Any]) -> dict[str, Any]: + # Strip before merging, since merging groups rows by their object ids. + for field in [field for field in record if field not in _INGESTION_ROW_FIELDS]: + if field not in self._stripped_fields: + self._stripped_fields.add(field) + self.logger.warning(f"Ingestion keys cannot write the {field!r} field. Dropping it from logged rows.") + del record[field] + return record + + def _submit_logs_request(self, items: Sequence[LogItemWithMeta], max_request_size_result: dict[str, Any]): + data_str = construct_logs3_data(items) + payload = data_str.encode("utf-8") + if self.all_publish_payloads_dir: + _HTTPBackgroundLogger._write_payload_to_dir(payload_dir=self.all_publish_payloads_dir, payload=data_str) + + start_time = time.time() + try: + if len(payload) > max_request_size_result["max_request_size"]: + reference = self._upload("logs3_overflow", payload, "application/json") + self._request("post", "/v1/logs", json={"api_version": DATA_API_VERSION, "rows": reference}) + self._overflow_upload_count += 1 + else: + self._request("post", "/v1/logs", data=payload, headers={"Content-Type": "application/json"}) + except Exception as e: + if self.failed_publish_payloads_dir: + _HTTPBackgroundLogger._write_payload_to_dir( + payload_dir=self.failed_publish_payloads_dir, payload=data_str + ) + self._log_failed_payloads_dir() + errmsg = f"log request failed. Elapsed time: {time.time() - start_time} seconds. Payload size: {len(payload)}. Error: {e}" + if self.sync_flush: + raise Exception(errmsg) from e + print(errmsg, file=self.outfile) + print("Dropping batch", file=self.outfile) + + def _upload(self, purpose: str, data: bytes, content_type: str, filename: str | None = None) -> Any: + """Upload `data` in the chunks the data plane asks for and return the committed reference.""" + sha256 = hashlib.sha256(data).hexdigest() + body: dict[str, Any] = { + "purpose": purpose, + "size_bytes": len(data), + "content_type": content_type, + "sha256": sha256, + } + if filename is not None: + body["filename"] = filename + + for _ in range(self.num_tries): + # The grant's lifetime is relative, so measure it from before the request goes out. + requested_at = time.monotonic() + grant = self._request("post", "/v1/uploads", json=body) + upload_id = str(uuid.UUID(grant["upload_id"])) + chunk_bytes = grant["chunk_bytes"] + if chunk_bytes <= 0 or grant["num_chunks"] != -(-len(data) // chunk_bytes): + raise RuntimeError(f"Invalid upload grant from the data plane: {grant}") + deadline = requested_at + grant["expires_in_ms"] / 1000 + + try: + for index in range(grant["num_chunks"]): + chunk = data[index * chunk_bytes : (index + 1) * chunk_bytes] + ack = self._request( + "put", + f"/v1/uploads/{upload_id}/chunks/{index}", + deadline=deadline, + data=chunk, + headers={"Content-Type": "application/octet-stream"}, + ) + if ack.get("index") != index or ack.get("size_bytes") != len(chunk): + raise RuntimeError(f"Invalid chunk acknowledgement from the data plane: {ack}") + result = self._request("post", f"/v1/uploads/{upload_id}/complete", deadline=deadline, json={}) + except _UploadGrantExpired: + print(f"Upload {upload_id} expired before it completed. Starting a new upload", file=self.outfile) + continue + + reference = result.get("reference") + if ( + result.get("size_bytes") != len(data) + or result.get("sha256") != sha256 + or not isinstance(reference, dict) + or reference.get("type") != _UPLOAD_REFERENCE_TYPES[purpose] + ): + raise RuntimeError(f"Upload {upload_id} completed with unexpected content: {result}") + return reference + + raise RuntimeError(f"Uploads kept expiring before they completed after {self.num_tries} attempts") + + def _request(self, method: str, path: str, deadline: float | None = None, **kwargs: Any) -> Any: + """Send one request, retrying transient failures, and return its JSON body.""" + error = "" + for i in range(self.num_tries): + if deadline is not None and time.monotonic() >= deadline: + raise _UploadGrantExpired() + retry_after = None + try: + resp = getattr(self.api_conn.get(), method)(path, timeout=self.timeout, **kwargs) + except requests_exceptions.RequestException as e: + if not is_retryable_request_exception(e): + raise + error = str(e) + else: + if resp.ok: + return resp.json() + error = f"{resp.status_code}: {resp.text}".replace(self.endpoint.key, "") + if resp.status_code not in DEFAULT_RETRYABLE_STATUSES: + raise RuntimeError(f"{method.upper()} {path} failed with {error}") + retry_after = _parse_retry_after(resp.headers.get("Retry-After"), time.time()) + + if i + 1 < self.num_tries: + sleep_time_s = max( + BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * (2**i), min(retry_after or 0, DEFAULT_MAX_ELAPSED_TIME) + ) + print(f"{method.upper()} {path} failed with {error}. Retrying in {sleep_time_s}s", file=self.outfile) + time.sleep(sleep_time_s) + + raise RuntimeError(f"{method.upper()} {path} failed after {self.num_tries} attempts with {error}") + + def _internal_reset_global_state() -> None: global _state _state = BraintrustState() @@ -1927,6 +2238,7 @@ def init_logger( state: BraintrustState | None = None, environment: SpanOriginEnvironment | None = None, _create_in_project_group: str | None = None, + ingestion_key: str | None = None, ) -> "Logger": """ Create a new logger in a specified project. If the project does not exist, it will be created. @@ -1941,11 +2253,24 @@ def init_logger( :param force_login: Login again, even if you have already logged in (by default, the logger will not login if you are already logged in) :param set_current: If true (the default), set the global current-experiment to the newly-created one. :param _create_in_project_group: Experimental: the name of an existing project group to create the project in. Existing projects must already belong to the group, otherwise registration fails with a 409. Ignored when project_id is provided. + :param ingestion_key: (Optional) The ingestion key URL of a project, for example `https://api.braintrust.dev/ingest?ingestKey=bt-ik-...`. + If the parameter is not specified and `api_key` isn't either, will try to use the `BRAINTRUST_INGESTION_KEY` environment variable. + Ingestion keys can only write traces to the project they were issued for, so the logger never logs in, ignores `BRAINTRUST_API_KEY` + and any previous login, and doesn't register or look up the project. `project` and `project_id` are optional and only used to + label exported spans. A `project_id` is also sent with each row, and the data plane rejects rows for a different project. :returns: The newly created Logger. """ state = state or _state state.span_origin_environment = detect_environment(environment) + + if api_key is not None and ingestion_key is not None: + raise ValueError("Pass either api_key or ingestion_key to init_logger, not both") + if api_key is None: + ingestion_key = ingestion_key or BraintrustEnv.INGESTION_KEY.get(None, use_dotenv=True) + if ingestion_key and _create_in_project_group is not None: + raise ValueError("Loggers using an ingestion key cannot create projects") + compute_metadata_args = dict(project_name=project, project_id=project_id) if project_id is None and _create_in_project_group is not None: compute_metadata_args["project_group_name"] = _create_in_project_group @@ -1957,20 +2282,36 @@ def init_logger( "project_id": project_id, } - def compute_metadata(): - state.login(org_name=org_name, api_key=api_key, app_url=app_url, force_login=force_login) - return _compute_logger_metadata(**compute_metadata_args, state=state) + if ingestion_key: + # The data plane resolves the org and project from the key, so there is nothing to look up. + public_metadata = OrgProjectMetadata( + org_id=cast(str, None), + project=ObjectMetadata(id=cast(str, project_id), name=project or "", full_info={}), + ) + ret = Logger( + lazy_metadata=LazyValue(lambda: public_metadata, use_mutex=False), + async_flush=async_flush, + compute_metadata_args=compute_metadata_args, + link_args=link_args, + state=state, + public_bg_logger=state.public_bg_logger(_parse_ingestion_url(ingestion_key)), + ) + else: + + def compute_metadata(): + state.login(org_name=org_name, api_key=api_key, app_url=app_url, force_login=force_login) + return _compute_logger_metadata(**compute_metadata_args, state=state) - # For loggers, enable queue size limit enforcement (bounded queue) - state.enforce_queue_size_limit(True) + # For loggers, enable queue size limit enforcement (bounded queue) + state.enforce_queue_size_limit(True) - ret = Logger( - lazy_metadata=LazyValue(compute_metadata, use_mutex=True), - async_flush=async_flush, - compute_metadata_args=compute_metadata_args, - link_args=link_args, - state=state, - ) + ret = Logger( + lazy_metadata=LazyValue(compute_metadata, use_mutex=True), + async_flush=async_flush, + compute_metadata_args=compute_metadata_args, + link_args=link_args, + state=state, + ) if set_current: if _state is None: raise RuntimeError("_state is None in init_logger. This should never happen.") @@ -2568,6 +2909,15 @@ def current_logger() -> "Logger | None": return _state._cv_logger.get() or _state._local_logger +def _current_public_logger() -> "Logger | None": + """Return the current logger if it uses an ingestion key and no experiment takes precedence.""" + + logger = current_logger() + if logger is None or logger._public_bg_logger is None or current_experiment() is not None: + return None + return logger + + def current_span() -> Span: """Return the currently-active span for logging (set by running a span under a context manager). If there is no active span, returns a no-op span object, which supports the same interface as spans but does no logging. @@ -2851,18 +3201,26 @@ def _resolve_w3c_parent(context: dict) -> "tuple[str | None, PropagatedState | N braintrust_parent = parse_baggage(baggage_value).get(BRAINTRUST_PARENT_KEY) if not braintrust_parent: braintrust_parent = _current_braintrust_parent() - if not braintrust_parent: + public_logger = _current_public_logger() + if not braintrust_parent and public_logger is not None: + # Ingestion key spans always go to the key's project, so the trace only needs span ids. + object_type, object_id, compute_args = ( + SpanObjectTypeV3.PROJECT_LOGS, + None, + public_logger._compute_metadata_args, + ) + elif not braintrust_parent: logging.warning( "Received traceparent without a braintrust.parent and no active logger/experiment; " "cannot route the trace. Starting a fresh local span instead." ) return None, None - - parsed_parent = _braintrust_parent_to_components(braintrust_parent) - if parsed_parent is None: - logging.warning(f"Invalid braintrust.parent: {braintrust_parent!r}") - return None, None - object_type, object_id, compute_args = parsed_parent + else: + parsed_parent = _braintrust_parent_to_components(braintrust_parent) + if parsed_parent is None: + logging.warning(f"Invalid braintrust.parent: {braintrust_parent!r}") + return None, None + object_type, object_id, compute_args = parsed_parent tracestate = get_header(context, TRACESTATE_HEADER) or None @@ -3103,10 +3461,25 @@ def start_span( # No row id (object-only slug): there is no span to link to, so start # a fresh root, keeping the object routing from the slug. parent_span_ids = None + public_logger = _current_public_logger() + if public_logger is not None: + # Resolving the parent's object needs a private login. Ingestion key + # spans stay in the key's project anyway, so the parent only + # contributes its span ids. + object_args = dict( + parent_object_type=public_logger._parent_object_type(), + parent_object_id=public_logger._lazy_id, + parent_compute_object_metadata_args=public_logger._compute_metadata_args, + public_bg_logger=public_logger._public_bg_logger, + ) + else: + object_args = dict( + parent_object_type=parent_obj.object_type, + parent_object_id=LazyValue(_span_components_to_object_id_lambda(parent_obj), use_mutex=False), + parent_compute_object_metadata_args=parent_obj.compute_object_metadata_args, + ) return SpanImpl( - parent_object_type=parent_obj.object_type, - parent_object_id=LazyValue(_span_components_to_object_id_lambda(parent_obj), use_mutex=False), - parent_compute_object_metadata_args=parent_obj.compute_object_metadata_args, + **object_args, parent_span_ids=parent_span_ids, name=name, type=type, @@ -3138,6 +3511,10 @@ def flush(): """Flush any pending rows to the server.""" _state.global_bg_logger().flush() + with _state._public_bg_loggers_lock: + public_bg_loggers = list(_state._public_bg_loggers.values()) + for public_bg_logger in public_bg_loggers: + public_bg_logger.flush() def _internal_start_span_with_initial_merge( @@ -3219,7 +3596,11 @@ def validate_tags(tags: Sequence[str]) -> None: seen.add(tag) -def _extract_attachments(event: dict[str, Any], attachments: list["BaseAttachment"]) -> None: +def _extract_attachments( + event: dict[str, Any], + attachments: list["BaseAttachment"], + get_reference: "Callable[[BaseAttachment], AttachmentReference] | None" = None, +) -> None: """ Helper function for uploading attachments. Recursively extracts `Attachment` and `ExternalAttachment` values and replaces them with their associated @@ -3227,13 +3608,15 @@ def _extract_attachments(event: dict[str, Any], attachments: list["BaseAttachmen :param event: The event to filter. Will be modified in-place. :param attachments: Flat array of extracted attachments (output parameter). + :param get_reference: Returns the reference that replaces an attachment. Defaults to `attachment.reference`. """ def _helper(v: Any) -> Any: # Base case: Attachment or ExternalAttachment. if isinstance(v, BaseAttachment): attachments.append(v) - return v.reference # Attachment cannot be nested. + # Attachment cannot be nested. + return get_reference(v) if get_reference else v.reference # Recursive case: object. if isinstance(v, dict): @@ -3909,6 +4292,13 @@ def __str__(self) -> str: return f"data:{self.reference['content_type']};base64,{b64_content}" +def _object_id_fields(object_type: SpanObjectTypeV3, object_id: str | None) -> dict[str, str]: + # Ingestion key loggers may not know their project. The data plane resolves it from the key. + if object_id is None and object_type == SpanObjectTypeV3.PROJECT_LOGS: + return {"log_id": "g"} + return _get_exporter()(object_type=object_type, object_id=object_id).object_id_fields() + + def _log_feedback_impl( parent_object_type: SpanObjectTypeV3, parent_object_id: LazyValue[str], @@ -3919,7 +4309,12 @@ def _log_feedback_impl( comment: str | None = None, metadata: Metadata | None = None, source: Literal["external", "app", "api", None] = None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, ): + if public_bg_logger is not None and (comment is not None or metadata is not None or source is not None): + # Comments and audit fields are writable provenance, which ingestion keys cannot write. + raise ValueError("Feedback logged with an ingestion key can only include scores, expected, and tags") + if source is None: source = "external" elif source not in VALID_SOURCES: @@ -3943,13 +4338,11 @@ def _log_feedback_impl( update_event = {k: v for k, v in update_event.items() if v is not None} update_event = bt_safe_deep_copy(update_event) + audit_fields = {} if public_bg_logger else {AUDIT_SOURCE_FIELD: source, AUDIT_METADATA_FIELD: metadata} + bg_logger = public_bg_logger or _state.global_bg_logger() def parent_ids(): - exporter = _get_exporter() - return exporter( - object_type=parent_object_type, - object_id=parent_object_id.get(), - ).object_id_fields() + return _object_id_fields(parent_object_type, parent_object_id.get()) if len(update_event) > 0: @@ -3958,14 +4351,11 @@ def compute_update_record(): id=id, **update_event, **parent_ids(), - **{ - AUDIT_SOURCE_FIELD: source, - AUDIT_METADATA_FIELD: metadata, - IS_MERGE_FIELD: True, - }, + **audit_fields, + **{IS_MERGE_FIELD: True}, ) - _state.global_bg_logger().log(LazyValue(compute_update_record, use_mutex=False)) + bg_logger.log(LazyValue(compute_update_record, use_mutex=False)) if comment is not None: # pylint: disable=function-redefined @@ -3985,7 +4375,7 @@ def compute_comment_record(): **{AUDIT_SOURCE_FIELD: source, AUDIT_METADATA_FIELD: metadata}, ) - _state.global_bg_logger().log(LazyValue(compute_comment_record, use_mutex=False)) + bg_logger.log(LazyValue(compute_comment_record, use_mutex=False)) def _update_span_impl( @@ -3994,6 +4384,7 @@ def _update_span_impl( id: str, root_span_id: str | None, span_id: str | None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, **event: Any, ): if (root_span_id is None) != (span_id is None): @@ -4011,25 +4402,18 @@ def _update_span_impl( update_event = bt_safe_deep_copy(update_event) customizers = _get_span_customizers() - def parent_ids(): - exporter = _get_exporter() - return exporter( - object_type=parent_object_type, - object_id=parent_object_id.get(), - ).object_id_fields() - def compute_record() -> dict[str, Any] | None: record = dict( id=id, **update_event, - **parent_ids(), + **_object_id_fields(parent_object_type, parent_object_id.get()), **{ IS_MERGE_FIELD: True, }, ) return _customize_span_export(record, customizers) if customizers else record - _state.global_bg_logger().log(LazyValue(compute_record, use_mutex=False)) + (public_bg_logger or _state.global_bg_logger()).log(LazyValue(compute_record, use_mutex=False)) def update_span(exported: str, **event: Any) -> None: @@ -4054,6 +4438,17 @@ def update_span(exported: str, **event: Any) -> None: event_without_span_ids.pop("span_id", None) event_without_span_ids.pop("root_span_id", None) + public_logger = _current_public_logger() + if public_logger is not None: + # Resolving the slug's object needs a private login, and ingestion key + # rows stay in the key's project anyway. + return public_logger.update_span( + id=components.row_id, + root_span_id=components.root_span_id, + span_id=components.span_id, + **event_without_span_ids, + ) + return _update_span_impl( parent_object_type=components.object_type, parent_object_id=LazyValue(_span_components_to_object_id_lambda(components), use_mutex=False), @@ -4247,6 +4642,7 @@ def _start_span_parent_args( parent_span_ids: ParentSpanIds | None, propagated_event: dict[str, Any] | None, propagated_state: "PropagatedState | None" = None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, ) -> dict[str, Any]: # `parent` may be an exported slug string or an opaque W3C trace-context dict. parent_slug, parent_propagated_state = _normalize_parent(parent) @@ -4257,16 +4653,22 @@ def _start_span_parent_args( f"Mismatch between expected span parent object type {parent_object_type} and provided type {parent_components.object_type}" ) - parent_components_object_id_lambda = _span_components_to_object_id_lambda(parent_components) + if public_bg_logger is not None: + # Resolving the parent's object needs a private login. Ingestion key + # spans stay in the key's project anyway, so the parent only + # contributes its span ids. + arg_parent_object_id = parent_object_id + else: + parent_components_object_id_lambda = _span_components_to_object_id_lambda(parent_components) - def compute_parent_object_id(): - parent_components_object_id = parent_components_object_id_lambda() - assert parent_object_id.get() == parent_components_object_id, ( - f"Mismatch between expected span parent object id {parent_object_id.get()} and provided id {parent_components_object_id}" - ) - return parent_object_id.get() + def compute_parent_object_id(): + parent_components_object_id = parent_components_object_id_lambda() + assert parent_object_id.get() == parent_components_object_id, ( + f"Mismatch between expected span parent object id {parent_object_id.get()} and provided id {parent_components_object_id}" + ) + return parent_object_id.get() - arg_parent_object_id = LazyValue(compute_parent_object_id, use_mutex=False) + arg_parent_object_id = LazyValue(compute_parent_object_id, use_mutex=False) if parent_components.row_id and _parent_span_ids_usable( parent_components.span_id, parent_components.root_span_id ): @@ -4292,6 +4694,7 @@ def compute_parent_object_id(): parent_span_ids=arg_parent_span_ids, propagated_event=arg_propagated_event, propagated_state=arg_propagated_state, + public_bg_logger=public_bg_logger, ) @@ -4777,6 +5180,7 @@ def __init__( state: BraintrustState | None = None, lookup_span_parent: bool = True, internal: SpanInternalOptions | None = None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, ): if span_attributes is None: span_attributes = SpanAttributes() @@ -4797,6 +5201,7 @@ def __init__( self.parent_object_type = parent_object_type self.parent_object_id = parent_object_id self.parent_compute_object_metadata_args = parent_compute_object_metadata_args + self.public_bg_logger = public_bg_logger # Merge propagated_event into event. The propagated_event data will get # propagated-and-merged into every subspan. @@ -4967,14 +5372,10 @@ def write_span_cache(record: dict[str, Any]) -> None: write_span_cache(serializable_partial_record) def compute_record() -> dict[str, Any] | None: - exporter = _get_exporter() record = dict( **serializable_partial_record, **{k: v.get() for k, v in lazy_partial_record.items()}, - **exporter( - object_type=self.parent_object_type, - object_id=self.parent_object_id.get(), - ).object_id_fields(), + **_object_id_fields(self.parent_object_type, self.parent_object_id.get()), ) # Resolve and customize inside the cached LazyValue: every incremental # span record is transformed once, before the background @@ -4993,13 +5394,14 @@ def compute_record() -> dict[str, Any] | None: lazy_record = LazyValue(compute_record, use_mutex=pending_cache_key is not None) if pending_cache_key is not None: self.state.span_cache._track_pending_record(self.root_span_id, pending_cache_key, lazy_record) - self.state.global_bg_logger().log(lazy_record) + (self.public_bg_logger or self.state.global_bg_logger()).log(lazy_record) def log_feedback(self, **event: Any) -> None: return _log_feedback_impl( parent_object_type=self.parent_object_type, parent_object_id=self.parent_object_id, id=self.id, + public_bg_logger=self.public_bg_logger, **event, ) @@ -5034,6 +5436,7 @@ def start_span( parent_span_ids=parent_span_ids, propagated_event=coalesce(propagated_event, self.propagated_event), propagated_state=self.propagated_state, + public_bg_logger=self.public_bg_logger, ), name=name, type=type, @@ -5059,10 +5462,10 @@ def end(self, end_time: float | None = None) -> float: def export(self) -> str: if self.parent_compute_object_metadata_args and not self.parent_object_id.has_succeeded: object_id = None - compute_object_metadata_args = self.parent_compute_object_metadata_args else: object_id = self.parent_object_id.get() - compute_object_metadata_args = None + # Ingestion key loggers may not know their project id. + compute_object_metadata_args = None if object_id else self.parent_compute_object_metadata_args # Choose SpanComponents version based on the active ID format (hex -> V4, # legacy UUID -> V3). Coupled via _get_exporter() so the two never desync. @@ -5142,6 +5545,9 @@ def link(self) -> str: return NOOP_SPAN_PERMALINK def permalink(self) -> str: + if self.public_bg_logger is not None: + # Permalinks resolve the org through a private login. + return self.link() try: return permalink(self.export()) except Exception as e: @@ -5156,7 +5562,7 @@ def close(self, end_time=None) -> float: def flush(self) -> None: """Flush any pending rows to the server.""" - self.state.global_bg_logger().flush() + (self.public_bg_logger or self.state.global_bg_logger()).flush() def set_current(self): if self.can_set_current: @@ -5942,8 +6348,10 @@ def __init__( compute_metadata_args: dict | None = None, link_args: dict | None = None, state: BraintrustState | None = None, + public_bg_logger: "_PublicHTTPBackgroundLogger | None" = None, ): self._lazy_metadata = lazy_metadata + self._public_bg_logger = public_bg_logger self.async_flush = async_flush self._compute_metadata_args = compute_metadata_args self.last_start_time = time.time() @@ -6186,6 +6594,7 @@ def log_feedback( comment=comment, metadata=metadata, source=source, + public_bg_logger=self._public_bg_logger, ) def start_span( @@ -6237,6 +6646,7 @@ def update_span(self, id: str, **event: Any) -> None: id=id, root_span_id=root_span_id, span_id=span_id, + public_bg_logger=self._public_bg_logger, **event, ) @@ -6262,6 +6672,7 @@ def _start_span_impl( parent_compute_object_metadata_args=self._compute_metadata_args, parent_span_ids=None, propagated_event=propagated_event, + public_bg_logger=self._public_bg_logger, ) return SpanImpl( **parent_args, @@ -6287,10 +6698,10 @@ def export(self) -> str: # this `_lazy_id` object specifically will also be marked as computed. if self._compute_metadata_args and not self._lazy_id.has_succeeded: object_id = None - compute_object_metadata_args = self._compute_metadata_args else: object_id = self._lazy_id.get() - compute_object_metadata_args = None + # Ingestion key loggers may not know their project id. + compute_object_metadata_args = None if object_id else self._compute_metadata_args exporter = _get_exporter() return exporter( @@ -6322,7 +6733,7 @@ def flush(self) -> None: """ Flush any pending logs to the server. """ - self.state.global_bg_logger().flush() + (self._public_bg_logger or self.state.global_bg_logger()).flush() @dataclasses.dataclass diff --git a/py/src/braintrust/test_ingestion_key.py b/py/src/braintrust/test_ingestion_key.py new file mode 100644 index 000000000..2c6eccc29 --- /dev/null +++ b/py/src/braintrust/test_ingestion_key.py @@ -0,0 +1,569 @@ +"""Tests for loggers that publish with a public ingestion key. + +The fake data plane below implements the ingestion key contract over real HTTP, +so these tests exercise the actual logger, batcher, and upload flow. +""" + +import hashlib +import http.server +import json +import re +import threading +import uuid +from contextlib import contextmanager + +import braintrust +import pytest +from braintrust import logger +from braintrust.logger import Attachment, ExternalAttachment, _parse_ingestion_url +from braintrust.test_helpers import memory_logger, simulate_login # noqa: F401 + + +KEY = "bt-ik-" + "a" * 48 +OTHER_KEY = "bt-ik-" + "b" * 48 +BASE_PATH = "/deployment/base/ingest" +UPLOAD_PATH = re.compile(rf"^{BASE_PATH}/v1/uploads/([0-9a-f-]+)/(chunks/(\d+)|complete)$") + + +class FakeDataPlane: + """Records requests and answers them like a data plane serving one ingestion key.""" + + def __init__(self, key: str = KEY, chunk_bytes: int = 4): + self.key = key + self.chunk_bytes = chunk_bytes + self.requests: list[dict] = [] + self.rows: list[dict] = [] + self.uploads: dict[str, dict] = {} + # Responses to return before handling requests normally, keyed by (method, route). + self.scripted: dict[tuple[str, str], list[tuple[int, dict, bytes]]] = {} + self.grant_lifetimes_ms: list[int] = [] + self.project_id: str | None = None + self.lock = threading.Lock() + + def script(self, method: str, route: str, *responses: tuple[int, dict, bytes]) -> None: + self.scripted.setdefault((method, route), []).extend(responses) + + def paths(self) -> list[tuple[str, str]]: + return [(r["method"], r["path"]) for r in self.requests] + + def handle(self, method: str, path: str, headers, body: bytes) -> tuple[int, dict, bytes]: + with self.lock: + self.requests.append({"method": method, "path": path, "headers": dict(headers), "body": body}) + route = self._route(path) + scripted = self.scripted.get((method, route)) + if scripted: + return scripted.pop(0) + if headers.get("Authorization") != f"Bearer {self.key}": + return 401, {}, b'{"error": "invalid ingestion key"}' + if route == "logs" and method == "POST": + return self._logs(json.loads(body)) + if route == "create" and method == "POST": + return self._create(json.loads(body)) + if route == "chunk" and method == "PUT": + return self._chunk(path, headers, body) + if route == "complete" and method == "POST": + return self._complete(path) + return 404, {}, b'{"error": "not found"}' + + @staticmethod + def _route(path: str) -> str: + if path == f"{BASE_PATH}/v1/logs": + return "logs" + if path == f"{BASE_PATH}/v1/uploads": + return "create" + match = UPLOAD_PATH.match(path) + if match: + return "chunk" if match.group(3) is not None else "complete" + return "unknown" + + def _logs(self, body: dict) -> tuple[int, dict, bytes]: + assert body["api_version"] == 2 + rows = body["rows"] + if isinstance(rows, dict): + assert rows == {"type": "logs3_overflow", "key": rows["key"]} + upload = self.uploads[rows["key"]] + assert upload["purpose"] == "logs3_overflow" and upload["completed"] + rows = json.loads(upload["data"])["rows"] + for row in rows: + if self.project_id is not None and row.get("project_id", self.project_id) != self.project_id: + return 403, {}, b'{"error": "row project does not match the ingestion key"}' + self.rows.extend(rows) + return 200, {}, json.dumps({"ids": [row["id"] for row in rows]}).encode() + + def _create(self, body: dict) -> tuple[int, dict, bytes]: + upload_id = str(uuid.uuid4()) + num_chunks = -(-body["size_bytes"] // self.chunk_bytes) + self.uploads[upload_id] = {**body, "chunks": {}, "completed": False} + lifetime = self.grant_lifetimes_ms.pop(0) if self.grant_lifetimes_ms else 300000 + response = { + "upload_id": upload_id, + "chunk_bytes": self.chunk_bytes, + "num_chunks": num_chunks, + "expires_in_ms": lifetime, + } + return 201, {}, json.dumps(response).encode() + + def _chunk(self, path: str, headers, body: bytes) -> tuple[int, dict, bytes]: + match = UPLOAD_PATH.match(path) + assert match is not None + upload = self.uploads[match.group(1)] + index = int(match.group(3)) + assert headers.get("Content-Type") == "application/octet-stream" + assert "Content-Encoding" not in headers + upload["chunks"][index] = body + return 200, {}, json.dumps({"index": index, "size_bytes": len(body)}).encode() + + def _complete(self, path: str) -> tuple[int, dict, bytes]: + match = UPLOAD_PATH.match(path) + assert match is not None + upload_id = match.group(1) + upload = self.uploads[upload_id] + data = b"".join(upload["chunks"][i] for i in sorted(upload["chunks"])) + assert len(data) == upload["size_bytes"] + assert hashlib.sha256(data).hexdigest() == upload["sha256"] + upload["data"] = data + upload["completed"] = True + if upload["purpose"] == "attachment": + reference = { + "type": "braintrust_attachment", + "key": upload_id, + "filename": upload["filename"], + "content_type": upload["content_type"], + } + else: + reference = {"type": "logs3_overflow", "key": upload_id} + response = {"reference": reference, "size_bytes": len(data), "sha256": upload["sha256"]} + return 200, {}, json.dumps(response).encode() + + +@contextmanager +def serve(data_plane: FakeDataPlane): + class Handler(http.server.BaseHTTPRequestHandler): + def log_message(self, format, *args): + pass + + def _handle(self): + body = self.rfile.read(int(self.headers.get("Content-Length", "0"))) + status, headers, response_body = data_plane.handle(self.command, self.path, self.headers, body) + self.send_response(status) + for name, value in headers.items(): + self.send_header(name, value) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(response_body))) + self.end_headers() + self.wfile.write(response_body) + + do_GET = do_POST = do_PUT = do_PATCH = do_DELETE = _handle + + server = http.server.ThreadingHTTPServer(("127.0.0.1", 0), Handler) + thread = threading.Thread(target=server.serve_forever, kwargs={"poll_interval": 0.01}, daemon=True) + thread.start() + try: + yield f"http://127.0.0.1:{server.server_address[1]}{BASE_PATH}" + finally: + server.shutdown() + server.server_close() + + +def ingestion_url(root: str, key: str = KEY) -> str: + return f"{root}?ingestKey={key}" + + +@pytest.fixture(autouse=True) +def deterministic_flush(monkeypatch): + # Rows only publish on explicit flushes, failures raise, and retries don't sleep. + monkeypatch.setenv("BRAINTRUST_SYNC_FLUSH", "true") + monkeypatch.setenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH", "true") + monkeypatch.delenv("BRAINTRUST_INGESTION_KEY", raising=False) + monkeypatch.setattr(logger, "BACKGROUND_LOGGER_BASE_SLEEP_TIME_S", 0) + logger._state.reset_parent_state() + yield + logger._state.reset_parent_state() + + +@pytest.fixture +def data_plane(): + plane = FakeDataPlane() + with serve(plane) as root: + plane.root = root + yield plane + + +@pytest.fixture +def no_private_login(monkeypatch): + def fail(*args, **kwargs): + raise AssertionError("ingestion key loggers must not log in") + + monkeypatch.setattr(logger, "login_to_state", fail) + monkeypatch.setattr(logger, "_compute_logger_metadata", fail) + + +def flush_error(target) -> str: + """Flush and return the error, which publishing raises inside an exception group.""" + with pytest.raises(BaseException) as exc_info: + target.flush() + return "; ".join(str(e) for e in exc_info.value.exceptions) + + +def assert_no_key_in_urls(plane: FakeDataPlane) -> None: + for request in plane.requests: + assert "ingestKey" not in request["path"] + assert plane.key not in request["path"] + + +@pytest.mark.parametrize( + "url,expected", + [ + (f"https://dp.example/ingest?ingestKey={KEY}", "https://dp.example/ingest"), + (f"https://dp.example/deployment/base/ingest/?ingestKey={KEY}", "https://dp.example/deployment/base/ingest"), + (f"http://localhost:8000/ingest?ingestKey={KEY}\n", "http://localhost:8000/ingest"), + ], +) +def test_parse_ingestion_url(url, expected): + endpoint = _parse_ingestion_url(url) + + assert endpoint.url == expected + assert endpoint.key == KEY + assert KEY not in repr(endpoint) + + +@pytest.mark.parametrize( + "url", + [ + KEY, + f"ftp://dp.example/ingest?ingestKey={KEY}", + f"https://user:pass@dp.example/ingest?ingestKey={KEY}", + f"https://dp.example/ingest?ingestKey={KEY}#fragment", + f"https://dp.example/ingest/v1/logs?ingestKey={KEY}", + f"https://dp.example/not-ingest?ingestKey={KEY}", + f"https://dp.example/ingest?ingestKey={KEY}&ingestKey={KEY}", + f"https://dp.example/ingest?ingestKey={KEY}&project=foo", + f"https://dp.example/ingest?key={KEY}", + f"https://dp.example/ingest?ingestKey={KEY}x", + "https://dp.example/ingest?ingestKey=sk-" + "a" * 48, + f"https://dp.example:notaport/ingest?ingestKey={KEY}", + ], +) +def test_parse_ingestion_url_rejects_invalid_urls_without_leaking_the_key(url): + with pytest.raises(ValueError, match="Invalid Braintrust ingestion key URL") as exc_info: + _parse_ingestion_url(url) + + assert "a" * 48 not in str(exc_info.value) + assert exc_info.value.__cause__ is None and exc_info.value.__context__ is None + + +def test_first_flush_only_posts_rows_with_the_key(data_plane, monkeypatch, no_private_login): + # Ambient private credentials must be ignored. + monkeypatch.setenv("BRAINTRUST_API_KEY", "sk-private") + + public_logger = braintrust.init_logger(project="my-project", ingestion_key=ingestion_url(data_plane.root)) + row_id = public_logger.log(input="question", output="answer", scores={"good": 1}, metadata={"a": 1}) + public_logger.flush() + + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] + request = data_plane.requests[0] + assert request["headers"]["Authorization"] == f"Bearer {KEY}" + assert request["headers"]["Content-Type"] == "application/json" + [row] = data_plane.rows + assert row["id"] == row_id + assert row["log_id"] == "g" + assert "project_id" not in row + assert { + "input": "question", + "output": "answer", + "scores": {"good": 1}, + "metadata": {"a": 1}, + }.items() <= row.items() + assert set(row) <= set(logger._INGESTION_ROW_FIELDS) + assert public_logger.id is None + assert not logger._state.logged_in + assert_no_key_in_urls(data_plane) + + +def test_trace_context_round_trips_without_project_metadata(data_plane, no_private_login): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + + with public_logger.start_span(name="root") as root: + exported = root.export() + headers = root.inject() + assert root.permalink() + # A parent exported elsewhere only contributes its span ids. + with braintrust.start_span(name="from-slug", parent=exported) as from_slug: + pass + with braintrust.start_span(name="from-headers", parent=braintrust.extract_trace_context(headers)) as from_headers: + pass + public_logger.flush() + + rows = {row["span_attributes"]["name"]: row for row in data_plane.rows} + assert rows["from-slug"]["root_span_id"] == root.root_span_id + assert rows["from-slug"]["span_parents"] == [root.span_id] + assert rows["from-headers"]["root_span_id"] == root.root_span_id + assert rows["from-headers"]["span_parents"] == [root.span_id] + assert all("project_id" not in row for row in data_plane.rows) + assert from_slug.public_bg_logger is public_logger._public_bg_logger + assert from_headers.public_bg_logger is public_logger._public_bg_logger + + +def test_explicit_project_id_is_sent_and_mismatches_are_not_retried(data_plane, no_private_login): + data_plane.project_id = "key-project" + public_logger = braintrust.init_logger(project_id="other-project", ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + + assert "403" in flush_error(public_logger) + + [request] = data_plane.requests + [row] = json.loads(request["body"])["rows"] + assert row["project_id"] == "other-project" + assert data_plane.rows == [] + + +def test_retries_transient_failures_with_the_same_key(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_NUM_RETRIES", "2") + data_plane.script( + "POST", + "logs", + (429, {"Retry-After": "0"}, b'{"error": "rate limited"}'), + (503, {}, b'{"error": "unavailable"}'), + ) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + public_logger.flush() + + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] * 3 + assert {r["headers"]["Authorization"] for r in data_plane.requests} == {f"Bearer {KEY}"} + assert len(data_plane.rows) == 1 + + +def test_rejected_writes_are_not_retried_or_sent_elsewhere(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_NUM_RETRIES", "2") + data_plane.script("POST", "logs", (401, {}, f'{{"error": "bad key {KEY}"}}'.encode())) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + + error = flush_error(public_logger) + + assert "401" in error + assert KEY not in error + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] + + +def test_attachments_upload_in_advertised_chunks_before_rows(data_plane, no_private_login): + data_plane.chunk_bytes = 4 + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + attachment = Attachment(data=b"hello world", filename="greeting.txt", content_type="text/plain") + empty = Attachment(data=b"", filename="empty.txt", content_type="text/plain") + external = ExternalAttachment(url="s3://bucket/file.txt", filename="file.txt", content_type="text/plain") + public_logger.log(input={"file": attachment, "empty": empty, "external": external, "again": attachment}) + public_logger.flush() + + [attachment_id, empty_id] = list(data_plane.uploads) + assert data_plane.paths() == [ + ("POST", f"{BASE_PATH}/v1/uploads"), + ("PUT", f"{BASE_PATH}/v1/uploads/{attachment_id}/chunks/0"), + ("PUT", f"{BASE_PATH}/v1/uploads/{attachment_id}/chunks/1"), + ("PUT", f"{BASE_PATH}/v1/uploads/{attachment_id}/chunks/2"), + ("POST", f"{BASE_PATH}/v1/uploads/{attachment_id}/complete"), + ("POST", f"{BASE_PATH}/v1/uploads"), + ("POST", f"{BASE_PATH}/v1/uploads/{empty_id}/complete"), + ("POST", f"{BASE_PATH}/v1/logs"), + ] + assert [len(r["body"]) for r in data_plane.requests if r["method"] == "PUT"] == [4, 4, 3] + assert json.loads(data_plane.requests[0]["body"]) == { + "purpose": "attachment", + "size_bytes": 11, + "content_type": "text/plain", + "sha256": hashlib.sha256(b"hello world").hexdigest(), + "filename": "greeting.txt", + } + assert json.loads(data_plane.requests[4]["body"]) == {} + assert data_plane.uploads[empty_id]["size_bytes"] == 0 + [row] = data_plane.rows + file_reference = { + "type": "braintrust_attachment", + "key": attachment_id, + "filename": "greeting.txt", + "content_type": "text/plain", + } + assert row["input"] == { + "file": file_reference, + "empty": { + "type": "braintrust_attachment", + "key": empty_id, + "filename": "empty.txt", + "content_type": "text/plain", + }, + "external": external.reference, + "again": file_reference, + } + + +def test_rows_with_failed_attachments_are_dropped(data_plane, no_private_login): + data_plane.script("POST", "create", (400, {}, b'{"error": "uploads disabled"}')) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abc", filename="a.txt", content_type="text/plain")) + kept_id = public_logger.log(input="no attachment") + + with pytest.raises(RuntimeError, match="400"): + public_logger.flush() + + assert [row["id"] for row in data_plane.rows] == [kept_id] + assert not any("attachment" in path for _, path in data_plane.paths()) + + +def test_expired_upload_grant_starts_a_new_upload(data_plane, no_private_login): + data_plane.grant_lifetimes_ms = [1] + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + public_logger.flush() + + [expired_id, upload_id] = list(data_plane.uploads) + assert not data_plane.uploads[expired_id]["chunks"] + assert data_plane.uploads[upload_id]["completed"] + assert data_plane.rows[0]["input"]["key"] == upload_id + + +def test_oversized_batches_overflow_through_uploads(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_MAX_REQUEST_SIZE", "2000") + data_plane.chunk_bytes = 1024 + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + row_id = public_logger.log(input="x" * 5000) + public_logger.flush() + + [upload_id] = list(data_plane.uploads) + upload = data_plane.uploads[upload_id] + assert upload["purpose"] == "logs3_overflow" + assert upload["content_type"] == "application/json" + overflow_payload = json.loads(upload["data"]) + assert set(overflow_payload) == {"api_version", "rows"} + assert overflow_payload["api_version"] == 2 + assert [row["id"] for row in overflow_payload["rows"]] == [row_id] + [_, *chunk_requests, _, logs_request] = data_plane.requests + assert [len(r["body"]) for r in chunk_requests] == [1024] * (len(chunk_requests) - 1) + [ + len(upload["data"]) - 1024 * (len(chunk_requests) - 1) + ] + assert json.loads(logs_request["body"]) == { + "api_version": 2, + "rows": {"type": "logs3_overflow", "key": upload_id}, + } + assert [row["id"] for row in data_plane.rows] == [row_id] + + +def test_feedback_and_disallowed_fields(data_plane, no_private_login): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + with public_logger.start_span(name="root") as span: + span.log(input="x", experiment_id="someone-elses-experiment", dataset_id="d") + public_logger.log_feedback(id=span.id, scores={"thumbs_up": 1}, tags=["a"]) + braintrust.update_span(span.export(), output="updated") + with pytest.raises(ValueError, match="ingestion key"): + public_logger.log_feedback(id=span.id, comment="nice") + with pytest.raises(ValueError, match="ingestion key"): + span.log_feedback(scores={"thumbs_up": 0}, metadata={"user": "u"}) + public_logger.flush() + + [row] = data_plane.rows + assert row["id"] == span.id + assert row["scores"] == {"thumbs_up": 1} + assert row["output"] == "updated" + assert set(row) <= set(logger._INGESTION_ROW_FIELDS) + + +def test_masking_and_global_flush_cover_public_loggers(data_plane, no_private_login): + braintrust.set_masking_function(lambda value: "masked" if value == "secret" else value) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="secret") + braintrust.flush() + + assert data_plane.rows[0]["input"] == "masked" + + +def test_env_ingestion_key_takes_precedence_over_ambient_private_credentials( + data_plane, monkeypatch, no_private_login +): + monkeypatch.setenv("BRAINTRUST_API_KEY", "sk-private") + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(data_plane.root)) + + public_logger = braintrust.init_logger(project="p") + public_logger.log(input="x") + public_logger.flush() + + assert public_logger._public_bg_logger is not None + assert len(data_plane.rows) == 1 + + +def test_env_ingestion_key_ignores_previous_login(data_plane, monkeypatch, memory_logger): + simulate_login() + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(data_plane.root)) + + public_logger = braintrust.init_logger(project="p") + public_logger.log(input="x") + public_logger.flush() + + assert len(data_plane.rows) == 1 + assert memory_logger.pop() == [] + + +def test_explicit_ingestion_key_overrides_env(data_plane, monkeypatch, no_private_login): + other_plane = FakeDataPlane(key=OTHER_KEY) + with serve(other_plane) as other_root: + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(other_root, OTHER_KEY)) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + public_logger.flush() + + assert len(data_plane.rows) == 1 + assert other_plane.requests == [] + + +def test_explicit_api_key_overrides_env_ingestion_key(data_plane, monkeypatch, memory_logger): + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(data_plane.root)) + + private_logger = braintrust.init_logger(project="p", project_id="private-project", api_key=logger.TEST_API_KEY) + private_logger.log(input="x") + private_logger.flush() + + assert private_logger._public_bg_logger is None + assert [row["input"] for row in memory_logger.pop()] == ["x"] + assert data_plane.requests == [] + + +def test_explicit_api_key_and_ingestion_key_conflict(data_plane): + with pytest.raises(ValueError, match="not both"): + braintrust.init_logger(api_key="sk-private", ingestion_key=ingestion_url(data_plane.root)) + + +def test_public_and_private_loggers_do_not_mix(data_plane, memory_logger): + simulate_login() + other_plane = FakeDataPlane(key=OTHER_KEY) + with serve(other_plane) as other_root: + private_logger = braintrust.init_logger(project="private", project_id="private-project", set_current=False) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root), set_current=False) + other_logger = braintrust.init_logger(ingestion_key=ingestion_url(other_root, OTHER_KEY), set_current=False) + + def log_many(target, name): + for i in range(20): + target.log(input=f"{name}-{i}") + + threads = [ + threading.Thread(target=log_many, args=(target, name)) + for target, name in [(public_logger, "public"), (other_logger, "other")] + ] + for thread in threads: + thread.start() + # The memory logger override is thread local, so private rows log from this thread. + log_many(private_logger, "private") + for thread in threads: + thread.join() + braintrust.flush() + + assert sorted(row["input"] for row in data_plane.rows) == sorted(f"public-{i}" for i in range(20)) + assert sorted(row["input"] for row in other_plane.rows) == sorted(f"other-{i}" for i in range(20)) + assert {r["headers"]["Authorization"] for r in other_plane.requests} == {f"Bearer {OTHER_KEY}"} + assert sorted(row["input"] for row in memory_logger.pop()) == sorted(f"private-{i}" for i in range(20)) + + +def test_key_does_not_leak_through_repr(data_plane): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + + assert KEY not in repr(public_logger) + assert KEY not in repr(public_logger._public_bg_logger.endpoint) + assert KEY not in public_logger.export() From 63fbc35f21610fb44661179d1449b8dbd0226bf0 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:12:57 +0000 Subject: [PATCH 2/5] fix: Align ingestion key logger with the corrected data plane contract The data plane rejects `external_attachment` references from ingestion keys, since they can point at arbitrary object store content. Rows containing an `ExternalAttachment` are now dropped with a clear error instead of being published. Batches above 512 KiB now overflow through an upload, since Lambda base64-encodes binary bodies and ingress accepts 1 MiB by default. An explicit `ingestion_key=""` or malformed URL now raises instead of falling back to `BRAINTRUST_INGESTION_KEY` or the private login. Upload grants with malformed lifetimes, chunk counts, or ids are rejected, a 410 on a chunk or commit starts a new upload, and request timeouts and retry waits never outlive the grant. Requests don't follow redirects, so the key can't be forwarded elsewhere, and transport errors are redacted. The README now says that anyone holding a key can update any project row whose ID they know. --- py/README.md | 5 + py/src/braintrust/logger.py | 90 +++++-- py/src/braintrust/test_ingestion_key.py | 297 +++++++++++++++++++++++- 3 files changed, 364 insertions(+), 28 deletions(-) diff --git a/py/README.md b/py/README.md index 4007b660e..c4a3b2afb 100644 --- a/py/README.md +++ b/py/README.md @@ -57,6 +57,11 @@ The logger sends rows and attachments straight to the data plane in the URL, wit It never logs in, so it ignores `BRAINTRUST_API_KEY` and any earlier `braintrust.login()`, and it doesn't register or look up the project. Passing an explicit `api_key` to `init_logger` uses that key instead of `BRAINTRUST_INGESTION_KEY`, and passing both `api_key` and `ingestion_key` is an error. Feedback logged with an ingestion key can only include scores, expected values, and tags. +`Attachment` works as usual, but `ExternalAttachment` isn't supported, since its URL can point at content outside the project. + +Ingestion keys don't own the rows they write. +Anyone holding the key can update any row of the project whose ID they know, including rows logged with an API key or with another ingestion key. +Row IDs aren't secrets, so treat every row in the project as writable by whoever holds the key. To export OpenTelemetry spans with an ingestion key, point a standard OTLP exporter at the `/otel/v1/traces` path under the ingestion URL and drop the query string: diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index e1964c845..24a30f13d 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -1558,6 +1558,10 @@ def set_masking_function(self, masking_function: Callable[[Any], Any] | None): } ) _UPLOAD_REFERENCE_TYPES = {"attachment": "braintrust_attachment", "logs3_overflow": LOGS3_OVERFLOW_REFERENCE_TYPE} +_MAX_UPLOAD_GRANT_LIFETIME_MS = 300_000 +# Requests above this go through an overflow upload. Lambda deployments +# base64-encode binary bodies, and data plane ingress accepts 1 MiB by default. +DEFAULT_INGESTION_MAX_REQUEST_SIZE = 512 * 1024 @dataclasses.dataclass(frozen=True) @@ -1622,7 +1626,7 @@ def __init__(self, endpoint: _IngestionEndpoint): def _get_max_request_size(self) -> dict[str, Any]: # Any batch can overflow through the upload API, so there is no server limit to look up. - max_request_size = coalesce(self._max_request_size_override, DEFAULT_MAX_REQUEST_SIZE) + max_request_size = coalesce(self._max_request_size_override, DEFAULT_INGESTION_MAX_REQUEST_SIZE) return {"max_request_size": max_request_size, "can_use_overflow": True} def flush(self, batch_size: int | None = None): @@ -1645,7 +1649,7 @@ def flush(self, batch_size: int | None = None): except Exception as e: attachment_errors.append(e) print( - f"Dropping row {item.get(ID_FIELD)} because an attachment failed to upload: {e}", + f"Dropping row {item.get(ID_FIELD)} because of an attachment that can't be uploaded: {e}", file=self.outfile, ) continue @@ -1666,7 +1670,8 @@ def _upload_attachment( self, attachment: "BaseAttachment", uploads: "dict[int, AttachmentReference | Exception]" ) -> AttachmentReference: if isinstance(attachment, ExternalAttachment): - return attachment.reference + # The data plane rejects these, since their URL can point at any object store content. + raise ValueError("ExternalAttachment is not supported with ingestion keys. Use Attachment instead") # The same attachment object can appear in several rows of one flush. if id(attachment) not in uploads: reference = attachment.reference @@ -1733,14 +1738,26 @@ def _upload(self, purpose: str, data: bytes, content_type: str, filename: str | # The grant's lifetime is relative, so measure it from before the request goes out. requested_at = time.monotonic() grant = self._request("post", "/v1/uploads", json=body) - upload_id = str(uuid.UUID(grant["upload_id"])) - chunk_bytes = grant["chunk_bytes"] - if chunk_bytes <= 0 or grant["num_chunks"] != -(-len(data) // chunk_bytes): - raise RuntimeError(f"Invalid upload grant from the data plane: {grant}") - deadline = requested_at + grant["expires_in_ms"] / 1000 + chunk_bytes, num_chunks, expires_in_ms = ( + grant.get(k) for k in ("chunk_bytes", "num_chunks", "expires_in_ms") + ) + try: + # Normalizing the id keeps a malformed response from changing the request path. + upload_id = str(uuid.UUID(grant.get("upload_id"))) + except (TypeError, ValueError, AttributeError): + upload_id = None + if ( + upload_id is None + or not all(type(v) is int for v in (chunk_bytes, num_chunks, expires_in_ms)) + or chunk_bytes <= 0 + or num_chunks != -(-len(data) // chunk_bytes) + or not 0 < expires_in_ms <= _MAX_UPLOAD_GRANT_LIFETIME_MS + ): + raise RuntimeError(f"Invalid upload grant from the data plane: {self._redact(str(grant))}") + deadline = requested_at + expires_in_ms / 1000 try: - for index in range(grant["num_chunks"]): + for index in range(num_chunks): chunk = data[index * chunk_bytes : (index + 1) * chunk_bytes] ack = self._request( "put", @@ -1750,7 +1767,9 @@ def _upload(self, purpose: str, data: bytes, content_type: str, filename: str | headers={"Content-Type": "application/octet-stream"}, ) if ack.get("index") != index or ack.get("size_bytes") != len(chunk): - raise RuntimeError(f"Invalid chunk acknowledgement from the data plane: {ack}") + raise RuntimeError( + f"Invalid chunk acknowledgement from the data plane: {self._redact(str(ack))}" + ) result = self._request("post", f"/v1/uploads/{upload_id}/complete", deadline=deadline, json={}) except _UploadGrantExpired: print(f"Upload {upload_id} expired before it completed. Starting a new upload", file=self.outfile) @@ -1763,28 +1782,42 @@ def _upload(self, purpose: str, data: bytes, content_type: str, filename: str | or not isinstance(reference, dict) or reference.get("type") != _UPLOAD_REFERENCE_TYPES[purpose] ): - raise RuntimeError(f"Upload {upload_id} completed with unexpected content: {result}") + raise RuntimeError( + f"Upload {upload_id} completed with unexpected content: {self._redact(str(result))}" + ) return reference raise RuntimeError(f"Uploads kept expiring before they completed after {self.num_tries} attempts") def _request(self, method: str, path: str, deadline: float | None = None, **kwargs: Any) -> Any: - """Send one request, retrying transient failures, and return its JSON body.""" + """Send one request, retrying transient failures, and return its JSON body. + + Requests against an upload grant pass its `deadline`. They never wait + past it, and raise `_UploadGrantExpired` once it passes or the data + plane reports the grant as gone. + """ error = "" for i in range(self.num_tries): - if deadline is not None and time.monotonic() >= deadline: - raise _UploadGrantExpired() + timeout = self.timeout + if deadline is not None: + remaining = deadline - time.monotonic() + if remaining <= 0: + raise _UploadGrantExpired() + timeout = min(timeout, remaining) retry_after = None try: - resp = getattr(self.api_conn.get(), method)(path, timeout=self.timeout, **kwargs) + # Redirects could forward the key somewhere else. + resp = getattr(self.api_conn.get(), method)(path, timeout=timeout, allow_redirects=False, **kwargs) except requests_exceptions.RequestException as e: if not is_retryable_request_exception(e): - raise - error = str(e) + raise RuntimeError(f"{method.upper()} {path} failed with {self._redact(str(e))}") from None + error = self._redact(str(e)) else: - if resp.ok: + if 200 <= resp.status_code < 300: return resp.json() - error = f"{resp.status_code}: {resp.text}".replace(self.endpoint.key, "") + if deadline is not None and resp.status_code == 410: + raise _UploadGrantExpired() + error = self._redact(f"{resp.status_code}: {resp.text}") if resp.status_code not in DEFAULT_RETRYABLE_STATUSES: raise RuntimeError(f"{method.upper()} {path} failed with {error}") retry_after = _parse_retry_after(resp.headers.get("Retry-After"), time.time()) @@ -1793,11 +1826,16 @@ def _request(self, method: str, path: str, deadline: float | None = None, **kwar sleep_time_s = max( BACKGROUND_LOGGER_BASE_SLEEP_TIME_S * (2**i), min(retry_after or 0, DEFAULT_MAX_ELAPSED_TIME) ) + if deadline is not None and time.monotonic() + sleep_time_s >= deadline: + raise _UploadGrantExpired() print(f"{method.upper()} {path} failed with {error}. Retrying in {sleep_time_s}s", file=self.outfile) time.sleep(sleep_time_s) raise RuntimeError(f"{method.upper()} {path} failed after {self.num_tries} attempts with {error}") + def _redact(self, text: str) -> str: + return text.replace(self.endpoint.key, "") + def _internal_reset_global_state() -> None: global _state @@ -2253,11 +2291,13 @@ def init_logger( :param force_login: Login again, even if you have already logged in (by default, the logger will not login if you are already logged in) :param set_current: If true (the default), set the global current-experiment to the newly-created one. :param _create_in_project_group: Experimental: the name of an existing project group to create the project in. Existing projects must already belong to the group, otherwise registration fails with a 409. Ignored when project_id is provided. - :param ingestion_key: (Optional) The ingestion key URL of a project, for example `https://api.braintrust.dev/ingest?ingestKey=bt-ik-...`. - If the parameter is not specified and `api_key` isn't either, will try to use the `BRAINTRUST_INGESTION_KEY` environment variable. + :param ingestion_key: (Optional) The ingestion key URL of a project, for example `https://dp.example.com/ingest?ingestKey=bt-ik-...`. + If neither this nor `api_key` is specified, will try to use the `BRAINTRUST_INGESTION_KEY` environment variable. Ingestion keys can only write traces to the project they were issued for, so the logger never logs in, ignores `BRAINTRUST_API_KEY` and any previous login, and doesn't register or look up the project. `project` and `project_id` are optional and only used to label exported spans. A `project_id` is also sent with each row, and the data plane rejects rows for a different project. + Anyone holding the key can update any row of the project whose id they know, including rows logged with an API key or another + ingestion key. `ExternalAttachment` isn't supported, and rows containing one are dropped. :returns: The newly created Logger. """ @@ -2266,9 +2306,9 @@ def init_logger( if api_key is not None and ingestion_key is not None: raise ValueError("Pass either api_key or ingestion_key to init_logger, not both") - if api_key is None: - ingestion_key = ingestion_key or BraintrustEnv.INGESTION_KEY.get(None, use_dotenv=True) - if ingestion_key and _create_in_project_group is not None: + if api_key is None and ingestion_key is None: + ingestion_key = BraintrustEnv.INGESTION_KEY.get(None, use_dotenv=True) + if ingestion_key is not None and _create_in_project_group is not None: raise ValueError("Loggers using an ingestion key cannot create projects") compute_metadata_args = dict(project_name=project, project_id=project_id) @@ -2282,7 +2322,7 @@ def init_logger( "project_id": project_id, } - if ingestion_key: + if ingestion_key is not None: # The data plane resolves the org and project from the key, so there is nothing to look up. public_metadata = OrgProjectMetadata( org_id=cast(str, None), diff --git a/py/src/braintrust/test_ingestion_key.py b/py/src/braintrust/test_ingestion_key.py index 2c6eccc29..39491055d 100644 --- a/py/src/braintrust/test_ingestion_key.py +++ b/py/src/braintrust/test_ingestion_key.py @@ -7,8 +7,10 @@ import hashlib import http.server import json +import logging import re import threading +import time import uuid from contextlib import contextmanager @@ -37,6 +39,8 @@ def __init__(self, key: str = KEY, chunk_bytes: int = 4): # Responses to return before handling requests normally, keyed by (method, route). self.scripted: dict[tuple[str, str], list[tuple[int, dict, bytes]]] = {} self.grant_lifetimes_ms: list[int] = [] + # Fields to override in the next grants, to simulate malformed responses. + self.grant_overrides: list[dict] = [] self.project_id: str | None = None self.lock = threading.Lock() @@ -87,6 +91,9 @@ def _logs(self, body: dict) -> tuple[int, dict, bytes]: for row in rows: if self.project_id is not None and row.get("project_id", self.project_id) != self.project_id: return 403, {}, b'{"error": "row project does not match the ingestion key"}' + # External references can point at any object store content, so ingestion keys can't write them. + if '"external_attachment"' in json.dumps(row): + return 400, {}, b'{"error": "external_attachment references are not allowed"}' self.rows.extend(rows) return 200, {}, json.dumps({"ids": [row["id"] for row in rows]}).encode() @@ -100,6 +107,7 @@ def _create(self, body: dict) -> tuple[int, dict, bytes]: "chunk_bytes": self.chunk_bytes, "num_chunks": num_chunks, "expires_in_ms": lifetime, + **(self.grant_overrides.pop(0) if self.grant_overrides else {}), } return 201, {}, json.dumps(response).encode() @@ -352,8 +360,14 @@ def test_attachments_upload_in_advertised_chunks_before_rows(data_plane, no_priv public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) attachment = Attachment(data=b"hello world", filename="greeting.txt", content_type="text/plain") empty = Attachment(data=b"", filename="empty.txt", content_type="text/plain") - external = ExternalAttachment(url="s3://bucket/file.txt", filename="file.txt", content_type="text/plain") - public_logger.log(input={"file": attachment, "empty": empty, "external": external, "again": attachment}) + # A reference committed earlier is plain data, so it passes through unchanged. + committed = { + "type": "braintrust_attachment", + "key": str(uuid.uuid4()), + "filename": "old.txt", + "content_type": "text/plain", + } + public_logger.log(input={"file": attachment, "empty": empty, "committed": committed, "again": attachment}) public_logger.flush() [attachment_id, empty_id] = list(data_plane.uploads) @@ -392,11 +406,34 @@ def test_attachments_upload_in_advertised_chunks_before_rows(data_plane, no_priv "filename": "empty.txt", "content_type": "text/plain", }, - "external": external.reference, + "committed": committed, "again": file_reference, } +def test_external_attachments_are_rejected_before_publishing(data_plane, no_private_login): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + external = ExternalAttachment(url="s3://private-bucket/file.txt", filename="file.txt", content_type="text/plain") + public_logger.log(input={"file": external}) + kept_id = public_logger.log(input="no attachment") + + with pytest.raises(ValueError, match="ExternalAttachment is not supported with ingestion keys"): + public_logger.flush() + + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] + assert [row["id"] for row in data_plane.rows] == [kept_id] + assert b"private-bucket" not in data_plane.requests[0]["body"] + + +def test_raw_external_attachment_references_are_not_a_silent_success(data_plane, no_private_login): + # Plain dicts aren't validated client side, so the data plane's rejection must surface. + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input={"type": "external_attachment", "url": "s3://private-bucket/file.txt"}) + + assert "400" in flush_error(public_logger) + assert data_plane.rows == [] + + def test_rows_with_failed_attachments_are_dropped(data_plane, no_private_login): data_plane.script("POST", "create", (400, {}, b'{"error": "uploads disabled"}')) public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) @@ -567,3 +604,257 @@ def test_key_does_not_leak_through_repr(data_plane): assert KEY not in repr(public_logger) assert KEY not in repr(public_logger._public_bg_logger.endpoint) assert KEY not in public_logger.export() + + +def test_explicit_empty_or_malformed_ingestion_key_never_falls_back(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_API_KEY", "sk-private") + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", ingestion_url(data_plane.root)) + + for value in ["", " ", "https://dp.example/ingest", KEY]: + with pytest.raises(ValueError, match="Invalid Braintrust ingestion key URL"): + braintrust.init_logger(ingestion_key=value) + + assert braintrust.current_logger() is None + assert data_plane.requests == [] + + +def test_malformed_env_ingestion_key_never_falls_back(monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_API_KEY", "sk-private") + monkeypatch.setenv("BRAINTRUST_INGESTION_KEY", "https://dp.example/ingest?ingestKey=not-a-key") + + with pytest.raises(ValueError, match="Invalid Braintrust ingestion key URL"): + braintrust.init_logger(project="p") + + +def test_default_overflow_threshold_is_512_kib(data_plane, monkeypatch, no_private_login): + monkeypatch.delenv("BRAINTRUST_MAX_REQUEST_SIZE", raising=False) + data_plane.chunk_bytes = 512 * 1024 + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + + small_id = public_logger.log(input="x" * (200 * 1024)) + public_logger.flush() + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/logs")] + assert len(data_plane.requests[0]["body"]) < 512 * 1024 + + large_id = public_logger.log(input="x" * (600 * 1024)) + public_logger.flush() + [upload_id] = list(data_plane.uploads) + assert [method for method, _ in data_plane.paths()[1:]] == ["POST", "PUT", "PUT", "POST", "POST"] + assert json.loads(data_plane.requests[-1]["body"]) == { + "api_version": 2, + "rows": {"type": "logs3_overflow", "key": upload_id}, + } + assert [row["id"] for row in data_plane.rows] == [small_id, large_id] + assert not any(path.endswith("/version") for _, path in data_plane.paths()) + + +@pytest.mark.parametrize( + "override", + [ + {"expires_in_ms": 0}, + {"expires_in_ms": -1}, + {"expires_in_ms": 300001}, + {"expires_in_ms": "300000"}, + {"expires_in_ms": True}, + {"expires_in_ms": None}, + {"num_chunks": 3}, + {"num_chunks": 1}, + {"chunk_bytes": 0}, + {"chunk_bytes": 2.5}, + {"upload_id": "../../v1/logs"}, + {"upload_id": None}, + ], +) +def test_malformed_upload_grants_are_rejected(data_plane, override, no_private_login): + # 6 bytes in 4 byte chunks is 2 chunks. + data_plane.grant_overrides = [override] + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + + with pytest.raises(RuntimeError, match="Invalid upload grant") as exc_info: + public_logger.flush() + + assert KEY not in str(exc_info.value) + assert data_plane.paths() == [("POST", f"{BASE_PATH}/v1/uploads")] + assert data_plane.rows == [] + + +def test_maximum_grant_lifetime_is_accepted(data_plane, no_private_login): + data_plane.grant_lifetimes_ms = [300000] + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + public_logger.flush() + + assert len(data_plane.rows) == 1 + + +def test_gone_upload_grant_starts_a_new_upload(data_plane, no_private_login): + data_plane.script("PUT", "chunk", (410, {}, b'{"error": "upload grant expired"}')) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + public_logger.flush() + + [gone_id, upload_id] = list(data_plane.uploads) + assert not data_plane.uploads[gone_id]["completed"] + assert data_plane.rows[0]["input"]["key"] == upload_id + + +def test_gone_upload_grants_are_retried_a_bounded_number_of_times(data_plane, monkeypatch, no_private_login): + monkeypatch.setenv("BRAINTRUST_NUM_RETRIES", "1") + data_plane.script("PUT", "chunk", *[(410, {}, b'{"error": "upload grant expired"}')] * 5) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + + with pytest.raises(RuntimeError, match="kept expiring"): + public_logger.flush() + + assert len(data_plane.uploads) == 2 + assert data_plane.rows == [] + + +def test_retry_waits_never_outlive_the_upload_grant(data_plane, no_private_login): + data_plane.grant_lifetimes_ms = [2000] + data_plane.script("PUT", "chunk", (503, {"Retry-After": "30"}, b'{"error": "unavailable"}')) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) + + started_at = time.monotonic() + public_logger.flush() + + assert time.monotonic() - started_at < 5 + [expired_id, upload_id] = list(data_plane.uploads) + assert not data_plane.uploads[expired_id]["completed"] + assert data_plane.rows[0]["input"]["key"] == upload_id + + +def test_transport_failures_never_reveal_the_key(data_plane, monkeypatch, capsys, caplog, no_private_login): + monkeypatch.setenv("BRAINTRUST_NUM_RETRIES", "1") + # Reach the closed port directly, even if the environment configures a proxy. + monkeypatch.setenv("NO_PROXY", "127.0.0.1") + monkeypatch.setenv("no_proxy", "127.0.0.1") + caplog.set_level(logging.DEBUG) + with serve(FakeDataPlane()) as closed_root: + pass + errors = [] + for root in [closed_root, data_plane.root]: + data_plane.script("POST", "logs", *[(500, {}, f'{{"error": "bad {KEY}"}}'.encode())] * 2) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(root)) + public_logger.log(input="x") + errors.append(flush_error(public_logger)) + + assert "Connection refused" in errors[0] + assert "500" in errors[1] + output = capsys.readouterr() + for text in [*errors, output.out, output.err, caplog.text]: + assert KEY not in text + assert "ingestKey" not in text + + +def test_redirects_are_not_followed(data_plane, no_private_login): + other_plane = FakeDataPlane() + with serve(other_plane) as other_root: + data_plane.script("POST", "logs", (307, {"Location": f"{other_root}/v1/logs"}, b"")) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + public_logger.log(input="x") + error = flush_error(public_logger) + + assert "307" in error + assert other_plane.requests == [] + + +def test_background_thread_publishes_without_explicit_flush(data_plane, monkeypatch, no_private_login): + monkeypatch.delenv("BRAINTRUST_SYNC_FLUSH") + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + row_id = public_logger.log(input="x") + + deadline = time.monotonic() + 10 + while not data_plane.rows and time.monotonic() < deadline: + time.sleep(0.01) + + assert [row["id"] for row in data_plane.rows] == [row_id] + + +def test_atexit_flush_publishes_queued_rows(data_plane, monkeypatch, no_private_login): + monkeypatch.delenv("BRAINTRUST_DISABLE_ATEXIT_FLUSH") + registered = [] + monkeypatch.setattr(logger.atexit, "register", registered.append) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + row_id = public_logger.log(input="x") + + bg_logger = public_logger._public_bg_logger + assert registered == [bg_logger._finalize] + bg_logger._finalize() + + assert [row["id"] for row in data_plane.rows] == [row_id] + + +def test_public_and_private_spans_nest_without_crossing_loggers(data_plane, memory_logger): + simulate_login() + private_logger = braintrust.init_logger(project="private", project_id="private-project", set_current=False) + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root), set_current=False) + + with private_logger.start_span(name="private-root") as private_root: + with public_logger.start_span(name="public-child") as public_child: + # Global helpers follow the current span. + with braintrust.start_span(name="public-grandchild"): + pass + with private_logger.start_span(name="private-grandchild"): + pass + braintrust.flush() + + public_rows = {row["span_attributes"]["name"]: row for row in data_plane.rows} + private_rows = {row["span_attributes"]["name"]: row for row in memory_logger.pop()} + assert set(public_rows) == {"public-child", "public-grandchild"} + assert set(private_rows) == {"private-root", "private-grandchild"} + assert public_rows["public-child"]["span_parents"] == [private_root.span_id] + assert public_rows["public-grandchild"]["span_parents"] == [public_child.span_id] + assert private_rows["private-grandchild"]["span_parents"] == [public_child.span_id] + assert {row["root_span_id"] for row in [*public_rows.values(), *private_rows.values()]} == { + private_root.root_span_id + } + assert all(row["project_id"] == "private-project" for row in private_rows.values()) + assert all("project_id" not in row for row in public_rows.values()) + + +def test_public_trace_context_continues_in_a_private_service(data_plane, memory_logger): + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root), set_current=False) + with public_logger.start_span(name="client") as client_span: + headers = client_span.inject() + public_logger.flush() + + simulate_login() + # Resolving the id up front keeps the private logger from looking up its project by name. + braintrust.init_logger(project="server", project_id="server-project")._lazy_id.get() + with braintrust.start_span(name="handler", parent=braintrust.extract_trace_context(headers)): + pass + braintrust.flush() + + [handler] = memory_logger.pop() + assert handler["project_id"] == "server-project" + assert handler["root_span_id"] == client_span.root_span_id + assert handler["span_parents"] == [client_span.span_id] + assert [row["span_attributes"]["name"] for row in data_plane.rows] == ["client"] + + +def test_traced_functions_under_a_public_logger_stay_public(data_plane, memory_logger): + simulate_login() + public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) + + span_ids = [] + + @braintrust.traced + def answer(question): + span_ids.append(braintrust.current_span().id) + with braintrust.start_span(name="lookup") as lookup: + lookup.log(metadata={"q": question}) + return "42" + + answer("why") + public_logger.log_feedback(id=span_ids[0], scores={"ok": 1}) + braintrust.flush() + + rows = {row["span_attributes"]["name"]: row for row in data_plane.rows} + assert set(rows) == {"answer", "lookup"} + assert rows["answer"]["output"] == "42" + assert rows["answer"]["scores"] == {"ok": 1} + assert memory_logger.pop() == [] From 3389f40fb064c86488442ce25865026680e58df7 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:21:21 +0000 Subject: [PATCH 3/5] fix: Parse ingestion key queries like URLSearchParams Empty `&` segments are now skipped instead of rejecting the URL, which matches how the other SDKs and the data plane read the query. --- py/src/braintrust/logger.py | 3 ++- py/src/braintrust/test_ingestion_key.py | 8 ++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/py/src/braintrust/logger.py b/py/src/braintrust/logger.py index 24a30f13d..1f7996b6c 100644 --- a/py/src/braintrust/logger.py +++ b/py/src/braintrust/logger.py @@ -1577,7 +1577,8 @@ def _parse_ingestion_url(value: str) -> _IngestionEndpoint: try: parsed = urlsplit(value.strip()) parsed.port # Raises for malformed ports. - query = parse_qsl(parsed.query, keep_blank_values=True, strict_parsing=True) + # Matches URLSearchParams: empty segments are skipped and values are percent-decoded. + query = parse_qsl(parsed.query, keep_blank_values=True) reason = None except ValueError: reason = "it is not a valid URL" diff --git a/py/src/braintrust/test_ingestion_key.py b/py/src/braintrust/test_ingestion_key.py index 39491055d..d49ee5bf6 100644 --- a/py/src/braintrust/test_ingestion_key.py +++ b/py/src/braintrust/test_ingestion_key.py @@ -225,6 +225,9 @@ def assert_no_key_in_urls(plane: FakeDataPlane) -> None: (f"https://dp.example/ingest?ingestKey={KEY}", "https://dp.example/ingest"), (f"https://dp.example/deployment/base/ingest/?ingestKey={KEY}", "https://dp.example/deployment/base/ingest"), (f"http://localhost:8000/ingest?ingestKey={KEY}\n", "http://localhost:8000/ingest"), + (f"https://dp.example/base/ingest///?ingestKey={KEY}", "https://dp.example/base/ingest"), + (f"https://dp.example/ingest?&ingestKey={KEY}&&", "https://dp.example/ingest"), + (f"https://dp.example/ingest?ingest%4Bey=bt%2Dik%2D{'a' * 48}", "https://dp.example/ingest"), ], ) def test_parse_ingestion_url(url, expected): @@ -242,6 +245,11 @@ def test_parse_ingestion_url(url, expected): f"ftp://dp.example/ingest?ingestKey={KEY}", f"https://user:pass@dp.example/ingest?ingestKey={KEY}", f"https://dp.example/ingest?ingestKey={KEY}#fragment", + f"https://dp.example/ingest?ingestKey={KEY}#", + f"https://dp.example/ingest#?ingestKey={KEY}", + "https://dp.example/ingest?ingestKey", + f"https://dp.example/ingest?ingestKey={KEY}&ingestKey=", + f"https://dp.example/ingest?ingestKey={KEY}=x", f"https://dp.example/ingest/v1/logs?ingestKey={KEY}", f"https://dp.example/not-ingest?ingestKey={KEY}", f"https://dp.example/ingest?ingestKey={KEY}&ingestKey={KEY}", From f517b8633d39cfe2ea595f4ba705d3531841f962 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:33:20 +0000 Subject: [PATCH 4/5] test: Make ingestion key tests independent of platform timing Windows has a coarser monotonic clock and words refused connections differently, so expire the grant with a delayed response and match the urllib3 error text instead. --- py/src/braintrust/test_ingestion_key.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/py/src/braintrust/test_ingestion_key.py b/py/src/braintrust/test_ingestion_key.py index d49ee5bf6..85d000bb0 100644 --- a/py/src/braintrust/test_ingestion_key.py +++ b/py/src/braintrust/test_ingestion_key.py @@ -41,6 +41,8 @@ def __init__(self, key: str = KEY, chunk_bytes: int = 4): self.grant_lifetimes_ms: list[int] = [] # Fields to override in the next grants, to simulate malformed responses. self.grant_overrides: list[dict] = [] + # Delay before answering upload creation, to expire short grants regardless of clock resolution. + self.create_delay_s = 0.0 self.project_id: str | None = None self.lock = threading.Lock() @@ -98,6 +100,8 @@ def _logs(self, body: dict) -> tuple[int, dict, bytes]: return 200, {}, json.dumps({"ids": [row["id"] for row in rows]}).encode() def _create(self, body: dict) -> tuple[int, dict, bytes]: + time.sleep(self.create_delay_s) + self.create_delay_s = 0.0 upload_id = str(uuid.uuid4()) num_chunks = -(-body["size_bytes"] // self.chunk_bytes) self.uploads[upload_id] = {**body, "chunks": {}, "completed": False} @@ -456,7 +460,9 @@ def test_rows_with_failed_attachments_are_dropped(data_plane, no_private_login): def test_expired_upload_grant_starts_a_new_upload(data_plane, no_private_login): - data_plane.grant_lifetimes_ms = [1] + # The grant expires while its creation response is in flight. + data_plane.grant_lifetimes_ms = [50] + data_plane.create_delay_s = 0.2 public_logger = braintrust.init_logger(ingestion_key=ingestion_url(data_plane.root)) public_logger.log(input=Attachment(data=b"abcdef", filename="a.txt", content_type="text/plain")) public_logger.flush() @@ -750,7 +756,7 @@ def test_transport_failures_never_reveal_the_key(data_plane, monkeypatch, capsys public_logger.log(input="x") errors.append(flush_error(public_logger)) - assert "Connection refused" in errors[0] + assert "Failed to establish a new connection" in errors[0] assert "500" in errors[1] output = capsys.readouterr() for text in [*errors, output.out, output.err, caplog.text]: From 27662d0dbadf8dd65053ce3600f31bd344e83124 Mon Sep 17 00:00:00 2001 From: lforst <8118419+lforst@users.noreply.github.com> Date: Fri, 2 Oct 2026 13:46:18 +0000 Subject: [PATCH 5/5] test: Assert ingestion key transport failures with SDK-owned messages The closed-port check now relies on the SDK's retry message instead of urllib3 wording, and the grant expiry test pins the exact request sequence so nothing can be sent against the expired grant. --- py/src/braintrust/test_ingestion_key.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/py/src/braintrust/test_ingestion_key.py b/py/src/braintrust/test_ingestion_key.py index 85d000bb0..7e49fd0ee 100644 --- a/py/src/braintrust/test_ingestion_key.py +++ b/py/src/braintrust/test_ingestion_key.py @@ -468,6 +468,15 @@ def test_expired_upload_grant_starts_a_new_upload(data_plane, no_private_login): public_logger.flush() [expired_id, upload_id] = list(data_plane.uploads) + # Nothing is sent against the expired grant, so its id never shows up in a path. + assert data_plane.paths() == [ + ("POST", f"{BASE_PATH}/v1/uploads"), + ("POST", f"{BASE_PATH}/v1/uploads"), + ("PUT", f"{BASE_PATH}/v1/uploads/{upload_id}/chunks/0"), + ("PUT", f"{BASE_PATH}/v1/uploads/{upload_id}/chunks/1"), + ("POST", f"{BASE_PATH}/v1/uploads/{upload_id}/complete"), + ("POST", f"{BASE_PATH}/v1/logs"), + ] assert not data_plane.uploads[expired_id]["chunks"] assert data_plane.uploads[upload_id]["completed"] assert data_plane.rows[0]["input"]["key"] == upload_id @@ -756,8 +765,10 @@ def test_transport_failures_never_reveal_the_key(data_plane, monkeypatch, capsys public_logger.log(input="x") errors.append(flush_error(public_logger)) - assert "Failed to establish a new connection" in errors[0] - assert "500" in errors[1] + # The closed port fails without an HTTP response, and is retried like any transport error. + assert "POST /v1/logs failed after 2 attempts with " in errors[0] + assert re.search(r"failed after 2 attempts with \d{3}: ", errors[0]) is None + assert "POST /v1/logs failed after 2 attempts with 500: " in errors[1] output = capsys.readouterr() for text in [*errors, output.out, output.err, caplog.text]: assert KEY not in text