diff --git a/.github/scripts/emit_litefuse_otel_io.py b/.github/scripts/emit_litefuse_otel_io.py index 4ef45f088797de..ab53461caa4e87 100644 --- a/.github/scripts/emit_litefuse_otel_io.py +++ b/.github/scripts/emit_litefuse_otel_io.py @@ -19,6 +19,7 @@ import argparse import base64 from datetime import datetime, timedelta, timezone +import hashlib import json import os import secrets @@ -924,7 +925,7 @@ def build_subagent_session_payload(args, session_path, session_events): "name": "codex.subagent.review", "startTime": first_timestamp, "endTime": root_end, - "input": {"session_file": session_path, "thread_id": thread_id}, + "input": {"prompt": trace_input}, "output": {"final_message": trace_output}, "environment": args.environment, "metadata": trace_metadata, @@ -1024,11 +1025,16 @@ def compact_context_event(event, max_chars): return compact -def shrink_event_for_payload(event, max_payload_bytes): +def shrink_event_for_payload(event, max_payload_bytes, payload_bytes=None): shrunk = json.loads(json.dumps(event, ensure_ascii=False)) body = shrunk.get("body") if isinstance(shrunk.get("body"), dict) else {} event_name = body.get("name") + def measured_payload_bytes(candidate): + if payload_bytes is not None: + return payload_bytes(candidate) + return json_payload_bytes({"batch": [candidate]}) + def candidate_with_limits(max_chars, max_context_events=None): candidate = json.loads(json.dumps(shrunk, ensure_ascii=False)) candidate_body = candidate.get("body") if isinstance(candidate.get("body"), dict) else {} @@ -1078,26 +1084,40 @@ def candidate_with_limits(max_chars, max_context_events=None): candidate_body["metadata"] = truncate_json(metadata, max_chars) return candidate + def maximize_candidate(max_chars, max_context_events=None): + best = candidate_with_limits(max_chars, max_context_events) + lower = max_chars + 1 + upper = max_payload_bytes + while lower <= upper: + middle = (lower + upper) // 2 + candidate = candidate_with_limits(middle, max_context_events) + if measured_payload_bytes(candidate) <= max_payload_bytes: + best = candidate + lower = middle + 1 + else: + upper = middle - 1 + return best + for max_chars in (2_000, 1_000, 500, 200, 80): candidate = candidate_with_limits(max_chars) - if json_payload_bytes({"batch": [candidate]}) <= max_payload_bytes: - return candidate + if measured_payload_bytes(candidate) <= max_payload_bytes: + return maximize_candidate(max_chars) for max_context_events in (50, 20, 10, 5, 2, 1, 0): for max_chars in (80, 40, 20, 10): candidate = candidate_with_limits(max_chars, max_context_events) - if json_payload_bytes({"batch": [candidate]}) <= max_payload_bytes: - return candidate + if measured_payload_bytes(candidate) <= max_payload_bytes: + return maximize_candidate(max_chars, max_context_events) raise RuntimeError( "Litefuse ingestion event is too large after truncation: " - f"{json_payload_bytes({'batch': [event]})} bytes > {max_payload_bytes} bytes; " + f"{measured_payload_bytes(event)} bytes > {max_payload_bytes} bytes; " f"type={event.get('type')}, name={event_name}" ) -def chunk_payload(payload, max_payload_bytes): +def chunk_payload(payload, max_payload_bytes, shrink_oversized=True): batch = payload.get("batch") or [] chunks = [] current = [] @@ -1106,7 +1126,7 @@ def chunk_payload(payload, max_payload_bytes): for event in batch: event_payload = {"batch": [event]} event_bytes = json_payload_bytes(event_payload) - if event_bytes > max_payload_bytes: + if event_bytes > max_payload_bytes and shrink_oversized: event = shrink_event_for_payload(event, max_payload_bytes) event_payload = {"batch": [event]} event_bytes = json_payload_bytes(event_payload) @@ -1126,6 +1146,249 @@ def chunk_payload(payload, max_payload_bytes): return chunks +def otel_id(value, byte_count): + expected_length = byte_count * 2 + normalized = str(value or "").lower() + if len(normalized) == expected_length and all( + char in "0123456789abcdef" for char in normalized + ): + return normalized + return hashlib.blake2b(normalized.encode(), digest_size=byte_count).hexdigest() + + +def unix_nanos(timestamp): + parsed = datetime.fromisoformat(timestamp.replace("Z", "+00:00")) + if parsed.tzinfo is None: + raise ValueError(f"OpenTelemetry timestamp has no timezone: {timestamp}") + delta = parsed.astimezone(timezone.utc) - datetime(1970, 1, 1, tzinfo=timezone.utc) + seconds = delta.days * 86_400 + delta.seconds + return str(seconds * 1_000_000_000 + delta.microseconds * 1_000) + + +def otel_any_value(value): + if isinstance(value, bool): + return {"boolValue": value} + if isinstance(value, int): + return {"intValue": str(value)} + if isinstance(value, float): + return {"doubleValue": value} + if isinstance(value, str): + return {"stringValue": value} + if isinstance(value, list) and all( + isinstance(item, (bool, int, float, str)) for item in value + ): + return {"arrayValue": {"values": [otel_any_value(item) for item in value]}} + return {"stringValue": json_attr(value)} + + +def otel_attributes(values): + return [ + {"key": key, "value": otel_any_value(value)} + for key, value in values.items() + if value is not None + ] + + +def serialized_otel_value(value): + if isinstance(value, str): + return value + return json_attr(value) + + +def metadata_otel_attributes(prefix, metadata): + if not isinstance(metadata, dict): + return {prefix: serialized_otel_value(metadata)} if metadata is not None else {} + return { + f"{prefix}.{key}": ( + value if isinstance(value, (str, int)) else serialized_otel_value(value) + ) + for key, value in metadata.items() + if value is not None + } + + +def trace_body_from_payload(payload): + for event in payload.get("batch") or []: + if event.get("type") == "trace-create" and isinstance(event.get("body"), dict): + return event["body"] + raise RuntimeError("Litefuse payload is missing its trace-create context event") + + +def legacy_event_to_otel_span(event, trace_body): + event_type = event.get("type") + if event_type not in ("span-create", "generation-create"): + raise RuntimeError(f"Unsupported trace event for OTLP conversion: {event_type}") + body = event.get("body") if isinstance(event.get("body"), dict) else {} + trace_id = otel_id(body.get("traceId"), 16) + span_id = otel_id(body.get("id"), 8) + parent_id = body.get("parentObservationId") + start_time = body.get("startTime") or event.get("timestamp") + end_time = body.get("endTime") or start_time + + attributes = { + "langfuse.trace.name": trace_body.get("name"), + "session.id": trace_body.get("sessionId"), + "langfuse.trace.tags": trace_body.get("tags"), + "langfuse.environment": body.get("environment") + or trace_body.get("environment"), + "langfuse.observation.type": ( + "generation" if event_type == "generation-create" else "span" + ), + "langfuse.observation.input": ( + serialized_otel_value(body["input"]) + if body.get("input") is not None + else None + ), + "langfuse.observation.output": ( + serialized_otel_value(body["output"]) + if body.get("output") is not None + else None + ), + "langfuse.observation.level": body.get("level"), + "langfuse.observation.status_message": body.get("statusMessage"), + } + attributes.update( + metadata_otel_attributes( + "langfuse.trace.metadata", trace_body.get("metadata") or {} + ) + ) + attributes.update( + metadata_otel_attributes( + "langfuse.observation.metadata", body.get("metadata") or {} + ) + ) + if event_type == "generation-create": + attributes["langfuse.observation.model.name"] = body.get("model") + if body.get("usageDetails") is not None: + attributes["langfuse.observation.usage_details"] = serialized_otel_value( + body["usageDetails"] + ) + if not parent_id: + attributes["langfuse.internal.is_app_root"] = True + + span = { + "traceId": trace_id, + "spanId": span_id, + "name": body.get("name") or "codex.unknown", + "kind": 1, + "startTimeUnixNano": unix_nanos(start_time), + "endTimeUnixNano": unix_nanos(end_time), + "attributes": otel_attributes(attributes), + "status": { + "code": 2 if body.get("level") == "ERROR" else 1, + **( + {"message": body["statusMessage"]} + if body.get("statusMessage") + else {} + ), + }, + "flags": 1, + } + if parent_id: + span["parentSpanId"] = otel_id(parent_id, 8) + return span + + +def otlp_payload(trace_body, events): + spans = [legacy_event_to_otel_span(event, trace_body) for event in events] + return { + "resourceSpans": [ + { + "resource": { + "attributes": otel_attributes( + { + "service.name": "doris-code-review", + "langfuse.environment": trace_body.get("environment"), + } + ) + }, + "scopeSpans": [ + { + "scope": {"name": "doris-litefuse-exporter", "version": "2"}, + "spans": spans, + } + ], + } + ] + } + + +def otlp_span_count(payload): + return sum( + len(scope_spans.get("spans") or []) + for resource_spans in payload.get("resourceSpans") or [] + for scope_spans in resource_spans.get("scopeSpans") or [] + ) + + +def otlp_chunks(payload, max_payload_bytes, trace_body=None): + trace_body = trace_body or trace_body_from_payload(payload) + events = [ + event + for event in payload.get("batch") or [] + if event.get("type") in ("span-create", "generation-create") + ] + chunks = [] + + def add_chunk(candidate_events): + candidate_payload = otlp_payload(trace_body, candidate_events) + request_size = json_payload_bytes(candidate_payload) + if request_size <= max_payload_bytes: + chunks.append( + ( + {"batch": candidate_events}, + candidate_payload, + request_size, + ) + ) + return + if len(candidate_events) > 1: + middle = len(candidate_events) // 2 + add_chunk(candidate_events[:middle]) + add_chunk(candidate_events[middle:]) + return + + event = candidate_events[0] + shrunk_event = shrink_event_for_payload( + event, + max_payload_bytes, + payload_bytes=lambda candidate: json_payload_bytes( + otlp_payload(trace_body, [candidate]) + ), + ) + shrunk_payload = otlp_payload(trace_body, [shrunk_event]) + chunks.append( + ( + {"batch": [shrunk_event]}, + shrunk_payload, + json_payload_bytes(shrunk_payload), + ) + ) + + if not events: + raise RuntimeError("Litefuse payload contains no spans for OTLP ingestion") + prechunk_limit = max(1_000, max_payload_bytes // 2) + # Bound multi-event OTLP encodes without truncating an individual event before + # add_chunk measures its encoded span against the full request limit. + for legacy_chunk, _legacy_size in chunk_payload( + {"batch": events}, prechunk_limit, shrink_oversized=False + ): + add_chunk(legacy_chunk["batch"]) + return chunks + + +def split_otlp_chunk(chunk, max_payload_bytes, trace_body): + events = chunk.get("batch") or [] + if len(events) < 2: + raise RuntimeError("Cannot split an OTLP chunk with fewer than two spans") + middle = len(events) // 2 + return otlp_chunks( + {"batch": events[:middle]}, max_payload_bytes, trace_body + ) + otlp_chunks( + {"batch": events[middle:]}, max_payload_bytes, trace_body + ) + + def post_payload_once(endpoint, public_key, secret_key, payload, timeout_seconds): auth = base64.b64encode(f"{public_key}:{secret_key}".encode()).decode() request = urllib.request.Request( @@ -1134,46 +1397,36 @@ def post_payload_once(endpoint, public_key, secret_key, payload, timeout_seconds headers={ "Content-Type": "application/json", "Authorization": f"Basic {auth}", + "x-langfuse-ingestion-version": "4", + "x-langfuse-sdk-name": "doris-code-review", + "x-langfuse-sdk-version": "2", }, method="POST", ) with urllib.request.urlopen(request, timeout=timeout_seconds) as response: body = response.read().decode() detail = json.loads(body) if body else {} - errors = detail.get("errors") if isinstance(detail, dict) else None - if errors: - raise RuntimeError(f"Litefuse ingestion returned errors: {json_attr(errors)}") + partial_success = ( + detail.get("partialSuccess") or detail.get("partial_success") or {} + if isinstance(detail, dict) + else {} + ) + rejected_spans = int( + partial_success.get("rejectedSpans") + or partial_success.get("rejected_spans") + or 0 + ) + if rejected_spans: + raise RuntimeError( + "Litefuse OTLP ingestion partially rejected " + f"{rejected_spans} spans: {json_attr(partial_success)}" + ) return { "status": response.status, - "success_count": len(detail.get("successes") or []) - if isinstance(detail, dict) - else 0, + "success_count": otlp_span_count(payload), } -def retry_payload_chunks_after_413(payload, request_size, max_payload_bytes): - batch = payload.get("batch") or [] - if not batch: - raise RuntimeError( - "Litefuse ingestion returned 413 for an empty payload chunk" - ) - - next_limit = max(1_000, min(max_payload_bytes - 1, request_size // 2)) - if len(batch) == 1: - event = shrink_event_for_payload(batch[0], next_limit) - return [({"batch": [event]}, json_payload_bytes({"batch": [event]}))] - - return chunk_payload(payload, next_limit) - - -def retry_payload_chunks_after_transport_error(payload, request_size, max_payload_bytes): - batch = payload.get("batch") or [] - if len(batch) <= 1: - return [(payload, request_size)] - next_limit = max(1_000, min(max_payload_bytes - 1, request_size // 2)) - return chunk_payload(payload, next_limit) - - def post_payload( endpoint, public_key, @@ -1189,38 +1442,52 @@ def post_payload( request_sizes = [] payload_too_large_retry_count = 0 transport_retry_count = 0 - chunks = chunk_payload(payload, max_payload_bytes) + trace_body = trace_body_from_payload(payload) + chunks = otlp_chunks(payload, max_payload_bytes, trace_body) while chunks: - chunk, request_size = chunks.pop(0) + chunk, otlp_chunk, request_size = chunks.pop(0) try: status = post_payload_once( - endpoint, public_key, secret_key, chunk, timeout_seconds + endpoint, public_key, secret_key, otlp_chunk, timeout_seconds ) except urllib.error.HTTPError as exc: if exc.code != 413: - raise + error_body = exc.read().decode("utf-8", errors="replace") + raise RuntimeError( + "Litefuse OTLP ingestion returned " + f"HTTP {exc.code}: {truncate_text(error_body, 4_000)}" + ) from exc payload_too_large_retry_count += 1 - chunks = ( - retry_payload_chunks_after_413( - chunk, request_size, max_payload_bytes - ) - + chunks - ) + chunk_events = chunk.get("batch") or [] + if len(chunk_events) > 1: + chunks = split_otlp_chunk( + chunk, max_payload_bytes, trace_body + ) + chunks + continue + # Retry a strictly smaller complete OTLP envelope. Using half of the + # rejected size can be smaller than the fixed resource/trace attributes + # and prevent a viable reduced observation from being attempted. + next_limit = request_size - 1 + chunks = otlp_chunks(chunk, next_limit, trace_body) + chunks continue except (TimeoutError, urllib.error.URLError) as exc: transport_retry_count += 1 if transport_retry_count > retry_attempts: raise RuntimeError( - "Litefuse ingestion failed after transport retries: " + "Litefuse OTLP ingestion failed after transport retries: " f"{type(exc).__name__}: {exc}" ) from exc time.sleep(retry_sleep_seconds) - chunks = ( - retry_payload_chunks_after_transport_error( - chunk, request_size, max_payload_bytes + chunk_events = chunk.get("batch") or [] + if len(chunk_events) > 1: + chunks = split_otlp_chunk( + chunk, max_payload_bytes, trace_body + ) + chunks + else: + chunks.insert( + 0, + (chunk, otlp_chunk, request_size), ) - + chunks - ) continue statuses.append(status["status"]) success_count += int(status.get("success_count") or 0) @@ -1247,25 +1514,40 @@ def fetch_trace(base_url, public_key, secret_key, trace_id): return json.loads(response.read().decode()) -def fetch_observations_v2(base_url, public_key, secret_key, trace_id): +def fetch_observations_v2( + base_url, public_key, secret_key, trace_id, max_pages=10 +): auth = base64.b64encode(f"{public_key}:{secret_key}".encode()).decode() now = datetime.now(timezone.utc) - params = urllib.parse.urlencode( - { - "traceId": trace_id, - "fromStartTime": (now - timedelta(hours=1)).isoformat().replace("+00:00", "Z"), - "toStartTime": (now + timedelta(minutes=5)).isoformat().replace("+00:00", "Z"), - "fields": "core,basic,io,trace_context,model,usage", - "limit": "100", - } - ) - request = urllib.request.Request( - f"{base_url.rstrip('/')}/api/public/v2/observations?{params}", - headers={"Authorization": f"Basic {auth}"}, - method="GET", + query = { + "traceId": trace_id, + "fromStartTime": (now - timedelta(hours=1)).isoformat().replace("+00:00", "Z"), + "toStartTime": (now + timedelta(minutes=5)).isoformat().replace("+00:00", "Z"), + "fields": "core,basic,io,trace_context,model,usage", + "limit": "1000", + } + rows = [] + cursor = "" + for _ in range(max_pages): + if cursor: + query["cursor"] = cursor + params = urllib.parse.urlencode(query) + request = urllib.request.Request( + f"{base_url.rstrip('/')}/api/public/v2/observations?{params}", + headers={"Authorization": f"Basic {auth}"}, + method="GET", + ) + with urllib.request.urlopen(request, timeout=30) as response: + payload = json.loads(response.read().decode()) + rows.extend(observation_rows_from_v2(payload)) + meta = payload.get("meta") if isinstance(payload, dict) else {} + cursor = meta.get("cursor") if isinstance(meta, dict) else "" + if not cursor: + return {**payload, "data": rows} + raise RuntimeError( + "Litefuse v2 observations remained paginated after " + f"{max_pages} pages for trace {trace_id}" ) - with urllib.request.urlopen(request, timeout=30) as response: - return json.loads(response.read().decode()) def fetch_observations_legacy( @@ -1536,7 +1818,7 @@ def parse_args(): def main(): args = parse_args() - endpoint = args.endpoint or f"{args.base_url.rstrip('/')}/api/public/ingestion" + endpoint = args.endpoint or f"{args.base_url.rstrip('/')}/api/public/otel/v1/traces" if args.max_context_json_chars <= 0: args.max_context_json_chars = args.max_json_chars @@ -1570,16 +1852,18 @@ def main(): if args.dry_run: result["batch_count"] = len(payload["batch"]) result["event_types"] = [event["type"] for event in payload["batch"][:10]] - chunks = chunk_payload(payload, args.max_payload_bytes) + chunks = otlp_chunks(payload, args.max_payload_bytes) result["request_count"] = len(chunks) - result["request_sizes"] = [request_size for _, request_size in chunks] + result["request_sizes"] = [request_size for _, _, request_size in chunks] result["max_request_size"] = ( max(result["request_sizes"]) if result["request_sizes"] else 0 ) result["subagent_traces"] = [] for subagent_payload in subagent_payloads: - chunks = chunk_payload(subagent_payload["payload"], args.max_payload_bytes) - request_sizes = [request_size for _chunk, request_size in chunks] + chunks = otlp_chunks( + subagent_payload["payload"], args.max_payload_bytes + ) + request_sizes = [request_size for _chunk, _otel, request_size in chunks] result["subagent_traces"].append( { "trace_id": subagent_payload["trace_id"], diff --git a/.github/scripts/test_emit_litefuse_otel_io.py b/.github/scripts/test_emit_litefuse_otel_io.py new file mode 100644 index 00000000000000..8af9d509f09602 --- /dev/null +++ b/.github/scripts/test_emit_litefuse_otel_io.py @@ -0,0 +1,502 @@ +#!/usr/bin/env python3 +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import importlib.util +import io +import json +from pathlib import Path +from types import SimpleNamespace +import unittest +from unittest import mock +import urllib.error +import urllib.parse + + +MODULE_PATH = Path(__file__).with_name("emit_litefuse_otel_io.py") +SPEC = importlib.util.spec_from_file_location("emit_litefuse_otel_io", MODULE_PATH) +MODULE = importlib.util.module_from_spec(SPEC) +SPEC.loader.exec_module(MODULE) + + +def attribute_values(span): + values = {} + for attribute in span["attributes"]: + value = attribute["value"] + if "arrayValue" in value: + values[attribute["key"]] = [ + next(iter(item.values())) + for item in value["arrayValue"]["values"] + ] + else: + values[attribute["key"]] = next(iter(value.values())) + return values + + +class FakeResponse: + status = 200 + + def __enter__(self): + return self + + def __exit__(self, _exc_type, _exc, _traceback): + return False + + def read(self): + return b"{}" + + +class PartialSuccessResponse(FakeResponse): + def read(self): + return b'{"partialSuccess":{"rejectedSpans":"1","errorMessage":"bad span"}}' + + +class JsonResponse(FakeResponse): + def __init__(self, payload): + self.payload = payload + + def read(self): + return json.dumps(self.payload).encode() + + +class LitefuseOtelExporterTest(unittest.TestCase): + def trace_body(self): + return { + "id": "1" * 32, + "name": "doris-ai-review", + "sessionId": "run-123", + "environment": "github-actions", + "metadata": {"repository": "apache/doris", "codex_jsonl": True}, + "tags": ["doris-ai-review", "codex-jsonl"], + } + + def span_event(self, span_id, parent_id=None, output_size=0): + body = { + "id": span_id, + "traceId": "1" * 32, + "name": "codex.command", + "startTime": "2026-09-01T00:00:00.000Z", + "endTime": "2026-09-01T00:00:01.000Z", + "input": {"command": "git status"}, + "output": {"status": "completed", "text": "x" * output_size}, + "environment": "github-actions", + "metadata": {"item_type": "command_execution"}, + "level": "DEFAULT", + } + if parent_id: + body["parentObservationId"] = parent_id + return {"type": "span-create", "body": body} + + def test_converts_legacy_events_to_otlp_hierarchy_and_attributes(self): + root_id = "2" * 32 + child_id = "3" * 32 + root = self.span_event(root_id) + child = self.span_event(child_id, root_id) + + payload = MODULE.otlp_payload(self.trace_body(), [root, child]) + spans = payload["resourceSpans"][0]["scopeSpans"][0]["spans"] + + self.assertEqual(len(spans), 2) + self.assertEqual(spans[0]["traceId"], "1" * 32) + self.assertEqual(len(spans[0]["spanId"]), 16) + self.assertNotIn("parentSpanId", spans[0]) + self.assertEqual(spans[1]["parentSpanId"], spans[0]["spanId"]) + attributes = attribute_values(spans[0]) + self.assertEqual(attributes["langfuse.trace.name"], "doris-ai-review") + self.assertEqual(attributes["session.id"], "run-123") + self.assertEqual(attributes["langfuse.observation.type"], "span") + self.assertTrue(attributes["langfuse.internal.is_app_root"]) + self.assertEqual( + attributes["langfuse.trace.tags"], + ["doris-ai-review", "codex-jsonl"], + ) + + def test_subagent_root_observation_preserves_task_input(self): + args = SimpleNamespace( + max_input_chars=200_000, + max_output_chars=200_000, + max_json_chars=40_000, + repository="apache/doris", + workflow="Code Review", + run_id="run-123", + pr_number="67413", + head_sha="a" * 40, + base_sha="b" * 40, + reasoning_effort="xhigh", + session_id="run-123", + subagent_trace_name="doris-ai-review-subagent", + environment="github-actions", + model="gpt-5.6-sol", + ) + session_path = "/tmp/thread-123.jsonl" + events = [ + { + "type": "session_meta", + "timestamp": "2026-09-01T00:00:00.000Z", + "payload": {"id": "thread-123"}, + }, + { + "type": "response_item", + "timestamp": "2026-09-01T00:00:01.000Z", + "payload": { + "type": "message", + "role": "user", + "content": [{"type": "input_text", "text": "review task"}], + }, + }, + { + "type": "response_item", + "timestamp": "2026-09-01T00:00:02.000Z", + "payload": { + "type": "message", + "role": "assistant", + "content": [{"type": "output_text", "text": "review result"}], + }, + }, + ] + + result = MODULE.build_subagent_session_payload(args, session_path, events) + chunks = MODULE.otlp_chunks(result["payload"], max_payload_bytes=20_000) + spans = [ + span + for _legacy, otel, _size in chunks + for resource in otel["resourceSpans"] + for scope in resource["scopeSpans"] + for span in scope["spans"] + ] + root = next(span for span in spans if span["name"] == "codex.subagent.review") + attributes = attribute_values(root) + + self.assertEqual( + json.loads(attributes["langfuse.observation.input"]), + {"prompt": "review task"}, + ) + self.assertEqual( + attributes["langfuse.observation.metadata.session_file"], session_path + ) + self.assertEqual( + attributes["langfuse.observation.metadata.thread_id"], "thread-123" + ) + + def test_chunks_encoded_otlp_payloads_to_requested_size(self): + trace_event = {"type": "trace-create", "body": self.trace_body()} + span_events = [ + self.span_event(f"{index:032x}", output_size=2_000) + for index in range(1, 9) + ] + + chunks = MODULE.otlp_chunks( + {"batch": [trace_event, *span_events]}, max_payload_bytes=6_000 + ) + + self.assertGreater(len(chunks), 1) + self.assertEqual( + sum(MODULE.otlp_span_count(otel) for _legacy, otel, _size in chunks), + len(span_events), + ) + self.assertTrue(all(size <= 6_000 for _legacy, _otel, size in chunks)) + + def test_truncates_one_oversized_span(self): + trace_event = {"type": "trace-create", "body": self.trace_body()} + span_event = self.span_event("2" * 32, output_size=50_000) + + chunks = MODULE.otlp_chunks( + {"batch": [trace_event, span_event]}, max_payload_bytes=6_000 + ) + + self.assertEqual(len(chunks), 1) + _legacy, otel, size = chunks[0] + self.assertLessEqual(size, 6_000) + output = attribute_values( + otel["resourceSpans"][0]["scopeSpans"][0]["spans"][0] + )["langfuse.observation.output"] + self.assertIn("truncated_json", output) + + def test_does_not_truncate_span_below_full_otlp_limit(self): + trace_event = {"type": "trace-create", "body": self.trace_body()} + span_event = self.span_event("2" * 32, output_size=3_500) + original = MODULE.otlp_payload(self.trace_body(), [span_event]) + original_size = MODULE.json_payload_bytes(original) + + self.assertGreater(original_size, 3_000) + self.assertLessEqual(original_size, 6_000) + chunks = MODULE.otlp_chunks( + {"batch": [trace_event, span_event]}, max_payload_bytes=6_000 + ) + + self.assertEqual(len(chunks), 1) + _legacy, otel, size = chunks[0] + self.assertEqual(size, original_size) + output = attribute_values( + otel["resourceSpans"][0]["scopeSpans"][0]["spans"][0] + )["langfuse.observation.output"] + self.assertNotIn("truncated_json", output) + self.assertEqual(len(json.loads(output)["text"]), 3_500) + + def test_preserves_near_limit_single_span_payload(self): + trace_event = {"type": "trace-create", "body": self.trace_body()} + span_event = self.span_event("2" * 32, output_size=4_509) + original = MODULE.otlp_payload(self.trace_body(), [span_event]) + original_size = MODULE.json_payload_bytes(original) + + self.assertGreater(original_size, 6_000) + self.assertLess(original_size, 6_100) + chunks = MODULE.otlp_chunks( + {"batch": [trace_event, span_event]}, max_payload_bytes=6_000 + ) + + self.assertEqual(len(chunks), 1) + _legacy, otel, size = chunks[0] + self.assertGreater(size, 5_500) + self.assertLessEqual(size, 6_000) + output = attribute_values( + otel["resourceSpans"][0]["scopeSpans"][0]["spans"][0] + )["langfuse.observation.output"] + self.assertGreater(len(json.loads(output)["truncated_json"]), 4_000) + + def test_prechunks_before_otlp_encoding(self): + trace_event = {"type": "trace-create", "body": self.trace_body()} + span_events = [ + self.span_event(f"{index:032x}", output_size=1_000) + for index in range(1, 21) + ] + encoded_batch_sizes = [] + original_otlp_payload = MODULE.otlp_payload + + def recording_otlp_payload(trace_body, events): + encoded_batch_sizes.append(len(events)) + return original_otlp_payload(trace_body, events) + + with mock.patch.object( + MODULE, "otlp_payload", side_effect=recording_otlp_payload + ): + MODULE.otlp_chunks( + {"batch": [trace_event, *span_events]}, max_payload_bytes=6_000 + ) + + self.assertLess(max(encoded_batch_sizes), len(span_events)) + + def test_posts_otlp_v4_headers(self): + payload = MODULE.otlp_payload( + self.trace_body(), [self.span_event("2" * 32)] + ) + captured = {} + + def fake_urlopen(request, timeout): + captured["request"] = request + captured["timeout"] = timeout + return FakeResponse() + + with mock.patch.object(MODULE.urllib.request, "urlopen", fake_urlopen): + status = MODULE.post_payload_once( + "https://litefuse.example/api/public/otel/v1/traces", + "public", + "secret", + payload, + 30, + ) + + headers = {key.lower(): value for key, value in captured["request"].header_items()} + self.assertEqual(headers["content-type"], "application/json") + self.assertEqual(headers["x-langfuse-ingestion-version"], "4") + self.assertEqual(headers["x-langfuse-sdk-name"], "doris-code-review") + self.assertEqual(captured["timeout"], 30) + self.assertEqual(status["success_count"], 1) + + def test_paginates_v2_observations_until_root_is_visible(self): + first_page = [{"id": f"child-{index}"} for index in range(1_000)] + responses = [ + JsonResponse({"data": first_page, "meta": {"cursor": "next"}}), + JsonResponse({"data": [{"id": "root"}], "meta": {}}), + ] + requests = [] + + def fake_urlopen(request, timeout): + requests.append((request, timeout)) + return responses.pop(0) + + with mock.patch.object(MODULE.urllib.request, "urlopen", fake_urlopen): + payload = MODULE.fetch_observations_v2( + "https://litefuse.example", "public", "secret", "trace-id" + ) + + self.assertEqual(len(payload["data"]), 1_001) + self.assertEqual(payload["data"][-1], {"id": "root"}) + first_query = urllib.parse.parse_qs( + urllib.parse.urlparse(requests[0][0].full_url).query + ) + second_query = urllib.parse.parse_qs( + urllib.parse.urlparse(requests[1][0].full_url).query + ) + self.assertEqual(first_query["limit"], ["1000"]) + self.assertNotIn("cursor", first_query) + self.assertEqual(second_query["cursor"], ["next"]) + + def test_rejects_incomplete_v2_observation_pagination(self): + response = JsonResponse( + {"data": [{"id": "newest"}], "meta": {"cursor": "still-more"}} + ) + + with mock.patch.object( + MODULE.urllib.request, "urlopen", return_value=response + ): + with self.assertRaisesRegex(RuntimeError, "remained paginated"): + MODULE.fetch_observations_v2( + "https://litefuse.example", + "public", + "secret", + "trace-id", + max_pages=1, + ) + + def test_rejects_otlp_partial_success(self): + payload = MODULE.otlp_payload( + self.trace_body(), [self.span_event("2" * 32)] + ) + + with mock.patch.object( + MODULE.urllib.request, "urlopen", return_value=PartialSuccessResponse() + ): + with self.assertRaisesRegex(RuntimeError, "partially rejected 1 spans"): + MODULE.post_payload_once( + "https://litefuse.example/api/public/otel/v1/traces", + "public", + "secret", + payload, + 30, + ) + + def test_splits_multi_span_chunk_after_transport_error(self): + trace_event = {"type": "trace-create", "body": self.trace_body()} + payload = { + "batch": [ + trace_event, + self.span_event("2" * 32), + self.span_event("3" * 32, "2" * 32), + ] + } + + with mock.patch.object( + MODULE.urllib.request, + "urlopen", + side_effect=[ + urllib.error.URLError("connection reset"), + FakeResponse(), + FakeResponse(), + ], + ): + status = MODULE.post_payload( + "https://litefuse.example/api/public/otel/v1/traces", + "public", + "secret", + payload, + 10_000, + 30, + 3, + 0, + ) + + self.assertEqual(status["transport_retries"], 1) + self.assertEqual(status["request_count"], 2) + self.assertEqual(status["success_count"], 2) + + def test_splits_multi_span_chunk_after_http_413(self): + trace_event = {"type": "trace-create", "body": self.trace_body()} + payload = { + "batch": [ + trace_event, + self.span_event("2" * 32), + self.span_event("3" * 32, "2" * 32), + ] + } + payload_too_large = urllib.error.HTTPError( + "https://litefuse.example/api/public/otel/v1/traces", + 413, + "Payload Too Large", + {}, + io.BytesIO(b"payload too large"), + ) + + with mock.patch.object( + MODULE.urllib.request, + "urlopen", + side_effect=[payload_too_large, FakeResponse(), FakeResponse()], + ): + status = MODULE.post_payload( + "https://litefuse.example/api/public/otel/v1/traces", + "public", + "secret", + payload, + 10_000, + 30, + 3, + 0, + ) + + self.assertEqual(status["payload_too_large_retries"], 1) + self.assertEqual(status["request_count"], 2) + self.assertEqual(status["success_count"], 2) + + def test_retries_single_span_413_without_half_size_ceiling(self): + trace_body = { + **self.trace_body(), + "metadata": {"repository": "apache/doris", "fixed": "m" * 3_000}, + } + trace_event = {"type": "trace-create", "body": trace_body} + span_event = self.span_event("2" * 32, output_size=1_500) + server_limit = ( + MODULE.json_payload_bytes(MODULE.otlp_payload(trace_body, [span_event])) - 1 + ) + request_sizes = [] + + def reject_first_request(request, timeout): + request_sizes.append(len(request.data)) + if len(request.data) > server_limit: + raise urllib.error.HTTPError( + request.full_url, + 413, + "Payload Too Large", + {}, + io.BytesIO(b"payload too large"), + ) + return FakeResponse() + + with mock.patch.object( + MODULE.urllib.request, "urlopen", side_effect=reject_first_request + ): + status = MODULE.post_payload( + "https://litefuse.example/api/public/otel/v1/traces", + "public", + "secret", + {"batch": [trace_event, span_event]}, + 10_000, + 30, + 3, + 0, + ) + + self.assertEqual(status["payload_too_large_retries"], 1) + self.assertEqual(status["request_count"], 1) + self.assertEqual(status["success_count"], 1) + self.assertEqual(len(request_sizes), 2) + self.assertLessEqual(request_sizes[1], server_limit) + self.assertLess(request_sizes[1], request_sizes[0]) + self.assertGreater(request_sizes[1], request_sizes[0] // 2) + + +if __name__ == "__main__": + unittest.main()