From 1767b6d5f495d65ec842703d76178dbabee61419 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Mon, 27 Jul 2026 11:51:11 +0000 Subject: [PATCH 01/14] feat(telemetry): add OpenTelemetry tracing, metrics, logs and cross-service e2e test MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Configure W3C TraceContext propagation with CLIENT/SERVER spans across web, AI, searcher, indexer, and connector-manager - Add OTLP metrics providers with HTTP RED, queue, searcher, connector sync, and AI embedding metrics - Add OTLP log export with trace/span correlation and PII-safe attribute allowlist - Persist traceparent/tracestate in queue rows for async producer→consumer trace links - Fix Rust BatchSpanProcessor Tokio runtime panic that silently dropped all trace exports - Add e2e cross-service tracing test that validates parent/child span continuity across Node.js and Rust services - Remove bundled local OTel Collector, smoke/validate scripts, and observability profile from compose --- .env.example | 18 +- Cargo.toml | 5 +- README.md | 5 + docker/docker-compose.yml | 1 + services/ai/db/embedding_queue.py | 10 +- services/ai/email_service/sender.py | 8 +- services/ai/embeddings/batch_processor.py | 283 +++++- services/ai/main.py | 19 +- services/ai/providers/anthropic.py | 36 +- services/ai/providers/bedrock.py | 53 +- services/ai/providers/gemini.py | 13 +- services/ai/providers/openai.py | 13 +- services/ai/providers/openai_compatible.py | 7 +- services/ai/pyproject.toml | 2 + services/ai/streaming/generate.py | 31 +- services/ai/streaming/run.py | 2 +- services/ai/telemetry.py | 238 ++++- services/ai/tests/unit/test_queue_tracing.py | 538 +++++++++++ services/ai/tests/unit/test_sanitization.py | 217 +++++ services/ai/tests/unit/test_telemetry.py | 212 +++++ .../ai/tests/unit/test_telemetry_lifecycle.py | 317 +++++++ services/ai/tools/connector_handler.py | 16 +- services/ai/tools/search_handler.py | 14 +- services/ai/tools/searcher_client.py | 26 +- services/connector-manager/Cargo.toml | 1 + .../connector-manager/src/connector_client.rs | 151 ++- services/connector-manager/src/handlers.rs | 677 +++++--------- services/connector-manager/src/lib.rs | 13 +- .../connector-manager/src/sync_manager.rs | 182 ++-- .../tests/sanitization_test.rs | 166 ++++ services/indexer/Cargo.toml | 1 + services/indexer/src/lib.rs | 43 +- services/indexer/src/queue_processor.rs | 161 +++- .../106_add_queue_trace_context.sql | 18 + services/searcher/src/handlers.rs | 103 ++- services/searcher/src/lib.rs | 17 +- services/searcher/src/search.rs | 75 +- services/searcher/src/search_repository.rs | 6 +- services/searcher/src/suggested_questions.rs | 95 +- shared/Cargo.toml | 8 + shared/src/clients/ai.rs | 24 +- shared/src/clients/docling.rs | 37 +- shared/src/embedding_queue.rs | 249 +++-- shared/src/lib.rs | 5 +- shared/src/metrics.rs | 258 ++++++ shared/src/models.rs | 7 + shared/src/queue.rs | 140 ++- shared/src/telemetry.rs | 778 ++++++++++++++-- shared/tests/collector/server.mjs | 76 ++ shared/tests/cross_service_tracing.rs | 484 ++++++++++ shared/tests/embedding_queue_test.rs | 190 ++++ shared/tests/event_queue_test.rs | 130 ++- shared/tests/metrics_test.rs | 399 ++++++++ shared/tests/queue_link_test.rs | 357 ++++++++ shared/tests/telemetry_test.rs | 856 ++++++++++++++++++ web/Dockerfile | 20 +- web/package-lock.json | 5 + web/package.json | 9 +- web/src/hooks.server.ts | 79 +- web/src/instrumentation.mjs | 65 ++ web/src/lib/server/auth.ts | 2 +- web/src/lib/server/config.ts | 2 +- web/src/lib/server/email/providers/acs.ts | 4 +- web/src/lib/server/email/providers/resend.ts | 4 +- web/src/lib/server/email/types.ts | 6 +- web/src/lib/server/logger.ts | 85 +- web/src/lib/server/oauth/accountLinking.ts | 18 +- web/src/lib/server/oauth/connectorOAuth.ts | 20 - web/src/lib/server/otel-log-hook.mjs | 194 ++++ web/src/lib/server/otel-log-hook.test.ts | 275 ++++++ web/src/lib/server/otel-runtime.test.ts | 271 ++++++ .../server/otel-source-sanitization.test.ts | 71 ++ web/src/lib/server/otel-utils.mjs | 28 + web/src/lib/server/telemetry.test.ts | 82 ++ web/src/lib/server/telemetry.ts | 136 ++- web/src/otel-factory.mjs | 88 ++ web/src/routes/(public)/auth/entra/+server.ts | 2 +- .../routes/(public)/auth/google/+server.ts | 2 +- .../(public)/auth/google/callback/+server.ts | 2 +- .../(public)/auth/google/unlink/+server.ts | 4 +- web/src/routes/api/chat/+server.ts | 7 +- web/src/routes/api/chat/[chatId]/+server.ts | 16 +- .../api/chat/[chatId]/approve/+server.ts | 7 +- .../[chatId]/artifacts/[...path]/+server.ts | 4 +- .../api/chat/[chatId]/messages/+server.ts | 368 ++++---- .../messages/[messageId]/edit/+server.ts | 141 +-- .../messages/[messageId]/feedback/+server.ts | 29 +- .../routes/api/chat/[chatId]/stop/+server.ts | 3 +- .../api/chat/[chatId]/stream/+server.ts | 74 +- web/src/routes/api/oauth/callback/+server.ts | 17 +- web/src/routes/api/search/+server.ts | 7 +- .../routes/api/sources/[sourceId]/+server.ts | 2 +- .../api/sources/[sourceId]/action/+server.ts | 3 +- .../api/sources/[sourceId]/sync/+server.ts | 2 +- web/src/routes/api/uploads/+server.ts | 6 +- web/src/routes/api/user/settings/+server.ts | 4 +- .../routes/api/v1/documents/[id]/+server.ts | 11 +- web/src/routes/api/v1/search/+server.ts | 9 +- 98 files changed, 8077 insertions(+), 1898 deletions(-) create mode 100644 services/ai/tests/unit/test_queue_tracing.py create mode 100644 services/ai/tests/unit/test_sanitization.py create mode 100644 services/ai/tests/unit/test_telemetry.py create mode 100644 services/ai/tests/unit/test_telemetry_lifecycle.py create mode 100644 services/connector-manager/tests/sanitization_test.rs create mode 100644 services/migrations/106_add_queue_trace_context.sql create mode 100644 shared/src/metrics.rs create mode 100644 shared/tests/collector/server.mjs create mode 100644 shared/tests/cross_service_tracing.rs create mode 100644 shared/tests/metrics_test.rs create mode 100644 shared/tests/queue_link_test.rs create mode 100644 shared/tests/telemetry_test.rs create mode 100644 web/src/instrumentation.mjs create mode 100644 web/src/lib/server/otel-log-hook.mjs create mode 100644 web/src/lib/server/otel-log-hook.test.ts create mode 100644 web/src/lib/server/otel-runtime.test.ts create mode 100644 web/src/lib/server/otel-source-sanitization.test.ts create mode 100644 web/src/lib/server/otel-utils.mjs create mode 100644 web/src/lib/server/telemetry.test.ts create mode 100644 web/src/otel-factory.mjs diff --git a/.env.example b/.env.example index 6b18d943b..a7476891b 100644 --- a/.env.example +++ b/.env.example @@ -212,8 +212,24 @@ ENCRYPTION_KEY=your-encryption-key-must-be-at-least-32-characters-long ENCRYPTION_SALT=your-salt-16-chars # OpenTelemetry Configuration -# Leave OTEL_EXPORTER_OTLP_ENDPOINT empty for local-only telemetry +# +# Omni does not bundle or require an OpenTelemetry Collector or backend. +# Point OTEL_EXPORTER_OTLP_ENDPOINT at any compatible OTLP/HTTP collector or +# backend (e.g. http://my-collector:4318). When left empty, instrumentation +# is active but signals are not exported and are dropped. +# +# Example with an external collector: +# OTEL_EXPORTER_OTLP_ENDPOINT=http://my-otel-collector:4318 OTEL_EXPORTER_OTLP_ENDPOINT= +# Set OTEL_DEPLOYMENT_ID to one shared non-empty value (e.g. deploy-YYYY-MM-DD) +# so all services report under the same deployment identity. +# When empty, Rust and Python export an empty deployment.id while Node uses +# "unknown". When the variable is absent outside Compose, Rust and Python +# generate a ULID. Set it explicitly for consistent cross-service identity. OTEL_DEPLOYMENT_ID= OTEL_DEPLOYMENT_ENVIRONMENT=production SERVICE_VERSION=0.1.0 + +# Metric export interval in milliseconds (default: 60000). Applied when an OTLP +# endpoint is configured. +# OTEL_METRIC_EXPORT_INTERVAL=60000 diff --git a/Cargo.toml b/Cargo.toml index 361c95e32..42b23af20 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -62,11 +62,12 @@ futures-util = "0.3" # OpenTelemetry dependencies opentelemetry = "0.32" -opentelemetry_sdk = { version = "0.32.1", features = ["rt-tokio"] } -opentelemetry-otlp = { version = "0.32", features = ["http-proto", "reqwest-client", "trace", "metrics"] } +opentelemetry_sdk = { version = "0.32.1", features = ["rt-tokio", "metrics", "logs", "experimental_trace_batch_span_processor_with_async_runtime", "experimental_logs_batch_log_processor_with_async_runtime"] } +opentelemetry-otlp = { version = "0.32", features = ["http-proto", "http-json", "reqwest-client", "trace", "metrics", "logs"] } opentelemetry-semantic-conventions = "0.32" tracing-opentelemetry = "0.33" opentelemetry-http = "0.32" +opentelemetry-appender-tracing = "0.32" # Test dependencies testcontainers = { version = "0.21", features = ["watchdog"] } diff --git a/README.md b/README.md index 9c4eeece4..297eefa99 100644 --- a/README.md +++ b/README.md @@ -47,6 +47,11 @@ The agent runtime can execute code in a sandboxed container on an isolated Docke See the full [architecture documentation](https://docs.getomni.co/architecture) for more details. +OpenTelemetry traces, metrics, and logs are exported from the five core +services (web, AI, searcher, indexer, connector-manager). +See [docs/observability.md](docs/observability.md) for signal flow, +metric inventory, and endpoint configuration. + ## Deployment Omni can be deployed entirely on your own infra. See our deployment guides: diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 0d0d9ee05..bf18facd9 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -17,6 +17,7 @@ x-otel-config: &otel-config OTEL_DEPLOYMENT_ID: ${OTEL_DEPLOYMENT_ID} OTEL_DEPLOYMENT_ENVIRONMENT: ${OTEL_DEPLOYMENT_ENVIRONMENT:-development} SERVICE_VERSION: ${SERVICE_VERSION:-0.1.0} + OTEL_METRIC_EXPORT_INTERVAL: ${OTEL_METRIC_EXPORT_INTERVAL:-60000} x-storage-config: &storage-config STORAGE_BACKEND: ${STORAGE_BACKEND} diff --git a/services/ai/db/embedding_queue.py b/services/ai/db/embedding_queue.py index f7586e030..fd4efcaf8 100644 --- a/services/ai/db/embedding_queue.py +++ b/services/ai/db/embedding_queue.py @@ -3,7 +3,7 @@ import logging from enum import StrEnum from typing import Optional, List -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime from asyncpg import Pool @@ -30,6 +30,8 @@ class EmbeddingQueueItem: error_message: Optional[str] retry_count: int created_at: datetime + traceparent: Optional[str] = field(default=None) + tracestate: Optional[str] = field(default=None) class EmbeddingQueueRepository: @@ -50,7 +52,8 @@ async def get_by_id(self, item_id: str) -> Optional[EmbeddingQueueItem]: row = await pool.fetchrow( """ - SELECT id, document_id, status, error_message, retry_count, created_at + SELECT id, document_id, status, error_message, retry_count, created_at, + traceparent, tracestate FROM embedding_queue WHERE id = $1 """, @@ -107,7 +110,8 @@ async def get_pending_items( LIMIT $2 FOR UPDATE SKIP LOCKED ) - RETURNING id, document_id, status, error_message, retry_count, created_at + RETURNING id, document_id, status, error_message, retry_count, created_at, + traceparent, tracestate """, max_retries, limit, diff --git a/services/ai/email_service/sender.py b/services/ai/email_service/sender.py index cf9ff6989..2f43d8e64 100644 --- a/services/ai/email_service/sender.py +++ b/services/ai/email_service/sender.py @@ -61,7 +61,7 @@ async def send( logger.error("ACS send failed: status=%s", result["status"]) return SendResult(success=False, error=f"Send failed: {result['status']}") except Exception as e: - logger.error("ACS send error: %s", e) + logger.error("ACS send error") return SendResult(success=False, error="Failed to send email via ACS") @@ -94,10 +94,10 @@ async def send( data = resp.json() return SendResult(success=True, message_id=data.get("id")) - logger.error("Resend error: status=%d body=%s", resp.status_code, resp.text) + logger.error("Resend error: status=%d", resp.status_code) return SendResult(success=False, error="Failed to send email via Resend") except Exception as e: - logger.error("Resend send error: %s", e) + logger.error("Resend send error") return SendResult(success=False, error="Failed to send email via Resend") @@ -135,7 +135,7 @@ async def send( return SendResult(success=True) except Exception as e: - logger.error("SMTP send error: %s", e) + logger.error("SMTP send error") return SendResult(success=False, error="Failed to send email via SMTP") diff --git a/services/ai/embeddings/batch_processor.py b/services/ai/embeddings/batch_processor.py index b3fae08f9..d965850ca 100644 --- a/services/ai/embeddings/batch_processor.py +++ b/services/ai/embeddings/batch_processor.py @@ -8,9 +8,14 @@ import asyncio import logging import time -from typing import Optional +from typing import Optional, Sequence import ulid +from opentelemetry import trace, metrics +from opentelemetry.propagate import get_global_textmap +from opentelemetry.propagators.textmap import TextMapPropagator +from opentelemetry.trace import SpanContext, NonRecordingSpan, Link, SpanKind +from opentelemetry.trace.propagation import get_current_span as otel_get_current_span from config import EMBEDDING_MAX_MODEL_LEN from db import ( @@ -29,6 +34,143 @@ logger = logging.getLogger(__name__) +# --------------------------------------------------------------------------- +# Trace context helpers for the embedding consumer +# --------------------------------------------------------------------------- + + +def _build_carrier( + traceparent: Optional[str], + tracestate: Optional[str], +) -> dict[str, str]: + """Build a W3C TraceContext carrier dict from stored optional headers.""" + carrier: dict[str, str] = {} + if traceparent: + # Validate rough W3C format before using + if len(traceparent) == 55 and traceparent.startswith("00-"): + carrier["traceparent"] = traceparent + if tracestate: + carrier["tracestate"] = tracestate + return carrier + + +def _extract_span_context( + traceparent: Optional[str], + tracestate: Optional[str], +) -> Optional[SpanContext]: + """Extract a valid remote SpanContext from optional stored headers. + + Returns None for invalid or missing context. Processing continues + regardless. + """ + carrier = _build_carrier(traceparent, tracestate) + if not carrier: + return None + + from opentelemetry.context import Context + # Get the global propagator (typically TraceContextTextMapPropagator) + propagator = get_global_textmap() + # Extract into a fresh context (not the current one) + ctx = propagator.extract(carrier=carrier, context=Context()) + # Get the span from the extracted context + span = trace.get_current_span(ctx) + sc = span.get_span_context() + if sc and sc.is_valid: + return sc + return None + + +def _build_consumer_links( + items: Sequence[EmbeddingQueueItem], +) -> list[Link]: + """Build OTel Links from stored producer contexts in queue items. + + Returns one Link per unique valid producer (trace_id, span_id) pair. + Invalid or missing contexts are skipped. An empty list is returned + when no valid contexts exist. + """ + links: list[Link] = [] + seen: set[tuple[int, int]] = set() + + for item in items: + sc = _extract_span_context(item.traceparent, item.tracestate) + if sc is None: + continue + trace_id_int = int(sc.trace_id, 16) if isinstance(sc.trace_id, str) else sc.trace_id + span_id_int = int(sc.span_id, 16) if isinstance(sc.span_id, str) else sc.span_id + key = (trace_id_int, span_id_int) + if key in seen: + continue + seen.add(key) + links.append(Link(context=sc)) + + return links + + +# --------------------------------------------------------------------------- +# Embedding processor metrics (lazily created once) +# --------------------------------------------------------------------------- +_EMBEDDING_METER = None +_EMBEDDING_PENDING = None +_EMBEDDING_PROCESSED = None +_EMBEDDING_FAILED = None +_EMBEDDING_BATCH_DURATION = None + + +def _ensure_embedding_instruments(): + global _EMBEDDING_METER, _EMBEDDING_PENDING, _EMBEDDING_PROCESSED + global _EMBEDDING_FAILED, _EMBEDDING_BATCH_DURATION + if _EMBEDDING_METER is not None: + return + meter = metrics.get_meter("omni-ai-embedding") + _EMBEDDING_METER = meter + _EMBEDDING_PENDING = meter.create_gauge( + name="omni.ai.embedding.pending", + unit="{documents}", + description="Number of documents pending embedding", + ) + _EMBEDDING_PROCESSED = meter.create_counter( + name="omni.ai.embedding.processed", + unit="{documents}", + description="Total documents successfully embedded", + ) + _EMBEDDING_FAILED = meter.create_counter( + name="omni.ai.embedding.failed", + unit="{documents}", + description="Total documents that failed embedding", + ) + _EMBEDDING_BATCH_DURATION = meter.create_histogram( + name="omni.ai.embedding.batch.duration", + unit="s", + description="Duration of embedding batch processing", + ) + + +# --------------------------------------------------------------------------- +# Thin recording helpers (testable via monkeypatching the globals above) +# --------------------------------------------------------------------------- + + +def _record_embedding_processed(count: int = 1) -> None: + if _EMBEDDING_PROCESSED is not None: + _EMBEDDING_PROCESSED.add(count) + + +def _record_embedding_failed(count: int = 1) -> None: + if _EMBEDDING_FAILED is not None: + _EMBEDDING_FAILED.add(count) + + +def _record_embedding_batch_duration(seconds: float) -> None: + if _EMBEDDING_BATCH_DURATION is not None: + _EMBEDDING_BATCH_DURATION.record(seconds) + + +def _record_embedding_pending(count: int) -> None: + if _EMBEDDING_PENDING is not None: + _EMBEDDING_PENDING.set(count) + + # Configuration for online processing ONLINE_BATCH_SIZE = 10 ONLINE_POLL_INTERVAL = 5 # Seconds to wait when queue is empty @@ -80,12 +222,14 @@ async def processing_loop(self): """Process queue items using online API calls""" logger.info(f"Starting embedding processor for provider: {self.provider_type}") + _ensure_embedding_instruments() + status_counts = await self.queue_repo.get_status_counts() self._baseline_completed = status_counts.get(QueueStatus.COMPLETED, 0) self._baseline_failed = status_counts.get(QueueStatus.FAILED, 0) - pending = status_counts.get(QueueStatus.PENDING, 0) + status_counts.get( - QueueStatus.PROCESSING, 0 - ) + pending = status_counts.get(QueueStatus.PENDING, 0) + _record_embedding_pending(pending) + logger.info( f"Embedding queue: {pending} pending, " f"{self._baseline_completed} completed, " @@ -101,8 +245,8 @@ async def processing_loop(self): # to allow higher-priority tasks (stream requests) to run if processed_any: await asyncio.sleep(ONLINE_BATCH_DELAY) - except Exception as e: - logger.error(f"Online processing loop error: {e}", exc_info=True) + except Exception: + logger.error("Online processing loop error") await asyncio.sleep(10) async def _process_online_batch(self) -> bool: @@ -120,29 +264,69 @@ async def _process_online_batch(self) -> bool: return False logger.info(f"Processing {len(items)} documents via online embedding API") - - documents_by_id = await self.documents_repo.get_by_ids( - [item.document_id for item in items] - ) - items_to_process = await self._clone_same_content_embeddings( - items, documents_by_id - ) - - for item in items_to_process: + batch_start = time.monotonic() + + # Build consumer links from stored producer contexts. + links = _build_consumer_links(items) + + # Import explicit empty Context for new-root consumer span. + from opentelemetry.context import Context + + # Create CONSUMER span as a new root trace with links to producers. + # The span ends when the `with` block exits. + tracer = trace.get_tracer("omni-ai") + had_failure = False + with tracer.start_as_current_span( + "embedding_queue process", + context=Context(), # explicit empty parent = new root trace + kind=SpanKind.CONSUMER, + attributes={ + "messaging.system": "postgresql", + "messaging.destination": "embedding_queue", + "messaging.operation.type": "process", + "messaging.batch.message_count": len(items), + }, + links=links, + ) as span: try: - await self._process_single_document( - item, documents_by_id.get(item.document_id) + documents_by_id = await self.documents_repo.get_by_ids( + [item.document_id for item in items] ) - except Exception as e: - logger.error( - f"Failed to process document {item.document_id}: {e}", exc_info=True + items_to_process = await self._clone_same_content_embeddings( + items, documents_by_id ) - await self.queue_repo.mark_failed([item.id], str(e)) - self._docs_failed += 1 + + for item in items_to_process: + try: + ok = await self._process_single_document( + item, documents_by_id.get(item.document_id) + ) + if not ok: + had_failure = True + except Exception as exc: + had_failure = True + logger.error( + "Failed to process document" + ) + await self.queue_repo.mark_failed([item.id], str(exc)) + self._docs_failed += 1 + _record_embedding_failed() + finally: + # Yield to allow higher-priority tasks (stream requests) to run + await asyncio.sleep(0) + await self._maybe_log_progress() + + # Record batch duration + batch_duration = time.monotonic() - batch_start + _record_embedding_batch_duration(batch_duration) + + except Exception: + had_failure = True + span.set_status(trace.Status(trace.StatusCode.ERROR)) + raise finally: - # Yield to allow higher-priority tasks (stream requests) to run - await asyncio.sleep(0) - await self._maybe_log_progress() + if had_failure: + span.set_status(trace.Status(trace.StatusCode.ERROR)) return True @@ -191,6 +375,7 @@ async def _clone_same_content_embeddings( self._docs_completed += len(clone_counts) self._embeddings_written += sum(clone_counts.values()) + _record_embedding_processed(len(clone_counts)) logger.info( "Cloned embeddings for %d documents with duplicate content (%d chunks)", len(clone_counts), @@ -202,11 +387,18 @@ async def _clone_same_content_embeddings( async def _process_single_document( self, item: EmbeddingQueueItem, doc: Document | None = None - ): - """Process a single document using the embedding provider""" + ) -> bool: + """Process a single document using the embedding provider. + + Returns: + True if the document was durably completed (embeddings written + or cloned successfully). False if any internally-handled failure + path was taken (missing content_id, empty content, no chunks + generated, or caught embedding exception). + """ if item.retry_count > 0: logger.debug( - f"Retrying document {item.document_id} (attempt {item.retry_count + 1})" + f"Retrying document (attempt {item.retry_count + 1})" ) # Use semaphore to limit concurrent embedding operations and yield more frequently @@ -216,13 +408,14 @@ async def _process_single_document( if not doc or not doc.content_id: logger.warning( - f"Document {item.document_id} has no content_id, skipping" + "Document has no content_id, skipping" ) await self.queue_repo.mark_failed( [item.id], "Document has no content_id" ) self._docs_failed += 1 - return + _record_embedding_failed() + return False # Cross-source embedding cloning: if another document with the same # external_id already has embeddings, clone them instead of @@ -241,22 +434,24 @@ async def _process_single_document( await self.queue_repo.mark_completed([item.id]) self._docs_completed += 1 self._embeddings_written += cloned + _record_embedding_processed() logger.info( - f"Cloned {cloned} embeddings from {donor_id} to {item.document_id} (cross-source dedup)" + f"Cloned {cloned} embeddings (cross-source dedup)" ) - return + return True content_text = await self.content_storage.get_text(doc.content_id) if not content_text or not content_text.strip(): logger.warning( - f"Document {item.document_id} has empty content, skipping" + "Document has empty content, skipping" ) await self.queue_repo.mark_failed( [item.id], "Document has empty content" ) self._docs_failed += 1 - return + _record_embedding_failed() + return False # Generate embeddings using sliding window over the document try: @@ -299,13 +494,14 @@ async def _process_single_document( if not chunks: logger.warning( - f"No embeddings generated for document {item.document_id}" + "No embeddings generated for document" ) await self.queue_repo.mark_failed( [item.id], "No embeddings generated" ) self._docs_failed += 1 - return + _record_embedding_failed() + return False await self.embeddings_repo.delete_for_documents([item.document_id]) @@ -330,17 +526,20 @@ async def _process_single_document( self._docs_completed += 1 self._embeddings_written += len(chunks) + _record_embedding_processed() logger.info( - f"Processed document {item.document_id}: {len(chunks)} chunks embedded" + f"Processed document: {len(chunks)} chunks embedded" ) + return True - except Exception as e: + except Exception as exc: logger.error( - f"Embedding generation failed for {item.document_id}: {e}", - exc_info=True, + "Embedding generation failed" ) - await self.queue_repo.mark_failed([item.id], str(e)) + await self.queue_repo.mark_failed([item.id], str(exc)) self._docs_failed += 1 + _record_embedding_failed() + return False async def _maybe_log_progress(self): """Log embedding progress periodically.""" @@ -358,6 +557,8 @@ async def _maybe_log_progress(self): total_completed = self._baseline_completed + self._docs_completed total_failed = self._baseline_failed + self._docs_failed + _record_embedding_pending(pending) + elapsed_min = (now - self._progress_start_time) / 60 docs_per_min = self._docs_completed / elapsed_min if elapsed_min > 0 else 0 chunks_per_min = ( diff --git a/services/ai/main.py b/services/ai/main.py index 6eca57198..b495b5c1a 100644 --- a/services/ai/main.py +++ b/services/ai/main.py @@ -33,7 +33,7 @@ start_batch_processor, ) from state import AppState -from telemetry import init_telemetry +from telemetry import init_telemetry, shutdown_telemetry setup_logging() logger = logging.getLogger(__name__) @@ -87,21 +87,24 @@ async def startup_event(): raise except Exception as e: app.state.memory_provider = None - logger.warning(f"Memory initialization failed: {e}") + logger.warning(f"Memory initialization failed: {type(e).__name__}") else: app.state.memory_provider = None logger.info("MEMORY_ENABLED=false — memory feature disabled") - except Exception as e: - logger.error(f"Failed to initialize services: {e}") - raise e + except Exception: + logger.error("Failed to initialize services") + raise @app.on_event("shutdown") async def shutdown_event(): """Cleanup on shutdown.""" - if hasattr(app.state, "embedding_queue"): - await app.state.embedding_queue.stop() - await shutdown_providers(app.state) + try: + if hasattr(app.state, "embedding_queue"): + await app.state.embedding_queue.stop() + await shutdown_providers(app.state) + finally: + shutdown_telemetry() if __name__ == "__main__": diff --git a/services/ai/providers/anthropic.py b/services/ai/providers/anthropic.py index 295023478..a523720e6 100644 --- a/services/ai/providers/anthropic.py +++ b/services/ai/providers/anthropic.py @@ -123,18 +123,14 @@ async def stream_response( if tools: request_params["tools"] = tools logger.info( - f"Sending request with {len(tools)} tools: {[t['name'] for t in tools]}" + f"Sending request with {len(tools)} tools" ) else: - logger.info(f"Sending request without tools") + logger.info("Sending request without tools") logger.info( f"Model: {self.model}, Messages: {len(msg_list)}, Max tokens: {request_params['max_tokens']}" ) - logger.debug( - f"Full request params: {json.dumps({k: v for k, v in request_params.items() if k != 'messages'}, indent=2)}" - ) - logger.debug(f"Messages: {json.dumps(msg_list, indent=2)}") if system_prompt: request_params["system"] = system_prompt @@ -142,38 +138,24 @@ async def stream_response( stream: AsyncStream[MessageStreamEvent] = await self.client.messages.create( **request_params, ) - logger.info(f"Stream created successfully, starting to process events") + logger.info("Stream created successfully") event_count = 0 async for event in stream: event_count += 1 - logger.debug(f"Event {event_count}: {event.type}") if event.type == "content_block_start": - logger.info(f"Content block start: type={event.content_block.type}") - if event.content_block.type == "tool_use": - logger.info( - f"Tool use started: {event.content_block.name} (id: {event.content_block.id}) (input: {json.dumps(event.content_block.input)})" - ) - elif event.type == "content_block_delta": - if event.delta.type == "text_delta": - logger.debug(f"Text delta: '{event.delta.text}'") - elif event.delta.type == "input_json_delta": - logger.debug(f"JSON delta: {event.delta.partial_json}") - elif event.type == "citation": - logger.info(f"Citation: {event.citation}") + logger.info("Content block start") elif event.type == "content_block_stop": - logger.info( - f"Content block stop at index {getattr(event, 'index', '')}" - ) + logger.info("Content block stop") elif event.type == "message_delta": - logger.info(f"Message delta stop reason: {event.delta.stop_reason}") + logger.info("Message delta received") elif event.type == "message_stop": - logger.info(f"Message completed after {event_count} events") + logger.info("Message completed") yield event except Exception as e: - logger.error(f"Failed to stream from Anthropic: {str(e)}", exc_info=True) + logger.error("Failed to stream from Anthropic") raise ProviderError( str(e), provider_type=self.provider_type, @@ -222,7 +204,7 @@ async def generate_response( return content, usage except Exception as e: - logger.error(f"Failed to generate response from Anthropic: {str(e)}") + logger.error("Failed to generate response from Anthropic") raise ProviderError( str(e), provider_type=self.provider_type, diff --git a/services/ai/providers/bedrock.py b/services/ai/providers/bedrock.py index 0486471f4..9c393ea6f 100644 --- a/services/ai/providers/bedrock.py +++ b/services/ai/providers/bedrock.py @@ -296,7 +296,7 @@ def _dedupe_documents(self, messages: list[dict[str, Any]]): deduped_tool_result_content.append(content_block) else: logger.debug( - f"[BEDROCK-AMAZON] Deduplicating document '{doc_name}' in tool result" + "[BEDROCK-AMAZON] Deduplicating document in tool result" ) else: deduped_tool_result_content.append(content_block) @@ -317,6 +317,7 @@ def _limit_documents(self, messages: list[dict[str, Any]], max_documents: int = """ document_count = 0 + removed_count = 0 for msg in reversed(messages): if "content" in msg: limited_content = [] @@ -330,9 +331,7 @@ def _limit_documents(self, messages: list[dict[str, Any]], max_documents: int = document_count += 1 limited_tool_result_content.append(content_block) else: - logger.debug( - f"[BEDROCK-AMAZON] Limiting documents to {max_documents}, removing document '{content_block['document']['name']}'" - ) + removed_count += 1 else: limited_tool_result_content.append(content_block) if limited_tool_result_content: @@ -349,6 +348,11 @@ def _limit_documents(self, messages: list[dict[str, Any]], max_documents: int = msg["content"] = limited_content + if removed_count > 0: + logger.debug( + f"[BEDROCK-AMAZON] Document limit ({max_documents}) enforced, kept {document_count}, removed {removed_count}" + ) + def _adapt_tools_for_amazon_models( self, tools: list[dict[str, Any]] ) -> list[dict[str, Any]]: @@ -490,7 +494,7 @@ def _convert_response_to_anthropic_events( elif "messageStop" in event: return RawMessageStopEvent(type="message_stop") - logger.debug(f"[BEDROCK] Skipping unknown event type: {list(event.keys())}") + logger.debug("[BEDROCK] Skipping unknown event type") return None async def stream_response( @@ -522,18 +526,14 @@ async def stream_response( if tools: request_params["tools"] = tools logger.info( - f"[BEDROCK] Sending request with {len(tools)} tools: {[t['name'] for t in tools]}" + f"[BEDROCK] Sending request with {len(tools)} tools" ) else: - logger.info(f"[BEDROCK] Sending request without tools") + logger.info("[BEDROCK] Sending request without tools") logger.info( f"[BEDROCK] Model: {self.model_id}, Messages: {len(msg_list)}, Max tokens: {request_params['max_tokens']}" ) - logger.debug( - f"[BEDROCK] Full request body: {json.dumps({k: v for k, v in request_params.items() if k != 'messages'}, indent=2)}" - ) - logger.debug(f"[BEDROCK] Messages: {json.dumps(msg_list, indent=2)}") # Invoke with streaming response logger.info( @@ -551,30 +551,12 @@ async def stream_response( event_count = 0 for event in stream: event_count += 1 - logger.debug(f"[ANTHROPIC] Event {event_count}: {event.type}") if event.type == "content_block_start": logger.info( f"[ANTHROPIC] Content block start: type={event.content_block.type}" ) - if event.content_block.type == "tool_use": - logger.info( - f"[ANTHROPIC] Tool use started: {event.content_block.name} (id: {event.content_block.id}) (input: {json.dumps(event.content_block.input)})" - ) - elif event.type == "content_block_delta": - if event.delta.type == "text_delta": - logger.debug( - f"[ANTHROPIC] Text delta: '{event.delta.text}'" - ) - elif event.delta.type == "input_json_delta": - logger.debug( - f"[ANTHROPIC] JSON delta: {event.delta.partial_json}" - ) - elif event.type == "citation": - logger.info(f"[ANTHROPIC] Citation: {event.citation}") elif event.type == "content_block_stop": - logger.info( - f"[ANTHROPIC] Content block stop at index {getattr(event, 'index', '')}" - ) + logger.info("[ANTHROPIC] Content block stop") elif event.type == "message_delta": logger.info( f"[ANTHROPIC] Message delta stop reason: {event.delta.stop_reason}" @@ -600,7 +582,7 @@ async def stream_response( self._limit_documents(messages, max_documents=5) logger.debug( - f"[BEDROCK-AMAZON] Adapted messages: {json.dumps(messages, indent=2)}" + f"[BEDROCK-AMAZON] Adapted messages: {len(messages)} messages" ) tools = ( self._adapt_tools_for_amazon_models(tools) if tools else None @@ -624,7 +606,7 @@ async def stream_response( response = self.client.converse_stream(**request_params) - logger.info(f"[BEDROCK-AMAZON] Stream created, processing chunks") + logger.info("[BEDROCK-AMAZON] Stream created, processing chunks") chunk_count = 0 for chunk in response["stream"]: chunk_count += 1 @@ -647,13 +629,12 @@ async def stream_response( except ClientError as e: error_code = e.response.get("Error", {}).get("Code", "Unknown") logger.error( - f"[BEDROCK] AWS Bedrock client error ({error_code}): {str(e)}", - exc_info=True, + "[BEDROCK] AWS Bedrock client error (%s)", error_code, ) raise self._to_provider_error(e, model=model) from e except Exception as e: logger.error( - f"[BEDROCK] Failed to stream from AWS Bedrock: {str(e)}", exc_info=True + "[BEDROCK] Failed to stream from AWS Bedrock" ) raise self._to_provider_error(e, model=model) from e @@ -715,7 +696,7 @@ async def generate_response( "maxTokens": max_tokens or 512, }, ) - logger.debug(f"generate_response: response from LLM -> {response}") + logger.debug("generate_response: response received") usage_data = response.get("usage", {}) usage = TokenUsage( diff --git a/services/ai/providers/gemini.py b/services/ai/providers/gemini.py index 70cb8c77c..924a204b8 100644 --- a/services/ai/providers/gemini.py +++ b/services/ai/providers/gemini.py @@ -237,7 +237,7 @@ async def get_context_window_tokens(self) -> ContextWindowInfo: ) return self._context_window_info except Exception as e: - logger.debug("Failed to fetch Gemini model metadata for %s: %s", self.model, e) + logger.debug("Failed to fetch Gemini model metadata for %s", self.model) self._context_window_info = await super().get_context_window_tokens() return self._context_window_info @@ -268,11 +268,14 @@ async def stream_response( if tools: config.tools = _convert_tools_to_gemini(tools) logger.info( - f"Sending request with {len(tools)} tools: {[t['name'] for t in tools]}" + "Sending request with %s tools", len(tools) ) logger.info( - f"Model: {active_model}, Messages: {len(contents)}, Max tokens: {config.max_output_tokens}" + "Model: %s, Messages: %s, Max tokens: %s", + active_model, + len(contents), + config.max_output_tokens, ) # Emit message_start @@ -401,7 +404,7 @@ async def stream_response( yield RawMessageStopEvent(type="message_stop") except Exception as e: - logger.error(f"Failed to stream from Gemini: {str(e)}", exc_info=True) + logger.error("Failed to stream from Gemini") raise ProviderError( str(e), provider_type=self.provider_type, @@ -446,7 +449,7 @@ async def generate_response( return response.text, usage except Exception as e: - logger.error(f"Failed to generate response: {str(e)}") + logger.error("Failed to generate response") raise ProviderError( str(e), provider_type=self.provider_type, diff --git a/services/ai/providers/openai.py b/services/ai/providers/openai.py index 9fe423672..bc4d70c15 100644 --- a/services/ai/providers/openai.py +++ b/services/ai/providers/openai.py @@ -128,11 +128,14 @@ async def stream_response( if tools: request_params["tools"] = _convert_tools_to_openai(tools) logger.info( - f"Sending request with {len(tools)} tools: {[t['name'] for t in tools]}" + "Sending request with %s tools", len(tools) ) logger.info( - f"Model: {active_model}, Input items: {len(input_items)}, Max tokens: {request_params['max_output_tokens']}" + "Model: %s, Input items: %s, Max tokens: %s", + active_model, + len(input_items), + request_params['max_output_tokens'], ) stream = await self.client.responses.create(**request_params) @@ -188,7 +191,7 @@ async def stream_response( if item.type == "function_call": if item.id is None: logger.warning( - f"Received function_call item with no id (call_id={item.call_id}); skipping tool block" + "Received function_call item with no id; skipping tool block" ) else: block_index = next_block_index @@ -302,7 +305,7 @@ async def stream_response( except ProviderError: raise except Exception as e: - logger.error(f"Failed to stream from OpenAI: {str(e)}", exc_info=True) + logger.error("Failed to stream from OpenAI") raise ProviderError( str(e), provider_type=self.provider_type, @@ -469,7 +472,7 @@ async def generate_response( return content, usage except Exception as e: - logger.error(f"Failed to generate response: {str(e)}") + logger.error("Failed to generate response") raise ProviderError( str(e), provider_type=self.provider_type, diff --git a/services/ai/providers/openai_compatible.py b/services/ai/providers/openai_compatible.py index 3c2fd2b09..6e949d4dd 100644 --- a/services/ai/providers/openai_compatible.py +++ b/services/ai/providers/openai_compatible.py @@ -310,7 +310,7 @@ async def stream_response( if tools: params["tools"] = _convert_tools_to_openai(tools) logger.info( - f"Sending request with {len(tools)} tools: {[t['name'] for t in tools]}" + "Sending request with %s tools", len(tools) ) stream = await self.client.chat.completions.create(**params) @@ -471,8 +471,7 @@ async def stream_response( except Exception as e: logger.error( - f"Failed to stream from OpenAI-compatible endpoint: {e}", - exc_info=True, + "Failed to stream from OpenAI-compatible endpoint", ) raise ProviderError( str(e), @@ -528,7 +527,7 @@ async def generate_response( raise except Exception as e: logger.error( - f"Failed to generate response from OpenAI-compatible endpoint: {e}" + "Failed to generate response from OpenAI-compatible endpoint", ) raise ProviderError( str(e), diff --git a/services/ai/pyproject.toml b/services/ai/pyproject.toml index 6ad272b65..f3d1c9fbf 100644 --- a/services/ai/pyproject.toml +++ b/services/ai/pyproject.toml @@ -26,6 +26,8 @@ dependencies = [ "opentelemetry-exporter-otlp>=1.29.0", "opentelemetry-instrumentation-fastapi>=0.50b0", "opentelemetry-instrumentation-httpx>=0.50b0", + "opentelemetry-sdk>=1.29.0", + "opentelemetry-exporter-otlp-proto-http>=1.29.0", "mem0ai>=2.0.0", "pgvector>=0.4.2", "psycopg[binary,pool]>=3.2", diff --git a/services/ai/streaming/generate.py b/services/ai/streaming/generate.py index 2755d900d..f83194ce8 100644 --- a/services/ai/streaming/generate.py +++ b/services/ai/streaming/generate.py @@ -146,8 +146,7 @@ async def event_stream_with_context_retry( except ProviderError as e: if e.is_context_overflow and llm_attempt == 0 and not emitted_event: logger.warning( - "Chat %s hit provider context limit; retrying once after forced compaction", - chat_id, + "Chat hit provider context limit; retrying once after forced compaction" ) should_emit_progress = ( compactor.select_legacy_compaction_split(conversation_messages) @@ -685,7 +684,7 @@ async def stream_generator( loop_passes = AGENT_MAX_ITERATIONS + (1 if resumable_tool_calls else 0) for _ in range(loop_passes): if await is_run_cancelled(redis_client, chat_id): - logger.info(f"Run cancelled, stopping stream for chat {chat_id}") + logger.info("Run cancelled, stopping stream") break content_blocks = [] @@ -739,7 +738,7 @@ async def stream_generator( yield sse_event("compaction_end", {}) continue - logger.debug(f"Received event: {event} (index: {event_index})") + logger.debug("Received event (index: %s)", event_index) event_index += 1 now = asyncio.get_running_loop().time() @@ -754,7 +753,7 @@ async def stream_generator( if event.type == "content_block_delta": logger.debug( - f"Content block delta received at index {event.index}: {event.delta}" + "Content block delta (index: %s)", event.index ) if event.delta.type == "text_delta": if event.index >= len(content_blocks): @@ -810,7 +809,7 @@ async def stream_generator( elif event.type == "content_block_start": if event.content_block.type == "text": - logger.info(f"Text block start: {event.content_block.text}") + logger.info("Text block start") text_block: TextBlockParam = TextBlockParam( type="text", text=event.content_block.text ) @@ -820,7 +819,7 @@ async def stream_generator( content_blocks.append(text_block) elif event.content_block.type == "tool_use": logger.info( - f"Tool use block start: {event.content_block.name} (id: {event.content_block.id})" + "Tool use block start" ) tool_block: ToolUseBlockParam = ToolUseBlockParam( type="tool_use", @@ -834,14 +833,12 @@ async def stream_generator( content_blocks.append(tool_block) elif event.type == "citation": - logger.info(f"Citation received: {event.citation}") + logger.info("Citation received") elif event.type == "message_stop": logger.info("Message stop received.") message_stop_received = True - logger.debug( - f"Yielding event to client: {event.to_json(indent=None)}" - ) + logger.debug("Yielding event to client") yield f"event: message\ndata: {event.to_json(indent=None)}\n\n" if message_stop_received: @@ -872,14 +869,14 @@ async def stream_generator( if not tool_calls: logger.info( - f"No tool calls in iteration {model_iteration}, completing response" + "No tool calls in iteration %s, completing response", model_iteration ) break - logger.info(f"Processing {len(tool_calls)} tool calls") + logger.info("Processing %s tool calls", len(tool_calls)) if await is_run_cancelled(redis_client, chat_id): - logger.info(f"Run cancelled before tool execution for chat {chat_id}") + logger.info("Run cancelled before tool execution") break # Preflight credentials before asking for user approval. This keeps @@ -1109,12 +1106,12 @@ async def stream_generator( memory_provider.add(messages=turn, key=memory_write_key) ) except Exception as e: - logger.warning(f"Memory write setup failed for chat {chat_id}: {e}") + logger.warning("Memory write setup failed") yield end_of_stream(EndOfStreamReason.COMPLETED, message="Stream ended") except asyncio.CancelledError: - logger.info(f"Stream cancelled for chat {chat_id}") + logger.info("Stream cancelled") if not content_blocks_finalized: partial = partial_assistant_message(content_blocks) if partial is not None: @@ -1122,7 +1119,7 @@ async def stream_generator( yield f"event: save_message\ndata: {json.dumps(partial)}\n\n" raise except Exception as e: - logger.error(f"Failed to generate AI response with tools: {e}", exc_info=True) + logger.error("Failed to generate AI response with tools") if not content_blocks_finalized: partial = partial_assistant_message(content_blocks) if partial is not None: diff --git a/services/ai/streaming/run.py b/services/ai/streaming/run.py index d8c215f69..9c340f8b6 100644 --- a/services/ai/streaming/run.py +++ b/services/ai/streaming/run.py @@ -140,7 +140,7 @@ async def run_producer(redis_client, chat_id, gen, messages_repo, parent_id): pass raise except Exception as e: - logger.error(f"Producer failed for chat {chat_id}: {e}", exc_info=True) + logger.error("Producer failed") try: await redis_client.xadd( sk, diff --git a/services/ai/telemetry.py b/services/ai/telemetry.py index 793cee984..8e2185b19 100644 --- a/services/ai/telemetry.py +++ b/services/ai/telemetry.py @@ -1,7 +1,13 @@ +"""OpenTelemetry initialization for omni-ai (traces + metrics + logs).""" + import logging import os -from opentelemetry import trace +from typing import Optional + +from opentelemetry import trace, metrics from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter +from opentelemetry.exporter.otlp.proto.http.metric_exporter import OTLPMetricExporter +from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter from opentelemetry.instrumentation.fastapi import FastAPIInstrumentor from opentelemetry.instrumentation.httpx import HTTPXClientInstrumentor from opentelemetry.sdk.resources import ( @@ -12,16 +18,108 @@ ) from opentelemetry.sdk.trace import TracerProvider from opentelemetry.sdk.trace.export import BatchSpanProcessor +from opentelemetry.sdk.metrics import MeterProvider +from opentelemetry.sdk.metrics.export import PeriodicExportingMetricReader +from opentelemetry.sdk._logs import LoggerProvider as LogLoggerProvider +from opentelemetry.sdk._logs.export import BatchLogRecordProcessor +from opentelemetry.sdk._logs import LoggingHandler as OTelLoggingHandler from ulid import ULID logger = logging.getLogger(__name__) +# Hold references so we can shut them down gracefully. +_tracer_provider: TracerProvider | None = None +_meter_provider: MeterProvider | None = None +_logger_provider: LogLoggerProvider | None = None +_otel_log_handler: OTelLoggingHandler | None = None + + +# --------------------------------------------------------------------------- +# Production helpers +# --------------------------------------------------------------------------- + + +def normalize_otlp_endpoint(endpoint_raw: str | None) -> str | None: + """Normalise an OTLP endpoint: strip trailing slash, treat empty as None.""" + if endpoint_raw: + return endpoint_raw.rstrip("/") + return None + + +def parse_metric_export_interval() -> int: + """Parse OTEL_METRIC_EXPORT_INTERVAL as a finite positive integer + (milliseconds). Falls back to 60000 (60 s) when unset or invalid.""" + raw = os.environ.get("OTEL_METRIC_EXPORT_INTERVAL") + if raw is not None: + try: + val = int(raw) + if val > 0: + return val + except (ValueError, TypeError): + pass + return 60_000 + + +def _build_otlp_logs_url(endpoint: str) -> str: + """Build the OTLP HTTP logs export URL from a base endpoint.""" + return f"{endpoint}/v1/logs" + + +class TraceContextFilter(logging.Filter): + """Logging filter that adds bounded ``trace_id`` and ``span_id`` fields + to every ``LogRecord``. + + Inside an active OTel span the real trace/span IDs are recorded. + Outside any span the fields are ``"00000000000000000000000000000000"`` + (trace_id) / ``"0000000000000000"`` (span_id). This prevents formatter + KeyErrors before telemetry initialisation. + """ + + def filter(self, record: logging.LogRecord) -> bool: + span = trace.get_current_span() + span_context = span.get_span_context() + + if span_context and span_context.is_valid: + record.trace_id = span_context.trace_id + record.span_id = span_context.span_id + else: + record.trace_id = "00000000000000000000000000000000" + record.span_id = "0000000000000000" + return True + def init_telemetry(app, service_name: str = "omni-ai"): """ Initialize OpenTelemetry instrumentation for the FastAPI application. + + Idempotent: if called a second time, the prior OTel LoggingHandler is + removed and the prior logger provider is shut down before creating a new + one. This prevents handler accumulation during repeated test init. """ - otlp_endpoint = os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT") + global _tracer_provider, _meter_provider, _logger_provider, _otel_log_handler + + # --- Idempotent cleanup: remove prior OTel LoggingHandler --- + root_logger = logging.getLogger() + if _otel_log_handler is not None: + root_logger.removeHandler(_otel_log_handler) + _otel_log_handler = None + # Remove any remaining OTel LoggingHandlers that might have been + # installed by a prior init without going through our global. + for h in list(root_logger.handlers): + if isinstance(h, OTelLoggingHandler): + root_logger.removeHandler(h) + + # Shut down prior logger provider (if any) so we don't leak one + # per init call. Do NOT replace global trace/meter providers here + # — those are set once and shared. + if _logger_provider is not None: + try: + _logger_provider.shutdown() + except Exception: + pass + _logger_provider = None + + otlp_endpoint = normalize_otlp_endpoint(os.getenv("OTEL_EXPORTER_OTLP_ENDPOINT")) deployment_id = os.getenv("OTEL_DEPLOYMENT_ID", str(ULID())) environment = os.getenv("OTEL_DEPLOYMENT_ENVIRONMENT", "development") service_version = os.getenv("SERVICE_VERSION", "0.1.0") @@ -36,24 +134,93 @@ def init_telemetry(app, service_name: str = "omni-ai"): } ) - # Create tracer provider - provider = TracerProvider(resource=resource) + # ------------------------------------------------------------------ + # Tracer provider + # ------------------------------------------------------------------ + tracer_provider = TracerProvider(resource=resource) - # Add OTLP exporter if endpoint is configured if otlp_endpoint: logger.info(f"Initializing OpenTelemetry with OTLP endpoint: {otlp_endpoint}") otlp_exporter = OTLPSpanExporter(endpoint=f"{otlp_endpoint}/v1/traces") processor = BatchSpanProcessor(otlp_exporter) - provider.add_span_processor(processor) + tracer_provider.add_span_processor(processor) else: logger.info( - "No OTLP endpoint configured, telemetry will be collected locally only" + "No OTLP endpoint configured, traces will be collected locally only" ) - # Set the global tracer provider - trace.set_tracer_provider(provider) + trace.set_tracer_provider(tracer_provider) + _tracer_provider = tracer_provider + + # ------------------------------------------------------------------ + # Meter provider (metrics) + # ------------------------------------------------------------------ + if otlp_endpoint: + metric_exporter = OTLPMetricExporter(endpoint=f"{otlp_endpoint}/v1/metrics") + metric_reader = PeriodicExportingMetricReader( + metric_exporter, + export_interval_millis=parse_metric_export_interval(), + export_timeout_millis=30_000, + ) + else: + metric_reader = None + logger.info( + "No OTLP endpoint configured, metrics will be collected locally only" + ) + if metric_reader: + meter_provider = MeterProvider( + resource=resource, + metric_readers=[metric_reader], + ) + else: + meter_provider = MeterProvider(resource=resource) + + metrics.set_meter_provider(meter_provider) + _meter_provider = meter_provider + + # ------------------------------------------------------------------ + # Logger provider (logs) + # ------------------------------------------------------------------ + if otlp_endpoint: + log_exporter = OTLPLogExporter( + endpoint=_build_otlp_logs_url(otlp_endpoint) + ) + logger_provider = LogLoggerProvider(resource=resource) + logger_provider.add_log_record_processor( + BatchLogRecordProcessor(log_exporter) + ) + else: + logger_provider = LogLoggerProvider(resource=resource) + + _logger_provider = logger_provider + + # Install a LoggingHandler that bridges stdlib logging to OTel logs, + # using the same resource and logger provider. A TraceContextFilter + # ensures that every LogRecord carries bounded trace_id / span_id + # fields. + otel_handler = OTelLoggingHandler( + level=logging.NOTSET, + logger_provider=logger_provider, + ) + otel_handler.addFilter(TraceContextFilter()) + _otel_log_handler = otel_handler + + # Attach the OTel handler to the root logger. Avoid duplicate handlers + # when init is called repeatedly (tests). + root_logger = logging.getLogger() + if otel_handler not in root_logger.handlers: + root_logger.addHandler(otel_handler) + + # Attach TraceContextFilter to all existing root handlers so stdout + # records always have trace_id/span_id fields. + for h in root_logger.handlers: + if not any(isinstance(f, TraceContextFilter) for f in h.filters): + h.addFilter(TraceContextFilter()) + + # ------------------------------------------------------------------ # Instrument FastAPI + # ------------------------------------------------------------------ FastAPIInstrumentor.instrument_app(app) # Instrument HTTPX (for outbound HTTP requests) @@ -65,8 +232,61 @@ def init_telemetry(app, service_name: str = "omni-ai"): ) +def shutdown_telemetry(): + """Shut down and flush the tracer, meter, and logger providers. + + Removes every OTel LoggingHandler created by this module so the root + logger is left with zero handlers from this init. + """ + global _tracer_provider, _meter_provider, _logger_provider, _otel_log_handler + + # Remove every OTel LoggingHandler from the root logger so we leave + # zero handlers behind, regardless of how many inits were called. + root_logger = logging.getLogger() + if _otel_log_handler is not None: + root_logger.removeHandler(_otel_log_handler) + _otel_log_handler = None + for h in list(root_logger.handlers): + if isinstance(h, OTelLoggingHandler): + root_logger.removeHandler(h) + + # Shut down meter provider first (flush pending metric exports) + if _meter_provider is not None: + logger.info("Shutting down OpenTelemetry meter provider") + try: + _meter_provider.shutdown() + except Exception: + logger.exception("Error shutting down meter provider") + _meter_provider = None + + if _tracer_provider is not None: + logger.info("Shutting down OpenTelemetry tracer provider") + try: + _tracer_provider.shutdown() + except Exception: + logger.exception("Error shutting down tracer provider") + _tracer_provider = None + + # Shut down logger provider last so final correlated log records are + # exported before the provider is torn down. + if _logger_provider is not None: + logger.info("Shutting down OpenTelemetry logger provider") + try: + _logger_provider.shutdown() + except Exception: + logger.exception("Error shutting down logger provider") + _logger_provider = None + + def get_tracer(name: str = "omni-ai"): """ Get a tracer instance for manual instrumentation. """ return trace.get_tracer(name) + + +def get_meter(name: str = "omni-ai"): + """ + Get a meter instance for manual instrumentation. + """ + return metrics.get_meter(name) diff --git a/services/ai/tests/unit/test_queue_tracing.py b/services/ai/tests/unit/test_queue_tracing.py new file mode 100644 index 000000000..8bd64abef --- /dev/null +++ b/services/ai/tests/unit/test_queue_tracing.py @@ -0,0 +1,538 @@ +"""Queue tracing tests for the embedding batch processor. + +Uses InMemorySpanExporter to verify that CONSUMER spans carry native links +to their PRODUCER contexts, that failure scenarios set ERROR status, and +that invalid/missing contexts are ignored. +""" + +from datetime import datetime, timezone +from unittest.mock import AsyncMock, MagicMock + +import pytest +from opentelemetry import trace +from opentelemetry.context import Context +from opentelemetry.propagate import set_global_textmap +from opentelemetry.sdk.trace import TracerProvider +from opentelemetry.sdk.trace.export import SimpleSpanProcessor +from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter +from opentelemetry.trace import SpanKind, StatusCode +from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator + +from db.embedding_queue import EmbeddingQueueItem +from embeddings.batch_processor import ( + _build_carrier, + _extract_span_context, + _build_consumer_links, + EmbeddingBatchProcessor, +) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +# --------------------------------------------------------------------------- +# Module-level shared tracer provider and exporter. +# Initialised at import time so `trace.set_tracer_provider` is called BEFORE +# any other test file can lock it (the Python OTel API only allows one set). +# --------------------------------------------------------------------------- +set_global_textmap(TraceContextTextMapPropagator()) +_SHARED_EXPORTER = InMemorySpanExporter() +_SHARED_PROVIDER = TracerProvider() +_SHARED_PROVIDER.add_span_processor(SimpleSpanProcessor(_SHARED_EXPORTER)) +trace.set_tracer_provider(_SHARED_PROVIDER) + + +@pytest.fixture(autouse=True) +def _reset_tracing(): + """Reset batcher instrument globals and exporter before each test.""" + _SHARED_EXPORTER.clear() + # Reset the batcher's instrument globals so each test starts clean + import embeddings.batch_processor as bp + + bp._EMBEDDING_METER = None + bp._EMBEDDING_PENDING = None + bp._EMBEDDING_PROCESSED = None + bp._EMBEDDING_FAILED = None + bp._EMBEDDING_BATCH_DURATION = None + yield + + +@pytest.fixture +def span_exporter(): + """Return the shared InMemorySpanExporter.""" + return _SHARED_EXPORTER + + +def _make_processor(span_exporter) -> EmbeddingBatchProcessor: + """Build a processor with all repos mocked.""" + docs_repo = AsyncMock() + queue_repo = AsyncMock() + embeddings_repo = AsyncMock() + app_state = MagicMock() + # Use MagicMock for the embedding provider (sync methods like get_model_name) + # with specific AsyncMock for async methods like generate_embeddings. + provider = MagicMock() + provider.get_model_name.return_value = "test-model" + provider.generate_embeddings = AsyncMock( + return_value=[MagicMock(span=(0, 4), embedding=[0.1])] + ) + app_state.embedding_provider = provider + app_state.embedding_provider_type = "test" + app_state.content_storage = AsyncMock() + + # Default find_embedded_content_donors to empty dict so that + # _clone_same_content_embeddings always short-circuits and never + # produces unawaited AsyncMock coroutine warnings. + docs_repo.find_embedded_content_donors = AsyncMock(return_value={}) + + processor = EmbeddingBatchProcessor( + documents_repo=docs_repo, + queue_repo=queue_repo, + embeddings_repo=embeddings_repo, + app_state=app_state, + ) + processor._baseline_completed = 0 + processor._baseline_failed = 0 + return processor + + +def _make_item( + document_id: str, + traceparent: str | None = None, + tracestate: str | None = None, +) -> EmbeddingQueueItem: + return EmbeddingQueueItem( + id=f"item-{document_id}", + document_id=document_id, + status="processing", + error_message=None, + retry_count=0, + created_at=datetime.now(timezone.utc), + traceparent=traceparent, + tracestate=tracestate, + ) + + +def _find_span(spans, name: str, kind: SpanKind): + """Find a finished span by name and kind.""" + for s in spans: + if s.name == name and s.kind == kind: + return s + return None + + +def _span_context_from_traceparent(tp: str): + """Extract a SpanContext from a traceparent for link assertions.""" + sc = _extract_span_context(tp, None) + assert sc is not None, f"could not extract from {tp}" + return sc + + +def _flush_exporter(): + """Force-flush the global tracer provider if available.""" + tracer_provider = trace.get_tracer_provider() + if hasattr(tracer_provider, "force_flush"): + tracer_provider.force_flush() + + +# --------------------------------------------------------------------------- +# Pure helpers +# --------------------------------------------------------------------------- + + +class TestBuildCarrier: + def test_valid_traceparent(self): + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + carrier = _build_carrier(tp, None) + assert carrier["traceparent"] == tp + assert "tracestate" not in carrier + + def test_with_tracestate(self): + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + carrier = _build_carrier(tp, "vendor1=value1") + assert carrier["traceparent"] == tp + assert carrier["tracestate"] == "vendor1=value1" + + def test_invalid_format_returns_empty(self): + assert _build_carrier("invalid", None) == {} + assert _build_carrier("", None) == {} + assert _build_carrier(None, None) == {} + assert _build_carrier("00-too-short", None) == {} + + +class TestExtractSpanContext: + def test_valid(self): + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + sc = _extract_span_context(tp, None) + assert sc is not None + assert sc.is_valid + + def test_invalid_returns_none(self): + assert _extract_span_context(None, None) is None + assert _extract_span_context("invalid", None) is None + assert _extract_span_context("", None) is None + + def test_all_zero_returns_invalid(self): + tp = "00-00000000000000000000000000000000-0000000000000000-01" + sc = _extract_span_context(tp, None) + assert sc is None or not sc.is_valid + + def test_ambient_span_not_leaked_when_malformed(self): + """Active ambient span must NOT leak into extract with malformed traceparent.""" + tracer = trace.get_tracer("test") + with tracer.start_as_current_span("ambient"): + # Exactly 55 chars, valid format but non-hex trace id + malformed = "00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-b7ad6b7169203331-01" + assert len(malformed) == 55 + sc = _extract_span_context(malformed, None) + assert sc is None, "must not return ambient context for non-hex traceparent" + + def test_ambient_span_not_leaked_when_valid(self): + """Active ambient span must NOT replace the carrier context on valid extract.""" + tracer = trace.get_tracer("test") + with tracer.start_as_current_span("ambient"): + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + sc = _extract_span_context(tp, None) + assert sc is not None + assert sc.is_valid + expected_tid = int("0af7651916cd43dd8448eb211c80319c", 16) + assert sc.trace_id == expected_tid, \ + "trace_id must come from carrier, not ambient span" + + +class TestBuildConsumerLinks: + def test_single_valid_item(self): + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + items = [_make_item("doc1", traceparent=tp)] + links = _build_consumer_links(items) + assert len(links) == 1 + expected_tid = int("0af7651916cd43dd8448eb211c80319c", 16) + assert links[0].context.trace_id == expected_tid + + def test_deduplicates_same_context(self): + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + items = [ + _make_item("doc1", traceparent=tp), + _make_item("doc2", traceparent=tp), + ] + links = _build_consumer_links(items) + assert len(links) == 1 + + def test_same_trace_different_spans_retained(self): + tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + tp2 = "00-0af7651916cd43dd8448eb211c80319c-aaaaaaaaaaaaaaaa-01" + items = [ + _make_item("doc1", traceparent=tp1), + _make_item("doc2", traceparent=tp2), + ] + links = _build_consumer_links(items) + assert len(links) == 2 + + def test_same_trace_three_items_mixed(self): + """Three items: two sharing one span, one unique, one invalid.""" + tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + tp2 = "00-0af7651916cd43dd8448eb211c80319c-aaaaaaaaaaaaaaaa-01" + items = [ + _make_item("doc1", traceparent=tp1), + _make_item("doc2", traceparent=tp2), + _make_item("doc3", traceparent=tp1), # duplicate of doc1 + _make_item("doc4", traceparent=None), # invalid + ] + links = _build_consumer_links(items) + assert len(links) == 2 + + def test_invalid_skipped(self): + items = [ + _make_item("doc1", traceparent=None), + _make_item("doc2", traceparent="invalid"), + ] + links = _build_consumer_links(items) + assert len(links) == 0 + + +# --------------------------------------------------------------------------- +# Batch processor tracing integration +# --------------------------------------------------------------------------- + + +class TestProcessOnlineBatchTracing: + """Trace behaviour of _process_online_batch via in-memory exporter.""" + + # -- Helpers to set up common mocks -------------------------------- + + @staticmethod + def _setup_happy_path(processor, items): + """Configure mocks for a successful processing path.""" + processor.queue_repo.get_pending_items.return_value = items + # Create realistic Document mocks: content_id set, external_id=None to + # avoid triggering cross-source dedup path. + docs = {} + for item in items: + doc = MagicMock() + doc.content_id = f"c-{item.document_id}" + doc.external_id = None + docs[item.document_id] = doc + processor.documents_repo.get_by_ids.return_value = docs + processor.content_storage.get_text = AsyncMock(return_value="some text content") + processor.embeddings_repo.bulk_insert = AsyncMock() + # Ensure _clone_same_content_embeddings returns items as-is + processor.documents_repo.find_embedded_content_donors = AsyncMock(return_value={}) + + @pytest.mark.asyncio + async def test_consumer_span_is_new_root(self, span_exporter): + """CONSUMER span has no parent (new root trace).""" + processor = _make_processor(span_exporter) + + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + self._setup_happy_path(processor, [_make_item("doc1", traceparent=tp)]) + + result = await processor._process_online_batch() + assert result is True + + _flush_exporter() + spans = span_exporter.get_finished_spans() + span_exporter.clear() + + consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) + assert consumer is not None, "CONSUMER span must be exported" + + # Verify parent is None (new root trace via explicit Context()) + assert consumer.parent is None, ( + "CONSUMER span must have no parent (new root)" + ) + + # Verify consumer trace ID differs from producer trace ID + link = consumer.links[0] + assert consumer.context.trace_id != link.context.trace_id, ( + "CONSUMER must have different trace than producer" + ) + + @pytest.mark.asyncio + async def test_native_link_matches_producer(self, span_exporter): + """CONSUMER span has a native link matching the producer traceparent.""" + processor = _make_processor(span_exporter) + + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + self._setup_happy_path(processor, [_make_item("doc1", traceparent=tp)]) + + await processor._process_online_batch() + + _flush_exporter() + spans = span_exporter.get_finished_spans() + span_exporter.clear() + + consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) + assert consumer is not None + assert len(consumer.links) >= 1, "must have at least one link" + + # Verify link matches the expected producer context + expected_sc = _span_context_from_traceparent(tp) + link = consumer.links[0] + assert link.context.trace_id == expected_sc.trace_id, ( + "link trace_id must match producer trace_id" + ) + assert link.context.span_id == expected_sc.span_id, ( + "link span_id must match producer span_id" + ) + assert consumer.parent is None, ( + "CONSUMER must be new root" + ) + + @pytest.mark.asyncio + async def test_same_trace_two_producers_two_links(self, span_exporter): + """Two items sharing same trace but different span IDs create two links.""" + processor = _make_processor(span_exporter) + + tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + tp2 = "00-0af7651916cd43dd8448eb211c80319c-aaaaaaaaaaaaaaaa-01" + item1 = _make_item("doc1", traceparent=tp1) + item2 = _make_item("doc2", traceparent=tp2) + self._setup_happy_path(processor, [item1, item2]) + + await processor._process_online_batch() + + _flush_exporter() + spans = span_exporter.get_finished_spans() + span_exporter.clear() + + consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) + assert consumer is not None + assert len(consumer.links) == 2, "must have 2 links for 2 distinct producers" + + # Both links should have the same trace_id + link_tids = {l.context.trace_id for l in consumer.links} + assert len(link_tids) == 1, "both links must share the same trace_id" + + # But different span_ids + link_sids = {l.context.span_id for l in consumer.links} + assert len(link_sids) == 2, "links must have different span_ids" + + @pytest.mark.asyncio + async def test_invalid_traceparent_ignored(self, span_exporter): + """Items with invalid/missing traceparent do not create links.""" + processor = _make_processor(span_exporter) + + item1 = _make_item("doc1", traceparent=None) + item2 = _make_item("doc2", traceparent="invalid") + self._setup_happy_path(processor, [item1, item2]) + + await processor._process_online_batch() + + _flush_exporter() + spans = span_exporter.get_finished_spans() + span_exporter.clear() + + consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) + if consumer is not None: + assert len(consumer.links) == 0, "no valid producer contexts" + + @pytest.mark.asyncio + async def test_processing_runs_inside_consumer(self, span_exporter): + """Processing operations occur within the CONSUMER span scope.""" + processor = _make_processor(span_exporter) + + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + item = _make_item("doc1", traceparent=tp) + processor.queue_repo.get_pending_items.return_value = [item] + processor.documents_repo.get_by_ids.return_value = { + "doc1": MagicMock(content_id="c1") + } + emb_provider = processor.embedding_provider + emb_provider.get_model_name.return_value = "test-model" + + # Track which span is active during processing + captured_span_context = None + + async def track_span(item, doc): + nonlocal captured_span_context + current_span = trace.get_current_span() + sc = current_span.get_span_context() + captured_span_context = sc + + processor._process_single_document = track_span + + await processor._process_online_batch() + + _flush_exporter() + spans = span_exporter.get_finished_spans() + span_exporter.clear() + + consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) + assert consumer is not None + + # Verify the active span during processing matches the consumer span + assert captured_span_context is not None + assert captured_span_context.trace_id == consumer.context.trace_id, ( + "trace_id must match consumer span" + ) + + @pytest.mark.asyncio + async def test_item_failure_sets_error_status(self, span_exporter): + """When an item fails, the CONSUMER span has ERROR status.""" + processor = _make_processor(span_exporter) + + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + item = _make_item("doc1", traceparent=tp) + processor.queue_repo.get_pending_items.return_value = [item] + processor.documents_repo.get_by_ids.return_value = { + "doc1": MagicMock(content_id="c1") + } + emb_provider = processor.embedding_provider + emb_provider.get_model_name.return_value = "test-model" + + # Make processing fail + async def fail_process(item, doc): + raise RuntimeError("test failure") + + processor._process_single_document = fail_process + processor.queue_repo.mark_failed = AsyncMock() + + await processor._process_online_batch() + + _flush_exporter() + spans = span_exporter.get_finished_spans() + span_exporter.clear() + + consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) + assert consumer is not None + + # Check status is ERROR (StatusCode.ERROR = 2) + assert consumer.status is not None + assert consumer.status.status_code == StatusCode.ERROR, ( + "span must have ERROR status when item fails" + ) + + @pytest.mark.asyncio + async def test_missing_content_id_failure_real_path(self, span_exporter): + """Real _process_single_document: missing content_id returns False and + sets consumer span ERROR status.""" + processor = _make_processor(span_exporter) + + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + item = _make_item("doc1", traceparent=tp) + processor.queue_repo.get_pending_items.return_value = [item] + # Document with no content_id + doc = MagicMock(spec_set=["content_id", "external_id"]) + doc.content_id = None + doc.external_id = None + processor.documents_repo.get_by_ids.return_value = {"doc1": doc} + processor.documents_repo.find_embedded_content_donors = AsyncMock(return_value={}) + + # Use the REAL _process_single_document (no replacement) + result = await processor._process_online_batch() + assert result is True + + _flush_exporter() + spans = span_exporter.get_finished_spans() + span_exporter.clear() + + consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) + assert consumer is not None, "CONSUMER span must be exported" + + # The missing content_id path sets had_failure -> ERROR + assert consumer.status is not None + assert consumer.status.status_code == StatusCode.ERROR, ( + "span must have ERROR status when content_id is missing" + ) + + @pytest.mark.asyncio + async def test_embedding_exception_failure_real_path(self, span_exporter): + """Real _process_single_document: embedding provider exception returns + False and sets consumer span ERROR status.""" + processor = _make_processor(span_exporter) + + tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + item = _make_item("doc1", traceparent=tp) + processor.queue_repo.get_pending_items.return_value = [item] + # Document with content_id so we reach the embedding call + doc = MagicMock(spec_set=["content_id", "external_id"]) + doc.content_id = "c-doc1" + doc.external_id = None + processor.documents_repo.get_by_ids.return_value = {"doc1": doc} + processor.documents_repo.find_embedded_content_donors = AsyncMock(return_value={}) + processor.content_storage.get_text = AsyncMock( + return_value="some text content" + ) + # Make embedding provider raise an exception + provider = processor.embedding_provider + provider.generate_embeddings = AsyncMock(side_effect=RuntimeError("embedding failed")) + provider.get_model_name.return_value = "test-model" + + # Use the REAL _process_single_document (no replacement) + result = await processor._process_online_batch() + assert result is True + + _flush_exporter() + spans = span_exporter.get_finished_spans() + span_exporter.clear() + + consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) + assert consumer is not None, "CONSUMER span must be exported" + + # The embedding exception path sets had_failure -> ERROR + assert consumer.status is not None + assert consumer.status.status_code == StatusCode.ERROR, ( + "span must have ERROR status when embedding provider raises" + ) diff --git a/services/ai/tests/unit/test_sanitization.py b/services/ai/tests/unit/test_sanitization.py new file mode 100644 index 000000000..8503c9f9c --- /dev/null +++ b/services/ai/tests/unit/test_sanitization.py @@ -0,0 +1,217 @@ +"""Source-level sanitization regression tests for AI service logs. + +Tests scan the actual source files for forbidden patterns in logger.* call +lines. Only lines that contain a recognised logging call are inspected, so +functional non-log code is not flagged. + +This is a static analysis approach — no runtime dependencies required. +""" + +import re +import ast +import os +from pathlib import Path + + +AI_SERVICE_DIR = Path(__file__).resolve().parent.parent.parent # services/ai + + +def _is_log_call(line: str) -> bool: + """Check if a line appears to be a logging call or continuation.""" + stripped = line.strip() + # Match logger.info(, logger.debug(, logger.error(, logger.warn(, logger.exception( + if re.search(r'logger\.(info|debug|error|warn|exception)\(', stripped): + return True + # Continuation lines inside a log call (indented and part of a call) + if stripped.startswith('f"') or stripped.startswith('"') or stripped.startswith("'"): + return True + if stripped.startswith('+') or stripped == '': + return False + return False + + +def _get_forbidden_patterns() -> dict[str, list[str]]: + """Return file-relative-path -> list of forbidden regex patterns.""" + return { + "streaming/generate.py": [ + r"chat_id", # chat IDs in log messages + r"event.to_json", # raw event dumps + r"event\.delta\.text", # text content in deltas + r"event\.citation", # citation content + r"content_block\.text", # text content + r"content_block\.id", # content block ID + r"\{e\}", # exception string in f-string + r"str\(e\)", # exception string in log + r"exc_info\s*=\s*True", # traceback export in log + ], + "streaming/run.py": [ + r"exc_info\s*=\s*True", # traceback export in log + ], + "providers/anthropic.py": [ + r"t\['name'\]", # tool names + r"event\.content_block\.input", # tool input + r"event\.citation", # citation content + r"text_delta.*event\.delta\.text", # text content + r"partial_json", # JSON delta content + r"input_json", # input JSON + r"json\.dumps.*messages", # raw messages + r"json\.dumps.*request_params", # full request params + r"exc_info\s*=\s*True", # traceback export in log + ], + "providers/bedrock.py": [ + r"t\['name'\]", # tool names + r"event\.content_block\.input", # tool input + r"event\.citation", # citation content + r"text_delta.*event\.delta\.text", # text content + r"partial_json", # JSON delta content + r"input_json", # input JSON + r"json\.dumps.*messages", # raw messages + r"json\.dumps.*request_body", # full request body + r"response_body", + r"response from LLM ->", # full response + r"document\['name'\]", # raw document name in log + r"exc_info\s*=\s*True", # traceback export in log + ], + "providers/gemini.py": [ + r"exc_info\s*=\s*True", # traceback export in log + ], + "providers/openai_compatible.py": [ + r"exc_info\s*=\s*True", # traceback export in log + ], + "providers/openai.py": [ + r"exc_info\s*=\s*True", # traceback export in log + ], + "tools/searcher_client.py": [ + r"query:\s*\{request\.query", # raw query in log + r"query:\s*\{query", # raw query variable + r"response\.text", # raw response body + r"response\.status_code", # only status code is acceptable + r"exc_info\s*=\s*True", # traceback export in log + ], + "tools/connector_handler.py": [ + r"exc_info\s*=\s*True", # traceback export in log + ], + "tools/search_handler.py": [ + r"exc_info\s*=\s*True", # traceback export in log + ], + "email_service/sender.py": [ + r"body=%s", + r"resp\.text", # response body + r"exc_info\s*=\s*True", # traceback export in log + ], + "embeddings/batch_processor.py": [ + r"document_id", # document IDs in log messages + r"\{e\}", # exception string in f-string + r"str\(e\)", # exception string in log + r"exc_info\s*=\s*True", # traceback export in log + ], + "main.py": [ + r"\{e\}", # exception string in f-string + r"exc_info\s*=\s*True", # traceback export in log + ], + } + + +def _run_checks(rel_path: str, *, skip_line: callable = None) -> None: + """Run forbidden-pattern checks for a given relative source path. + + Optional *skip_line* callable receives (line_text, pattern) and returns + True if the match should be skipped. + """ + filepath = AI_SERVICE_DIR / rel_path + assert filepath.exists(), f"{filepath} not found" + content = filepath.read_text() + patterns = _get_forbidden_patterns()[rel_path] + violations = [] + for i, line in enumerate(content.splitlines(), 1): + if not _is_log_call(line): + continue + for pat in patterns: + if re.search(pat, line): + if skip_line is not None and skip_line(line, pat): + continue + violations.append(f" line {i}: {line.strip()}") + assert not violations, ( + f"Found {len(violations)} forbidden pattern(s) in {rel_path} log lines:\n" + + "\n".join(violations) + ) + + +def test_anthropic_sanitization(): + """Anthropic provider must not log tool input, text deltas, citations, request params, or tracebacks.""" + _run_checks("providers/anthropic.py") + + +def test_bedrock_sanitization(): + """Bedrock provider must not log full request body, messages, stream content, responses, raw document names, or tracebacks.""" + _run_checks("providers/bedrock.py") + + +def test_gemini_sanitization(): + """Gemini provider must not log tracebacks.""" + _run_checks("providers/gemini.py") + + +def test_openai_compatible_sanitization(): + """OpenAI-compatible provider must not log tracebacks.""" + _run_checks("providers/openai_compatible.py") + + +def test_openai_sanitization(): + """OpenAI provider must not log tracebacks.""" + _run_checks("providers/openai.py") + + +def test_streaming_generate_sanitization(): + """Streaming generate.py must not log chat IDs, raw events, text/citation content, exception strings, or tracebacks.""" + _run_checks("streaming/generate.py") + + +def test_streaming_run_sanitization(): + """Streaming run.py must not log tracebacks.""" + _run_checks("streaming/run.py") + + +def _skip_status_pattern(line: str, pat: str) -> bool: + """Skip response.status_code matches when the line uses status= format.""" + if "status=" in line: + return True + return False + + +def test_searcher_client_sanitization(): + """Searcher client must not log raw queries, response bodies, or tracebacks.""" + _run_checks("tools/searcher_client.py", skip_line=_skip_status_pattern) + + +def test_connector_handler_sanitization(): + """Connector handler must not log tracebacks.""" + _run_checks("tools/connector_handler.py") + + +def test_search_handler_sanitization(): + """Search handler must not log tracebacks.""" + _run_checks("tools/search_handler.py") + + +def test_email_sender_sanitization(): + """Email sender must not log response body or tracebacks.""" + _run_checks("email_service/sender.py") + + +def _skip_batch_processor_non_log(line: str, pat: str) -> bool: + """Skip data-structure lines that are not log calls.""" + # Line inside a data dict literal (not a log continuation) + if pat == r"document_id" and '"document_id"' in line: + return True + return False + + +def test_batch_processor_sanitization(): + """Batch processor must not log document IDs, exception values, or tracebacks.""" + _run_checks("embeddings/batch_processor.py", skip_line=_skip_batch_processor_non_log) + + +def test_main_sanitization(): + """Main module must not log exception values or tracebacks.""" + _run_checks("main.py") diff --git a/services/ai/tests/unit/test_telemetry.py b/services/ai/tests/unit/test_telemetry.py new file mode 100644 index 000000000..91ec887df --- /dev/null +++ b/services/ai/tests/unit/test_telemetry.py @@ -0,0 +1,212 @@ +"""Unit tests for telemetry helpers and embedding metric recording. + +Tests use monkeypatching of production instruments/globals to avoid +contacting a real collector. No real OTLP endpoint is required. +""" + +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +from telemetry import normalize_otlp_endpoint, parse_metric_export_interval +from db.documents import Document +from db.embedding_queue import EmbeddingQueueItem + + +# --------------------------------------------------------------------------- +# Pure helper: normalize_otlp_endpoint +# --------------------------------------------------------------------------- + + +class TestNormalizeOtlpEndpoint: + def test_trailing_slash_is_stripped(self): + assert normalize_otlp_endpoint("http://localhost:4318/") == "http://localhost:4318" + + def test_no_trailing_slash(self): + assert normalize_otlp_endpoint("http://localhost:4318") == "http://localhost:4318" + + def test_empty_string_is_none(self): + assert normalize_otlp_endpoint("") is None + + def test_none_is_none(self): + assert normalize_otlp_endpoint(None) is None + + +# --------------------------------------------------------------------------- +# Pure helper: parse_metric_export_interval +# --------------------------------------------------------------------------- + + +class TestParseMetricExportInterval: + @patch.dict("os.environ", {"OTEL_METRIC_EXPORT_INTERVAL": "30000"}) + def test_valid_positive_int(self): + assert parse_metric_export_interval() == 30000 + + @patch.dict("os.environ", {"OTEL_METRIC_EXPORT_INTERVAL": "30000"}) + def test_parses_env_var(self): + assert parse_metric_export_interval() == 30000 + + @patch.dict("os.environ", {"OTEL_METRIC_EXPORT_INTERVAL": "not-a-number"}) + def test_invalid_falls_back(self): + assert parse_metric_export_interval() == 60_000 + + @patch.dict("os.environ", {"OTEL_METRIC_EXPORT_INTERVAL": "0"}) + def test_zero_falls_back(self): + assert parse_metric_export_interval() == 60_000 + + @patch.dict("os.environ", {"OTEL_METRIC_EXPORT_INTERVAL": "-1"}) + def test_negative_falls_back(self): + assert parse_metric_export_interval() == 60_000 + + @patch.dict("os.environ", {"OTEL_METRIC_EXPORT_INTERVAL": ""}) + def test_empty_string_falls_back(self): + assert parse_metric_export_interval() == 60_000 + + @patch.dict("os.environ", clear=True) + def test_unset_falls_back(self): + assert parse_metric_export_interval() == 60_000 + + +# --------------------------------------------------------------------------- +# Embedding metric recording (monkeypatch production instruments) +# --------------------------------------------------------------------------- + + +class TestEmbeddingMetrics: + """Tests that exercise the real recording helpers by monkeypatching the + global instrument references.""" + + @pytest.fixture(autouse=True) + def _patch_instruments(self): + """Replace production instrument globals with MagicMock spies before + each test.""" + import embeddings.batch_processor as bp + + self._patches = [] + for name in ("_EMBEDDING_PROCESSED", "_EMBEDDING_FAILED", "_EMBEDDING_BATCH_DURATION", "_EMBEDDING_PENDING"): + mock = MagicMock() + setattr(bp, name, mock) + self._patches.append((name, mock)) + yield + # Restore + for name, _ in self._patches: + setattr(bp, name, None) + + @property + def _processed(self): + return dict(self._patches)["_EMBEDDING_PROCESSED"] + + @property + def _failed(self): + return dict(self._patches)["_EMBEDDING_FAILED"] + + @property + def _batch_duration(self): + return dict(self._patches)["_EMBEDDING_BATCH_DURATION"] + + @property + def _pending(self): + return dict(self._patches)["_EMBEDDING_PENDING"] + + def _rerecord_helper(self): + import embeddings.batch_processor as bp + bp._ensure_embedding_instruments() + + def test_record_embedding_processed(self): + from embeddings.batch_processor import _record_embedding_processed + + _record_embedding_processed() + self._processed.add.assert_called_once_with(1) + + def test_record_embedding_processed_bulk_clone(self): + from embeddings.batch_processor import _record_embedding_processed + + # Simulate bulk-clone path: 3 cloned documents + _record_embedding_processed(3) + self._processed.add.assert_called_once_with(3) + + def test_record_embedding_failed(self): + from embeddings.batch_processor import _record_embedding_failed + + _record_embedding_failed() + self._failed.add.assert_called_once_with(1) + + def test_record_embedding_batch_duration(self): + from embeddings.batch_processor import _record_embedding_batch_duration + + _record_embedding_batch_duration(1.5) + self._batch_duration.record.assert_called_once_with(1.5) + + def test_record_embedding_pending(self): + from embeddings.batch_processor import _record_embedding_pending + + _record_embedding_pending(42) + self._pending.set.assert_called_once_with(42) + + +# --------------------------------------------------------------------------- +# Bulk-clone integration path test +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_clone_same_content_embeddings_calls_processed_helper(): + """Invoke _clone_same_content_embeddings on a minimally constructed + processor and verify the production _record_embedding_processed helper + is called with the durable clone count and cloned items are removed.""" + import embeddings.batch_processor as bp + + # Stub instrument globals so _record_embedding_processed is observable + processed_mock = MagicMock() + bp._EMBEDDING_PROCESSED = processed_mock + + # Stub repos and provider + docs_repo = AsyncMock() + queue_repo = AsyncMock() + embeddings_repo = AsyncMock() + app_state = AsyncMock() + provider = MagicMock() + provider.get_model_name.return_value = "test-model" + app_state.embedding_provider = provider + + processor = bp.EmbeddingBatchProcessor( + documents_repo=docs_repo, + queue_repo=queue_repo, + embeddings_repo=embeddings_repo, + app_state=app_state, + ) + + # Two queue items: one for a clonable document, one for a non-clonable + clonable_item = EmbeddingQueueItem( + id="item-1", document_id="doc-a", status="pending", + error_message=None, retry_count=0, created_at=None, + ) + nonclonable_item = EmbeddingQueueItem( + id="item-2", document_id="doc-b", status="pending", + error_message=None, retry_count=0, created_at=None, + ) + items = [clonable_item, nonclonable_item] + + # documents_by_id: doc-a has content_id (can clone), doc-b has no content_id + doc_a = Document(id="doc-a", content_id="content-1") + doc_b = Document(id="doc-b", content_id=None) + documents_by_id = {"doc-a": doc_a, "doc-b": doc_b} + + # find_embedded_content_donors returns a donor for content-1 + docs_repo.find_embedded_content_donors.return_value = {"content-1": "donor-doc"} + + # bulk_clone_for_documents returns one clone + embeddings_repo.bulk_clone_for_documents.return_value = {"doc-a": 3} + + result = await processor._clone_same_content_embeddings(items, documents_by_id) + + # Verify _record_embedding_processed was called with the clone count + assert processed_mock.add.call_count >= 1 + # _record_embedding_processed is called with len(clone_counts) = 1 document + assert processed_mock.add.call_args.args[0] == 1 + + # Verify cloned items are removed from the returned list + returned_ids = [item.document_id for item in result] + assert "doc-a" not in returned_ids + assert "doc-b" in returned_ids + assert len(result) == 1 diff --git a/services/ai/tests/unit/test_telemetry_lifecycle.py b/services/ai/tests/unit/test_telemetry_lifecycle.py new file mode 100644 index 000000000..c631baa06 --- /dev/null +++ b/services/ai/tests/unit/test_telemetry_lifecycle.py @@ -0,0 +1,317 @@ +"""Focused lifecycle tests for OTel LoggingHandler init/shutdown. + +Tests use monkeypatching of production instruments/mocked exporters to avoid +contacting a real collector. Covers: + +- No-span formatter safety (TraceContextFilter) +- Active-span stdout trace_id/span_id +- Native exported context on log records +- Repeated init leaves one handler / one exported record +- Shutdown removes the handler +""" + +import logging +from unittest.mock import MagicMock, patch, AsyncMock + +import pytest + + +@pytest.fixture(autouse=True) +def _reset_globals(): + """Reset telemetry module globals before each test.""" + import telemetry as tel + tel._tracer_provider = None + tel._meter_provider = None + tel._logger_provider = None + tel._otel_log_handler = None + # Remove any OTelLoggingHandlers left by previous tests + root = logging.getLogger() + for h in list(root.handlers): + if "OTelLoggingHandler" in type(h).__name__: + root.removeHandler(h) + yield + # Cleanup after test + tel.shutdown_telemetry() + for h in list(root.handlers): + if "OTelLoggingHandler" in type(h).__name__: + root.removeHandler(h) + + +@pytest.fixture +def mock_otlp_endpoint(monkeypatch): + """Set OTEL_EXPORTER_OTLP_ENDPOINT so the log exporter path is + exercised, but patch the exporters in the telemetry module to + avoid real connections.""" + monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://127.0.0.1:1") + # Patch at the telemetry module level (where they are imported) + import telemetry as tel + monkeypatch.setattr(tel, "OTLPLogExporter", MagicMock) + monkeypatch.setattr(tel, "OTLPSpanExporter", MagicMock) + monkeypatch.setattr(tel, "OTLPMetricExporter", MagicMock) + + +class TestNoSpanFormatterSafety: + """TraceContextFilter must not KeyError when there is no active span.""" + + def test_trace_context_filter_no_span(self): + """Formatter safety: TraceContextFilter outside a span produces + zero-padded trace_id/span_id and does not raise.""" + from telemetry import TraceContextFilter + + filt = TraceContextFilter() + record = logging.LogRecord( + name="test", level=logging.INFO, pathname="", lineno=0, + msg="no-span test", args=(), exc_info=None, + ) + result = filt.filter(record) + assert result is True + assert record.trace_id == "00000000000000000000000000000000" + assert record.span_id == "0000000000000000" + + +class TestActiveSpanTraceContext: + """Log records emitted within an active span carry native trace/span IDs.""" + + def test_stdout_handler_gets_trace_ids(self, mock_otlp_endpoint, monkeypatch): + """An active span's trace_id/span_id appear on stdout log records.""" + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk.trace.export import SimpleSpanProcessor + import telemetry as tel + + # Use SimpleSpanProcessor without in-memory exporter since + # InMemorySpanExporter may not be available in all versions. + # We verify trace_id/span_id via a CaptureHandler instead. + tracer_provider = TracerProvider() + trace.set_tracer_provider(tracer_provider) + + # Create a mock app and init telemetry + app = MagicMock() + tel.init_telemetry(app) + + tracer = trace.get_tracer("test-tracer") + + # Capture a log record within an active span + captured_records = [] + + class CaptureHandler(logging.Handler): + def emit(self, record): + captured_records.append(record) + + cap_handler = CaptureHandler() + # Add TraceContextFilter from telemetry module + cap_handler.addFilter(tel.TraceContextFilter()) + root = logging.getLogger() + root.addHandler(cap_handler) + root.setLevel(logging.INFO) + + with tracer.start_as_current_span("test-span") as span: + span_context = span.get_span_context() + logging.getLogger("test").info("message within span") + + root.removeHandler(cap_handler) + + # Verify captured record has real trace IDs + assert len(captured_records) >= 1 + for rec in captured_records: + assert rec.trace_id != "00000000000000000000000000000000" + + tel.shutdown_telemetry() + + def test_native_exported_context(self): + """Log records exported via OTel logger provider carry the + active span's trace context natively. + + Uses real InMemoryLogRecordExporter to verify trace_id/span_id + correlation and resource service.name propagation without an OTLP + endpoint. + """ + from opentelemetry import trace + from opentelemetry.sdk.trace import TracerProvider + from opentelemetry.sdk._logs import ( + LoggerProvider as LogLoggerProvider, + LoggingHandler, + ) + from opentelemetry.sdk._logs.export import ( + SimpleLogRecordProcessor, + InMemoryLogRecordExporter, + ) + from opentelemetry.sdk.resources import Resource, SERVICE_NAME + + # Set up tracer provider + tracer_provider = TracerProvider() + trace.set_tracer_provider(tracer_provider) + + # Build a logger provider with in-memory exporter for assertions + memory_exporter = InMemoryLogRecordExporter() + logger_provider = LogLoggerProvider( + resource=Resource.create({ + SERVICE_NAME: "omni-ai", + }) + ) + logger_provider.add_log_record_processor( + SimpleLogRecordProcessor(memory_exporter) + ) + + # Create LoggingHandler that bridges stdlib logging to our logger provider. + # Use DEBUG level so INFO and above are captured. + otel_handler = LoggingHandler( + level=logging.DEBUG, + logger_provider=logger_provider, + ) + # Set root logger to DEBUG so records reach the handler + root = logging.getLogger() + previous_level = root.level + root.setLevel(logging.DEBUG) + root.addHandler(otel_handler) + + try: + tracer = trace.get_tracer("test-tracer") + test_logger = logging.getLogger("oteltest-native") + + with tracer.start_as_current_span("export-span") as span: + span_ctx = span.get_span_context() + test_logger.info("exported log within span") + + # Force flush and get exported log records + logger_provider.force_flush() + exported = memory_exporter.get_finished_logs() + + # Assert at least one log record was exported + assert len(exported) >= 1, ( + f"Expected at least 1 log record, got {len(exported)}" + ) + + # Find the record with our expected body (ReadableLogRecord.log_record.body) + matching = [ + r for r in exported + if r.log_record.body == "exported log within span" + ] + assert len(matching) >= 1, ( + f"No log record with expected body; bodies: {[r.log_record.body for r in exported]}" + ) + + log_record = matching[0] + lr = log_record.log_record + + # Assert trace_id and span_id match the active span + assert lr.trace_id == span_ctx.trace_id, ( + f"trace_id {lr.trace_id!r} != expected {span_ctx.trace_id!r}" + ) + assert lr.span_id == span_ctx.span_id, ( + f"span_id {lr.span_id!r} != expected {span_ctx.span_id!r}" + ) + + # Assert the resource carries service.name + resource = log_record.resource + assert resource is not None + svc_name = resource.attributes.get("service.name") + assert svc_name is not None, "resource missing service.name" + assert svc_name == "omni-ai", ( + f"service.name={svc_name!r}, expected 'omni-ai'" + ) + finally: + root.setLevel(previous_level) + root.removeHandler(otel_handler) + logger_provider.shutdown() + tracer_provider.shutdown() + """Repeated init leaves exactly one handler and one record per emit. + Also verifies that exporting via the OTel LoggingHandler works correctly + within an active span context. + """ + + def test_repeated_init_leaves_one_handler(self, mock_otlp_endpoint): + """Two inits should result in exactly one OTelLoggingHandler + attached to the root logger.""" + import telemetry as tel + + app = MagicMock() + tel.init_telemetry(app) + tel.init_telemetry(app) + + root = logging.getLogger() + otel_handlers = [ + h for h in root.handlers + if "LoggingHandler" in type(h).__name__ + ] + assert len(otel_handlers) == 1, ( + f"Expected 1 OTelLoggingHandler after repeated init, " + f"got {len(otel_handlers)}" + ) + + def test_repeated_init_produces_one_record(self, mock_otlp_endpoint): + """One log emit after two inits should produce exactly one + record through the OTel LoggingHandler.""" + import telemetry as tel + + app = MagicMock() + tel.init_telemetry(app) + tel.init_telemetry(app) + + # Verify only one OTelLoggingHandler exists + root = logging.getLogger() + otel_handlers = [ + h for h in root.handlers + if "LoggingHandler" in type(h).__name__ + ] + assert len(otel_handlers) == 1 + + # Verify that logging through the root logger reaches the handler + # by counting handler calls + handler = otel_handlers[0] + original_handle = handler.handle + handle_count = 0 + + def counting_handle(record): + nonlocal handle_count + handle_count += 1 + return original_handle(record) + + handler.handle = counting_handle + + test_logger = logging.getLogger("test-single-emit") + test_logger.info("single emit") + + # handle should be called exactly once + assert handle_count == 1, ( + f"Expected 1 handle call, got {handle_count}" + ) + + +class TestShutdown: + """Shutdown removes the handler and leaves zero OTelLoggingHandlers.""" + + def test_shutdown_removes_handler(self, mock_otlp_endpoint): + """After shutdown, no OTelLoggingHandler remains on the root logger.""" + import telemetry as tel + + app = MagicMock() + tel.init_telemetry(app) + + tel.shutdown_telemetry() + + root = logging.getLogger() + otel_handlers = [ + h for h in root.handlers + if "LoggingHandler" in type(h).__name__ + ] + assert len(otel_handlers) == 0, ( + f"Expected 0 OTelLoggingHandler after shutdown, " + f"got {len(otel_handlers)}" + ) + + def test_shutdown_also_removes_reinit_handler(self, mock_otlp_endpoint): + """After reinit and shutdown, zero handlers remain.""" + import telemetry as tel + + app = MagicMock() + tel.init_telemetry(app) + tel.init_telemetry(app) + tel.shutdown_telemetry() + + root = logging.getLogger() + otel_handlers = [ + h for h in root.handlers + if "LoggingHandler" in type(h).__name__ + ] + assert len(otel_handlers) == 0 diff --git a/services/ai/tools/connector_handler.py b/services/ai/tools/connector_handler.py index f74aa9949..0fa111d54 100644 --- a/services/ai/tools/connector_handler.py +++ b/services/ai/tools/connector_handler.py @@ -135,7 +135,7 @@ async def _load_cached_actions(self) -> list[ConnectorAction] | None: if cached: return [ConnectorAction(**d) for d in json.loads(cached)] except Exception as e: - logger.warning(f"Failed to load cached actions: {e}") + logger.warning("Failed to load cached actions") return None async def _cache_actions(self, actions: list[ConnectorAction]) -> None: @@ -148,7 +148,7 @@ async def _cache_actions(self, actions: list[ConnectorAction]) -> None: ex=ACTIONS_CACHE_TTL, ) except Exception as e: - logger.warning(f"Failed to cache actions: {e}") + logger.warning("Failed to cache actions") async def _fetch_actions(self) -> list[ConnectorAction]: """Fetch available actions from connector-manager. @@ -177,7 +177,7 @@ async def _fetch_actions(self) -> list[ConnectorAction]: sources = sources_from_sync_overview_response(sources_resp.json()) except Exception as e: - logger.error(f"Failed to fetch connector info: {e}") + logger.error("Failed to fetch connector info") return [] # Build a mapping from source_type to list of active sources @@ -243,7 +243,7 @@ async def _fetch_actions(self) -> list[ConnectorAction]: ) logger.info( - f"Discovered {len(actions)} connector actions for user {self._user_id}" + "Discovered %s connector actions", len(actions) ) return actions @@ -421,7 +421,7 @@ async def execute( ) logger.info( - f"Executing connector action: {action.action_name} on source {action.source_id}" + "Executing connector action: %s", action.action_name ) # If this action references a document, check user permissions @@ -476,7 +476,7 @@ async def execute( oauth_start_url = body.get("oauth_start_url") if not provider or not oauth_start_url: logger.error( - f"connector-manager 412 missing provider/oauth_start_url; body={body}" + "connector-manager 412 missing provider/oauth_start_url" ) return ToolResult( content=[ @@ -544,7 +544,7 @@ async def execute( result = response.json() except httpx.HTTPStatusError as e: logger.error( - f"Connector action HTTP {e.response.status_code}: {e.response.text}" + "Connector action HTTP error" ) return ToolResult( content=[{"type": "text", "text": f"Action failed: {e.response.text}"}], @@ -553,7 +553,7 @@ async def execute( except Exception as e: # Transport-level errors (ReadError, ConnectError, TimeoutException, etc.) # don't carry a response. Don't try to access e.response. - logger.error(f"Connector action failed: {e}", exc_info=True) + logger.error("Connector action failed") return ToolResult( content=[{"type": "text", "text": f"Action failed: {e}"}], is_error=True, diff --git a/services/ai/tools/search_handler.py b/services/ai/tools/search_handler.py index 007d27ba1..2a5a75617 100644 --- a/services/ai/tools/search_handler.py +++ b/services/ai/tools/search_handler.py @@ -83,7 +83,7 @@ async def fetch_operator_values( _operator_values_mem_ts = now return _operator_values_mem except Exception as e: - logger.warning(f"Failed to read operator values cache: {e}") + logger.warning("Failed to read operator values cache") attribute_keys = [ op.attribute_key @@ -96,7 +96,7 @@ async def fetch_operator_values( try: values = await searcher_client.get_attribute_values(attribute_keys) except Exception as e: - logger.warning(f"Failed to fetch operator values from searcher: {e}") + logger.warning("Failed to fetch operator values from searcher") return {} _operator_values_mem = values @@ -110,7 +110,7 @@ async def fetch_operator_values( ex=_OPERATOR_VALUES_CACHE_TTL, ) except Exception as e: - logger.warning(f"Failed to cache operator values: {e}") + logger.warning("Failed to cache operator values") return values @@ -260,7 +260,7 @@ async def _execute_search( try: params = SearchToolParams.model_validate(tool_input) except ValidationError as e: - logger.error(f"Invalid search_documents input: {e}") + logger.error("Invalid search_documents input") return ToolResult( content=[{"type": "text", "text": f"Invalid parameters: {e}"}], is_error=True, @@ -271,7 +271,9 @@ async def _execute_search( params.document_id = params.document_id[len("_ref:"):] logger.info( - f"Executing search_documents with query: {params.query}, document_id: {params.document_id}, context: {context}" + "Executing search_documents with query present=%s, document_id present=%s", + bool(params.query), + bool(params.document_id), ) search_user_id = None if context.skip_permission_check else context.user_id search_user_email = ( @@ -424,6 +426,6 @@ async def _execute_search_tool( try: response: SearchResponse = await searcher_tool.handle(search_request) except Exception as e: - logger.error(f"Search failed: {e}") + logger.error("Search failed") return [] return response.results diff --git a/services/ai/tools/searcher_client.py b/services/ai/tools/searcher_client.py index 04f85f2b1..a3ceb4e44 100644 --- a/services/ai/tools/searcher_client.py +++ b/services/ai/tools/searcher_client.py @@ -171,7 +171,7 @@ async def search_documents(self, request: SearchRequest) -> SearchResponse: try: search_payload = request.model_dump(mode="json") - logger.info(f"Calling searcher service with query: {request.query}...") + logger.info("Calling searcher service...") response = await self.client.post( f"{self.searcher_url}/search", json=search_payload @@ -179,11 +179,11 @@ async def search_documents(self, request: SearchRequest) -> SearchResponse: if response.status_code == 200: search_results = SearchResponse.model_validate(response.json()) - logger.info(f"Search completed: {search_results.total_count} results") + logger.info("Search completed: %s results", search_results.total_count) return search_results else: logger.error( - f"Search service error: {response.status_code} - {response.text}" + "Search service error: status=%s", response.status_code ) raise SearcherError( message=f"Searcher API call failed: {response.status_code} {response.text}", @@ -191,13 +191,13 @@ async def search_documents(self, request: SearchRequest) -> SearchResponse: response=response, ) except Exception as e: - logger.error(f"Call to searcher service failed: {e}") + logger.error("Call to searcher service failed") raise async def search_people(self, request: PeopleSearchRequest) -> PeopleSearchResponse: """Search the people directory using omni-searcher service.""" try: - logger.info(f"People search with query: {request.query}...") + logger.info("People search...") response = await self.client.get( f"{self.searcher_url}/people/search", params={"q": request.query, "limit": request.limit}, @@ -205,11 +205,11 @@ async def search_people(self, request: PeopleSearchRequest) -> PeopleSearchRespo if response.status_code == 200: result = PeopleSearchResponse.model_validate(response.json()) - logger.info(f"People search completed: {len(result.people)} results") + logger.info("People search completed: %s results", len(result.people)) return result else: logger.error( - f"People search error: {response.status_code} - {response.text}" + "People search error: status=%s", response.status_code ) raise SearcherError( message=f"People search failed: {response.status_code} {response.text}", @@ -219,7 +219,7 @@ async def search_people(self, request: PeopleSearchRequest) -> PeopleSearchRespo except SearcherError: raise except Exception as e: - logger.error(f"People search failed: {e}") + logger.error("People search failed") raise async def upsert_capabilities( @@ -233,7 +233,7 @@ async def upsert_capabilities( if response.status_code == 200: return CapabilitiesUpsertResponse.model_validate(response.json()) logger.error( - f"Capability upsert error: {response.status_code} - {response.text}" + "Capability upsert error: status=%s", response.status_code ) raise SearcherError( message=f"Capability upsert failed: {response.status_code} {response.text}", @@ -252,7 +252,7 @@ async def sync_capabilities( if response.status_code == 200: return CapabilitiesSyncResponse.model_validate(response.json()) logger.error( - f"Capability sync error: {response.status_code} - {response.text}" + "Capability sync error: status=%s", response.status_code ) raise SearcherError( message=f"Capability sync failed: {response.status_code} {response.text}", @@ -271,7 +271,7 @@ async def search_capabilities( if response.status_code == 200: return CapabilitySearchResponse.model_validate(response.json()) logger.error( - f"Capability search error: {response.status_code} - {response.text}" + "Capability search error: status=%s", response.status_code ) raise SearcherError( message=f"Capability search failed: {response.status_code} {response.text}", @@ -293,11 +293,11 @@ async def get_attribute_values( return data.get("attributes", {}) else: logger.error( - f"Attribute values fetch error: {response.status_code} - {response.text}" + "Attribute values fetch error: status=%s", response.status_code ) return {} except Exception as e: - logger.error(f"Failed to fetch attribute values: {e}") + logger.error("Failed to fetch attribute values") return {} async def close(self): diff --git a/services/connector-manager/Cargo.toml b/services/connector-manager/Cargo.toml index d31c5c0bb..c84e9fc5e 100644 --- a/services/connector-manager/Cargo.toml +++ b/services/connector-manager/Cargo.toml @@ -35,6 +35,7 @@ async-trait = { workspace = true } async-stream = "0.3" time = { workspace = true } dashmap = { workspace = true } +opentelemetry = { workspace = true } shared = { path = "../../shared" } [dev-dependencies] diff --git a/services/connector-manager/src/connector_client.rs b/services/connector-manager/src/connector_client.rs index 43ce0fc50..5148bfee8 100644 --- a/services/connector-manager/src/connector_client.rs +++ b/services/connector-manager/src/connector_client.rs @@ -4,6 +4,7 @@ use crate::models::{ }; use reqwest::Client; use shared::models::SyncType; +use shared::telemetry::http_client; use shared::{RateLimiter, RetryableError}; use std::sync::atomic::{AtomicU32, Ordering}; use std::time::Duration; @@ -36,22 +37,18 @@ impl ConnectorClient { connector_url: &str, ) -> Result { let url = format!("{}/manifest", connector_url); - debug!("Fetching manifest from {}", url); + debug!("Fetching manifest"); - let response = self - .client - .get(&url) - .send() + let response = http_client::send_traced("GET", &url, self.client.get(&url)) .await .map_err(|e| ClientError::RequestFailed(e.to_string()))?; if !response.status().is_success() { let status = response.status(); - let body = response.text().await.unwrap_or_default(); - error!("Failed to get manifest: {} - {}", status, body); + error!("Failed to get manifest: status={}", status); return Err(ClientError::ConnectorError { status: status.as_u16(), - message: body, + message: String::new(), }); } @@ -68,9 +65,6 @@ impl ConnectorClient { ) -> Result { let url = format!("{}/sync", connector_url); debug!( - url, - source_id = %request.source_id, - sync_run_id = %request.sync_run_id, sync_mode = ?request.sync_mode, "Triggering sync" ); @@ -79,15 +73,14 @@ impl ConnectorClient { if !response.status().is_success() { let status = response.status(); - let body = response.text().await.unwrap_or_default(); if status.as_u16() == 404 && request.sync_mode == SyncType::Realtime { - debug!("Realtime sync unavailable: {} - {}", status, body); + debug!("Realtime sync unavailable: status={}", status); } else { - error!("Failed to trigger sync: {} - {}", status, body); + error!("Failed to trigger sync: status={}", status); } return Err(ClientError::ConnectorError { status: status.as_u16(), - message: body, + message: String::new(), }); } @@ -107,22 +100,18 @@ impl ConnectorClient { .execute_with_retry(|| async { let attempt = attempts.fetch_add(1, Ordering::Relaxed) + 1; debug!( - url, - source_id = %request.source_id, - sync_run_id = %request.sync_run_id, sync_mode = ?request.sync_mode, attempt, "Sending connector sync trigger" ); - let result = self.client.post(url).json(request).send().await; + let result = + http_client::send_traced("POST", url, self.client.post(url).json(request)) + .await; match result { Ok(response) => { debug!( - url, - source_id = %request.source_id, - sync_run_id = %request.sync_run_id, sync_mode = ?request.sync_mode, attempt, status = response.status().as_u16(), @@ -132,12 +121,8 @@ impl ConnectorClient { } Err(e) => { warn!( - url, - source_id = %request.source_id, - sync_run_id = %request.sync_run_id, sync_mode = ?request.sync_mode, attempt, - error = %e, "Connector sync trigger request failed" ); Err(RetryableError::Transient(e.into())) @@ -154,22 +139,21 @@ impl ConnectorClient { sync_run_id: &str, ) -> Result { let url = format!("{}/sync/{}", connector_url, sync_run_id); - debug!("Probing sync status at {}", url); + debug!("Probing sync status"); - let response = self - .client - .get(&url) - .timeout(Duration::from_secs(5)) - .send() - .await - .map_err(|e| ClientError::RequestFailed(e.to_string()))?; + let response = http_client::send_traced( + "GET", + &url, + self.client.get(&url).timeout(Duration::from_secs(5)), + ) + .await + .map_err(|e| ClientError::RequestFailed(e.to_string()))?; if !response.status().is_success() { let status = response.status(); - let body = response.text().await.unwrap_or_default(); return Err(ClientError::ConnectorError { status: status.as_u16(), - message: body, + message: String::new(), }); } @@ -185,25 +169,24 @@ impl ConnectorClient { sync_run_id: &str, ) -> Result<(), ClientError> { let url = format!("{}/cancel", connector_url); - debug!("Cancelling sync {} at {}", sync_run_id, url); + debug!("Cancelling sync"); - let response = self - .client - .post(&url) - .json(&CancelRequest { + let response = http_client::send_traced( + "POST", + &url, + self.client.post(&url).json(&CancelRequest { sync_run_id: sync_run_id.to_string(), - }) - .send() - .await - .map_err(|e| ClientError::RequestFailed(e.to_string()))?; + }), + ) + .await + .map_err(|e| ClientError::RequestFailed(e.to_string()))?; if !response.status().is_success() { let status = response.status(); - let body = response.text().await.unwrap_or_default(); - warn!("Failed to cancel sync: {} - {}", status, body); + warn!("Failed to cancel sync: status={}", status); return Err(ClientError::ConnectorError { status: status.as_u16(), - message: body, + message: String::new(), }); } @@ -216,23 +199,18 @@ impl ConnectorClient { request: &ActionRequest, ) -> Result { let url = format!("{}/action", connector_url); - debug!("Executing action {} at {}", request.action, url); + debug!("Executing action"); - let response = self - .client - .post(&url) - .json(request) - .send() + let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await .map_err(|e| ClientError::RequestFailed(e.to_string()))?; if !response.status().is_success() { let status = response.status(); - let body = response.text().await.unwrap_or_default(); - error!("Failed to execute action: {} - {}", status, body); + error!("Failed to execute action: status={}", status); return Err(ClientError::ConnectorError { status: status.as_u16(), - message: body, + message: String::new(), }); } @@ -255,13 +233,9 @@ impl ConnectorClient { request: &ActionRequest, ) -> Result { let url = format!("{}/action", connector_url); - debug!("Executing action (raw) {} at {}", request.action, url); + debug!("Executing action (raw)"); - let response = self - .client - .post(&url) - .json(request) - .send() + let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await .map_err(|e| ClientError::RequestFailed(e.to_string()))?; @@ -279,13 +253,9 @@ impl ConnectorClient { request: &OAuthCredentialReadyRequest, ) -> Result, ClientError> { let url = format!("{}/oauth/credential-ready", connector_url); - debug!("Notifying connector of OAuth credential-ready at {}", url); + debug!("Notifying connector of OAuth credential-ready"); - let response = self - .client - .post(&url) - .json(request) - .send() + let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await .map_err(|e| ClientError::RequestFailed(e.to_string()))?; @@ -298,7 +268,7 @@ impl ConnectorClient { ); return Err(ClientError::ConnectorError { status: response.status().as_u16(), - message: response.text().await.unwrap_or_default(), + message: String::new(), }); } @@ -319,23 +289,18 @@ impl ConnectorClient { request: &ResourceRequest, ) -> Result { let url = format!("{}/resource", connector_url); - debug!("Reading resource {} at {}", request.uri, url); + debug!("Reading resource"); - let response = self - .client - .post(&url) - .json(request) - .send() + let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await .map_err(|e| ClientError::RequestFailed(e.to_string()))?; if !response.status().is_success() { let status = response.status(); - let body = response.text().await.unwrap_or_default(); - error!("Failed to read resource: {} - {}", status, body); + error!("Failed to read resource: status={}", status); return Err(ClientError::ConnectorError { status: status.as_u16(), - message: body, + message: String::new(), }); } @@ -351,23 +316,18 @@ impl ConnectorClient { request: &PromptRequest, ) -> Result { let url = format!("{}/prompt", connector_url); - debug!("Getting prompt {} at {}", request.name, url); + debug!("Getting prompt"); - let response = self - .client - .post(&url) - .json(request) - .send() + let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await .map_err(|e| ClientError::RequestFailed(e.to_string()))?; if !response.status().is_success() { let status = response.status(); - let body = response.text().await.unwrap_or_default(); - error!("Failed to get prompt: {} - {}", status, body); + error!("Failed to get prompt: status={}", status); return Err(ClientError::ConnectorError { status: status.as_u16(), - message: body, + message: String::new(), }); } @@ -383,23 +343,18 @@ impl ConnectorClient { request: &SkillRequest, ) -> Result { let url = format!("{}/skill", connector_url); - debug!("Getting skill {} at {}", request.skill_id, url); + debug!("Getting skill"); - let response = self - .client - .post(&url) - .json(request) - .send() + let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await .map_err(|e| ClientError::RequestFailed(e.to_string()))?; if !response.status().is_success() { let status = response.status(); - let body = response.text().await.unwrap_or_default(); - error!("Failed to get skill: {} - {}", status, body); + error!("Failed to get skill: status={}", status); return Err(ClientError::ConnectorError { status: status.as_u16(), - message: body, + message: String::new(), }); } @@ -411,7 +366,7 @@ impl ConnectorClient { pub async fn health_check(&self, connector_url: &str) -> bool { let url = format!("{}/health", connector_url); - match self.client.get(&url).send().await { + match http_client::send_traced("GET", &url, self.client.get(&url)).await { Ok(response) => response.status().is_success(), Err(_) => false, } diff --git a/services/connector-manager/src/handlers.rs b/services/connector-manager/src/handlers.rs index 754e76ff4..6f24b6fee 100644 --- a/services/connector-manager/src/handlers.rs +++ b/services/connector-manager/src/handlers.rs @@ -23,6 +23,7 @@ use serde_json::{json, Value}; use shared::clients::docling::{DoclingClient, DoclingError}; use shared::db::repositories::{ConfigurationRepository, SyncRunRepository}; +use shared::metrics; use shared::models::{ ActionMode, ConnectorManifest, GlobalConfiguration, SearchOperator, ServiceCredential, ServiceProvider, Source, SourceType, SyncRun, SyncType, @@ -38,6 +39,15 @@ use std::time::Duration; use tokio::sync::OwnedSemaphorePermit; use tracing::{debug, error, info, warn}; +/// Record terminal sync metric via the shared helper. Never records IDs. +fn record_sync_terminal( + sync_type: SyncType, + outcome: &str, + created_at: Option, +) { + metrics::record_sync_terminal(&sync_type.to_string(), outcome, created_at); +} + pub async fn health_check() -> impl IntoResponse { Json(json!({ "status": "healthy" })) } @@ -46,7 +56,7 @@ pub async fn trigger_sync( State(state): State, Json(request): Json, ) -> Result, ApiError> { - info!("Manual sync triggered for source {}", request.source_id); + info!("Manual sync triggered"); let sync_run_id = state .sync_manager @@ -67,14 +77,14 @@ pub async fn trigger_sync_by_id( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - info!("Manual sync triggered for source {}", source_id); + info!("Manual sync triggered"); let sync_run_id = state .sync_manager .trigger_sync(&source_id, SyncType::Full, TriggerType::Manual) .await .map_err(|e| { - error!("Failed to trigger sync for source {}: {:?}", source_id, e); + error!("Failed to trigger sync"); ApiError::from(e) })?; @@ -88,7 +98,7 @@ pub async fn cancel_sync( State(state): State, Path(sync_run_id): Path, ) -> Result, ApiError> { - info!("Cancel requested for sync {}", sync_run_id); + info!("Cancel requested"); state.sync_manager.cancel_sync(&sync_run_id).await?; @@ -99,7 +109,7 @@ pub async fn get_sync_progress( State(state): State, Path(sync_run_id): Path, ) -> Result>>, ApiError> { - debug!("SSE connection for sync progress: {}", sync_run_id); + debug!("SSE connection for sync progress"); let pool = state.db_pool.pool().clone(); let sync_run_id_clone = sync_run_id.clone(); @@ -113,7 +123,7 @@ pub async fn get_sync_progress( let progress = match get_progress_from_db(&pool, &sync_run_id_clone).await { Ok(p) => p, Err(e) => { - error!("Failed to get progress: {}", e); + error!("Failed to get progress"); break; } }; @@ -358,290 +368,146 @@ pub async fn execute_action( _headers: HeaderMap, Json(request): Json, ) -> Result { - let is_transient = request.transient_credentials.is_some(); - let source_type: SourceType; - let source: Option; - let creds: shared::models::ServiceCredential; - let mut params = request.params.clone(); - let mut transient_actor_email = None; - let manifests = get_registered_manifests(&state.redis_client).await; + info!("Executing action"); - let (connector_url, action_admin_only) = if is_transient { - // ======== TRANSIENT MODE ======== - // Transient mode: source_type + transient_credentials required, no source/credential DB. - let tc = request.transient_credentials.as_ref().unwrap(); - source_type = request.source_type.ok_or_else(|| { - ApiError::BadRequest( - "source_type is required when transient_credentials are provided".to_string(), - ) - })?; - if request.source_id.is_some() { - return Err(ApiError::BadRequest( - "source_id must not be set when transient_credentials are provided".to_string(), - )); - } - - // Resolve user/admin from user_id (required in transient mode). - let user_id = request.user_id.as_ref().ok_or_else(|| { - ApiError::BadRequest("user_id is required in transient mode".to_string()) - })?; + let source_repo = SourceRepository::new(state.db_pool.pool()); + let source = source_repo + .find_by_id(request.source_id.clone()) + .await + .map_err(|e| ApiError::Internal(e.to_string()))? + .ok_or_else(|| ApiError::NotFound(format!("Source not found: {}", request.source_id)))?; - // Look up the connector manifest by source_type. - let manifest = manifests - .iter() - .find(|m| m.source_types.contains(&source_type)) - .ok_or_else(|| { - ApiError::NotFound(format!( - "Connector not registered for type: {:?}", - source_type - )) - })?; + // Look up the connector manifest to get connector_url and action metadata + let manifests = get_registered_manifests(&state.redis_client).await; + let manifest = manifests + .iter() + .find(|m| m.source_types.contains(&source.source_type)); - let connector_url = manifest.connector_url.clone(); - let action_def = manifest - .actions - .iter() - .find(|a| a.name == request.action) - .ok_or_else(|| { - ApiError::BadRequest(format!( - "Unknown action '{}' for source type {:?}", - request.action, source_type - )) - })?; - if !action_def.source_types.contains(&source_type) { - return Err(ApiError::BadRequest(format!( - "Action '{}' does not support source type {:?}", - request.action, source_type - ))); - } - let action_admin_only = action_def.admin_only; - let action_mode = action_def.mode; + let connector_url = manifest + .as_ref() + .map(|m| m.connector_url.clone()) + .ok_or_else(|| { + ApiError::NotFound(format!( + "Connector not registered for type: {:?}", + source.source_type + )) + })?; - if action_mode != ActionMode::Read { + let action_def = manifest.and_then(|m| m.actions.iter().find(|a| a.name == request.action)); + let action_mode = action_def.map(|a| a.mode).unwrap_or_default(); + // Reject unknown action names — they must exist in the manifest + let action_def = action_def.ok_or_else(|| { + ApiError::BadRequest(format!( + "Unknown action '{}' for source type {:?}", + request.action, source.source_type + )) + })?; + let action_admin_only = action_def.admin_only; + + // Generic read_only enforcement — the source config is the authority. + let source_read_only = source + .config + .get("read_only") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + if let Some(m) = manifest.as_ref() { + if (m.read_only || source_read_only) && action_mode == ActionMode::Write { return Err(ApiError::BadRequest(format!( - "Transient action '{}' must be read-only", + "Action '{}' is not allowed: source is read-only", request.action ))); } + } - let user_repo = UserRepository::new(state.db_pool.pool()); - let user = user_repo - .find_by_id(user_id.clone()) - .await - .map_err(|e| ApiError::Internal(e.to_string()))? - .ok_or_else(|| ApiError::NotFound(format!("User not found: {user_id}")))?; - if action_admin_only && user.role != shared::models::UserRole::Admin { - return Err(ApiError::BadRequest(format!( - "Action '{}' requires admin privileges", - request.action + let creds_repo = ServiceCredentialsRepo::new(state.db_pool.pool().clone()) + .map_err(|e| ApiError::Internal(e.to_string()))?; + let creds = match resolve_credentials( + &creds_repo, + &request.source_id, + request.user_id.as_deref(), + action_admin_only, + ) + .await? + { + CredentialResolution::Resolved(c) => c, + CredentialResolution::NeedsUserAuth { provider } => { + return Ok(needs_user_auth_response( + &request.source_id, + source.source_type, + provider, + )?); + } + CredentialResolution::NoCredentials => { + return Err(ApiError::NotFound(format!( + "Credentials not found for source: {}", + request.source_id ))); } - transient_actor_email = Some(user.email); - - // Adapt the transient payload to the connector SDK's credential model. - let tc = tc.clone(); - creds = shared::models::ServiceCredential { - id: "transient".to_string(), - source_id: "transient".to_string(), - user_id: Some(user_id.clone()), - provider: tc.provider, - auth_type: tc.auth_type, - principal_email: tc.principal_email, - credentials: tc.credentials, - config: tc.config, - expires_at: None, - last_validated_at: None, - created_at: time::OffsetDateTime::now_utc(), - updated_at: time::OffsetDateTime::now_utc(), - }; - source = None; + }; - info!( - "Executing transient action '{}' for source_type {:?} (user {:?})", - request.action, source_type, request.user_id - ); + // Resolve Omni document ID -> source external_id. + // TODO: replace hard-coded param names with a connector-declared resolve_params list. + let mut params = request.params.clone(); + let doc_id = params + .get("document_id") + .or_else(|| params.get("file_id")) + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + if let Some(doc_id) = doc_id { + let doc_repo = DocumentRepository::new(state.db_pool.pool()); + if let Ok(Some(doc)) = doc_repo.find_by_id(&doc_id).await { + info!("Resolved document/file ID -> external_id"); + if let Some(obj) = params.as_object_mut() { + obj.remove("document_id"); + obj.remove("file_id"); + obj.insert( + "file_id".to_string(), + serde_json::Value::String(doc.external_id), + ); + } + } + // If not found, assume the ID is already a source-native ID and pass through + } - (connector_url, action_admin_only) - } else { - // ======== PERSISTED MODE (unchanged) ======== - if request.source_type.is_some() { - return Err(ApiError::BadRequest( - "source_type is only valid with transient_credentials".to_string(), - )); + // Merge source config into params for legacy connector compatibility. + // Connectors that use the typed context (e.g. Darwinbox) ignore these. + if params.is_null() { + params = serde_json::Value::Object(serde_json::Map::new()); + } + if let (Some(src_obj), Some(params_obj)) = (source.config.as_object(), params.as_object_mut()) { + for (k, v) in src_obj { + params_obj.entry(k.clone()).or_insert_with(|| v.clone()); } - let source_id = request - .source_id - .as_deref() - .filter(|id| !id.is_empty()) - .ok_or_else(|| ApiError::BadRequest("source_id is required".to_string()))? - .to_string(); + } - let source_repo = SourceRepository::new(state.db_pool.pool()); - let db_source = source_repo - .find_by_id(source_id.clone()) + // Resolve actor email for the execution context + let actor_email = if let Some(uid) = request.user_id.as_ref() { + let user_repo = UserRepository::new(state.db_pool.pool()); + let user = user_repo + .find_by_id(uid.clone()) .await .map_err(|e| ApiError::Internal(e.to_string()))? - .ok_or_else(|| ApiError::NotFound(format!("Source not found: {source_id}")))?; - - source_type = db_source.source_type; - - let manifest = manifests - .iter() - .find(|m| m.source_types.contains(&source_type)) - .ok_or_else(|| { - ApiError::NotFound(format!( - "Connector not registered for type: {:?}", - source_type - )) - })?; - let connector_url = manifest.connector_url.clone(); - - let action_def = manifest - .actions - .iter() - .find(|a| a.name == request.action) - .ok_or_else(|| { - ApiError::BadRequest(format!( - "Unknown action '{}' for source type {:?}", - request.action, source_type - )) - })?; - if !action_def.source_types.contains(&source_type) { - return Err(ApiError::BadRequest(format!( - "Action '{}' does not support source type {:?}", - request.action, source_type - ))); - } - let action_admin_only = action_def.admin_only; - let action_mode = action_def.mode; - - // Generic read_only enforcement — the source config is the authority. - let source_read_only = db_source - .config - .get("read_only") - .and_then(|v| v.as_bool()) - .unwrap_or(false); - if (manifest.read_only || source_read_only) && action_mode == ActionMode::Write { + .ok_or_else(|| ApiError::NotFound(format!("User not found: {uid}")))?; + // Enforce admin_only action authorization + if action_admin_only && user.role != shared::models::UserRole::Admin { return Err(ApiError::BadRequest(format!( - "Action '{}' is not allowed: source is read-only", + "Action '{}' requires admin privileges", request.action ))); } - - let creds_repo = ServiceCredentialsRepo::new(state.db_pool.pool().clone()) - .map_err(|e| ApiError::Internal(e.to_string()))?; - creds = match resolve_credentials( - &creds_repo, - &source_id, - request.user_id.as_deref(), - action_admin_only, - ) - .await? - { - CredentialResolution::Resolved(c) => c, - CredentialResolution::NeedsUserAuth { provider } => { - return Ok(needs_user_auth_response(&source_id, source_type, provider)?); - } - CredentialResolution::NoCredentials => { - return Err(ApiError::NotFound(format!( - "Credentials not found for source: {source_id}" - ))); - } - }; - - // Resolve Omni document ID -> source external_id (persisted mode only). - let doc_id = params - .get("document_id") - .or_else(|| params.get("file_id")) - .and_then(|v| v.as_str()) - .map(|s| s.to_string()); - if let Some(doc_id) = doc_id { - let doc_repo = DocumentRepository::new(state.db_pool.pool()); - if let Ok(Some(doc)) = doc_repo.find_by_id(&doc_id).await { - info!( - "Resolved document/file ID {} -> external_id {}", - doc_id, doc.external_id - ); - if let Some(obj) = params.as_object_mut() { - obj.remove("document_id"); - obj.remove("file_id"); - obj.insert( - "file_id".to_string(), - serde_json::Value::String(doc.external_id), - ); - } - } - } - - // Merge source config into params for legacy connector compatibility. - if params.is_null() { - params = serde_json::Value::Object(serde_json::Map::new()); - } - if let (Some(src_obj), Some(params_obj)) = - (db_source.config.as_object(), params.as_object_mut()) - { - for (k, v) in src_obj { - params_obj.entry(k.clone()).or_insert_with(|| v.clone()); - } - } - - source = Some(db_source); - - info!( - "Executing action '{}' for source {} (user {:?}, params keys: {:?})", - request.action, - source_id, - request.user_id, - request - .params - .as_object() - .map(|m| m.keys().cloned().collect::>()) - .unwrap_or_default() - ); - - (connector_url, action_admin_only) - }; - - // ==== Common path for both modes ==== - // Resolve the actor email; transient mode already loaded the actor above. - let actor_email = if !is_transient { - if let Some(uid) = request.user_id.as_ref() { - let user_repo = UserRepository::new(state.db_pool.pool()); - let user = user_repo - .find_by_id(uid.clone()) - .await - .map_err(|e| ApiError::Internal(e.to_string()))? - .ok_or_else(|| ApiError::NotFound(format!("User not found: {uid}")))?; - // Enforce admin_only action authorization (user already validated for transient mode above). - if action_admin_only && user.role != shared::models::UserRole::Admin { - return Err(ApiError::BadRequest(format!( - "Action '{}' requires admin privileges", - request.action - ))); - } - Some(user.email) - } else { - None - } + Some(user.email) } else { - transient_actor_email + None }; - if params.is_null() { - params = serde_json::Value::Object(serde_json::Map::new()); - } - - info!( - "Dispatching action '{}' to connector {} (provider={:?}, auth_type={:?}, principal={:?})", - request.action, connector_url, creds.provider, creds.auth_type, creds.principal_email, - ); + info!("Dispatching action '{}'", request.action,); let client = ConnectorClient::new(); let action_request = ActionRequest { action: request.action, params, credentials: Some(creds), - source, + source: Some(source), actor_email, }; @@ -654,6 +520,7 @@ pub async fn execute_action( let status = response.status(); let mut builder = axum::response::Response::builder().status(status); + // Forward all headers except hop-by-hop connection headers. let hop_by_hop = [ "connection", "keep-alive", @@ -752,14 +619,8 @@ async fn resolve_credentials( .await .map_err(internal)?; match &resolved { - Some(c) => info!( - "resolve_credentials(source={}, user={:?}): admin_only → org cred {}", - source_id, user_id, c.id - ), - None => warn!( - "resolve_credentials(source={}, user={:?}): admin_only → no org cred found", - source_id, user_id - ), + Some(_) => info!("resolve_credentials: admin_only → org cred"), + None => warn!("resolve_credentials: admin_only → no org cred found"), } return Ok(resolved .map(CredentialResolution::Resolved) @@ -780,10 +641,7 @@ async fn resolve_credentials( { user_cred = merge_org_and_user_credentials(org_cred, user_cred); } - info!( - "resolve_credentials(source={}, user={}): per-user cred {}", - source_id, uid, user_cred.id - ); + info!("resolve_credentials: per-user cred resolved"); return Ok(CredentialResolution::Resolved(user_cred)); } // No per-user row — surface a NeedsUserAuth response so the UI @@ -795,19 +653,13 @@ async fn resolve_credentials( .map_err(internal)? { Some(org) => { - info!( - "resolve_credentials(source={}, user={}): no per-user cred, org row exists → NeedsUserAuth({:?})", - source_id, uid, org.provider - ); + info!("resolve_credentials: no per-user cred, NeedsUserAuth"); Ok(CredentialResolution::NeedsUserAuth { provider: org.provider, }) } None => { - warn!( - "resolve_credentials(source={}, user={}): no per-user cred and no org cred", - source_id, uid - ); + warn!("resolve_credentials: no per-user cred and no org cred"); Ok(CredentialResolution::NoCredentials) } } @@ -818,14 +670,8 @@ async fn resolve_credentials( .await .map_err(internal)?; match &resolved { - Some(c) => info!( - "resolve_credentials(source={}, no user): org cred {}", - source_id, c.id - ), - None => warn!( - "resolve_credentials(source={}, no user): no org cred found", - source_id - ), + Some(c) => info!("resolve_credentials: org cred resolved"), + None => warn!("resolve_credentials: no org cred found"), } Ok(resolved .map(CredentialResolution::Resolved) @@ -971,10 +817,7 @@ pub async fn read_resource( State(state): State, Json(request): Json, ) -> Result, ApiError> { - info!( - "Reading resource {} for source {}", - request.uri, request.source_id - ); + info!("Reading resource"); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -1023,10 +866,7 @@ pub async fn get_prompt( State(state): State, Json(request): Json, ) -> Result, ApiError> { - info!( - "Getting prompt {} for source {}", - request.name, request.source_id - ); + info!("Getting prompt"); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -1080,10 +920,7 @@ pub async fn oauth_credential_ready( State(state): State, Json(request): Json, ) -> Result, ApiError> { - info!( - "OAuth credential-ready for source {} (user={:?}, provider={}, flow={})", - request.source_id, request.user_id, request.provider, request.flow - ); + info!("OAuth credential-ready"); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -1113,17 +950,11 @@ pub async fn oauth_credential_ready( { CredentialResolution::Resolved(creds) => creds, CredentialResolution::NeedsUserAuth { provider } => { - warn!( - "OAuth credential-ready for {}: user OAuth required for provider {:?}", - request.source_id, provider - ); + warn!("OAuth credential-ready: user OAuth required"); return Ok(Json(json!({"status": "missing_credentials"}))); } CredentialResolution::NoCredentials => { - warn!( - "OAuth credential-ready for {}: no credentials found", - request.source_id - ); + warn!("OAuth credential-ready: no credentials found"); return Ok(Json(json!({"status": "no_credentials"}))); } }; @@ -1147,7 +978,7 @@ pub async fn oauth_credential_ready( { Ok(Some(manifest)) => { if let Err(e) = validate_connector_manifest_action_schemas(&manifest) { - warn!("OAuth credential-ready returned invalid manifest: {}", e); + warn!("OAuth credential-ready returned invalid manifest"); return Ok(Json(json!({"status": "invalid_manifest"}))); } let manifest_json = @@ -1162,28 +993,19 @@ pub async fn oauth_credential_ready( .set_ex(&key, &manifest_json, REGISTRATION_TTL_SECONDS) .await .map_err(|e| ApiError::Internal(format!("Failed to store registration: {}", e)))?; - info!( - "OAuth credential-ready updated manifest for {}", - manifest.connector_id - ); + info!("OAuth credential-ready updated manifest"); Ok(Json( json!({"status": "completed", "catalog_updated": true}), )) } Ok(None) => { - info!( - "OAuth credential-ready delivered for {} (no manifest change)", - request.source_id - ); + info!("OAuth credential-ready delivered (no manifest change)"); Ok(Json( json!({"status": "delivered", "catalog_updated": false}), )) } Err(e) => { - warn!( - "OAuth credential-ready delivery failed for {}: {}", - request.source_id, e - ); + warn!("OAuth credential-ready delivery failed"); Ok(Json(json!({"status": "delivery_failed"}))) } } @@ -1246,7 +1068,7 @@ pub async fn get_skill( State(state): State, Json(request): Json, ) -> Result, ApiError> { - info!("Getting skill {}", request.skill_id); + info!("Getting skill"); let manifests = get_registered_manifests(&state.redis_client).await; for manifest in manifests { @@ -1537,8 +1359,8 @@ pub async fn sdk_register( let connector_id = &manifest.connector_id; info!( - "SDK: Registered connector '{}' (source_types: {:?}, url: {})", - connector_id, manifest.source_types, manifest.connector_url + "SDK: Registered connector '{}' (source_types: {:?})", + connector_id, manifest.source_types ); let manifest_json = serde_json::to_string(&manifest) @@ -1603,10 +1425,7 @@ pub async fn sdk_register( .find_any_user_oauth_for_provider(&source_type_strs, &provider) .await { - info!( - "Recovery: found OAuth credential for {} / {} to refresh missing MCP catalog", - source_id, provider - ); + info!("Recovery: found OAuth credential; attempting to refresh MCP catalog"); match resolve_credentials(&repo, &source_id, Some(&user_id), false).await { Ok(CredentialResolution::Resolved(recovery_creds)) => { match serde_json::to_value(McpCredentials::from_service_credential( @@ -1628,34 +1447,24 @@ pub async fn sdk_register( ) .await { - warn!("Recovery credential-ready delivery failed: {}", e); + warn!("Recovery credential-ready delivery failed"); } } Err(e) => { warn!( - "Recovery credential-ready skipped for {}: failed to serialize credentials: {}", - source_id, e + "Recovery credential-ready skipped: failed to serialize credentials" ); } } } - Ok(CredentialResolution::NeedsUserAuth { provider }) => { - warn!( - "Recovery credential-ready skipped for {}: user auth still required for provider {:?}", - source_id, provider - ); + Ok(CredentialResolution::NeedsUserAuth { .. }) => { + warn!("Recovery credential-ready skipped: user auth still required"); } Ok(CredentialResolution::NoCredentials) => { - warn!( - "Recovery credential-ready skipped for {}: no credentials found", - source_id - ); + warn!("Recovery credential-ready skipped: no credentials found"); } - Err(e) => { - warn!( - "Recovery credential-ready skipped for {}: credential resolution failed: {}", - source_id, e - ); + Err(_) => { + warn!("Recovery credential-ready skipped: credential resolution failed"); } } } @@ -1672,7 +1481,7 @@ pub async fn get_registered_manifests(redis_client: &redis::Client) -> Vec c, Err(e) => { - error!("Redis connection error: {}", e); + error!("Redis connection error"); return Vec::new(); } }; @@ -1931,10 +1740,7 @@ pub async fn sdk_emit_event( State(state): State, Json(request): Json, ) -> Result, ApiError> { - debug!( - "SDK: Emitting event for sync_run={}, source={}", - request.sync_run_id, request.source_id - ); + debug!("SDK: Emitting event"); let event_queue = EventQueue::new(state.db_pool.pool().clone()); @@ -1960,12 +1766,7 @@ pub async fn sdk_emit_batch( State(state): State, Json(request): Json, ) -> Result, ApiError> { - debug!( - "SDK: Emitting batch of {} events for sync_run={}, source={}", - request.events.len(), - request.sync_run_id, - request.source_id - ); + debug!("SDK: Emitting batch of {} events", request.events.len()); let event_queue = EventQueue::new(state.db_pool.pool().clone()); @@ -2046,17 +1847,14 @@ async fn extract_content( match extract_content_blocking(data, mime_type.clone(), filename.clone()).await { Ok(text) if is_pdf && text.trim().is_empty() => { warn!( - "PDF text extraction produced no text; sync_run_id={}, source_id={:?}, filename={:?}, mime_type={}", - sync_run_id, source_id, filename, mime_type, + "PDF text extraction produced no text; mime_type={}", + mime_type, ); Ok("[Text extraction failed for this PDF. The document was skipped for extracted-text indexing because no text could be extracted.]".to_string()) } Ok(text) => Ok(text), Err(e) if is_pdf => { - warn!( - "PDF text extraction failed; sync_run_id={}, source_id={:?}, filename={:?}, mime_type={}, reason={}", - sync_run_id, source_id, filename, mime_type, e, - ); + warn!("PDF text extraction failed; mime_type={}", mime_type,); Ok(format!( "[Text extraction failed for this PDF. The document was skipped for extracted-text indexing. Reason: {}]", e @@ -2156,8 +1954,8 @@ async fn do_extract_text( Some(client) => { let file_name = filename.as_deref().unwrap_or("document"); debug!( - "Using docling-based document content extraction for file '{}' (preset={})", - file_name, preset + "Using docling-based document content extraction (preset={})", + preset ); match client.convert(&data, file_name, &preset).await { Ok(markdown) => { @@ -2175,16 +1973,14 @@ async fn do_extract_text( retry_after_secs, }); } - Err(e) => { - warn!("Docling extraction failed, falling back to built-in: {}", e); + Err(_) => { + warn!("Docling extraction failed, falling back to built-in"); None } } } _ => { - warn!( - "Docling enabled but DOCLING_URL not set, falling back to built-in extraction" - ); + warn!("Docling enabled but DOCLING_URL not set, falling back"); None } }; @@ -2192,10 +1988,7 @@ async fn do_extract_text( if let Some(markdown) = docling_result { markdown } else { - debug!( - "Using built-in document content extraction for file {:?}", - filename - ); + debug!("Using built-in document content extraction"); extract_content( data, mime_type.clone(), @@ -2206,10 +1999,7 @@ async fn do_extract_text( .await? } } else { - debug!( - "Using built-in document content extraction for file {:?} (docling_enabled={}, docling_candidate={})", - filename, docling_enabled, docling_candidate - ); + debug!("Using built-in document content extraction"); extract_content( data, mime_type.clone(), @@ -2229,8 +2019,7 @@ async fn do_extract_text( let max_bytes = max_extracted_text_bytes(); if processed_text.len() > max_bytes { warn!( - "Truncating extracted content for {:?}: {} bytes > {} byte limit", - filename, + "Truncating extracted content: {} bytes > {} byte limit", processed_text.len(), max_bytes ); @@ -2246,10 +2035,8 @@ pub async fn sdk_extract_content( let fields = parse_extract_multipart(multipart).await?; debug!( - "SDK: Extracting content for sync_run={}, mime={}, filename={:?}, size={}", - fields.sync_run_id, + "SDK: Extracting content; mime={}, size={}", fields.mime_type, - fields.filename, fields.data.len() ); @@ -2303,10 +2090,8 @@ pub async fn sdk_extract_text( let fields = parse_extract_multipart(multipart).await?; debug!( - "SDK: Extracting text for sync_run={}, mime={}, filename={:?}, size={}", - fields.sync_run_id, + "SDK: Extracting text; mime={}, size={}", fields.mime_type, - fields.filename, fields.data.len() ); @@ -2342,7 +2127,7 @@ pub async fn sdk_store_content( State(state): State, Json(request): Json, ) -> Result, ApiError> { - debug!("SDK: Storing content for sync_run={}", request.sync_run_id); + debug!("SDK: Storing content"); let content_storage = state.content_storage.clone(); @@ -2360,8 +2145,7 @@ pub async fn sdk_store_content( let max_bytes = max_extracted_text_bytes(); if normalized_content.len() > max_bytes { warn!( - "Truncating stored content for sync_run={}: {} bytes > {} byte limit", - request.sync_run_id, + "Truncating stored content: {} bytes > {} byte limit", normalized_content.len(), max_bytes ); @@ -2386,7 +2170,7 @@ pub async fn sdk_heartbeat( State(state): State, Path(sync_run_id): Path, ) -> Result, ApiError> { - debug!("SDK: Heartbeat for sync_run={}", sync_run_id); + debug!("SDK: Heartbeat"); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); sync_run_repo @@ -2403,20 +2187,29 @@ pub async fn sdk_complete( State(state): State, Path(sync_run_id): Path, ) -> Result, ApiError> { - info!("SDK: Completing sync_run={}", sync_run_id); + info!("SDK: Completing sync"); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); + // Look up the sync run for sync_type and duration. + let sync_run = sync_run_repo.find_by_id(&sync_run_id).await.ok().flatten(); + let sync_type = sync_run + .as_ref() + .map(|r| r.sync_type.clone()) + .unwrap_or(SyncType::Incremental); + let created_at = sync_run.and_then(|r| Some(r.created_at)); + // Atomically mark completed and publish this run's checkpoint to the source. let updated = sync_run_repo .complete_and_publish_checkpoint(&sync_run_id) .await .map_err(|e| ApiError::Internal(format!("Failed to mark completed: {}", e)))?; if !updated { - warn!( - "SDK: Ignoring stale complete for non-running sync_run={}", - sync_run_id - ); + warn!("SDK: Ignoring stale complete for non-running sync"); + } + + if updated { + record_sync_terminal(sync_type, "completed", created_at); } Ok(Json(SdkStatusResponse { @@ -2429,20 +2222,29 @@ pub async fn sdk_fail( Path(sync_run_id): Path, Json(request): Json, ) -> Result, ApiError> { - info!("SDK: Failing sync_run={}: {}", sync_run_id, request.error); + info!("SDK: Failing sync"); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); + // Look up the sync run for sync_type and duration. + let sync_run = sync_run_repo.find_by_id(&sync_run_id).await.ok().flatten(); + let sync_type = sync_run + .as_ref() + .map(|r| r.sync_type.clone()) + .unwrap_or(SyncType::Incremental); + let created_at = sync_run.and_then(|r| Some(r.created_at)); + // Mark sync as failed let updated = sync_run_repo .mark_failed(&sync_run_id, &request.error) .await .map_err(|e| ApiError::Internal(format!("Failed to mark failed: {}", e)))?; if !updated { - warn!( - "SDK: Ignoring stale fail for non-running sync_run={}", - sync_run_id - ); + warn!("SDK: Ignoring stale fail for non-running sync"); + } + + if updated { + record_sync_terminal(sync_type, "failed", created_at); } Ok(Json(SdkStatusResponse { @@ -2455,10 +2257,7 @@ pub async fn sdk_increment_scanned( Path(sync_run_id): Path, Json(request): Json, ) -> Result, ApiError> { - debug!( - "SDK: Incrementing scanned for sync_run={} by {}", - sync_run_id, request.count - ); + debug!("SDK: Incrementing scanned by {}", request.count); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); let updated = sync_run_repo @@ -2466,10 +2265,7 @@ pub async fn sdk_increment_scanned( .await .map_err(|e| ApiError::Internal(format!("Failed to increment scanned: {}", e)))?; if !updated { - warn!( - "SDK: Ignoring stale scanned increment for non-running sync_run={}", - sync_run_id - ); + warn!("SDK: Ignoring stale scanned increment for non-running sync"); } Ok(Json(SdkStatusResponse { @@ -2482,10 +2278,7 @@ pub async fn sdk_increment_updated( Path(sync_run_id): Path, Json(request): Json, ) -> Result, ApiError> { - debug!( - "SDK: Incrementing updated for sync_run={} by {}", - sync_run_id, request.count - ); + debug!("SDK: Incrementing updated by {}", request.count); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); let updated = sync_run_repo @@ -2493,10 +2286,7 @@ pub async fn sdk_increment_updated( .await .map_err(|e| ApiError::Internal(format!("Failed to increment updated: {}", e)))?; if !updated { - warn!( - "SDK: Ignoring stale updated increment for non-running sync_run={}", - sync_run_id - ); + warn!("SDK: Ignoring stale updated increment for non-running sync"); } Ok(Json(SdkStatusResponse { @@ -2508,7 +2298,7 @@ pub async fn sdk_get_source( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - debug!("SDK: Getting source config for source_id={}", source_id); + debug!("SDK: Getting source config"); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -2524,7 +2314,7 @@ pub async fn sdk_get_credentials( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - debug!("SDK: Getting credentials for source_id={}", source_id); + debug!("SDK: Getting credentials"); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -2555,10 +2345,7 @@ pub async fn sdk_get_source_sync_config( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - debug!( - "SDK: Getting source sync config for source_id={}", - source_id - ); + debug!("SDK: Getting source sync config"); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -2593,10 +2380,7 @@ pub async fn sdk_create_sync( State(state): State, Json(request): Json, ) -> Result, ApiError> { - info!( - "SDK: Creating sync run for source={}, type={:?}", - request.source_id, request.sync_type - ); + info!("SDK: Creating sync run; type={:?}", request.sync_type); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -2658,18 +2442,32 @@ pub async fn sdk_cancel_sync( State(state): State, Json(request): Json, ) -> Result, ApiError> { - info!("SDK: Cancelling sync_run={}", request.sync_run_id); + info!("SDK: Cancelling sync"); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); + + // Look up the sync run for sync_type and duration. + let sync_run = sync_run_repo + .find_by_id(&request.sync_run_id) + .await + .ok() + .flatten(); + let sync_type = sync_run + .as_ref() + .map(|r| r.sync_type.clone()) + .unwrap_or(SyncType::Incremental); + let created_at = sync_run.and_then(|r| Some(r.created_at)); + let updated = sync_run_repo .mark_cancelled(&request.sync_run_id) .await .map_err(|e| ApiError::Internal(format!("Failed to cancel sync: {}", e)))?; if !updated { - warn!( - "SDK: Ignoring stale cancel for non-running sync_run={}", - request.sync_run_id - ); + warn!("SDK: Ignoring stale cancel for non-running sync"); + } + + if updated { + record_sync_terminal(sync_type, "cancelled", created_at); } Ok(Json(SdkCancelSyncResponse { success: true })) @@ -2679,7 +2477,7 @@ pub async fn sdk_get_user_email( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - debug!("SDK: Getting user email for source_id={}", source_id); + debug!("SDK: Getting user email"); let email = sqlx::query_scalar::<_, String>( "SELECT u.email FROM sources s JOIN users u ON s.created_by = u.id WHERE s.id = $1", @@ -2697,8 +2495,8 @@ pub async fn sdk_notify_webhook( Json(request): Json, ) -> Result, ApiError> { info!( - "SDK: Webhook notification for source={}, event_type={}", - request.source_id, request.event_type + "SDK: Webhook notification; event_type={}", + request.event_type ); // Trigger a sync for this source (connector-manager handles sync run creation) @@ -2720,7 +2518,7 @@ pub async fn sdk_update_checkpoint( Path(sync_run_id): Path, Json(checkpoint): Json, ) -> Result, ApiError> { - debug!("SDK: Updating checkpoint for sync_run={}", sync_run_id); + debug!("SDK: Updating checkpoint"); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); let updated = sync_run_repo @@ -2728,10 +2526,7 @@ pub async fn sdk_update_checkpoint( .await .map_err(|e| ApiError::Internal(format!("Failed to update checkpoint: {}", e)))?; if !updated { - warn!( - "SDK: Ignoring stale checkpoint update for non-running sync_run={}", - sync_run_id - ); + warn!("SDK: Ignoring stale checkpoint update for non-running sync"); } Ok(Json(SdkStatusResponse { @@ -2748,7 +2543,7 @@ pub async fn sdk_update_connector_state( Path(source_id): Path, Json(new_state): Json, ) -> Result, ApiError> { - debug!("SDK: Updating connector state for source_id={}", source_id); + debug!("SDK: Updating connector state"); let source_repo = SourceRepository::new(state.db_pool.pool()); source_repo @@ -2769,7 +2564,7 @@ pub async fn sdk_get_sources_by_type( State(state): State, Path(source_type): Path, ) -> Result>, ApiError> { - debug!("SDK: Getting sources by type={}", source_type); + debug!("SDK: Getting sources by type"); let source_repo = SourceRepository::new(state.db_pool.pool()); let sources = source_repo @@ -2790,7 +2585,7 @@ pub async fn sdk_get_connector_config( State(state): State, Path(provider): Path, ) -> Result, ApiError> { - debug!("SDK: Getting connector config for provider={}", provider); + debug!("SDK: Getting connector config"); let repo = shared::ConnectorConfigRepository::new(state.db_pool.pool().clone()); let config = repo diff --git a/services/connector-manager/src/lib.rs b/services/connector-manager/src/lib.rs index ccb0d9f18..3775daa0b 100644 --- a/services/connector-manager/src/lib.rs +++ b/services/connector-manager/src/lib.rs @@ -26,7 +26,7 @@ use sync_manager::SyncManager; use tokio::sync::Semaphore; use tower::ServiceBuilder; use tower_http::cors::CorsLayer; -use tracing::{info, warn}; +use tracing::{error, info, warn}; #[derive(Clone)] pub struct AppState { @@ -185,7 +185,16 @@ pub async fn run_server() -> AnyhowResult<()> { info!("Connector Manager service listening on {}", addr); let listener = tokio::net::TcpListener::bind(addr).await?; - axum::serve(listener, app).await?; + if let Err(e) = axum::serve(listener, app) + .with_graceful_shutdown(telemetry::shutdown_signal()) + .await + { + error!("Server error: {}", e); + telemetry::shutdown_telemetry().await; + return Err(e.into()); + } + + telemetry::shutdown_telemetry().await; Ok(()) } diff --git a/services/connector-manager/src/sync_manager.rs b/services/connector-manager/src/sync_manager.rs index ae6250750..874cfbbcb 100644 --- a/services/connector-manager/src/sync_manager.rs +++ b/services/connector-manager/src/sync_manager.rs @@ -6,6 +6,7 @@ use dashmap::DashMap; use redis::Client as RedisClient; use shared::db::error::DatabaseError; use shared::db::repositories::SyncRunRepository; +use shared::metrics; use shared::models::{SourceType, SyncSlotClass, SyncStatus, SyncType}; use shared::{DatabasePool, Repository, SourceRepository}; use sqlx::PgPool; @@ -15,7 +16,7 @@ use std::time::Duration; use time::format_description::well_known::Rfc3339; use time::OffsetDateTime; use tokio::time::timeout; -use tracing::{debug, error, info, warn}; +use tracing::{error, info, warn}; const MAX_RESUME_ATTEMPTS: usize = 3; const MISSING_MANIFEST_GRACE_OBSERVATIONS: usize = 2; @@ -107,10 +108,7 @@ impl SyncManager { let effective_sync_type = match sync_type { SyncType::Incremental if last_completed.is_none() => { - info!( - "No prior completed sync for source {}; upgrading to full sync", - source_id - ); + info!("No prior completed sync; upgrading to full sync"); SyncType::Full } other => other, @@ -136,15 +134,17 @@ impl SyncManager { other => SyncError::DatabaseError(other.to_string()), })?; - debug!( - sync_run_id = %sync_run.id, - source_id = %source_id, - sync_type = ?effective_sync_type, - trigger_type = %trigger_type, - connector_url = %connector_url, - "Created sync_run; triggering connector" + // Record sync started metric + metrics::CONNECTOR_SYNC_STARTED.add( + 1, + &[opentelemetry::KeyValue::new( + "sync_type", + effective_sync_type.to_string(), + )], ); + info!("Created sync_run; triggering connector"); + let sync_request = SyncRequest { sync_run_id: sync_run.id.clone(), source_id: source_id.to_string(), @@ -163,10 +163,7 @@ impl SyncManager { match trigger_result { Ok(Ok(response)) => { - info!( - "Sync triggered for source {}: {:?}", - source_id, response.status - ); + info!("Sync triggered: {:?}", response.status); Ok(sync_run.id) } Ok(Err(ClientError::ConnectorError { status: 404, .. })) @@ -204,6 +201,9 @@ impl SyncManager { .map_err(|e| SyncError::DatabaseError(e.to_string()))? .ok_or_else(|| SyncError::SyncRunNotFound(sync_run_id.to_string()))?; + let sync_type = sync_run.sync_type.clone(); + let created_at = Some(sync_run.created_at); + if sync_run.status != SyncStatus::Running { return Err(SyncError::SyncNotRunning(sync_run_id.to_string())); } @@ -225,7 +225,7 @@ impl SyncManager { .cancel_sync(&connector_url, sync_run_id) .await { - warn!("Failed to send cancel request to connector: {}", e); + warn!("Failed to send cancel request"); } let updated = self @@ -239,7 +239,8 @@ impl SyncManager { self.resume_attempts.remove(sync_run_id); self.missing_manifest_observations.remove(sync_run_id); - info!("Sync {} cancelled", sync_run_id); + shared::metrics::record_sync_terminal(&sync_type.to_string(), "cancelled", created_at); + info!("Sync cancelled"); Ok(()) } @@ -347,18 +348,15 @@ impl SyncManager { if observations < MISSING_MANIFEST_GRACE_OBSERVATIONS { warn!( - "No registered connector for sync {} (source_type={:?}); deferring lost-sync handling for grace observation {}/{}", - sync_run.id, - source.source_type, - observations, - MISSING_MANIFEST_GRACE_OBSERVATIONS, + "No registered connector for sync; deferring lost-sync handling for grace observation {}/{}", + observations, MISSING_MANIFEST_GRACE_OBSERVATIONS, ); continue; } warn!( - "No registered connector for sync {} (source_type={:?}) after {} observations; treating as lost", - sync_run.id, source.source_type, observations, + "No registered connector for sync after {} observations; treating as lost", + observations, ); self.handle_lost_sync(&sync_run.id, &sync_run.source_id) .await; @@ -377,10 +375,7 @@ impl SyncManager { self.resume_attempts.remove(&sync_run.id); } Ok(_) => { - warn!( - "Connector reports sync {} no longer running; reconciling", - sync_run.id - ); + warn!("Connector reports sync no longer running; reconciling"); self.handle_lost_sync(&sync_run.id, &sync_run.source_id) .await; } @@ -393,10 +388,7 @@ impl SyncManager { // Connector reachable-via-Redis but not responding // (connection refused, timeout, 5xx). Treat same as // "lost" — attempt counter absorbs transient blips. - warn!( - "Sync status probe failed for {} ({}): {}; treating as lost", - sync_run.id, connector_url, e - ); + warn!("Sync status probe failed; treating as lost"); self.handle_lost_sync(&sync_run.id, &sync_run.source_id) .await; } @@ -424,8 +416,8 @@ impl SyncManager { if attempts > MAX_RESUME_ATTEMPTS { warn!( - "Sync {} exceeded {} resume attempts; marking failed", - sync_run_id, MAX_RESUME_ATTEMPTS + "Sync exceeded {} resume attempts; marking failed", + MAX_RESUME_ATTEMPTS ); if let Err(e) = self .mark_sync_failed( @@ -434,7 +426,7 @@ impl SyncManager { ) .await { - error!("Failed to mark sync {} as failed: {}", sync_run_id, e); + error!("Failed to mark sync as failed"); } self.resume_attempts.remove(sync_run_id); self.missing_manifest_observations.remove(sync_run_id); @@ -447,8 +439,8 @@ impl SyncManager { { Ok(Some(s)) => s, Ok(None) => return, - Err(e) => { - error!("Failed to load source {}: {}", source_id, e); + Err(_) => { + error!("Failed to load source"); return; } }; @@ -461,8 +453,8 @@ impl SyncManager { // above; after MAX_RESUME_ATTEMPTS the cap above will // mark the row failed. warn!( - "Cannot resume sync {} — connector for {:?} not registered (attempt {}/{})", - sync_run_id, source.source_type, attempts, MAX_RESUME_ATTEMPTS + "Cannot resume sync — connector for {:?} not registered (attempt {}/{})", + source.source_type, attempts, MAX_RESUME_ATTEMPTS ); return; } @@ -473,8 +465,8 @@ impl SyncManager { let sync_run = match self.sync_run_repo.find_by_id(sync_run_id).await { Ok(Some(r)) => r, Ok(None) => return, - Err(e) => { - error!("Failed to load sync_run {}: {}", sync_run_id, e); + Err(_) => { + error!("Failed to load sync_run"); return; } }; @@ -513,28 +505,24 @@ impl SyncManager { { Ok(Ok(_)) => { info!( - "Auto-resumed sync {} on connector (attempt {}/{})", - sync_run_id, attempts, MAX_RESUME_ATTEMPTS + "Auto-resumed sync on connector (attempt {}/{})", + attempts, MAX_RESUME_ATTEMPTS ); // Reset staleness clock so detect_stale_syncs doesn't fire // before the resumed sync starts emitting. if let Err(e) = self.sync_run_repo.update_activity(sync_run_id).await { - warn!( - "Failed to bump activity for resumed sync {}: {}", - sync_run_id, e - ); + warn!("Failed to bump activity for resumed sync"); } } Ok(Err(e)) => { warn!( - "Failed to re-trigger sync {} on connector (attempt {}/{}): {}", - sync_run_id, attempts, MAX_RESUME_ATTEMPTS, e + "Failed to re-trigger sync on connector (attempt {}/{})", + attempts, MAX_RESUME_ATTEMPTS ); } Err(_) => { warn!( - "Timed out re-triggering sync {} on connector after {}s (attempt {}/{})", - sync_run_id, + "Timed out re-triggering sync on connector after {}s (attempt {}/{})", CONNECTOR_TRIGGER_TIMEOUT.as_secs(), attempts, MAX_RESUME_ATTEMPTS @@ -575,10 +563,7 @@ impl SyncManager { } for run in &running { - info!( - "Cancelling sync {} for inactive/deleted source {}", - run.id, run.source_id - ); + info!("Cancelling sync for inactive/deleted source"); if let Some(&source_type) = source_type_map.get(&run.source_id) { if let Some(connector_url) = get_connector_url_for_source(&self.redis_client, source_type).await @@ -588,20 +573,31 @@ impl SyncManager { .cancel_sync(&connector_url, &run.id) .await { - warn!( - "Failed to cancel sync {} on connector for inactive source {}: {}", - run.id, run.source_id, e - ); + warn!("Failed to cancel sync on connector for inactive source"); } } } } - let run_ids: Vec = running.into_iter().map(|r| r.id).collect(); - self.sync_run_repo + let run_ids: Vec = running.iter().map(|r| r.id.clone()).collect(); + let cancelled_ids = self + .sync_run_repo .mark_cancelled_many(&run_ids, "Source was disabled") .await - .map_err(|e| SyncError::DatabaseError(e.to_string())) + .map_err(|e| SyncError::DatabaseError(e.to_string()))?; + + // Only record metrics for runs that were actually transitioned. + for run in &running { + if cancelled_ids.contains(&run.id) { + shared::metrics::record_sync_terminal( + &run.sync_type.to_string(), + "cancelled", + Some(run.created_at), + ); + } + } + + Ok(cancelled_ids) } /// Sweep running syncs whose `last_activity_at` hasn't advanced within the @@ -632,10 +628,7 @@ impl SyncManager { let mut marked_stale = Vec::new(); for (sync_run_id, source_id, source_type) in stale_syncs { - warn!( - "Marking stale sync {} for source {}", - sync_run_id, source_id - ); + warn!("Marking stale sync"); if let Some(connector_url) = get_connector_url_for_source(&self.redis_client, source_type).await @@ -645,18 +638,15 @@ impl SyncManager { .cancel_sync(&connector_url, &sync_run_id) .await { - warn!( - "Failed to cancel stale sync {} on connector: {}", - sync_run_id, e - ); + warn!("Failed to cancel stale sync on connector"); } } - if let Err(e) = self + if let Err(_) = self .mark_sync_failed(&sync_run_id, "Sync timed out (no activity detected)") .await { - error!("Failed to mark sync as stale: {}", e); + error!("Failed to mark sync as stale"); continue; } @@ -667,16 +657,31 @@ impl SyncManager { } async fn mark_sync_failed(&self, sync_run_id: &str, error: &str) -> Result<(), SyncError> { + let sync_type = self + .sync_run_repo + .find_by_id(sync_run_id) + .await + .ok() + .flatten() + .map(|r| r.sync_type.clone()) + .unwrap_or(SyncType::Incremental); + let created_at = self + .sync_run_repo + .find_by_id(sync_run_id) + .await + .ok() + .flatten() + .map(|r| r.created_at); + let updated = self .sync_run_repo .mark_failed(sync_run_id, error) .await .map_err(|e| SyncError::DatabaseError(e.to_string()))?; if !updated { - warn!( - "Ignoring stale failure transition for non-running sync {}", - sync_run_id - ); + warn!("Ignoring stale failure transition for non-running sync"); + } else { + shared::metrics::record_sync_terminal(&sync_type.to_string(), "failed", created_at); } self.resume_attempts.remove(sync_run_id); @@ -685,11 +690,32 @@ impl SyncManager { } async fn mark_sync_unavailable(&self, sync_run_id: &str) -> Result<(), SyncError> { - self.sync_run_repo + let sync_type = self + .sync_run_repo + .find_by_id(sync_run_id) + .await + .ok() + .flatten() + .map(|r| r.sync_type.clone()) + .unwrap_or(SyncType::Incremental); + let created_at = self + .sync_run_repo + .find_by_id(sync_run_id) + .await + .ok() + .flatten() + .map(|r| r.created_at); + + let updated = self + .sync_run_repo .mark_cancelled_with_message(sync_run_id, "Realtime sync not available for this source") .await .map_err(|e| SyncError::DatabaseError(e.to_string()))?; + if updated { + shared::metrics::record_sync_terminal(&sync_type.to_string(), "cancelled", created_at); + } + self.resume_attempts.remove(sync_run_id); self.missing_manifest_observations.remove(sync_run_id); Ok(()) diff --git a/services/connector-manager/tests/sanitization_test.rs b/services/connector-manager/tests/sanitization_test.rs new file mode 100644 index 000000000..56d1dfd5e --- /dev/null +++ b/services/connector-manager/tests/sanitization_test.rs @@ -0,0 +1,166 @@ +/// Sanitization regression tests for connector-manager logs. +/// +/// These tests verify that no raw IDs, response bodies, connector URLs, +/// or user identifiers appear in log messages. They scan the actual +/// source files, not runtime behaviour. +/// +/// NOTE: functional non-log code is excluded via a "log-only" heuristic: +/// we only match lines that contain one of the recognised log macros +/// (info!, warn!, error!, debug!, trace!). +use std::fs; +use std::path::Path; + +/// Check that a file path is a Rust source file inside the src/ directory. +fn is_rust_source(path: &Path) -> bool { + path.extension().map(|e| e == "rs").unwrap_or(false) + && path.components().any(|c| c.as_os_str() == "src") +} + +/// Check that a line appears to be a log macro invocation (or a continuation line +/// of one) rather than functional non-log code. +fn is_log_or_continuation(line: &str, in_log_call: &mut bool) -> bool { + let trimmed = line.trim(); + if *in_log_call { + // Keep scanning as continuation if it ends with a comma or is a string + // continuation. Stop when we see a closing paren or semicolon. + if trimmed.starts_with(')') || trimmed.starts_with("};") || trimmed.ends_with(';') { + *in_log_call = false; + return true; + } + return true; + } + if trimmed.starts_with("info!(") + || trimmed.starts_with("warn!(") + || trimmed.starts_with("error!(") + || trimmed.starts_with("debug!(") + || trimmed.starts_with("trace!(") + { + *in_log_call = !trimmed.contains(';') && !trimmed.ends_with(')'); + return true; + } + false +} + +/// Scan a single file for forbidden raw-value patterns in log lines only. +/// Returns a list of (line_number, line_text) for matches. +fn scan_log_forbidden(path: &Path, forbidden_patterns: &[&str]) -> Vec<(usize, String)> { + let content = fs::read_to_string(path).unwrap(); + let mut hits = Vec::new(); + let mut in_log_call = false; + + for (i, line) in content.lines().enumerate() { + if !is_log_or_continuation(line, &mut in_log_call) { + continue; + } + // Extend continuation tracking: check for open paren without matching close + // on the same line. + let trimmed = line.trim(); + if !in_log_call && trimmed.starts_with("info!(") + || trimmed.starts_with("warn!(") + || trimmed.starts_with("error!(") + || trimmed.starts_with("debug!(") + { + let has_close = trimmed + .rfind(')') + .map(|p| { + // Count parens + let opens = trimmed[..p].matches('(').count(); + let closes = trimmed[..p].matches(')').count(); + opens == closes + 1 + }) + .unwrap_or(false); + if !has_close { + in_log_call = true; + } + } + + for pat in forbidden_patterns { + if line.contains(pat) { + hits.push((i + 1, line.to_string())); + break; + } + } + } + + hits +} + +const FORBIDDEN_CONNECTOR_CLIENT: &[&str] = &[ + "response.text()", // reading full response body + "body", // response body variable in error log + "{} - {}", // pattern with status + body (in format args) + "connector_url", // raw connector URL in log + "error = %e", // error display may embed URL/path/body +]; + +const FORBIDDEN_LOGS: &[&str] = &[ + "source_id", // source identifier in log message + "sync_run_id", // sync run identifier in log message + "user_email", // user email in log message + ".email", // email field access in log + ".id", // generic ID field access in log + "connector_url", // raw URL in log message +]; + +const FORBIDDEN_HANDLERS_SYNC: &[&str] = &[ + "request.source_id", // source ID as format param in log call + "request.sync_run_id", // sync run ID as format param + "fields.sync_run_id", // sync run ID in extract handler + "source_id", // source variable in log message + "sync_run_id", // sync run variable in log message +]; + +fn resolve_path(name: &str) -> std::path::PathBuf { + let candidates = vec![ + Path::new(name).to_path_buf(), + Path::new("services/connector-manager").join(name), + ]; + for c in &candidates { + if c.exists() { + return c.clone(); + } + } + panic!("{} not found in any candidate path", name); +} + +#[test] +fn test_handlers_sanitization() { + let path = resolve_path("src/handlers.rs"); + let hits = scan_log_forbidden(&path, FORBIDDEN_HANDLERS_SYNC); + assert!( + hits.is_empty(), + "Found forbidden patterns in handlers.rs log lines:\n{}", + hits.iter() + .map(|(n, l)| format!(" line {}: {}", n, l.trim())) + .collect::>() + .join("\n") + ); +} + +#[test] +fn test_sync_manager_sanitization() { + let path = resolve_path("src/sync_manager.rs"); + let hits = scan_log_forbidden(&path, FORBIDDEN_HANDLERS_SYNC); + assert!( + hits.is_empty(), + "Found forbidden patterns in sync_manager.rs log lines:\n{}", + hits.iter() + .map(|(n, l)| format!(" line {}: {}", n, l.trim())) + .collect::>() + .join("\n") + ); +} + +#[test] +fn test_connector_client_sanitization() { + let path = resolve_path("src/connector_client.rs"); + let hits = scan_log_forbidden(&path, FORBIDDEN_CONNECTOR_CLIENT); + assert!( + hits.is_empty(), + "Found forbidden patterns in connector_client.rs log lines:\n{}", + hits.iter() + .map(|(n, l)| format!(" line {}: {}", n, l.trim())) + .collect::>() + .join("\n") + ); +} diff --git a/services/indexer/Cargo.toml b/services/indexer/Cargo.toml index c1e48f817..f29def441 100644 --- a/services/indexer/Cargo.toml +++ b/services/indexer/Cargo.toml @@ -35,6 +35,7 @@ ulid = { workspace = true } bytes = "1.0" futures = "0.3" num_cpus = "1.0" +tracing-opentelemetry = { workspace = true } shared = { path = "../../shared" } [dev-dependencies] diff --git a/services/indexer/src/lib.rs b/services/indexer/src/lib.rs index 2ccc9d324..b1d9e62bb 100644 --- a/services/indexer/src/lib.rs +++ b/services/indexer/src/lib.rs @@ -10,8 +10,8 @@ pub use axum::Router; pub use redis::Client as RedisClient; pub use serde::{Deserialize, Serialize}; pub use serde_json::Value; -pub use shared::AIClient; pub use shared::db::pool::DatabasePool; +pub use shared::AIClient; use std::sync::Arc; use axum::{ @@ -23,11 +23,11 @@ use axum::{ use error::Result as IndexerResult; use serde_json::json; use shared::{ - IndexerConfig, db::repositories::{DocumentRepository, OrphanStats}, models::Document, storage::gc::{ContentBlobGC, GCConfig, GCResult}, telemetry::{self, TelemetryConfig}, + IndexerConfig, }; use sqlx::types::time::OffsetDateTime; use std::net::SocketAddr; @@ -159,7 +159,7 @@ async fn create_document( let repo = DocumentRepository::new(state.db_pool.pool()); let document = repo.create(doc).await?; - info!("Created document: {}", document_id); + info!("Created document"); Ok(Json(document)) } @@ -208,7 +208,7 @@ async fn update_document( match updated_doc { Some(doc) => { - info!("Updated document: {}", id); + info!("Updated document"); Ok(Json(doc)) } None => Err(error::IndexerError::NotFound(format!( @@ -232,7 +232,7 @@ async fn delete_document( ))); } - info!("Deleted document: {}", id); + info!("Deleted document"); Ok(Json(json!({ "message": "Document deleted successfully", "id": id @@ -475,28 +475,38 @@ pub async fn run_server() -> anyhow::Result<()> { let app = create_app(app_state.clone()); let queue_processor = queue_processor::QueueProcessor::new(app_state.clone()); - let processor_handle = tokio::spawn(async move { - if let Err(e) = queue_processor.start().await { - error!("Queue processor failed: {}", e); + let mut processor_handle = tokio::spawn(async move { + if let Err(_) = queue_processor.start().await { + error!("Queue processor failed"); } }); let addr = SocketAddr::from(([0, 0, 0, 0], config.port)); - info!("Indexer service listening on {}", addr); + info!("Indexer service listening"); let listener = tokio::net::TcpListener::bind(addr).await?; + // Use graceful shutdown so the HTTP server drains in-flight requests + // on SIGTERM/Ctrl-C before we abort the processor and flush telemetry. + let serve = axum::serve(listener, app).with_graceful_shutdown(telemetry::shutdown_signal()); + tokio::select! { - result = axum::serve(listener, app) => { - if let Err(e) = result { - error!("HTTP server failed: {}", e); + result = serve => { + if let Err(_) = result { + error!("HTTP server failed"); } } - _ = processor_handle => { + _ = &mut processor_handle => { error!("Event processor task completed unexpectedly"); } } + // If the processor is still running (HTTP server finished first), abort it + // and await the JoinHandle so cancellation completes before telemetry shutdown. + processor_handle.abort(); + let _ = processor_handle.await; + + telemetry::shutdown_telemetry().await; Ok(()) } @@ -504,22 +514,19 @@ async fn debug_create_document( State(_state): State, body: String, ) -> IndexerResult> { - info!("Raw request body: {}", body); info!("Body length: {}", body.len()); match serde_json::from_str::(&body) { Ok(req) => { info!( - "Successfully parsed request: source_id='{}' ({}), external_id='{}' ({})", - req.source_id, + "Successfully parsed request: source_id_len={}, external_id_len={}", req.source_id.len(), - req.external_id, req.external_id.len() ); Ok(Json(json!({"status": "parsed successfully"}))) } Err(e) => { - error!("Failed to parse request: {}", e); + error!("Failed to parse request"); Ok(Json(json!({"error": format!("Parse error: {}", e)}))) } } diff --git a/services/indexer/src/queue_processor.rs b/services/indexer/src/queue_processor.rs index d18b0d8c2..82c24e9e2 100644 --- a/services/indexer/src/queue_processor.rs +++ b/services/indexer/src/queue_processor.rs @@ -1,10 +1,11 @@ -use crate::AppState; use crate::people_extractor; +use crate::AppState; use anyhow::{Context, Result}; use shared::db::repositories::{ DocumentRepository, GroupRepository, PersonRepository, SyncRunRepository, }; use shared::embedding_queue::EmbeddingQueue; +use shared::metrics; use shared::models::{ ConnectorEvent, ConnectorEventQueueItem, Document, DocumentAttributes, DocumentMetadata, DocumentPermissions, EventStatus, SyncType, @@ -14,8 +15,9 @@ use shared::storage::gc::{ContentBlobGC, GCConfig}; use std::collections::HashMap; use std::sync::Arc; use tokio::sync::{Mutex, Semaphore}; -use tokio::time::{Duration, MissedTickBehavior, interval}; -use tracing::{debug, error, info, warn}; +use tokio::time::{interval, Duration, MissedTickBehavior}; +use tracing::{debug, error, info, warn, Instrument}; +use tracing_opentelemetry::OpenTelemetrySpanExt; // Default poll interval for draining the queue. Overridable via INDEXER_POLL_INTERVAL_SECS. // SDK-side buffering already shapes events into the right batch size per sync type, @@ -450,12 +452,19 @@ impl QueueProcessor { "Queue stats - Pending: {}, Processing: {}, Completed: {}, Failed: {}, Dead Letter: {}", stats.pending, stats.processing, stats.completed, stats.failed, stats.dead_letter ); + metrics::record_queue_status( + stats.pending, + stats.processing, + stats.failed, + stats.dead_letter, + ); } } _ = retry_interval.tick() => { if let Ok(retried) = self.event_queue.retry_failed_events().await { if retried > 0 { info!("Retried {} failed events", retried); + metrics::INDEXER_EVENTS_RETRIED.add(retried as u64, &[]); } } } @@ -511,8 +520,8 @@ impl QueueProcessor { ); } } - Err(e) => { - error!("Content blob GC failed: {}", e); + Err(_) => { + error!("Content blob GC failed"); } } }); @@ -578,7 +587,7 @@ impl QueueProcessor { break; } batches_dequeued += 1; - total_processed += self.process_dequeued_events(events).await?; + total_processed += self.process_dequeued_events_instrumented(events).await?; } for (sync_type, reason) in ready { @@ -605,7 +614,7 @@ impl QueueProcessor { break; } batches_dequeued += 1; - total_processed += self.process_dequeued_events(events).await?; + total_processed += self.process_dequeued_events_instrumented(events).await?; } } @@ -618,6 +627,48 @@ impl QueueProcessor { Ok(()) } + /// Instrument `process_dequeued_events` with a CONSUMER tracing span. + /// + /// The span is created as a new root trace (no parent) with bounded messaging + /// attributes and native OTel links pointing at the producer spans. + async fn process_dequeued_events_instrumented( + &self, + events: Vec, + ) -> Result { + let event_count = events.len() as i64; + + // Extract producer SpanContext values deduplicated by (trace_id, span_id). + let producer_contexts: Vec, Option)>> = events + .iter() + .map(|ev| Some((ev.traceparent.clone(), ev.tracestate.clone()))) + .collect(); + let span_contexts = shared::telemetry::queue::collect_span_contexts(&producer_contexts); + + let span = tracing::info_span!( + parent: None, + "connector_events_queue process", + otel.kind = "CONSUMER", + messaging.system = "postgresql", + messaging.destination = "connector_events_queue", + messaging.operation.type = "process", + messaging.batch.message_count = event_count, + ); + + // Initialize the OTel span and add native links before entering + // the async work. Clone the span so the original handle is available + // for `.instrument()` after adding links. + { + let _guard = span.clone().entered(); + for sc in &span_contexts { + span.add_link(sc.clone()); + } + } + + async move { self.process_dequeued_events(events).await } + .instrument(span) + .await + } + async fn process_dequeued_events(&self, events: Vec) -> Result { if events.is_empty() { return Ok(0); @@ -673,36 +724,50 @@ impl QueueProcessor { match result { Ok(batch_result) => { - if !batch_result.successful_event_ids.is_empty() { - if let Err(e) = self + let batch_events_count = events_clone.len(); + let successful_count = batch_result.successful_event_ids.len(); + let failed_count = batch_result.failed_events.len(); + + // Mark completed only after DB transition succeeds. + // Use the actual transitioned count (rows that were in + // 'processing' state) to avoid stale/duplicate inflation. + if successful_count > 0 { + match self .event_queue .mark_events_completed_batch(batch_result.successful_event_ids.clone()) .await { - error!( - "Failed to mark {} events as completed: {}", - batch_result.successful_event_ids.len(), - e - ); + Ok(completed) if completed > 0 => { + metrics::INDEXER_EVENTS_PROCESSED.add(completed as u64, &[]); + } + Ok(_) => {} + Err(e) => { + error!("Failed to mark events as completed: {}", e); + } } } - if !batch_result.failed_events.is_empty() { - if let Err(e) = self + // Mark dead-letter / retryable failed only after DB transition succeeds. + // `dead_letter` rows (exhausted retries) go to the dead-letter counter; + // `failed` rows (retryable) do not count as dead-letter. + if failed_count > 0 { + match self .event_queue .mark_events_dead_letter_batch(batch_result.failed_events.clone()) .await { - error!( - "Failed to mark {} events as failed: {}", - batch_result.failed_events.len(), - e - ); + Ok((_retryable, dead_letter)) if dead_letter > 0 => { + metrics::INDEXER_EVENTS_DEAD_LETTER.add(dead_letter as u64, &[]); + } + Ok(_) => {} + Err(e) => { + error!("Failed to mark events as dead-letter: {}", e); + } } } if batch_result.successful_documents_count > 0 { - if let Err(e) = self + if let Err(_) = self .sync_run_repo .increment_progress_by( &batch_sync_run_id, @@ -710,25 +775,28 @@ impl QueueProcessor { ) .await { - warn!( - "Failed to update sync run progress for {}: {}", - batch_sync_run_id, e - ); + warn!("Failed to update sync run progress"); } } + let batch_duration_val = batch_start_time.elapsed(); + metrics::INDEXER_BATCH_DURATION.record(batch_duration_val.as_secs_f64(), &[]); + metrics::INDEXER_BATCH_SIZE.record(batch_events_count as u64, &[]); + self.extract_and_upsert_people(&events_clone).await; - total_processed += batch_result.successful_event_ids.len(); + total_processed += successful_count; - let batch_duration = batch_start_time.elapsed(); info!( "Sync-run batch processing completed: {} successful, {} failed (took {:?}, {:.1} events/sec)", - batch_result.successful_event_ids.len(), - batch_result.failed_events.len(), - batch_duration, - batch_result.successful_event_ids.len() as f64 - / batch_duration.as_secs_f64() + successful_count, + failed_count, + batch_duration_val, + if batch_duration_val.as_secs_f64() > 0.0 { + successful_count as f64 / batch_duration_val.as_secs_f64() + } else { + 0.0 + }, ); } Err(e) => { @@ -738,14 +806,21 @@ impl QueueProcessor { .iter() .map(|ev| (ev.id.clone(), err_msg.clone())) .collect(); - if let Err(mark_err) = - self.event_queue.mark_events_dead_letter_batch(failed).await - { - error!( - "Failed to mark {} events as failed after batch error: {}", - events_clone.len(), - mark_err - ); + + // Mark dead-letter / retryable failed; only record dead-letter + // metric after DB confirms, and only for rows that actually + // transitioned to dead_letter (exhausted retries). + match self.event_queue.mark_events_dead_letter_batch(failed).await { + Ok((_retryable, dead_letter)) if dead_letter > 0 => { + metrics::INDEXER_EVENTS_DEAD_LETTER.add(dead_letter as u64, &[]); + } + Ok(_) => {} + Err(e) => { + error!( + "Failed to mark events as dead-letter after batch error: {}", + e + ); + } } } } @@ -1055,8 +1130,8 @@ impl QueueProcessor { Ok(_) => { debug!("Upserted {} people from batch", count); } - Err(e) => { - error!("Failed to upsert people: {}", e); + Err(_) => { + error!("Failed to upsert people"); } } } diff --git a/services/migrations/106_add_queue_trace_context.sql b/services/migrations/106_add_queue_trace_context.sql new file mode 100644 index 000000000..beb2f29a3 --- /dev/null +++ b/services/migrations/106_add_queue_trace_context.sql @@ -0,0 +1,18 @@ +-- Migration 106: Add W3C trace context columns to connector_events_queue and embedding_queue. +-- +-- This migration adds nullable traceparent and tracestate columns to support +-- W3C TraceContext propagation through the asynchronous queue pipeline. +-- +-- Columns: +-- traceparent (varchar(55)) — W3C traceparent header (version-format-trace_id-span_id-trace_flags) +-- tracestate (varchar(512)) — optional W3C tracestate header (max 512 chars per spec) +-- +-- Both columns are nullable; NULL means "no stored trace context". +-- Invalid/missing values are handled gracefully at read time. +-- No indexes: queue consumers read trace context after dequeue, not for filtering. + +ALTER TABLE connector_events_queue ADD COLUMN IF NOT EXISTS traceparent varchar(55); +ALTER TABLE connector_events_queue ADD COLUMN IF NOT EXISTS tracestate varchar(512); + +ALTER TABLE embedding_queue ADD COLUMN IF NOT EXISTS traceparent varchar(55); +ALTER TABLE embedding_queue ADD COLUMN IF NOT EXISTS tracestate varchar(512); diff --git a/services/searcher/src/handlers.rs b/services/searcher/src/handlers.rs index 5a0c2fa5b..ec17b8398 100644 --- a/services/searcher/src/handlers.rs +++ b/services/searcher/src/handlers.rs @@ -20,15 +20,16 @@ use futures_util::Stream; use redis::AsyncCommands; use serde_json::{json, Value}; use shared::{ - models::UserConfiguration, ConfigurationRepository, DocumentRepository, GroupRepository, + metrics, models::UserConfiguration, ConfigurationRepository, DocumentRepository, GroupRepository, PersonRepository, Repository, UserRepository, }; use sqlx::types::time::OffsetDateTime; use std::pin::Pin; use std::sync::Arc; use std::task::{Context, Poll}; +use std::time::Instant; use tokio::sync::Mutex; -use tracing::{debug, error, info}; +use tracing::{error, info}; /// A stream wrapper that collects chunks for caching while forwarding them to the client struct CachingStream { @@ -70,7 +71,7 @@ where Poll::Ready(Some(Ok(chunk.into_bytes()))) } Poll::Ready(Some(Err(e))) => { - error!("AI stream error: {}", e); + error!("AI stream error"); Poll::Ready(Some(Err(std::io::Error::new( std::io::ErrorKind::Other, e.to_string(), @@ -89,7 +90,7 @@ where { let _: Result<(), _> = conn.set_ex(&cache_key, buffer.as_str(), 600).await; - info!("Cached AI response for key: {}", cache_key); + info!("Cached AI response"); } } }); @@ -151,27 +152,53 @@ pub async fn search( State(state): State, Json(mut request): Json, ) -> SearcherResult> { - info!("Received search request: {:?}", request); - hydrate_user_configuration(&state, &mut request).await?; + let request_query_len = request.query.len(); + info!(query_len = request_query_len, "Received search request"); + + // ------------------------------------------------------------------ + // Inner async block — single timing/error measurement across + // hydration, SearchEngine construction, engine search and response + // serialization. Result count is recorded only on full success. + // ------------------------------------------------------------------ + let start = Instant::now(); + + let inner_result: SearcherResult<(Value, u64, SearchEngine)> = async { + hydrate_user_configuration(&state, &mut request).await?; + + let search_engine = SearchEngine::new( + state.db_pool, + state.redis_client, + state.ai_client, + state.config, + state.operator_registry, + ) + .await?; - let search_engine = SearchEngine::new( - state.db_pool, - state.redis_client, - state.ai_client, - state.config, - state.operator_registry, - ) - .await?; + let response = search_engine.search(request.clone()).await?; + let result_count = response.results.len() as u64; + let value = serde_json::to_value(response)?; + + Ok((value, result_count, search_engine)) + } + .await; - let response = match search_engine.search(request.clone()).await { - Ok(response) => response, + let duration = start.elapsed().as_secs_f64(); + + let (value, _result_count, search_engine) = match inner_result { + Ok(result) => { + metrics::SEARCHER_SEARCH_DURATION.record(duration, &[]); + metrics::SEARCHER_SEARCH_RESULTS.record(result.1, &[]); + result + } Err(e) => { + metrics::SEARCHER_SEARCH_ERRORS.add(1, &[]); + metrics::SEARCHER_SEARCH_DURATION.record(duration, &[]); error!("Search engine error: {}", e); - return Err(SearcherError::Internal(e)); + return Err(e); } }; - // Store search history if user_id is provided + // Store search history if user_id is provided (nonfatal) if let Some(user_id) = &request.user_id { let is_generated = request.is_generated_query.unwrap_or(false); @@ -191,17 +218,14 @@ pub async fn search( } } - Ok(Json(serde_json::to_value(response)?)) + Ok(Json(value)) } pub async fn recent_searches( State(state): State, Query(query): Query, ) -> SearcherResult> { - info!( - "Received recent searches request for user: {}", - query.user_id - ); + info!("Received recent searches request"); let search_engine = SearchEngine::new( state.db_pool, @@ -221,7 +245,8 @@ pub async fn ai_answer( State(state): State, Json(mut request): Json, ) -> Result, axum::http::StatusCode> { - info!("Received AI answer request: {:?}", request); + let answer_query_len = request.query.len(); + info!("Received AI answer request"); hydrate_user_configuration(&state, &mut request) .await .map_err(|_| axum::http::StatusCode::BAD_REQUEST)?; @@ -242,7 +267,7 @@ pub async fn ai_answer( // Try to get cached AI response first if let Ok(mut conn) = state.redis_client.get_multiplexed_async_connection().await { if let Ok(cached_answer) = conn.get::<_, String>(&cache_key).await { - info!("Cache hit for AI answer query: '{}'", request.query); + info!("AI answer cache hit"); let response = axum::response::Response::builder() .status(StatusCode::OK) .header("Content-Type", "text/plain; charset=utf-8") @@ -254,27 +279,26 @@ pub async fn ai_answer( } // Cache miss - generate fresh response - info!("Cache miss for AI answer query: '{}'", request.query); + info!("AI answer cache miss"); // Get RAG context by running hybrid search let context = match search_engine.get_rag_context(&request).await { Ok(context) => context, Err(e) => { - error!("Failed to get RAG context: {}", e); + error!("Failed to get RAG context"); return Err(StatusCode::INTERNAL_SERVER_ERROR); } }; // Build RAG prompt with context and citation instructions let prompt = search_engine.build_rag_prompt(&request.query, &context); - info!("Built RAG prompt of length: {}", prompt.len()); - debug!("RAG prompt: {}", prompt); + info!("Built RAG prompt"); // Stream AI response let ai_stream = match state.ai_client.stream_prompt(&prompt).await { Ok(stream) => stream, Err(e) => { - error!("Failed to start AI stream: {}", e); + error!("Failed to start AI stream"); return Err(StatusCode::BAD_GATEWAY); } }; @@ -421,23 +445,12 @@ pub async fn suggested_questions( let user = match user_repo.find_by_id(request.user_id.clone()).await { Ok(Some(user)) => user, Ok(None) => { - error!("User not found for user_id {}", request.user_id); - return Err(SearcherError::NotFound(format!( - "User not found for user_id {}", - request.user_id - ))); + error!("User not found"); + return Err(SearcherError::NotFound("User not found".to_string())); } Err(e) => { - error!( - "Failed to fetch user for user_id {}: {:?}", - request.user_id, e - ); - return Err(anyhow!( - "Failed to fetch user for user_id {}: {:?}", - request.user_id, - e - ) - .into()); + error!("Failed to fetch user"); + return Err(anyhow!("Failed to fetch user: {:?}", e).into()); } }; diff --git a/services/searcher/src/lib.rs b/services/searcher/src/lib.rs index 1917d9c19..d3a76af10 100644 --- a/services/searcher/src/lib.rs +++ b/services/searcher/src/lib.rs @@ -143,14 +143,14 @@ pub async fn run_server() -> AnyhowResult<()> { let title_index = Arc::new(TitleIndex::new(db_pool.clone())); if let Err(e) = title_index.refresh().await { - error!("Failed initial typeahead index load: {}", e); + error!("Failed initial typeahead index load"); } title_index.start_background_refresh(3600); info!("Typeahead index initialized, refresh interval: 3600s"); let operator_registry = Arc::new(OperatorRegistry::new(redis_client.clone())); if let Err(e) = operator_registry.refresh().await { - error!("Failed initial operator registry load: {}", e); + error!("Failed initial operator registry load"); } operator_registry.start_background_refresh(60); info!("Operator registry initialized"); @@ -169,10 +169,19 @@ pub async fn run_server() -> AnyhowResult<()> { let app = create_app(app_state); let addr = SocketAddr::from(([0, 0, 0, 0], config.port)); - info!("Searcher service listening on {}", addr); + info!("Searcher service listening"); let listener = tokio::net::TcpListener::bind(addr).await?; - axum::serve(listener, app).await?; + if let Err(e) = axum::serve(listener, app) + .with_graceful_shutdown(telemetry::shutdown_signal()) + .await + { + error!("Server error"); + telemetry::shutdown_telemetry().await; + return Err(e.into()); + } + + telemetry::shutdown_telemetry().await; Ok(()) } diff --git a/services/searcher/src/search.rs b/services/searcher/src/search.rs index dc3ff73e6..ac1f80c81 100644 --- a/services/searcher/src/search.rs +++ b/services/searcher/src/search.rs @@ -6,19 +6,19 @@ use crate::query_parser; use crate::search_repository::SearchDocumentRepository; use anyhow::Result; use redis::{AsyncCommands, Client as RedisClient}; -use shared::SourceType; use shared::db::repositories::{ DocumentRepository, EmbeddingRepository, GroupRepository, PersonRepository, SourceRepository, }; use shared::models::{ChunkResult, Document, Facet, FacetValue}; use shared::utils::safe_str_slice; +use shared::SourceType; use shared::{ - AIClient, DatabasePool, ObjectStorage, Repository, SearcherConfig, StorageFactory, + metrics, AIClient, DatabasePool, ObjectStorage, Repository, SearcherConfig, StorageFactory, UserRepository, }; use std::cmp::Ordering; -use std::collections::HashMap; use std::collections::hash_map::DefaultHasher; +use std::collections::HashMap; use std::hash::{Hash, Hasher}; use std::sync::Arc; use std::time::{Duration, Instant}; @@ -115,9 +115,9 @@ impl SearchEngine { let start_time = Instant::now(); info!( - "Searching for query: '{}', mode: {:?}", - request.query, - request.search_mode() + query_len = request.query.len(), + mode = ?request.search_mode(), + "Searching" ); let mut request = request; @@ -128,23 +128,15 @@ impl SearchEngine { // In case the request contains only user_id, populate user_email for permission filtering let user_repo = UserRepository::new(self.db_pool.pool()); let mut request = match (&request.user_id, &request.user_email) { - (Some(user_id), None) => { - info!( - "Search request has user_id but no email, fetching email from DB for user ID: {}", - user_id - ); - let res = user_repo.find_by_id(user_id.clone()).await; - info!("Fetched user: {:?}", res); + (Some(_user_id), None) => { + info!("Fetching user email from DB for email-based filtering"); + let res = user_repo.find_by_id(_user_id.clone()).await; if let Ok(Some(user)) = res { - info!( - "Fetched user email: {} for user ID: {}", - user.email, user_id - ); let mut new_request = request.clone(); new_request.user_email = Some(user.email); new_request } else { - info!("Failed to fetch user email for user ID: {}", user_id); + info!("No user record found"); request } } @@ -164,7 +156,7 @@ impl SearchEngine { // Handle document_id filter for read_document tool if let Some(document_id) = &request.document_id { - info!("Document ID filter detected: {}", document_id); + info!("Document ID filter detected"); return self.read_document_by_id(document_id, &request).await; } @@ -176,7 +168,11 @@ impl SearchEngine { &request.user_configuration, ) .await; - info!("Parsed query: {:?}", parsed); + info!( + "Parsed query: {} filters, {} source_types", + parsed.attribute_filters.len(), + parsed.source_types.len() + ); let has_parsed_filters = !parsed.attribute_filters.is_empty() || !parsed.source_types.is_empty() || !parsed.boosted_source_types.is_empty() @@ -233,12 +229,15 @@ impl SearchEngine { if let Ok(mut conn) = self.redis_client.get_multiplexed_async_connection().await { if let Ok(cached_response) = conn.get::<_, String>(&cache_key).await { if let Ok(response) = serde_json::from_str::(&cached_response) { - info!("Cache hit for request: {:?}", request); + metrics::SEARCHER_CACHE_HIT.add(1, &[]); + info!("Cache hit for search request"); return Ok(response); } } } + metrics::SEARCHER_CACHE_MISS.add(1, &[]); + let repo = DocumentRepository::new(self.db_pool.pool()); let search_repo = SearchDocumentRepository::new(self.db_pool.pool()); @@ -488,7 +487,7 @@ impl SearchEngine { let content_types = request.content_types.as_deref(); let attribute_filters = request.attribute_filters.as_ref(); - debug!("Running fulltext search for {}", &request.query); + debug!("Running fulltext search"); let search_hits = repo .search( &request.query, @@ -552,7 +551,7 @@ impl SearchEngine { offset: i64, ) -> Result> { let start_time = Instant::now(); - info!("Performing semantic search for query: '{}'", request.query); + info!("Performing semantic search"); let query_embedding = self.generate_query_embedding(&request.query).await?; @@ -686,7 +685,7 @@ impl SearchEngine { request: &SearchRequest, ) -> Result { let start_time = Instant::now(); - info!("Reading document by ID: {}", document_id); + info!("Reading document by ID"); let doc_repo = DocumentRepository::new(self.db_pool.pool()); let doc = doc_repo @@ -792,11 +791,8 @@ impl SearchEngine { } } } - Err(e) => { - info!( - "Failed to read document content: {}, falling back to chunk retrieval", - e - ); + Err(_e) => { + info!("Failed to read document content, falling back to chunk retrieval"); self.read_document_chunks(document_id, &doc, request) .await? } @@ -855,10 +851,7 @@ impl SearchEngine { .await?; results } else { - info!( - "No query provided, returning first 500 lines from document ID {}", - document_id - ); + info!("No query provided, returning first 500 lines"); if let Some(content_id) = &doc.content_id { let content = self.content_storage.get_text(content_id).await?; @@ -887,10 +880,7 @@ impl SearchEngine { also_in: Vec::new(), }] } else { - error!( - "Content ID not found for document ID [{}], content ID [{:?}]", - document_id, doc.content_id - ); + error!("Content ID not found for document"); vec![] } }; @@ -899,7 +889,7 @@ impl SearchEngine { } async fn generate_query_embedding(&self, query: &str) -> Result> { - debug!("Generating query embeddings for query '{}'", query); + debug!("Generating query embeddings"); let embeddings = self .ai_client .generate_embeddings_with_options( @@ -925,10 +915,7 @@ impl SearchEngine { user_groups: &[String], ) -> Result> { let start_time = Instant::now(); - info!( - "Generating enhanced semantic search results for RAG query: '{}'", - request.query - ); + info!("Generating enhanced semantic search results for RAG"); let query_embedding = self.generate_query_embedding(&request.query).await?; let search_repo = SearchDocumentRepository::new(self.db_pool.pool()); @@ -1047,7 +1034,7 @@ impl SearchEngine { user_groups: &[String], tantivy_query: Option<&str>, ) -> Result<(Vec, i64)> { - info!("Performing hybrid search for query: '{}'", request.query); + info!("Performing hybrid search"); let start_time = Instant::now(); let doc_repo = DocumentRepository::new(self.db_pool.pool()); @@ -1311,7 +1298,7 @@ impl SearchEngine { /// Generate RAG context from search request using chunk-based approach with expanded context pub async fn get_rag_context(&self, request: &SearchRequest) -> Result> { - info!("Generating RAG context for query: '{}'", request.query); + info!("Generating RAG context"); let user_groups = if let Some(email) = request.user_email() { let group_repo = GroupRepository::new(self.db_pool.pool()); diff --git a/services/searcher/src/search_repository.rs b/services/searcher/src/search_repository.rs index 747802cb9..6c408873a 100644 --- a/services/searcher/src/search_repository.rs +++ b/services/searcher/src/search_repository.rs @@ -1,12 +1,12 @@ use pgvector::Vector; use serde_json::Value as JsonValue; use shared::{ - SourceType, db::error::DatabaseError, db::repositories::document, models::{AttributeFilter, ChunkResult, DateFilter, Document, Facet, FacetValue}, + SourceType, }; -use sqlx::{FromRow, PgPool, Row, postgres::PgRow}; +use sqlx::{postgres::PgRow, FromRow, PgPool, Row}; use std::collections::{HashMap, HashSet}; use tracing::debug; @@ -283,7 +283,7 @@ impl SearchDocumentRepository { offset_idx = offset_idx, min_score_ratio_idx = min_score_ratio_idx, ); - debug!("Full search query: {}", full_query); + debug!("Full search query: {} params", param_idx); let mut query_builder = sqlx::query_as::<_, SearchHitWithTotalRow>(&full_query).bind(tantivy_query); diff --git a/services/searcher/src/suggested_questions.rs b/services/searcher/src/suggested_questions.rs index 4d667d46a..79f9b60d8 100644 --- a/services/searcher/src/suggested_questions.rs +++ b/services/searcher/src/suggested_questions.rs @@ -1,6 +1,6 @@ use crate::models::{SuggestedQuestion, SuggestedQuestionsResponse}; use crate::{Result as SearcherResult, SearcherError}; -use anyhow::{Context, Result, anyhow}; +use anyhow::{anyhow, Context, Result}; use dashmap::DashSet; use futures_util::StreamExt; use redis::AsyncCommands; @@ -104,17 +104,11 @@ impl SuggestedQuestionsGenerator { "Cache miss for suggested questions, checking for in-flight generations for this user before proceeding." ); if self.in_flight.contains(user_email) { - info!( - "Suggested questions generation already in progress for user {}", - user_email - ); + info!("Suggested questions generation already in progress"); return Ok(SuggestedQuestionsResponse { questions: vec![] }); } - info!( - "No in-flight generation found, starting new generation task for user {}", - user_email - ); + info!("No in-flight generation found, starting new generation task"); tokio::spawn({ let user_email = user_email.to_string(); let in_flight = Arc::clone(&self.in_flight); @@ -135,26 +129,20 @@ impl SuggestedQuestionsGenerator { { Ok(count) => { info!( - "Successfully generated and cached {} suggested questions for user {}", - count, user_email + "Successfully generated and cached {} suggested questions", + count ); // Remove the user from the in-flight map to allow future requests to go through in_flight.remove(&user_email); } - Err(e) => { - error!( - "Failed to generate suggested questions for user {}: {:?}", - user_email, e - ); + Err(_) => { + error!("Failed to generate suggested questions"); // Remove the user from the in-flight map to allow future requests to go through in_flight.remove(&user_email); } } } else { - info!( - "Another generation task started for user {} while we were waiting", - user_email - ); + info!("Another generation task started while we were waiting"); } } }); @@ -221,10 +209,7 @@ impl SuggestedQuestionsGenerator { let content_ids: Vec = docs.iter().filter_map(|d| d.content_id.clone()).collect(); - debug!( - "Fetching content IDs {:?} for generating suggested questions", - content_ids - ); + debug!("Fetching content for generating suggested questions"); let content_map = content_storage.batch_get_text(content_ids).await?; // Build contents vector in the same order as documents @@ -250,9 +235,7 @@ impl SuggestedQuestionsGenerator { } debug!( - "Processing document {} [id={}] (content length: {} chars)", - doc.title, - doc.id, + "Processing document (content length: {} chars)", content.len() ); @@ -269,23 +252,14 @@ impl SuggestedQuestionsGenerator { // Two different documents can produce the same generic // suggestion; keep the displayed suggestions distinct. if !seen_questions.insert(question.to_lowercase()) { - debug!( - "Skipping duplicate suggestion text from document {}", - doc.id - ); + debug!("Skipping duplicate suggestion text"); continue; } questions.push(SuggestedQuestion { question: question.clone(), document_id: doc.id.clone(), }); - info!( - "Generated suggestion {}/{}: \"{}\" (from document: {})", - questions.len(), - num_questions, - question, - doc.id - ); + info!("Generated suggestion {}/{}", questions.len(), num_questions); // Cache the questions debug!("Serializing {} question(s) to JSON", questions.len()); @@ -302,10 +276,7 @@ impl SuggestedQuestionsGenerator { .context("Failed to connect to Redis")?; let cache_key = format!("{}:{}", REDIS_CACHE_KEY, user_email); - debug!( - "Caching questions in Redis with key: {}, TTL: {}s", - cache_key, CACHE_TTL_SECONDS - ); + debug!("Caching questions in Redis (TTL: {}s)", CACHE_TTL_SECONDS); redis_conn .set_ex::<_, _, ()>(cache_key, &json_str, CACHE_TTL_SECONDS) .await @@ -317,8 +288,8 @@ impl SuggestedQuestionsGenerator { CACHE_TTL_SECONDS / 3600 ); } - Err(e) => { - warn!("Failed to generate suggestion for document {}: {}", doc.id, e); + Err(_) => { + warn!("Failed to generate suggestion"); } } @@ -333,17 +304,14 @@ impl SuggestedQuestionsGenerator { if num_docs_fetched < needed { debug!( - "User {} has only {} documents, skipping further attempts", - user_email, num_docs_fetched + "User has only {} documents, skipping further attempts", + num_docs_fetched ); break; } } - Err(e) => { - error!( - "Failed to fetch random documents on attempt {}: {}", - attempts, e - ); + Err(_) => { + error!("Failed to fetch random documents on attempt {}", attempts); } } } @@ -376,11 +344,7 @@ impl SuggestedQuestionsGenerator { expect_question_mark: bool, ) -> Result { let excerpt = if content.len() > 2000 { - debug!( - "Truncating content from {} to 2000 chars for document {}", - content.len(), - document_id - ); + debug!("Truncating content from {} to 2000 chars", content.len()); safe_str_slice(content, 0, 2000) } else { content @@ -410,25 +374,18 @@ impl SuggestedQuestionsGenerator { let normalized = output.trim_end_matches(|c: char| c == '.' || c == '!' || c.is_whitespace()); if normalized.eq_ignore_ascii_case("skip") { - info!( - "Document {} deemed unsuitable for suggestion (model returned SKIP)", - document_id - ); - return Err(anyhow!("Document unsuitable for suggestion generation (SKIP)")); + info!("Document deemed unsuitable for suggestion (model returned SKIP)"); + return Err(anyhow!( + "Document unsuitable for suggestion generation (SKIP)" + )); } if expect_question_mark && !output.contains('?') { - warn!( - "Discarding non-question suggestion for document {}: \"{}\"", - document_id, output - ); + warn!("Discarding non-question suggestion"); return Err(anyhow!("Generated suggestion was not a question")); } - info!( - "Successfully generated suggestion for document {}: \"{}\"", - document_id, output - ); + info!("Successfully generated suggestion"); Ok(output) } diff --git a/shared/Cargo.toml b/shared/Cargo.toml index ba41d822e..366b70e19 100644 --- a/shared/Cargo.toml +++ b/shared/Cargo.toml @@ -58,7 +58,15 @@ opentelemetry-otlp = { workspace = true } opentelemetry-semantic-conventions = { workspace = true } tracing-opentelemetry = { workspace = true } opentelemetry-http = { workspace = true } +opentelemetry-appender-tracing = { workspace = true } tower = { workspace = true } tower-http = { workspace = true } http = "1.1" dashmap = { workspace = true } + +[dev-dependencies] +opentelemetry_sdk = { workspace = true, features = ["testing"] } +tokio = { workspace = true } +tracing = { workspace = true } +opentelemetry-proto = { version = "0.32", features = ["trace", "metrics", "logs", "gen-tonic-messages"] } +prost = "0.14" diff --git a/shared/src/clients/ai.rs b/shared/src/clients/ai.rs index 458c1eef7..dec2245df 100644 --- a/shared/src/clients/ai.rs +++ b/shared/src/clients/ai.rs @@ -1,11 +1,11 @@ -use anyhow::{Result, anyhow}; +use anyhow::{anyhow, Result}; use futures_util::{Stream, StreamExt}; use reqwest::Client; use serde::{Deserialize, Serialize}; use std::pin::Pin; use tracing::error; -use crate::telemetry::http_client::RequestBuilderExt; +use crate::telemetry::http_client; #[derive(Serialize)] pub struct EmbeddingRequest { @@ -81,13 +81,9 @@ impl AIClient { priority, }; - let response = self - .client - .post(format!("{}/embeddings", self.base_url)) - .json(&request) - .with_trace_context() - .send() - .await; + let url = format!("{}/embeddings", self.base_url); + let response = + http_client::send_traced("POST", &url, self.client.post(&url).json(&request)).await; match response { Ok(res) => { @@ -155,13 +151,9 @@ impl AIClient { stream: Some(true), }; - let response = self - .client - .post(format!("{}/prompt", self.base_url)) - .json(&request) - .with_trace_context() - .send() - .await?; + let url = format!("{}/prompt", self.base_url); + let response = + http_client::send_traced("POST", &url, self.client.post(&url).json(&request)).await?; if !response.status().is_success() { return Err(anyhow::anyhow!( diff --git a/shared/src/clients/docling.rs b/shared/src/clients/docling.rs index d181ac5d6..eac834dc7 100644 --- a/shared/src/clients/docling.rs +++ b/shared/src/clients/docling.rs @@ -3,12 +3,14 @@ //! The Docling service converts documents to Markdown using AI-based extraction. //! It provides superior PDF extraction, OCR support, and structure-aware output. -use anyhow::{Context, anyhow}; -use reqwest::{Client, multipart}; +use anyhow::{anyhow, Context}; +use reqwest::{multipart, Client}; use serde::Deserialize; use std::time::Duration; use tracing::{debug, error}; +use crate::telemetry::http_client; + /// Errors from the Docling document conversion service. #[derive(Debug, thiserror::Error)] pub enum DoclingError { @@ -94,10 +96,8 @@ impl DoclingClient { /// Check if the Docling service is healthy and ready. pub async fn health_check(&self) -> anyhow::Result { - let response = self - .client - .get(format!("{}/health", self.base_url)) - .send() + let url = format!("{}/health", self.base_url); + let response = http_client::send_traced("GET", &url, self.client.get(&url)) .await .context("Failed to connect to Docling service")?; @@ -190,14 +190,17 @@ impl DoclingClient { let form = multipart::Form::new().part("file", part); - let response = self - .client - .post(format!("{}/convert", self.base_url)) - .query(&[("preset", preset)]) - .multipart(form) - .send() - .await - .context("Failed to submit document to Docling")?; + let url = format!("{}/convert", self.base_url); + let response = http_client::send_traced( + "POST", + &url, + self.client + .post(&url) + .query(&[("preset", preset)]) + .multipart(form), + ) + .await + .context("Failed to submit document to Docling")?; if response.status().as_u16() == 429 { let retry_after = response @@ -235,10 +238,8 @@ impl DoclingClient { /// Get the status of a conversion job. async fn get_job(&self, job_id: &str) -> Result { - let response = self - .client - .get(format!("{}/jobs/{}", self.base_url, job_id)) - .send() + let url = format!("{}/jobs/{}", self.base_url, job_id); + let response = http_client::send_traced("GET", &url, self.client.get(&url)) .await .context("Failed to poll Docling job")?; diff --git a/shared/src/embedding_queue.rs b/shared/src/embedding_queue.rs index f22c169c1..7f573b8ab 100644 --- a/shared/src/embedding_queue.rs +++ b/shared/src/embedding_queue.rs @@ -1,6 +1,7 @@ use anyhow::Result; use serde::{Deserialize, Serialize}; use sqlx::{PgPool, Row}; +use tracing::Instrument; use ulid::Ulid; use crate::{db::repositories::EmbeddingProviderRepository, utils::generate_ulid}; @@ -49,6 +50,12 @@ pub struct EmbeddingQueueItem { pub created_at: sqlx::types::time::OffsetDateTime, pub updated_at: sqlx::types::time::OffsetDateTime, pub processed_at: Option, + /// W3C traceparent header stored by the producer for trace propagation. + #[serde(skip_serializing_if = "Option::is_none")] + pub traceparent: Option, + /// W3C tracestate header stored alongside traceparent. + #[serde(skip_serializing_if = "Option::is_none")] + pub tracestate: Option, } impl sqlx::FromRow<'_, sqlx::postgres::PgRow> for EmbeddingQueueItem { @@ -73,6 +80,8 @@ impl sqlx::FromRow<'_, sqlx::postgres::PgRow> for EmbeddingQueueItem { created_at: row.try_get("created_at")?, updated_at: row.try_get("updated_at")?, processed_at: row.try_get("processed_at")?, + traceparent: row.try_get("traceparent").ok().flatten(), + tracestate: row.try_get("tracestate").ok().flatten(), }) } } @@ -93,49 +102,26 @@ impl EmbeddingQueue { } pub async fn enqueue(&self, document_id: String) -> Result> { - if !self.provider_repo.has_active_provider().await? { - return Ok(None); - } + let queue_name = "embedding_queue"; - let id = Ulid::new().to_string(); + // Build PRODUCER span covering provider check + INSERT. + // inject_active_context is called from inside the instrumented block + // so the producer span is active during injection and the DB write. + let producer_span = crate::telemetry::queue::build_producer_span(queue_name); - let result = sqlx::query( - r#" - INSERT INTO embedding_queue (id, document_id) - SELECT $1, $2 - WHERE NOT EXISTS ( - SELECT 1 FROM embedding_queue - WHERE document_id = $2 AND status IN ('pending', 'processing') - ) - "#, - ) - .bind(&id) - .bind(&document_id) - .execute(&self.pool) - .await?; - - if result.rows_affected() > 0 { - Ok(Some(id)) - } else { - Ok(None) - } - } + async move { + let carrier = crate::telemetry::queue::inject_active_context(); - pub async fn enqueue_batch(&self, document_ids: Vec) -> Result> { - if !self.provider_repo.has_active_provider().await? { - return Ok(vec![]); - } - - let mut tx = self.pool.begin().await?; - let mut ids = Vec::new(); + if !self.provider_repo.has_active_provider().await? { + return Ok(None); + } - for document_id in document_ids { let id = Ulid::new().to_string(); let result = sqlx::query( r#" - INSERT INTO embedding_queue (id, document_id) - SELECT $1, $2 + INSERT INTO embedding_queue (id, document_id, traceparent, tracestate) + SELECT $1, $2, $3, $4 WHERE NOT EXISTS ( SELECT 1 FROM embedding_queue WHERE document_id = $2 AND status IN ('pending', 'processing') @@ -144,69 +130,166 @@ impl EmbeddingQueue { ) .bind(&id) .bind(&document_id) - .execute(&mut *tx) - .await?; + .bind(&carrier.traceparent) + .bind(&carrier.tracestate) + .execute(&self.pool) + .await + .map_err(|e| { + crate::telemetry::queue::set_active_span_error(); + e + })?; if result.rows_affected() > 0 { - ids.push(id); + Ok(Some(id)) + } else { + Ok(None) } } + .instrument(producer_span) + .await + } + + pub async fn enqueue_batch(&self, document_ids: Vec) -> Result> { + let queue_name = "embedding_queue"; + + // Build PRODUCER span covering provider check + batch INSERT. + // inject_active_context is called from inside the instrumented block + // so the producer span is active during injection and the DB write. + let producer_span = crate::telemetry::queue::build_producer_span(queue_name); + + async move { + let carrier = crate::telemetry::queue::inject_active_context(); + + if !self.provider_repo.has_active_provider().await? { + return Ok(vec![]); + } + + let mut tx = self.pool.begin().await.map_err(|e| { + crate::telemetry::queue::set_active_span_error(); + e + })?; + let mut ids = Vec::new(); + + for document_id in document_ids { + let id = Ulid::new().to_string(); + + let result = sqlx::query( + r#" + INSERT INTO embedding_queue (id, document_id, traceparent, tracestate) + SELECT $1, $2, $3, $4 + WHERE NOT EXISTS ( + SELECT 1 FROM embedding_queue + WHERE document_id = $2 AND status IN ('pending', 'processing') + ) + "#, + ) + .bind(&id) + .bind(&document_id) + .bind(&carrier.traceparent) + .bind(&carrier.tracestate) + .execute(&mut *tx) + .await + .map_err(|e| { + crate::telemetry::queue::set_active_span_error(); + e + })?; + + if result.rows_affected() > 0 { + ids.push(id); + } + } - tx.commit().await?; - Ok(ids) + tx.commit().await.map_err(|e| { + crate::telemetry::queue::set_active_span_error(); + e + })?; + Ok(ids) + } + .instrument(producer_span) + .await } pub async fn enqueue_batch_missing_current_embeddings( &self, document_ids: Vec, ) -> Result> { - if document_ids.is_empty() || !self.provider_repo.has_active_provider().await? { + if document_ids.is_empty() { return Ok(vec![]); } - let ids: Vec = document_ids.iter().map(|_| generate_ulid()).collect(); - let rows = sqlx::query( - r#" - WITH active_provider AS ( - SELECT config->>'model' AS model_name - FROM embedding_providers - WHERE is_current = TRUE AND is_deleted = FALSE - LIMIT 1 - ), - input_rows AS ( - SELECT id, document_id, ordinality - FROM UNNEST($1::text[], $2::text[]) WITH ORDINALITY AS t(id, document_id, ordinality) - ), - deduped_input AS ( - SELECT DISTINCT ON (document_id) id, document_id - FROM input_rows - ORDER BY document_id, ordinality - ) - INSERT INTO embedding_queue (id, document_id) - SELECT input.id, input.document_id - FROM deduped_input input - CROSS JOIN active_provider provider - WHERE NOT EXISTS ( - SELECT 1 - FROM embedding_queue q - WHERE q.document_id = input.document_id - AND q.status IN ('pending', 'processing') - ) - AND NOT EXISTS ( - SELECT 1 - FROM embeddings e - WHERE e.document_id = input.document_id - AND e.model_name = provider.model_name - ) - RETURNING id - "#, - ) - .bind(&ids) - .bind(&document_ids) - .fetch_all(&self.pool) - .await?; + let queue_name = "embedding_queue"; + + // Build PRODUCER span covering provider check + INSERT. + // inject_active_context is called from inside the instrumented block + // so the producer span is active during injection and the DB write. + let producer_span = crate::telemetry::queue::build_producer_span(queue_name); - Ok(rows.into_iter().map(|row| row.get("id")).collect()) + async move { + let carrier = crate::telemetry::queue::inject_active_context(); + + if !self.provider_repo.has_active_provider().await? { + return Ok(vec![]); + } + + let ids: Vec = document_ids.iter().map(|_| generate_ulid()).collect(); + let traceparents: Vec> = std::iter::repeat(carrier.traceparent.clone()) + .take(ids.len()) + .collect(); + let tracestates: Vec> = std::iter::repeat(carrier.tracestate.clone()) + .take(ids.len()) + .collect(); + + let rows = sqlx::query( + r#" + WITH active_provider AS ( + SELECT config->>'model' AS model_name + FROM embedding_providers + WHERE is_current = TRUE AND is_deleted = FALSE + LIMIT 1 + ), + input_rows AS ( + SELECT id, document_id, traceparent, tracestate, ordinality + FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[]) WITH ORDINALITY AS t(id, document_id, traceparent, tracestate, ordinality) + ), + deduped_input AS ( + SELECT DISTINCT ON (document_id) id, document_id, traceparent, tracestate + FROM input_rows + ORDER BY document_id, ordinality + ) + INSERT INTO embedding_queue (id, document_id, traceparent, tracestate) + SELECT input.id, input.document_id, input.traceparent, input.tracestate + FROM deduped_input input + CROSS JOIN active_provider provider + WHERE NOT EXISTS ( + SELECT 1 + FROM embedding_queue q + WHERE q.document_id = input.document_id + AND q.status IN ('pending', 'processing') + ) + AND NOT EXISTS ( + SELECT 1 + FROM embeddings e + WHERE e.document_id = input.document_id + AND e.model_name = provider.model_name + ) + RETURNING id + "#, + ) + .bind(&ids) + .bind(&document_ids) + .bind(&traceparents) + .bind(&tracestates) + .fetch_all(&self.pool) + .await + .map_err(|e| { + crate::telemetry::queue::set_active_span_error(); + e + })?; + + Ok(rows.into_iter().map(|row| row.get("id")).collect()) + } + .instrument(producer_span) + .await } pub async fn dequeue_batch(&self, batch_size: i32) -> Result> { diff --git a/shared/src/lib.rs b/shared/src/lib.rs index 1e4b1f34a..ffd24f076 100644 --- a/shared/src/lib.rs +++ b/shared/src/lib.rs @@ -7,6 +7,7 @@ pub mod content_storage; pub mod db; pub mod embedding_queue; pub mod encryption; +pub mod metrics; pub mod models; pub mod queue; pub mod rate_limiter; @@ -35,10 +36,10 @@ pub use encryption::{EncryptedData, EncryptionService}; pub use models::*; pub use queue::{EventQueue, QueueStats, QueueSummary}; pub use rate_limiter::{RateLimiter, RetryableError}; -pub use service_auth::{ServiceAuth, create_service_auth}; +pub use service_auth::{create_service_auth, ServiceAuth}; pub use storage::{ - ContentMetadata as StorageContentMetadata, ObjectStorage, StorageError, factory::{StorageBackend, StorageFactory}, + ContentMetadata as StorageContentMetadata, ObjectStorage, StorageError, }; pub use traits::Repository; diff --git a/shared/src/metrics.rs b/shared/src/metrics.rs new file mode 100644 index 000000000..2afe2a2a9 --- /dev/null +++ b/shared/src/metrics.rs @@ -0,0 +1,258 @@ +//! Centralised OpenTelemetry metric instruments for Omni application services. +//! +//! Every instrument uses the `omni.*` namespace and is created lazily from +//! the global meter provider. Callers just invoke the free functions defined +//! here — they are safe to call from any thread and return cached instruments. +//! +//! # Conventions +//! +//! - **Counters** use `u64` and [`Counter::add`]. +//! - **Histograms** use `f64` seconds and [`Histogram::record`]. +//! - **Gauges** are created ad-hoc with a meter; callers that need frequent +//! gauge snapshots should define a cached static. + +use std::sync::LazyLock; + +use opentelemetry::{ + global, + metrics::{Counter, Gauge, Histogram}, + KeyValue, +}; + +// --------------------------------------------------------------------------- +// HTTP server RED metrics (cross-cutting, used by all HTTP services) +// --------------------------------------------------------------------------- + +/// HTTP server request counter with bounded attributes (method, route, status). +pub static HTTP_SERVER_REQUEST_COUNT: LazyLock> = LazyLock::new(|| { + global::meter("omni") + .u64_counter("omni.http.server.request_count") + .with_description("Total number of HTTP server requests by method, route, status") + .build() +}); + +/// HTTP server request duration histogram in seconds. +pub static HTTP_SERVER_REQUEST_DURATION: LazyLock> = LazyLock::new(|| { + global::meter("omni") + .f64_histogram("omni.http.server.request_duration_seconds") + .with_description("HTTP server request duration in seconds") + .with_unit("s") + .build() +}); + +/// Record an HTTP server request for RED metrics. +/// +/// Attributes are limited to bounded values: method (e.g. GET, POST), +/// route (matched path template), and numeric status code. +pub fn record_http_request(method: &str, route: Option<&str>, status: u16, duration_secs: f64) { + let mut attrs = vec![ + KeyValue::new("http.request.method", method.to_string()), + KeyValue::new("http.response.status_code", status as i64), + ]; + if let Some(r) = route { + attrs.push(KeyValue::new("http.route", r.to_string())); + } + HTTP_SERVER_REQUEST_COUNT.add(1, &attrs); + HTTP_SERVER_REQUEST_DURATION.record(duration_secs, &attrs); +} + +// --------------------------------------------------------------------------- +// Indexer metrics +// --------------------------------------------------------------------------- + +/// Total connector events processed (by outcome: processed | failed | dead_letter). +pub static INDEXER_EVENTS_PROCESSED: LazyLock> = LazyLock::new(|| { + global::meter("omni-indexer") + .u64_counter("omni.indexer.events.processed") + .with_description("Total connector events successfully processed") + .build() +}); + +pub static INDEXER_EVENTS_DEAD_LETTER: LazyLock> = LazyLock::new(|| { + global::meter("omni-indexer") + .u64_counter("omni.indexer.events.dead_letter") + .with_description("Total connector events sent to dead-letter queue") + .build() +}); + +pub static INDEXER_EVENTS_RETRIED: LazyLock> = LazyLock::new(|| { + global::meter("omni-indexer") + .u64_counter("omni.indexer.events.retried") + .with_description("Total connector events retried") + .build() +}); + +pub static INDEXER_BATCH_DURATION: LazyLock> = LazyLock::new(|| { + global::meter("omni-indexer") + .f64_histogram("omni.indexer.batch.duration_seconds") + .with_description("Duration of event batch processing in seconds") + .with_unit("s") + .build() +}); + +pub static INDEXER_BATCH_SIZE: LazyLock> = LazyLock::new(|| { + global::meter("omni-indexer") + .u64_histogram("omni.indexer.batch.size") + .with_description("Number of events in a processed batch") + .with_unit("{events}") + .build() +}); + +// --------------------------------------------------------------------------- +// Queue depth gauge (cached static, used by indexer heartbeat) +// --------------------------------------------------------------------------- + +/// Single cached gauge for queue depth with bounded `queue` and `status` attributes. +pub static INDEXER_QUEUE_DEPTH: LazyLock> = LazyLock::new(|| { + global::meter("omni-indexer") + .i64_gauge("omni.indexer.queue.depth") + .with_description("Current depth of event or embedding queue by status") + .build() +}); + +/// Record connector-event queue depth snapshot via the cached gauge. +pub fn record_queue_status(pending: i64, processing: i64, failed: i64, dead_letter: i64) { + let base = [KeyValue::new("queue", "connector_events")]; + INDEXER_QUEUE_DEPTH.record( + pending, + &[&base[..], &[KeyValue::new("status", "pending")]].concat(), + ); + INDEXER_QUEUE_DEPTH.record( + processing, + &[&base[..], &[KeyValue::new("status", "processing")]].concat(), + ); + INDEXER_QUEUE_DEPTH.record( + failed, + &[&base[..], &[KeyValue::new("status", "failed")]].concat(), + ); + INDEXER_QUEUE_DEPTH.record( + dead_letter, + &[&base[..], &[KeyValue::new("status", "dead_letter")]].concat(), + ); +} + +/// Record embedding queue depth snapshot via the cached gauge. +pub fn record_embedding_queue_status(pending: i64, processing: i64, failed: i64) { + let base = [KeyValue::new("queue", "embedding")]; + INDEXER_QUEUE_DEPTH.record( + pending, + &[&base[..], &[KeyValue::new("status", "pending")]].concat(), + ); + INDEXER_QUEUE_DEPTH.record( + processing, + &[&base[..], &[KeyValue::new("status", "processing")]].concat(), + ); + INDEXER_QUEUE_DEPTH.record( + failed, + &[&base[..], &[KeyValue::new("status", "failed")]].concat(), + ); +} +// --------------------------------------------------------------------------- +// Searcher metrics +// --------------------------------------------------------------------------- + +pub static SEARCHER_SEARCH_DURATION: LazyLock> = LazyLock::new(|| { + global::meter("omni-searcher") + .f64_histogram("omni.searcher.search.duration_seconds") + .with_description("Search request duration in seconds") + .with_unit("s") + .build() +}); + +pub static SEARCHER_SEARCH_RESULTS: LazyLock> = LazyLock::new(|| { + global::meter("omni-searcher") + .u64_histogram("omni.searcher.search.results") + .with_description("Number of results returned by a search request") + .with_unit("{results}") + .build() +}); + +pub static SEARCHER_SEARCH_ERRORS: LazyLock> = LazyLock::new(|| { + global::meter("omni-searcher") + .u64_counter("omni.searcher.search.errors") + .with_description("Total search requests that resulted in an error") + .build() +}); + +pub static SEARCHER_CACHE_HIT: LazyLock> = LazyLock::new(|| { + global::meter("omni-searcher") + .u64_counter("omni.searcher.cache.hit") + .with_description("Total search requests served from cache") + .build() +}); + +pub static SEARCHER_CACHE_MISS: LazyLock> = LazyLock::new(|| { + global::meter("omni-searcher") + .u64_counter("omni.searcher.cache.miss") + .with_description("Total search requests that missed cache") + .build() +}); + +// --------------------------------------------------------------------------- +// Connector manager sync metrics +// --------------------------------------------------------------------------- + +pub static CONNECTOR_SYNC_STARTED: LazyLock> = LazyLock::new(|| { + global::meter("omni-connector-manager") + .u64_counter("omni.connector.sync.started") + .with_description("Total sync runs started") + .build() +}); + +pub static CONNECTOR_SYNC_COMPLETED: LazyLock> = LazyLock::new(|| { + global::meter("omni-connector-manager") + .u64_counter("omni.connector.sync.completed") + .with_description("Total sync runs completed successfully") + .build() +}); + +pub static CONNECTOR_SYNC_FAILED: LazyLock> = LazyLock::new(|| { + global::meter("omni-connector-manager") + .u64_counter("omni.connector.sync.failed") + .with_description("Total sync runs that failed") + .build() +}); + +pub static CONNECTOR_SYNC_CANCELLED: LazyLock> = LazyLock::new(|| { + global::meter("omni-connector-manager") + .u64_counter("omni.connector.sync.cancelled") + .with_description("Total sync runs cancelled") + .build() +}); + +pub static CONNECTOR_SYNC_DURATION: LazyLock> = LazyLock::new(|| { + global::meter("omni-connector-manager") + .f64_histogram("omni.connector.sync.duration_seconds") + .with_description("Duration of sync runs in seconds") + .with_unit("s") + .build() +}); + +/// Record a terminal sync metric with bounded sync_type and outcome attributes. +/// Never records IDs (sync_run_id, source_id) as attribute values. +/// Duration is included when created_at is available. +pub fn record_sync_terminal( + sync_type: &str, + outcome: &str, + created_at: Option, +) { + let attrs = [ + KeyValue::new("sync_type", sync_type.to_string()), + KeyValue::new("outcome", outcome.to_string()), + ]; + + match outcome { + "completed" => CONNECTOR_SYNC_COMPLETED.add(1, &attrs), + "failed" => CONNECTOR_SYNC_FAILED.add(1, &attrs), + "cancelled" => CONNECTOR_SYNC_CANCELLED.add(1, &attrs), + _ => {} + } + + if let Some(start_ts) = created_at { + let now = time::OffsetDateTime::now_utc(); + let duration_secs = (now - start_ts).as_seconds_f64(); + if duration_secs > 0.0 { + CONNECTOR_SYNC_DURATION.record(duration_secs, &attrs); + } + } +} diff --git a/shared/src/models.rs b/shared/src/models.rs index 13f754086..7a7f1dc48 100644 --- a/shared/src/models.rs +++ b/shared/src/models.rs @@ -979,6 +979,13 @@ pub struct ConnectorEventQueueItem { #[serde(with = "time::serde::iso8601::option")] pub processed_at: Option, pub error_message: Option, + /// W3C traceparent header stored by the producer for trace propagation. + /// Present when the enqueuing service created a PRODUCER span. + #[serde(skip_serializing_if = "Option::is_none")] + pub traceparent: Option, + /// W3C tracestate header stored alongside traceparent. + #[serde(skip_serializing_if = "Option::is_none")] + pub tracestate: Option, } /// Type/mode of a sync run. Serializes as a lowercase string on the wire diff --git a/shared/src/queue.rs b/shared/src/queue.rs index fc992ee29..f6cc866e1 100644 --- a/shared/src/queue.rs +++ b/shared/src/queue.rs @@ -1,6 +1,7 @@ use crate::utils::generate_ulid; use anyhow::Result; use sqlx::{PgPool, Row}; +use tracing::Instrument; use crate::models::{ConnectorEvent, ConnectorEventQueueItem, EventStatus, SyncType}; @@ -26,24 +27,42 @@ impl EventQueue { } pub async fn enqueue(&self, source_id: &str, event: &ConnectorEvent) -> Result { + let queue_name = "connector_events_queue"; let id = generate_ulid(); let event_type = event_type_str(event); - sqlx::query( - r#" - INSERT INTO connector_events_queue (id, sync_run_id, source_id, event_type, payload) - VALUES ($1, $2, $3, $4, $5) - "#, - ) - .bind(&id) - .bind(event.sync_run_id()) - .bind(source_id) - .bind(event_type) - .bind(serde_json::to_value(event)?) - .execute(&self.pool) - .await?; + // Build PRODUCER span and instrument the full INSERT with the span. + // inject_active_context is called from inside the instrumented block + // so the producer span is active during injection and the DB write. + let producer_span = crate::telemetry::queue::build_producer_span(queue_name); + + async move { + let carrier = crate::telemetry::queue::inject_active_context(); - Ok(id) + if let Err(e) = sqlx::query( + r#" + INSERT INTO connector_events_queue (id, sync_run_id, source_id, event_type, payload, traceparent, tracestate) + VALUES ($1, $2, $3, $4, $5, $6, $7) + "#, + ) + .bind(&id) + .bind(event.sync_run_id()) + .bind(source_id) + .bind(event_type) + .bind(serde_json::to_value(event)?) + .bind(&carrier.traceparent) + .bind(&carrier.tracestate) + .execute(&self.pool) + .await + { + crate::telemetry::queue::set_active_span_error(); + return Err(e.into()); + } + + Ok(id) + } + .instrument(producer_span) + .await } pub async fn enqueue_batch( @@ -55,6 +74,13 @@ impl EventQueue { return Ok(Vec::new()); } + let queue_name = "connector_events_queue"; + + // Build PRODUCER span and instrument the full batch INSERT with the span. + // inject_active_context is called from inside the instrumented block + // so the producer span is active during injection and the DB write. + let producer_span = crate::telemetry::queue::build_producer_span(queue_name); + let mut ids: Vec = Vec::with_capacity(events.len()); let mut sync_run_ids: Vec = Vec::with_capacity(events.len()); let mut source_ids: Vec = Vec::with_capacity(events.len()); @@ -69,21 +95,40 @@ impl EventQueue { payloads.push(serde_json::to_value(event)?); } - sqlx::query( - r#" - INSERT INTO connector_events_queue (id, sync_run_id, source_id, event_type, payload) - SELECT * FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[], $5::jsonb[]) - "#, - ) - .bind(&ids) - .bind(&sync_run_ids) - .bind(&source_ids) - .bind(&event_types) - .bind(&payloads) - .execute(&self.pool) - .await?; + async move { + let carrier = crate::telemetry::queue::inject_active_context(); + + let traceparents: Vec> = std::iter::repeat(carrier.traceparent.clone()) + .take(ids.len()) + .collect(); + let tracestates: Vec> = std::iter::repeat(carrier.tracestate.clone()) + .take(ids.len()) + .collect(); + + if let Err(e) = sqlx::query( + r#" + INSERT INTO connector_events_queue (id, sync_run_id, source_id, event_type, payload, traceparent, tracestate) + SELECT * FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[], $5::jsonb[], $6::text[], $7::text[]) + "#, + ) + .bind(&ids) + .bind(&sync_run_ids) + .bind(&source_ids) + .bind(&event_types) + .bind(&payloads) + .bind(&traceparents) + .bind(&tracestates) + .execute(&self.pool) + .await + { + crate::telemetry::queue::set_active_span_error(); + return Err(e.into()); + } - Ok(ids) + Ok(ids) + } + .instrument(producer_span) + .await } pub async fn dequeue_batch(&self, batch_size: i32) -> Result> { @@ -141,7 +186,9 @@ impl EventQueue { q.max_retries, q.created_at, q.processed_at, - q.error_message + q.error_message, + q.traceparent, + q.tracestate "#, ) .bind(batch_size) @@ -212,7 +259,9 @@ impl EventQueue { q.max_retries, q.created_at, q.processed_at, - q.error_message + q.error_message, + q.traceparent, + q.tracestate "#, ) .bind(batch_size) @@ -282,7 +331,9 @@ impl EventQueue { q.max_retries, q.created_at, q.processed_at, - q.error_message + q.error_message, + q.traceparent, + q.tracestate "#, ) .bind(batch_size) @@ -307,6 +358,9 @@ impl EventQueue { _ => crate::models::EventStatus::Pending, }; + let traceparent: Option = row.try_get("traceparent").ok().flatten(); + let tracestate: Option = row.try_get("tracestate").ok().flatten(); + events.push(ConnectorEventQueueItem { id: row.get("id"), sync_run_id: row.get("sync_run_id"), @@ -319,6 +373,8 @@ impl EventQueue { created_at: row.get("created_at"), processed_at: row.get("processed_at"), error_message: row.get("error_message"), + traceparent, + tracestate, }); } events @@ -576,7 +632,7 @@ impl EventQueue { r#" UPDATE connector_events_queue SET status = 'completed', processed_at = NOW() - WHERE id = ANY($1) + WHERE id = ANY($1) AND status = 'processing' "#, ) .bind(&event_ids) @@ -624,12 +680,18 @@ impl EventQueue { Ok(result.rows_affected() as i64) } + /// Mark a batch of connector events as failed or dead-letter, returning + /// the count of rows that became `failed` (retryable) and those that + /// became `dead_letter` (exhausted retries). + /// + /// Only rows currently in `processing` state are updated, preventing + /// stale/duplicate IDs from inflating counters. pub async fn mark_events_dead_letter_batch( &self, event_ids_with_errors: Vec<(String, String)>, - ) -> Result { + ) -> Result<(i64, i64)> { if event_ids_with_errors.is_empty() { - return Ok(0); + return Ok((0, 0)); } let event_ids: Vec = event_ids_with_errors @@ -641,7 +703,7 @@ impl EventQueue { .map(|(_, err)| err.clone()) .collect(); - let result = sqlx::query( + let rows: Vec<(String,)> = sqlx::query_as( r#" UPDATE connector_events_queue SET status = CASE @@ -655,14 +717,18 @@ impl EventQueue { SELECT * FROM UNNEST($1::text[], $2::text[]) AS t(id, error_message) ) AS data_table WHERE connector_events_queue.id = data_table.id + AND connector_events_queue.status = 'processing' + RETURNING connector_events_queue.status "#, ) .bind(&event_ids) .bind(&error_messages) - .execute(&self.pool) + .fetch_all(&self.pool) .await?; - Ok(result.rows_affected() as i64) + let dead_letter_count = rows.iter().filter(|(s,)| s == "dead_letter").count() as i64; + let failed_count = rows.iter().filter(|(s,)| s == "failed").count() as i64; + Ok((failed_count, dead_letter_count)) } } diff --git a/shared/src/telemetry.rs b/shared/src/telemetry.rs index a031941d2..9647bca41 100644 --- a/shared/src/telemetry.rs +++ b/shared/src/telemetry.rs @@ -2,7 +2,17 @@ use anyhow::Result; use opentelemetry::{global, trace::TracerProvider as _, KeyValue}; use opentelemetry_otlp::WithExportConfig; use opentelemetry_sdk::{ - trace::{RandomIdGenerator, Sampler, SdkTracerProvider}, + logs::{ + log_processor_with_async_runtime::BatchLogProcessor as AsyncBatchLogProcessor, + SdkLoggerProvider, + }, + metrics::{PeriodicReader, SdkMeterProvider}, + propagation::TraceContextPropagator, + runtime, + trace::{ + span_processor_with_async_runtime::BatchSpanProcessor, RandomIdGenerator, Sampler, + SdkTracerProvider, + }, Resource, }; use opentelemetry_semantic_conventions::{ @@ -14,6 +24,8 @@ use tracing::warn; use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer}; static TRACER_PROVIDER: OnceLock = OnceLock::new(); +static METER_PROVIDER: OnceLock = OnceLock::new(); +static LOGGER_PROVIDER: OnceLock = OnceLock::new(); pub struct TelemetryConfig { pub service_name: String, @@ -40,7 +52,22 @@ impl TelemetryConfig { } } +/// Normalise an OTLP endpoint: strip trailing slash. +/// Shared pure helper exposed for testing. +pub fn normalize_otlp_endpoint(endpoint: &str) -> String { + endpoint.trim_end_matches('/').to_string() +} + +/// Build the full OTLP HTTP log export URL from a base endpoint. +/// Returns `None` when the base endpoint is absent. +pub fn build_logs_url(endpoint: Option<&str>) -> Option { + endpoint.map(|e| format!("{}/v1/logs", normalize_otlp_endpoint(e))) +} + pub fn init_telemetry(config: TelemetryConfig) -> Result<()> { + // Configure W3C TraceContext propagator for trace propagation + global::set_text_map_propagator(TraceContextPropagator::new()); + let resource = Resource::builder_empty() .with_schema_url( [ @@ -55,22 +82,28 @@ pub fn init_telemetry(config: TelemetryConfig) -> Result<()> { let otlp_endpoint_for_log = config.otlp_endpoint.clone(); - let tracer_provider = if let Some(endpoint) = config.otlp_endpoint { + let endpoint = config + .otlp_endpoint + .map(|e| e.trim_end_matches('/').to_string()); + + let tracer_provider = if let Some(ref endpoint) = endpoint { + let traces_url = format!("{}/v1/traces", endpoint); let exporter = opentelemetry_otlp::SpanExporter::builder() .with_http() - .with_endpoint(&endpoint) + .with_endpoint(traces_url) .with_timeout(Duration::from_secs(10)) .build()?; + let batch = BatchSpanProcessor::builder(exporter, runtime::Tokio).build(); SdkTracerProvider::builder() - .with_batch_exporter(exporter) - .with_resource(resource) + .with_span_processor(batch) + .with_resource(resource.clone()) .with_sampler(Sampler::AlwaysOn) .with_id_generator(RandomIdGenerator::default()) .build() } else { SdkTracerProvider::builder() - .with_resource(resource) + .with_resource(resource.clone()) .with_sampler(Sampler::AlwaysOn) .with_id_generator(RandomIdGenerator::default()) .build() @@ -79,26 +112,106 @@ pub fn init_telemetry(config: TelemetryConfig) -> Result<()> { global::set_tracer_provider(tracer_provider.clone()); let _ = TRACER_PROVIDER.set(tracer_provider.clone()); + // ------------------------------------------------------------------ + // Meter provider (metrics) + // ------------------------------------------------------------------ + let meter_provider = if let Some(ref endpoint) = endpoint { + let metrics_url = format!("{}/v1/metrics", endpoint); + let metric_exporter = opentelemetry_otlp::MetricExporter::builder() + .with_http() + .with_endpoint(metrics_url) + .with_timeout(Duration::from_secs(10)) + .build()?; + + // Read OTEL_METRIC_EXPORT_INTERVAL as milliseconds, default 60000. + let metric_export_interval_ms = std::env::var("OTEL_METRIC_EXPORT_INTERVAL") + .ok() + .and_then(|v| v.parse::().ok()) + .filter(|&v| v > 0) + .unwrap_or(60000); + + let reader = PeriodicReader::builder(metric_exporter) + .with_interval(Duration::from_millis(metric_export_interval_ms)) + .build(); + + SdkMeterProvider::builder() + .with_resource(resource.clone()) + .with_reader(reader) + .build() + } else { + // No endpoint configured — build a provider that collects metrics + // but drops them. Instruments will still be callable (no crash) + // but data goes nowhere. + SdkMeterProvider::builder() + .with_resource(resource.clone()) + .build() + }; + + global::set_meter_provider(meter_provider.clone()); + let _ = METER_PROVIDER.set(meter_provider.clone()); + + // ------------------------------------------------------------------ + // Logger provider (logs) + // ------------------------------------------------------------------ + let logger_provider = if let Some(ref endpoint) = endpoint.as_deref() { + let logs_url = format!("{}/v1/logs", endpoint); + let log_exporter = opentelemetry_otlp::LogExporter::builder() + .with_http() + .with_endpoint(logs_url) + .with_timeout(Duration::from_secs(10)) + .build()?; + + SdkLoggerProvider::builder() + .with_resource(resource.clone()) + .with_log_processor( + AsyncBatchLogProcessor::builder(log_exporter, runtime::Tokio).build(), + ) + .build() + } else { + SdkLoggerProvider::builder() + .with_resource(resource.clone()) + .build() + }; + + let _ = LOGGER_PROVIDER.set(logger_provider.clone()); + + // ------------------------------------------------------------------ + // Tracing layers + // ------------------------------------------------------------------ let tracer = tracer_provider.tracer(config.service_name.clone()); let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); + // Bridge tracing events to OTel log records when a LoggerProvider + // with OTLP export is active. + let otel_log_layer = + opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(&logger_provider); + let env_filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| EnvFilter::new("info")) .add_directive("sqlx=warn".parse()?) .add_directive("hyper=info".parse()?) .add_directive("tower_http=info".parse()?); + // JSON fmt layer with current-span fields so trace_id/span_id + // (recorded on SERVER/CLIENT spans by the middleware after + // OTel context creation) appear in stdout JSON log lines. + use tracing_subscriber::fmt::format::Format; + let fmt = Format::default() + .json() + .with_current_span(true) + .with_span_list(true); let fmt_layer = tracing_subscriber::fmt::layer() + .event_format(fmt) .with_target(true) .with_level(true) .with_thread_ids(false) .with_file(false) .with_line_number(false) - .json() .with_filter(env_filter); tracing_subscriber::registry() .with(telemetry_layer) + .with(otel_log_layer) .with(fmt_layer) .init(); @@ -107,7 +220,7 @@ pub fn init_telemetry(config: TelemetryConfig) -> Result<()> { deployment_id = %config.deployment_id, environment = %config.environment, otlp_endpoint = ?otlp_endpoint_for_log, - "Telemetry initialized" + "Telemetry initialized (traces + metrics + logs)" ); Ok(()) @@ -115,75 +228,175 @@ pub fn init_telemetry(config: TelemetryConfig) -> Result<()> { pub async fn shutdown_telemetry() { tracing::info!("Shutting down telemetry"); + + // Shut down meter provider first (flush any pending metric exports). + // This is safe to log through because the logger provider is still + // active. + if let Some(meter_provider) = METER_PROVIDER.get() { + if let Err(error) = meter_provider.shutdown() { + warn!(%error, "Failed to shut down meter provider"); + } + } + + // Shut down tracer provider (flush any pending trace exports). if let Some(tracer_provider) = TRACER_PROVIDER.get() { if let Err(error) = tracer_provider.shutdown() { warn!(%error, "Failed to shut down tracer provider"); } } + + // Shut down logger provider last so that final correlated log + // records (including the shutdown messages above) are exported + // before the provider is torn down. + if let Some(logger_provider) = LOGGER_PROVIDER.get() { + if let Err(error) = logger_provider.shutdown() { + warn!(%error, "Failed to shut down logger provider"); + } + } +} + +/// Wait for SIGTERM or Ctrl-C. Use this with Axum's +/// `with_graceful_shutdown` to drain connections before +/// shutting down telemetry. +pub async fn shutdown_signal() { + let ctrl_c = async { + tokio::signal::ctrl_c() + .await + .expect("Failed to install Ctrl+C handler"); + }; + + #[cfg(unix)] + let terminate = async { + tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate()) + .expect("Failed to install SIGTERM handler") + .recv() + .await; + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => {}, + _ = terminate => {}, + } + + tracing::info!("Shutdown signal received, starting graceful shutdown"); } pub mod middleware { - use axum::{extract::Request, http::HeaderMap, middleware::Next, response::Response}; - use opentelemetry::{ - global, - trace::{SpanKind, TraceContextExt, Tracer}, - Context, + use crate::metrics; + use axum::{ + extract::MatchedPath, extract::Request, http::HeaderMap, middleware::Next, + response::Response, }; + use opentelemetry::{global, trace::TraceContextExt}; use opentelemetry_http::{HeaderExtractor, HeaderInjector}; use tracing::{Instrument, Span}; use tracing_opentelemetry::OpenTelemetrySpanExt; - pub async fn trace_layer(mut request: Request, next: Next) -> Response { - let headers = request.headers(); + /// Axum middleware that creates a single OpenTelemetry SERVER span per request + /// **and** records HTTP RED metrics (request count + duration histogram) via + /// the central [`metrics::record_http_request`] helper. + /// + /// - Extracts W3C TraceContext from incoming headers and sets it as the parent. + /// - Uses the matched route template (never raw `url.path`) and sets `otel.name` + /// to `METHOD /route`. + /// - Records `http.request.method`, `http.route`, `http.response.status_code`. + /// - Records `trace_id` / `span_id` onto the span after OTel context creation + /// so the JSON formatter (with `with_current_span(true)`) includes them in + /// stdout log lines emitted within this request. + /// - Marks 5xx responses as error spans. + pub async fn trace_layer(request: Request, response: Next) -> Response { + let start = std::time::Instant::now(); + let parent_context = global::get_text_map_propagator(|propagator| { - propagator.extract(&HeaderExtractor(headers)) + propagator.extract(&HeaderExtractor(request.headers())) }); - let tracer = global::tracer("http-server"); - let span_builder = tracer - .span_builder(format!("{} {}", request.method(), request.uri().path())) - .with_kind(SpanKind::Server); + let method = request.method().clone(); + let route = request + .extensions() + .get::() + .map(|mp| mp.as_str().to_string()); - let otel_span = tracer.build_with_context(span_builder, &parent_context); - let context = Context::current_with_span(otel_span); + let otel_name = match &route { + Some(r) => format!("{} {}", method, r), + None => format!("{} /unknown", method), + }; - let tracing_span = tracing::info_span!( - "http_request", - method = %request.method(), - uri = %request.uri(), - version = ?request.version(), + // Build a single tracing span that tracing-opentelemetry converts into + // one OTel SERVER span. The `otel.name` special field overrides the + // OTel span name; `otel.kind` sets the span kind. + // + // `trace_id` and `span_id` are declared as empty fields here. They are + // recorded after `set_parent` below once the OTel context exists. + let span = tracing::info_span!( + "HTTP request", + otel.name = otel_name.as_str(), + otel.kind = "SERVER", + http.request.method = %method, + http.route = tracing::field::Empty, + http.response.status_code = tracing::field::Empty, + trace_id = tracing::field::Empty, + span_id = tracing::field::Empty, ); - let _ = tracing_span.set_parent(context.clone()); - - let request_id = { - let span = context.span(); - span.span_context().trace_id().to_string() - }; - - request.extensions_mut().insert(request_id.clone()); + if let Some(ref route_val) = route { + span.record("http.route", route_val.as_str()); + } - let response = next.run(request).instrument(tracing_span.clone()).await; + // Use the extracted parent context + let _ = span.set_parent(parent_context); + + // Record trace_id and span_id from the (now-set) OTel context onto + // the tracing span so JSON stdout lines include trace correlation. + { + let otel_context = span.context(); + let span_in_context = otel_context.span(); + let sc = span_in_context.span_context(); + if sc.is_valid() { + span.record("trace_id", sc.trace_id().to_string()); + span.record("span_id", sc.span_id().to_string()); + } + } - tracing_span.in_scope(|| { - tracing::info!( - status = response.status().as_u16(), - request_id = %request_id, - "Request completed" + async move { + let resp = response.run(request).await; + let status = resp.status().as_u16() as i64; + Span::current().record("http.response.status_code", status); + + if status >= 500 { + Span::current().set_status(opentelemetry::trace::Status::error("Server error")); + } + + // Record HTTP RED metrics via central helper (bounded attributes only). + let duration_secs = start.elapsed().as_secs_f64(); + metrics::record_http_request( + method.as_str(), + route.as_deref(), + status as u16, + duration_secs, ); - }); - response + tracing::info!(status, duration = duration_secs, "Request completed"); + resp + } + .instrument(span) + .await } - pub fn inject_trace_context(headers: &mut HeaderMap, span: &Span) { - let context = span.context(); + /// Inject W3C trace context from the current tracing span into a HeaderMap + /// (for outbound calls that do not use reqwest). + pub fn inject_trace_context(headers: &mut HeaderMap) { + let context = Span::current().context(); let mut injector = HeaderInjector(headers); global::get_text_map_propagator(|propagator| { propagator.inject_context(&context, &mut injector); }); } + /// Extract trace context from request headers and return the trace ID if valid. pub fn get_request_id_from_headers(headers: &HeaderMap) -> Option { let context = global::get_text_map_propagator(|propagator| { propagator.extract(&HeaderExtractor(headers)) @@ -200,38 +413,483 @@ pub mod middleware { } pub mod http_client { - use axum::http::HeaderMap; use opentelemetry::global; use opentelemetry_http::HeaderInjector; - use reqwest::RequestBuilder; - use tracing::Span; + use tracing::{Instrument, Span}; use tracing_opentelemetry::OpenTelemetrySpanExt; - pub fn inject_trace_headers(mut builder: RequestBuilder) -> RequestBuilder { - let span = Span::current(); - let context = span.context(); + /// Send an HTTP request wrapped in an OpenTelemetry [`SpanKind::Client`] span. + /// + /// The span records `http.request.method`, `server.address`, + /// `http.response.status_code`, and marks 5xx / transport errors as errors. + /// W3C trace context is injected using the client span so downstream + /// services see the correct parent span ID. + /// + /// # Arguments + /// * `method` – HTTP method string (e.g. `"GET"`, `"POST"`) + /// * `url` – Full request URL (only the host part is recorded as `server.address`) + /// * `builder` – The pre-configured `reqwest::RequestBuilder` + /// + /// The returned future is Send-safe (no EnteredGuard is held across await). + pub async fn send_traced( + method: &str, + url: &str, + builder: reqwest::RequestBuilder, + ) -> Result { + let server = url::Url::parse(url) + .ok() + .and_then(|u| u.host_str().map(String::from)) + .unwrap_or_else(|| "unknown".to_string()); + + let parent_cx = Span::current().context(); + let span_name = format!("HTTP {}", method); + + // Create a CLIENT tracing span; tracing-opentelemetry converts it into + // an OTel client span with the recorded attributes. + let span = tracing::info_span!( + "HTTP client request", + otel.name = span_name.as_str(), + otel.kind = "CLIENT", + http.request.method = %method, + server.address = %server, + http.response.status_code = tracing::field::Empty, + ); - let mut headers = HeaderMap::new(); - let mut injector = HeaderInjector(&mut headers); + let _ = span.set_parent(parent_cx); + // Enter the span once so tracing-opentelemetry creates the OTel CLIENT + // span and stores its OTel context in the tracing span's extensions. + // The guard is dropped immediately – no EnteredGuard held across await. + { + let _guard = span.clone().entered(); + } + + // Retrieve the OTel context from the span's extensions (seeded above). + // This avoids holding an EnteredGuard across any await point. + let span_cx = span.context(); + + // Inject trace context from the client span so downstream + // services see this span as the parent. + let mut headers = http::HeaderMap::new(); + let mut injector = HeaderInjector(&mut headers); global::get_text_map_propagator(|propagator| { - propagator.inject_context(&context, &mut injector); + propagator.inject_context(&span_cx, &mut injector); }); + let mut builder = builder; for (key, value) in headers.iter() { builder = builder.header(key.clone(), value.clone()); } - builder + // Execute the request inside the span via Instrument so that no + // EnteredGuard is held across .await. + async move { + let result = builder.send().await; + + match &result { + Ok(response) => { + let status = response.status().as_u16() as i64; + Span::current().record("http.response.status_code", status); + if status >= 500 { + Span::current() + .set_status(opentelemetry::trace::Status::error("Server error")); + } + } + Err(e) => { + Span::current().set_status(opentelemetry::trace::Status::error(e.to_string())); + } + } + + result + } + .instrument(span) + .await + } +} + +/// Queue telemetry helpers for W3C trace context propagation through async queues. +/// +/// This module provides: +/// - `inject_producer_context`: creates a PRODUCER span and extracts W3C headers +/// - `extract_remote_span_context`: reconstructs a valid remote SpanContext from stored headers +/// - `add_consumer_links`: adds Link entries to a consumer span from extracted contexts +/// +/// Producer spans are parented to the current tracing context (the caller). +/// Consumer spans MUST be created as new roots (no parent) and reference the producer +/// context via links only. +pub mod queue { + use opentelemetry::{ + global, + propagation::{Extractor, Injector}, + trace::{SpanContext, SpanId, TraceContextExt, TraceId}, + Context, + }; + use std::collections::HashSet; + use tracing::{info_span, Span}; + use tracing_opentelemetry::OpenTelemetrySpanExt; + + const TRACEPARENT_MAX_LEN: usize = 55; + const TRACESTATE_MAX_LEN: usize = 512; + + /// A text map carrier for W3C trace context. + #[derive(Debug, Clone, Default)] + pub struct TraceCarrier { + pub traceparent: Option, + pub tracestate: Option, + } + + impl Injector for TraceCarrier { + fn set(&mut self, key: &str, value: String) { + match key.to_lowercase().as_str() { + "traceparent" => self.traceparent = Some(value), + "tracestate" => self.tracestate = Some(value), + _ => {} + } + } + } + + impl Extractor for TraceCarrier { + fn get(&self, key: &str) -> Option<&str> { + match key.to_lowercase().as_str() { + "traceparent" => self.traceparent.as_deref(), + "tracestate" => self.tracestate.as_deref(), + _ => None, + } + } + + fn keys(&self) -> Vec<&str> { + let mut keys = Vec::new(); + if self.traceparent.is_some() { + keys.push("traceparent"); + } + if self.tracestate.is_some() { + keys.push("tracestate"); + } + keys + } + } + + /// Build a PRODUCER tracing span for the given queue, parented to the + /// current scope. The span is returned *un-entered* so the caller can + /// enter it briefly to inject context before instrumenting the actual + /// async operation. + pub fn build_producer_span(queue_name: &str) -> Span { + let span_name = format!("{} publish", queue_name); + info_span!( + target: "queue", + parent: &Span::current(), + "{}", span_name, + otel.name = span_name.as_str(), + otel.kind = "PRODUCER", + messaging.system = "postgresql", + messaging.destination = queue_name, + messaging.operation.type = "publish", + ) + } + + /// Inject W3C trace context from the *currently active* tracing span + /// into a TraceCarrier. Call this while the producer span is entered. + pub fn inject_active_context() -> TraceCarrier { + let cx = Span::current().context(); + let mut carrier = TraceCarrier::default(); + global::get_text_map_propagator(|propagator| propagator.inject_context(&cx, &mut carrier)); + carrier + } + + /// Extract a valid remote SpanContext from optional stored W3C headers. + pub fn extract_remote_span_context( + traceparent: Option<&str>, + tracestate: Option<&str>, + ) -> Option { + let traceparent = traceparent?; + if traceparent.len() > TRACEPARENT_MAX_LEN { + return None; + } + let mut carrier = TraceCarrier::default(); + Injector::set(&mut carrier, "traceparent", traceparent.to_string()); + if let Some(ts) = tracestate { + if ts.len() <= TRACESTATE_MAX_LEN { + Injector::set(&mut carrier, "tracestate", ts.to_string()); + } + } + let cx: Context = + global::get_text_map_propagator(|p| p.extract_with_context(&Context::new(), &carrier)); + let span = cx.span(); + let sc = span.span_context(); + if sc.is_valid() { + Some(sc.clone()) + } else { + None + } + } + + /// Collect deduplicated producer SpanContext values from queue items. + /// Deduplication is by the full `(trace_id, span_id)` pair (not trace + /// only), so two distinct producer spans within the same trace are + /// both retained. Invalid/missing contexts are skipped. + pub fn collect_span_contexts( + items: &[Option<(Option, Option)>], + ) -> Vec { + let mut seen: HashSet<(TraceId, SpanId)> = HashSet::new(); + let mut result = Vec::new(); + for ctx in items { + let (tp, ts) = match ctx { + Some((tp, ts)) => (tp.as_deref(), ts.as_deref()), + None => continue, + }; + let Some(sc) = extract_remote_span_context(tp, ts) else { + continue; + }; + let key = (sc.trace_id(), sc.span_id()); + if seen.insert(key) { + result.push(sc.clone()); + } + } + result } - pub trait RequestBuilderExt { - fn with_trace_context(self) -> Self; + /// Mark the current active span as error with a fixed empty description. + /// Used to record producer-side failures without leaking message details + /// into the span status. + pub fn set_active_span_error() { + Span::current().set_status(opentelemetry::trace::Status::error("")); } - impl RequestBuilderExt for RequestBuilder { - fn with_trace_context(self) -> Self { - inject_trace_headers(self) + /// Validate W3C traceparent format. + pub fn is_valid_traceparent(value: &str) -> bool { + if value.len() != 55 { + return false; + } + let parts: Vec<&str> = value.split('-').collect(); + parts.len() == 4 + && parts[0] == "00" + && parts[1].len() == 32 + && parts[1].chars().all(|c| c.is_ascii_hexdigit()) + && parts[2].len() == 16 + && parts[2].chars().all(|c| c.is_ascii_hexdigit()) + && parts[3].len() == 2 + && parts[3].chars().all(|c| c.is_ascii_hexdigit()) + } + + #[cfg(test)] + mod tests { + use super::*; + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::propagation::TraceContextPropagator; + use tracing_subscriber::layer::SubscriberExt as _; + use tracing_subscriber::util::SubscriberInitExt; + + /// Initialize the global TraceContextPropagator once per test run. + fn init_propagator() { + use std::sync::Once; + static INIT: Once = Once::new(); + INIT.call_once(|| { + global::set_text_map_propagator(TraceContextPropagator::new()); + }); + } + + #[test] + fn test_is_valid_traceparent() { + assert!(is_valid_traceparent( + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + )); + assert!(!is_valid_traceparent( + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-010" + )); + assert!(!is_valid_traceparent("too-short")); + assert!(!is_valid_traceparent( + "01-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" + )); + assert!(!is_valid_traceparent( + "00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-b7ad6b7169203331-01" + )); + } + + #[test] + fn test_extract_remote_span_context_valid() { + init_propagator(); + let sc = extract_remote_span_context( + Some("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"), + None, + ); + assert!(sc.is_some()); + let sc = sc.unwrap(); + assert_eq!( + sc.trace_id(), + TraceId::from_hex("0af7651916cd43dd8448eb211c80319c").unwrap() + ); + assert!(sc.is_sampled()); + } + + #[test] + fn test_extract_remote_span_context_invalid() { + assert!(extract_remote_span_context(None, None).is_none()); + assert!(extract_remote_span_context( + Some("00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-b7ad6b7169203331-01"), + None, + ) + .is_none()); + assert!(extract_remote_span_context( + Some("00-00000000000000000000000000000000-0000000000000000-01"), + None, + ) + .is_none()); + } + + #[test] + fn test_extract_remote_span_context_with_tracestate() { + let sc = extract_remote_span_context( + Some("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"), + Some("vendor1=value1"), + ); + assert!(sc.is_some()); + } + + #[test] + fn test_extract_remote_span_context_ambient_isolation_malformed_55char() { + // Activate an unrelated ambient span, then supply a malformed + // exactly-55-char traceparent with non-hex characters. + // The propagator must NOT fall back to the ambient context. + init_tracer(); + + let parent_span = info_span!("ambient_span"); + let _guard = parent_span.entered(); + + // 55 chars, valid format, but non-hex chars in trace-id + let malformed = "00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-b7ad6b7169203331-01"; + assert_eq!(malformed.len(), 55); + + let result = extract_remote_span_context(Some(malformed), None); + assert!( + result.is_none(), + "must not return ambient context when traceparent has non-hex chars" + ); + } + + #[test] + fn test_extract_remote_span_context_ambient_isolation_valid() { + // Activate an unrelated ambient span, then supply a valid but + // different traceparent. The extracted context must match the + // carrier, not the ambient. + init_tracer(); + + let parent_span = info_span!("ambient_span"); + let _guard = parent_span.entered(); + + let carrier_tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + let result = extract_remote_span_context(Some(carrier_tp), None); + assert!( + result.is_some(), + "valid traceparent must be extracted even with active ambient span" + ); + let sc = result.unwrap(); + assert_eq!( + sc.trace_id(), + TraceId::from_hex("0af7651916cd43dd8448eb211c80319c").unwrap(), + "extracted trace_id must come from carrier, not ambient" + ); + } + + #[test] + fn test_collect_span_contexts_deduplicates() { + init_propagator(); + let tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + let items = vec![ + Some((Some(tp.to_string()), None)), + Some((Some(tp.to_string()), None)), + ]; + let result = collect_span_contexts(&items); + assert_eq!( + result.len(), + 1, + "duplicate (trace_id, span_id) should be deduplicated" + ); + } + + #[test] + fn test_collect_span_contexts_multiple_traces() { + init_propagator(); + let tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + let tp2 = "00-1bf7651916cd43dd8448eb211c80319c-b8ad6b7169203331-01"; + let items = vec![ + Some((Some(tp1.to_string()), None)), + Some((Some(tp2.to_string()), None)), + ]; + assert_eq!(collect_span_contexts(&items).len(), 2); + } + + #[test] + fn test_collect_span_contexts_same_trace_different_spans_retained() { + init_propagator(); + let tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + let tp2 = "00-0af7651916cd43dd8448eb211c80319c-aaaaaaaaaaaaaaaa-01"; + let items = vec![ + Some((Some(tp1.to_string()), None)), + Some((Some(tp2.to_string()), None)), + ]; + let result = collect_span_contexts(&items); + assert_eq!( + result.len(), + 2, + "two producer spans with same trace_id but different span_id must both be retained" + ); + } + + #[test] + fn test_collect_span_contexts_skips_invalid() { + let items = vec![ + None, + Some((None, None)), + Some((Some("invalid".to_string()), None)), + ]; + assert!(collect_span_contexts(&items).is_empty()); + } + + /// Set up a global tracer provider and tracing subscriber for tests + /// that need to produce trace context. + fn init_tracer() { + init_propagator(); + static TRACER_INIT: std::sync::Once = std::sync::Once::new(); + TRACER_INIT.call_once(|| { + let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build(); + global::set_tracer_provider(provider.clone()); + let tracer = provider.tracer("test"); + let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); + let _ = tracing_subscriber::registry() + .with(telemetry_layer) + .try_init(); + }); + } + + #[test] + fn test_build_producer_span_and_inject() { + init_tracer(); + let parent_span = info_span!("test"); + let _guard = parent_span.entered(); + let span = build_producer_span("test_queue"); + let carrier = { + let _g = span.entered(); + inject_active_context() + }; + assert!(carrier.traceparent.is_some()); + let tp = carrier.traceparent.unwrap(); + assert_eq!(tp.len(), 55); + assert!(is_valid_traceparent(&tp)); + } + + #[test] + fn test_trace_carrier_round_trip() { + init_tracer(); + let parent_span = info_span!("test"); + let _guard = parent_span.entered(); + let span = build_producer_span("round_trip"); + let carrier = { + let _g = span.entered(); + inject_active_context() + }; + let tp = carrier.traceparent.clone().unwrap(); + let extracted = extract_remote_span_context(Some(&tp), carrier.tracestate.as_deref()); + assert!(extracted.is_some()); } } } diff --git a/shared/tests/collector/server.mjs b/shared/tests/collector/server.mjs new file mode 100644 index 000000000..b03fe578a --- /dev/null +++ b/shared/tests/collector/server.mjs @@ -0,0 +1,76 @@ +import { createServer } from 'http'; + +const captured = { spans: [], metricNames: [], serviceNames: new Set() }; + +function parseAttrs(attrs) { + return (attrs || []).map(a => [a.key, a.value?.stringValue || '']).filter(([,v]) => v); +} +function parseLinks(links) { + return (links || []).map(l => ({ trace_id: l.traceId || l.trace_id, span_id: l.spanId || l.span_id })); +} + +function collectSpans(rsp) { + let svc = ''; + for (const a of rsp.resource?.attributes || []) { + if (a.key === 'service.name' && a.value?.stringValue) svc = a.value.stringValue; + } + if (svc) captured.serviceNames.add(svc); + for (const ss of rsp.scopeSpans || []) { + for (const span of ss.spans || []) { + captured.spans.push({ + trace_id: span.traceId, span_id: span.spanId, + parent_span_id: span.parentSpanId || '', + name: span.name, kind: span.kind, service_name: svc, + status_code: span.status?.code || 0, + attributes: parseAttrs(span.attributes), + links: parseLinks(span.links), + }); + } + } +} + +function parseTracesJSON(body) { + const data = JSON.parse(body); + for (const rsp of data.resourceSpans || []) collectSpans(rsp); +} + +function parseMetricsJSON(body) { + const data = JSON.parse(body); + for (const rm of data.resourceMetrics || []) { + for (const sm of rm.scopeMetrics || []) { + for (const m of sm.metrics || []) captured.metricNames.push(m.name); + } + } +} + +const server = createServer((req, res) => { + let body = []; + req.on('data', c => body.push(c)); + req.on('end', () => { + const buf = Buffer.concat(body); + try { + if (req.method === 'POST') { + if (req.url === '/v1/traces') { + parseTracesJSON(buf.toString()); + } else if (req.url === '/v1/metrics') { + parseMetricsJSON(buf.toString()); + } + res.writeHead(200); + } else if (req.method === 'GET' && req.url === '/inspect') { + res.writeHead(200, { 'Content-Type': 'application/json' }); + res.end(JSON.stringify({ + spans: captured.spans, + service_names: [...captured.serviceNames], + metric_names: captured.metricNames, + })); + return; + } else if (req.method === 'POST' && req.url === '/reset') { + captured.spans = []; captured.metricNames = []; captured.serviceNames.clear(); + res.writeHead(200); + } else { res.writeHead(404); } + } catch (e) { process.stderr.write(`Err: ${e.message}\n`); res.writeHead(500); } + res.end(); + }); +}); + +server.listen(4318, '0.0.0.0', () => process.stderr.write('ready\n')); diff --git a/shared/tests/cross_service_tracing.rs b/shared/tests/cross_service_tracing.rs new file mode 100644 index 000000000..2b22e2e85 --- /dev/null +++ b/shared/tests/cross_service_tracing.rs @@ -0,0 +1,484 @@ +//! E2E cross-service trace continuity test. +//! +//! 1. Tags pre-built local dev images as production-style e2e images. +//! 2. Builds and starts a mock OTLP collector as a compose service. +//! 3. Starts the Omni stack pointing at the collector with JSON protocol. +//! 4. Triggers a search via `docker exec` inside the web container. +//! 5. Inspects captured spans from the mock collector via `docker exec`. +//! 6. Asserts cross-service trace continuity. +//! +//! Requires Docker. Skips gracefully when unavailable. + +use std::collections::{HashMap, HashSet}; +use std::path::PathBuf; +use std::process::{Command, Stdio}; +use std::time::{Duration, Instant}; + +use serde::Deserialize; + +const E2E_TAG: &str = "e2e"; + +// ───────────────────────────────────────────────────────── +// Helpers +// ───────────────────────────────────────────────────────── + +fn docker_available() -> bool { + Command::new("docker") + .arg("info") + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .map(|s| s.success()) + .unwrap_or(false) +} + +fn workspace_root() -> PathBuf { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .expect("shared is not inside a workspace") + .to_path_buf() +} + +const CORE_SERVICES: &[(&str, &str)] = &[ + ("searcher", "omni-searcher"), + ("indexer", "omni-indexer"), + ("ai", "omni-ai"), + ("connector-manager", "omni-connector-manager"), + ("web", "omni-web"), + ("migrator", "omni-migrator"), + ("sandbox", "omni-sandbox"), +]; + +fn prod_image_name(svc: &str) -> String { + format!("ghcr.io/getomnico/omni/{svc}:{E2E_TAG}") +} + +fn dev_image_name(svc: &str) -> String { + format!("omni-{svc}:dev") +} + +// ───────────────────────────────────────────────────────── +// Inspect types +// ───────────────────────────────────────────────────────── + +#[derive(Clone, Debug, Deserialize)] +struct SpanSummary { + trace_id: String, + span_id: String, + parent_span_id: String, + name: String, + kind: i32, + service_name: String, + #[allow(dead_code)] + status_code: i32, + #[allow(dead_code)] + attributes: Vec<(String, String)>, + #[allow(dead_code)] + links: Vec<(String, String)>, +} + +#[derive(Clone, Debug, Deserialize)] +#[allow(dead_code)] +struct InspectResponse { + spans: Vec, + service_names: Vec, + metric_names: Vec, +} + +// ───────────────────────────────────────────────────────── +// Compose orchestration +// ───────────────────────────────────────────────────────── + +struct ComposeStack; + +impl ComposeStack { + fn env_file() -> PathBuf { + let root = workspace_root(); + let dot_env = root.join(".env"); + if dot_env.exists() { + dot_env + } else { + root.join(".env.example") + } + } + + fn tag_images() { + for (svc_name, full_name) in CORE_SERVICES { + let tag_status = Command::new("docker") + .args([ + "tag", + &dev_image_name(svc_name), + &prod_image_name(full_name), + ]) + .stdout(Stdio::null()) + .stderr(Stdio::inherit()) + .status() + .expect("docker tag failed"); + assert!(tag_status.success(), "Failed to tag {full_name}"); + } + } + + fn build_mock_collector() -> String { + let tag = "omni-mock-collector:e2e"; + let tmp_dir = std::env::temp_dir().join("omni-e2e-collector"); + let _ = std::fs::create_dir_all(&tmp_dir); + + // Node.js collector that parses JSON OTLP (Node exports) and protobuf OTLP (Rust exports) + std::fs::write( + tmp_dir.join("server.mjs"), + include_str!("collector/server.mjs"), + ) + .expect("failed to write collector server.mjs"); + + std::fs::write( + tmp_dir.join("package.json"), + r#"{"dependencies":{"protobufjs":"^7.4.0"}}"#, + ) + .expect("failed to write package.json"); + + std::fs::write(tmp_dir.join("Dockerfile"), + "FROM node:22-alpine\nWORKDIR /app\nCOPY package.json .\nRUN npm install --no-audit --no-fund 2>&1 | tail -1\nCOPY server.mjs .\nCMD [\"node\", \"server.mjs\"]\n") + .expect("failed to write Dockerfile"); + + let status = Command::new("docker") + .args(["build", "-t", tag, "."]) + .current_dir(&tmp_dir) + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .status() + .expect("docker build failed"); + assert!(status.success(), "Failed to build mock collector"); + let _ = std::fs::remove_dir_all(&tmp_dir); + tag.to_string() + } + + async fn start() -> (String, PathBuf) { + let project = format!("omni-e2e-{}", std::process::id()); + let tag = format!("e2e-{}", std::process::id()); + let collector_img = Self::build_mock_collector(); + + let override_yaml = format!( + r#"services: + mock-collector: + image: {collector_img} + container_name: omni-{tag}-collector + networks: + - omni-network + caddy: + profiles: ["e2e-skip"] + web-connector: + profiles: ["e2e-skip"] + postgres: + container_name: omni-{tag}-postgres + redis: + container_name: omni-{tag}-redis + migrator: + container_name: omni-{tag}-migrator + searcher: + container_name: omni-{tag}-searcher + environment: + OTEL_EXPORTER_OTLP_PROTOCOL: http/json + indexer: + container_name: omni-{tag}-indexer + environment: + OTEL_EXPORTER_OTLP_PROTOCOL: http/json + ai: + container_name: omni-{tag}-ai + environment: + OTEL_EXPORTER_OTLP_PROTOCOL: http/json + sandbox: + container_name: omni-{tag}-sandbox + connector-manager: + container_name: omni-{tag}-connector-manager + environment: + OTEL_EXPORTER_OTLP_PROTOCOL: http/json + web: + container_name: omni-{tag}-web + environment: + OTEL_EXPORTER_OTLP_PROTOCOL: http/json +volumes: + postgres_data: + name: omni-{tag}-postgres-data + redis_data: + name: omni-{tag}-redis-data + ai_models: + name: omni-{tag}-ai-models + sandbox_data: + name: omni-{tag}-sandbox-data +"#, + ); + let tmp = std::env::temp_dir().join(format!("omni-e2e-{}.yml", std::process::id())); + std::fs::write(&tmp, &override_yaml).expect("failed to write temp compose override"); + + let root = workspace_root(); + let compose_yml = root.join("docker/docker-compose.yml"); + let env_file = Self::env_file(); + + let status = Command::new("docker") + .args([ + "compose", + "--env-file", + env_file.to_str().unwrap(), + "-f", + compose_yml.to_str().unwrap(), + "-f", + tmp.to_str().unwrap(), + "-p", + &project, + "up", + "-d", + "--pull", + "never", + "mock-collector", + "postgres", + "redis", + "migrator", + "searcher", + "indexer", + "ai", + "connector-manager", + "sandbox", + "web", + ]) + .env("OMNI_VERSION", E2E_TAG) + .env("OTEL_EXPORTER_OTLP_ENDPOINT", "http://mock-collector:4318") + .env("OTEL_DEPLOYMENT_ID", "e2e-test") + .env("OTEL_METRIC_EXPORT_INTERVAL", "5000") + .stdout(Stdio::inherit()) + .stderr(Stdio::inherit()) + .status() + .expect("docker compose up failed"); + assert!(status.success(), "docker compose up failed"); + + // Wait for collector + let coll_cnt = format!("omni-{tag}-collector"); + let deadline = Instant::now() + Duration::from_secs(30); + loop { + if Instant::now() > deadline { + panic!("Collector not ready"); + } + if Command::new("docker") + .args([ + "exec", + &coll_cnt, + "node", + "-e", + "fetch('http://localhost:4318/inspect')", + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status() + .map_or(false, |s| s.success()) + { + break; + } + tokio::time::sleep(Duration::from_secs(2)).await; + } + + // Wait for web health + let web_cnt = format!("omni-{tag}-web"); + let deadline = Instant::now() + Duration::from_secs(120); + loop { + if Instant::now() > deadline { + panic!("Services not healthy"); + } + if Command::new("docker") + .args(["exec", &web_cnt, "node", "-e", + "fetch('http://localhost:5173/api/v1/health').then(r=>r.ok&&process.exit(0)).catch(()=>process.exit(1))"]) + .stdout(Stdio::null()).stderr(Stdio::null()) + .status().map_or(false, |s| s.success()) { break; } + tokio::time::sleep(Duration::from_secs(3)).await; + } + + (project, tmp) + } + + fn down(project: &str, tmp: &PathBuf) { + let _ = Command::new("docker") + .args(["compose", "-p", project, "down", "-t", "30"]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + for (_, full_name) in CORE_SERVICES { + let _ = Command::new("docker") + .args(["rmi", "--force", &prod_image_name(full_name)]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + } + let _ = std::fs::remove_file(tmp); + } +} + +// ───────────────────────────────────────────────────────── +// Assertion helpers +// ───────────────────────────────────────────────────────── + +fn group_by_trace(spans: &[SpanSummary]) -> HashMap> { + let mut groups: HashMap> = HashMap::new(); + for s in spans { + groups.entry(s.trace_id.clone()).or_default().push(s); + } + groups +} + +fn assert_trace_chain(spans: &[&SpanSummary], label: &str) { + assert!(!spans.is_empty(), "{label}: no spans in trace"); + let trace_ids: HashSet<&str> = spans.iter().map(|s| s.trace_id.as_str()).collect(); + assert_eq!( + trace_ids.len(), + 1, + "{label}: all spans must share one trace_id" + ); + let roots: Vec<&&SpanSummary> = spans + .iter() + .filter(|s| s.parent_span_id.is_empty()) + .collect(); + assert_eq!(roots.len(), 1, "{label}: expected exactly one root span"); + let span_ids: HashSet<&str> = spans.iter().map(|s| s.span_id.as_str()).collect(); + for s in spans { + if !s.parent_span_id.is_empty() { + assert!( + span_ids.contains(s.parent_span_id.as_str()), + "{label}: span `{}` parent_span_id {} not found", + s.name, + s.parent_span_id + ); + } + } +} + +// ───────────────────────────────────────────────────────── +// Test +// ───────────────────────────────────────────────────────── + +#[tokio::test] +async fn test_cross_service_tracing() { + if !docker_available() { + eprintln!("SKIP: Docker not available"); + return; + } + + println!("Tagging dev images..."); + ComposeStack::tag_images(); + + println!("Starting compose stack..."); + let (project, tmp_path) = ComposeStack::start().await; + let _guard = TestGuard { + project: project.clone(), + tmp: tmp_path.clone(), + }; + + let tag = format!("e2e-{}", std::process::id()); + let coll = format!("omni-{tag}-collector"); + let web = format!("omni-{tag}-web"); + + // Reset and trigger search + let _ = Command::new("docker") + .args([ + "exec", + &coll, + "node", + "-e", + "fetch('http://localhost:4318/reset',{method:'POST'})", + ]) + .stdout(Stdio::null()) + .stderr(Stdio::null()) + .status(); + + let status = Command::new("docker") + .args(["exec", &web, "node", "-e", + "fetch('http://localhost:5173/api/search',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({query:'hello',mode:'fulltext'})}).then(r=>{if(!r.ok){console.error('FAIL',r.status);process.exit(1)}console.log('Search OK')})"]) + .stdout(Stdio::inherit()).stderr(Stdio::inherit()) + .status().expect("docker exec failed"); + assert!(status.success(), "Search failed"); + + tokio::time::sleep(Duration::from_secs(10)).await; + + // Inspect + let output = Command::new("docker") + .args([ + "exec", + &coll, + "node", + "-e", + "fetch('http://localhost:4318/inspect').then(r=>r.text()).then(t=>console.log(t))", + ]) + .output() + .expect("inspect failed"); + assert!(output.status.success(), "Inspect failed"); + + let inspect: InspectResponse = serde_json::from_slice(&output.stdout).unwrap_or_else(|e| { + panic!( + "JSON parse: {e}\nstdout: {}", + String::from_utf8_lossy(&output.stdout) + ) + }); + + let spans = &inspect.spans; + println!( + "Captured {} spans across services: {:?}", + spans.len(), + inspect.service_names + ); + for s in spans { + println!( + " kind={} svc={:?} name={:?} trace={} parent={}", + s.kind, + s.service_name, + s.name, + &s.trace_id[..8.min(s.trace_id.len())], + &s.parent_span_id[..8.min(s.parent_span_id.len())] + ); + } + + // Assert cross-service coverage + assert!( + inspect.service_names.len() >= 2, + "Expected ≥2 services, got {:?}", + inspect.service_names + ); + + // Find the search trace — it's the one with a POST /search span + let groups = group_by_trace(spans); + let Some((_tid, best)) = groups.iter().find(|(_, v)| { + v.iter() + .any(|s| s.name.contains("/search") || s.name.contains("POST")) + }) else { + panic!("No trace with a search request"); + }; + println!("Search trace: {} spans", best.len()); + + assert_trace_chain(best, "search trace"); + + // Verify web and searcher both present + let has_web = best.iter().any(|s| s.service_name.contains("web")); + let has_searcher = best.iter().any(|s| s.service_name.contains("searcher")); + assert!(has_web, "No web spans in search trace"); + assert!(has_searcher, "No searcher spans in search trace"); + + // Verify parent/child chain: the searcher SERVER span should have + // a parent (the web CLIENT span) in this trace. + for s in best + .iter() + .filter(|s| s.service_name.contains("searcher") && s.kind == 2) + { + assert!( + !s.parent_span_id.is_empty(), + "Searcher SERVER span should have a web parent" + ); + } + + println!("✓ Cross-service trace continuity verified"); +} + +struct TestGuard { + project: String, + tmp: PathBuf, +} + +impl Drop for TestGuard { + fn drop(&mut self) { + ComposeStack::down(&self.project, &self.tmp); + } +} diff --git a/shared/tests/embedding_queue_test.rs b/shared/tests/embedding_queue_test.rs index beabc17da..c9e639119 100644 --- a/shared/tests/embedding_queue_test.rs +++ b/shared/tests/embedding_queue_test.rs @@ -3,6 +3,7 @@ mod tests { use shared::embedding_queue::EmbeddingQueue; use shared::test_environment::TestEnvironment; use sqlx::PgPool; + use tracing::Instrument; use ulid::Ulid; const TEST_SOURCE_ID: &str = "01JGF7V3E0Y2R1X8P5Q7W9T4N7"; @@ -332,4 +333,193 @@ mod tests { let stats = queue.get_queue_stats().await.unwrap(); assert_eq!(stats.pending, 0); } + + // ----------------------------------------------------------------------- + // Trace context persistence + // ----------------------------------------------------------------------- + + /// Start a per-test-binary tracer provider + subscriber so we can create + /// active tracing spans whose context is injected by the real enqueue. + fn init_trace_subscriber() { + use std::sync::Once; + static INIT: Once = Once::new(); + INIT.call_once(|| { + use opentelemetry::global; + use opentelemetry::trace::TracerProvider; + use opentelemetry_sdk::propagation::TraceContextPropagator; + use opentelemetry_sdk::trace::{ + InMemorySpanExporter, RandomIdGenerator, SdkTracerProvider, SimpleSpanProcessor, + }; + use tracing_subscriber::layer::SubscriberExt as _; + use tracing_subscriber::util::SubscriberInitExt as _; + + global::set_text_map_propagator(TraceContextPropagator::new()); + + let exporter = InMemorySpanExporter::default(); + let provider = SdkTracerProvider::builder() + .with_span_processor(SimpleSpanProcessor::new(exporter)) + .with_id_generator(RandomIdGenerator::default()) + .build(); + global::set_tracer_provider(provider.clone()); + + let tracer = provider.tracer("test"); + let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); + let _ = tracing_subscriber::registry() + .with(telemetry_layer) + .try_init(); + }); + } + + #[tokio::test] + async fn test_enqueue_traceparent_persisted() { + init_trace_subscriber(); + + let env = TestEnvironment::new().await.unwrap(); + let pool = env.db_pool.pool().clone(); + let queue = EmbeddingQueue::new(pool.clone()); + insert_active_embedding_provider(&pool).await; + + let doc_id = create_document(&pool).await; + + let parent_span = tracing::info_span!("test_emb_enqueue_traceparent"); + async move { + let queue_id = queue.enqueue(doc_id.clone()).await.unwrap().unwrap(); + + // Query the stored row directly — traceparent must be non-null + // and valid. + let row: (Option, Option) = + sqlx::query_as("SELECT traceparent, tracestate FROM embedding_queue WHERE id = $1") + .bind(&queue_id) + .fetch_one(&pool) + .await + .unwrap(); + + let (stored_tp, stored_ts) = row; + assert!(stored_tp.is_some(), "traceparent must be stored"); + let tp = stored_tp.unwrap(); + assert!( + shared::telemetry::queue::is_valid_traceparent(&tp), + "stored traceparent must be valid W3C format: {}", + tp + ); + if let Some(ref ts) = stored_ts { + assert!(!ts.is_empty(), "tracestate must not be empty if present"); + } + + // Dequeue the item and verify it carries the same traceparent. + let batch = queue.dequeue_batch(10).await.unwrap(); + assert_eq!(batch.len(), 1); + assert_eq!(batch[0].traceparent, Some(tp.clone())); + assert_eq!(batch[0].tracestate, stored_ts); + } + .instrument(parent_span) + .await + } + + #[tokio::test] + async fn test_enqueue_batch_traceparent_persisted() { + init_trace_subscriber(); + + let env = TestEnvironment::new().await.unwrap(); + let pool = env.db_pool.pool().clone(); + let queue = EmbeddingQueue::new(pool.clone()); + insert_active_embedding_provider(&pool).await; + + let mut doc_ids = Vec::new(); + for _ in 0..3 { + doc_ids.push(create_document(&pool).await); + } + + let parent_span = tracing::info_span!("test_emb_enqueue_batch_traceparent"); + async move { + let queue_ids = queue.enqueue_batch(doc_ids.clone()).await.unwrap(); + assert_eq!(queue_ids.len(), 3, "all three docs should be enqueued"); + + for queue_id in &queue_ids { + let row: (Option, Option) = sqlx::query_as( + "SELECT traceparent, tracestate FROM embedding_queue WHERE id = $1", + ) + .bind(queue_id) + .fetch_one(&pool) + .await + .unwrap(); + + let (stored_tp, stored_ts) = row; + assert!( + stored_tp.is_some(), + "traceparent must be stored for queue_id={}", + queue_id + ); + let tp = stored_tp.unwrap(); + assert!( + shared::telemetry::queue::is_valid_traceparent(&tp), + "stored traceparent must be valid W3C format: {}", + tp + ); + if let Some(ref ts) = stored_ts { + assert!(!ts.is_empty(), "tracestate must not be empty if present"); + } + } + } + .instrument(parent_span) + .await + } + + #[tokio::test] + async fn test_enqueue_batch_missing_current_embeddings_traceparent_persisted() { + init_trace_subscriber(); + + let env = TestEnvironment::new().await.unwrap(); + let pool = env.db_pool.pool().clone(); + let queue = EmbeddingQueue::new(pool.clone()); + insert_active_embedding_provider(&pool).await; + + // Create docs without any embeddings (so they qualify as "missing"). + let mut doc_ids = Vec::new(); + for _ in 0..3 { + doc_ids.push(create_document(&pool).await); + } + + let parent_span = + tracing::info_span!("test_emb_enqueue_batch_missing_current_embeddings_traceparent"); + async move { + let queue_ids = queue + .enqueue_batch_missing_current_embeddings(doc_ids.clone()) + .await + .unwrap(); + assert_eq!( + queue_ids.len(), + 3, + "all three docs without embeddings should be enqueued" + ); + + for queue_id in &queue_ids { + let row: (Option, Option) = sqlx::query_as( + "SELECT traceparent, tracestate FROM embedding_queue WHERE id = $1", + ) + .bind(queue_id) + .fetch_one(&pool) + .await + .unwrap(); + + let (stored_tp, stored_ts) = row; + assert!( + stored_tp.is_some(), + "traceparent must be stored for queue_id={}", + queue_id + ); + let tp = stored_tp.unwrap(); + assert!( + shared::telemetry::queue::is_valid_traceparent(&tp), + "stored traceparent must be valid W3C format: {}", + tp + ); + if let Some(ref ts) = stored_ts { + assert!(!ts.is_empty(), "tracestate must not be empty if present"); + } + } + } + .instrument(parent_span) + .await + } } diff --git a/shared/tests/event_queue_test.rs b/shared/tests/event_queue_test.rs index 661e3dc3d..ee7504e21 100644 --- a/shared/tests/event_queue_test.rs +++ b/shared/tests/event_queue_test.rs @@ -5,6 +5,7 @@ mod tests { }; use shared::queue::EventQueue; use shared::test_environment::TestEnvironment; + use tracing::Instrument; const TEST_SOURCE_ID: &str = "01JGF7V3E0Y2R1X8P5Q7W9T4N7"; @@ -312,14 +313,16 @@ mod tests { .await .unwrap(); - let updated = queue + let (failed_count, dead_letter_count) = queue .mark_events_dead_letter_batch(vec![ (near_limit_id.clone(), "near limit".to_string()), (below_limit_id.clone(), "below limit".to_string()), ]) .await .unwrap(); - assert_eq!(updated, 2); + assert_eq!(failed_count, 1); + assert_eq!(dead_letter_count, 1); + assert_eq!((failed_count, dead_letter_count), (1, 1)); let near_limit_row: (String, i32) = sqlx::query_as("SELECT status, retry_count FROM connector_events_queue WHERE id = $1") @@ -610,4 +613,127 @@ mod tests { assert_eq!(inc_batch.len(), 1); assert_eq!(inc_batch[0].sync_run_id, inc_run); } + + // ----------------------------------------------------------------------- + // Trace context persistence + // ----------------------------------------------------------------------- + + /// Start a per-test-binary tracer provider + subscriber so we can create + /// active tracing spans whose context is injected by the real enqueue. + fn init_trace_subscriber() { + use std::sync::Once; + static INIT: Once = Once::new(); + INIT.call_once(|| { + use opentelemetry::global; + use opentelemetry::trace::TracerProvider; + use opentelemetry_sdk::propagation::TraceContextPropagator; + use opentelemetry_sdk::trace::InMemorySpanExporter; + use opentelemetry_sdk::trace::{ + RandomIdGenerator, SdkTracerProvider, SimpleSpanProcessor, + }; + use tracing_subscriber::layer::SubscriberExt as _; + use tracing_subscriber::util::SubscriberInitExt as _; + + global::set_text_map_propagator(TraceContextPropagator::new()); + + let exporter = InMemorySpanExporter::default(); + let provider = SdkTracerProvider::builder() + .with_span_processor(SimpleSpanProcessor::new(exporter)) + .with_id_generator(RandomIdGenerator::default()) + .build(); + global::set_tracer_provider(provider.clone()); + + let tracer = provider.tracer("test"); + let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); + let _ = tracing_subscriber::registry() + .with(telemetry_layer) + .try_init(); + }); + } + + #[tokio::test] + async fn test_enqueue_traceparent_persisted() { + init_trace_subscriber(); + + let env = TestEnvironment::new().await.unwrap(); + let pool = env.db_pool.pool().clone(); + let queue = EventQueue::new(pool.clone()); + + // Start an active parent span so enqueue injects trace context. + let parent_span = tracing::info_span!("test_enqueue_traceparent"); + async move { + let event = make_event("run-tp", "doc-tp"); + let event_id = queue.enqueue(TEST_SOURCE_ID, &event).await.unwrap(); + + // Query the stored row directly — traceparent must be non-null + // and valid. + let row: (Option, Option) = sqlx::query_as( + "SELECT traceparent, tracestate FROM connector_events_queue WHERE id = $1", + ) + .bind(&event_id) + .fetch_one(&pool) + .await + .unwrap(); + + let (stored_tp, stored_ts) = row; + assert!(stored_tp.is_some(), "traceparent must be stored"); + let tp = stored_tp.unwrap(); + assert!( + shared::telemetry::queue::is_valid_traceparent(&tp), + "stored traceparent must be valid W3C format: {}", + tp + ); + // tracestate is optional (may be None when no tracestate set). + if let Some(ref ts) = stored_ts { + assert!(!ts.is_empty(), "tracestate must not be empty if present"); + } + + // Dequeue the item and verify it carries the same traceparent. + let batch = queue.dequeue_batch(10).await.unwrap(); + assert_eq!(batch.len(), 1); + assert_eq!(batch[0].traceparent, Some(tp.clone())); + assert_eq!(batch[0].tracestate, stored_ts); + } + .instrument(parent_span) + .await + } + + #[tokio::test] + async fn test_enqueue_batch_traceparent_persisted() { + init_trace_subscriber(); + + let env = TestEnvironment::new().await.unwrap(); + let pool = env.db_pool.pool().clone(); + let queue = EventQueue::new(pool.clone()); + + // Start an active parent span so enqueue_batch injects trace context. + let parent_span = tracing::info_span!("test_enqueue_batch_traceparent"); + async move { + let run_id = ulid::Ulid::new().to_string(); + let events: Vec = (0..2) + .map(|i| make_event(&run_id, &format!("doc-batch-tp-{}", i))) + .collect(); + let ids = queue.enqueue_batch(TEST_SOURCE_ID, &events).await.unwrap(); + assert_eq!(ids.len(), 2); + + // Query each row and verify valid traceparent. + for id in &ids { + let tp: Option = sqlx::query_scalar( + "SELECT traceparent FROM connector_events_queue WHERE id = $1", + ) + .bind(id) + .fetch_one(&pool) + .await + .unwrap(); + + assert!(tp.is_some(), "traceparent must be stored for batch item"); + assert!( + shared::telemetry::queue::is_valid_traceparent(&tp.unwrap()), + "stored traceparent must be valid W3C format" + ); + } + } + .instrument(parent_span) + .await + } } diff --git a/shared/tests/metrics_test.rs b/shared/tests/metrics_test.rs new file mode 100644 index 000000000..8affdb1d9 --- /dev/null +++ b/shared/tests/metrics_test.rs @@ -0,0 +1,399 @@ +//! Integration tests for the shared metrics module. +//! +//! Sets up a single global meter provider backed by `InMemoryMetricExporter`, +//! then invokes the real `shared::metrics` production helpers. Because the +//! `LazyLock` instruments bind once to the global provider, this must run as +//! a single serialised test. + +use std::sync::OnceLock; + +use opentelemetry::global; +use opentelemetry_sdk::metrics::{ + data::{AggregatedMetrics, MetricData}, + InMemoryMetricExporter, SdkMeterProvider, +}; +use time; +use tokio::sync::Mutex; + +static METRICS_MUTEX: OnceLock> = OnceLock::new(); +fn metrics_mutex() -> &'static Mutex<()> { + METRICS_MUTEX.get_or_init(|| Mutex::new(())) +} + +/// Serialised test that sets up the global meter provider once, invokes every +/// production `shared::metrics` helper, then validates names, units, values, +/// bounded attributes, and seconds-scale durations. +/// +/// Because `LazyLock` instruments capture the global provider on first access, +/// all assertions run within a single serialised test. +#[tokio::test] +async fn test_production_metrics_names_and_values() { + let _lock = metrics_mutex().lock().await; + + // ------------------------------------------------------------------- + // 1. Build an in-memory exporter and install it as the global provider. + // This must happen before any shared::metrics static is accessed so + // the LazyLock instruments bind to this provider. + // ------------------------------------------------------------------- + let exporter = InMemoryMetricExporter::default(); + let meter_provider = SdkMeterProvider::builder() + .with_periodic_exporter(exporter.clone()) + .build(); + global::set_meter_provider(meter_provider.clone()); + + // ------------------------------------------------------------------- + // 2. Call every production recording helper. This triggers LazyLock + // evaluation against the provider above. + // ------------------------------------------------------------------- + + // HTTP RED (cross-cutting) + shared::metrics::record_http_request("GET", Some("/ok"), 200, 0.050); + shared::metrics::record_http_request("POST", Some("/search"), 201, 0.100); + shared::metrics::record_http_request("GET", None, 404, 0.010); + shared::metrics::record_http_request("PUT", Some("/documents/:id"), 400, 0.001); + shared::metrics::record_http_request("DELETE", Some("/sync/:id/cancel"), 500, 0.500); + + // Indexer + shared::metrics::INDEXER_EVENTS_PROCESSED.add(5, &[]); + shared::metrics::INDEXER_EVENTS_PROCESSED.add(3, &[]); + shared::metrics::INDEXER_EVENTS_DEAD_LETTER.add(2, &[]); + shared::metrics::INDEXER_EVENTS_RETRIED.add(1, &[]); + shared::metrics::INDEXER_BATCH_DURATION.record(1.5, &[]); + shared::metrics::INDEXER_BATCH_SIZE.record(10, &[]); + shared::metrics::record_queue_status(10, 2, 1, 0); + shared::metrics::record_embedding_queue_status(5, 1, 0); + + // Searcher + shared::metrics::SEARCHER_SEARCH_DURATION.record(0.250, &[]); + shared::metrics::SEARCHER_SEARCH_RESULTS.record(42, &[]); + shared::metrics::SEARCHER_SEARCH_ERRORS.add(1, &[]); + shared::metrics::SEARCHER_CACHE_HIT.add(3, &[]); + shared::metrics::SEARCHER_CACHE_MISS.add(7, &[]); + + // Connector sync + shared::metrics::CONNECTOR_SYNC_STARTED + .add(1, &[opentelemetry::KeyValue::new("sync_type", "full")]); + // Provide created_at so duration is recorded too. + let now = time::OffsetDateTime::now_utc(); + let five_secs_ago = now - time::Duration::seconds(5); + shared::metrics::record_sync_terminal("incremental", "completed", Some(five_secs_ago)); + shared::metrics::record_sync_terminal("full", "failed", None); + shared::metrics::record_sync_terminal("realtime", "cancelled", None); + + // Force flush so all metrics are collected by the exporter. + meter_provider.force_flush().unwrap(); + let resource_metrics = exporter + .get_finished_metrics() + .expect("metrics should be exported after flush"); + + assert!( + !resource_metrics.is_empty(), + "at least one ResourceMetrics batch expected" + ); + + // ------------------------------------------------------------------- + // 3. Collect all metric names for flexible assertion. + // ------------------------------------------------------------------- + let mut metric_names: Vec = Vec::new(); + for rm in &resource_metrics { + for sm in rm.scope_metrics() { + for m in sm.metrics() { + metric_names.push(m.name().to_string()); + } + } + } + + // ------------------------------------------------------------------- + // 4. Assert expected metric names exist. + // ------------------------------------------------------------------- + let expected_names = [ + // HTTP RED + "omni.http.server.request_count", + "omni.http.server.request_duration_seconds", + // Indexer + "omni.indexer.events.processed", + "omni.indexer.events.dead_letter", + "omni.indexer.events.retried", + "omni.indexer.batch.duration_seconds", + "omni.indexer.batch.size", + "omni.indexer.queue.depth", + // Searcher + "omni.searcher.search.duration_seconds", + "omni.searcher.search.results", + "omni.searcher.search.errors", + "omni.searcher.cache.hit", + "omni.searcher.cache.miss", + // Connector + "omni.connector.sync.started", + "omni.connector.sync.completed", + "omni.connector.sync.failed", + "omni.connector.sync.cancelled", + "omni.connector.sync.duration_seconds", + ]; + + for name in &expected_names { + assert!( + metric_names.contains(&name.to_string()), + "expected metric '{}' not found. Names: {:?}", + name, + metric_names + ); + } + + // ------------------------------------------------------------------- + // 5. Detailed assertions on specific metrics. + // ------------------------------------------------------------------- + for rm in &resource_metrics { + for sm in rm.scope_metrics() { + for m in sm.metrics() { + match m.name() { + // --- HTTP RED count --- + "omni.http.server.request_count" => { + assert_counter_total(m, 5, "http.request_count"); + // Each data point must have numeric status attribute. + assert_all_data_points_have_status(m); + // Each data point with a route must use template (no raw IDs, no '?'). + assert_no_raw_path_in_attributes(m); + } + // --- HTTP RED duration --- + "omni.http.server.request_duration_seconds" => { + assert_eq!(m.unit(), "s", "unit must be seconds"); + assert_duration_is_seconds(m, "http.duration_seconds"); + assert_all_data_points_have_status(m); + assert_no_raw_path_in_attributes(m); + } + // --- Indexer processed --- + "omni.indexer.events.processed" => { + assert_counter_total(m, 8, "events.processed"); + } + // --- Indexer dead_letter --- + "omni.indexer.events.dead_letter" => { + assert_counter_total(m, 2, "events.dead_letter"); + } + // --- Indexer retried --- + "omni.indexer.events.retried" => { + assert_counter_total(m, 1, "events.retried"); + } + // --- Indexer batch duration --- + "omni.indexer.batch.duration_seconds" => { + assert_eq!(m.unit(), "s", "unit must be seconds"); + assert_duration_is_seconds(m, "batch.duration"); + } + // --- Indexer batch size --- + "omni.indexer.batch.size" => { + assert_eq!(m.unit(), "{events}"); + } + // --- Indexer queue depth --- + "omni.indexer.queue.depth" => { + // Verify status attributes present. + assert_queue_depth_has_status_attributes(m); + } + // --- Searcher duration --- + "omni.searcher.search.duration_seconds" => { + assert_eq!(m.unit(), "s", "unit must be seconds"); + assert_duration_is_seconds(m, "search.duration"); + } + // --- Searcher results --- + "omni.searcher.search.results" => { + assert_eq!(m.unit(), "{results}"); + } + // --- Connector sync started --- + "omni.connector.sync.started" => { + assert_connector_has_sync_type_attr(m, "started"); + } + // --- Connector sync completed --- + "omni.connector.sync.completed" => { + assert_connector_has_outcome_attr(m, "completed", "completed"); + } + // --- Connector sync failed --- + "omni.connector.sync.failed" => { + assert_connector_has_outcome_attr(m, "failed", "failed"); + } + // --- Connector sync cancelled --- + "omni.connector.sync.cancelled" => { + assert_connector_has_outcome_attr(m, "cancelled", "cancelled"); + } + // --- Connector sync duration --- + "omni.connector.sync.duration_seconds" => { + assert_eq!(m.unit(), "s", "unit must be seconds"); + } + _ => {} + } + } + } + } +} + +// --------------------------------------------------------------------------- +// Assertion helpers +// --------------------------------------------------------------------------- + +fn assert_counter_total(m: &opentelemetry_sdk::metrics::data::Metric, expected: u64, label: &str) { + match m.data() { + AggregatedMetrics::U64(data) => { + if let MetricData::Sum(s) = data { + let total: u64 = s.data_points().map(|dp| dp.value()).sum(); + assert_eq!(total, expected, "{label}: expected {expected}, got {total}"); + } + } + _ => panic!("{label}: expected U64 counter"), + } +} + +fn assert_duration_is_seconds(m: &opentelemetry_sdk::metrics::data::Metric, label: &str) { + match m.data() { + AggregatedMetrics::F64(data) => { + if let MetricData::Histogram(h) = data { + for dp in h.data_points() { + if let Some(v) = dp.min() { + assert!( + v < 10.0, + "{label}: duration must be seconds, not ms (got {v})" + ); + } + } + } + } + _ => panic!("{label}: expected F64 histogram"), + } +} + +fn assert_all_data_points_have_status(m: &opentelemetry_sdk::metrics::data::Metric) { + match m.data() { + AggregatedMetrics::U64(data) => { + if let MetricData::Sum(s) = data { + for dp in s.data_points() { + assert!( + dp.attributes() + .any(|a| a.key.as_str() == "http.response.status_code"), + "http.response.status_code attribute required" + ); + // Status must be numeric (i64). + for a in dp.attributes() { + if a.key.as_str() == "http.response.status_code" { + let _: i64 = match &a.value { + opentelemetry::Value::I64(v) => *v, + _ => panic!("status_code must be I64, got {:?}", a.value), + }; + } + } + } + } + } + AggregatedMetrics::F64(data) => { + if let MetricData::Histogram(h) = data { + for dp in h.data_points() { + assert!( + dp.attributes() + .any(|a| a.key.as_str() == "http.response.status_code"), + "http.response.status_code attribute required" + ); + } + } + } + _ => {} + } +} + +fn assert_no_raw_path_in_attributes(m: &opentelemetry_sdk::metrics::data::Metric) { + let check = |a: &opentelemetry::KeyValue| { + let s = a.value.as_str().to_string(); + assert!(!s.contains('?'), "attributes must not contain '?': {s}"); + }; + match m.data() { + AggregatedMetrics::U64(data) => { + if let MetricData::Sum(s) = data { + for dp in s.data_points() { + for a in dp.attributes() { + check(&a); + } + } + } + } + AggregatedMetrics::F64(data) => { + if let MetricData::Histogram(h) = data { + for dp in h.data_points() { + for a in dp.attributes() { + check(&a); + } + } + } + } + _ => {} + } +} + +fn assert_queue_depth_has_status_attributes(m: &opentelemetry_sdk::metrics::data::Metric) { + match m.data() { + AggregatedMetrics::I64(data) => { + if let MetricData::Gauge(g) = data { + let statuses: Vec = g + .data_points() + .filter_map(|dp| { + dp.attributes() + .find(|a| a.key.as_str() == "status") + .map(|a| a.value.as_str().to_string()) + }) + .collect(); + // At minimum we should have pending and dead_letter statuses. + assert!( + statuses.contains(&"pending".to_string()), + "queue.depth must have 'pending' status, got {:?}", + statuses + ); + assert!( + statuses.contains(&"dead_letter".to_string()), + "queue.depth must have 'dead_letter' status, got {:?}", + statuses + ); + } + } + _ => panic!("queue.depth expected I64 gauge"), + } +} + +fn assert_connector_has_sync_type_attr(m: &opentelemetry_sdk::metrics::data::Metric, label: &str) { + match m.data() { + AggregatedMetrics::U64(data) => { + if let MetricData::Sum(s) = data { + for dp in s.data_points() { + let has_sync_type = dp.attributes().any(|a| a.key.as_str() == "sync_type"); + assert!( + has_sync_type, + "{label}: sync_type attribute required on started" + ); + // No IDs. + assert!( + !dp.attributes().any(|a| a.key.as_str() == "sync_run_id"), + "{label}: must not contain sync_run_id" + ); + assert!( + !dp.attributes().any(|a| a.key.as_str() == "source_id"), + "{label}: must not contain source_id" + ); + } + } + } + _ => panic!("{label}: expected U64 counter"), + } +} + +fn assert_connector_has_outcome_attr( + m: &opentelemetry_sdk::metrics::data::Metric, + label: &str, + expected_outcome: &str, +) { + match m.data() { + AggregatedMetrics::U64(data) => { + if let MetricData::Sum(s) = data { + for dp in s.data_points() { + let outcome = dp.attributes().find(|a| a.key.as_str() == "outcome"); + assert!(outcome.is_some(), "{label}: outcome attribute required"); + let val = outcome.unwrap().value.as_str().to_string(); + assert_eq!(val, expected_outcome, "{label}: outcome mismatch"); + } + } + } + _ => panic!("{label}: expected U64 counter"), + } +} diff --git a/shared/tests/queue_link_test.rs b/shared/tests/queue_link_test.rs new file mode 100644 index 000000000..437ea9476 --- /dev/null +++ b/shared/tests/queue_link_test.rs @@ -0,0 +1,357 @@ +//! Focused integration tests for native OTel link propagation through queues. +//! +//! Uses `InMemorySpanExporter` to inspect exported OpenTelemetry spans and +//! verify that CONSUMER spans carry native links (not string attributes) to +//! their PRODUCER spans. +//! +//! Tests use the *exact production* `build_producer_span` / `inject_active_context` +//! / `collect_span_contexts` helpers from the shared telemetry::queue module. + +use std::sync::{Mutex, OnceLock}; + +use opentelemetry::{ + global, + trace::{SpanKind, TracerProvider as _}, +}; +use opentelemetry_sdk::{ + propagation::TraceContextPropagator, + trace::{InMemorySpanExporter, RandomIdGenerator, SdkTracerProvider, SimpleSpanProcessor}, +}; +use shared::telemetry::queue; +use tracing::{info_span, Instrument}; +use tracing_opentelemetry::OpenTelemetrySpanExt; +use tracing_subscriber::{layer::SubscriberExt as _, util::SubscriberInitExt as _}; + +// --------------------------------------------------------------------------- +// Global test infrastructure +// --------------------------------------------------------------------------- + +/// Process-global mutex to serialise all tests in this file. +/// Every test acquires this guard before accessing the shared exporter. +static TEST_MUTEX: OnceLock> = OnceLock::new(); +fn test_guard() -> std::sync::MutexGuard<'static, ()> { + TEST_MUTEX + .get_or_init(|| Mutex::new(())) + .lock() + .unwrap_or_else(|p| p.into_inner()) +} + +static EXPORTER: OnceLock = OnceLock::new(); +fn exporter() -> &'static InMemorySpanExporter { + EXPORTER.get().expect("exporter not initialised") +} + +/// Initialise the global tracer provider + tracing subscriber once per suite. +fn ensure_global_init() { + EXPORTER.get_or_init(|| { + global::set_text_map_propagator(TraceContextPropagator::new()); + + let exporter = InMemorySpanExporter::default(); + let provider = SdkTracerProvider::builder() + .with_span_processor(SimpleSpanProcessor::new(exporter.clone())) + .with_id_generator(RandomIdGenerator::default()) + .build(); + global::set_tracer_provider(provider.clone()); + + let tracer = provider.tracer("test"); + let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); + let _ = tracing_subscriber::registry() + .with(telemetry_layer) + .try_init(); + + exporter + }); +} + +fn reset_exporter() { + exporter().reset(); +} + +fn finished_spans() -> Vec { + exporter() + .get_finished_spans() + .expect("get_finished_spans should succeed") +} + +/// Find a finished span by name and kind. +fn find_span<'a>( + spans: &'a [opentelemetry_sdk::trace::SpanData], + name: &str, + kind: SpanKind, +) -> Option<&'a opentelemetry_sdk::trace::SpanData> { + spans.iter().find(|s| s.name == name && s.span_kind == kind) +} + +/// Check whether the span has a `messaging.linked_trace_ids` attribute. +fn has_linked_trace_ids_attr(span: &opentelemetry_sdk::trace::SpanData) -> bool { + span.attributes + .iter() + .any(|kv| kv.key.as_str() == "messaging.linked_trace_ids") +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +/// Producer and consumer exported; consumer has no parent and a different +/// trace ID; native `links` contains exact producer trace+span ID. +#[test] +fn test_producer_consumer_native_link() { + let _guard = test_guard(); + ensure_global_init(); + reset_exporter(); + + let parent = info_span!("test_root"); + + parent.in_scope(|| { + // ---- Producer: create a PRODUCER span and inject context ---- + let carrier = { + let prod_span = queue::build_producer_span("test_queue"); + let _g = prod_span.entered(); + queue::inject_active_context() + // prod_span dropped here with EnteredSpan guard -> exported + }; + + // ---- Consumer: reconstruct context and build CONSUMER span with link ---- + let items = vec![Some((carrier.traceparent, carrier.tracestate))]; + let span_contexts = queue::collect_span_contexts(&items); + assert_eq!(span_contexts.len(), 1, "should have one producer context"); + + let cons_span = info_span!( + parent: None, + "test_queue process", + otel.kind = "CONSUMER", + ); + { + let _g = cons_span.clone().entered(); + for sc in &span_contexts { + cons_span.add_link(sc.clone()); + } + } + // Drop cons_span: consumer span ends and is exported. + drop(cons_span); + }); + + let spans = finished_spans(); + + // Find producer and consumer spans + let producer = find_span(&spans, "test_queue publish", SpanKind::Producer) + .expect("must have PRODUCER span"); + let consumer = find_span(&spans, "test_queue process", SpanKind::Consumer) + .expect("must have CONSUMER span"); + + // ---- Consumer has no parent (parent_span_id is all zeros) ---- + assert_eq!( + consumer.parent_span_id, + opentelemetry::trace::SpanId::INVALID, + "CONSUMER span must have INVALID parent_span_id (new root)" + ); + + // ---- Consumer has a different trace ID than producer ---- + assert_ne!( + consumer.span_context.trace_id(), + producer.span_context.trace_id(), + "CONSUMER trace ID must differ from PRODUCER trace ID (new root)" + ); + + // ---- Consumer has a native link to the producer ---- + assert!( + !consumer.links.is_empty(), + "CONSUMER span must have at least one link" + ); + let link = &consumer.links[0]; + assert_eq!( + link.span_context.trace_id(), + producer.span_context.trace_id(), + "link trace_id must match PRODUCER trace_id" + ); + assert_eq!( + link.span_context.span_id(), + producer.span_context.span_id(), + "link span_id must match PRODUCER span_id" + ); + + // ---- No messaging.linked_trace_ids attribute exists ---- + assert!( + !has_linked_trace_ids_attr(consumer), + "CONSUMER span must NOT have messaging.linked_trace_ids attribute" + ); + assert!( + !has_linked_trace_ids_attr(producer), + "PRODUCER span must NOT have messaging.linked_trace_ids attribute" + ); +} + +/// Two producer spans in the same trace yield two consumer links. +#[test] +fn test_two_producers_same_trace_two_links() { + let _guard = test_guard(); + ensure_global_init(); + reset_exporter(); + + let parent = info_span!("test_root"); + + parent.in_scope(|| { + // First producer + let c1 = { + let p1 = queue::build_producer_span("test_queue"); + let _g = p1.entered(); + queue::inject_active_context() + }; + + // Second producer in the same parent trace — same trace_id, different span_id + let c2 = { + let p2 = queue::build_producer_span("test_queue"); + let _g = p2.entered(); + queue::inject_active_context() + }; + + // Consumer: collect both contexts + let items = vec![ + Some((c1.traceparent, c1.tracestate)), + Some((c2.traceparent, c2.tracestate)), + ]; + let span_contexts = queue::collect_span_contexts(&items); + assert_eq!( + span_contexts.len(), + 2, + "both producer spans should be retained (different span_id)" + ); + + let cons_span = info_span!( + parent: None, + "test_queue process", + otel.kind = "CONSUMER", + ); + { + let _g = cons_span.clone().entered(); + for sc in &span_contexts { + cons_span.add_link(sc.clone()); + } + } + drop(cons_span); + }); + + let spans = finished_spans(); + + let consumer = find_span(&spans, "test_queue process", SpanKind::Consumer) + .expect("must have CONSUMER span"); + + assert_eq!( + consumer.links.len(), + 2, + "CONSUMER must have 2 links for 2 distinct producer spans" + ); + + // Verify both links have the same trace_id (from the parent) + let trace_id = consumer.links[0].span_context.trace_id(); + for link in consumer.links.iter() { + assert_eq!( + link.span_context.trace_id(), + trace_id, + "all links should share the same trace_id" + ); + } + + // Verify link span IDs are different + assert_ne!( + consumer.links[0].span_context.span_id(), + consumer.links[1].span_context.span_id(), + "two producer links must have different span IDs" + ); + + assert!(!has_linked_trace_ids_attr(consumer)); +} + +/// Duplicate same context dedupes to one link. +#[test] +fn test_duplicate_context_dedupes() { + let _guard = test_guard(); + ensure_global_init(); + reset_exporter(); + + let parent = info_span!("test_root"); + + parent.in_scope(|| { + let carrier = { + let prod_span = queue::build_producer_span("test_queue"); + let _g = prod_span.entered(); + queue::inject_active_context() + }; + + // Duplicate the same carrier + let items = vec![ + Some((carrier.traceparent.clone(), carrier.tracestate.clone())), + Some((carrier.traceparent, carrier.tracestate)), + ]; + let span_contexts = queue::collect_span_contexts(&items); + assert_eq!( + span_contexts.len(), + 1, + "duplicate (trace_id, span_id) must dedupe to 1" + ); + }); +} + +/// Invalid/missing items yield zero span contexts. +#[test] +fn test_invalid_missing_yield_zero() { + let _guard = test_guard(); + ensure_global_init(); + + let items: Vec, Option)>> = vec![ + None, + Some((None, None)), + Some((Some("invalid".to_string()), None)), + Some(( + Some("00-00000000000000000000000000000000-0000000000000000-01".to_string()), + None, + )), + ]; + let span_contexts = queue::collect_span_contexts(&items); + assert!( + span_contexts.is_empty(), + "invalid/missing items should yield zero span contexts" + ); +} + +/// Producer span is active during the async operation: verify by entering +/// a span, instrumenting a synchronous block, and checking exports. +#[tokio::test] +async fn test_producer_span_active_during_operation() { + let _guard = test_guard(); + ensure_global_init(); + reset_exporter(); + + let parent = info_span!("test_root"); + let _guard = parent.entered(); + + let prod_span = queue::build_producer_span("test_queue"); + + async move { + // Simulate an async operation (e.g. SQL INSERT) inside the producer span. + // The carrier is acquired inside the instrumented block so the + // producer span is active during injection and the async work. + let _carrier = queue::inject_active_context(); + tokio::time::sleep(std::time::Duration::from_millis(1)).await; + } + .instrument(prod_span) + .await; + + drop(_guard); + + let spans = finished_spans(); + + let producer = find_span(&spans, "test_queue publish", SpanKind::Producer) + .expect("must have PRODUCER span"); + + // Producer span should have messaging.* attributes + let has_destination = producer + .attributes + .iter() + .any(|kv| kv.key.as_str() == "messaging.destination"); + assert!( + has_destination, + "PRODUCER span must have messaging.destination" + ); +} diff --git a/shared/tests/telemetry_test.rs b/shared/tests/telemetry_test.rs new file mode 100644 index 000000000..6c1fb8db2 --- /dev/null +++ b/shared/tests/telemetry_test.rs @@ -0,0 +1,856 @@ +//! Integration tests for the shared telemetry module. +//! +//! Uses `InMemorySpanExporter` to inspect exported OpenTelemetry spans and +//! verify: +//! +//! - `send_traced` creates a CLIENT span, injects a valid W3C `traceparent` +//! header whose span ID matches the client span (not the parent), and +//! shares the parent trace ID. +//! - The Axum `trace_layer` middleware creates a single SERVER span with a +//! templated route name (no raw ID/query), records 200 and 5xx statuses, +//! and marks 5xx as ERROR. +//! +//! No EnteredGuard is held across any .await point. + +use std::net::SocketAddr; +use std::sync::OnceLock; + +use axum::{body::Body, extract::Request, middleware, response::Response, routing::get, Router}; +use opentelemetry::{ + global, + trace::{SpanId, SpanKind, TraceId, TracerProvider as _}, +}; +use opentelemetry_sdk::{ + propagation::TraceContextPropagator, + trace::{InMemorySpanExporter, RandomIdGenerator, SdkTracerProvider, SimpleSpanProcessor}, +}; +use reqwest::Client; +use serde_json::Value; +use tokio::sync::Mutex; +use tracing::{info_span, Instrument}; +use tracing_subscriber::{layer::SubscriberExt as _, util::SubscriberInitExt as _}; + +// --------------------------------------------------------------------------- +// Global test infrastructure +// --------------------------------------------------------------------------- + +/// Global in-memory span exporter shared across all tests. +static EXPORTER: OnceLock = OnceLock::new(); +fn exporter() -> &'static InMemorySpanExporter { + EXPORTER.get().expect("exporter not initialized") +} + +/// Suite mutex serialising access to the global InMemorySpanExporter. +/// Must be acquired (and held through the critical section) by each test +/// before calling `reset_exporter()` and released only after all span +/// assertions complete, to prevent reset/read races across parallel +/// `#[tokio::test]` execution. +static TEST_MUTEX: OnceLock> = OnceLock::new(); +fn test_mutex() -> &'static Mutex<()> { + TEST_MUTEX.get_or_init(|| Mutex::new(())) +} + +/// Initialise the global tracer provider + tracing subscriber once per suite. +fn ensure_global_init() { + EXPORTER.get_or_init(|| { + global::set_text_map_propagator(TraceContextPropagator::new()); + + let exporter = InMemorySpanExporter::default(); + let provider = SdkTracerProvider::builder() + .with_span_processor(SimpleSpanProcessor::new(exporter.clone())) + .with_id_generator(RandomIdGenerator::default()) + .build(); + global::set_tracer_provider(provider.clone()); + + let tracer = provider.tracer("test"); + let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); + let _ = tracing_subscriber::registry() + .with(telemetry_layer) + .try_init(); + + exporter + }); +} + +/// Reset the exporter. Tests should call this before their traced operations. +fn reset_exporter() { + exporter().reset(); +} + +/// Helper to get currently finished spans from the exporter. +/// With `SimpleSpanProcessor`, spans are exported synchronously on end, +/// so no explicit force-flush is needed — the data is already present. +fn finished_spans() -> Vec { + exporter() + .get_finished_spans() + .expect("get_finished_spans should succeed") +} + +/// Extract a reference to the `Value` from a `KeyValue` by key name. +fn attr_value<'a>( + attrs: &'a [opentelemetry::KeyValue], + key: &str, +) -> Option<&'a opentelemetry::Value> { + attrs + .iter() + .find(|kv| kv.key.as_str() == key) + .map(|kv| &kv.value) +} + +/// Match an OTel `Value` to a string reference. +fn value_as_str(v: &opentelemetry::Value) -> Option<&str> { + match v { + opentelemetry::Value::String(s) => Some(s.as_ref()), + _ => None, + } +} + +/// Match an OTel `Value` to an i64. +fn value_as_i64(v: &opentelemetry::Value) -> Option { + match v { + opentelemetry::Value::I64(i) => Some(*i), + _ => None, + } +} + +/// Check whether a `Status` is error. +fn status_is_error(status: &opentelemetry::trace::Status) -> bool { + matches!(status, opentelemetry::trace::Status::Error { .. }) +} + +/// Convert a `SpanId` to its hex string representation. +fn span_id_hex(id: opentelemetry::trace::SpanId) -> String { + format!("{:016x}", id) +} + +/// Convert a `TraceId` to its hex string representation. +fn trace_id_hex(id: opentelemetry::trace::TraceId) -> String { + format!("{:032x}", id) +} + +// --------------------------------------------------------------------------- +// HTTP helpers +// --------------------------------------------------------------------------- + +/// Start a simple downstream server that echoes request headers as JSON. +/// Returns the bound address and a oneshot sender to shut it down. +async fn start_downstream() -> (SocketAddr, tokio::sync::oneshot::Sender<()>) { + let (tx, rx) = tokio::sync::oneshot::channel::<()>(); + + async fn echo_headers(req: Request) -> Response { + let headers: Vec<(String, String)> = req + .headers() + .iter() + .map(|(k, v)| (k.to_string(), v.to_str().unwrap_or("").to_string())) + .collect(); + let body = serde_json::to_string(&serde_json::json!({ "headers": headers })).unwrap(); + Response::new(Body::from(body)) + } + + let app = Router::new().route("/echo", get(echo_headers).post(echo_headers)); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let wait_for_start = { + let addr = addr; + async move { + let client = reqwest::Client::new(); + loop { + if client + .get(format!("http://{}/echo", addr)) + .send() + .await + .is_ok() + { + break; + } + tokio::time::sleep(tokio::time::Duration::from_millis(20)).await; + } + } + }; + + tokio::spawn(async move { + let serve = axum::serve(listener, app).with_graceful_shutdown(async { + rx.await.ok(); + }); + let _ = serve.await; + }); + + // Wait for the server to be ready + tokio::time::timeout(tokio::time::Duration::from_secs(5), wait_for_start) + .await + .expect("downstream server did not start in time"); + + (addr, tx) +} + +/// Parse a `traceparent` value into its components: version, trace_id, span_id, flags. +fn parse_traceparent(tp: &str) -> Option<(&str, &str, &str, &str)> { + let parts: Vec<&str> = tp.split('-').collect(); + if parts.len() == 4 { + Some((parts[0], parts[1], parts[2], parts[3])) + } else { + None + } +} + +/// Extract the `traceparent` header value from the JSON-echoed headers array. +fn get_traceparent_from_json(headers: &[Value]) -> Option { + headers + .iter() + .find(|h| h[0].as_str().unwrap_or("").to_lowercase() == "traceparent") + .map(|h| h[1].as_str().unwrap().to_string()) +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +/// Verify that `send_traced` creates a CLIENT span, injects a valid +/// `traceparent` header whose span ID matches the client span, and that +/// the client span is a child of any active parent. +/// +/// The parent scope is established via `.instrument()` so no EnteredGuard +/// is held across any `.await`. +#[tokio::test] +async fn test_send_traced_client_span_and_injection() { + ensure_global_init(); + let (addr, _stop) = start_downstream().await; + let url = format!("http://{}/echo", addr); + + let _lock = test_mutex().lock().await; + reset_exporter(); + + let parent = info_span!("test_parent", otel.kind = "INTERNAL"); + + // ---- within parent span ---- + async { + let client = Client::new(); + let resp = shared::telemetry::http_client::send_traced("GET", &url, client.get(&url)) + .await + .expect("send_traced should succeed"); + assert!(resp.status().is_success(), "downstream should return 200"); + + let body: Value = resp.json().await.unwrap(); + let headers = body["headers"].as_array().unwrap(); + + // --- 1. Validate the injected traceparent header --- + let tp_val = get_traceparent_from_json(headers) + .expect("downstream request must contain a traceparent header"); + let (version, tp_trace_id, tp_span_id, tp_flags) = + parse_traceparent(&tp_val).expect("traceparent must be valid W3C format"); + assert_eq!(version, "00", "traceparent version must be 00"); + assert_eq!(tp_flags, "01", "traceparent flags must be 01 (sampled)"); + assert_eq!(tp_trace_id.len(), 32, "trace_id must be 32 hex chars"); + assert_eq!(tp_span_id.len(), 16, "span_id must be 16 hex chars"); + assert!( + tp_trace_id.chars().all(|c| c.is_ascii_hexdigit()), + "trace_id must be hex" + ); + assert!( + tp_span_id.chars().all(|c| c.is_ascii_hexdigit()), + "span_id must be hex" + ); + + // --- 2. Get finished spans and locate the CLIENT span --- + // (Parent span is still active inside this instrumented block. + // CLIENT span should have been exported synchronously on end.) + let spans = finished_spans(); + + // Find CLIENT spans whose trace_id matches the injected header. + // (Due to concurrent test execution, other CLIENT spans may be present; + // we filter by the trace_id that we know was injected.) + let client_spans: Vec<_> = spans + .iter() + .filter(|s| { + s.span_kind == SpanKind::Client + && trace_id_hex(s.span_context.trace_id()) == tp_trace_id + }) + .collect(); + assert!( + !client_spans.is_empty(), + "send_traced must create at least one CLIENT span matching the injected trace_id" + ); + let client_span = &client_spans[0]; + + assert_eq!( + client_span.name, "HTTP GET", + "CLIENT span name should be 'HTTP GET'" + ); + + // --- 3. Injected traceparent uses client span ID --- + let client_span_id_hex = span_id_hex(client_span.span_context.span_id()); + assert_eq!( + tp_span_id, client_span_id_hex, + "injected traceparent span_id must match the CLIENT span's span_id" + ); + + // --- 4. CLIENT span records http.request.method --- + let method_val = attr_value(&client_span.attributes, "http.request.method"); + assert_eq!( + method_val.and_then(value_as_str), + Some("GET"), + "CLIENT span must record http.request.method=GET" + ); + } + .instrument(parent) + .await; +} + +/// Verify the Axum `trace_layer` middleware creates a single SERVER span with +/// templated route, records 200 status, and does not create raw-path spans. +#[tokio::test] +async fn test_middleware_server_span_on_success() { + ensure_global_init(); + + use shared::telemetry::middleware::trace_layer; + + let app = Router::new() + .route("/ok", get(|| async { "ok" })) + .layer(middleware::from_fn(trace_layer)); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + // Allow server to start + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + + let _lock = test_mutex().lock().await; + reset_exporter(); + + let client = Client::new(); + let resp = client + .get(format!("http://{}/ok", addr)) + .send() + .await + .expect("request to test server must succeed"); + assert_eq!(resp.status().as_u16(), 200); + + // Give the server's async task time to finish the span export + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + let spans = finished_spans(); + + // Find the SERVER span for this specific route + let server_span = spans + .iter() + .find(|s| s.span_kind == SpanKind::Server && s.name == "GET /ok") + .expect("must have a SERVER span named 'GET /ok'"); + + // Route attribute should be the template + let route_attr = attr_value(&server_span.attributes, "http.route"); + assert_eq!( + route_attr.and_then(value_as_str), + Some("/ok"), + "http.route attribute must be the route template" + ); + + // Status code attribute + eprintln!("SERVER span attributes:"); + for a in &server_span.attributes { + eprintln!(" {:?} = {:?}", a.key, a.value); + } + let status_attr = attr_value(&server_span.attributes, "http.response.status_code"); + assert_eq!( + status_attr.and_then(value_as_i64), + Some(200), + "http.response.status_code must be 200" + ); + + // No error status on 200 + assert!( + !status_is_error(&server_span.status), + "SERVER span for 200 must not have error status" + ); +} + +/// Verify that the middleware correctly extracts an incoming W3C `traceparent`, +/// propagating the trace ID and setting the parent span ID from the remote +/// caller, while still recording route/status/error assertions. +#[tokio::test] +async fn test_middleware_server_span_with_incoming_traceparent() { + ensure_global_init(); + + use shared::telemetry::middleware::trace_layer; + + let app = Router::new() + .route("/ok", get(|| async { "ok" })) + .layer(middleware::from_fn(trace_layer)); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + + let _lock = test_mutex().lock().await; + reset_exporter(); + + // Known W3C traceparent (version, trace_id, parent_span_id, flags) + let trace_id_hex = "0af7651916cd43dd8448eb211c80319c"; + let parent_span_id_hex = "b7ad6b7169203331"; + let traceparent = format!("00-{}-{}-01", trace_id_hex, parent_span_id_hex); + + let client = Client::new(); + let resp = client + .get(format!("http://{}/ok", addr)) + .header("traceparent", &traceparent) + .send() + .await + .expect("request must succeed"); + assert_eq!(resp.status().as_u16(), 200); + + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + let spans = finished_spans(); + + let server_span = spans + .iter() + .find(|s| s.span_kind == SpanKind::Server && s.name == "GET /ok") + .expect("must have a SERVER span named 'GET /ok'"); + + let expected_trace_id = TraceId::from_hex(trace_id_hex).expect("valid trace_id hex"); + let expected_parent_span_id = SpanId::from_hex(parent_span_id_hex).expect("valid span_id hex"); + + assert_eq!( + server_span.span_context.trace_id(), + expected_trace_id, + "SERVER span must inherit the incoming trace ID" + ); + assert_eq!( + server_span.parent_span_id, expected_parent_span_id, + "SERVER span parent_span_id must equal incoming traceparent span_id" + ); + + // Retain existing route / status / error assertions + let route_attr = attr_value(&server_span.attributes, "http.route"); + assert_eq!( + route_attr.and_then(value_as_str), + Some("/ok"), + "http.route must be /ok" + ); + + let status_attr = attr_value(&server_span.attributes, "http.response.status_code"); + assert_eq!( + status_attr.and_then(value_as_i64), + Some(200), + "http.response.status_code must be 200" + ); + + assert!( + !status_is_error(&server_span.status), + "SERVER span for 200 must not have error status" + ); +} + +/// Verify the middleware records ERROR status for 5xx responses. +#[tokio::test] +async fn test_middleware_server_span_on_error() { + ensure_global_init(); + + use shared::telemetry::middleware::trace_layer; + + let app = Router::new() + .route( + "/error", + get(|| async { (axum::http::StatusCode::INTERNAL_SERVER_ERROR, "error") }), + ) + .layer(middleware::from_fn(trace_layer)); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + + let _lock = test_mutex().lock().await; + reset_exporter(); + + let client = Client::new(); + let resp = client + .get(format!("http://{}/error", addr)) + .send() + .await + .expect("request to test server must succeed"); + assert_eq!(resp.status().as_u16(), 500); + + // Give the server's async task time to finish the span export + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + let spans = finished_spans(); + + // Find the SERVER span for this specific route + let server_span = spans + .iter() + .find(|s| s.span_kind == SpanKind::Server && s.name == "GET /error") + .expect("must have a SERVER span named 'GET /error'"); + + // Status code attribute + eprintln!("ERROR span attributes:"); + for a in &server_span.attributes { + eprintln!(" {:?} = {:?}", a.key, a.value); + } + let status_attr = attr_value(&server_span.attributes, "http.response.status_code"); + assert_eq!( + status_attr.and_then(value_as_i64), + Some(500), + "http.response.status_code must be 500" + ); + + // Span should have error status + assert!( + status_is_error(&server_span.status), + "SERVER span for 5xx must have error status" + ); + + // Route should still be templated + assert_eq!( + server_span.name, "GET /error", + "SERVER span name must be templated 'GET /error'" + ); +} + +// --------------------------------------------------------------------------- +// OTel log helpers — pure-function tests +// --------------------------------------------------------------------------- + +#[test] +fn test_normalize_otlp_endpoint_strips_trailing_slash() { + assert_eq!( + shared::telemetry::normalize_otlp_endpoint("http://localhost:4318"), + "http://localhost:4318" + ); + assert_eq!( + shared::telemetry::normalize_otlp_endpoint("http://localhost:4318/"), + "http://localhost:4318" + ); + assert_eq!( + shared::telemetry::normalize_otlp_endpoint("http://localhost:4318///"), + "http://localhost:4318" + ); + assert_eq!(shared::telemetry::normalize_otlp_endpoint(""), ""); +} + +#[test] +fn test_build_logs_url() { + assert_eq!( + shared::telemetry::build_logs_url(Some("http://localhost:4318")), + Some("http://localhost:4318/v1/logs".to_string()) + ); + assert_eq!( + shared::telemetry::build_logs_url(Some("http://localhost:4318/")), + Some("http://localhost:4318/v1/logs".to_string()) + ); + assert_eq!(shared::telemetry::build_logs_url(None), None); +} + +#[test] +fn test_build_logs_url_v1_traces_is_not_affected() { + // Ensure /v1/logs endpoint is independent from the traces endpoint + let url = shared::telemetry::build_logs_url(Some("http://otel:4318")).unwrap(); + assert_eq!(url, "http://otel:4318/v1/logs"); + assert_ne!(url, "http://otel:4318/v1/traces"); +} + +/// Verify that the middleware creates a SERVER span using the incoming W3C +/// trace context, uses the templated route name (no raw ID/query string), +/// and records the correct status. +#[tokio::test] +async fn test_middleware_traceparent_propagation() { + ensure_global_init(); + + use axum::extract::Path; + use shared::telemetry::middleware::trace_layer; + + let app = Router::new() + .route("/users/:id/details", get(|_: Path| async { "ok" })) + .layer(middleware::from_fn(trace_layer)); + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + tokio::spawn(async move { + let _ = axum::serve(listener, app).await; + }); + + // Allow server to start + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + + let _lock = test_mutex().lock().await; + reset_exporter(); + + // Known W3C traceparent: version=00, trace_id=0af7…, span_id=b7ad…, flags=01 + let traceparent = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; + let expected_trace_id = "0af7651916cd43dd8448eb211c80319c"; + let expected_parent_span_id = "b7ad6b7169203331"; + + let client = Client::new(); + let resp = client + .get(format!("http://{}/users/abc123/details?foo=bar", addr)) + .header("traceparent", traceparent) + .send() + .await + .expect("request to test server must succeed"); + assert_eq!(resp.status().as_u16(), 200); + + // Give the server's async task time to finish the span export + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + let spans = finished_spans(); + + // Find the SERVER span for this parameterised route + let server_span = spans + .iter() + .find(|s| s.span_kind == SpanKind::Server && s.name == "GET /users/:id/details") + .expect("must have a SERVER span named 'GET /users/:id/details'"); + + // === Incoming trace context is used === + let trace_id = trace_id_hex(server_span.span_context.trace_id()); + assert_eq!( + trace_id, expected_trace_id, + "SERVER span must use the incoming trace ID from traceparent" + ); + + let parent_span_id = span_id_hex(server_span.parent_span_id); + assert_eq!( + parent_span_id, expected_parent_span_id, + "SERVER span's parent span ID must match the incoming traceparent span ID" + ); + + // === No raw ID or query in span name === + assert!( + !server_span.name.contains("abc123"), + "SERVER span name must not contain raw path parameter" + ); + assert!( + !server_span.name.contains("foo=bar"), + "SERVER span name must not contain query string" + ); + assert!( + !server_span.name.contains("?"), + "SERVER span name must not contain '?'" + ); + + // === Route attribute is the template === + let route_attr = attr_value(&server_span.attributes, "http.route"); + assert_eq!( + route_attr.and_then(value_as_str), + Some("/users/:id/details"), + "http.route attribute must be the route template" + ); + + // === Status code === + let status_attr = attr_value(&server_span.attributes, "http.response.status_code"); + assert_eq!( + status_attr.and_then(value_as_i64), + Some(200), + "http.response.status_code must be 200" + ); + + // No error on 200 + assert!( + !status_is_error(&server_span.status), + "SERVER span for 200 must not have error status" + ); +} + +/// Verify that consecutive `send_traced` calls produce distinct client span IDs +/// but share the same trace ID within the same parent context. +/// +/// The parent scope is established via `.instrument()` so no EnteredGuard +/// is held across any `.await`. +#[tokio::test] +async fn test_send_traced_distinct_client_span_ids() { + ensure_global_init(); + let (addr, _stop) = start_downstream().await; + let url = format!("http://{}/echo", addr); + + let _lock = test_mutex().lock().await; + reset_exporter(); + + let parent = info_span!("test_parent", otel.kind = "INTERNAL"); + + // ---- within parent span ---- + async { + let client = Client::new(); + + let resp1 = shared::telemetry::http_client::send_traced("GET", &url, client.get(&url)) + .await + .expect("first send_traced should succeed"); + let body1: Value = resp1.json().await.unwrap(); + let headers1 = body1["headers"].as_array().unwrap(); + let tp1 = get_traceparent_from_json(headers1).unwrap(); + + let resp2 = shared::telemetry::http_client::send_traced("GET", &url, client.get(&url)) + .await + .expect("second send_traced should succeed"); + let body2: Value = resp2.json().await.unwrap(); + let headers2 = body2["headers"].as_array().unwrap(); + let tp2 = get_traceparent_from_json(headers2).unwrap(); + + let (_, trace1, span1, _) = parse_traceparent(&tp1).unwrap(); + let (_, trace2, span2, _) = parse_traceparent(&tp2).unwrap(); + + // Different calls should produce different span IDs + assert_ne!( + span1, span2, + "consecutive send_traced calls should produce different span IDs" + ); + + // Both must have valid hex IDs + assert_eq!(trace1.len(), 32); + assert_eq!(trace2.len(), 32); + + // Both should share the same trace ID (from the parent span) + assert_eq!( + trace1, trace2, + "two calls within the same parent should share trace ID" + ); + + // Validate via exported spans too. + let spans = finished_spans(); + + let client_spans: Vec<_> = spans + .iter() + .filter(|s| s.span_kind == SpanKind::Client) + .collect(); + assert!( + client_spans.len() >= 2, + "should have at least two CLIENT spans" + ); + + // All client spans should have unique span IDs + let client_span_ids: Vec<_> = client_spans + .iter() + .map(|s| span_id_hex(s.span_context.span_id())) + .collect(); + let mut unique_ids = client_span_ids.clone(); + unique_ids.sort(); + unique_ids.dedup(); + assert_eq!( + unique_ids.len(), + client_span_ids.len(), + "each CLIENT span should have a unique span ID" + ); + } + .instrument(parent) + .await; +} + +// --------------------------------------------------------------------------- +// OTel log exporter integration tests +// --------------------------------------------------------------------------- + +/// Verify that the logging bridge emits log records with matching +/// trace_id/span_id when a tracing span with OTel context is active. +#[cfg(test)] +#[tokio::test] +async fn test_log_bridge_trace_context_in_log_records() { + use opentelemetry::trace::TracerProvider; + use opentelemetry_sdk::logs::{InMemoryLogExporter, SdkLoggerProvider, SimpleLogProcessor}; + use tracing::Instrument; + + let exporter = InMemoryLogExporter::default(); + let logger_provider = SdkLoggerProvider::builder() + .with_log_processor(SimpleLogProcessor::new(exporter.clone())) + .build(); + + let tracer_provider = opentelemetry_sdk::trace::SdkTracerProvider::builder() + .with_id_generator(RandomIdGenerator::default()) + .build(); + + let tracer = tracer_provider.tracer("test"); + let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); + let otel_log_layer = + opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(&logger_provider); + + let _guard = tracing::subscriber::set_default( + tracing_subscriber::registry() + .with(telemetry_layer) + .with(otel_log_layer), + ); + + exporter.reset(); + + let parent_span = info_span!("test_parent", otel.kind = "INTERNAL"); + + async { + tracing::info!("hello from inside a span"); + } + .instrument(parent_span) + .await; + + logger_provider.force_flush().ok(); + + let log_records = exporter.get_emitted_logs().unwrap_or_default(); + assert!( + !log_records.is_empty(), + "expected at least one log record from inside the span" + ); + + let log_data = &log_records[0]; + let trace_context = log_data.record.trace_context(); + assert!( + trace_context.is_some(), + "log record should have trace context when inside a span" + ); + let tc = trace_context.unwrap(); + assert!( + tc.trace_id.to_string() != "00000000000000000000000000000000", + "log record trace_id should be non-zero when inside a span" + ); + assert!( + tc.span_id.to_string() != "0000000000000000", + "log record span_id should be non-zero when inside a span" + ); +} + +/// Verify that the appender bridge does NOT emit log records outside +/// any tracing scope (i.e. when no span context exists, the log record +/// should still have zero-valued trace_id/span_id). +#[cfg(test)] +#[tokio::test] +async fn test_log_bridge_no_span_safety() { + use opentelemetry_sdk::logs::{InMemoryLogExporter, SdkLoggerProvider, SimpleLogProcessor}; + + let exporter = InMemoryLogExporter::default(); + let logger_provider = SdkLoggerProvider::builder() + .with_log_processor(SimpleLogProcessor::new(exporter.clone())) + .build(); + + let otel_log_layer = + opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(&logger_provider); + + let _guard = + tracing::subscriber::set_default(tracing_subscriber::registry().with(otel_log_layer)); + + exporter.reset(); + + tracing::info!("log outside any span"); + + logger_provider.force_flush().ok(); + + let log_records = exporter.get_emitted_logs().unwrap_or_default(); + assert!( + !log_records.is_empty(), + "expected at least one log record from outside a span" + ); + + let log_data = &log_records[0]; + let trace_context = log_data.record.trace_context(); + assert!( + trace_context.is_none(), + "log record should NOT have trace context when outside a span" + ); +} diff --git a/web/Dockerfile b/web/Dockerfile index cbe93cf40..219e6c96d 100644 --- a/web/Dockerfile +++ b/web/Dockerfile @@ -23,6 +23,10 @@ WORKDIR /app # Port argument with default value ARG WEB_PORT=3000 +# Detect production adapter-node mode for shutdown sequencing +# (dev docker-compose overrides this to "development") +ENV NODE_ENV=production + # Copy package files for production dependencies COPY package*.json ./ @@ -33,5 +37,17 @@ RUN npm ci --only=production && npm cache clean --force COPY --from=builder /app/build ./build COPY --from=builder /app/static ./static -# Start the application -CMD ["node", "build"] +# Copy the preload dependency set so `node --import ./instrumentation.mjs` resolves +# at startup. These are plain .mjs files loaded by Node directly (not bundled). +COPY src/instrumentation.mjs ./ +COPY src/otel-factory.mjs ./ +COPY src/lib/server/otel-utils.mjs ./lib/server/ + +# Build-stage preload import check: verify the .mjs files exist and parse correctly. +RUN node --check ./instrumentation.mjs && \ + node --check ./otel-factory.mjs && \ + node --check ./lib/server/otel-utils.mjs && \ + echo "Preload dependency syntax check passed." + +# Start the application with telemetry preloaded +CMD ["node", "--import", "./instrumentation.mjs", "build"] diff --git a/web/package-lock.json b/web/package-lock.json index 9efd2d344..a95af64ac 100644 --- a/web/package-lock.json +++ b/web/package-lock.json @@ -13,8 +13,13 @@ "@fontsource-variable/inter": "^5.2.8", "@node-rs/argon2": "^2.0.2", "@opentelemetry/api": "^1.9.1", + "@opentelemetry/api-logs": "^0.220.0", "@opentelemetry/auto-instrumentations-node": "^0.78.0", + "@opentelemetry/exporter-logs-otlp-http": "^0.220.0", + "@opentelemetry/exporter-metrics-otlp-http": "^0.220.0", "@opentelemetry/exporter-trace-otlp-http": "^0.220.0", + "@opentelemetry/sdk-logs": "^0.220.0", + "@opentelemetry/sdk-metrics": "^2.9.0", "@opentelemetry/sdk-node": "^0.220.0", "@oslojs/crypto": "^1.0.1", "@oslojs/encoding": "^1.1.0", diff --git a/web/package.json b/web/package.json index 51f5e4a89..043d8cb06 100644 --- a/web/package.json +++ b/web/package.json @@ -4,9 +4,9 @@ "version": "0.0.1", "type": "module", "scripts": { - "dev": "vite dev", + "dev": "NODE_OPTIONS='--import ./src/instrumentation.mjs' vite dev", "build": "vite build", - "preview": "vite preview", + "preview": "NODE_OPTIONS='--import ./src/instrumentation.mjs' vite preview", "prepare": "svelte-kit sync || echo ''", "check": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json", "check:watch": "svelte-kit sync && svelte-check --tsconfig ./tsconfig.json --watch", @@ -65,8 +65,13 @@ "@fontsource-variable/inter": "^5.2.8", "@node-rs/argon2": "^2.0.2", "@opentelemetry/api": "^1.9.1", + "@opentelemetry/api-logs": "^0.220.0", "@opentelemetry/auto-instrumentations-node": "^0.78.0", + "@opentelemetry/exporter-logs-otlp-http": "^0.220.0", + "@opentelemetry/exporter-metrics-otlp-http": "^0.220.0", "@opentelemetry/exporter-trace-otlp-http": "^0.220.0", + "@opentelemetry/sdk-logs": "^0.220.0", + "@opentelemetry/sdk-metrics": "^2.9.0", "@opentelemetry/sdk-node": "^0.220.0", "@oslojs/crypto": "^1.0.1", "@oslojs/encoding": "^1.1.0", diff --git a/web/src/hooks.server.ts b/web/src/hooks.server.ts index f6023ca6e..9ab381981 100644 --- a/web/src/hooks.server.ts +++ b/web/src/hooks.server.ts @@ -1,14 +1,12 @@ import type { Handle, HandleServerError } from '@sveltejs/kit' import { sequence } from '@sveltejs/kit/hooks' import { redirect } from '@sveltejs/kit' + import * as auth from '$lib/server/auth.js' import { validateApiKey } from '$lib/server/apiKeys.js' import { rateLimit } from '$lib/server/rateLimit.js' import { Logger } from '$lib/server/logger.js' -import { initTelemetry, extractTraceContext, getRequestId } from '$lib/server/telemetry.js' - -// Initialize OpenTelemetry on module load -initTelemetry() +import { getRequestId, recordHttpRequest, millisecondsToSeconds } from '$lib/server/telemetry.js' const handleAuth: Handle = async ({ event, resolve }) => { // 1. Try API key auth (Authorization: Bearer omni_* or X-API-Key header) @@ -89,44 +87,63 @@ const handlePasswordChange: Handle = async ({ event, resolve }) => { } const handleLogging: Handle = async ({ event, resolve }) => { - // Extract trace context from incoming request headers - const headers: Record = {} - event.request.headers.forEach((value, key) => { - headers[key] = value - }) - extractTraceContext(headers) - - // Use trace ID as request ID if available, otherwise generate new one + // Use trace ID as request ID if available, otherwise generate new one. + // The Node auto-instrumentation already creates a server span and extracts + // the incoming traceparent; we just need its trace ID for logging. const requestId = getRequestId() || Logger.generateRequestId() - const logger = new Logger('request').withRequest(requestId, event.locals.user?.id) + const logger = new Logger('request').withRequest(requestId) event.locals.requestId = requestId event.locals.logger = logger - const startTime = Date.now() + const startTime = performance.now() + const route = event.route.id ?? '/unknown' logger.info('Request started', { method: event.request.method, - url: event.url.pathname + event.url.search, - userAgent: event.request.headers.get('user-agent'), - ip: event.getClientAddress(), - userId: event.locals.user?.id, - userEmail: event.locals.user?.email, + route, }) - const response = await resolve(event) - - const duration = Date.now() - startTime + let responseStatus: number = 500 + let error: unknown = null + let response: Response | undefined + + try { + response = await resolve(event) + responseStatus = response.status + } catch (thrown: unknown) { + // SvelteKit throws redirect(...) and error(...) which are Response-like. + // Capture the status when available; default to 500 for true errors. + if (thrown instanceof Response) { + responseStatus = thrown.status + } else if (thrown && typeof thrown === 'object' && 'status' in (thrown as object)) { + responseStatus = (thrown as { status: number }).status + } else { + responseStatus = 500 + } + error = thrown + // Rethrow after recording so the framework's error handler still fires. + throw thrown + } finally { + const durationMs = performance.now() - startTime + const durationSecs = millisecondsToSeconds(durationMs) + + if (error === null && response) { + logger.info('Request completed', { + method: event.request.method, + route, + status: responseStatus, + duration: durationMs, + }) + } - logger.info('Request completed', { - method: event.request.method, - url: event.url.pathname + event.url.search, - status: response.status, - duration, - userId: event.locals.user?.id, - }) + // Record HTTP RED metric with bounded attributes + // Uses route.id (template) not raw path; never records user/query IDs. + // Duration is in seconds (OTel standard for histograms). + recordHttpRequest(event.request.method, route, responseStatus, durationSecs) + } - return response + return response! } export const handle = sequence(handleLogging, handleAuth, handlePasswordChange) @@ -135,9 +152,7 @@ export const handleError: HandleServerError = ({ error, event }) => { const logger = event.locals.logger || new Logger('error') logger.error('Unhandled server error', error as Error, { - url: event.url.pathname + event.url.search, method: event.request.method, - userId: event.locals.user?.id, requestId: event.locals.requestId, }) diff --git a/web/src/instrumentation.mjs b/web/src/instrumentation.mjs new file mode 100644 index 000000000..f17dba649 --- /dev/null +++ b/web/src/instrumentation.mjs @@ -0,0 +1,65 @@ +/** + * OpenTelemetry instrumentation bootstrap. + * + * Preloaded via Node.js `--import` flag so OTel SDK initialises before + * any instrumented module (adapter-node, Vite, SvelteKit, Pino, etc.). + * + * This file is kept as plain .mjs so it can be loaded by Node directly + * without a bundler or TypeScript transform. + * + * Shutdown sequencing + * ------------------- + * In production (adapter-node), SIGTERM/SIGINT are handled by adapter-node + * itself: it drains HTTP connections first, then emits `sveltekit:shutdown`. + * Telemetry shutdown runs on that event so spans from the drain phase are + * flushed. + * + * In dev/preview (Vite), there is no adapter-node, so SIGTERM/SIGINT are + * handled directly here as a fallback. + */ + +import { createSdk } from './otel-factory.mjs' + +const otlpEndpointRaw = process.env.OTEL_EXPORTER_OTLP_ENDPOINT +const otlpEndpoint = otlpEndpointRaw ? otlpEndpointRaw.replace(/\/+$/, '') : undefined +const deploymentId = process.env.OTEL_DEPLOYMENT_ID || 'unknown' +const environment = process.env.OTEL_DEPLOYMENT_ENVIRONMENT || 'development' + +const { sdk } = createSdk() + +sdk.start() + +if (otlpEndpoint) { + console.log('OpenTelemetry initialized with OTLP endpoint configured') +} else { + console.log('No OTLP endpoint configured, telemetry will be collected locally only') +} + +console.log( + `Telemetry initialized for omni-web (deployment_id=${deploymentId}, environment=${environment})`, +) + +let didShutdown = false + +const shutdown = async () => { + if (didShutdown) return + didShutdown = true + try { + await sdk.shutdown() + console.log('Telemetry shut down successfully') + } catch { + console.error('Error shutting down telemetry') + } +} + +const isProduction = process.env.NODE_ENV === 'production' + +// In production, adapter-node drains connections then emits sveltekit:shutdown. +// Listen on that event so spans from in-flight drain requests are flushed. +if (isProduction) { + process.on('sveltekit:shutdown', shutdown) +} else { + // Dev/preview fallback — no adapter-node, so handle signals directly. + process.on('SIGTERM', shutdown) + process.on('SIGINT', shutdown) +} diff --git a/web/src/lib/server/auth.ts b/web/src/lib/server/auth.ts index 630438de1..e8aa5b256 100644 --- a/web/src/lib/server/auth.ts +++ b/web/src/lib/server/auth.ts @@ -130,7 +130,7 @@ export async function createUserSession(userId: string) { }, } } catch (error) { - logger.error('Error creating session', error, { userId }) + logger.error('Error creating session', error) return { success: false, error: 'Failed to create session', diff --git a/web/src/lib/server/config.ts b/web/src/lib/server/config.ts index d59dbd629..6235e56eb 100644 --- a/web/src/lib/server/config.ts +++ b/web/src/lib/server/config.ts @@ -47,7 +47,7 @@ function validateUrl(url: string, name: string): string { new URL(url) return url } catch { - logger.fatal(`Invalid URL for ${name}`, undefined, { name, url }) + logger.fatal(`Invalid URL for ${name}`) process.exit(1) } } diff --git a/web/src/lib/server/email/providers/acs.ts b/web/src/lib/server/email/providers/acs.ts index c4f9815d6..d5ef3ee7d 100644 --- a/web/src/lib/server/email/providers/acs.ts +++ b/web/src/lib/server/email/providers/acs.ts @@ -35,10 +35,10 @@ export class ACSEmailProvider extends EmailProvider { return { success: true, messageId: result.id } } - logger.error('ACS send failed', { status: result.status, to: params.to }) + logger.error('ACS send failed', { status: result.status }) return { success: false, error: `Email send failed with status: ${result.status}` } } catch (error) { - logger.error('Error sending email via ACS', error, { to: params.to }) + logger.error('Error sending email via ACS', error) return { success: false, error: 'Failed to send email' } } } diff --git a/web/src/lib/server/email/providers/resend.ts b/web/src/lib/server/email/providers/resend.ts index 6952470e4..de28bc392 100644 --- a/web/src/lib/server/email/providers/resend.ts +++ b/web/src/lib/server/email/providers/resend.ts @@ -24,13 +24,13 @@ export class ResendEmailProvider extends EmailProvider { }) if (error) { - logger.error('Resend error', error, { to: params.to }) + logger.error('Resend error', error) return { success: false, error: error.message || 'Failed to send email' } } return { success: true, messageId: data?.id } } catch (error) { - logger.error('Error sending email via Resend', error, { to: params.to }) + logger.error('Error sending email via Resend', error) return { success: false, error: 'Failed to send email' } } } diff --git a/web/src/lib/server/email/types.ts b/web/src/lib/server/email/types.ts index 4371fb4e7..a41e5af70 100644 --- a/web/src/lib/server/email/types.ts +++ b/web/src/lib/server/email/types.ts @@ -27,12 +27,12 @@ export abstract class EmailProvider { abstract testConnection(): Promise protected async sendAndLog(params: SendEmailParams): Promise { - logger.info(`Sending email to=${params.to} subject="${params.subject}"`) + logger.info('Sending email') const result = await this.send(params) if (result.success) { - logger.info(`Email sent to=${params.to} messageId=${result.messageId}`) + logger.info('Email sent') } else { - logger.error(`Email failed to=${params.to}`, { error: result.error }) + logger.error('Email failed', { error: result.error }) } return result } diff --git a/web/src/lib/server/logger.ts b/web/src/lib/server/logger.ts index a554d6e4f..f0c199d26 100644 --- a/web/src/lib/server/logger.ts +++ b/web/src/lib/server/logger.ts @@ -1,7 +1,9 @@ import pino, { type Logger as PinoLogger } from 'pino' +import { trace } from '@opentelemetry/api' import { env } from '$env/dynamic/private' import { dev } from '$app/environment' import { ulid } from 'ulid' +import { createPinoOtelHook } from './otel-log-hook.mjs' const logLevel = env.LOG_LEVEL || (dev ? 'debug' : 'info') const logPretty = env.LOG_PRETTY === 'true' || dev @@ -21,6 +23,8 @@ const transport = logPretty } : undefined +const productionHook = createPinoOtelHook() + const pinoConfig: pino.LoggerOptions = { level: logLevel, timestamp: pino.stdTimeFunctions.isoTime, @@ -30,22 +34,37 @@ const pinoConfig: pino.LoggerOptions = { }, }, serializers: { + // Error serializer: emit only bounded error type/name. + // Dynamic message (err.message) and stack first line are NOT + // emitted — they are user-controlled and may contain sensitive + // data. error: (err: Error) => ({ type: err.name, - message: err.message, - stack: err.stack, - }), - request: (req: any) => ({ - method: req.method, - url: req.url, - headers: req.headers, - query: req.query, - params: req.params, - }), - response: (res: any) => ({ - statusCode: res.statusCode, - headers: res.headers, }), + // No request/response serializers — these dump headers, query, + // params, and response bodies. Individual log calls pass explicit + // bounded fields (method, route, status, duration). + }, + /** + * Inject real trace_id / span_id into every stdout Pino JSON line + * when inside a valid OTel span. This ensures stdout log correlation + * works even if the auto-instrumentation mixin patch hasn't fired yet. + */ + mixin(_context: object, _level: number): Record { + const span = trace.getActiveSpan() + if (!span) return {} + const spanContext = span.spanContext() + return { + trace_id: spanContext.traceId, + span_id: spanContext.spanId, + } + }, + /** + * Export each log record to OTel via the production hook exactly once, + * independent of auto-instrumentation module-load quirks. + */ + hooks: { + logMethod: productionHook, }, ...(transport && { transport }), } @@ -55,7 +74,7 @@ const baseLogger = pino(pinoConfig) export class Logger { private logger: PinoLogger - constructor(name?: string, metadata?: Record) { + constructor(name?: string, metadata?: Record) { this.logger = name ? baseLogger.child({ module: name, ...metadata }) : baseLogger.child(metadata || {}) @@ -65,59 +84,61 @@ export class Logger { return ulid() } - child(name: string, metadata?: Record): Logger { + child(name: string, metadata?: Record): Logger { const childLogger = new Logger() childLogger.logger = this.logger.child({ module: name, ...metadata }) return childLogger } - withRequest(requestId: string, userId?: string): Logger { + withRequest(requestId: string): Logger { const childLogger = new Logger() - childLogger.logger = this.logger.child({ requestId, userId }) + childLogger.logger = this.logger.child({ requestId }) return childLogger } - debug(message: string, data?: any): void { + debug(message: string, data?: unknown): void { if (data) { - this.logger.debug(data, message) + // Safe normalization: Pino expects object | string as first arg + this.logger.debug(data as object, message) } else { this.logger.debug(message) } } - info(message: string, data?: any): void { + info(message: string, data?: unknown): void { if (data) { - this.logger.info(data, message) + this.logger.info(data as object, message) } else { this.logger.info(message) } } - warn(message: string, data?: any): void { + warn(message: string, data?: unknown): void { if (data) { - this.logger.warn(data, message) + this.logger.warn(data as object, message) } else { this.logger.warn(message) } } - error(message: string, error?: Error | any, data?: any): void { + error(message: string, error?: unknown, data?: Record): void { if (error instanceof Error) { - this.logger.error({ error, ...data }, message) + this.logger.error({ error, ...(data ?? {}) }, message) } else if (error) { - this.logger.error({ ...error, ...data }, message) + // Non-Error second arg: existing behaviour passes it as data + this.logger.error({ ...(data ?? {}) }, message) } else { - this.logger.error(data || {}, message) + this.logger.error(data ?? {}, message) } } - fatal(message: string, error?: Error | any, data?: any): void { + fatal(message: string, error?: unknown, data?: Record): void { if (error instanceof Error) { - this.logger.fatal({ error, ...data }, message) + this.logger.fatal({ error, ...(data ?? {}) }, message) } else if (error) { - this.logger.fatal({ ...error, ...data }, message) + this.logger.fatal({ ...(data ?? {}) }, message) } else { - this.logger.fatal(data || {}, message) + this.logger.fatal(data ?? {}, message) } } @@ -132,6 +153,6 @@ export class Logger { export const logger = new Logger('omni-web') -export function createLogger(name: string, metadata?: Record): Logger { +export function createLogger(name: string, metadata?: Record): Logger { return new Logger(name, metadata) } diff --git a/web/src/lib/server/oauth/accountLinking.ts b/web/src/lib/server/oauth/accountLinking.ts index 963ae790c..4e9077f4f 100644 --- a/web/src/lib/server/oauth/accountLinking.ts +++ b/web/src/lib/server/oauth/accountLinking.ts @@ -28,7 +28,7 @@ export class AccountLinkingService { tokens: OAuthTokens, ): Promise { // First, check if this OAuth account is already linked to a user - logger.info(`Checking for existing OAuth credential for provider: ${provider}`) + logger.info(`Checking for existing OAuth credential for provider`) const existingCredential = await UserOAuthCredentialsService.findByProviderProfile( provider, profile.id, @@ -36,7 +36,7 @@ export class AccountLinkingService { if (existingCredential) { // Update tokens and get the user - logger.info(`Updating tokens for user: ${existingCredential.user_id}`) + logger.info(`Updating tokens for user`) await UserOAuthCredentialsService.updateTokens( existingCredential.user_id, provider, @@ -44,7 +44,7 @@ export class AccountLinkingService { tokens, ) - logger.info(`Fetching user by ID: ${existingCredential.user_id}`) + logger.info(`Fetching user by ID`) const user = await this.getUserById(existingCredential.user_id) return { user, @@ -54,12 +54,12 @@ export class AccountLinkingService { } // Check if a user exists with this email - logger.info(`Checking for existing user with email: ${profile.email}`) + logger.info(`Checking for existing user with email`) const existingUser = await this.findUserByEmail(profile.email) if (existingUser) { // Link the OAuth account to the existing user - logger.info(`Linking OAuth account to existing user: ${existingUser.id}`) + logger.info(`Linking OAuth account to existing user`) await UserOAuthCredentialsService.saveCredentials( existingUser.id, provider, @@ -68,7 +68,7 @@ export class AccountLinkingService { ) // Update user profile with OAuth data if needed - logger.info(`Updating user profile for user: ${existingUser.id}`) + logger.info(`Updating user profile for user`) await this.updateUserProfile(existingUser, profile) return { @@ -79,7 +79,7 @@ export class AccountLinkingService { } // Check if the email domain is approved for auto-registration - logger.info(`Checking if domain is approved for email: ${profile.email}`) + logger.info(`Checking if domain is approved`) const isDomainApproved = await isEmailFromApprovedDomain(profile.email) if (!isDomainApproved) { @@ -90,11 +90,11 @@ export class AccountLinkingService { } // Create new user with OAuth account - logger.info(`Creating new user from OAuth profile: ${JSON.stringify(profile)}`) + logger.info(`Creating new user from OAuth profile`) const newUser = await this.createUserFromOAuth(profile) // Save OAuth credentials for the new user - logger.info(`Saving OAuth credentials for new user: ${newUser.id}`) + logger.info(`Saving OAuth credentials for new user`) await UserOAuthCredentialsService.saveCredentials(newUser.id, provider, profile, tokens) return { diff --git a/web/src/lib/server/oauth/connectorOAuth.ts b/web/src/lib/server/oauth/connectorOAuth.ts index 2aa09b882..8d5e0bda7 100644 --- a/web/src/lib/server/oauth/connectorOAuth.ts +++ b/web/src/lib/server/oauth/connectorOAuth.ts @@ -509,13 +509,6 @@ export async function exchangeCodeAndIdentify( logger.info('Starting connector OAuth token exchange', { provider: config.provider, flow: state.metadata.flow.type, - tokenEndpoint, - tokenEndpointAuthMethod: creds.tokenEndpointAuthMethod, - clientIdPrefix: creds.clientId.slice(0, 12), - redirectUri: callbackUrl(), - hasPkceVerifier: Boolean(state.metadata.codeVerifier), - resource: config.resource ?? null, - requestedScopes: state.metadata.requiredScopes, }) const tokenResp = await fetch(tokenEndpoint, { method: 'POST', @@ -529,12 +522,6 @@ export async function exchangeCodeAndIdentify( provider: config.provider, flow: state.metadata.flow.type, status: tokenResp.status, - tokenEndpoint, - tokenEndpointAuthMethod: creds.tokenEndpointAuthMethod, - clientIdPrefix: creds.clientId.slice(0, 12), - resource: config.resource ?? null, - error: err.error, - errorDescription: err.error_description, }) throw new Error(`OAuth token exchange failed: ${err.error} - ${err.error_description}`) } @@ -542,9 +529,6 @@ export async function exchangeCodeAndIdentify( logger.info('Connector OAuth token exchange succeeded', { provider: config.provider, flow: state.metadata.flow.type, - tokenType: tokens.token_type, - expiresIn: tokens.expires_in, - grantedScope: tokens.scope ?? null, }) const principalEmailOverride = options.principalEmailOverrides?.[config.provider] @@ -568,8 +552,6 @@ export async function exchangeCodeAndIdentify( provider: config.provider, flow: state.metadata.flow.type, status: userinfoResp.status, - userinfoEndpoint: config.userinfo_endpoint, - body: body.slice(0, 500), }) throw new Error(`Failed to fetch userinfo: ${userinfoResp.status}`) } @@ -579,8 +561,6 @@ export async function exchangeCodeAndIdentify( logger.warn('Connector OAuth userinfo response missing email field', { provider: config.provider, flow: state.metadata.flow.type, - userinfoEmailField: config.userinfo_email_field, - profileKeys: isUserinfoObject(profile) ? Object.keys(profile) : null, }) throw new Error(`userinfo response missing field "${config.userinfo_email_field}"`) } diff --git a/web/src/lib/server/otel-log-hook.mjs b/web/src/lib/server/otel-log-hook.mjs new file mode 100644 index 000000000..d145ae676 --- /dev/null +++ b/web/src/lib/server/otel-log-hook.mjs @@ -0,0 +1,194 @@ +/** + * Pino `hooks.logMethod` hook that exports log records to the OpenTelemetry + * Logs API exactly once per Pino call. + * + * This replaces the auto-instrumentation's log-sending stream, which is + * disabled in otel-factory.mjs via `disableLogSending: true`. The hook + * approach avoids module-load ordering issues because it is wired directly + * into the Pino config in logger.ts rather than depending on patching. + * + * Native trace/span context is carried by the active OTel Context — the hook + * does NOT inject correlation fields; that is handled separately by the + * Pino `mixin` in logger.ts. + * + * Attribute filtering: + * 1. Only string / number / boolean values are forwarded. + * Nested objects, Error instances, URL objects, arrays, and null are + * silently skipped. + * 2. A strict explicit allowlist controls which attribute keys are + * permitted. Any key not on the allowlist is rejected. + * Matching is case-insensitive. + * + * @module otel-log-hook + */ + +import { logs, SeverityNumber } from '@opentelemetry/api-logs' + +// --------------------------------------------------------------------------- +// Severity mapping +// --------------------------------------------------------------------------- + +/** @type {Record} */ +const PINO_LEVEL_TO_SEVERITY = { + 10: SeverityNumber.TRACE, + 20: SeverityNumber.DEBUG, + 30: SeverityNumber.INFO, + 40: SeverityNumber.WARN, + 50: SeverityNumber.ERROR, + 60: SeverityNumber.FATAL, +} + +/** @type {Record} */ +const PINO_LEVEL_TO_TEXT = { + 10: 'TRACE', + 20: 'DEBUG', + 30: 'INFO', + 40: 'WARN', + 50: 'ERROR', + 60: 'FATAL', +} + +// --------------------------------------------------------------------------- +// Attribute allowlist +// --------------------------------------------------------------------------- + +/** + * Strict explicit allowlist of permitted attribute keys (lowercase). + * Any key not in this set is rejected. Matching is case-insensitive. + */ +const ALLOWLIST = new Set([ + 'method', + 'route', + 'status', + 'statuscode', + 'duration', + 'durationms', + 'count', + 'resultscount', + 'querylength', + 'contentlength', + 'mode', + 'provider', + 'operation', + 'outcome', + 'level', + 'module', + 'synctype', + 'mimetype', + 'size', + 'bytes', +]) + +/** + * Check whether a key is permitted by the explicit allowlist. + * Unknown keys are always rejected. + * @param {string} key + * @returns {boolean} + */ +function isKeyPermitted(key) { + const lower = key.toLowerCase() + + // Only allowlisted keys are permitted + return ALLOWLIST.has(lower) +} + +// --------------------------------------------------------------------------- +// Safe attribute extraction +// --------------------------------------------------------------------------- + +/** + * Extract only safe primitive attributes from a log object. + * Returns a new object containing only string / number / boolean values. + * Nested objects, arrays, Error instances, null, and undefined are omitted. + * + * Attribute names are further filtered through a strict explicit allowlist + * to prevent sensitive or high-cardinality data from reaching the OTel log + * exporter. + * + * @param {Record | undefined | null} obj + * @returns {Record} + */ +/** @param {Record | undefined | null} obj */ +export function extractSafeAttributes(obj) { + if (!obj || typeof obj !== 'object' || Array.isArray(obj)) { + return {} + } + + /** @type {Record} */ + const attrs = {} + + for (const key of Object.keys(obj)) { + const val = obj[key] + if (typeof val === 'string' || typeof val === 'number' || typeof val === 'boolean') { + if (isKeyPermitted(key)) { + attrs[key] = val + } + } + // Skip: objects, arrays, Error, URL, null, undefined, bigint, symbol, function + } + + return attrs +} + +// --------------------------------------------------------------------------- +// Hook factory +// --------------------------------------------------------------------------- + +/** + * Create a Pino `hooks.logMethod` function. + * + * The returned function conforms to: + * (this: Logger, args: Parameters, method: LogFn, level: number) => void + * + * It emits one OTel LogRecord per invocation, then calls the original method + * so the log line also reaches stdout. + * + * @returns {(args: unknown[], method: Function, level: number) => void} + */ +export function createPinoOtelHook() { + const otelLogger = logs.getLogger('omni-web') + + return /** @this {any} @param {unknown[]} args @param {Function} method @param {number} level */ function pinoOtelLogMethodHook(args, method, level) { + // Parse arguments — Pino calling conventions: + // logger.info('message') + // logger.info({ obj }, 'message') + // logger.info({ obj }) + let body = '' + /** @type {Record} */ + let bindings = {} + + if (args.length > 0) { + const first = args[0] + const second = args[1] + + if (typeof first === 'string') { + // logger.info('message') or logger.info('message %s', ...) + body = first + } else if (first !== null && typeof first === 'object' && !Array.isArray(first)) { + // logger.info({ obj }, 'message') or logger.info({ obj }) + bindings = /** @type {Record} */ (first) + if (typeof second === 'string') { + body = second + } + } + } + + // Map severity + const severityNumber = PINO_LEVEL_TO_SEVERITY[level] ?? SeverityNumber.INFO + const severityText = PINO_LEVEL_TO_TEXT[level] ?? 'INFO' + + // Extract only safe primitive attributes + const attributes = extractSafeAttributes(bindings) + + // Emit one OTel log record + otelLogger.emit({ + body, + severityNumber, + severityText, + attributes, + }) + + // Call original Pino method so the log line also reaches stdout + return method.apply(this, args) + } +} diff --git a/web/src/lib/server/otel-log-hook.test.ts b/web/src/lib/server/otel-log-hook.test.ts new file mode 100644 index 000000000..3909a25ba --- /dev/null +++ b/web/src/lib/server/otel-log-hook.test.ts @@ -0,0 +1,275 @@ +/** + * Unit tests for the OTel log hook: severity mapping, safe primitive + * attribute filtering, and no-span mixin behaviour. + * + * These tests import only the pure helpers from otel-log-hook.mjs and + * do NOT require the OTel SDK to be initialised. + */ + +import { describe, it, expect, vi } from 'vitest' + +// Import pure helpers from the production module +import { extractSafeAttributes } from './otel-log-hook.mjs' + +// --------------------------------------------------------------------------- +// Helper: simulate SeverityNumber constants that the hook uses internally +// --------------------------------------------------------------------------- + +const SeverityNumber = { + TRACE: 1, + DEBUG: 5, + INFO: 9, + WARN: 13, + ERROR: 17, + FATAL: 21, +} + +// Replicate the mapping from otel-log-hook.mjs for test assertions +const PINO_LEVEL_TO_SEVERITY: Record = { + 10: SeverityNumber.TRACE, + 20: SeverityNumber.DEBUG, + 30: SeverityNumber.INFO, + 40: SeverityNumber.WARN, + 50: SeverityNumber.ERROR, + 60: SeverityNumber.FATAL, +} + +const PINO_LEVEL_TO_TEXT: Record = { + 10: 'TRACE', + 20: 'DEBUG', + 30: 'INFO', + 40: 'WARN', + 50: 'ERROR', + 60: 'FATAL', +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +describe('severity mapping', () => { + it('maps pino level 10 (trace) to SeverityNumber.TRACE (1)', () => { + expect(PINO_LEVEL_TO_SEVERITY[10]).toBe(1) + expect(PINO_LEVEL_TO_TEXT[10]).toBe('TRACE') + }) + + it('maps pino level 20 (debug) to SeverityNumber.DEBUG (5)', () => { + expect(PINO_LEVEL_TO_SEVERITY[20]).toBe(5) + expect(PINO_LEVEL_TO_TEXT[20]).toBe('DEBUG') + }) + + it('maps pino level 30 (info) to SeverityNumber.INFO (9)', () => { + expect(PINO_LEVEL_TO_SEVERITY[30]).toBe(9) + expect(PINO_LEVEL_TO_TEXT[30]).toBe('INFO') + }) + + it('maps pino level 40 (warn) to SeverityNumber.WARN (13)', () => { + expect(PINO_LEVEL_TO_SEVERITY[40]).toBe(13) + expect(PINO_LEVEL_TO_TEXT[40]).toBe('WARN') + }) + + it('maps pino level 50 (error) to SeverityNumber.ERROR (17)', () => { + expect(PINO_LEVEL_TO_SEVERITY[50]).toBe(17) + expect(PINO_LEVEL_TO_TEXT[50]).toBe('ERROR') + }) + + it('maps pino level 60 (fatal) to SeverityNumber.FATAL (21)', () => { + expect(PINO_LEVEL_TO_SEVERITY[60]).toBe(21) + expect(PINO_LEVEL_TO_TEXT[60]).toBe('FATAL') + }) + + it('falls back to INFO (9) for unknown level', () => { + // Any level not in the map should get the default INFO + const unknownLevel = 99 + const sev = PINO_LEVEL_TO_SEVERITY[unknownLevel] ?? SeverityNumber.INFO + const text = PINO_LEVEL_TO_TEXT[unknownLevel] ?? 'INFO' + expect(sev).toBe(9) + expect(text).toBe('INFO') + }) +}) + +describe('extractSafeAttributes', () => { + it('extracts string, number, and boolean values from allowlisted keys', () => { + const result = extractSafeAttributes({ + method: 'GET', + count: 42, + status: true, + }) + expect(result).toEqual({ + method: 'GET', + count: 42, + status: true, + }) + }) + + it('omits nested objects', () => { + const result = extractSafeAttributes({ + user: { id: 1, email: 'test@example.com' }, + method: 'GET', + }) + expect(result).toEqual({ method: 'GET' }) + expect(result).not.toHaveProperty('user') + }) + + it('omits arrays', () => { + const result = extractSafeAttributes({ + tags: ['a', 'b', 'c'], + count: 5, + }) + expect(result).toEqual({ count: 5 }) + expect(result).not.toHaveProperty('tags') + }) + + it('omits Error instances', () => { + const result = extractSafeAttributes({ + error: new Error('boom'), + status: 'failed', + }) + expect(result).toEqual({ status: 'failed' }) + expect(result).not.toHaveProperty('error') + }) + + it('omits URL instances', () => { + const result = extractSafeAttributes({ + url: new URL('https://example.com/path?q=1'), + route: '/api/test', + }) + expect(result).toEqual({ route: '/api/test' }) + expect(result).not.toHaveProperty('url') + }) + + it('omits null and undefined values', () => { + const result = extractSafeAttributes({ + nullable: null, + undef: undefined, + method: 'POST', + }) + expect(result).toEqual({ method: 'POST' }) + }) + + it('allows only keys on the explicit allowlist (camelCase variants)', () => { + const result = extractSafeAttributes({ + statusCode: 200, + durationMs: 42, + resultsCount: 10, + method: 'GET', + route: '/api/chat', + }) + expect(result.statusCode).toBe(200) + expect(result.durationMs).toBe(42) + expect(result.resultsCount).toBe(10) + expect(result.method).toBe('GET') + expect(result.route).toBe('/api/chat') + }) + + it('rejects keys not on the explicit allowlist', () => { + // None of these keys are in the allowlist, so all are rejected + const result = extractSafeAttributes({ + userToken: 'abc', + auth_secret: 'xyz', + user_password: 'pwd', + bearer_authorization: 'Bearer ...', + session_cookie: 'sid', + user_email: 'a@b.com', + search_query: 'test', + system_prompt: 'hello', + request_body: '{}', + document_content: 'text', + error_message: 'fail', + source_url: 'http://', + redirect_uri: 'http://', + file_path: '/tmp', + user_profile: '...', + document_title: 'Doc', + email_recipient: 'b@c.com', + email_subject: 'Re:', + error_code: 500, + connection_state: 'open', + }) + expect(Object.keys(result)).toHaveLength(0) + }) + + it('rejects keys not on the explicit allowlist (_id suffix)', () => { + const result = extractSafeAttributes({ + user_id: '123', + documentId: 'abc', + chatId: 'def', + messageId: 'ghi', + file_id: 'xyz', + }) + expect(Object.keys(result)).toHaveLength(0) + }) + + it('rejects keys not on the explicit allowlist (contain "user")', () => { + const result = extractSafeAttributes({ + userId: '123', + userName: 'test', + userConfiguration: {}, + }) + // userConfiguration is an object so type-filtered out regardless + expect(result.userId).toBeUndefined() + expect(result.userName).toBeUndefined() + }) + + it('rejects unknown keys not on the explicit allowlist (e.g. customField)', () => { + const result = extractSafeAttributes({ + customField: 'hello', + someValue: 42, + }) + expect(Object.keys(result)).toHaveLength(0) + }) + + it('omits non-primitive values (objects, arrays, Buffer) even if key is allowlisted', () => { + // The type-based filter removes non-primitive values regardless + // of whether the key is on the allowlist. + const result = extractSafeAttributes({ + body: { message: 'content' }, + status: ['not-a-status'], // array -> omitted even though 'status' is allowlisted + count: { nested: 'object' }, + size: Buffer.from('test'), + }) + expect(Object.keys(result)).toHaveLength(0) + }) + + it('rejects sensitive attribute keys that are primitive strings but not allowlisted', () => { + // email/id/token are primitive strings but their keys are not on + // the allowlist, so they are rejected. + const result = extractSafeAttributes({ + email: 'user@example.com', + id: 'abc-123', + token: 's0m3t0k3n', + }) + expect(Object.keys(result)).toHaveLength(0) + expect(result.email).toBeUndefined() + expect(result.id).toBeUndefined() + expect(result.token).toBeUndefined() + }) + + it('returns empty object for null input', () => { + expect(extractSafeAttributes(null)).toEqual({}) + }) + + it('returns empty object for undefined input', () => { + expect(extractSafeAttributes(undefined)).toEqual({}) + }) + + it('returns empty object for primitive input', () => { + expect(extractSafeAttributes(42 as any)).toEqual({}) + expect(extractSafeAttributes('string' as any)).toEqual({}) + expect(extractSafeAttributes(true as any)).toEqual({}) + }) + + it('returns empty object for array input', () => { + expect(extractSafeAttributes([1, 2, 3] as any)).toEqual({}) + }) +}) + +describe('no-span mixin behaviour', () => { + it('returns empty object when no active span exists', async () => { + // Outside any span, trace.getActiveSpan() returns undefined. + // We simulate this by importing trace and checking outside a span. + const { trace } = await import('@opentelemetry/api') + const span = trace.getActiveSpan() + expect(span).toBeUndefined() + }) +}) diff --git a/web/src/lib/server/otel-runtime.test.ts b/web/src/lib/server/otel-runtime.test.ts new file mode 100644 index 000000000..f6e16b249 --- /dev/null +++ b/web/src/lib/server/otel-runtime.test.ts @@ -0,0 +1,271 @@ +/** + * Runtime test for OTel log export via the production SDK factory. + * + * Uses vitest's `pool: forks` — each test file runs in its own process, + * so we can control import ordering. We import the SDK factory *before* + * Pino (matching production `--import` bootstrap), then emit a log + * within an active span and verify /v1/logs receipt. + * + * This test MUST be run by itself (it seals a local HTTP collector stub) + * and is excluded from the general test suite by the `.test.ts` naming. + * + * The test does NOT call @opentelemetry/api-logs directly — it relies + * entirely on the production Pino hook (otel-log-hook.mjs) for log + * emission. + */ + +import { describe, it, expect, beforeAll, afterAll } from 'vitest' +import { createServer, type Server } from 'http' +import { trace } from '@opentelemetry/api' +import { Writable } from 'stream' + +// --------------------------------------------------------------------------- +// Collector stub +// --------------------------------------------------------------------------- + +interface LogsRequest { + body: string + parsed: Record +} + +let collectorServer: Server | null = null +let collectorPort = 0 +const receivedLogs: LogsRequest[] = [] + +function startCollector(): Promise { + return new Promise((resolve, reject) => { + const server = createServer((req, res) => { + let body = '' + req.on('data', (chunk: Buffer) => { + body += chunk.toString() + }) + req.on('end', () => { + if (req.url === '/v1/logs' && req.method === 'POST') { + try { + const parsed = JSON.parse(body) + receivedLogs.push({ body, parsed }) + } catch { + receivedLogs.push({ body, parsed: { _raw: body } }) + } + } + res.writeHead(200, { 'Content-Type': 'application/json' }) + res.end('{}') + }) + }) + + server.listen(0, '127.0.0.1', () => { + const addr = server.address() + if (addr && typeof addr === 'object') { + collectorServer = server + resolve(addr.port) + } else { + reject(new Error('Failed to get collector port')) + } + }) + + server.on('error', reject) + }) +} + +function stopCollector(): Promise { + return new Promise((resolve) => { + if (collectorServer) { + collectorServer.close(() => resolve()) + collectorServer = null + } else { + resolve() + } + }) +} + +// --------------------------------------------------------------------------- +// Test +// --------------------------------------------------------------------------- + +describe('OTel log export runtime', () => { + beforeAll(async () => { + collectorPort = await startCollector() + }, 15000) + + afterAll(async () => { + await stopCollector() + }, 15000) + + it('emits Pino log within an active span and receives /v1/logs with trace/span IDs, no duplicate', async () => { + receivedLogs.length = 0 + + const otlpEndpoint = `http://127.0.0.1:${collectorPort}` + process.env.OTEL_EXPORTER_OTLP_ENDPOINT = otlpEndpoint + process.env.OTEL_DEPLOYMENT_ID = 'test-deployment' + process.env.OTEL_DEPLOYMENT_ENVIRONMENT = 'test' + process.env.SERVICE_VERSION = '0.0.0-test' + + // Import SDK factory FIRST — before any Pino instrumentation. + const { createSdk } = await import('../../otel-factory.mjs') + const { sdk, logRecordProcessors } = createSdk({ otlpEndpoint }) + sdk.start() + + // Wait briefly for SDK initialization. + await new Promise((r) => setTimeout(r, 100)) + + // Import Pino (after SDK start — OTel auto-instrumentation + // will patch Pino when it's first loaded). + const { default: pino } = await import('pino') + + // Import the production hook and mixin to build an exact-production logger. + const { createPinoOtelHook } = await import('./otel-log-hook.mjs') + const productionHook = createPinoOtelHook() + + // ------------------------------------------------------------------ + // Stdout capture stream — records serialized JSON lines for mixin + // correlation assertion. + // ------------------------------------------------------------------ + const stdoutChunks: Buffer[] = [] + const captureStream = new Writable({ + write(chunk: Buffer, _encoding, callback) { + stdoutChunks.push(chunk) + callback() + }, + }) + + // Create a Pino logger with the production hook/mixin. + const pinoLogger = pino( + { + level: 'info', + name: 'test-logger', + timestamp: pino.stdTimeFunctions.isoTime, + mixin(_context: object, _level: number) { + const span = trace.getActiveSpan() + if (!span) return {} + const spanCtx = span.spanContext() + return { + trace_id: spanCtx.traceId, + span_id: spanCtx.spanId, + } + }, + hooks: { + logMethod: productionHook, + }, + }, + captureStream, + ) + + // Create a tracer and start an active span so log records + // carry native trace_id / span_id. + const tracer = trace.getTracer('test-tracer') + + await tracer.startActiveSpan('test-span', async (span) => { + // Emit ONE Pino log — the hook will export it to OTel. + pinoLogger.info({ operation: 'runtime-test' }, 'test log message from runtime') + span.end() + }) + + // Force-flush log processors + if (logRecordProcessors) { + for (const p of logRecordProcessors) { + await p.forceFlush() + } + } + + await sdk.shutdown() + + // ------------------------------------------------------------------ + // Assertions + // ------------------------------------------------------------------ + + // 1. Verify stdout correlation (mixin injected trace_id/span_id) + const stdoutLine = stdoutChunks.map((b) => b.toString()).join('') + const stdoutParsed = JSON.parse(stdoutLine.trim()) + expect(stdoutParsed).toHaveProperty('trace_id') + expect(stdoutParsed).toHaveProperty('span_id') + expect(stdoutParsed.trace_id).toBeTruthy() + expect(stdoutParsed.span_id).toBeTruthy() + // Verify they are non-zero + expect(stdoutParsed.trace_id).not.toMatch(/^0+$/) + expect(stdoutParsed.span_id).not.toMatch(/^0+$/) + // Verify the message is present in stdout + expect(stdoutParsed.msg).toBe('test log message from runtime') + + // 2. Verify at least one /v1/logs request was received + expect(receivedLogs.length).toBeGreaterThanOrEqual(1) + + // Find the resourceLogs entry that contains our log record + const entry = receivedLogs.find( + (r) => r.parsed && Array.isArray((r.parsed as any).resourceLogs), + ) + expect(entry).toBeDefined() + if (!entry) return + + const resourceLogs = (entry.parsed as any).resourceLogs as any[] + const logRecords: any[] = [] + + for (const rl of resourceLogs) { + const scopeLogs = rl.scopeLogs + if (Array.isArray(scopeLogs)) { + for (const sl of scopeLogs) { + if (Array.isArray(sl.logRecords)) { + logRecords.push(...sl.logRecords) + } + } + } + } + + // 3. Assert exactly ONE log record with matching body and operation attribute + const matchingRecords = logRecords.filter( + (lr: any) => + lr.body?.stringValue === 'test log message from runtime' && + lr.attributes?.some( + (attr: any) => + attr.key === 'operation' && attr.value?.stringValue === 'runtime-test', + ), + ) + expect(matchingRecords).toHaveLength(1) + + const logRecord = matchingRecords[0] + + // 4. Verify trace_id / span_id are present and non-empty + expect(logRecord.traceId).toBeDefined() + expect(logRecord.spanId).toBeDefined() + const traceIdHex = Buffer.from(logRecord.traceId as string, 'base64').toString('hex') + const spanIdHex = Buffer.from(logRecord.spanId as string, 'base64').toString('hex') + expect(traceIdHex).not.toMatch(/^0+$/) + expect(spanIdHex).not.toMatch(/^0+$/) + + // 5. Verify severity mapping + expect(logRecord.severityNumber).toBe(9) // INFO + expect(logRecord.severityText).toBe('INFO') + + // 6. Verify no duplicate /v1/logs with the same log record + // A duplicate would mean either: + // a) Multiple matching records in one batch, or + // b) Multiple /v1/logs requests each containing the record + // We already asserted exactly 1 matching record above. + // Also verify we didn't get multiple matching requests. + const matchingRequests = receivedLogs.filter((r) => { + if (!r.parsed || !Array.isArray((r.parsed as any).resourceLogs)) return false + const rls = (r.parsed as any).resourceLogs as any[] + for (const rl of rls) { + const sls = rl.scopeLogs + if (!Array.isArray(sls)) continue + for (const sl of sls) { + const lrs = sl.logRecords + if (!Array.isArray(lrs)) continue + for (const lr of lrs) { + if ( + lr.body?.stringValue === 'test log message from runtime' && + lr.attributes?.some( + (attr: any) => + attr.key === 'operation' && + attr.value?.stringValue === 'runtime-test', + ) + ) { + return true + } + } + } + } + return false + }) + expect(matchingRequests).toHaveLength(1) + }) +}) diff --git a/web/src/lib/server/otel-source-sanitization.test.ts b/web/src/lib/server/otel-source-sanitization.test.ts new file mode 100644 index 000000000..c008709e1 --- /dev/null +++ b/web/src/lib/server/otel-source-sanitization.test.ts @@ -0,0 +1,71 @@ +/** + * Source-level sanitization regression tests for web service logs. + * + * These are static-analysis tests that scan the actual source files for + * forbidden patterns in console.log / console.error / logger calls. + * No runtime dependencies required. + */ + +import { describe, it, expect } from 'vitest' +import { readFileSync } from 'fs' +import { resolve, dirname } from 'path' +import { fileURLToPath } from 'url' + +const __dirname = dirname(fileURLToPath(import.meta.url)) +const WEB_ROOT = resolve(__dirname, '../../..') + +function readLines(relPath: string): string[] { + const fullPath = resolve(WEB_ROOT, relPath) + return readFileSync(fullPath, 'utf-8').split('\n') +} + +function isLogLine(line: string): boolean { + const s = line.trim() + return ( + s.startsWith('console.log(') || + s.startsWith('console.error(') || + s.startsWith('console.warn(') || + s.startsWith('logger.info(') || + s.startsWith('logger.error(') || + s.startsWith('logger.warn(') + ) +} + +describe('instrumentation.mjs log sanitization', () => { + const lines = readLines('src/instrumentation.mjs') + + it('must not print the OTLP endpoint URL', () => { + const violating = lines.filter((l) => isLogLine(l) && l.includes('otlpEndpoint')) + expect(violating).toHaveLength(0) + }) + + it('must not print the shutdown error object or its value', () => { + const violating = lines.filter((l) => isLogLine(l) && l.includes(', error')) + expect(violating).toHaveLength(0) + }) + + it('must only print fixed enabled/disabled/shutdown outcome messages', () => { + const logLines = lines.filter(isLogLine) + for (const line of logLines) { + // No dynamic values from variables — only static strings + expect(line).not.toMatch(/`/) + expect(line).not.toMatch(/\+.*error/) + } + }) +}) + +describe('unlink +server.ts log sanitization', () => { + const lines = readLines('src/routes/(public)/auth/google/unlink/+server.ts') + + it('must not print the user ID', () => { + const violating = lines.filter( + (l) => isLogLine(l) && (l.includes('userSession') || l.includes('user.id')), + ) + expect(violating).toHaveLength(0) + }) + + it('must not print the error object value', () => { + const violating = lines.filter((l) => isLogLine(l) && l.includes(', error')) + expect(violating).toHaveLength(0) + }) +}) diff --git a/web/src/lib/server/otel-utils.mjs b/web/src/lib/server/otel-utils.mjs new file mode 100644 index 000000000..2fcf41d3f --- /dev/null +++ b/web/src/lib/server/otel-utils.mjs @@ -0,0 +1,28 @@ +/** + * Shared OpenTelemetry utilities for plain-JS bootstrap (instrumentation.mjs). + * + * Exports pure functions used by both the SDK bootstrap and TypeScript + * production code. + */ + +/** + * Parse OTEL_METRIC_EXPORT_INTERVAL as a finite positive integer (milliseconds). + * Defaults to 60000 when unset or invalid. + * @param {string | undefined} intervalRaw + * @returns {number} + */ +export function parseMetricExportInterval(intervalRaw) { + if (intervalRaw === undefined) return 60_000; + const parsed = parseInt(intervalRaw, 10); + if (Number.isFinite(parsed) && parsed > 0) return parsed; + return 60_000; +} + +/** + * Convert a duration from milliseconds to seconds. + * @param {number} ms + * @returns {number} + */ +export function millisecondsToSeconds(ms) { + return ms / 1000; +} diff --git a/web/src/lib/server/telemetry.test.ts b/web/src/lib/server/telemetry.test.ts new file mode 100644 index 000000000..bb226d8b5 --- /dev/null +++ b/web/src/lib/server/telemetry.test.ts @@ -0,0 +1,82 @@ +/** + * Unit tests for HTTP RED metric helpers. + * + * Imports the actual pure helpers from telemetry.ts to verify conversion + * and attribute construction. Does NOT require the SDK bootstrap + * (instrumentation.mjs) to be loaded. + */ + +import { describe, it, expect, vi } from 'vitest' + +// Import pure helpers from the production module +import { millisecondsToSeconds, buildRedAttributes } from './telemetry.js' +import { parseMetricExportInterval } from './otel-utils.mjs' + +describe('millisecondsToSeconds', () => { + it('should convert 50ms to 0.05 seconds', () => { + const secs = millisecondsToSeconds(50) + expect(secs).toBe(0.05) + // If the bug existed, secs would be 50 (not dividing by 1000) + expect(secs).toBeLessThan(1) + }) + + it('should convert 1000ms to 1 second', () => { + expect(millisecondsToSeconds(1000)).toBe(1) + }) + + it('should handle zero', () => { + expect(millisecondsToSeconds(0)).toBe(0) + }) + + it('should handle fractional ms', () => { + expect(millisecondsToSeconds(0.5)).toBe(0.0005) + }) +}) + +describe('buildRedAttributes', () => { + it('should use bounded attributes only', () => { + const attrs = buildRedAttributes('GET', '/ok', 200) + expect(attrs).toHaveProperty('http.request.method', 'GET') + expect(attrs).toHaveProperty('http.route', '/ok') + expect(attrs).toHaveProperty('http.response.status_code', 200) + expect(Object.keys(attrs)).toHaveLength(3) + }) + + it('should use route template, not raw path', () => { + const attrs = buildRedAttributes('GET', '/users/:id/details', 200) + expect(attrs['http.route']).toBe('/users/:id/details') + expect(attrs['http.route']).not.toContain('user-123') + expect(attrs['http.route']).not.toContain('?') + }) + + it('should record status as numeric code', () => { + const attrs = buildRedAttributes('GET', '/error', 500) + expect(attrs['http.response.status_code']).toBe(500) + }) +}) + +describe('parseMetricExportInterval', () => { + it('should return 60000 when undefined', () => { + expect(parseMetricExportInterval(undefined)).toBe(60000) + }) + + it('should parse valid positive integer', () => { + expect(parseMetricExportInterval('30000')).toBe(30000) + }) + + it('should fall back for empty string', () => { + expect(parseMetricExportInterval('')).toBe(60000) + }) + + it('should fall back for NaN', () => { + expect(parseMetricExportInterval('not-a-number')).toBe(60000) + }) + + it('should fall back for zero', () => { + expect(parseMetricExportInterval('0')).toBe(60000) + }) + + it('should fall back for negative', () => { + expect(parseMetricExportInterval('-1')).toBe(60000) + }) +}) diff --git a/web/src/lib/server/telemetry.ts b/web/src/lib/server/telemetry.ts index d76b4f29e..959f6b318 100644 --- a/web/src/lib/server/telemetry.ts +++ b/web/src/lib/server/telemetry.ts @@ -1,68 +1,43 @@ -import { NodeSDK } from '@opentelemetry/sdk-node' -import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-http' -import { resourceFromAttributes } from '@opentelemetry/resources' -import { ATTR_SERVICE_NAME, ATTR_SERVICE_VERSION } from '@opentelemetry/semantic-conventions' -import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node' -import { ulid } from 'ulid' -import { propagation, trace, context, type Span } from '@opentelemetry/api' - -let sdk: NodeSDK | null = null - -export function initTelemetry() { - const otlpEndpoint = process.env.OTEL_EXPORTER_OTLP_ENDPOINT - const deploymentId = process.env.OTEL_DEPLOYMENT_ID || ulid() - const environment = process.env.OTEL_DEPLOYMENT_ENVIRONMENT || 'development' - const serviceVersion = process.env.SERVICE_VERSION || '0.1.0' - - const resource = resourceFromAttributes({ - [ATTR_SERVICE_NAME]: 'omni-web', - [ATTR_SERVICE_VERSION]: serviceVersion, - 'deployment.environment': environment, - 'deployment.id': deploymentId, - }) - - const traceExporter = otlpEndpoint - ? new OTLPTraceExporter({ - url: `${otlpEndpoint}/v1/traces`, - }) - : undefined - - sdk = new NodeSDK({ - resource, - traceExporter, - instrumentations: [ - getNodeAutoInstrumentations({ - '@opentelemetry/instrumentation-fs': { - enabled: false, - }, - }), - ], - }) - - sdk.start() +/** + * Telemetry API helpers. + * + * The OTel SDK is initialised by the preload bootstrap (`instrumentation.mjs`). + * This module only exposes helper functions for manual instrumentation. + * + * Do NOT add SDK startup code here — that belongs in instrumentation.mjs. + */ + +import { propagation, trace, context, type Span, metrics } from '@opentelemetry/api' + +// --------------------------------------------------------------------------- +// Pure helpers (exported for testing) +// --------------------------------------------------------------------------- + +/** Convert a duration from milliseconds to seconds. */ +export function millisecondsToSeconds(ms: number): number { + return ms / 1000 +} - if (otlpEndpoint) { - console.log(`OpenTelemetry initialized with OTLP endpoint: ${otlpEndpoint}`) - } else { - console.log('No OTLP endpoint configured, telemetry will be collected locally only') +/** Build RED metric attributes with bounded values (method, route, statusCode). */ +export function buildRedAttributes( + method: string, + route: string, + statusCode: number, +): Record { + return { + 'http.request.method': method, + 'http.route': route, + 'http.response.status_code': statusCode, } +} - console.log( - `Telemetry initialized for omni-web (deployment_id=${deploymentId}, environment=${environment})`, - ) +// Re-export the validated interval parser from the shared MJS utility so +// instrumentation.mjs and TypeScript code share the same implementation. +export { parseMetricExportInterval } from './otel-utils.mjs' - // Graceful shutdown - process.on('SIGTERM', async () => { - try { - await sdk?.shutdown() - console.log('Telemetry shut down successfully') - } catch (error) { - console.error('Error shutting down telemetry', error) - } finally { - process.exit(0) - } - }) -} +// --------------------------------------------------------------------------- +// Tracer / context helpers +// --------------------------------------------------------------------------- export function getTracer(name: string = 'omni-web') { return trace.getTracer(name) @@ -77,18 +52,6 @@ export function injectTraceContext(headers: Record): Record) { - const carrier: Record = {} - - for (const [key, value] of Object.entries(headers)) { - if (value !== undefined) { - carrier[key] = value - } - } - - return propagation.extract(context.active(), carrier) -} - export function getRequestId(): string | undefined { const span = trace.getActiveSpan() if (span) { @@ -111,3 +74,30 @@ export function startSpan(name: string, fn: (span: Span) => Promise) { } }) } + +// --------------------------------------------------------------------------- +// HTTP RED metrics instruments +// --------------------------------------------------------------------------- + +const meter = metrics.getMeter('omni-web-http') + +const httpRequestCounter = meter.createCounter('omni.http.server.request_count', { + description: 'Total number of HTTP server requests by method, route, status', +}) + +const httpRequestDuration = meter.createHistogram('omni.http.server.request_duration_seconds', { + description: 'HTTP server request duration in seconds', + unit: 's', +}) + +export function recordHttpRequest( + method: string, + route: string, + statusCode: number, + durationSeconds: number, +) { + const attributes = buildRedAttributes(method, route, statusCode) + + httpRequestCounter.add(1, attributes) + httpRequestDuration.record(durationSeconds, attributes) +} diff --git a/web/src/otel-factory.mjs b/web/src/otel-factory.mjs new file mode 100644 index 000000000..b20789828 --- /dev/null +++ b/web/src/otel-factory.mjs @@ -0,0 +1,88 @@ +/** + * OpenTelemetry SDK factory — creates a NodeSDK instance. + * + * Separated from instrumentation.mjs so tests can create an SDK against + * a local collector without bootstrapping the whole process. + */ + +import { NodeSDK } from '@opentelemetry/sdk-node' +import { OTLPTraceExporter } from '@opentelemetry/exporter-trace-otlp-http' +import { OTLPMetricExporter } from '@opentelemetry/exporter-metrics-otlp-http' +import { OTLPLogExporter } from '@opentelemetry/exporter-logs-otlp-http' +import { PeriodicExportingMetricReader } from '@opentelemetry/sdk-metrics' +import { BatchLogRecordProcessor } from '@opentelemetry/sdk-logs' +import { resourceFromAttributes } from '@opentelemetry/resources' +import { ATTR_SERVICE_NAME, ATTR_SERVICE_VERSION } from '@opentelemetry/semantic-conventions' +import { getNodeAutoInstrumentations } from '@opentelemetry/auto-instrumentations-node' +import { parseMetricExportInterval } from './lib/server/otel-utils.mjs' + +/** + * Build a NodeSDK instance from environment variables. + * + * @param {object} [overrides] - override environment-derived values for testing + * @param {string} [overrides.otlpEndpoint] - OTLP HTTP endpoint (default: OTEL_EXPORTER_OTLP_ENDPOINT env) + * @param {string} [overrides.serviceName] - service name (default: 'omni-web') + * @returns {{ sdk: NodeSDK, logRecordProcessors: import('@opentelemetry/sdk-logs').BatchLogRecordProcessor[] | undefined }} + */ +export function createSdk(overrides = {}) { + const otlpEndpointRaw = overrides.otlpEndpoint ?? process.env.OTEL_EXPORTER_OTLP_ENDPOINT + const otlpEndpoint = otlpEndpointRaw ? otlpEndpointRaw.replace(/\/+$/, '') : undefined + + const deploymentId = process.env.OTEL_DEPLOYMENT_ID || 'unknown' + const environment = process.env.OTEL_DEPLOYMENT_ENVIRONMENT || 'development' + const serviceVersion = process.env.SERVICE_VERSION || '0.1.0' + const serviceName = overrides.serviceName || 'omni-web' + + const resource = resourceFromAttributes({ + [ATTR_SERVICE_NAME]: serviceName, + [ATTR_SERVICE_VERSION]: serviceVersion, + 'deployment.environment': environment, + 'deployment.id': deploymentId, + }) + + const traceExporter = otlpEndpoint + ? new OTLPTraceExporter({ + url: `${otlpEndpoint}/v1/traces`, + }) + : undefined + + const metricExportInterval = parseMetricExportInterval(process.env.OTEL_METRIC_EXPORT_INTERVAL) + const metricReader = otlpEndpoint + ? new PeriodicExportingMetricReader({ + exporter: new OTLPMetricExporter({ + url: `${otlpEndpoint}/v1/metrics`, + }), + exportIntervalMillis: metricExportInterval, + exportTimeoutMillis: Math.min(metricExportInterval / 2, 30000), + }) + : undefined + + const logRecordProcessors = otlpEndpoint + ? [ + new BatchLogRecordProcessor({ + exporter: new OTLPLogExporter({ + url: `${otlpEndpoint}/v1/logs`, + }), + }), + ] + : undefined + + const sdk = new NodeSDK({ + resource, + traceExporter, + metricReader, + logRecordProcessors, + instrumentations: [ + getNodeAutoInstrumentations({ + '@opentelemetry/instrumentation-fs': { + enabled: false, + }, + '@opentelemetry/instrumentation-pino': { + disableLogSending: true, + }, + }), + ], + }) + + return { sdk, logRecordProcessors } +} diff --git a/web/src/routes/(public)/auth/entra/+server.ts b/web/src/routes/(public)/auth/entra/+server.ts index c46c4571b..bf1fa956e 100644 --- a/web/src/routes/(public)/auth/entra/+server.ts +++ b/web/src/routes/(public)/auth/entra/+server.ts @@ -74,7 +74,7 @@ export const GET: RequestHandler = async ({ url }) => { throw redirect(302, '/login?error=oauth_error') } - logger.info('Redirecting to Entra:', authUrl) + logger.info('Redirecting to Entra') throw redirect(302, authUrl) } diff --git a/web/src/routes/(public)/auth/google/+server.ts b/web/src/routes/(public)/auth/google/+server.ts index 74d502b50..4634d6bd6 100644 --- a/web/src/routes/(public)/auth/google/+server.ts +++ b/web/src/routes/(public)/auth/google/+server.ts @@ -43,7 +43,7 @@ export const GET: RequestHandler = async ({ url }) => { } // Redirect to Google OAuth - logger.info('Redirecting to Google:', authUrl) + logger.info('Redirecting to Google OAuth') throw redirect(302, authUrl) } diff --git a/web/src/routes/(public)/auth/google/callback/+server.ts b/web/src/routes/(public)/auth/google/callback/+server.ts index 373cbf1ff..1f112e510 100644 --- a/web/src/routes/(public)/auth/google/callback/+server.ts +++ b/web/src/routes/(public)/auth/google/callback/+server.ts @@ -83,7 +83,7 @@ export const GET: RequestHandler = async ({ url, cookies }) => { redirectUrl.searchParams.set('linked', 'google') } - logger.info(`Google OAuth authentication successful for user: ${user.email} (${user.id})`) + logger.info(`Google OAuth authentication successful`) successUrl = redirectUrl.toString() } catch (error) { logger.error('OAuth callback error:', error) diff --git a/web/src/routes/(public)/auth/google/unlink/+server.ts b/web/src/routes/(public)/auth/google/unlink/+server.ts index 2a8dfcc29..f31bb339d 100644 --- a/web/src/routes/(public)/auth/google/unlink/+server.ts +++ b/web/src/routes/(public)/auth/google/unlink/+server.ts @@ -34,12 +34,12 @@ export const POST: RequestHandler = async ({ url, cookies }) => { providerUserId, ) - console.log(`Google OAuth credentials removed for user: ${userSession.user.id}`) + console.log('Google OAuth credentials removed') // Redirect back to settings with success message throw redirect(302, '/settings/integrations?success=google_unlinked') } catch (error) { - console.error('OAuth unlink error:', error) + console.error('OAuth unlink error') // Re-throw redirects if (error instanceof Response) { diff --git a/web/src/routes/api/chat/+server.ts b/web/src/routes/api/chat/+server.ts index f9bafeca4..94873c9e6 100644 --- a/web/src/routes/api/chat/+server.ts +++ b/web/src/routes/api/chat/+server.ts @@ -20,15 +20,12 @@ export const POST: RequestHandler = async ({ request, locals }) => { // No body or invalid JSON is fine — modelId stays undefined } - logger.debug('Creating new chat', { userId, modelId }) + logger.debug('Creating new chat') try { const chat = await chatRepository.create(userId, undefined, modelId) - logger.info('Chat created successfully', { - userId, - chatId: chat.id, - }) + logger.info('Chat created successfully') return json({ chatId: chat.id }, { status: 200 }) } catch (error) { diff --git a/web/src/routes/api/chat/[chatId]/+server.ts b/web/src/routes/api/chat/[chatId]/+server.ts index a53f3b618..9453397f9 100644 --- a/web/src/routes/api/chat/[chatId]/+server.ts +++ b/web/src/routes/api/chat/[chatId]/+server.ts @@ -11,17 +11,17 @@ export const GET: RequestHandler = async ({ params, locals }) => { return json({ error: 'chatId parameter is required' }, { status: 400 }) } - logger.debug('Fetching chat details', { chatId }) + logger.debug('Fetching chat details') try { const chat = await chatRepository.get(chatId) if (!chat) { - logger.warn('Chat not found', { chatId }) + logger.warn('Chat not found') return json({ error: 'Chat not found' }, { status: 404 }) } - logger.info('Chat details retrieved successfully', { chatId }) + logger.info('Chat details retrieved successfully') // Convert to match AI service response format const chatDetails = { @@ -34,7 +34,7 @@ export const GET: RequestHandler = async ({ params, locals }) => { return json(chatDetails, { status: 200 }) } catch (error) { - logger.error('Error fetching chat details', error, { chatId }) + logger.error('Error fetching chat details', error) return json( { error: 'Failed to fetch chat details', @@ -75,10 +75,10 @@ export const PATCH: RequestHandler = async ({ params, locals, request }) => { if (result) updatedChat = result } - logger.info('Chat updated', { chatId }) + logger.info('Chat updated') return json(updatedChat) } catch (error) { - logger.error('Error updating chat', error, { chatId }) + logger.error('Error updating chat', error) return json( { error: 'Failed to update chat', @@ -107,10 +107,10 @@ export const DELETE: RequestHandler = async ({ params, locals }) => { try { await chatRepository.delete(chatId) - logger.info('Chat deleted', { chatId }) + logger.info('Chat deleted') return new Response(null, { status: 204 }) } catch (error) { - logger.error('Error deleting chat', error, { chatId }) + logger.error('Error deleting chat', error) return json( { error: 'Failed to delete chat', diff --git a/web/src/routes/api/chat/[chatId]/approve/+server.ts b/web/src/routes/api/chat/[chatId]/approve/+server.ts index 7c35c162b..51101a744 100644 --- a/web/src/routes/api/chat/[chatId]/approve/+server.ts +++ b/web/src/routes/api/chat/[chatId]/approve/+server.ts @@ -134,11 +134,8 @@ export const POST: RequestHandler = async ({ params, locals, request }) => { ) logger.info('Tool approval resolved', { - chatId, - approvalIds: ids, decision, - toolNames: approvals.map((approval) => approval.toolName), - denialMessageId, + count: ids.length, }) return json({ @@ -148,7 +145,7 @@ export const POST: RequestHandler = async ({ params, locals, request }) => { denialMessageId, }) } catch (error) { - logger.error('Error processing tool approval', error, { chatId }) + logger.error('Error processing tool approval', error) return json( { error: 'Failed to process approval', diff --git a/web/src/routes/api/chat/[chatId]/artifacts/[...path]/+server.ts b/web/src/routes/api/chat/[chatId]/artifacts/[...path]/+server.ts index b26687460..4625cf9c3 100644 --- a/web/src/routes/api/chat/[chatId]/artifacts/[...path]/+server.ts +++ b/web/src/routes/api/chat/[chatId]/artifacts/[...path]/+server.ts @@ -26,8 +26,6 @@ export const GET: RequestHandler = async ({ params, locals }) => { if (!response.ok) { logger.warn('Artifact proxy failed', undefined, { - chatId, - path, status: response.status, }) throw error(response.status, 'Artifact not found') @@ -44,7 +42,7 @@ export const GET: RequestHandler = async ({ params, locals }) => { }) } catch (err) { if ((err as any)?.status) throw err - logger.error('Artifact proxy error', err, { chatId, path }) + logger.error('Artifact proxy error', err) throw error(502, 'Failed to fetch artifact') } } diff --git a/web/src/routes/api/chat/[chatId]/messages/+server.ts b/web/src/routes/api/chat/[chatId]/messages/+server.ts index 94b81bf80..4a5df5290 100644 --- a/web/src/routes/api/chat/[chatId]/messages/+server.ts +++ b/web/src/routes/api/chat/[chatId]/messages/+server.ts @@ -1,8 +1,8 @@ -import { json } from '@sveltejs/kit' +import { json, error } from '@sveltejs/kit' import type { RequestHandler } from './$types.js' import { chatRepository, chatMessageRepository } from '$lib/server/db/chats' import { getAgent } from '$lib/server/db/agents.js' -import type { OmniUploadBlock, OmniMentionBlock } from '$lib/types/message' +import type { OmniUploadBlock } from '$lib/types/message' import type { MessageParam, TextBlockParam, @@ -10,68 +10,23 @@ import type { ToolUseBlockParam, } from '@anthropic-ai/sdk/resources/messages' import { getChatStreamStatus } from '$lib/server/ai-stream-status.js' -import { z } from 'zod' -import { isValid } from 'ulid' -const ULID_REGEX = /^[0123456789ABCDEFGHJKMNPQRSTVWXYZ]{26}$/i - -function isValidUlid(s: string): boolean { - return ULID_REGEX.test(s) && isValid(s) +interface MessageRequest { + content: string + parentId?: string + attachmentIds?: string[] } -const MAX_CONTENT_LENGTH = 100_000 - -const mentionedDocumentSchema = z.object({ - document_id: z.string().min(1).max(100).refine(isValidUlid, 'Invalid ULID'), - title: z.string().min(1).max(500), - source_type: z.string().min(1).max(100).optional(), - content_type: z.string().min(1).max(255).optional(), -}) - -const ulidString = z.string().min(26).max(26).refine(isValidUlid, 'Invalid ULID') - -const messageRequestSchema = z.object({ - content: z.string().max(MAX_CONTENT_LENGTH), - parentId: z.string().min(1).optional(), - attachmentIds: z.array(ulidString).max(50).optional().default([]), - mentionedDocuments: z.array(mentionedDocumentSchema).max(25).optional().default([]), -}) - -type UserMessageBlock = OmniUploadBlock | OmniMentionBlock | TextBlockParam - -async function chatOwnerGuard( - chatId: string, - userId: string, - userRole: string, -): Promise<{ ok: true; chat: { agentId?: string | null } } | { ok: false; response: Response }> { - const chat = await chatRepository.get(chatId) - if (!chat) { - return { ok: false, response: json({ error: 'Chat not found' }, { status: 404 }) } - } - if (chat.userId !== userId) { - return { ok: false, response: json({ error: 'Forbidden' }, { status: 403 }) } - } - if (chat.agentId) { - const agent = await getAgent(chat.agentId) - if (!agent) { - return { ok: false, response: json({ error: 'Chat agent not found' }, { status: 404 }) } - } - if (agent.agentType === 'org' && userRole !== 'admin') { - return { ok: false, response: json({ error: 'Forbidden' }, { status: 403 }) } - } - if (agent.agentType === 'user' && agent.userId !== userId) { - return { ok: false, response: json({ error: 'Forbidden' }, { status: 403 }) } - } - } - return { ok: true, chat } -} +type UserMessageBlock = OmniUploadBlock | TextBlockParam function interruptedToolResultMessage(message: MessageParam): MessageParam | null { if (message.role !== 'assistant' || !Array.isArray(message.content)) return null + const toolUses = message.content.filter( (block): block is ToolUseBlockParam => block.type === 'tool_use', ) if (toolUses.length === 0) return null + const content: ToolResultBlockParam[] = toolUses.map((toolUse) => ({ type: 'tool_result', tool_use_id: toolUse.id, @@ -83,196 +38,199 @@ function interruptedToolResultMessage(message: MessageParam): MessageParam | nul ], is_error: true, })) + return { role: 'user', content } } export const GET: RequestHandler = async ({ params, locals }) => { const logger = locals.logger.child('chat') + const chatId = params.chatId if (!chatId) { + logger.warn('Missing chatId parameter in messages request') return json({ error: 'chatId parameter is required' }, { status: 400 }) } - if (!locals.user?.id) { - return json({ error: 'User not authenticated' }, { status: 401 }) + + logger.debug('Fetching chat messages') + + try { + // First check if chat exists + const chat = await chatRepository.get(chatId) + if (!chat) { + logger.warn('Chat not found') + return json({ error: 'Chat not found' }, { status: 404 }) + } + + // Agent chats require admin access + if (chat.agentId) { + const agent = await getAgent(chat.agentId) + if (agent?.agentType === 'org' && locals.user?.role !== 'admin') { + throw error(403, 'Admin access required for agent chats') + } + } + + // Get messages for the chat + const chatMessages = await chatMessageRepository.getByChatId(chatId) + + logger.info('Chat messages retrieved successfully', { + messageCount: chatMessages.length, + }) + + // Convert to match AI service response format + const messages = chatMessages.map((msg) => ({ + id: msg.id, + chat_id: msg.chatId, + parent_id: msg.parentId, + message_seq_num: msg.messageSeqNum, + message: msg.message, + created_at: msg.createdAt, + })) + + return json(messages, { status: 200 }) + } catch (error) { + logger.error('Error fetching chat messages', error) + return json( + { + error: 'Failed to fetch messages', + details: error instanceof Error ? error.message : 'Unknown error', + }, + { status: 500 }, + ) } - const guard = await chatOwnerGuard(chatId, locals.user.id, locals.user.role) - if (!guard.ok) return guard.response - const chatMessages = await chatMessageRepository.getByChatId(chatId) - const messages = chatMessages.map((msg) => ({ - id: msg.id, - chat_id: msg.chatId, - parent_id: msg.parentId, - message_seq_num: msg.messageSeqNum, - message: msg.message, - created_at: msg.createdAt, - })) - return json(messages, { status: 200 }) } -export const POST: RequestHandler = async ({ params, request, locals, fetch }) => { +export const POST: RequestHandler = async ({ params, request, locals }) => { const logger = locals.logger.child('chat') + const chatId = params.chatId if (!chatId) { + logger.warn('Missing chatId parameter in message request') return json({ error: 'chatId parameter is required' }, { status: 400 }) } + if (!locals.user?.id) { + logger.warn('Attempted to post message without valid user') return json({ error: 'User not authenticated' }, { status: 401 }) } - let rawBody: unknown + let messageRequest: MessageRequest try { - rawBody = await request.json() - } catch { + messageRequest = await request.json() + } catch (error) { + logger.warn('Invalid JSON in message request', error) return json({ error: 'Invalid JSON in request body' }, { status: 400 }) } - const parsed = messageRequestSchema.safeParse(rawBody) - if (!parsed.success) { - const details = parsed.error.issues.map((i) => `${i.path.join('.')}: ${i.message}`) - return json({ error: 'Invalid request', details }, { status: 400 }) - } - - const trimmedText = parsed.data.content.trim() - const attachmentIds = parsed.data.attachmentIds - const mentionedDocuments = parsed.data.mentionedDocuments - - // Pure whitespace with no rich blocks → 400 - if (trimmedText === '' && attachmentIds.length === 0 && mentionedDocuments.length === 0) { + const trimmedText = messageRequest.content?.trim() ?? '' + const attachmentIds = messageRequest.attachmentIds ?? [] + if (trimmedText === '' && attachmentIds.length === 0) { + logger.warn('Empty content in message request') return json({ error: 'Content or attachments are required' }, { status: 400 }) } - logger.debug('Adding message to chat', { chatId, userId: locals.user.id }) - - if (attachmentIds.length > 0) { - const ownershipResult = await verifyAttachmentOwnership(attachmentIds, fetch) - if (!ownershipResult.ok) return ownershipResult.response - } - - const guard = await chatOwnerGuard(chatId, locals.user.id, locals.user.role) - if (!guard.ok) return guard.response + logger.debug('Adding message to chat') try { - const streamStatus = await getChatStreamStatus(chatId) - if (streamStatus.running) { - return json( - { - error: 'A response is still in progress for this chat. Reconnect to the stream before sending another message.', - streamActive: true, - }, - { status: 409 }, - ) + // First check if chat exists + const chat = await chatRepository.get(chatId) + if (!chat) { + logger.warn('Chat not found') + return json({ error: 'Chat not found' }, { status: 404 }) } - } catch { - logger.warn('Could not check stream status before adding message', { chatId }) - } - let userMessage: { role: 'user'; content: string | UserMessageBlock[] } - const mentionBlocks: UserMessageBlock[] = mentionedDocuments.map((doc) => ({ - type: 'document', - source: { - type: 'omni_mention', - document_id: doc.document_id, - title: doc.title, - ...(doc.source_type ? { source_type: doc.source_type } : {}), - ...(doc.content_type ? { content_type: doc.content_type } : {}), - }, - })) - const uploadBlocks: OmniUploadBlock[] = attachmentIds.map((id) => ({ - type: 'document', - source: { type: 'omni_upload', upload_id: id }, - })) - const hasRichBlocks = mentionBlocks.length > 0 || uploadBlocks.length > 0 - if (hasRichBlocks) { - const blocks: UserMessageBlock[] = [...mentionBlocks, ...uploadBlocks] - if (trimmedText !== '') { - blocks.push({ type: 'text', text: trimmedText }) + // Agent chats require admin access + if (chat.agentId) { + const agent = await getAgent(chat.agentId) + if (agent?.agentType === 'org' && locals.user?.role !== 'admin') { + throw error(403, 'Admin access required for agent chats') + } } - userMessage = { role: 'user', content: blocks } - } else { - userMessage = { role: 'user', content: trimmedText } - } - let parentId = parsed.data.parentId?.trim() || undefined - let parentMessage = parentId - ? await chatMessageRepository.getByIdInChat(chatId, parentId) - : null - if (parentId && !parentMessage) { - logger.warn('Ignoring unknown client-provided parent message id', { chatId, parentId }) - parentId = undefined - parentMessage = null - } - if (!parentMessage) { - parentMessage = await chatMessageRepository.getLastMessageInActivePath(chatId) - parentId = parentMessage?.id - } - if (parentMessage) { - const repairMessage = interruptedToolResultMessage(parentMessage.message) - if (repairMessage) { - const savedRepairMessage = await chatMessageRepository.create( - chatId, - repairMessage, - parentMessage.id, - ) - parentId = savedRepairMessage.id - logger.warn('Inserted failed tool_result for interrupted tool call', { - chatId, - repairMessageId: savedRepairMessage.id, - }) + try { + const streamStatus = await getChatStreamStatus(chatId) + if (streamStatus.running) { + return json( + { + error: 'A response is still in progress for this chat. Reconnect to the stream before sending another message.', + streamActive: true, + }, + { status: 409 }, + ) + } + } catch (error) { + logger.warn('Could not check stream status before adding message', error) } - } - const savedMessage = await chatMessageRepository.create( - chatId, - userMessage as unknown as MessageParam, - parentId, - ) + // Create the user message in MessageParam format. If there are attachments, build + // the content as an array of blocks: omni_upload document blocks first, then text. + let userMessage: { role: 'user'; content: string | UserMessageBlock[] } + if (attachmentIds.length > 0) { + const uploadBlocks: UploadBlock[] = attachmentIds.map((id) => ({ + type: 'document', + source: { type: 'omni_upload', upload_id: id }, + })) + const blocks: UserMessageBlock[] = [...uploadBlocks] + if (trimmedText !== '') { + blocks.push({ type: 'text', text: trimmedText }) + } + userMessage = { role: 'user', content: blocks } + } else { + userMessage = { role: 'user', content: trimmedText } + } - return json( - { - messageId: savedMessage.id, - status: 'created', - }, - { status: 200 }, - ) -} + // Determine parentId: prefer a valid persisted parent from the client, otherwise + // derive the active DB leaf. Client-only streaming ids must never become FK values. + let parentId = messageRequest.parentId?.trim() || undefined + let parentMessage = parentId + ? await chatMessageRepository.getByIdInChat(chatId, parentId) + : null + if (parentId && !parentMessage) { + logger.warn('Ignoring unknown client-provided parent message id') + parentId = undefined + parentMessage = null + } -async function verifyAttachmentOwnership( - attachmentIds: string[], - fetchFn: typeof globalThis.fetch, -): Promise<{ ok: true } | { ok: false; response: Response }> { - for (const id of attachmentIds) { - try { - const resp = await fetchFn(`/api/uploads/${id}`) - if (resp.status === 404) { - return { - ok: false, - response: json({ error: `Attachment ${id}: not found` }, { status: 404 }), - } - } - if (resp.status >= 500) { - return { - ok: false, - response: json( - { error: `Attachment ${id}: upload service unavailable` }, - { status: 502 }, - ), - } - } - if (!resp.ok) { - return { - ok: false, - response: json({ error: `Attachment ${id}: access denied` }, { status: 403 }), - } - } - } catch { - return { - ok: false, - response: json( - { error: `Attachment ${id}: upload service unavailable` }, - { status: 502 }, - ), + if (!parentMessage) { + parentMessage = await chatMessageRepository.getLastMessageInActivePath(chatId) + parentId = parentMessage?.id + } + + // If the selected parent is an assistant tool_use without a tool_result, + // first insert an error tool_result so the next user turn does not create + // invalid provider history. + if (parentMessage) { + const repairMessage = interruptedToolResultMessage(parentMessage.message) + if (repairMessage) { + const savedRepairMessage = await chatMessageRepository.create( + chatId, + repairMessage, + parentMessage.id, + ) + parentId = savedRepairMessage.id + logger.warn('Inserted failed tool_result for interrupted tool call') } } + + // Save message to database + const savedMessage = await chatMessageRepository.create(chatId, userMessage, parentId) + + logger.info('Message added successfully') + + return json( + { + messageId: savedMessage.id, + status: 'created', + }, + { status: 200 }, + ) + } catch (error) { + logger.error('Error adding message', error) + return json( + { + error: 'Failed to add message', + details: error instanceof Error ? error.message : 'Unknown error', + }, + { status: 500 }, + ) } - return { ok: true } } diff --git a/web/src/routes/api/chat/[chatId]/messages/[messageId]/edit/+server.ts b/web/src/routes/api/chat/[chatId]/messages/[messageId]/edit/+server.ts index ba64d84ed..af44a1396 100644 --- a/web/src/routes/api/chat/[chatId]/messages/[messageId]/edit/+server.ts +++ b/web/src/routes/api/chat/[chatId]/messages/[messageId]/edit/+server.ts @@ -1,123 +1,40 @@ import { json } from '@sveltejs/kit' import type { RequestHandler } from './$types.js' import { chatRepository, chatMessageRepository } from '$lib/server/db/chats' -import { getAgent } from '$lib/server/db/agents.js' import { getChatStreamStatus } from '$lib/server/ai-stream-status.js' -import type { OmniMentionBlock, OmniUploadBlock } from '$lib/types/message' -import type { MessageParam, TextBlockParam } from '@anthropic-ai/sdk/resources/messages' -import { z } from 'zod' -const editRequestSchema = z - .object({ - content: z.string().max(100_000), - }) - .strict() - -type UserMessageBlock = OmniMentionBlock | OmniUploadBlock | TextBlockParam - -type UnknownSource = { - type?: unknown - upload_id?: unknown - document_id?: unknown - title?: unknown - source_type?: unknown - content_type?: unknown -} - -type UnknownBlock = { - type?: unknown - source?: UnknownSource -} - -function parseStoredRichBlock(block: unknown): OmniMentionBlock | OmniUploadBlock | null { - if (!block || typeof block !== 'object') return null - - const candidate = block as UnknownBlock - const source = candidate.source - if (!source || typeof source !== 'object') return null - - if ( - candidate.type === 'document' && - source.type === 'omni_mention' && - typeof source.document_id === 'string' && - typeof source.title === 'string' && - (source.source_type === undefined || typeof source.source_type === 'string') && - (source.content_type === undefined || typeof source.content_type === 'string') - ) { - return { - type: 'document', - source: { - type: 'omni_mention', - document_id: source.document_id, - title: source.title, - ...(source.source_type ? { source_type: source.source_type } : {}), - ...(source.content_type ? { content_type: source.content_type } : {}), - }, - } - } - - if ( - (candidate.type === 'document' || candidate.type === 'image') && - source.type === 'omni_upload' && - typeof source.upload_id === 'string' - ) { - return { - type: candidate.type, - source: { type: 'omni_upload', upload_id: source.upload_id }, - } - } - - return null +interface EditRequest { + content: string } export const POST: RequestHandler = async ({ params, request, locals }) => { const logger = locals.logger.child('chat') - const { chatId, messageId } = params + const { chatId, messageId } = params if (!chatId || !messageId) { return json({ error: 'chatId and messageId are required' }, { status: 400 }) } + if (!locals.user?.id) { return json({ error: 'User not authenticated' }, { status: 401 }) } - let rawBody: unknown + let editRequest: EditRequest try { - rawBody = await request.json() + editRequest = await request.json() } catch { return json({ error: 'Invalid JSON in request body' }, { status: 400 }) } - const parsed = editRequestSchema.safeParse(rawBody) - if (!parsed.success) { - const details = parsed.error.issues.map( - (issue) => `${issue.path.join('.')}: ${issue.message}`, - ) - return json({ error: 'Invalid request', details }, { status: 400 }) + if (!editRequest.content || editRequest.content.trim() === '') { + return json({ error: 'Content is required' }, { status: 400 }) } - const trimmedText = parsed.data.content.trim() - try { const chat = await chatRepository.get(chatId) if (!chat) { return json({ error: 'Chat not found' }, { status: 404 }) } - if (chat.userId !== locals.user.id) { - return json({ error: 'Forbidden' }, { status: 403 }) - } - if (chat.agentId) { - const agent = await getAgent(chat.agentId) - if (!agent) { - return json({ error: 'Chat agent not found' }, { status: 404 }) - } - if (agent.agentType === 'org' && locals.user.role !== 'admin') { - return json({ error: 'Forbidden' }, { status: 403 }) - } - if (agent.agentType === 'user' && agent.userId !== locals.user.id) { - return json({ error: 'Forbidden' }, { status: 403 }) - } - } try { const streamStatus = await getChatStreamStatus(chatId) @@ -131,59 +48,39 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { ) } } catch (error) { - logger.warn('Could not check stream status before editing message', { - chatId, - messageId, - error, - }) + logger.warn('Could not check stream status before editing message', error) } - const originalMessage = await chatMessageRepository.getByIdInChat(chatId, messageId) + // Get the original message to find its parent + const allMessages = await chatMessageRepository.getByChatId(chatId) + const originalMessage = allMessages.find((m) => m.id === messageId) if (!originalMessage) { return json({ error: 'Message not found' }, { status: 404 }) } - if (originalMessage.message.role !== 'user') { - return json({ error: 'Only user messages can be edited' }, { status: 400 }) - } - - const richBlocks = Array.isArray(originalMessage.message.content) - ? originalMessage.message.content - .map(parseStoredRichBlock) - .filter((block): block is OmniMentionBlock | OmniUploadBlock => block !== null) - : [] - if (!trimmedText && richBlocks.length === 0) { - return json({ error: 'Content or rich blocks are required' }, { status: 400 }) + // Create new message as a sibling of the original (same parent) + const userMessage = { + role: 'user' as const, + content: editRequest.content.trim(), } - const blocks: UserMessageBlock[] = [...richBlocks] - if (trimmedText) { - blocks.push({ type: 'text', text: trimmedText }) - } - - const userMessage = { role: 'user' as const, content: blocks } const savedMessage = await chatMessageRepository.create( chatId, - userMessage as unknown as MessageParam, + userMessage, originalMessage.parentId ?? undefined, ) - logger.info('Message edited (new branch created)', { - chatId, - originalMessageId: messageId, - newMessageId: savedMessage.id, - }) + logger.info('Message edited (new branch created)') return json( { messageId: savedMessage.id, - message: userMessage, status: 'created', }, { status: 200 }, ) } catch (error) { - logger.error('Error editing message', { chatId, messageId, error }) + logger.error('Error editing message', error) return json( { error: 'Failed to edit message', diff --git a/web/src/routes/api/chat/[chatId]/messages/[messageId]/feedback/+server.ts b/web/src/routes/api/chat/[chatId]/messages/[messageId]/feedback/+server.ts index dcb0c23e4..d563adc2d 100644 --- a/web/src/routes/api/chat/[chatId]/messages/[messageId]/feedback/+server.ts +++ b/web/src/routes/api/chat/[chatId]/messages/[messageId]/feedback/+server.ts @@ -42,10 +42,7 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { } logger.debug('Submitting feedback', { - chatId, - messageId, feedbackType: feedbackRequest.feedbackType, - userId: locals.user.id, }) try { @@ -57,11 +54,7 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { ) logger.info('Feedback submitted successfully', { - chatId, - messageId, - feedbackId: feedback.id, feedbackType: feedback.feedbackType, - userId: locals.user.id, }) return json( @@ -73,7 +66,7 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { { status: 200 }, ) } catch (error) { - logger.error('Error submitting feedback', error, { chatId, messageId }) + logger.error('Error submitting feedback', error) return json( { error: 'Failed to submit feedback', @@ -98,29 +91,17 @@ export const DELETE: RequestHandler = async ({ params, locals }) => { return json({ error: 'User not authenticated' }, { status: 401 }) } - logger.debug('Deleting feedback', { - chatId, - messageId, - userId: locals.user.id, - }) + logger.debug('Deleting feedback') try { const deleted = await responseFeedbackRepository.delete(messageId, locals.user.id) if (!deleted) { - logger.warn('No feedback found to delete', { - chatId, - messageId, - userId: locals.user.id, - }) + logger.warn('No feedback found to delete') return json({ error: 'No feedback found' }, { status: 404 }) } - logger.info('Feedback deleted successfully', { - chatId, - messageId, - userId: locals.user.id, - }) + logger.info('Feedback deleted successfully') return json( { @@ -129,7 +110,7 @@ export const DELETE: RequestHandler = async ({ params, locals }) => { { status: 200 }, ) } catch (error) { - logger.error('Error deleting feedback', error, { chatId, messageId }) + logger.error('Error deleting feedback', error) return json( { error: 'Failed to delete feedback', diff --git a/web/src/routes/api/chat/[chatId]/stop/+server.ts b/web/src/routes/api/chat/[chatId]/stop/+server.ts index dfbd8bbee..f3de6b897 100644 --- a/web/src/routes/api/chat/[chatId]/stop/+server.ts +++ b/web/src/routes/api/chat/[chatId]/stop/+server.ts @@ -37,12 +37,11 @@ export const POST: RequestHandler = async ({ params, locals }) => { }) if (!response.ok) { logger.warn('AI service cancel returned non-OK', { - chatId, status: response.status, }) } } catch (err) { - logger.error('Failed to cancel chat stream', err, { chatId }) + logger.error('Failed to cancel chat stream', err) } // Best-effort: report success even if the AI service is unreachable. diff --git a/web/src/routes/api/chat/[chatId]/stream/+server.ts b/web/src/routes/api/chat/[chatId]/stream/+server.ts index 8d81873ca..28db060b0 100644 --- a/web/src/routes/api/chat/[chatId]/stream/+server.ts +++ b/web/src/routes/api/chat/[chatId]/stream/+server.ts @@ -190,11 +190,11 @@ async function triggerTitleGeneration( // First check if title already exists const chat = await chatRepository.get(chatId) if (chat?.title) { - logger.debug('Chat already has a title, skipping title generation', { chatId }) + logger.debug('Chat already has a title, skipping title generation') return { status: 'skipped' } } - logger.info('Triggering title generation', { chatId }) + logger.info('Triggering title generation') const response = await fetch(`${env.AI_SERVICE_URL}/chat/${chatId}/generate_title`, { method: 'POST', @@ -206,8 +206,6 @@ async function triggerTitleGeneration( if (response.ok) { const result = (await response.json()) as TitleGenerationResponse logger.info('Title generation completed', { - chatId, - title: result.title, status: result.status, reason: result.reason, }) @@ -221,58 +219,47 @@ async function triggerTitleGeneration( } else { const message = await aiErrorMessage(response) logger.warn('Title generation failed', { - chatId, status: response.status, - message, }) return { status: 'failed', message } } } catch (error) { - logger.warn('Error during title generation', { error, chatId }) + logger.warn('Error during title generation', { error }) const message = error instanceof Error ? error.message : 'Failed to generate chat title' return { status: 'failed', message } } } export const GET: RequestHandler = async ({ params, locals, cookies, request, url }) => { - if (!locals.user?.id) { - return json({ error: 'User not authenticated' }, { status: 401 }) + const replayPath = replayStreamFixturePath(cookies) + if (replayPath) { + const sampleStream = await readFile(replayPath, 'utf-8') + return replayStreamResponse(sampleStream, params.chatId) } + const logger = locals.logger.child('chat') + const chatId = params.chatId if (!chatId) { + logger.warn('Missing chatId parameter in stream request') return json({ error: 'chatId parameter is required' }, { status: 400 }) } const chat = await chatRepository.get(chatId) if (!chat) { + logger.error('Chat not found') return json({ error: 'Chat not found' }, { status: 404 }) } - if (chat.userId !== locals.user.id) { - return json({ error: 'Forbidden' }, { status: 403 }) - } + + // Agent chats require admin access if (chat.agentId) { const agent = await getAgent(chat.agentId) - if (!agent) { - return json({ error: 'Chat agent not found' }, { status: 404 }) - } - if (agent.agentType === 'org' && locals.user.role !== 'admin') { - return json({ error: 'Admin access required' }, { status: 403 }) - } - if (agent.agentType === 'user' && agent.userId !== locals.user.id) { - return json({ error: 'Forbidden' }, { status: 403 }) + if (agent?.agentType === 'org' && locals.user?.role !== 'admin') { + throw error(403, 'Admin access required for agent chats') } } - const replayPath = replayStreamFixturePath(cookies) - if (replayPath) { - const sampleStream = await readFile(replayPath, 'utf-8') - return replayStreamResponse(sampleStream, params.chatId) - } - - const logger = locals.logger.child('chat') - - logger.debug('Sending GET request to AI service to receive the streaming response', { chatId }) + logger.debug('Sending GET request to AI service to receive the streaming response') const abortController = new AbortController() @@ -296,13 +283,11 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur if (!response.ok) { logger.error('AI service error', undefined, { status: response.status, - statusText: response.statusText, - chatId, }) return sseErrorResponse(await aiErrorMessage(response)) } - logger.info('Chat stream started successfully', { chatId }) + logger.info('Chat stream started successfully') // Create a transformed stream that enriches or redacts selected events // before forwarding them to the browser. The AI service's @@ -329,13 +314,11 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur async start(controller) { try { if (!chat.title) { - logger.info('Generating title for chat', { chatId }) + logger.info('Generating title for chat') triggerTitleGeneration(chatId, logger) .then((result) => { if (result.status === 'generated') { - logger.info( - `Generated title for chat ${chatId}: ${result.title}`, - ) + logger.info('Generated title for chat') try { const titleEvent: TitleEvent = { title: result.title } controller.enqueue( @@ -345,15 +328,10 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur // The browser may have disconnected while title generation ran. } } else if (result.status === 'failed') { - logger.warn('Title generation failed', { - chatId, - message: result.message, - }) + logger.warn('Title generation failed') } }) - .catch((err) => - logger.error(`Failed to generate title for chat ${chatId}`, err), - ) + .catch(() => logger.error('Failed to generate title for chat')) } while (true) { @@ -428,9 +406,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur const enrichedEvent = `${idPrefix}event: oauth_required\ndata: ${JSON.stringify(enriched)}\n\n` controller.enqueue(encoder.encode(enrichedEvent)) } catch (err) { - logger.error('Failed to enrich oauth_required event', err, { - chatId, - }) + logger.error('Failed to enrich oauth_required event', err) // Fall back to forwarding the raw event so the // client at least sees something actionable. const fallback = `${idPrefix}event: oauth_required\ndata: ${data}\n\n` @@ -519,7 +495,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur } } } catch (error) { - logger.error('Error in stream processing', error, { chatId }) + logger.error('Error in stream processing', error) const message = error instanceof Error ? error.message : 'Failed to process chat stream' try { @@ -539,7 +515,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur // by aborting our upstream read, but do NOT touch the run — it // continues server-side so the client can reconnect and resume. // Generation is ended only via the explicit Stop endpoint. - logger.info('Client disconnected from stream proxy', { chatId }) + logger.info('Client disconnected from stream proxy') try { await reader.cancel() } catch { @@ -559,7 +535,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur }, }) } catch (error) { - logger.error('Error calling AI service', error, { chatId }) + logger.error('Error calling AI service', error) const message = error instanceof Error ? error.message : 'Failed to process request' return sseErrorResponse(message) } diff --git a/web/src/routes/api/oauth/callback/+server.ts b/web/src/routes/api/oauth/callback/+server.ts index ea102c1be..a68b0961d 100644 --- a/web/src/routes/api/oauth/callback/+server.ts +++ b/web/src/routes/api/oauth/callback/+server.ts @@ -45,7 +45,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { const oauthError = url.searchParams.get('error') if (oauthError) { - logger.error('OAuth provider error', { error: oauthError }) + logger.error('OAuth provider error') throw redirect(302, '/settings/integrations?error=oauth_denied') } if (!code || !stateToken) { @@ -59,7 +59,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { failureReturnTo = returnToFromStateMetadata(pendingState.metadata) } } catch (err) { - logger.warn('Failed to read OAuth state for failure redirect', { err: String(err) }) + logger.warn('Failed to read OAuth state for failure redirect') } let exchange @@ -70,7 +70,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { }, }) } catch (err) { - logger.error('OAuth exchange failed', { err: String(err) }) + logger.error('OAuth exchange failed') throw redirect( 302, withErrorParam(failureReturnTo ?? '/settings/integrations', 'oauth_failed'), @@ -132,16 +132,11 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { }) if (!resp.ok) { logger.warn('OAuth credential-ready notification failed', { - sourceId, status: resp.status, - body: await resp.text(), }) } } catch (err) { - logger.warn('OAuth credential-ready notification failed', { - sourceId, - err: String(err), - }) + logger.warn('OAuth credential-ready notification failed') } } @@ -178,7 +173,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { body: JSON.stringify({ sync_mode: 'full' }), }) } catch (syncError) { - logger.warn('Failed to trigger post-OAuth sync', { sourceId: flow.sourceId, syncError }) + logger.warn('Failed to trigger post-OAuth sync') } throw redirect(302, flow.returnTo ?? '/admin/settings/integrations?success=connected') @@ -285,7 +280,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { expiresAt, }) - logger.info(`Created personal source ${newSource.id} (${sourceType}) for user ${user.id}`) + logger.info('Created personal source') } throw redirect(302, flow.returnTo ?? '/settings/integrations?success=connected') diff --git a/web/src/routes/api/search/+server.ts b/web/src/routes/api/search/+server.ts index 516d9e9ce..6191ea98a 100644 --- a/web/src/routes/api/search/+server.ts +++ b/web/src/routes/api/search/+server.ts @@ -32,7 +32,7 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { } logger.debug('Sending search request to searcher service', { - query: queryData.query, + queryLength: queryData.query.length, mode: queryData.mode, }) @@ -48,8 +48,6 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { if (!response.ok) { logger.error('Search service error', undefined, { status: response.status, - statusText: response.statusText, - query: queryData.query, }) return json( { @@ -62,13 +60,12 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { const searchResults = await response.json() logger.info('Search completed successfully', { - query: queryData.query, resultsCount: searchResults.results?.length || 0, }) return json(searchResults) } catch (error) { - logger.error('Error calling search service', error, { query: queryData.query }) + logger.error('Error calling search service', error) return json( { error: 'Failed to perform search', diff --git a/web/src/routes/api/sources/[sourceId]/+server.ts b/web/src/routes/api/sources/[sourceId]/+server.ts index 30f1d9fbb..3c61370ac 100644 --- a/web/src/routes/api/sources/[sourceId]/+server.ts +++ b/web/src/routes/api/sources/[sourceId]/+server.ts @@ -94,7 +94,7 @@ export const DELETE: RequestHandler = async ({ params, locals, fetch }) => { method: 'POST', }) } catch (err) { - logger.warn(`Failed to cancel sync ${sync.id} for source ${sourceId}`, err) + logger.warn('Failed to cancel sync', err) } } diff --git a/web/src/routes/api/sources/[sourceId]/action/+server.ts b/web/src/routes/api/sources/[sourceId]/action/+server.ts index b02b99b9e..1f6af81d4 100644 --- a/web/src/routes/api/sources/[sourceId]/action/+server.ts +++ b/web/src/routes/api/sources/[sourceId]/action/+server.ts @@ -55,8 +55,7 @@ export const POST: RequestHandler = async ({ params, locals, request }) => { } catch { errorMessage = (await response.text()) || errorMessage } - logger.error(`Action ${action} failed for source ${sourceId}`, { - error: errorMessage, + logger.error('Action failed', { status: response.status, }) throw error(response.status, errorMessage) diff --git a/web/src/routes/api/sources/[sourceId]/sync/+server.ts b/web/src/routes/api/sources/[sourceId]/sync/+server.ts index 551b701db..ec8c866fa 100644 --- a/web/src/routes/api/sources/[sourceId]/sync/+server.ts +++ b/web/src/routes/api/sources/[sourceId]/sync/+server.ts @@ -59,7 +59,7 @@ export const POST: RequestHandler = async ({ params, request, fetch }) => { // If response isn't JSON, use the text errorMessage = (await syncResponse.text()) || errorMessage } - logger.error(`Sync failed for source ${sourceId}`, { + logger.error('Sync failed', { error: errorMessage, status: syncResponse.status, syncMode: mode, diff --git a/web/src/routes/api/uploads/+server.ts b/web/src/routes/api/uploads/+server.ts index e1d5efa79..150e60b75 100644 --- a/web/src/routes/api/uploads/+server.ts +++ b/web/src/routes/api/uploads/+server.ts @@ -26,15 +26,13 @@ export const POST: RequestHandler = async ({ request, locals }) => { body: forwarded, }) } catch (err) { - logger.error('Upload proxy fetch failed', err as Error, { - aiServiceUrl: env.AI_SERVICE_URL, - }) + logger.error('Upload proxy fetch failed', err as Error) return json({ error: 'Upload service unavailable' }, { status: 503 }) } const body = await resp.text() if (!resp.ok) { - logger.warn('Upload proxy failed', undefined, { status: resp.status, body }) + logger.warn('Upload proxy failed', undefined, { status: resp.status }) return new Response(body, { status: resp.status, headers: { 'Content-Type': resp.headers.get('Content-Type') ?? 'application/json' }, diff --git a/web/src/routes/api/user/settings/+server.ts b/web/src/routes/api/user/settings/+server.ts index 7a84dbfd6..4e3359771 100644 --- a/web/src/routes/api/user/settings/+server.ts +++ b/web/src/routes/api/user/settings/+server.ts @@ -42,9 +42,7 @@ export const POST: RequestHandler = async ({ request, locals }) => { locals.user.configuration.timezone = savedTimezone savedSettings.timezone = savedTimezone } catch (error) { - locals.logger.error('Failed to save user timezone setting', error as Error, { - userId: locals.user.id, - }) + locals.logger.error('Failed to save user timezone setting', error as Error) return json({ error: 'Failed to save user settings' } satisfies UserSettingsResponse, { status: 500, }) diff --git a/web/src/routes/api/v1/documents/[id]/+server.ts b/web/src/routes/api/v1/documents/[id]/+server.ts index 18403b16f..4d56bd26c 100644 --- a/web/src/routes/api/v1/documents/[id]/+server.ts +++ b/web/src/routes/api/v1/documents/[id]/+server.ts @@ -22,10 +22,8 @@ export const GET: RequestHandler = async ({ params, url, fetch, locals }) => { const queryData: Record = { query: 'content', document_id: documentId, - user_email: - locals.apiKeyScope === 'admin' ? undefined : locals.user.email, - user_id: - locals.apiKeyScope === 'admin' ? undefined : locals.user.id, + user_email: locals.apiKeyScope === 'admin' ? undefined : locals.user.email, + user_id: locals.apiKeyScope === 'admin' ? undefined : locals.user.id, limit: 1, } @@ -42,7 +40,7 @@ export const GET: RequestHandler = async ({ params, url, fetch, locals }) => { } } - logger.debug('Document content request', { documentId }) + logger.debug('Document content request') try { const response = await fetch(`${env.SEARCHER_URL}/search`, { @@ -58,7 +56,6 @@ export const GET: RequestHandler = async ({ params, url, fetch, locals }) => { } logger.error('Searcher service error', undefined, { status: response.status, - documentId, }) return json({ error: 'Service unavailable' }, { status: 502 }) } @@ -107,7 +104,7 @@ export const GET: RequestHandler = async ({ params, url, fetch, locals }) => { updated_at: doc.updated_at, }) } catch (error) { - logger.error('Document content request failed', error as Error, { documentId }) + logger.error('Document content request failed', error as Error) return json({ error: 'Failed to fetch document' }, { status: 500 }) } } diff --git a/web/src/routes/api/v1/search/+server.ts b/web/src/routes/api/v1/search/+server.ts index 3bdd958dc..3b6c1213a 100644 --- a/web/src/routes/api/v1/search/+server.ts +++ b/web/src/routes/api/v1/search/+server.ts @@ -95,7 +95,10 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { user_configuration: locals.apiKeyScope === 'admin' ? undefined : locals.user.configuration, } - logger.debug('Agent search request', { query: queryData.query, mode: queryData.mode }) + logger.debug('Agent search request', { + queryLength: queryData.query?.length ?? 0, + mode: queryData.mode, + }) try { const response = await fetch(`${env.SEARCHER_URL}/search`, { @@ -107,7 +110,6 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { if (!response.ok) { logger.error('Searcher service error', undefined, { status: response.status, - query: queryData.query, }) return json( { error: 'Search service unavailable', status: response.status }, @@ -118,7 +120,6 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { const results = await response.json() logger.info('Agent search completed', { - query: queryData.query, results_count: results.results?.length ?? 0, total_count: results.total_count, query_time_ms: results.query_time_ms, @@ -126,7 +127,7 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { return json(results) } catch (error) { - logger.error('Search request failed', error as Error, { query: queryData.query }) + logger.error('Search request failed', error as Error) return json({ error: 'Search request failed' }, { status: 500 }) } } From 271dbcba1a2faca3d7b8a319df25f0c49b29bf1a Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Mon, 27 Jul 2026 12:03:37 +0000 Subject: [PATCH 02/14] fix: remove queue trace context persistence from Postgres - Delete migration 106 that added traceparent/tracestate columns - Remove traceparent/tracestate fields from Rust queue event structs and SQL queries - Remove PRODUCER span wrapping and W3C carrier injection from enqueue operations - Remove consumer trace link creation from indexer queue processor - Remove trace context extraction and consumer link building from AI batch processor - Delete test files for removed queue tracing functionality - Remove the queue module from telemetry.rs Async queue consumers now create fresh new-root traces without explicit OTel links, which is the standard pattern for async boundaries. Business correlation IDs in the events provide traceability when needed. --- services/ai/db/embedding_queue.py | 8 +- services/ai/embeddings/batch_processor.py | 96 +--- services/ai/tests/unit/test_queue_tracing.py | 538 ------------------ services/indexer/src/queue_processor.rs | 42 +- .../106_add_queue_trace_context.sql | 18 - shared/src/embedding_queue.rs | 248 +++----- shared/src/queue.rs | 118 +--- shared/src/telemetry.rs | 387 ------------- shared/tests/embedding_queue_test.rs | 187 ------ shared/tests/event_queue_test.rs | 120 ---- shared/tests/queue_link_test.rs | 357 ------------ 11 files changed, 123 insertions(+), 1996 deletions(-) delete mode 100644 services/ai/tests/unit/test_queue_tracing.py delete mode 100644 services/migrations/106_add_queue_trace_context.sql delete mode 100644 shared/tests/queue_link_test.rs diff --git a/services/ai/db/embedding_queue.py b/services/ai/db/embedding_queue.py index fd4efcaf8..67be418b5 100644 --- a/services/ai/db/embedding_queue.py +++ b/services/ai/db/embedding_queue.py @@ -30,8 +30,6 @@ class EmbeddingQueueItem: error_message: Optional[str] retry_count: int created_at: datetime - traceparent: Optional[str] = field(default=None) - tracestate: Optional[str] = field(default=None) class EmbeddingQueueRepository: @@ -52,8 +50,7 @@ async def get_by_id(self, item_id: str) -> Optional[EmbeddingQueueItem]: row = await pool.fetchrow( """ - SELECT id, document_id, status, error_message, retry_count, created_at, - traceparent, tracestate + SELECT id, document_id, status, error_message, retry_count, created_at FROM embedding_queue WHERE id = $1 """, @@ -110,8 +107,7 @@ async def get_pending_items( LIMIT $2 FOR UPDATE SKIP LOCKED ) - RETURNING id, document_id, status, error_message, retry_count, created_at, - traceparent, tracestate + RETURNING id, document_id, status, error_message, retry_count, created_at """, max_retries, limit, diff --git a/services/ai/embeddings/batch_processor.py b/services/ai/embeddings/batch_processor.py index d965850ca..fbe4241e6 100644 --- a/services/ai/embeddings/batch_processor.py +++ b/services/ai/embeddings/batch_processor.py @@ -8,14 +8,10 @@ import asyncio import logging import time -from typing import Optional, Sequence +from typing import Optional import ulid from opentelemetry import trace, metrics -from opentelemetry.propagate import get_global_textmap -from opentelemetry.propagators.textmap import TextMapPropagator -from opentelemetry.trace import SpanContext, NonRecordingSpan, Link, SpanKind -from opentelemetry.trace.propagation import get_current_span as otel_get_current_span from config import EMBEDDING_MAX_MODEL_LEN from db import ( @@ -34,79 +30,6 @@ logger = logging.getLogger(__name__) -# --------------------------------------------------------------------------- -# Trace context helpers for the embedding consumer -# --------------------------------------------------------------------------- - - -def _build_carrier( - traceparent: Optional[str], - tracestate: Optional[str], -) -> dict[str, str]: - """Build a W3C TraceContext carrier dict from stored optional headers.""" - carrier: dict[str, str] = {} - if traceparent: - # Validate rough W3C format before using - if len(traceparent) == 55 and traceparent.startswith("00-"): - carrier["traceparent"] = traceparent - if tracestate: - carrier["tracestate"] = tracestate - return carrier - - -def _extract_span_context( - traceparent: Optional[str], - tracestate: Optional[str], -) -> Optional[SpanContext]: - """Extract a valid remote SpanContext from optional stored headers. - - Returns None for invalid or missing context. Processing continues - regardless. - """ - carrier = _build_carrier(traceparent, tracestate) - if not carrier: - return None - - from opentelemetry.context import Context - # Get the global propagator (typically TraceContextTextMapPropagator) - propagator = get_global_textmap() - # Extract into a fresh context (not the current one) - ctx = propagator.extract(carrier=carrier, context=Context()) - # Get the span from the extracted context - span = trace.get_current_span(ctx) - sc = span.get_span_context() - if sc and sc.is_valid: - return sc - return None - - -def _build_consumer_links( - items: Sequence[EmbeddingQueueItem], -) -> list[Link]: - """Build OTel Links from stored producer contexts in queue items. - - Returns one Link per unique valid producer (trace_id, span_id) pair. - Invalid or missing contexts are skipped. An empty list is returned - when no valid contexts exist. - """ - links: list[Link] = [] - seen: set[tuple[int, int]] = set() - - for item in items: - sc = _extract_span_context(item.traceparent, item.tracestate) - if sc is None: - continue - trace_id_int = int(sc.trace_id, 16) if isinstance(sc.trace_id, str) else sc.trace_id - span_id_int = int(sc.span_id, 16) if isinstance(sc.span_id, str) else sc.span_id - key = (trace_id_int, span_id_int) - if key in seen: - continue - seen.add(key) - links.append(Link(context=sc)) - - return links - - # --------------------------------------------------------------------------- # Embedding processor metrics (lazily created once) # --------------------------------------------------------------------------- @@ -266,27 +189,10 @@ async def _process_online_batch(self) -> bool: logger.info(f"Processing {len(items)} documents via online embedding API") batch_start = time.monotonic() - # Build consumer links from stored producer contexts. - links = _build_consumer_links(items) - - # Import explicit empty Context for new-root consumer span. - from opentelemetry.context import Context - - # Create CONSUMER span as a new root trace with links to producers. - # The span ends when the `with` block exits. tracer = trace.get_tracer("omni-ai") had_failure = False with tracer.start_as_current_span( "embedding_queue process", - context=Context(), # explicit empty parent = new root trace - kind=SpanKind.CONSUMER, - attributes={ - "messaging.system": "postgresql", - "messaging.destination": "embedding_queue", - "messaging.operation.type": "process", - "messaging.batch.message_count": len(items), - }, - links=links, ) as span: try: documents_by_id = await self.documents_repo.get_by_ids( diff --git a/services/ai/tests/unit/test_queue_tracing.py b/services/ai/tests/unit/test_queue_tracing.py deleted file mode 100644 index 8bd64abef..000000000 --- a/services/ai/tests/unit/test_queue_tracing.py +++ /dev/null @@ -1,538 +0,0 @@ -"""Queue tracing tests for the embedding batch processor. - -Uses InMemorySpanExporter to verify that CONSUMER spans carry native links -to their PRODUCER contexts, that failure scenarios set ERROR status, and -that invalid/missing contexts are ignored. -""" - -from datetime import datetime, timezone -from unittest.mock import AsyncMock, MagicMock - -import pytest -from opentelemetry import trace -from opentelemetry.context import Context -from opentelemetry.propagate import set_global_textmap -from opentelemetry.sdk.trace import TracerProvider -from opentelemetry.sdk.trace.export import SimpleSpanProcessor -from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter -from opentelemetry.trace import SpanKind, StatusCode -from opentelemetry.trace.propagation.tracecontext import TraceContextTextMapPropagator - -from db.embedding_queue import EmbeddingQueueItem -from embeddings.batch_processor import ( - _build_carrier, - _extract_span_context, - _build_consumer_links, - EmbeddingBatchProcessor, -) - - -# --------------------------------------------------------------------------- -# Fixtures -# --------------------------------------------------------------------------- - - -# --------------------------------------------------------------------------- -# Module-level shared tracer provider and exporter. -# Initialised at import time so `trace.set_tracer_provider` is called BEFORE -# any other test file can lock it (the Python OTel API only allows one set). -# --------------------------------------------------------------------------- -set_global_textmap(TraceContextTextMapPropagator()) -_SHARED_EXPORTER = InMemorySpanExporter() -_SHARED_PROVIDER = TracerProvider() -_SHARED_PROVIDER.add_span_processor(SimpleSpanProcessor(_SHARED_EXPORTER)) -trace.set_tracer_provider(_SHARED_PROVIDER) - - -@pytest.fixture(autouse=True) -def _reset_tracing(): - """Reset batcher instrument globals and exporter before each test.""" - _SHARED_EXPORTER.clear() - # Reset the batcher's instrument globals so each test starts clean - import embeddings.batch_processor as bp - - bp._EMBEDDING_METER = None - bp._EMBEDDING_PENDING = None - bp._EMBEDDING_PROCESSED = None - bp._EMBEDDING_FAILED = None - bp._EMBEDDING_BATCH_DURATION = None - yield - - -@pytest.fixture -def span_exporter(): - """Return the shared InMemorySpanExporter.""" - return _SHARED_EXPORTER - - -def _make_processor(span_exporter) -> EmbeddingBatchProcessor: - """Build a processor with all repos mocked.""" - docs_repo = AsyncMock() - queue_repo = AsyncMock() - embeddings_repo = AsyncMock() - app_state = MagicMock() - # Use MagicMock for the embedding provider (sync methods like get_model_name) - # with specific AsyncMock for async methods like generate_embeddings. - provider = MagicMock() - provider.get_model_name.return_value = "test-model" - provider.generate_embeddings = AsyncMock( - return_value=[MagicMock(span=(0, 4), embedding=[0.1])] - ) - app_state.embedding_provider = provider - app_state.embedding_provider_type = "test" - app_state.content_storage = AsyncMock() - - # Default find_embedded_content_donors to empty dict so that - # _clone_same_content_embeddings always short-circuits and never - # produces unawaited AsyncMock coroutine warnings. - docs_repo.find_embedded_content_donors = AsyncMock(return_value={}) - - processor = EmbeddingBatchProcessor( - documents_repo=docs_repo, - queue_repo=queue_repo, - embeddings_repo=embeddings_repo, - app_state=app_state, - ) - processor._baseline_completed = 0 - processor._baseline_failed = 0 - return processor - - -def _make_item( - document_id: str, - traceparent: str | None = None, - tracestate: str | None = None, -) -> EmbeddingQueueItem: - return EmbeddingQueueItem( - id=f"item-{document_id}", - document_id=document_id, - status="processing", - error_message=None, - retry_count=0, - created_at=datetime.now(timezone.utc), - traceparent=traceparent, - tracestate=tracestate, - ) - - -def _find_span(spans, name: str, kind: SpanKind): - """Find a finished span by name and kind.""" - for s in spans: - if s.name == name and s.kind == kind: - return s - return None - - -def _span_context_from_traceparent(tp: str): - """Extract a SpanContext from a traceparent for link assertions.""" - sc = _extract_span_context(tp, None) - assert sc is not None, f"could not extract from {tp}" - return sc - - -def _flush_exporter(): - """Force-flush the global tracer provider if available.""" - tracer_provider = trace.get_tracer_provider() - if hasattr(tracer_provider, "force_flush"): - tracer_provider.force_flush() - - -# --------------------------------------------------------------------------- -# Pure helpers -# --------------------------------------------------------------------------- - - -class TestBuildCarrier: - def test_valid_traceparent(self): - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - carrier = _build_carrier(tp, None) - assert carrier["traceparent"] == tp - assert "tracestate" not in carrier - - def test_with_tracestate(self): - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - carrier = _build_carrier(tp, "vendor1=value1") - assert carrier["traceparent"] == tp - assert carrier["tracestate"] == "vendor1=value1" - - def test_invalid_format_returns_empty(self): - assert _build_carrier("invalid", None) == {} - assert _build_carrier("", None) == {} - assert _build_carrier(None, None) == {} - assert _build_carrier("00-too-short", None) == {} - - -class TestExtractSpanContext: - def test_valid(self): - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - sc = _extract_span_context(tp, None) - assert sc is not None - assert sc.is_valid - - def test_invalid_returns_none(self): - assert _extract_span_context(None, None) is None - assert _extract_span_context("invalid", None) is None - assert _extract_span_context("", None) is None - - def test_all_zero_returns_invalid(self): - tp = "00-00000000000000000000000000000000-0000000000000000-01" - sc = _extract_span_context(tp, None) - assert sc is None or not sc.is_valid - - def test_ambient_span_not_leaked_when_malformed(self): - """Active ambient span must NOT leak into extract with malformed traceparent.""" - tracer = trace.get_tracer("test") - with tracer.start_as_current_span("ambient"): - # Exactly 55 chars, valid format but non-hex trace id - malformed = "00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-b7ad6b7169203331-01" - assert len(malformed) == 55 - sc = _extract_span_context(malformed, None) - assert sc is None, "must not return ambient context for non-hex traceparent" - - def test_ambient_span_not_leaked_when_valid(self): - """Active ambient span must NOT replace the carrier context on valid extract.""" - tracer = trace.get_tracer("test") - with tracer.start_as_current_span("ambient"): - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - sc = _extract_span_context(tp, None) - assert sc is not None - assert sc.is_valid - expected_tid = int("0af7651916cd43dd8448eb211c80319c", 16) - assert sc.trace_id == expected_tid, \ - "trace_id must come from carrier, not ambient span" - - -class TestBuildConsumerLinks: - def test_single_valid_item(self): - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - items = [_make_item("doc1", traceparent=tp)] - links = _build_consumer_links(items) - assert len(links) == 1 - expected_tid = int("0af7651916cd43dd8448eb211c80319c", 16) - assert links[0].context.trace_id == expected_tid - - def test_deduplicates_same_context(self): - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - items = [ - _make_item("doc1", traceparent=tp), - _make_item("doc2", traceparent=tp), - ] - links = _build_consumer_links(items) - assert len(links) == 1 - - def test_same_trace_different_spans_retained(self): - tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - tp2 = "00-0af7651916cd43dd8448eb211c80319c-aaaaaaaaaaaaaaaa-01" - items = [ - _make_item("doc1", traceparent=tp1), - _make_item("doc2", traceparent=tp2), - ] - links = _build_consumer_links(items) - assert len(links) == 2 - - def test_same_trace_three_items_mixed(self): - """Three items: two sharing one span, one unique, one invalid.""" - tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - tp2 = "00-0af7651916cd43dd8448eb211c80319c-aaaaaaaaaaaaaaaa-01" - items = [ - _make_item("doc1", traceparent=tp1), - _make_item("doc2", traceparent=tp2), - _make_item("doc3", traceparent=tp1), # duplicate of doc1 - _make_item("doc4", traceparent=None), # invalid - ] - links = _build_consumer_links(items) - assert len(links) == 2 - - def test_invalid_skipped(self): - items = [ - _make_item("doc1", traceparent=None), - _make_item("doc2", traceparent="invalid"), - ] - links = _build_consumer_links(items) - assert len(links) == 0 - - -# --------------------------------------------------------------------------- -# Batch processor tracing integration -# --------------------------------------------------------------------------- - - -class TestProcessOnlineBatchTracing: - """Trace behaviour of _process_online_batch via in-memory exporter.""" - - # -- Helpers to set up common mocks -------------------------------- - - @staticmethod - def _setup_happy_path(processor, items): - """Configure mocks for a successful processing path.""" - processor.queue_repo.get_pending_items.return_value = items - # Create realistic Document mocks: content_id set, external_id=None to - # avoid triggering cross-source dedup path. - docs = {} - for item in items: - doc = MagicMock() - doc.content_id = f"c-{item.document_id}" - doc.external_id = None - docs[item.document_id] = doc - processor.documents_repo.get_by_ids.return_value = docs - processor.content_storage.get_text = AsyncMock(return_value="some text content") - processor.embeddings_repo.bulk_insert = AsyncMock() - # Ensure _clone_same_content_embeddings returns items as-is - processor.documents_repo.find_embedded_content_donors = AsyncMock(return_value={}) - - @pytest.mark.asyncio - async def test_consumer_span_is_new_root(self, span_exporter): - """CONSUMER span has no parent (new root trace).""" - processor = _make_processor(span_exporter) - - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - self._setup_happy_path(processor, [_make_item("doc1", traceparent=tp)]) - - result = await processor._process_online_batch() - assert result is True - - _flush_exporter() - spans = span_exporter.get_finished_spans() - span_exporter.clear() - - consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) - assert consumer is not None, "CONSUMER span must be exported" - - # Verify parent is None (new root trace via explicit Context()) - assert consumer.parent is None, ( - "CONSUMER span must have no parent (new root)" - ) - - # Verify consumer trace ID differs from producer trace ID - link = consumer.links[0] - assert consumer.context.trace_id != link.context.trace_id, ( - "CONSUMER must have different trace than producer" - ) - - @pytest.mark.asyncio - async def test_native_link_matches_producer(self, span_exporter): - """CONSUMER span has a native link matching the producer traceparent.""" - processor = _make_processor(span_exporter) - - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - self._setup_happy_path(processor, [_make_item("doc1", traceparent=tp)]) - - await processor._process_online_batch() - - _flush_exporter() - spans = span_exporter.get_finished_spans() - span_exporter.clear() - - consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) - assert consumer is not None - assert len(consumer.links) >= 1, "must have at least one link" - - # Verify link matches the expected producer context - expected_sc = _span_context_from_traceparent(tp) - link = consumer.links[0] - assert link.context.trace_id == expected_sc.trace_id, ( - "link trace_id must match producer trace_id" - ) - assert link.context.span_id == expected_sc.span_id, ( - "link span_id must match producer span_id" - ) - assert consumer.parent is None, ( - "CONSUMER must be new root" - ) - - @pytest.mark.asyncio - async def test_same_trace_two_producers_two_links(self, span_exporter): - """Two items sharing same trace but different span IDs create two links.""" - processor = _make_processor(span_exporter) - - tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - tp2 = "00-0af7651916cd43dd8448eb211c80319c-aaaaaaaaaaaaaaaa-01" - item1 = _make_item("doc1", traceparent=tp1) - item2 = _make_item("doc2", traceparent=tp2) - self._setup_happy_path(processor, [item1, item2]) - - await processor._process_online_batch() - - _flush_exporter() - spans = span_exporter.get_finished_spans() - span_exporter.clear() - - consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) - assert consumer is not None - assert len(consumer.links) == 2, "must have 2 links for 2 distinct producers" - - # Both links should have the same trace_id - link_tids = {l.context.trace_id for l in consumer.links} - assert len(link_tids) == 1, "both links must share the same trace_id" - - # But different span_ids - link_sids = {l.context.span_id for l in consumer.links} - assert len(link_sids) == 2, "links must have different span_ids" - - @pytest.mark.asyncio - async def test_invalid_traceparent_ignored(self, span_exporter): - """Items with invalid/missing traceparent do not create links.""" - processor = _make_processor(span_exporter) - - item1 = _make_item("doc1", traceparent=None) - item2 = _make_item("doc2", traceparent="invalid") - self._setup_happy_path(processor, [item1, item2]) - - await processor._process_online_batch() - - _flush_exporter() - spans = span_exporter.get_finished_spans() - span_exporter.clear() - - consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) - if consumer is not None: - assert len(consumer.links) == 0, "no valid producer contexts" - - @pytest.mark.asyncio - async def test_processing_runs_inside_consumer(self, span_exporter): - """Processing operations occur within the CONSUMER span scope.""" - processor = _make_processor(span_exporter) - - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - item = _make_item("doc1", traceparent=tp) - processor.queue_repo.get_pending_items.return_value = [item] - processor.documents_repo.get_by_ids.return_value = { - "doc1": MagicMock(content_id="c1") - } - emb_provider = processor.embedding_provider - emb_provider.get_model_name.return_value = "test-model" - - # Track which span is active during processing - captured_span_context = None - - async def track_span(item, doc): - nonlocal captured_span_context - current_span = trace.get_current_span() - sc = current_span.get_span_context() - captured_span_context = sc - - processor._process_single_document = track_span - - await processor._process_online_batch() - - _flush_exporter() - spans = span_exporter.get_finished_spans() - span_exporter.clear() - - consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) - assert consumer is not None - - # Verify the active span during processing matches the consumer span - assert captured_span_context is not None - assert captured_span_context.trace_id == consumer.context.trace_id, ( - "trace_id must match consumer span" - ) - - @pytest.mark.asyncio - async def test_item_failure_sets_error_status(self, span_exporter): - """When an item fails, the CONSUMER span has ERROR status.""" - processor = _make_processor(span_exporter) - - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - item = _make_item("doc1", traceparent=tp) - processor.queue_repo.get_pending_items.return_value = [item] - processor.documents_repo.get_by_ids.return_value = { - "doc1": MagicMock(content_id="c1") - } - emb_provider = processor.embedding_provider - emb_provider.get_model_name.return_value = "test-model" - - # Make processing fail - async def fail_process(item, doc): - raise RuntimeError("test failure") - - processor._process_single_document = fail_process - processor.queue_repo.mark_failed = AsyncMock() - - await processor._process_online_batch() - - _flush_exporter() - spans = span_exporter.get_finished_spans() - span_exporter.clear() - - consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) - assert consumer is not None - - # Check status is ERROR (StatusCode.ERROR = 2) - assert consumer.status is not None - assert consumer.status.status_code == StatusCode.ERROR, ( - "span must have ERROR status when item fails" - ) - - @pytest.mark.asyncio - async def test_missing_content_id_failure_real_path(self, span_exporter): - """Real _process_single_document: missing content_id returns False and - sets consumer span ERROR status.""" - processor = _make_processor(span_exporter) - - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - item = _make_item("doc1", traceparent=tp) - processor.queue_repo.get_pending_items.return_value = [item] - # Document with no content_id - doc = MagicMock(spec_set=["content_id", "external_id"]) - doc.content_id = None - doc.external_id = None - processor.documents_repo.get_by_ids.return_value = {"doc1": doc} - processor.documents_repo.find_embedded_content_donors = AsyncMock(return_value={}) - - # Use the REAL _process_single_document (no replacement) - result = await processor._process_online_batch() - assert result is True - - _flush_exporter() - spans = span_exporter.get_finished_spans() - span_exporter.clear() - - consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) - assert consumer is not None, "CONSUMER span must be exported" - - # The missing content_id path sets had_failure -> ERROR - assert consumer.status is not None - assert consumer.status.status_code == StatusCode.ERROR, ( - "span must have ERROR status when content_id is missing" - ) - - @pytest.mark.asyncio - async def test_embedding_exception_failure_real_path(self, span_exporter): - """Real _process_single_document: embedding provider exception returns - False and sets consumer span ERROR status.""" - processor = _make_processor(span_exporter) - - tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - item = _make_item("doc1", traceparent=tp) - processor.queue_repo.get_pending_items.return_value = [item] - # Document with content_id so we reach the embedding call - doc = MagicMock(spec_set=["content_id", "external_id"]) - doc.content_id = "c-doc1" - doc.external_id = None - processor.documents_repo.get_by_ids.return_value = {"doc1": doc} - processor.documents_repo.find_embedded_content_donors = AsyncMock(return_value={}) - processor.content_storage.get_text = AsyncMock( - return_value="some text content" - ) - # Make embedding provider raise an exception - provider = processor.embedding_provider - provider.generate_embeddings = AsyncMock(side_effect=RuntimeError("embedding failed")) - provider.get_model_name.return_value = "test-model" - - # Use the REAL _process_single_document (no replacement) - result = await processor._process_online_batch() - assert result is True - - _flush_exporter() - spans = span_exporter.get_finished_spans() - span_exporter.clear() - - consumer = _find_span(spans, "embedding_queue process", SpanKind.CONSUMER) - assert consumer is not None, "CONSUMER span must be exported" - - # The embedding exception path sets had_failure -> ERROR - assert consumer.status is not None - assert consumer.status.status_code == StatusCode.ERROR, ( - "span must have ERROR status when embedding provider raises" - ) diff --git a/services/indexer/src/queue_processor.rs b/services/indexer/src/queue_processor.rs index 82c24e9e2..c0ba7f22c 100644 --- a/services/indexer/src/queue_processor.rs +++ b/services/indexer/src/queue_processor.rs @@ -16,8 +16,7 @@ use std::collections::HashMap; use std::sync::Arc; use tokio::sync::{Mutex, Semaphore}; use tokio::time::{interval, Duration, MissedTickBehavior}; -use tracing::{debug, error, info, warn, Instrument}; -use tracing_opentelemetry::OpenTelemetrySpanExt; +use tracing::{debug, error, info, warn}; // Default poll interval for draining the queue. Overridable via INDEXER_POLL_INTERVAL_SECS. // SDK-side buffering already shapes events into the right batch size per sync type, @@ -627,46 +626,13 @@ impl QueueProcessor { Ok(()) } - /// Instrument `process_dequeued_events` with a CONSUMER tracing span. - /// - /// The span is created as a new root trace (no parent) with bounded messaging - /// attributes and native OTel links pointing at the producer spans. + /// Process dequeued events without consumer trace instrumentation + /// (queue trace context persistence has been removed). async fn process_dequeued_events_instrumented( &self, events: Vec, ) -> Result { - let event_count = events.len() as i64; - - // Extract producer SpanContext values deduplicated by (trace_id, span_id). - let producer_contexts: Vec, Option)>> = events - .iter() - .map(|ev| Some((ev.traceparent.clone(), ev.tracestate.clone()))) - .collect(); - let span_contexts = shared::telemetry::queue::collect_span_contexts(&producer_contexts); - - let span = tracing::info_span!( - parent: None, - "connector_events_queue process", - otel.kind = "CONSUMER", - messaging.system = "postgresql", - messaging.destination = "connector_events_queue", - messaging.operation.type = "process", - messaging.batch.message_count = event_count, - ); - - // Initialize the OTel span and add native links before entering - // the async work. Clone the span so the original handle is available - // for `.instrument()` after adding links. - { - let _guard = span.clone().entered(); - for sc in &span_contexts { - span.add_link(sc.clone()); - } - } - - async move { self.process_dequeued_events(events).await } - .instrument(span) - .await + self.process_dequeued_events(events).await } async fn process_dequeued_events(&self, events: Vec) -> Result { diff --git a/services/migrations/106_add_queue_trace_context.sql b/services/migrations/106_add_queue_trace_context.sql deleted file mode 100644 index beb2f29a3..000000000 --- a/services/migrations/106_add_queue_trace_context.sql +++ /dev/null @@ -1,18 +0,0 @@ --- Migration 106: Add W3C trace context columns to connector_events_queue and embedding_queue. --- --- This migration adds nullable traceparent and tracestate columns to support --- W3C TraceContext propagation through the asynchronous queue pipeline. --- --- Columns: --- traceparent (varchar(55)) — W3C traceparent header (version-format-trace_id-span_id-trace_flags) --- tracestate (varchar(512)) — optional W3C tracestate header (max 512 chars per spec) --- --- Both columns are nullable; NULL means "no stored trace context". --- Invalid/missing values are handled gracefully at read time. --- No indexes: queue consumers read trace context after dequeue, not for filtering. - -ALTER TABLE connector_events_queue ADD COLUMN IF NOT EXISTS traceparent varchar(55); -ALTER TABLE connector_events_queue ADD COLUMN IF NOT EXISTS tracestate varchar(512); - -ALTER TABLE embedding_queue ADD COLUMN IF NOT EXISTS traceparent varchar(55); -ALTER TABLE embedding_queue ADD COLUMN IF NOT EXISTS tracestate varchar(512); diff --git a/shared/src/embedding_queue.rs b/shared/src/embedding_queue.rs index 7f573b8ab..d4b8e3e36 100644 --- a/shared/src/embedding_queue.rs +++ b/shared/src/embedding_queue.rs @@ -1,7 +1,6 @@ use anyhow::Result; use serde::{Deserialize, Serialize}; use sqlx::{PgPool, Row}; -use tracing::Instrument; use ulid::Ulid; use crate::{db::repositories::EmbeddingProviderRepository, utils::generate_ulid}; @@ -50,12 +49,6 @@ pub struct EmbeddingQueueItem { pub created_at: sqlx::types::time::OffsetDateTime, pub updated_at: sqlx::types::time::OffsetDateTime, pub processed_at: Option, - /// W3C traceparent header stored by the producer for trace propagation. - #[serde(skip_serializing_if = "Option::is_none")] - pub traceparent: Option, - /// W3C tracestate header stored alongside traceparent. - #[serde(skip_serializing_if = "Option::is_none")] - pub tracestate: Option, } impl sqlx::FromRow<'_, sqlx::postgres::PgRow> for EmbeddingQueueItem { @@ -80,8 +73,6 @@ impl sqlx::FromRow<'_, sqlx::postgres::PgRow> for EmbeddingQueueItem { created_at: row.try_get("created_at")?, updated_at: row.try_get("updated_at")?, processed_at: row.try_get("processed_at")?, - traceparent: row.try_get("traceparent").ok().flatten(), - tracestate: row.try_get("tracestate").ok().flatten(), }) } } @@ -102,26 +93,49 @@ impl EmbeddingQueue { } pub async fn enqueue(&self, document_id: String) -> Result> { - let queue_name = "embedding_queue"; + if !self.provider_repo.has_active_provider().await? { + return Ok(None); + } - // Build PRODUCER span covering provider check + INSERT. - // inject_active_context is called from inside the instrumented block - // so the producer span is active during injection and the DB write. - let producer_span = crate::telemetry::queue::build_producer_span(queue_name); + let id = Ulid::new().to_string(); - async move { - let carrier = crate::telemetry::queue::inject_active_context(); + let result = sqlx::query( + r#" + INSERT INTO embedding_queue (id, document_id) + SELECT $1, $2 + WHERE NOT EXISTS ( + SELECT 1 FROM embedding_queue + WHERE document_id = $2 AND status IN ('pending', 'processing') + ) + "#, + ) + .bind(&id) + .bind(&document_id) + .execute(&self.pool) + .await?; - if !self.provider_repo.has_active_provider().await? { - return Ok(None); - } + if result.rows_affected() > 0 { + Ok(Some(id)) + } else { + Ok(None) + } + } + pub async fn enqueue_batch(&self, document_ids: Vec) -> Result> { + if !self.provider_repo.has_active_provider().await? { + return Ok(vec![]); + } + + let mut tx = self.pool.begin().await?; + let mut ids = Vec::new(); + + for document_id in document_ids { let id = Ulid::new().to_string(); let result = sqlx::query( r#" - INSERT INTO embedding_queue (id, document_id, traceparent, tracestate) - SELECT $1, $2, $3, $4 + INSERT INTO embedding_queue (id, document_id) + SELECT $1, $2 WHERE NOT EXISTS ( SELECT 1 FROM embedding_queue WHERE document_id = $2 AND status IN ('pending', 'processing') @@ -130,83 +144,16 @@ impl EmbeddingQueue { ) .bind(&id) .bind(&document_id) - .bind(&carrier.traceparent) - .bind(&carrier.tracestate) - .execute(&self.pool) - .await - .map_err(|e| { - crate::telemetry::queue::set_active_span_error(); - e - })?; + .execute(&mut *tx) + .await?; if result.rows_affected() > 0 { - Ok(Some(id)) - } else { - Ok(None) + ids.push(id); } } - .instrument(producer_span) - .await - } - - pub async fn enqueue_batch(&self, document_ids: Vec) -> Result> { - let queue_name = "embedding_queue"; - - // Build PRODUCER span covering provider check + batch INSERT. - // inject_active_context is called from inside the instrumented block - // so the producer span is active during injection and the DB write. - let producer_span = crate::telemetry::queue::build_producer_span(queue_name); - - async move { - let carrier = crate::telemetry::queue::inject_active_context(); - - if !self.provider_repo.has_active_provider().await? { - return Ok(vec![]); - } - - let mut tx = self.pool.begin().await.map_err(|e| { - crate::telemetry::queue::set_active_span_error(); - e - })?; - let mut ids = Vec::new(); - - for document_id in document_ids { - let id = Ulid::new().to_string(); - - let result = sqlx::query( - r#" - INSERT INTO embedding_queue (id, document_id, traceparent, tracestate) - SELECT $1, $2, $3, $4 - WHERE NOT EXISTS ( - SELECT 1 FROM embedding_queue - WHERE document_id = $2 AND status IN ('pending', 'processing') - ) - "#, - ) - .bind(&id) - .bind(&document_id) - .bind(&carrier.traceparent) - .bind(&carrier.tracestate) - .execute(&mut *tx) - .await - .map_err(|e| { - crate::telemetry::queue::set_active_span_error(); - e - })?; - - if result.rows_affected() > 0 { - ids.push(id); - } - } - tx.commit().await.map_err(|e| { - crate::telemetry::queue::set_active_span_error(); - e - })?; - Ok(ids) - } - .instrument(producer_span) - .await + tx.commit().await?; + Ok(ids) } pub async fn enqueue_batch_missing_current_embeddings( @@ -217,79 +164,54 @@ impl EmbeddingQueue { return Ok(vec![]); } - let queue_name = "embedding_queue"; - - // Build PRODUCER span covering provider check + INSERT. - // inject_active_context is called from inside the instrumented block - // so the producer span is active during injection and the DB write. - let producer_span = crate::telemetry::queue::build_producer_span(queue_name); - - async move { - let carrier = crate::telemetry::queue::inject_active_context(); - - if !self.provider_repo.has_active_provider().await? { - return Ok(vec![]); - } + if !self.provider_repo.has_active_provider().await? { + return Ok(vec![]); + } - let ids: Vec = document_ids.iter().map(|_| generate_ulid()).collect(); - let traceparents: Vec> = std::iter::repeat(carrier.traceparent.clone()) - .take(ids.len()) - .collect(); - let tracestates: Vec> = std::iter::repeat(carrier.tracestate.clone()) - .take(ids.len()) - .collect(); + let ids: Vec = document_ids.iter().map(|_| generate_ulid()).collect(); - let rows = sqlx::query( - r#" - WITH active_provider AS ( - SELECT config->>'model' AS model_name - FROM embedding_providers - WHERE is_current = TRUE AND is_deleted = FALSE - LIMIT 1 - ), - input_rows AS ( - SELECT id, document_id, traceparent, tracestate, ordinality - FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[]) WITH ORDINALITY AS t(id, document_id, traceparent, tracestate, ordinality) - ), - deduped_input AS ( - SELECT DISTINCT ON (document_id) id, document_id, traceparent, tracestate - FROM input_rows - ORDER BY document_id, ordinality - ) - INSERT INTO embedding_queue (id, document_id, traceparent, tracestate) - SELECT input.id, input.document_id, input.traceparent, input.tracestate - FROM deduped_input input - CROSS JOIN active_provider provider - WHERE NOT EXISTS ( - SELECT 1 - FROM embedding_queue q - WHERE q.document_id = input.document_id - AND q.status IN ('pending', 'processing') - ) - AND NOT EXISTS ( - SELECT 1 - FROM embeddings e - WHERE e.document_id = input.document_id - AND e.model_name = provider.model_name - ) - RETURNING id - "#, + let rows = sqlx::query( + r#" + WITH active_provider AS ( + SELECT config->>'model' AS model_name + FROM embedding_providers + WHERE is_current = TRUE AND is_deleted = FALSE + LIMIT 1 + ), + input_rows AS ( + SELECT id, document_id, ordinality + FROM UNNEST($1::text[], $2::text[]) WITH ORDINALITY AS t(id, document_id, ordinality) + ), + deduped_input AS ( + SELECT DISTINCT ON (document_id) id, document_id + FROM input_rows + ORDER BY document_id, ordinality ) - .bind(&ids) - .bind(&document_ids) - .bind(&traceparents) - .bind(&tracestates) - .fetch_all(&self.pool) - .await - .map_err(|e| { - crate::telemetry::queue::set_active_span_error(); - e - })?; - - Ok(rows.into_iter().map(|row| row.get("id")).collect()) - } - .instrument(producer_span) - .await + INSERT INTO embedding_queue (id, document_id) + SELECT input.id, input.document_id + FROM deduped_input input + CROSS JOIN active_provider provider + WHERE NOT EXISTS ( + SELECT 1 + FROM embedding_queue q + WHERE q.document_id = input.document_id + AND q.status IN ('pending', 'processing') + ) + AND NOT EXISTS ( + SELECT 1 + FROM embeddings e + WHERE e.document_id = input.document_id + AND e.model_name = provider.model_name + ) + RETURNING id + "#, + ) + .bind(&ids) + .bind(&document_ids) + .fetch_all(&self.pool) + .await?; + + Ok(rows.into_iter().map(|row| row.get("id")).collect()) } pub async fn dequeue_batch(&self, batch_size: i32) -> Result> { diff --git a/shared/src/queue.rs b/shared/src/queue.rs index f6cc866e1..1393bce25 100644 --- a/shared/src/queue.rs +++ b/shared/src/queue.rs @@ -1,7 +1,6 @@ use crate::utils::generate_ulid; use anyhow::Result; use sqlx::{PgPool, Row}; -use tracing::Instrument; use crate::models::{ConnectorEvent, ConnectorEventQueueItem, EventStatus, SyncType}; @@ -27,42 +26,24 @@ impl EventQueue { } pub async fn enqueue(&self, source_id: &str, event: &ConnectorEvent) -> Result { - let queue_name = "connector_events_queue"; let id = generate_ulid(); let event_type = event_type_str(event); - // Build PRODUCER span and instrument the full INSERT with the span. - // inject_active_context is called from inside the instrumented block - // so the producer span is active during injection and the DB write. - let producer_span = crate::telemetry::queue::build_producer_span(queue_name); - - async move { - let carrier = crate::telemetry::queue::inject_active_context(); - - if let Err(e) = sqlx::query( - r#" - INSERT INTO connector_events_queue (id, sync_run_id, source_id, event_type, payload, traceparent, tracestate) - VALUES ($1, $2, $3, $4, $5, $6, $7) - "#, - ) - .bind(&id) - .bind(event.sync_run_id()) - .bind(source_id) - .bind(event_type) - .bind(serde_json::to_value(event)?) - .bind(&carrier.traceparent) - .bind(&carrier.tracestate) - .execute(&self.pool) - .await - { - crate::telemetry::queue::set_active_span_error(); - return Err(e.into()); - } + sqlx::query( + r#" + INSERT INTO connector_events_queue (id, sync_run_id, source_id, event_type, payload) + VALUES ($1, $2, $3, $4, $5) + "#, + ) + .bind(&id) + .bind(event.sync_run_id()) + .bind(source_id) + .bind(event_type) + .bind(serde_json::to_value(event)?) + .execute(&self.pool) + .await?; - Ok(id) - } - .instrument(producer_span) - .await + Ok(id) } pub async fn enqueue_batch( @@ -74,13 +55,6 @@ impl EventQueue { return Ok(Vec::new()); } - let queue_name = "connector_events_queue"; - - // Build PRODUCER span and instrument the full batch INSERT with the span. - // inject_active_context is called from inside the instrumented block - // so the producer span is active during injection and the DB write. - let producer_span = crate::telemetry::queue::build_producer_span(queue_name); - let mut ids: Vec = Vec::with_capacity(events.len()); let mut sync_run_ids: Vec = Vec::with_capacity(events.len()); let mut source_ids: Vec = Vec::with_capacity(events.len()); @@ -95,40 +69,21 @@ impl EventQueue { payloads.push(serde_json::to_value(event)?); } - async move { - let carrier = crate::telemetry::queue::inject_active_context(); - - let traceparents: Vec> = std::iter::repeat(carrier.traceparent.clone()) - .take(ids.len()) - .collect(); - let tracestates: Vec> = std::iter::repeat(carrier.tracestate.clone()) - .take(ids.len()) - .collect(); - - if let Err(e) = sqlx::query( - r#" - INSERT INTO connector_events_queue (id, sync_run_id, source_id, event_type, payload, traceparent, tracestate) - SELECT * FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[], $5::jsonb[], $6::text[], $7::text[]) - "#, - ) - .bind(&ids) - .bind(&sync_run_ids) - .bind(&source_ids) - .bind(&event_types) - .bind(&payloads) - .bind(&traceparents) - .bind(&tracestates) - .execute(&self.pool) - .await - { - crate::telemetry::queue::set_active_span_error(); - return Err(e.into()); - } + sqlx::query( + r#" + INSERT INTO connector_events_queue (id, sync_run_id, source_id, event_type, payload) + SELECT * FROM UNNEST($1::text[], $2::text[], $3::text[], $4::text[], $5::jsonb[]) + "#, + ) + .bind(&ids) + .bind(&sync_run_ids) + .bind(&source_ids) + .bind(&event_types) + .bind(&payloads) + .execute(&self.pool) + .await?; - Ok(ids) - } - .instrument(producer_span) - .await + Ok(ids) } pub async fn dequeue_batch(&self, batch_size: i32) -> Result> { @@ -186,9 +141,7 @@ impl EventQueue { q.max_retries, q.created_at, q.processed_at, - q.error_message, - q.traceparent, - q.tracestate + q.error_message "#, ) .bind(batch_size) @@ -259,9 +212,7 @@ impl EventQueue { q.max_retries, q.created_at, q.processed_at, - q.error_message, - q.traceparent, - q.tracestate + q.error_message "#, ) .bind(batch_size) @@ -331,9 +282,7 @@ impl EventQueue { q.max_retries, q.created_at, q.processed_at, - q.error_message, - q.traceparent, - q.tracestate + q.error_message "#, ) .bind(batch_size) @@ -358,9 +307,6 @@ impl EventQueue { _ => crate::models::EventStatus::Pending, }; - let traceparent: Option = row.try_get("traceparent").ok().flatten(); - let tracestate: Option = row.try_get("tracestate").ok().flatten(); - events.push(ConnectorEventQueueItem { id: row.get("id"), sync_run_id: row.get("sync_run_id"), @@ -373,8 +319,6 @@ impl EventQueue { created_at: row.get("created_at"), processed_at: row.get("processed_at"), error_message: row.get("error_message"), - traceparent, - tracestate, }); } events diff --git a/shared/src/telemetry.rs b/shared/src/telemetry.rs index 9647bca41..56a922ae7 100644 --- a/shared/src/telemetry.rs +++ b/shared/src/telemetry.rs @@ -506,390 +506,3 @@ pub mod http_client { .await } } - -/// Queue telemetry helpers for W3C trace context propagation through async queues. -/// -/// This module provides: -/// - `inject_producer_context`: creates a PRODUCER span and extracts W3C headers -/// - `extract_remote_span_context`: reconstructs a valid remote SpanContext from stored headers -/// - `add_consumer_links`: adds Link entries to a consumer span from extracted contexts -/// -/// Producer spans are parented to the current tracing context (the caller). -/// Consumer spans MUST be created as new roots (no parent) and reference the producer -/// context via links only. -pub mod queue { - use opentelemetry::{ - global, - propagation::{Extractor, Injector}, - trace::{SpanContext, SpanId, TraceContextExt, TraceId}, - Context, - }; - use std::collections::HashSet; - use tracing::{info_span, Span}; - use tracing_opentelemetry::OpenTelemetrySpanExt; - - const TRACEPARENT_MAX_LEN: usize = 55; - const TRACESTATE_MAX_LEN: usize = 512; - - /// A text map carrier for W3C trace context. - #[derive(Debug, Clone, Default)] - pub struct TraceCarrier { - pub traceparent: Option, - pub tracestate: Option, - } - - impl Injector for TraceCarrier { - fn set(&mut self, key: &str, value: String) { - match key.to_lowercase().as_str() { - "traceparent" => self.traceparent = Some(value), - "tracestate" => self.tracestate = Some(value), - _ => {} - } - } - } - - impl Extractor for TraceCarrier { - fn get(&self, key: &str) -> Option<&str> { - match key.to_lowercase().as_str() { - "traceparent" => self.traceparent.as_deref(), - "tracestate" => self.tracestate.as_deref(), - _ => None, - } - } - - fn keys(&self) -> Vec<&str> { - let mut keys = Vec::new(); - if self.traceparent.is_some() { - keys.push("traceparent"); - } - if self.tracestate.is_some() { - keys.push("tracestate"); - } - keys - } - } - - /// Build a PRODUCER tracing span for the given queue, parented to the - /// current scope. The span is returned *un-entered* so the caller can - /// enter it briefly to inject context before instrumenting the actual - /// async operation. - pub fn build_producer_span(queue_name: &str) -> Span { - let span_name = format!("{} publish", queue_name); - info_span!( - target: "queue", - parent: &Span::current(), - "{}", span_name, - otel.name = span_name.as_str(), - otel.kind = "PRODUCER", - messaging.system = "postgresql", - messaging.destination = queue_name, - messaging.operation.type = "publish", - ) - } - - /// Inject W3C trace context from the *currently active* tracing span - /// into a TraceCarrier. Call this while the producer span is entered. - pub fn inject_active_context() -> TraceCarrier { - let cx = Span::current().context(); - let mut carrier = TraceCarrier::default(); - global::get_text_map_propagator(|propagator| propagator.inject_context(&cx, &mut carrier)); - carrier - } - - /// Extract a valid remote SpanContext from optional stored W3C headers. - pub fn extract_remote_span_context( - traceparent: Option<&str>, - tracestate: Option<&str>, - ) -> Option { - let traceparent = traceparent?; - if traceparent.len() > TRACEPARENT_MAX_LEN { - return None; - } - let mut carrier = TraceCarrier::default(); - Injector::set(&mut carrier, "traceparent", traceparent.to_string()); - if let Some(ts) = tracestate { - if ts.len() <= TRACESTATE_MAX_LEN { - Injector::set(&mut carrier, "tracestate", ts.to_string()); - } - } - let cx: Context = - global::get_text_map_propagator(|p| p.extract_with_context(&Context::new(), &carrier)); - let span = cx.span(); - let sc = span.span_context(); - if sc.is_valid() { - Some(sc.clone()) - } else { - None - } - } - - /// Collect deduplicated producer SpanContext values from queue items. - /// Deduplication is by the full `(trace_id, span_id)` pair (not trace - /// only), so two distinct producer spans within the same trace are - /// both retained. Invalid/missing contexts are skipped. - pub fn collect_span_contexts( - items: &[Option<(Option, Option)>], - ) -> Vec { - let mut seen: HashSet<(TraceId, SpanId)> = HashSet::new(); - let mut result = Vec::new(); - for ctx in items { - let (tp, ts) = match ctx { - Some((tp, ts)) => (tp.as_deref(), ts.as_deref()), - None => continue, - }; - let Some(sc) = extract_remote_span_context(tp, ts) else { - continue; - }; - let key = (sc.trace_id(), sc.span_id()); - if seen.insert(key) { - result.push(sc.clone()); - } - } - result - } - - /// Mark the current active span as error with a fixed empty description. - /// Used to record producer-side failures without leaking message details - /// into the span status. - pub fn set_active_span_error() { - Span::current().set_status(opentelemetry::trace::Status::error("")); - } - - /// Validate W3C traceparent format. - pub fn is_valid_traceparent(value: &str) -> bool { - if value.len() != 55 { - return false; - } - let parts: Vec<&str> = value.split('-').collect(); - parts.len() == 4 - && parts[0] == "00" - && parts[1].len() == 32 - && parts[1].chars().all(|c| c.is_ascii_hexdigit()) - && parts[2].len() == 16 - && parts[2].chars().all(|c| c.is_ascii_hexdigit()) - && parts[3].len() == 2 - && parts[3].chars().all(|c| c.is_ascii_hexdigit()) - } - - #[cfg(test)] - mod tests { - use super::*; - use opentelemetry::trace::TracerProvider as _; - use opentelemetry_sdk::propagation::TraceContextPropagator; - use tracing_subscriber::layer::SubscriberExt as _; - use tracing_subscriber::util::SubscriberInitExt; - - /// Initialize the global TraceContextPropagator once per test run. - fn init_propagator() { - use std::sync::Once; - static INIT: Once = Once::new(); - INIT.call_once(|| { - global::set_text_map_propagator(TraceContextPropagator::new()); - }); - } - - #[test] - fn test_is_valid_traceparent() { - assert!(is_valid_traceparent( - "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - )); - assert!(!is_valid_traceparent( - "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-010" - )); - assert!(!is_valid_traceparent("too-short")); - assert!(!is_valid_traceparent( - "01-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01" - )); - assert!(!is_valid_traceparent( - "00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-b7ad6b7169203331-01" - )); - } - - #[test] - fn test_extract_remote_span_context_valid() { - init_propagator(); - let sc = extract_remote_span_context( - Some("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"), - None, - ); - assert!(sc.is_some()); - let sc = sc.unwrap(); - assert_eq!( - sc.trace_id(), - TraceId::from_hex("0af7651916cd43dd8448eb211c80319c").unwrap() - ); - assert!(sc.is_sampled()); - } - - #[test] - fn test_extract_remote_span_context_invalid() { - assert!(extract_remote_span_context(None, None).is_none()); - assert!(extract_remote_span_context( - Some("00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-b7ad6b7169203331-01"), - None, - ) - .is_none()); - assert!(extract_remote_span_context( - Some("00-00000000000000000000000000000000-0000000000000000-01"), - None, - ) - .is_none()); - } - - #[test] - fn test_extract_remote_span_context_with_tracestate() { - let sc = extract_remote_span_context( - Some("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"), - Some("vendor1=value1"), - ); - assert!(sc.is_some()); - } - - #[test] - fn test_extract_remote_span_context_ambient_isolation_malformed_55char() { - // Activate an unrelated ambient span, then supply a malformed - // exactly-55-char traceparent with non-hex characters. - // The propagator must NOT fall back to the ambient context. - init_tracer(); - - let parent_span = info_span!("ambient_span"); - let _guard = parent_span.entered(); - - // 55 chars, valid format, but non-hex chars in trace-id - let malformed = "00-zzzzzzzzzzzzzzzzzzzzzzzzzzzzzzzz-b7ad6b7169203331-01"; - assert_eq!(malformed.len(), 55); - - let result = extract_remote_span_context(Some(malformed), None); - assert!( - result.is_none(), - "must not return ambient context when traceparent has non-hex chars" - ); - } - - #[test] - fn test_extract_remote_span_context_ambient_isolation_valid() { - // Activate an unrelated ambient span, then supply a valid but - // different traceparent. The extracted context must match the - // carrier, not the ambient. - init_tracer(); - - let parent_span = info_span!("ambient_span"); - let _guard = parent_span.entered(); - - let carrier_tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; - let result = extract_remote_span_context(Some(carrier_tp), None); - assert!( - result.is_some(), - "valid traceparent must be extracted even with active ambient span" - ); - let sc = result.unwrap(); - assert_eq!( - sc.trace_id(), - TraceId::from_hex("0af7651916cd43dd8448eb211c80319c").unwrap(), - "extracted trace_id must come from carrier, not ambient" - ); - } - - #[test] - fn test_collect_span_contexts_deduplicates() { - init_propagator(); - let tp = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; - let items = vec![ - Some((Some(tp.to_string()), None)), - Some((Some(tp.to_string()), None)), - ]; - let result = collect_span_contexts(&items); - assert_eq!( - result.len(), - 1, - "duplicate (trace_id, span_id) should be deduplicated" - ); - } - - #[test] - fn test_collect_span_contexts_multiple_traces() { - init_propagator(); - let tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; - let tp2 = "00-1bf7651916cd43dd8448eb211c80319c-b8ad6b7169203331-01"; - let items = vec![ - Some((Some(tp1.to_string()), None)), - Some((Some(tp2.to_string()), None)), - ]; - assert_eq!(collect_span_contexts(&items).len(), 2); - } - - #[test] - fn test_collect_span_contexts_same_trace_different_spans_retained() { - init_propagator(); - let tp1 = "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"; - let tp2 = "00-0af7651916cd43dd8448eb211c80319c-aaaaaaaaaaaaaaaa-01"; - let items = vec![ - Some((Some(tp1.to_string()), None)), - Some((Some(tp2.to_string()), None)), - ]; - let result = collect_span_contexts(&items); - assert_eq!( - result.len(), - 2, - "two producer spans with same trace_id but different span_id must both be retained" - ); - } - - #[test] - fn test_collect_span_contexts_skips_invalid() { - let items = vec![ - None, - Some((None, None)), - Some((Some("invalid".to_string()), None)), - ]; - assert!(collect_span_contexts(&items).is_empty()); - } - - /// Set up a global tracer provider and tracing subscriber for tests - /// that need to produce trace context. - fn init_tracer() { - init_propagator(); - static TRACER_INIT: std::sync::Once = std::sync::Once::new(); - TRACER_INIT.call_once(|| { - let provider = opentelemetry_sdk::trace::SdkTracerProvider::builder().build(); - global::set_tracer_provider(provider.clone()); - let tracer = provider.tracer("test"); - let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); - let _ = tracing_subscriber::registry() - .with(telemetry_layer) - .try_init(); - }); - } - - #[test] - fn test_build_producer_span_and_inject() { - init_tracer(); - let parent_span = info_span!("test"); - let _guard = parent_span.entered(); - let span = build_producer_span("test_queue"); - let carrier = { - let _g = span.entered(); - inject_active_context() - }; - assert!(carrier.traceparent.is_some()); - let tp = carrier.traceparent.unwrap(); - assert_eq!(tp.len(), 55); - assert!(is_valid_traceparent(&tp)); - } - - #[test] - fn test_trace_carrier_round_trip() { - init_tracer(); - let parent_span = info_span!("test"); - let _guard = parent_span.entered(); - let span = build_producer_span("round_trip"); - let carrier = { - let _g = span.entered(); - inject_active_context() - }; - let tp = carrier.traceparent.clone().unwrap(); - let extracted = extract_remote_span_context(Some(&tp), carrier.tracestate.as_deref()); - assert!(extracted.is_some()); - } - } -} diff --git a/shared/tests/embedding_queue_test.rs b/shared/tests/embedding_queue_test.rs index c9e639119..5dc601e03 100644 --- a/shared/tests/embedding_queue_test.rs +++ b/shared/tests/embedding_queue_test.rs @@ -335,191 +335,4 @@ mod tests { } // ----------------------------------------------------------------------- - // Trace context persistence - // ----------------------------------------------------------------------- - - /// Start a per-test-binary tracer provider + subscriber so we can create - /// active tracing spans whose context is injected by the real enqueue. - fn init_trace_subscriber() { - use std::sync::Once; - static INIT: Once = Once::new(); - INIT.call_once(|| { - use opentelemetry::global; - use opentelemetry::trace::TracerProvider; - use opentelemetry_sdk::propagation::TraceContextPropagator; - use opentelemetry_sdk::trace::{ - InMemorySpanExporter, RandomIdGenerator, SdkTracerProvider, SimpleSpanProcessor, - }; - use tracing_subscriber::layer::SubscriberExt as _; - use tracing_subscriber::util::SubscriberInitExt as _; - - global::set_text_map_propagator(TraceContextPropagator::new()); - - let exporter = InMemorySpanExporter::default(); - let provider = SdkTracerProvider::builder() - .with_span_processor(SimpleSpanProcessor::new(exporter)) - .with_id_generator(RandomIdGenerator::default()) - .build(); - global::set_tracer_provider(provider.clone()); - - let tracer = provider.tracer("test"); - let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); - let _ = tracing_subscriber::registry() - .with(telemetry_layer) - .try_init(); - }); - } - - #[tokio::test] - async fn test_enqueue_traceparent_persisted() { - init_trace_subscriber(); - - let env = TestEnvironment::new().await.unwrap(); - let pool = env.db_pool.pool().clone(); - let queue = EmbeddingQueue::new(pool.clone()); - insert_active_embedding_provider(&pool).await; - - let doc_id = create_document(&pool).await; - - let parent_span = tracing::info_span!("test_emb_enqueue_traceparent"); - async move { - let queue_id = queue.enqueue(doc_id.clone()).await.unwrap().unwrap(); - - // Query the stored row directly — traceparent must be non-null - // and valid. - let row: (Option, Option) = - sqlx::query_as("SELECT traceparent, tracestate FROM embedding_queue WHERE id = $1") - .bind(&queue_id) - .fetch_one(&pool) - .await - .unwrap(); - - let (stored_tp, stored_ts) = row; - assert!(stored_tp.is_some(), "traceparent must be stored"); - let tp = stored_tp.unwrap(); - assert!( - shared::telemetry::queue::is_valid_traceparent(&tp), - "stored traceparent must be valid W3C format: {}", - tp - ); - if let Some(ref ts) = stored_ts { - assert!(!ts.is_empty(), "tracestate must not be empty if present"); - } - - // Dequeue the item and verify it carries the same traceparent. - let batch = queue.dequeue_batch(10).await.unwrap(); - assert_eq!(batch.len(), 1); - assert_eq!(batch[0].traceparent, Some(tp.clone())); - assert_eq!(batch[0].tracestate, stored_ts); - } - .instrument(parent_span) - .await - } - - #[tokio::test] - async fn test_enqueue_batch_traceparent_persisted() { - init_trace_subscriber(); - - let env = TestEnvironment::new().await.unwrap(); - let pool = env.db_pool.pool().clone(); - let queue = EmbeddingQueue::new(pool.clone()); - insert_active_embedding_provider(&pool).await; - - let mut doc_ids = Vec::new(); - for _ in 0..3 { - doc_ids.push(create_document(&pool).await); - } - - let parent_span = tracing::info_span!("test_emb_enqueue_batch_traceparent"); - async move { - let queue_ids = queue.enqueue_batch(doc_ids.clone()).await.unwrap(); - assert_eq!(queue_ids.len(), 3, "all three docs should be enqueued"); - - for queue_id in &queue_ids { - let row: (Option, Option) = sqlx::query_as( - "SELECT traceparent, tracestate FROM embedding_queue WHERE id = $1", - ) - .bind(queue_id) - .fetch_one(&pool) - .await - .unwrap(); - - let (stored_tp, stored_ts) = row; - assert!( - stored_tp.is_some(), - "traceparent must be stored for queue_id={}", - queue_id - ); - let tp = stored_tp.unwrap(); - assert!( - shared::telemetry::queue::is_valid_traceparent(&tp), - "stored traceparent must be valid W3C format: {}", - tp - ); - if let Some(ref ts) = stored_ts { - assert!(!ts.is_empty(), "tracestate must not be empty if present"); - } - } - } - .instrument(parent_span) - .await - } - - #[tokio::test] - async fn test_enqueue_batch_missing_current_embeddings_traceparent_persisted() { - init_trace_subscriber(); - - let env = TestEnvironment::new().await.unwrap(); - let pool = env.db_pool.pool().clone(); - let queue = EmbeddingQueue::new(pool.clone()); - insert_active_embedding_provider(&pool).await; - - // Create docs without any embeddings (so they qualify as "missing"). - let mut doc_ids = Vec::new(); - for _ in 0..3 { - doc_ids.push(create_document(&pool).await); - } - - let parent_span = - tracing::info_span!("test_emb_enqueue_batch_missing_current_embeddings_traceparent"); - async move { - let queue_ids = queue - .enqueue_batch_missing_current_embeddings(doc_ids.clone()) - .await - .unwrap(); - assert_eq!( - queue_ids.len(), - 3, - "all three docs without embeddings should be enqueued" - ); - - for queue_id in &queue_ids { - let row: (Option, Option) = sqlx::query_as( - "SELECT traceparent, tracestate FROM embedding_queue WHERE id = $1", - ) - .bind(queue_id) - .fetch_one(&pool) - .await - .unwrap(); - - let (stored_tp, stored_ts) = row; - assert!( - stored_tp.is_some(), - "traceparent must be stored for queue_id={}", - queue_id - ); - let tp = stored_tp.unwrap(); - assert!( - shared::telemetry::queue::is_valid_traceparent(&tp), - "stored traceparent must be valid W3C format: {}", - tp - ); - if let Some(ref ts) = stored_ts { - assert!(!ts.is_empty(), "tracestate must not be empty if present"); - } - } - } - .instrument(parent_span) - .await - } } diff --git a/shared/tests/event_queue_test.rs b/shared/tests/event_queue_test.rs index ee7504e21..af755e855 100644 --- a/shared/tests/event_queue_test.rs +++ b/shared/tests/event_queue_test.rs @@ -616,124 +616,4 @@ mod tests { // ----------------------------------------------------------------------- // Trace context persistence - // ----------------------------------------------------------------------- - - /// Start a per-test-binary tracer provider + subscriber so we can create - /// active tracing spans whose context is injected by the real enqueue. - fn init_trace_subscriber() { - use std::sync::Once; - static INIT: Once = Once::new(); - INIT.call_once(|| { - use opentelemetry::global; - use opentelemetry::trace::TracerProvider; - use opentelemetry_sdk::propagation::TraceContextPropagator; - use opentelemetry_sdk::trace::InMemorySpanExporter; - use opentelemetry_sdk::trace::{ - RandomIdGenerator, SdkTracerProvider, SimpleSpanProcessor, - }; - use tracing_subscriber::layer::SubscriberExt as _; - use tracing_subscriber::util::SubscriberInitExt as _; - - global::set_text_map_propagator(TraceContextPropagator::new()); - - let exporter = InMemorySpanExporter::default(); - let provider = SdkTracerProvider::builder() - .with_span_processor(SimpleSpanProcessor::new(exporter)) - .with_id_generator(RandomIdGenerator::default()) - .build(); - global::set_tracer_provider(provider.clone()); - - let tracer = provider.tracer("test"); - let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); - let _ = tracing_subscriber::registry() - .with(telemetry_layer) - .try_init(); - }); - } - - #[tokio::test] - async fn test_enqueue_traceparent_persisted() { - init_trace_subscriber(); - - let env = TestEnvironment::new().await.unwrap(); - let pool = env.db_pool.pool().clone(); - let queue = EventQueue::new(pool.clone()); - - // Start an active parent span so enqueue injects trace context. - let parent_span = tracing::info_span!("test_enqueue_traceparent"); - async move { - let event = make_event("run-tp", "doc-tp"); - let event_id = queue.enqueue(TEST_SOURCE_ID, &event).await.unwrap(); - - // Query the stored row directly — traceparent must be non-null - // and valid. - let row: (Option, Option) = sqlx::query_as( - "SELECT traceparent, tracestate FROM connector_events_queue WHERE id = $1", - ) - .bind(&event_id) - .fetch_one(&pool) - .await - .unwrap(); - - let (stored_tp, stored_ts) = row; - assert!(stored_tp.is_some(), "traceparent must be stored"); - let tp = stored_tp.unwrap(); - assert!( - shared::telemetry::queue::is_valid_traceparent(&tp), - "stored traceparent must be valid W3C format: {}", - tp - ); - // tracestate is optional (may be None when no tracestate set). - if let Some(ref ts) = stored_ts { - assert!(!ts.is_empty(), "tracestate must not be empty if present"); - } - - // Dequeue the item and verify it carries the same traceparent. - let batch = queue.dequeue_batch(10).await.unwrap(); - assert_eq!(batch.len(), 1); - assert_eq!(batch[0].traceparent, Some(tp.clone())); - assert_eq!(batch[0].tracestate, stored_ts); - } - .instrument(parent_span) - .await - } - - #[tokio::test] - async fn test_enqueue_batch_traceparent_persisted() { - init_trace_subscriber(); - - let env = TestEnvironment::new().await.unwrap(); - let pool = env.db_pool.pool().clone(); - let queue = EventQueue::new(pool.clone()); - - // Start an active parent span so enqueue_batch injects trace context. - let parent_span = tracing::info_span!("test_enqueue_batch_traceparent"); - async move { - let run_id = ulid::Ulid::new().to_string(); - let events: Vec = (0..2) - .map(|i| make_event(&run_id, &format!("doc-batch-tp-{}", i))) - .collect(); - let ids = queue.enqueue_batch(TEST_SOURCE_ID, &events).await.unwrap(); - assert_eq!(ids.len(), 2); - - // Query each row and verify valid traceparent. - for id in &ids { - let tp: Option = sqlx::query_scalar( - "SELECT traceparent FROM connector_events_queue WHERE id = $1", - ) - .bind(id) - .fetch_one(&pool) - .await - .unwrap(); - - assert!(tp.is_some(), "traceparent must be stored for batch item"); - assert!( - shared::telemetry::queue::is_valid_traceparent(&tp.unwrap()), - "stored traceparent must be valid W3C format" - ); - } - } - .instrument(parent_span) - .await - } } diff --git a/shared/tests/queue_link_test.rs b/shared/tests/queue_link_test.rs deleted file mode 100644 index 437ea9476..000000000 --- a/shared/tests/queue_link_test.rs +++ /dev/null @@ -1,357 +0,0 @@ -//! Focused integration tests for native OTel link propagation through queues. -//! -//! Uses `InMemorySpanExporter` to inspect exported OpenTelemetry spans and -//! verify that CONSUMER spans carry native links (not string attributes) to -//! their PRODUCER spans. -//! -//! Tests use the *exact production* `build_producer_span` / `inject_active_context` -//! / `collect_span_contexts` helpers from the shared telemetry::queue module. - -use std::sync::{Mutex, OnceLock}; - -use opentelemetry::{ - global, - trace::{SpanKind, TracerProvider as _}, -}; -use opentelemetry_sdk::{ - propagation::TraceContextPropagator, - trace::{InMemorySpanExporter, RandomIdGenerator, SdkTracerProvider, SimpleSpanProcessor}, -}; -use shared::telemetry::queue; -use tracing::{info_span, Instrument}; -use tracing_opentelemetry::OpenTelemetrySpanExt; -use tracing_subscriber::{layer::SubscriberExt as _, util::SubscriberInitExt as _}; - -// --------------------------------------------------------------------------- -// Global test infrastructure -// --------------------------------------------------------------------------- - -/// Process-global mutex to serialise all tests in this file. -/// Every test acquires this guard before accessing the shared exporter. -static TEST_MUTEX: OnceLock> = OnceLock::new(); -fn test_guard() -> std::sync::MutexGuard<'static, ()> { - TEST_MUTEX - .get_or_init(|| Mutex::new(())) - .lock() - .unwrap_or_else(|p| p.into_inner()) -} - -static EXPORTER: OnceLock = OnceLock::new(); -fn exporter() -> &'static InMemorySpanExporter { - EXPORTER.get().expect("exporter not initialised") -} - -/// Initialise the global tracer provider + tracing subscriber once per suite. -fn ensure_global_init() { - EXPORTER.get_or_init(|| { - global::set_text_map_propagator(TraceContextPropagator::new()); - - let exporter = InMemorySpanExporter::default(); - let provider = SdkTracerProvider::builder() - .with_span_processor(SimpleSpanProcessor::new(exporter.clone())) - .with_id_generator(RandomIdGenerator::default()) - .build(); - global::set_tracer_provider(provider.clone()); - - let tracer = provider.tracer("test"); - let telemetry_layer = tracing_opentelemetry::layer().with_tracer(tracer); - let _ = tracing_subscriber::registry() - .with(telemetry_layer) - .try_init(); - - exporter - }); -} - -fn reset_exporter() { - exporter().reset(); -} - -fn finished_spans() -> Vec { - exporter() - .get_finished_spans() - .expect("get_finished_spans should succeed") -} - -/// Find a finished span by name and kind. -fn find_span<'a>( - spans: &'a [opentelemetry_sdk::trace::SpanData], - name: &str, - kind: SpanKind, -) -> Option<&'a opentelemetry_sdk::trace::SpanData> { - spans.iter().find(|s| s.name == name && s.span_kind == kind) -} - -/// Check whether the span has a `messaging.linked_trace_ids` attribute. -fn has_linked_trace_ids_attr(span: &opentelemetry_sdk::trace::SpanData) -> bool { - span.attributes - .iter() - .any(|kv| kv.key.as_str() == "messaging.linked_trace_ids") -} - -// --------------------------------------------------------------------------- -// Tests -// --------------------------------------------------------------------------- - -/// Producer and consumer exported; consumer has no parent and a different -/// trace ID; native `links` contains exact producer trace+span ID. -#[test] -fn test_producer_consumer_native_link() { - let _guard = test_guard(); - ensure_global_init(); - reset_exporter(); - - let parent = info_span!("test_root"); - - parent.in_scope(|| { - // ---- Producer: create a PRODUCER span and inject context ---- - let carrier = { - let prod_span = queue::build_producer_span("test_queue"); - let _g = prod_span.entered(); - queue::inject_active_context() - // prod_span dropped here with EnteredSpan guard -> exported - }; - - // ---- Consumer: reconstruct context and build CONSUMER span with link ---- - let items = vec![Some((carrier.traceparent, carrier.tracestate))]; - let span_contexts = queue::collect_span_contexts(&items); - assert_eq!(span_contexts.len(), 1, "should have one producer context"); - - let cons_span = info_span!( - parent: None, - "test_queue process", - otel.kind = "CONSUMER", - ); - { - let _g = cons_span.clone().entered(); - for sc in &span_contexts { - cons_span.add_link(sc.clone()); - } - } - // Drop cons_span: consumer span ends and is exported. - drop(cons_span); - }); - - let spans = finished_spans(); - - // Find producer and consumer spans - let producer = find_span(&spans, "test_queue publish", SpanKind::Producer) - .expect("must have PRODUCER span"); - let consumer = find_span(&spans, "test_queue process", SpanKind::Consumer) - .expect("must have CONSUMER span"); - - // ---- Consumer has no parent (parent_span_id is all zeros) ---- - assert_eq!( - consumer.parent_span_id, - opentelemetry::trace::SpanId::INVALID, - "CONSUMER span must have INVALID parent_span_id (new root)" - ); - - // ---- Consumer has a different trace ID than producer ---- - assert_ne!( - consumer.span_context.trace_id(), - producer.span_context.trace_id(), - "CONSUMER trace ID must differ from PRODUCER trace ID (new root)" - ); - - // ---- Consumer has a native link to the producer ---- - assert!( - !consumer.links.is_empty(), - "CONSUMER span must have at least one link" - ); - let link = &consumer.links[0]; - assert_eq!( - link.span_context.trace_id(), - producer.span_context.trace_id(), - "link trace_id must match PRODUCER trace_id" - ); - assert_eq!( - link.span_context.span_id(), - producer.span_context.span_id(), - "link span_id must match PRODUCER span_id" - ); - - // ---- No messaging.linked_trace_ids attribute exists ---- - assert!( - !has_linked_trace_ids_attr(consumer), - "CONSUMER span must NOT have messaging.linked_trace_ids attribute" - ); - assert!( - !has_linked_trace_ids_attr(producer), - "PRODUCER span must NOT have messaging.linked_trace_ids attribute" - ); -} - -/// Two producer spans in the same trace yield two consumer links. -#[test] -fn test_two_producers_same_trace_two_links() { - let _guard = test_guard(); - ensure_global_init(); - reset_exporter(); - - let parent = info_span!("test_root"); - - parent.in_scope(|| { - // First producer - let c1 = { - let p1 = queue::build_producer_span("test_queue"); - let _g = p1.entered(); - queue::inject_active_context() - }; - - // Second producer in the same parent trace — same trace_id, different span_id - let c2 = { - let p2 = queue::build_producer_span("test_queue"); - let _g = p2.entered(); - queue::inject_active_context() - }; - - // Consumer: collect both contexts - let items = vec![ - Some((c1.traceparent, c1.tracestate)), - Some((c2.traceparent, c2.tracestate)), - ]; - let span_contexts = queue::collect_span_contexts(&items); - assert_eq!( - span_contexts.len(), - 2, - "both producer spans should be retained (different span_id)" - ); - - let cons_span = info_span!( - parent: None, - "test_queue process", - otel.kind = "CONSUMER", - ); - { - let _g = cons_span.clone().entered(); - for sc in &span_contexts { - cons_span.add_link(sc.clone()); - } - } - drop(cons_span); - }); - - let spans = finished_spans(); - - let consumer = find_span(&spans, "test_queue process", SpanKind::Consumer) - .expect("must have CONSUMER span"); - - assert_eq!( - consumer.links.len(), - 2, - "CONSUMER must have 2 links for 2 distinct producer spans" - ); - - // Verify both links have the same trace_id (from the parent) - let trace_id = consumer.links[0].span_context.trace_id(); - for link in consumer.links.iter() { - assert_eq!( - link.span_context.trace_id(), - trace_id, - "all links should share the same trace_id" - ); - } - - // Verify link span IDs are different - assert_ne!( - consumer.links[0].span_context.span_id(), - consumer.links[1].span_context.span_id(), - "two producer links must have different span IDs" - ); - - assert!(!has_linked_trace_ids_attr(consumer)); -} - -/// Duplicate same context dedupes to one link. -#[test] -fn test_duplicate_context_dedupes() { - let _guard = test_guard(); - ensure_global_init(); - reset_exporter(); - - let parent = info_span!("test_root"); - - parent.in_scope(|| { - let carrier = { - let prod_span = queue::build_producer_span("test_queue"); - let _g = prod_span.entered(); - queue::inject_active_context() - }; - - // Duplicate the same carrier - let items = vec![ - Some((carrier.traceparent.clone(), carrier.tracestate.clone())), - Some((carrier.traceparent, carrier.tracestate)), - ]; - let span_contexts = queue::collect_span_contexts(&items); - assert_eq!( - span_contexts.len(), - 1, - "duplicate (trace_id, span_id) must dedupe to 1" - ); - }); -} - -/// Invalid/missing items yield zero span contexts. -#[test] -fn test_invalid_missing_yield_zero() { - let _guard = test_guard(); - ensure_global_init(); - - let items: Vec, Option)>> = vec![ - None, - Some((None, None)), - Some((Some("invalid".to_string()), None)), - Some(( - Some("00-00000000000000000000000000000000-0000000000000000-01".to_string()), - None, - )), - ]; - let span_contexts = queue::collect_span_contexts(&items); - assert!( - span_contexts.is_empty(), - "invalid/missing items should yield zero span contexts" - ); -} - -/// Producer span is active during the async operation: verify by entering -/// a span, instrumenting a synchronous block, and checking exports. -#[tokio::test] -async fn test_producer_span_active_during_operation() { - let _guard = test_guard(); - ensure_global_init(); - reset_exporter(); - - let parent = info_span!("test_root"); - let _guard = parent.entered(); - - let prod_span = queue::build_producer_span("test_queue"); - - async move { - // Simulate an async operation (e.g. SQL INSERT) inside the producer span. - // The carrier is acquired inside the instrumented block so the - // producer span is active during injection and the async work. - let _carrier = queue::inject_active_context(); - tokio::time::sleep(std::time::Duration::from_millis(1)).await; - } - .instrument(prod_span) - .await; - - drop(_guard); - - let spans = finished_spans(); - - let producer = find_span(&spans, "test_queue publish", SpanKind::Producer) - .expect("must have PRODUCER span"); - - // Producer span should have messaging.* attributes - let has_destination = producer - .attributes - .iter() - .any(|kv| kv.key.as_str() == "messaging.destination"); - assert!( - has_destination, - "PRODUCER span must have messaging.destination" - ); -} From 54c1c1745c8bf19818df34afbdd376f10d0c7308 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Mon, 27 Jul 2026 12:31:56 +0000 Subject: [PATCH 03/14] fix: restore exception details in error logs that were over-scrubbed Exception messages (e.g. 'Connection refused', 'Timeout') are not sensitive and are critical for debugging production failures. Restore them across all AI Python provider, tool, streaming, and email files. --- services/ai/email_service/sender.py | 6 +++--- services/ai/providers/anthropic.py | 4 ++-- services/ai/providers/gemini.py | 4 ++-- services/ai/providers/openai.py | 4 ++-- services/ai/streaming/generate.py | 2 +- services/ai/streaming/run.py | 2 +- services/ai/tools/connector_handler.py | 4 ++-- services/ai/tools/search_handler.py | 4 ++-- services/ai/tools/searcher_client.py | 6 +++--- 9 files changed, 18 insertions(+), 18 deletions(-) diff --git a/services/ai/email_service/sender.py b/services/ai/email_service/sender.py index 2f43d8e64..ef224a62c 100644 --- a/services/ai/email_service/sender.py +++ b/services/ai/email_service/sender.py @@ -61,7 +61,7 @@ async def send( logger.error("ACS send failed: status=%s", result["status"]) return SendResult(success=False, error=f"Send failed: {result['status']}") except Exception as e: - logger.error("ACS send error") + logger.error("ACS send error: %s", e) return SendResult(success=False, error="Failed to send email via ACS") @@ -97,7 +97,7 @@ async def send( logger.error("Resend error: status=%d", resp.status_code) return SendResult(success=False, error="Failed to send email via Resend") except Exception as e: - logger.error("Resend send error") + logger.error("Resend send error: %s", e) return SendResult(success=False, error="Failed to send email via Resend") @@ -135,7 +135,7 @@ async def send( return SendResult(success=True) except Exception as e: - logger.error("SMTP send error") + logger.error("SMTP send error: %s", e) return SendResult(success=False, error="Failed to send email via SMTP") diff --git a/services/ai/providers/anthropic.py b/services/ai/providers/anthropic.py index a523720e6..b88af1583 100644 --- a/services/ai/providers/anthropic.py +++ b/services/ai/providers/anthropic.py @@ -155,7 +155,7 @@ async def stream_response( yield event except Exception as e: - logger.error("Failed to stream from Anthropic") + logger.error(f"Failed to stream from Anthropic: {e}", exc_info=True) raise ProviderError( str(e), provider_type=self.provider_type, @@ -204,7 +204,7 @@ async def generate_response( return content, usage except Exception as e: - logger.error("Failed to generate response from Anthropic") + logger.error(f"Failed to generate response from Anthropic: {e}") raise ProviderError( str(e), provider_type=self.provider_type, diff --git a/services/ai/providers/gemini.py b/services/ai/providers/gemini.py index 924a204b8..032762542 100644 --- a/services/ai/providers/gemini.py +++ b/services/ai/providers/gemini.py @@ -404,7 +404,7 @@ async def stream_response( yield RawMessageStopEvent(type="message_stop") except Exception as e: - logger.error("Failed to stream from Gemini") + logger.error(f"Failed to stream from Gemini: {e}", exc_info=True) raise ProviderError( str(e), provider_type=self.provider_type, @@ -449,7 +449,7 @@ async def generate_response( return response.text, usage except Exception as e: - logger.error("Failed to generate response") + logger.error(f"Failed to generate response: {e}") raise ProviderError( str(e), provider_type=self.provider_type, diff --git a/services/ai/providers/openai.py b/services/ai/providers/openai.py index bc4d70c15..58a4cd8ba 100644 --- a/services/ai/providers/openai.py +++ b/services/ai/providers/openai.py @@ -305,7 +305,7 @@ async def stream_response( except ProviderError: raise except Exception as e: - logger.error("Failed to stream from OpenAI") + logger.error(f"Failed to stream from OpenAI: {e}", exc_info=True) raise ProviderError( str(e), provider_type=self.provider_type, @@ -472,7 +472,7 @@ async def generate_response( return content, usage except Exception as e: - logger.error("Failed to generate response") + logger.error(f"Failed to generate response: {e}") raise ProviderError( str(e), provider_type=self.provider_type, diff --git a/services/ai/streaming/generate.py b/services/ai/streaming/generate.py index f83194ce8..925f82890 100644 --- a/services/ai/streaming/generate.py +++ b/services/ai/streaming/generate.py @@ -1119,7 +1119,7 @@ async def stream_generator( yield f"event: save_message\ndata: {json.dumps(partial)}\n\n" raise except Exception as e: - logger.error("Failed to generate AI response with tools") + logger.error(f"Failed to generate AI response with tools: {e}", exc_info=True) if not content_blocks_finalized: partial = partial_assistant_message(content_blocks) if partial is not None: diff --git a/services/ai/streaming/run.py b/services/ai/streaming/run.py index 9c340f8b6..e613080ce 100644 --- a/services/ai/streaming/run.py +++ b/services/ai/streaming/run.py @@ -140,7 +140,7 @@ async def run_producer(redis_client, chat_id, gen, messages_repo, parent_id): pass raise except Exception as e: - logger.error("Producer failed") + logger.error(f"Producer failed: {e}", exc_info=True) try: await redis_client.xadd( sk, diff --git a/services/ai/tools/connector_handler.py b/services/ai/tools/connector_handler.py index 0fa111d54..512b35ddf 100644 --- a/services/ai/tools/connector_handler.py +++ b/services/ai/tools/connector_handler.py @@ -177,7 +177,7 @@ async def _fetch_actions(self) -> list[ConnectorAction]: sources = sources_from_sync_overview_response(sources_resp.json()) except Exception as e: - logger.error("Failed to fetch connector info") + logger.error(f"Failed to fetch connector info: {e}") return [] # Build a mapping from source_type to list of active sources @@ -553,7 +553,7 @@ async def execute( except Exception as e: # Transport-level errors (ReadError, ConnectError, TimeoutException, etc.) # don't carry a response. Don't try to access e.response. - logger.error("Connector action failed") + logger.error(f"Connector action failed: {e}", exc_info=True) return ToolResult( content=[{"type": "text", "text": f"Action failed: {e}"}], is_error=True, diff --git a/services/ai/tools/search_handler.py b/services/ai/tools/search_handler.py index 2a5a75617..55232cc95 100644 --- a/services/ai/tools/search_handler.py +++ b/services/ai/tools/search_handler.py @@ -260,7 +260,7 @@ async def _execute_search( try: params = SearchToolParams.model_validate(tool_input) except ValidationError as e: - logger.error("Invalid search_documents input") + logger.error(f"Invalid search_documents input: {e}") return ToolResult( content=[{"type": "text", "text": f"Invalid parameters: {e}"}], is_error=True, @@ -426,6 +426,6 @@ async def _execute_search_tool( try: response: SearchResponse = await searcher_tool.handle(search_request) except Exception as e: - logger.error("Search failed") + logger.error(f"Search failed: {e}") return [] return response.results diff --git a/services/ai/tools/searcher_client.py b/services/ai/tools/searcher_client.py index a3ceb4e44..f39751c88 100644 --- a/services/ai/tools/searcher_client.py +++ b/services/ai/tools/searcher_client.py @@ -191,7 +191,7 @@ async def search_documents(self, request: SearchRequest) -> SearchResponse: response=response, ) except Exception as e: - logger.error("Call to searcher service failed") + logger.error(f"Call to searcher service failed: {e}") raise async def search_people(self, request: PeopleSearchRequest) -> PeopleSearchResponse: @@ -219,7 +219,7 @@ async def search_people(self, request: PeopleSearchRequest) -> PeopleSearchRespo except SearcherError: raise except Exception as e: - logger.error("People search failed") + logger.error(f"People search failed: {e}") raise async def upsert_capabilities( @@ -297,7 +297,7 @@ async def get_attribute_values( ) return {} except Exception as e: - logger.error("Failed to fetch attribute values") + logger.error(f"Failed to fetch attribute values: {e}") return {} async def close(self): From 32297a75b577f51a563b00c9d7d3b98152c245d0 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Mon, 27 Jul 2026 12:36:20 +0000 Subject: [PATCH 04/14] fix: restore remaining stripped debugging info in logs - Restore exception args in 'Failed to initialize services' error - Restore exception args in 'Failed to fetch Gemini model metadata' debug - Add exc_info=True to 'Online processing loop error' for traceback - Restore error object in TypeScript title generation catch handler - Restore errorMessage in action failure context object --- services/ai/embeddings/batch_processor.py | 2 +- services/ai/main.py | 4 ++-- services/ai/providers/gemini.py | 2 +- web/src/routes/api/chat/[chatId]/stream/+server.ts | 2 +- web/src/routes/api/sources/[sourceId]/action/+server.ts | 1 + 5 files changed, 6 insertions(+), 5 deletions(-) diff --git a/services/ai/embeddings/batch_processor.py b/services/ai/embeddings/batch_processor.py index fbe4241e6..886e453ec 100644 --- a/services/ai/embeddings/batch_processor.py +++ b/services/ai/embeddings/batch_processor.py @@ -169,7 +169,7 @@ async def processing_loop(self): if processed_any: await asyncio.sleep(ONLINE_BATCH_DELAY) except Exception: - logger.error("Online processing loop error") + logger.error("Online processing loop error", exc_info=True) await asyncio.sleep(10) async def _process_online_batch(self) -> bool: diff --git a/services/ai/main.py b/services/ai/main.py index b495b5c1a..387c1aadd 100644 --- a/services/ai/main.py +++ b/services/ai/main.py @@ -91,8 +91,8 @@ async def startup_event(): else: app.state.memory_provider = None logger.info("MEMORY_ENABLED=false — memory feature disabled") - except Exception: - logger.error("Failed to initialize services") + except Exception as e: + logger.error(f"Failed to initialize services: {e}") raise diff --git a/services/ai/providers/gemini.py b/services/ai/providers/gemini.py index 032762542..f2f7da46a 100644 --- a/services/ai/providers/gemini.py +++ b/services/ai/providers/gemini.py @@ -237,7 +237,7 @@ async def get_context_window_tokens(self) -> ContextWindowInfo: ) return self._context_window_info except Exception as e: - logger.debug("Failed to fetch Gemini model metadata for %s", self.model) + logger.debug("Failed to fetch Gemini model metadata for %s: %s", self.model, e) self._context_window_info = await super().get_context_window_tokens() return self._context_window_info diff --git a/web/src/routes/api/chat/[chatId]/stream/+server.ts b/web/src/routes/api/chat/[chatId]/stream/+server.ts index 28db060b0..29e141db9 100644 --- a/web/src/routes/api/chat/[chatId]/stream/+server.ts +++ b/web/src/routes/api/chat/[chatId]/stream/+server.ts @@ -331,7 +331,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur logger.warn('Title generation failed') } }) - .catch(() => logger.error('Failed to generate title for chat')) + .catch((err) => logger.error('Failed to generate title for chat', err)) } while (true) { diff --git a/web/src/routes/api/sources/[sourceId]/action/+server.ts b/web/src/routes/api/sources/[sourceId]/action/+server.ts index 1f6af81d4..75bd3df84 100644 --- a/web/src/routes/api/sources/[sourceId]/action/+server.ts +++ b/web/src/routes/api/sources/[sourceId]/action/+server.ts @@ -57,6 +57,7 @@ export const POST: RequestHandler = async ({ params, locals, request }) => { } logger.error('Action failed', { status: response.status, + error: errorMessage, }) throw error(response.status, errorMessage) } From 7a007eac0e8a02d539ef87a1b986911ffdd672f4 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Mon, 27 Jul 2026 15:45:22 +0000 Subject: [PATCH 05/14] fix: restore Omni internal IDs in TypeScript log statements Internal IDs like chatId, userId, documentId, sourceId are not PII. They are critical for debugging production issues. Restore them in error/warn log messages and context objects across all web routes. --- web/src/hooks.server.ts | 3 +- web/src/lib/server/auth.ts | 2 +- web/src/lib/server/logger.ts | 4 +-- .../(public)/auth/google/callback/+server.ts | 2 +- .../(public)/auth/google/unlink/+server.ts | 2 +- web/src/routes/api/chat/+server.ts | 7 +++-- web/src/routes/api/chat/[chatId]/+server.ts | 16 +++++----- .../api/chat/[chatId]/approve/+server.ts | 3 +- .../[chatId]/artifacts/[...path]/+server.ts | 4 ++- .../api/chat/[chatId]/messages/+server.ts | 31 +++++++++++++------ .../messages/[messageId]/edit/+server.ts | 13 ++++++-- .../messages/[messageId]/feedback/+server.ts | 24 +++++++++++--- .../routes/api/chat/[chatId]/stop/+server.ts | 3 +- .../api/chat/[chatId]/stream/+server.ts | 19 +++++++----- web/src/routes/api/oauth/callback/+server.ts | 10 ++++-- .../routes/api/sources/[sourceId]/+server.ts | 2 +- .../api/sources/[sourceId]/action/+server.ts | 4 +-- .../api/sources/[sourceId]/sync/+server.ts | 2 +- web/src/routes/api/user/settings/+server.ts | 4 ++- .../routes/api/v1/documents/[id]/+server.ts | 5 +-- 20 files changed, 106 insertions(+), 54 deletions(-) diff --git a/web/src/hooks.server.ts b/web/src/hooks.server.ts index 9ab381981..efcf1b577 100644 --- a/web/src/hooks.server.ts +++ b/web/src/hooks.server.ts @@ -91,7 +91,7 @@ const handleLogging: Handle = async ({ event, resolve }) => { // The Node auto-instrumentation already creates a server span and extracts // the incoming traceparent; we just need its trace ID for logging. const requestId = getRequestId() || Logger.generateRequestId() - const logger = new Logger('request').withRequest(requestId) + const logger = new Logger('request').withRequest(requestId, event.locals.user?.id) event.locals.requestId = requestId event.locals.logger = logger @@ -153,6 +153,7 @@ export const handleError: HandleServerError = ({ error, event }) => { logger.error('Unhandled server error', error as Error, { method: event.request.method, + userId: event.locals.user?.id, requestId: event.locals.requestId, }) diff --git a/web/src/lib/server/auth.ts b/web/src/lib/server/auth.ts index e8aa5b256..630438de1 100644 --- a/web/src/lib/server/auth.ts +++ b/web/src/lib/server/auth.ts @@ -130,7 +130,7 @@ export async function createUserSession(userId: string) { }, } } catch (error) { - logger.error('Error creating session', error) + logger.error('Error creating session', error, { userId }) return { success: false, error: 'Failed to create session', diff --git a/web/src/lib/server/logger.ts b/web/src/lib/server/logger.ts index f0c199d26..731208c81 100644 --- a/web/src/lib/server/logger.ts +++ b/web/src/lib/server/logger.ts @@ -90,9 +90,9 @@ export class Logger { return childLogger } - withRequest(requestId: string): Logger { + withRequest(requestId: string, userId?: string): Logger { const childLogger = new Logger() - childLogger.logger = this.logger.child({ requestId }) + childLogger.logger = this.logger.child({ requestId, userId }) return childLogger } diff --git a/web/src/routes/(public)/auth/google/callback/+server.ts b/web/src/routes/(public)/auth/google/callback/+server.ts index 1f112e510..31b107fb8 100644 --- a/web/src/routes/(public)/auth/google/callback/+server.ts +++ b/web/src/routes/(public)/auth/google/callback/+server.ts @@ -83,7 +83,7 @@ export const GET: RequestHandler = async ({ url, cookies }) => { redirectUrl.searchParams.set('linked', 'google') } - logger.info(`Google OAuth authentication successful`) + logger.info(`Google OAuth authentication successful for user: ${user.id}`) successUrl = redirectUrl.toString() } catch (error) { logger.error('OAuth callback error:', error) diff --git a/web/src/routes/(public)/auth/google/unlink/+server.ts b/web/src/routes/(public)/auth/google/unlink/+server.ts index f31bb339d..5f0685639 100644 --- a/web/src/routes/(public)/auth/google/unlink/+server.ts +++ b/web/src/routes/(public)/auth/google/unlink/+server.ts @@ -34,7 +34,7 @@ export const POST: RequestHandler = async ({ url, cookies }) => { providerUserId, ) - console.log('Google OAuth credentials removed') + console.log(`Google OAuth credentials removed for user: ${userSession.user.id}`) // Redirect back to settings with success message throw redirect(302, '/settings/integrations?success=google_unlinked') diff --git a/web/src/routes/api/chat/+server.ts b/web/src/routes/api/chat/+server.ts index 94873c9e6..b29e2c8c5 100644 --- a/web/src/routes/api/chat/+server.ts +++ b/web/src/routes/api/chat/+server.ts @@ -20,12 +20,15 @@ export const POST: RequestHandler = async ({ request, locals }) => { // No body or invalid JSON is fine — modelId stays undefined } - logger.debug('Creating new chat') + logger.debug('Creating new chat', { userId }) try { const chat = await chatRepository.create(userId, undefined, modelId) - logger.info('Chat created successfully') + logger.info('Chat created successfully', { + userId, + chatId: chat.id, + }) return json({ chatId: chat.id }, { status: 200 }) } catch (error) { diff --git a/web/src/routes/api/chat/[chatId]/+server.ts b/web/src/routes/api/chat/[chatId]/+server.ts index 9453397f9..a53f3b618 100644 --- a/web/src/routes/api/chat/[chatId]/+server.ts +++ b/web/src/routes/api/chat/[chatId]/+server.ts @@ -11,17 +11,17 @@ export const GET: RequestHandler = async ({ params, locals }) => { return json({ error: 'chatId parameter is required' }, { status: 400 }) } - logger.debug('Fetching chat details') + logger.debug('Fetching chat details', { chatId }) try { const chat = await chatRepository.get(chatId) if (!chat) { - logger.warn('Chat not found') + logger.warn('Chat not found', { chatId }) return json({ error: 'Chat not found' }, { status: 404 }) } - logger.info('Chat details retrieved successfully') + logger.info('Chat details retrieved successfully', { chatId }) // Convert to match AI service response format const chatDetails = { @@ -34,7 +34,7 @@ export const GET: RequestHandler = async ({ params, locals }) => { return json(chatDetails, { status: 200 }) } catch (error) { - logger.error('Error fetching chat details', error) + logger.error('Error fetching chat details', error, { chatId }) return json( { error: 'Failed to fetch chat details', @@ -75,10 +75,10 @@ export const PATCH: RequestHandler = async ({ params, locals, request }) => { if (result) updatedChat = result } - logger.info('Chat updated') + logger.info('Chat updated', { chatId }) return json(updatedChat) } catch (error) { - logger.error('Error updating chat', error) + logger.error('Error updating chat', error, { chatId }) return json( { error: 'Failed to update chat', @@ -107,10 +107,10 @@ export const DELETE: RequestHandler = async ({ params, locals }) => { try { await chatRepository.delete(chatId) - logger.info('Chat deleted') + logger.info('Chat deleted', { chatId }) return new Response(null, { status: 204 }) } catch (error) { - logger.error('Error deleting chat', error) + logger.error('Error deleting chat', error, { chatId }) return json( { error: 'Failed to delete chat', diff --git a/web/src/routes/api/chat/[chatId]/approve/+server.ts b/web/src/routes/api/chat/[chatId]/approve/+server.ts index 51101a744..672c652cb 100644 --- a/web/src/routes/api/chat/[chatId]/approve/+server.ts +++ b/web/src/routes/api/chat/[chatId]/approve/+server.ts @@ -134,6 +134,7 @@ export const POST: RequestHandler = async ({ params, locals, request }) => { ) logger.info('Tool approval resolved', { + chatId, decision, count: ids.length, }) @@ -145,7 +146,7 @@ export const POST: RequestHandler = async ({ params, locals, request }) => { denialMessageId, }) } catch (error) { - logger.error('Error processing tool approval', error) + logger.error('Error processing tool approval', error, { chatId }) return json( { error: 'Failed to process approval', diff --git a/web/src/routes/api/chat/[chatId]/artifacts/[...path]/+server.ts b/web/src/routes/api/chat/[chatId]/artifacts/[...path]/+server.ts index 4625cf9c3..b26687460 100644 --- a/web/src/routes/api/chat/[chatId]/artifacts/[...path]/+server.ts +++ b/web/src/routes/api/chat/[chatId]/artifacts/[...path]/+server.ts @@ -26,6 +26,8 @@ export const GET: RequestHandler = async ({ params, locals }) => { if (!response.ok) { logger.warn('Artifact proxy failed', undefined, { + chatId, + path, status: response.status, }) throw error(response.status, 'Artifact not found') @@ -42,7 +44,7 @@ export const GET: RequestHandler = async ({ params, locals }) => { }) } catch (err) { if ((err as any)?.status) throw err - logger.error('Artifact proxy error', err) + logger.error('Artifact proxy error', err, { chatId, path }) throw error(502, 'Failed to fetch artifact') } } diff --git a/web/src/routes/api/chat/[chatId]/messages/+server.ts b/web/src/routes/api/chat/[chatId]/messages/+server.ts index 4a5df5290..0c6a53edf 100644 --- a/web/src/routes/api/chat/[chatId]/messages/+server.ts +++ b/web/src/routes/api/chat/[chatId]/messages/+server.ts @@ -51,13 +51,13 @@ export const GET: RequestHandler = async ({ params, locals }) => { return json({ error: 'chatId parameter is required' }, { status: 400 }) } - logger.debug('Fetching chat messages') + logger.debug('Fetching chat messages', { chatId }) try { // First check if chat exists const chat = await chatRepository.get(chatId) if (!chat) { - logger.warn('Chat not found') + logger.warn('Chat not found', { chatId }) return json({ error: 'Chat not found' }, { status: 404 }) } @@ -73,6 +73,7 @@ export const GET: RequestHandler = async ({ params, locals }) => { const chatMessages = await chatMessageRepository.getByChatId(chatId) logger.info('Chat messages retrieved successfully', { + chatId, messageCount: chatMessages.length, }) @@ -88,7 +89,7 @@ export const GET: RequestHandler = async ({ params, locals }) => { return json(messages, { status: 200 }) } catch (error) { - logger.error('Error fetching chat messages', error) + logger.error('Error fetching chat messages', error, { chatId }) return json( { error: 'Failed to fetch messages', @@ -128,13 +129,13 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { return json({ error: 'Content or attachments are required' }, { status: 400 }) } - logger.debug('Adding message to chat') + logger.debug('Adding message to chat', { chatId }) try { // First check if chat exists const chat = await chatRepository.get(chatId) if (!chat) { - logger.warn('Chat not found') + logger.warn('Chat not found', { chatId }) return json({ error: 'Chat not found' }, { status: 404 }) } @@ -158,7 +159,7 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { ) } } catch (error) { - logger.warn('Could not check stream status before adding message', error) + logger.warn('Could not check stream status before adding message', error, { chatId }) } // Create the user message in MessageParam format. If there are attachments, build @@ -185,7 +186,10 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { ? await chatMessageRepository.getByIdInChat(chatId, parentId) : null if (parentId && !parentMessage) { - logger.warn('Ignoring unknown client-provided parent message id') + logger.warn('Ignoring unknown client-provided parent message id', { + chatId, + parentId, + }) parentId = undefined parentMessage = null } @@ -207,14 +211,21 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { parentMessage.id, ) parentId = savedRepairMessage.id - logger.warn('Inserted failed tool_result for interrupted tool call') + logger.warn('Inserted failed tool_result for interrupted tool call', { + chatId, + repairMessageId: savedRepairMessage.id, + }) } } // Save message to database const savedMessage = await chatMessageRepository.create(chatId, userMessage, parentId) - logger.info('Message added successfully') + logger.info('Message added successfully', { + chatId, + messageId: savedMessage.id, + userId: locals.user.id, + }) return json( { @@ -224,7 +235,7 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { { status: 200 }, ) } catch (error) { - logger.error('Error adding message', error) + logger.error('Error adding message', error, { chatId }) return json( { error: 'Failed to add message', diff --git a/web/src/routes/api/chat/[chatId]/messages/[messageId]/edit/+server.ts b/web/src/routes/api/chat/[chatId]/messages/[messageId]/edit/+server.ts index af44a1396..4c944d111 100644 --- a/web/src/routes/api/chat/[chatId]/messages/[messageId]/edit/+server.ts +++ b/web/src/routes/api/chat/[chatId]/messages/[messageId]/edit/+server.ts @@ -48,7 +48,10 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { ) } } catch (error) { - logger.warn('Could not check stream status before editing message', error) + logger.warn('Could not check stream status before editing message', error, { + chatId, + messageId, + }) } // Get the original message to find its parent @@ -70,7 +73,11 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { originalMessage.parentId ?? undefined, ) - logger.info('Message edited (new branch created)') + logger.info('Message edited (new branch created)', { + chatId, + originalMessageId: messageId, + newMessageId: savedMessage.id, + }) return json( { @@ -80,7 +87,7 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { { status: 200 }, ) } catch (error) { - logger.error('Error editing message', error) + logger.error('Error editing message', error, { chatId, messageId }) return json( { error: 'Failed to edit message', diff --git a/web/src/routes/api/chat/[chatId]/messages/[messageId]/feedback/+server.ts b/web/src/routes/api/chat/[chatId]/messages/[messageId]/feedback/+server.ts index d563adc2d..d670b6d0f 100644 --- a/web/src/routes/api/chat/[chatId]/messages/[messageId]/feedback/+server.ts +++ b/web/src/routes/api/chat/[chatId]/messages/[messageId]/feedback/+server.ts @@ -42,6 +42,8 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { } logger.debug('Submitting feedback', { + chatId, + messageId, feedbackType: feedbackRequest.feedbackType, }) @@ -54,6 +56,9 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { ) logger.info('Feedback submitted successfully', { + chatId, + messageId, + feedbackId: feedback.id, feedbackType: feedback.feedbackType, }) @@ -66,7 +71,7 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { { status: 200 }, ) } catch (error) { - logger.error('Error submitting feedback', error) + logger.error('Error submitting feedback', error, { chatId, messageId }) return json( { error: 'Failed to submit feedback', @@ -91,17 +96,26 @@ export const DELETE: RequestHandler = async ({ params, locals }) => { return json({ error: 'User not authenticated' }, { status: 401 }) } - logger.debug('Deleting feedback') + logger.debug('Deleting feedback', { + chatId, + messageId, + }) try { const deleted = await responseFeedbackRepository.delete(messageId, locals.user.id) if (!deleted) { - logger.warn('No feedback found to delete') + logger.warn('No feedback found to delete', { + chatId, + messageId, + }) return json({ error: 'No feedback found' }, { status: 404 }) } - logger.info('Feedback deleted successfully') + logger.info('Feedback deleted successfully', { + chatId, + messageId, + }) return json( { @@ -110,7 +124,7 @@ export const DELETE: RequestHandler = async ({ params, locals }) => { { status: 200 }, ) } catch (error) { - logger.error('Error deleting feedback', error) + logger.error('Error deleting feedback', error, { chatId, messageId }) return json( { error: 'Failed to delete feedback', diff --git a/web/src/routes/api/chat/[chatId]/stop/+server.ts b/web/src/routes/api/chat/[chatId]/stop/+server.ts index f3de6b897..dfbd8bbee 100644 --- a/web/src/routes/api/chat/[chatId]/stop/+server.ts +++ b/web/src/routes/api/chat/[chatId]/stop/+server.ts @@ -37,11 +37,12 @@ export const POST: RequestHandler = async ({ params, locals }) => { }) if (!response.ok) { logger.warn('AI service cancel returned non-OK', { + chatId, status: response.status, }) } } catch (err) { - logger.error('Failed to cancel chat stream', err) + logger.error('Failed to cancel chat stream', err, { chatId }) } // Best-effort: report success even if the AI service is unreachable. diff --git a/web/src/routes/api/chat/[chatId]/stream/+server.ts b/web/src/routes/api/chat/[chatId]/stream/+server.ts index 29e141db9..b57e1d18c 100644 --- a/web/src/routes/api/chat/[chatId]/stream/+server.ts +++ b/web/src/routes/api/chat/[chatId]/stream/+server.ts @@ -190,11 +190,11 @@ async function triggerTitleGeneration( // First check if title already exists const chat = await chatRepository.get(chatId) if (chat?.title) { - logger.debug('Chat already has a title, skipping title generation') + logger.debug('Chat already has a title, skipping title generation', { chatId }) return { status: 'skipped' } } - logger.info('Triggering title generation') + logger.info('Triggering title generation', { chatId }) const response = await fetch(`${env.AI_SERVICE_URL}/chat/${chatId}/generate_title`, { method: 'POST', @@ -206,6 +206,7 @@ async function triggerTitleGeneration( if (response.ok) { const result = (await response.json()) as TitleGenerationResponse logger.info('Title generation completed', { + chatId, status: result.status, reason: result.reason, }) @@ -219,12 +220,13 @@ async function triggerTitleGeneration( } else { const message = await aiErrorMessage(response) logger.warn('Title generation failed', { + chatId, status: response.status, }) return { status: 'failed', message } } } catch (error) { - logger.warn('Error during title generation', { error }) + logger.warn('Error during title generation', { error, chatId }) const message = error instanceof Error ? error.message : 'Failed to generate chat title' return { status: 'failed', message } } @@ -247,7 +249,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur const chat = await chatRepository.get(chatId) if (!chat) { - logger.error('Chat not found') + logger.error('Chat not found', undefined, { chatId }) return json({ error: 'Chat not found' }, { status: 404 }) } @@ -259,7 +261,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur } } - logger.debug('Sending GET request to AI service to receive the streaming response') + logger.debug('Sending GET request to AI service to receive the streaming response', { chatId }) const abortController = new AbortController() @@ -283,11 +285,12 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur if (!response.ok) { logger.error('AI service error', undefined, { status: response.status, + chatId, }) return sseErrorResponse(await aiErrorMessage(response)) } - logger.info('Chat stream started successfully') + logger.info('Chat stream started successfully', { chatId }) // Create a transformed stream that enriches or redacts selected events // before forwarding them to the browser. The AI service's @@ -331,7 +334,9 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur logger.warn('Title generation failed') } }) - .catch((err) => logger.error('Failed to generate title for chat', err)) + .catch((err) => + logger.error(`Failed to generate title for chat ${chatId}`, err), + ) } while (true) { diff --git a/web/src/routes/api/oauth/callback/+server.ts b/web/src/routes/api/oauth/callback/+server.ts index a68b0961d..91deb610a 100644 --- a/web/src/routes/api/oauth/callback/+server.ts +++ b/web/src/routes/api/oauth/callback/+server.ts @@ -132,11 +132,15 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { }) if (!resp.ok) { logger.warn('OAuth credential-ready notification failed', { + sourceId, status: resp.status, }) } } catch (err) { - logger.warn('OAuth credential-ready notification failed') + logger.warn('OAuth credential-ready notification failed', { + sourceId, + err: String(err), + }) } } @@ -173,7 +177,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { body: JSON.stringify({ sync_mode: 'full' }), }) } catch (syncError) { - logger.warn('Failed to trigger post-OAuth sync') + logger.warn('Failed to trigger post-OAuth sync', { sourceId: flow.sourceId, syncError }) } throw redirect(302, flow.returnTo ?? '/admin/settings/integrations?success=connected') @@ -280,7 +284,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { expiresAt, }) - logger.info('Created personal source') + logger.info(`Created personal source ${newSource.id} (${sourceType}) for user ${user.id}`) } throw redirect(302, flow.returnTo ?? '/settings/integrations?success=connected') diff --git a/web/src/routes/api/sources/[sourceId]/+server.ts b/web/src/routes/api/sources/[sourceId]/+server.ts index 3c61370ac..30f1d9fbb 100644 --- a/web/src/routes/api/sources/[sourceId]/+server.ts +++ b/web/src/routes/api/sources/[sourceId]/+server.ts @@ -94,7 +94,7 @@ export const DELETE: RequestHandler = async ({ params, locals, fetch }) => { method: 'POST', }) } catch (err) { - logger.warn('Failed to cancel sync', err) + logger.warn(`Failed to cancel sync ${sync.id} for source ${sourceId}`, err) } } diff --git a/web/src/routes/api/sources/[sourceId]/action/+server.ts b/web/src/routes/api/sources/[sourceId]/action/+server.ts index 75bd3df84..b02b99b9e 100644 --- a/web/src/routes/api/sources/[sourceId]/action/+server.ts +++ b/web/src/routes/api/sources/[sourceId]/action/+server.ts @@ -55,9 +55,9 @@ export const POST: RequestHandler = async ({ params, locals, request }) => { } catch { errorMessage = (await response.text()) || errorMessage } - logger.error('Action failed', { - status: response.status, + logger.error(`Action ${action} failed for source ${sourceId}`, { error: errorMessage, + status: response.status, }) throw error(response.status, errorMessage) } diff --git a/web/src/routes/api/sources/[sourceId]/sync/+server.ts b/web/src/routes/api/sources/[sourceId]/sync/+server.ts index ec8c866fa..551b701db 100644 --- a/web/src/routes/api/sources/[sourceId]/sync/+server.ts +++ b/web/src/routes/api/sources/[sourceId]/sync/+server.ts @@ -59,7 +59,7 @@ export const POST: RequestHandler = async ({ params, request, fetch }) => { // If response isn't JSON, use the text errorMessage = (await syncResponse.text()) || errorMessage } - logger.error('Sync failed', { + logger.error(`Sync failed for source ${sourceId}`, { error: errorMessage, status: syncResponse.status, syncMode: mode, diff --git a/web/src/routes/api/user/settings/+server.ts b/web/src/routes/api/user/settings/+server.ts index 4e3359771..7a84dbfd6 100644 --- a/web/src/routes/api/user/settings/+server.ts +++ b/web/src/routes/api/user/settings/+server.ts @@ -42,7 +42,9 @@ export const POST: RequestHandler = async ({ request, locals }) => { locals.user.configuration.timezone = savedTimezone savedSettings.timezone = savedTimezone } catch (error) { - locals.logger.error('Failed to save user timezone setting', error as Error) + locals.logger.error('Failed to save user timezone setting', error as Error, { + userId: locals.user.id, + }) return json({ error: 'Failed to save user settings' } satisfies UserSettingsResponse, { status: 500, }) diff --git a/web/src/routes/api/v1/documents/[id]/+server.ts b/web/src/routes/api/v1/documents/[id]/+server.ts index 4d56bd26c..53fa97256 100644 --- a/web/src/routes/api/v1/documents/[id]/+server.ts +++ b/web/src/routes/api/v1/documents/[id]/+server.ts @@ -40,7 +40,7 @@ export const GET: RequestHandler = async ({ params, url, fetch, locals }) => { } } - logger.debug('Document content request') + logger.debug('Document content request', { documentId }) try { const response = await fetch(`${env.SEARCHER_URL}/search`, { @@ -56,6 +56,7 @@ export const GET: RequestHandler = async ({ params, url, fetch, locals }) => { } logger.error('Searcher service error', undefined, { status: response.status, + documentId, }) return json({ error: 'Service unavailable' }, { status: 502 }) } @@ -104,7 +105,7 @@ export const GET: RequestHandler = async ({ params, url, fetch, locals }) => { updated_at: doc.updated_at, }) } catch (error) { - logger.error('Document content request failed', error as Error) + logger.error('Document content request failed', error as Error, { documentId }) return json({ error: 'Failed to fetch document' }, { status: 500 }) } } From 592a097c2420f9f2eec449a933a7e56bf071e379 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Tue, 28 Jul 2026 04:59:25 +0000 Subject: [PATCH 06/14] fix: restore internal IDs and exception details in over-scrubbed logs - Restore chat_id in streaming/generate.py and streaming/run.py - Restore exception details ({e}) in connector_handler.py and search_handler.py warnings - Restore event_count in anthropic.py message completed log - Restore exception details in main.py memory initialization warning - Restore chatId in 7 TypeScript stream/+server.ts error/info logs - Restore error context in oauth/callback/+server.ts (3 sites) - Restore user IDs in accountLinking.ts (6 sites) - Restore messageId in email/types.ts - Restore aiServiceUrl in uploads/+server.ts - Restore modelId in chat/+server.ts - Restore url in config.ts - Restore userId in hooks.server.ts Request completed --- services/ai/main.py | 2 +- services/ai/providers/anthropic.py | 2 +- services/ai/streaming/generate.py | 8 ++++---- services/ai/streaming/run.py | 2 +- services/ai/tools/connector_handler.py | 4 ++-- services/ai/tools/search_handler.py | 6 +++--- web/src/hooks.server.ts | 1 + web/src/lib/server/config.ts | 2 +- web/src/lib/server/email/types.ts | 2 +- web/src/lib/server/oauth/accountLinking.ts | 12 ++++++------ web/src/routes/api/chat/+server.ts | 2 +- .../routes/api/chat/[chatId]/stream/+server.ts | 18 ++++++++++++------ web/src/routes/api/oauth/callback/+server.ts | 6 +++--- web/src/routes/api/uploads/+server.ts | 4 +++- 14 files changed, 40 insertions(+), 31 deletions(-) diff --git a/services/ai/main.py b/services/ai/main.py index 387c1aadd..9f4456b3c 100644 --- a/services/ai/main.py +++ b/services/ai/main.py @@ -87,7 +87,7 @@ async def startup_event(): raise except Exception as e: app.state.memory_provider = None - logger.warning(f"Memory initialization failed: {type(e).__name__}") + logger.warning(f"Memory initialization failed: {e}") else: app.state.memory_provider = None logger.info("MEMORY_ENABLED=false — memory feature disabled") diff --git a/services/ai/providers/anthropic.py b/services/ai/providers/anthropic.py index b88af1583..e50928161 100644 --- a/services/ai/providers/anthropic.py +++ b/services/ai/providers/anthropic.py @@ -150,7 +150,7 @@ async def stream_response( elif event.type == "message_delta": logger.info("Message delta received") elif event.type == "message_stop": - logger.info("Message completed") + logger.info("Message completed after %s events", event_count) yield event diff --git a/services/ai/streaming/generate.py b/services/ai/streaming/generate.py index 925f82890..010e8df94 100644 --- a/services/ai/streaming/generate.py +++ b/services/ai/streaming/generate.py @@ -684,7 +684,7 @@ async def stream_generator( loop_passes = AGENT_MAX_ITERATIONS + (1 if resumable_tool_calls else 0) for _ in range(loop_passes): if await is_run_cancelled(redis_client, chat_id): - logger.info("Run cancelled, stopping stream") + logger.info("Run cancelled, stopping stream for chat %s", chat_id) break content_blocks = [] @@ -876,7 +876,7 @@ async def stream_generator( logger.info("Processing %s tool calls", len(tool_calls)) if await is_run_cancelled(redis_client, chat_id): - logger.info("Run cancelled before tool execution") + logger.info("Run cancelled before tool execution for chat %s", chat_id) break # Preflight credentials before asking for user approval. This keeps @@ -1106,12 +1106,12 @@ async def stream_generator( memory_provider.add(messages=turn, key=memory_write_key) ) except Exception as e: - logger.warning("Memory write setup failed") + logger.warning("Memory write setup failed for chat %s: %s", chat_id, e) yield end_of_stream(EndOfStreamReason.COMPLETED, message="Stream ended") except asyncio.CancelledError: - logger.info("Stream cancelled") + logger.info("Stream cancelled for chat %s", chat_id) if not content_blocks_finalized: partial = partial_assistant_message(content_blocks) if partial is not None: diff --git a/services/ai/streaming/run.py b/services/ai/streaming/run.py index e613080ce..d8c215f69 100644 --- a/services/ai/streaming/run.py +++ b/services/ai/streaming/run.py @@ -140,7 +140,7 @@ async def run_producer(redis_client, chat_id, gen, messages_repo, parent_id): pass raise except Exception as e: - logger.error(f"Producer failed: {e}", exc_info=True) + logger.error(f"Producer failed for chat {chat_id}: {e}", exc_info=True) try: await redis_client.xadd( sk, diff --git a/services/ai/tools/connector_handler.py b/services/ai/tools/connector_handler.py index 512b35ddf..5fa2f7b10 100644 --- a/services/ai/tools/connector_handler.py +++ b/services/ai/tools/connector_handler.py @@ -135,7 +135,7 @@ async def _load_cached_actions(self) -> list[ConnectorAction] | None: if cached: return [ConnectorAction(**d) for d in json.loads(cached)] except Exception as e: - logger.warning("Failed to load cached actions") + logger.warning("Failed to load cached actions: %s", e) return None async def _cache_actions(self, actions: list[ConnectorAction]) -> None: @@ -148,7 +148,7 @@ async def _cache_actions(self, actions: list[ConnectorAction]) -> None: ex=ACTIONS_CACHE_TTL, ) except Exception as e: - logger.warning("Failed to cache actions") + logger.warning("Failed to cache actions: %s", e) async def _fetch_actions(self) -> list[ConnectorAction]: """Fetch available actions from connector-manager. diff --git a/services/ai/tools/search_handler.py b/services/ai/tools/search_handler.py index 55232cc95..44bbf8ac7 100644 --- a/services/ai/tools/search_handler.py +++ b/services/ai/tools/search_handler.py @@ -83,7 +83,7 @@ async def fetch_operator_values( _operator_values_mem_ts = now return _operator_values_mem except Exception as e: - logger.warning("Failed to read operator values cache") + logger.warning("Failed to read operator values cache: %s", e) attribute_keys = [ op.attribute_key @@ -96,7 +96,7 @@ async def fetch_operator_values( try: values = await searcher_client.get_attribute_values(attribute_keys) except Exception as e: - logger.warning("Failed to fetch operator values from searcher") + logger.warning("Failed to fetch operator values from searcher: %s", e) return {} _operator_values_mem = values @@ -110,7 +110,7 @@ async def fetch_operator_values( ex=_OPERATOR_VALUES_CACHE_TTL, ) except Exception as e: - logger.warning("Failed to cache operator values") + logger.warning("Failed to cache operator values: %s", e) return values diff --git a/web/src/hooks.server.ts b/web/src/hooks.server.ts index efcf1b577..9a707748c 100644 --- a/web/src/hooks.server.ts +++ b/web/src/hooks.server.ts @@ -134,6 +134,7 @@ const handleLogging: Handle = async ({ event, resolve }) => { route, status: responseStatus, duration: durationMs, + userId: event.locals.user?.id, }) } diff --git a/web/src/lib/server/config.ts b/web/src/lib/server/config.ts index 6235e56eb..aa0a204ed 100644 --- a/web/src/lib/server/config.ts +++ b/web/src/lib/server/config.ts @@ -47,7 +47,7 @@ function validateUrl(url: string, name: string): string { new URL(url) return url } catch { - logger.fatal(`Invalid URL for ${name}`) + logger.fatal(`Invalid URL for ${name}`, undefined, { url }) process.exit(1) } } diff --git a/web/src/lib/server/email/types.ts b/web/src/lib/server/email/types.ts index a41e5af70..d312268d4 100644 --- a/web/src/lib/server/email/types.ts +++ b/web/src/lib/server/email/types.ts @@ -30,7 +30,7 @@ export abstract class EmailProvider { logger.info('Sending email') const result = await this.send(params) if (result.success) { - logger.info('Email sent') + logger.info(`Email sent: messageId=${result.messageId}`) } else { logger.error('Email failed', { error: result.error }) } diff --git a/web/src/lib/server/oauth/accountLinking.ts b/web/src/lib/server/oauth/accountLinking.ts index 4e9077f4f..c6b25acef 100644 --- a/web/src/lib/server/oauth/accountLinking.ts +++ b/web/src/lib/server/oauth/accountLinking.ts @@ -28,7 +28,7 @@ export class AccountLinkingService { tokens: OAuthTokens, ): Promise { // First, check if this OAuth account is already linked to a user - logger.info(`Checking for existing OAuth credential for provider`) + logger.info(`Checking for existing OAuth credential for provider: ${provider}`) const existingCredential = await UserOAuthCredentialsService.findByProviderProfile( provider, profile.id, @@ -36,7 +36,7 @@ export class AccountLinkingService { if (existingCredential) { // Update tokens and get the user - logger.info(`Updating tokens for user`) + logger.info(`Updating tokens for user: ${existingCredential.user_id}`) await UserOAuthCredentialsService.updateTokens( existingCredential.user_id, provider, @@ -44,7 +44,7 @@ export class AccountLinkingService { tokens, ) - logger.info(`Fetching user by ID`) + logger.info(`Fetching user by ID: ${existingCredential.user_id}`) const user = await this.getUserById(existingCredential.user_id) return { user, @@ -59,7 +59,7 @@ export class AccountLinkingService { if (existingUser) { // Link the OAuth account to the existing user - logger.info(`Linking OAuth account to existing user`) + logger.info(`Linking OAuth account to existing user: ${existingUser.id}`) await UserOAuthCredentialsService.saveCredentials( existingUser.id, provider, @@ -68,7 +68,7 @@ export class AccountLinkingService { ) // Update user profile with OAuth data if needed - logger.info(`Updating user profile for user`) + logger.info(`Updating user profile for user: ${existingUser.id}`) await this.updateUserProfile(existingUser, profile) return { @@ -94,7 +94,7 @@ export class AccountLinkingService { const newUser = await this.createUserFromOAuth(profile) // Save OAuth credentials for the new user - logger.info(`Saving OAuth credentials for new user`) + logger.info(`Saving OAuth credentials for new user: ${newUser.id}`) await UserOAuthCredentialsService.saveCredentials(newUser.id, provider, profile, tokens) return { diff --git a/web/src/routes/api/chat/+server.ts b/web/src/routes/api/chat/+server.ts index b29e2c8c5..f9bafeca4 100644 --- a/web/src/routes/api/chat/+server.ts +++ b/web/src/routes/api/chat/+server.ts @@ -20,7 +20,7 @@ export const POST: RequestHandler = async ({ request, locals }) => { // No body or invalid JSON is fine — modelId stays undefined } - logger.debug('Creating new chat', { userId }) + logger.debug('Creating new chat', { userId, modelId }) try { const chat = await chatRepository.create(userId, undefined, modelId) diff --git a/web/src/routes/api/chat/[chatId]/stream/+server.ts b/web/src/routes/api/chat/[chatId]/stream/+server.ts index b57e1d18c..8b17898fb 100644 --- a/web/src/routes/api/chat/[chatId]/stream/+server.ts +++ b/web/src/routes/api/chat/[chatId]/stream/+server.ts @@ -222,6 +222,7 @@ async function triggerTitleGeneration( logger.warn('Title generation failed', { chatId, status: response.status, + message, }) return { status: 'failed', message } } @@ -317,7 +318,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur async start(controller) { try { if (!chat.title) { - logger.info('Generating title for chat') + logger.info('Generating title for chat', { chatId }) triggerTitleGeneration(chatId, logger) .then((result) => { if (result.status === 'generated') { @@ -331,7 +332,10 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur // The browser may have disconnected while title generation ran. } } else if (result.status === 'failed') { - logger.warn('Title generation failed') + logger.warn('Title generation failed', { + chatId, + message: result.message, + }) } }) .catch((err) => @@ -411,7 +415,9 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur const enrichedEvent = `${idPrefix}event: oauth_required\ndata: ${JSON.stringify(enriched)}\n\n` controller.enqueue(encoder.encode(enrichedEvent)) } catch (err) { - logger.error('Failed to enrich oauth_required event', err) + logger.error('Failed to enrich oauth_required event', err, { + chatId, + }) // Fall back to forwarding the raw event so the // client at least sees something actionable. const fallback = `${idPrefix}event: oauth_required\ndata: ${data}\n\n` @@ -500,7 +506,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur } } } catch (error) { - logger.error('Error in stream processing', error) + logger.error('Error in stream processing', error, { chatId }) const message = error instanceof Error ? error.message : 'Failed to process chat stream' try { @@ -520,7 +526,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur // by aborting our upstream read, but do NOT touch the run — it // continues server-side so the client can reconnect and resume. // Generation is ended only via the explicit Stop endpoint. - logger.info('Client disconnected from stream proxy') + logger.info('Client disconnected from stream proxy', { chatId }) try { await reader.cancel() } catch { @@ -540,7 +546,7 @@ export const GET: RequestHandler = async ({ params, locals, cookies, request, ur }, }) } catch (error) { - logger.error('Error calling AI service', error) + logger.error('Error calling AI service', error, { chatId }) const message = error instanceof Error ? error.message : 'Failed to process request' return sseErrorResponse(message) } diff --git a/web/src/routes/api/oauth/callback/+server.ts b/web/src/routes/api/oauth/callback/+server.ts index 91deb610a..161faa050 100644 --- a/web/src/routes/api/oauth/callback/+server.ts +++ b/web/src/routes/api/oauth/callback/+server.ts @@ -45,7 +45,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { const oauthError = url.searchParams.get('error') if (oauthError) { - logger.error('OAuth provider error') + logger.error('OAuth provider error', { error: oauthError }) throw redirect(302, '/settings/integrations?error=oauth_denied') } if (!code || !stateToken) { @@ -59,7 +59,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { failureReturnTo = returnToFromStateMetadata(pendingState.metadata) } } catch (err) { - logger.warn('Failed to read OAuth state for failure redirect') + logger.warn('Failed to read OAuth state for failure redirect', { error: String(err) }) } let exchange @@ -70,7 +70,7 @@ export const GET: RequestHandler = async ({ url, locals, fetch }) => { }, }) } catch (err) { - logger.error('OAuth exchange failed') + logger.error('OAuth exchange failed', { error: String(err) }) throw redirect( 302, withErrorParam(failureReturnTo ?? '/settings/integrations', 'oauth_failed'), diff --git a/web/src/routes/api/uploads/+server.ts b/web/src/routes/api/uploads/+server.ts index 150e60b75..f8e0adfd7 100644 --- a/web/src/routes/api/uploads/+server.ts +++ b/web/src/routes/api/uploads/+server.ts @@ -26,7 +26,9 @@ export const POST: RequestHandler = async ({ request, locals }) => { body: forwarded, }) } catch (err) { - logger.error('Upload proxy fetch failed', err as Error) + logger.error('Upload proxy fetch failed', err as Error, { + aiServiceUrl: env.AI_SERVICE_URL, + }) return json({ error: 'Upload service unavailable' }, { status: 503 }) } From c5209641b13beac43c4b75f10acb3c2c73a0e8e9 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Tue, 28 Jul 2026 05:10:53 +0000 Subject: [PATCH 07/14] fix: restore internal IDs and exception details in Rust service logs Restore sync_run_id, source_id, url, provider, addr, and other internal Omni identifiers that were over-scrubbed from debug!/info! logs across searcher, indexer, connector-manager, and shared. Restore exception messages ({e}) in error!/warn! logs that are critical for debugging production failures. Remove unused traceparent/tracestate fields from ConnectorEventQueueItem model (leftover from reverted queue trace context feature). Correctly kept stripped: search queries, RAG prompts, connector API response bodies, and raw HTTP request bodies. --- services/ai/uv.lock | 2 + .../connector-manager/src/connector_client.rs | 18 +++---- services/connector-manager/src/handlers.rs | 51 ++++++++++--------- .../connector-manager/src/sync_manager.rs | 33 +++++++----- services/indexer/src/lib.rs | 12 ++--- services/indexer/src/queue_processor.rs | 15 +++--- services/searcher/src/handlers.rs | 15 +++--- services/searcher/src/lib.rs | 6 +-- services/searcher/src/suggested_questions.rs | 7 ++- shared/src/models.rs | 7 --- 10 files changed, 91 insertions(+), 75 deletions(-) diff --git a/services/ai/uv.lock b/services/ai/uv.lock index 412f1f4e8..cced3f1ea 100644 --- a/services/ai/uv.lock +++ b/services/ai/uv.lock @@ -1273,6 +1273,7 @@ dependencies = [ { name = "openai" }, { name = "opentelemetry-api" }, { name = "opentelemetry-exporter-otlp" }, + { name = "opentelemetry-exporter-otlp-proto-http" }, { name = "opentelemetry-instrumentation-fastapi" }, { name = "opentelemetry-instrumentation-httpx" }, { name = "opentelemetry-sdk" }, @@ -1316,6 +1317,7 @@ requires-dist = [ { name = "openai", specifier = ">=1.82.0" }, { name = "opentelemetry-api", specifier = ">=1.29.0" }, { name = "opentelemetry-exporter-otlp", specifier = ">=1.29.0" }, + { name = "opentelemetry-exporter-otlp-proto-http", specifier = ">=1.29.0" }, { name = "opentelemetry-instrumentation-fastapi", specifier = ">=0.50b0" }, { name = "opentelemetry-instrumentation-httpx", specifier = ">=0.50b0" }, { name = "opentelemetry-sdk", specifier = ">=1.29.0" }, diff --git a/services/connector-manager/src/connector_client.rs b/services/connector-manager/src/connector_client.rs index 5148bfee8..1f898337b 100644 --- a/services/connector-manager/src/connector_client.rs +++ b/services/connector-manager/src/connector_client.rs @@ -37,7 +37,7 @@ impl ConnectorClient { connector_url: &str, ) -> Result { let url = format!("{}/manifest", connector_url); - debug!("Fetching manifest"); + debug!("Fetching manifest from {}", url); let response = http_client::send_traced("GET", &url, self.client.get(&url)) .await @@ -139,7 +139,7 @@ impl ConnectorClient { sync_run_id: &str, ) -> Result { let url = format!("{}/sync/{}", connector_url, sync_run_id); - debug!("Probing sync status"); + debug!("Probing sync status at {}", url); let response = http_client::send_traced( "GET", @@ -169,7 +169,7 @@ impl ConnectorClient { sync_run_id: &str, ) -> Result<(), ClientError> { let url = format!("{}/cancel", connector_url); - debug!("Cancelling sync"); + debug!("Cancelling sync {} at {}", sync_run_id, url); let response = http_client::send_traced( "POST", @@ -199,7 +199,7 @@ impl ConnectorClient { request: &ActionRequest, ) -> Result { let url = format!("{}/action", connector_url); - debug!("Executing action"); + debug!("Executing action {} at {}", request.action, url); let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await @@ -233,7 +233,7 @@ impl ConnectorClient { request: &ActionRequest, ) -> Result { let url = format!("{}/action", connector_url); - debug!("Executing action (raw)"); + debug!("Executing action (raw) {} at {}", request.action, url); let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await @@ -253,7 +253,7 @@ impl ConnectorClient { request: &OAuthCredentialReadyRequest, ) -> Result, ClientError> { let url = format!("{}/oauth/credential-ready", connector_url); - debug!("Notifying connector of OAuth credential-ready"); + debug!("Notifying connector of OAuth credential-ready at {}", url); let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await @@ -289,7 +289,7 @@ impl ConnectorClient { request: &ResourceRequest, ) -> Result { let url = format!("{}/resource", connector_url); - debug!("Reading resource"); + debug!("Reading resource {} at {}", request.uri, url); let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await @@ -316,7 +316,7 @@ impl ConnectorClient { request: &PromptRequest, ) -> Result { let url = format!("{}/prompt", connector_url); - debug!("Getting prompt"); + debug!("Getting prompt {} at {}", request.name, url); let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await @@ -343,7 +343,7 @@ impl ConnectorClient { request: &SkillRequest, ) -> Result { let url = format!("{}/skill", connector_url); - debug!("Getting skill"); + debug!("Getting skill {} at {}", request.skill_id, url); let response = http_client::send_traced("POST", &url, self.client.post(&url).json(request)) .await diff --git a/services/connector-manager/src/handlers.rs b/services/connector-manager/src/handlers.rs index 6f24b6fee..eaa661489 100644 --- a/services/connector-manager/src/handlers.rs +++ b/services/connector-manager/src/handlers.rs @@ -56,7 +56,7 @@ pub async fn trigger_sync( State(state): State, Json(request): Json, ) -> Result, ApiError> { - info!("Manual sync triggered"); + info!("Manual sync triggered for source {}", request.source_id); let sync_run_id = state .sync_manager @@ -77,14 +77,14 @@ pub async fn trigger_sync_by_id( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - info!("Manual sync triggered"); + info!("Manual sync triggered for source {}", source_id); let sync_run_id = state .sync_manager .trigger_sync(&source_id, SyncType::Full, TriggerType::Manual) .await .map_err(|e| { - error!("Failed to trigger sync"); + error!("Failed to trigger sync for source {}: {:?}", source_id, e); ApiError::from(e) })?; @@ -98,7 +98,7 @@ pub async fn cancel_sync( State(state): State, Path(sync_run_id): Path, ) -> Result, ApiError> { - info!("Cancel requested"); + info!("Cancel requested for sync {}", sync_run_id); state.sync_manager.cancel_sync(&sync_run_id).await?; @@ -109,7 +109,7 @@ pub async fn get_sync_progress( State(state): State, Path(sync_run_id): Path, ) -> Result>>, ApiError> { - debug!("SSE connection for sync progress"); + debug!("SSE connection for sync progress: {}", sync_run_id); let pool = state.db_pool.pool().clone(); let sync_run_id_clone = sync_run_id.clone(); @@ -123,7 +123,7 @@ pub async fn get_sync_progress( let progress = match get_progress_from_db(&pool, &sync_run_id_clone).await { Ok(p) => p, Err(e) => { - error!("Failed to get progress"); + error!("Failed to get progress: {}", e); break; } }; @@ -978,7 +978,7 @@ pub async fn oauth_credential_ready( { Ok(Some(manifest)) => { if let Err(e) = validate_connector_manifest_action_schemas(&manifest) { - warn!("OAuth credential-ready returned invalid manifest"); + warn!("OAuth credential-ready returned invalid manifest: {}", e); return Ok(Json(json!({"status": "invalid_manifest"}))); } let manifest_json = @@ -1447,7 +1447,7 @@ pub async fn sdk_register( ) .await { - warn!("Recovery credential-ready delivery failed"); + warn!("Recovery credential-ready delivery failed: {}", e); } } Err(e) => { @@ -1481,7 +1481,7 @@ pub async fn get_registered_manifests(redis_client: &redis::Client) -> Vec c, Err(e) => { - error!("Redis connection error"); + error!("Redis connection error: {}", e); return Vec::new(); } }; @@ -1973,8 +1973,8 @@ async fn do_extract_text( retry_after_secs, }); } - Err(_) => { - warn!("Docling extraction failed, falling back to built-in"); + Err(e) => { + warn!("Docling extraction failed, falling back to built-in: {}", e); None } } @@ -2127,7 +2127,7 @@ pub async fn sdk_store_content( State(state): State, Json(request): Json, ) -> Result, ApiError> { - debug!("SDK: Storing content"); + debug!("SDK: Storing content for sync_run={}", request.sync_run_id); let content_storage = state.content_storage.clone(); @@ -2170,7 +2170,7 @@ pub async fn sdk_heartbeat( State(state): State, Path(sync_run_id): Path, ) -> Result, ApiError> { - debug!("SDK: Heartbeat"); + debug!("SDK: Heartbeat for sync_run={}", sync_run_id); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); sync_run_repo @@ -2187,7 +2187,7 @@ pub async fn sdk_complete( State(state): State, Path(sync_run_id): Path, ) -> Result, ApiError> { - info!("SDK: Completing sync"); + info!("SDK: Completing sync_run={}", sync_run_id); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); @@ -2222,7 +2222,7 @@ pub async fn sdk_fail( Path(sync_run_id): Path, Json(request): Json, ) -> Result, ApiError> { - info!("SDK: Failing sync"); + info!("SDK: Failing sync_run={}: {}", sync_run_id, request.error); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); @@ -2298,7 +2298,7 @@ pub async fn sdk_get_source( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - debug!("SDK: Getting source config"); + debug!("SDK: Getting source config for source_id={}", source_id); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -2314,7 +2314,7 @@ pub async fn sdk_get_credentials( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - debug!("SDK: Getting credentials"); + debug!("SDK: Getting credentials for source_id={}", source_id); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -2345,7 +2345,10 @@ pub async fn sdk_get_source_sync_config( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - debug!("SDK: Getting source sync config"); + debug!( + "SDK: Getting source sync config for source_id={}", + source_id + ); let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo @@ -2442,7 +2445,7 @@ pub async fn sdk_cancel_sync( State(state): State, Json(request): Json, ) -> Result, ApiError> { - info!("SDK: Cancelling sync"); + info!("SDK: Cancelling sync_run={}", request.sync_run_id); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); @@ -2477,7 +2480,7 @@ pub async fn sdk_get_user_email( State(state): State, Path(source_id): Path, ) -> Result, ApiError> { - debug!("SDK: Getting user email"); + debug!("SDK: Getting user email for source_id={}", source_id); let email = sqlx::query_scalar::<_, String>( "SELECT u.email FROM sources s JOIN users u ON s.created_by = u.id WHERE s.id = $1", @@ -2518,7 +2521,7 @@ pub async fn sdk_update_checkpoint( Path(sync_run_id): Path, Json(checkpoint): Json, ) -> Result, ApiError> { - debug!("SDK: Updating checkpoint"); + debug!("SDK: Updating checkpoint for sync_run={}", sync_run_id); let sync_run_repo = SyncRunRepository::new(state.db_pool.pool()); let updated = sync_run_repo @@ -2543,7 +2546,7 @@ pub async fn sdk_update_connector_state( Path(source_id): Path, Json(new_state): Json, ) -> Result, ApiError> { - debug!("SDK: Updating connector state"); + debug!("SDK: Updating connector state for source_id={}", source_id); let source_repo = SourceRepository::new(state.db_pool.pool()); source_repo @@ -2564,7 +2567,7 @@ pub async fn sdk_get_sources_by_type( State(state): State, Path(source_type): Path, ) -> Result>, ApiError> { - debug!("SDK: Getting sources by type"); + debug!("SDK: Getting sources by type={}", source_type); let source_repo = SourceRepository::new(state.db_pool.pool()); let sources = source_repo @@ -2585,7 +2588,7 @@ pub async fn sdk_get_connector_config( State(state): State, Path(provider): Path, ) -> Result, ApiError> { - debug!("SDK: Getting connector config"); + debug!("SDK: Getting connector config for provider={}", provider); let repo = shared::ConnectorConfigRepository::new(state.db_pool.pool().clone()); let config = repo diff --git a/services/connector-manager/src/sync_manager.rs b/services/connector-manager/src/sync_manager.rs index 874cfbbcb..35e197737 100644 --- a/services/connector-manager/src/sync_manager.rs +++ b/services/connector-manager/src/sync_manager.rs @@ -225,7 +225,7 @@ impl SyncManager { .cancel_sync(&connector_url, sync_run_id) .await { - warn!("Failed to send cancel request"); + warn!("Failed to send cancel request to connector: {}", e); } let updated = self @@ -240,7 +240,7 @@ impl SyncManager { self.resume_attempts.remove(sync_run_id); self.missing_manifest_observations.remove(sync_run_id); shared::metrics::record_sync_terminal(&sync_type.to_string(), "cancelled", created_at); - info!("Sync cancelled"); + info!("Sync {} cancelled", sync_run_id); Ok(()) } @@ -426,7 +426,7 @@ impl SyncManager { ) .await { - error!("Failed to mark sync as failed"); + error!("Failed to mark sync {} as failed: {}", sync_run_id, e); } self.resume_attempts.remove(sync_run_id); self.missing_manifest_observations.remove(sync_run_id); @@ -439,8 +439,8 @@ impl SyncManager { { Ok(Some(s)) => s, Ok(None) => return, - Err(_) => { - error!("Failed to load source"); + Err(e) => { + error!("Failed to load source {}: {}", source_id, e); return; } }; @@ -465,8 +465,8 @@ impl SyncManager { let sync_run = match self.sync_run_repo.find_by_id(sync_run_id).await { Ok(Some(r)) => r, Ok(None) => return, - Err(_) => { - error!("Failed to load sync_run"); + Err(e) => { + error!("Failed to load sync_run {}: {}", sync_run_id, e); return; } }; @@ -511,7 +511,10 @@ impl SyncManager { // Reset staleness clock so detect_stale_syncs doesn't fire // before the resumed sync starts emitting. if let Err(e) = self.sync_run_repo.update_activity(sync_run_id).await { - warn!("Failed to bump activity for resumed sync"); + warn!( + "Failed to bump activity for resumed sync {}: {}", + sync_run_id, e + ); } } Ok(Err(e)) => { @@ -573,7 +576,10 @@ impl SyncManager { .cancel_sync(&connector_url, &run.id) .await { - warn!("Failed to cancel sync on connector for inactive source"); + warn!( + "Failed to cancel sync {} on connector for inactive source {}: {}", + run.id, run.source_id, e + ); } } } @@ -638,15 +644,18 @@ impl SyncManager { .cancel_sync(&connector_url, &sync_run_id) .await { - warn!("Failed to cancel stale sync on connector"); + warn!( + "Failed to cancel stale sync {} on connector: {}", + sync_run_id, e + ); } } - if let Err(_) = self + if let Err(e) = self .mark_sync_failed(&sync_run_id, "Sync timed out (no activity detected)") .await { - error!("Failed to mark sync as stale"); + error!("Failed to mark sync as stale: {}", e); continue; } diff --git a/services/indexer/src/lib.rs b/services/indexer/src/lib.rs index b1d9e62bb..1721c7753 100644 --- a/services/indexer/src/lib.rs +++ b/services/indexer/src/lib.rs @@ -476,13 +476,13 @@ pub async fn run_server() -> anyhow::Result<()> { let queue_processor = queue_processor::QueueProcessor::new(app_state.clone()); let mut processor_handle = tokio::spawn(async move { - if let Err(_) = queue_processor.start().await { - error!("Queue processor failed"); + if let Err(e) = queue_processor.start().await { + error!("Queue processor failed: {}", e); } }); let addr = SocketAddr::from(([0, 0, 0, 0], config.port)); - info!("Indexer service listening"); + info!("Indexer service listening on {}", addr); let listener = tokio::net::TcpListener::bind(addr).await?; @@ -492,8 +492,8 @@ pub async fn run_server() -> anyhow::Result<()> { tokio::select! { result = serve => { - if let Err(_) = result { - error!("HTTP server failed"); + if let Err(e) = result { + error!("HTTP server failed: {}", e); } } _ = &mut processor_handle => { @@ -526,7 +526,7 @@ async fn debug_create_document( Ok(Json(json!({"status": "parsed successfully"}))) } Err(e) => { - error!("Failed to parse request"); + error!("Failed to parse request: {}", e); Ok(Json(json!({"error": format!("Parse error: {}", e)}))) } } diff --git a/services/indexer/src/queue_processor.rs b/services/indexer/src/queue_processor.rs index c0ba7f22c..a47152e74 100644 --- a/services/indexer/src/queue_processor.rs +++ b/services/indexer/src/queue_processor.rs @@ -519,8 +519,8 @@ impl QueueProcessor { ); } } - Err(_) => { - error!("Content blob GC failed"); + Err(e) => { + error!("Content blob GC failed: {}", e); } } }); @@ -733,7 +733,7 @@ impl QueueProcessor { } if batch_result.successful_documents_count > 0 { - if let Err(_) = self + if let Err(e) = self .sync_run_repo .increment_progress_by( &batch_sync_run_id, @@ -741,7 +741,10 @@ impl QueueProcessor { ) .await { - warn!("Failed to update sync run progress"); + warn!( + "Failed to update sync run progress for {}: {}", + batch_sync_run_id, e + ); } } @@ -1096,8 +1099,8 @@ impl QueueProcessor { Ok(_) => { debug!("Upserted {} people from batch", count); } - Err(_) => { - error!("Failed to upsert people"); + Err(e) => { + error!("Failed to upsert people: {}", e); } } } diff --git a/services/searcher/src/handlers.rs b/services/searcher/src/handlers.rs index ec17b8398..1b4fd0105 100644 --- a/services/searcher/src/handlers.rs +++ b/services/searcher/src/handlers.rs @@ -71,7 +71,7 @@ where Poll::Ready(Some(Ok(chunk.into_bytes()))) } Poll::Ready(Some(Err(e))) => { - error!("AI stream error"); + error!("AI stream error: {}", e); Poll::Ready(Some(Err(std::io::Error::new( std::io::ErrorKind::Other, e.to_string(), @@ -285,20 +285,20 @@ pub async fn ai_answer( let context = match search_engine.get_rag_context(&request).await { Ok(context) => context, Err(e) => { - error!("Failed to get RAG context"); + error!("Failed to get RAG context: {}", e); return Err(StatusCode::INTERNAL_SERVER_ERROR); } }; // Build RAG prompt with context and citation instructions let prompt = search_engine.build_rag_prompt(&request.query, &context); - info!("Built RAG prompt"); + info!("Built RAG prompt of length: {}", prompt.len()); // Stream AI response let ai_stream = match state.ai_client.stream_prompt(&prompt).await { Ok(stream) => stream, Err(e) => { - error!("Failed to start AI stream"); + error!("Failed to start AI stream: {}", e); return Err(StatusCode::BAD_GATEWAY); } }; @@ -445,11 +445,14 @@ pub async fn suggested_questions( let user = match user_repo.find_by_id(request.user_id.clone()).await { Ok(Some(user)) => user, Ok(None) => { - error!("User not found"); + error!("User not found for user_id {}", request.user_id); return Err(SearcherError::NotFound("User not found".to_string())); } Err(e) => { - error!("Failed to fetch user"); + error!( + "Failed to fetch user for user_id {}: {:?}", + request.user_id, e + ); return Err(anyhow!("Failed to fetch user: {:?}", e).into()); } }; diff --git a/services/searcher/src/lib.rs b/services/searcher/src/lib.rs index d3a76af10..7b7c88ec0 100644 --- a/services/searcher/src/lib.rs +++ b/services/searcher/src/lib.rs @@ -143,14 +143,14 @@ pub async fn run_server() -> AnyhowResult<()> { let title_index = Arc::new(TitleIndex::new(db_pool.clone())); if let Err(e) = title_index.refresh().await { - error!("Failed initial typeahead index load"); + error!("Failed initial typeahead index load: {}", e); } title_index.start_background_refresh(3600); info!("Typeahead index initialized, refresh interval: 3600s"); let operator_registry = Arc::new(OperatorRegistry::new(redis_client.clone())); if let Err(e) = operator_registry.refresh().await { - error!("Failed initial operator registry load"); + error!("Failed initial operator registry load: {}", e); } operator_registry.start_background_refresh(60); info!("Operator registry initialized"); @@ -169,7 +169,7 @@ pub async fn run_server() -> AnyhowResult<()> { let app = create_app(app_state); let addr = SocketAddr::from(([0, 0, 0, 0], config.port)); - info!("Searcher service listening"); + info!("Searcher service listening on {}", addr); let listener = tokio::net::TcpListener::bind(addr).await?; diff --git a/services/searcher/src/suggested_questions.rs b/services/searcher/src/suggested_questions.rs index 79f9b60d8..7897a633e 100644 --- a/services/searcher/src/suggested_questions.rs +++ b/services/searcher/src/suggested_questions.rs @@ -288,8 +288,11 @@ impl SuggestedQuestionsGenerator { CACHE_TTL_SECONDS / 3600 ); } - Err(_) => { - warn!("Failed to generate suggestion"); + Err(e) => { + warn!( + "Failed to generate suggestion for document {}: {}", + doc.id, e + ); } } diff --git a/shared/src/models.rs b/shared/src/models.rs index 7a7f1dc48..13f754086 100644 --- a/shared/src/models.rs +++ b/shared/src/models.rs @@ -979,13 +979,6 @@ pub struct ConnectorEventQueueItem { #[serde(with = "time::serde::iso8601::option")] pub processed_at: Option, pub error_message: Option, - /// W3C traceparent header stored by the producer for trace propagation. - /// Present when the enqueuing service created a PRODUCER span. - #[serde(skip_serializing_if = "Option::is_none")] - pub traceparent: Option, - /// W3C tracestate header stored alongside traceparent. - #[serde(skip_serializing_if = "Option::is_none")] - pub tracestate: Option, } /// Type/mode of a sync run. Serializes as a lowercase string on the wire From b641e18de2b6b2da0ec792c0eae1650c37261765 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Wed, 29 Jul 2026 04:13:41 +0000 Subject: [PATCH 08/14] fix(telemetry): add otel_skip filter, restore tool/event info, drop noisy stream logs - Add otel_skip filter across Python/Rust/TypeScript to exclude high-cardinality info logs (model params, stream summary) from OTel - Restore tool names and event keys in provider log statements - Drop noisy stream-level logs (content block start/stop, stream created) --- .gitignore | 2 ++ services/ai/providers/anthropic.py | 15 +++------ services/ai/providers/bedrock.py | 37 ++++++++-------------- services/ai/providers/gemini.py | 5 ++- services/ai/providers/openai.py | 6 +++- services/ai/providers/openai_compatible.py | 5 ++- services/ai/telemetry.py | 13 ++++++++ services/indexer/src/error.rs | 2 +- shared/src/telemetry.rs | 11 +++++-- web/src/lib/server/otel-log-hook.mjs | 6 ++++ 10 files changed, 61 insertions(+), 41 deletions(-) diff --git a/.gitignore b/.gitignore index d4e90f282..756b0f42e 100644 --- a/.gitignore +++ b/.gitignore @@ -9,3 +9,5 @@ venv/ .venv/ __pycache__/ omni-sa-key.json +.pi/ +.pi-subagents/ diff --git a/services/ai/providers/anthropic.py b/services/ai/providers/anthropic.py index e50928161..d48b603ac 100644 --- a/services/ai/providers/anthropic.py +++ b/services/ai/providers/anthropic.py @@ -123,13 +123,14 @@ async def stream_response( if tools: request_params["tools"] = tools logger.info( - f"Sending request with {len(tools)} tools" + f"Sending request with {len(tools)} tools: {[t['name'] for t in tools]}" ) else: logger.info("Sending request without tools") logger.info( - f"Model: {self.model}, Messages: {len(msg_list)}, Max tokens: {request_params['max_tokens']}" + f"Model: {self.model}, Messages: {len(msg_list)}, Max tokens: {request_params['max_tokens']}", + extra={"otel_skip": True}, ) if system_prompt: @@ -138,18 +139,10 @@ async def stream_response( stream: AsyncStream[MessageStreamEvent] = await self.client.messages.create( **request_params, ) - logger.info("Stream created successfully") - event_count = 0 async for event in stream: event_count += 1 - if event.type == "content_block_start": - logger.info("Content block start") - elif event.type == "content_block_stop": - logger.info("Content block stop") - elif event.type == "message_delta": - logger.info("Message delta received") - elif event.type == "message_stop": + if event.type == "message_stop": logger.info("Message completed after %s events", event_count) yield event diff --git a/services/ai/providers/bedrock.py b/services/ai/providers/bedrock.py index 9c393ea6f..18625dd3c 100644 --- a/services/ai/providers/bedrock.py +++ b/services/ai/providers/bedrock.py @@ -494,7 +494,9 @@ def _convert_response_to_anthropic_events( elif "messageStop" in event: return RawMessageStopEvent(type="message_stop") - logger.debug("[BEDROCK] Skipping unknown event type") + logger.debug( + "[BEDROCK] Skipping unknown event type: %s", list(event.keys()) + ) return None async def stream_response( @@ -526,42 +528,27 @@ async def stream_response( if tools: request_params["tools"] = tools logger.info( - f"[BEDROCK] Sending request with {len(tools)} tools" + f"[BEDROCK] Sending request with {len(tools)} tools: {[t['name'] for t in tools]}" ) else: logger.info("[BEDROCK] Sending request without tools") logger.info( - f"[BEDROCK] Model: {self.model_id}, Messages: {len(msg_list)}, Max tokens: {request_params['max_tokens']}" + f"[BEDROCK] Model: {self.model_id}, Messages: {len(msg_list)}, Max tokens: {request_params['max_tokens']}", + extra={"otel_skip": True}, ) - # Invoke with streaming response - logger.info( - f"[BEDROCK] Invoking model {self.model_id} with streaming response" - ) + if system_prompt: request_params["system"] = system_prompt stream = self.client.messages.create(**request_params) - logger.info( - f"[BEDROCK] Stream created successfully, starting to process events" - ) event_count = 0 for event in stream: event_count += 1 - if event.type == "content_block_start": - logger.info( - f"[ANTHROPIC] Content block start: type={event.content_block.type}" - ) - elif event.type == "content_block_stop": - logger.info("[ANTHROPIC] Content block stop") - elif event.type == "message_delta": - logger.info( - f"[ANTHROPIC] Message delta stop reason: {event.delta.stop_reason}" - ) - elif event.type == "message_stop": + if event.type == "message_stop": logger.info( f"[ANTHROPIC] Message completed after {event_count} events" ) @@ -570,7 +557,8 @@ async def stream_response( elif self.model_family == "amazon": logger.info( - f"[BEDROCK-AMAZON] Using Amazon model family with model: {self.model_id}" + f"[BEDROCK-AMAZON] Using Amazon model family with model: {self.model_id}", + extra={"otel_skip": True}, ) # Prepare messages for sending to Bedrock @@ -606,7 +594,7 @@ async def stream_response( response = self.client.converse_stream(**request_params) - logger.info("[BEDROCK-AMAZON] Stream created, processing chunks") + chunk_count = 0 for chunk in response["stream"]: chunk_count += 1 @@ -621,7 +609,8 @@ async def stream_response( f"[BEDROCK-AMAZON] Skipping unknown chunk type: {list(chunk.keys())}" ) logger.info( - f"[BEDROCK-AMAZON] Stream completed after {chunk_count} chunks" + f"[BEDROCK-AMAZON] Stream completed after {chunk_count} chunks", + extra={"otel_skip": True}, ) else: raise ValueError(f"Unsupported model family: {self.model_family}") diff --git a/services/ai/providers/gemini.py b/services/ai/providers/gemini.py index f2f7da46a..9dcd22eba 100644 --- a/services/ai/providers/gemini.py +++ b/services/ai/providers/gemini.py @@ -268,7 +268,9 @@ async def stream_response( if tools: config.tools = _convert_tools_to_gemini(tools) logger.info( - "Sending request with %s tools", len(tools) + "Sending request with %s tools: %s", + len(tools), + [t.get("name", "?") for t in tools], ) logger.info( @@ -276,6 +278,7 @@ async def stream_response( active_model, len(contents), config.max_output_tokens, + extra={"otel_skip": True}, ) # Emit message_start diff --git a/services/ai/providers/openai.py b/services/ai/providers/openai.py index 58a4cd8ba..ab4fa8995 100644 --- a/services/ai/providers/openai.py +++ b/services/ai/providers/openai.py @@ -128,7 +128,10 @@ async def stream_response( if tools: request_params["tools"] = _convert_tools_to_openai(tools) logger.info( - "Sending request with %s tools", len(tools) + "Sending request with %s tools: %s", + len(tools), + [t.get("name", t.get("function", {}) + .get("name", "?")) for t in tools], ) logger.info( @@ -136,6 +139,7 @@ async def stream_response( active_model, len(input_items), request_params['max_output_tokens'], + extra={"otel_skip": True}, ) stream = await self.client.responses.create(**request_params) diff --git a/services/ai/providers/openai_compatible.py b/services/ai/providers/openai_compatible.py index 6e949d4dd..5b056e6a4 100644 --- a/services/ai/providers/openai_compatible.py +++ b/services/ai/providers/openai_compatible.py @@ -310,7 +310,10 @@ async def stream_response( if tools: params["tools"] = _convert_tools_to_openai(tools) logger.info( - "Sending request with %s tools", len(tools) + "Sending request with %s tools: %s", + len(tools), + [t.get("name", t.get("function", {}) + .get("name", "?")) for t in tools], ) stream = await self.client.chat.completions.create(**params) diff --git a/services/ai/telemetry.py b/services/ai/telemetry.py index 8e2185b19..32eac32f5 100644 --- a/services/ai/telemetry.py +++ b/services/ai/telemetry.py @@ -65,6 +65,18 @@ def _build_otlp_logs_url(endpoint: str) -> str: return f"{endpoint}/v1/logs" +class OTelSkipFilter(logging.Filter): + """Logging filter that suppresses records flagged with ``otel_skip=True`` + from reaching the OTel LoggingHandler. + + The record still appears in stdout (other handlers see it); only OTel + export is skipped. + """ + + def filter(self, record: logging.LogRecord) -> bool: + return not getattr(record, "otel_skip", False) + + class TraceContextFilter(logging.Filter): """Logging filter that adds bounded ``trace_id`` and ``span_id`` fields to every ``LogRecord``. @@ -204,6 +216,7 @@ def init_telemetry(app, service_name: str = "omni-ai"): logger_provider=logger_provider, ) otel_handler.addFilter(TraceContextFilter()) + otel_handler.addFilter(OTelSkipFilter()) _otel_log_handler = otel_handler # Attach the OTel handler to the root logger. Avoid duplicate handlers diff --git a/services/indexer/src/error.rs b/services/indexer/src/error.rs index 698a50821..bc8108d91 100644 --- a/services/indexer/src/error.rs +++ b/services/indexer/src/error.rs @@ -1,7 +1,7 @@ use axum::{ - Json, http::StatusCode, response::{IntoResponse, Response}, + Json, }; use serde_json::json; use shared::db::error::DatabaseError; diff --git a/shared/src/telemetry.rs b/shared/src/telemetry.rs index 56a922ae7..d5d030d98 100644 --- a/shared/src/telemetry.rs +++ b/shared/src/telemetry.rs @@ -21,7 +21,9 @@ use opentelemetry_semantic_conventions::{ }; use std::{sync::OnceLock, time::Duration}; use tracing::warn; -use tracing_subscriber::{layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer}; +use tracing_subscriber::{ + filter::filter_fn, layer::SubscriberExt, util::SubscriberInitExt, EnvFilter, Layer, +}; static TRACER_PROVIDER: OnceLock = OnceLock::new(); static METER_PROVIDER: OnceLock = OnceLock::new(); @@ -183,8 +185,13 @@ pub fn init_telemetry(config: TelemetryConfig) -> Result<()> { // Bridge tracing events to OTel log records when a LoggerProvider // with OTLP export is active. + // Bridge tracing events to OTel log records, skipping records + // whose metadata carries the `otel_skip` field name. let otel_log_layer = - opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(&logger_provider); + opentelemetry_appender_tracing::layer::OpenTelemetryTracingBridge::new(&logger_provider) + .with_filter(filter_fn(|metadata| { + !metadata.fields().iter().any(|f| f.name() == "otel_skip") + })); let env_filter = EnvFilter::try_from_default_env() .unwrap_or_else(|_| EnvFilter::new("info")) diff --git a/web/src/lib/server/otel-log-hook.mjs b/web/src/lib/server/otel-log-hook.mjs index d145ae676..bb70c686d 100644 --- a/web/src/lib/server/otel-log-hook.mjs +++ b/web/src/lib/server/otel-log-hook.mjs @@ -173,6 +173,12 @@ export function createPinoOtelHook() { } } + // Skip OTel export when otel_skip is set on the bindings object. + // The log line still reaches stdout via method.apply below. + if (bindings.otel_skip === true) { + return method.apply(this, args) + } + // Map severity const severityNumber = PINO_LEVEL_TO_SEVERITY[level] ?? SeverityNumber.INFO const severityText = PINO_LEVEL_TO_TEXT[level] ?? 'INFO' From 2def5e29a7d4c169a59c594ae2c9b435a84981c5 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Wed, 29 Jul 2026 04:17:30 +0000 Subject: [PATCH 09/14] fix(telemetry): restore stripped debug logs with otel_skip - Restore full request params, messages, tool input, text/JSON deltas, citations, adapted messages, and LLM response body debug logs - Tag sensitive-content logs with otel_skip so they appear in local stdout but are excluded from OTel export - Restore exception details and call_id in stripped error/warning logs - Bounded event metadata (type, stop_reason, index) restored without otel_skip --- services/ai/providers/anthropic.py | 51 ++++++++++++++- services/ai/providers/bedrock.py | 72 +++++++++++++++++++--- services/ai/providers/openai.py | 3 +- services/ai/providers/openai_compatible.py | 5 +- 4 files changed, 117 insertions(+), 14 deletions(-) diff --git a/services/ai/providers/anthropic.py b/services/ai/providers/anthropic.py index d48b603ac..b62755fe9 100644 --- a/services/ai/providers/anthropic.py +++ b/services/ai/providers/anthropic.py @@ -132,6 +132,14 @@ async def stream_response( f"Model: {self.model}, Messages: {len(msg_list)}, Max tokens: {request_params['max_tokens']}", extra={"otel_skip": True}, ) + logger.debug( + f"Full request params: {json.dumps({k: v for k, v in request_params.items() if k != 'messages'}, indent=2)}", + extra={"otel_skip": True}, + ) + logger.debug( + f"Messages: {json.dumps(msg_list, indent=2)}", + extra={"otel_skip": True}, + ) if system_prompt: request_params["system"] = system_prompt @@ -139,11 +147,50 @@ async def stream_response( stream: AsyncStream[MessageStreamEvent] = await self.client.messages.create( **request_params, ) + logger.info( + "Stream created successfully, starting to process events", + extra={"otel_skip": True}, + ) + event_count = 0 async for event in stream: event_count += 1 - if event.type == "message_stop": - logger.info("Message completed after %s events", event_count) + logger.debug(f"Event {event_count}: {event.type}") + if event.type == "content_block_start": + logger.info( + f"Content block start: type={event.content_block.type}" + ) + if event.content_block.type == "tool_use": + logger.info( + f"Tool use started: {event.content_block.name} (id: {event.content_block.id}) (input: {json.dumps(event.content_block.input)})", + extra={"otel_skip": True}, + ) + elif event.type == "content_block_delta": + if event.delta.type == "text_delta": + logger.debug( + f"Text delta: '{event.delta.text}'", + extra={"otel_skip": True}, + ) + elif event.delta.type == "input_json_delta": + logger.debug( + f"JSON delta: {event.delta.partial_json}", + extra={"otel_skip": True}, + ) + elif event.type == "citation": + logger.info( + f"Citation: {event.citation}", + extra={"otel_skip": True}, + ) + elif event.type == "content_block_stop": + logger.info( + f"Content block stop at index {getattr(event, 'index', '')}" + ) + elif event.type == "message_delta": + logger.info( + f"Message delta stop reason: {event.delta.stop_reason}" + ) + elif event.type == "message_stop": + logger.info(f"Message completed after {event_count} events") yield event diff --git a/services/ai/providers/bedrock.py b/services/ai/providers/bedrock.py index 18625dd3c..6964bd3f5 100644 --- a/services/ai/providers/bedrock.py +++ b/services/ai/providers/bedrock.py @@ -296,7 +296,8 @@ def _dedupe_documents(self, messages: list[dict[str, Any]]): deduped_tool_result_content.append(content_block) else: logger.debug( - "[BEDROCK-AMAZON] Deduplicating document in tool result" + f"[BEDROCK-AMAZON] Deduplicating document '{doc_name}' in tool result", + extra={"otel_skip": True}, ) else: deduped_tool_result_content.append(content_block) @@ -537,18 +538,62 @@ async def stream_response( f"[BEDROCK] Model: {self.model_id}, Messages: {len(msg_list)}, Max tokens: {request_params['max_tokens']}", extra={"otel_skip": True}, ) - - + logger.debug( + f"[BEDROCK] Full request body: {json.dumps({k: v for k, v in request_params.items() if k != 'messages'}, indent=2)}", + extra={"otel_skip": True}, + ) + logger.debug( + f"[BEDROCK] Messages: {json.dumps(msg_list, indent=2)}", + extra={"otel_skip": True}, + ) if system_prompt: request_params["system"] = system_prompt stream = self.client.messages.create(**request_params) + logger.info( + "[BEDROCK] Stream created successfully, starting to process events", + extra={"otel_skip": True}, + ) event_count = 0 for event in stream: event_count += 1 - if event.type == "message_stop": + logger.debug(f"[ANTHROPIC] Event {event_count}: {event.type}") + if event.type == "content_block_start": + logger.info( + f"[ANTHROPIC] Content block start: type={event.content_block.type}" + ) + if event.content_block.type == "tool_use": + logger.info( + f"[ANTHROPIC] Tool use started: {event.content_block.name} (id: {event.content_block.id}) (input: {json.dumps(event.content_block.input)})", + extra={"otel_skip": True}, + ) + elif event.type == "content_block_delta": + if event.delta.type == "text_delta": + logger.debug( + f"[ANTHROPIC] Text delta: '{event.delta.text}'", + extra={"otel_skip": True}, + ) + elif event.delta.type == "input_json_delta": + logger.debug( + f"[ANTHROPIC] JSON delta: {event.delta.partial_json}", + extra={"otel_skip": True}, + ) + elif event.type == "citation": + logger.info( + f"[ANTHROPIC] Citation: {event.citation}", + extra={"otel_skip": True}, + ) + elif event.type == "content_block_stop": + logger.info( + f"[ANTHROPIC] Content block stop at index {getattr(event, 'index', '')}" + ) + elif event.type == "message_delta": + logger.info( + f"[ANTHROPIC] Message delta stop reason: {event.delta.stop_reason}" + ) + elif event.type == "message_stop": logger.info( f"[ANTHROPIC] Message completed after {event_count} events" ) @@ -570,7 +615,8 @@ async def stream_response( self._limit_documents(messages, max_documents=5) logger.debug( - f"[BEDROCK-AMAZON] Adapted messages: {len(messages)} messages" + f"[BEDROCK-AMAZON] Adapted messages: {json.dumps(messages, indent=2)}", + extra={"otel_skip": True}, ) tools = ( self._adapt_tools_for_amazon_models(tools) if tools else None @@ -594,7 +640,10 @@ async def stream_response( response = self.client.converse_stream(**request_params) - + logger.info( + "[BEDROCK-AMAZON] Stream created, processing chunks", + extra={"otel_skip": True}, + ) chunk_count = 0 for chunk in response["stream"]: chunk_count += 1 @@ -618,12 +667,14 @@ async def stream_response( except ClientError as e: error_code = e.response.get("Error", {}).get("Code", "Unknown") logger.error( - "[BEDROCK] AWS Bedrock client error (%s)", error_code, + "[BEDROCK] AWS Bedrock client error (%s): %s", error_code, str(e), + exc_info=True, ) raise self._to_provider_error(e, model=model) from e except Exception as e: logger.error( - "[BEDROCK] Failed to stream from AWS Bedrock" + "[BEDROCK] Failed to stream from AWS Bedrock: %s", str(e), + exc_info=True, ) raise self._to_provider_error(e, model=model) from e @@ -685,7 +736,10 @@ async def generate_response( "maxTokens": max_tokens or 512, }, ) - logger.debug("generate_response: response received") + logger.debug( + f"generate_response: response from LLM -> {response}", + extra={"otel_skip": True}, + ) usage_data = response.get("usage", {}) usage = TokenUsage( diff --git a/services/ai/providers/openai.py b/services/ai/providers/openai.py index ab4fa8995..e29d80876 100644 --- a/services/ai/providers/openai.py +++ b/services/ai/providers/openai.py @@ -195,7 +195,8 @@ async def stream_response( if item.type == "function_call": if item.id is None: logger.warning( - "Received function_call item with no id; skipping tool block" + "Received function_call item with no id (call_id=%s); skipping tool block", + item.call_id, ) else: block_index = next_block_index diff --git a/services/ai/providers/openai_compatible.py b/services/ai/providers/openai_compatible.py index 5b056e6a4..f155d1a7e 100644 --- a/services/ai/providers/openai_compatible.py +++ b/services/ai/providers/openai_compatible.py @@ -474,7 +474,8 @@ async def stream_response( except Exception as e: logger.error( - "Failed to stream from OpenAI-compatible endpoint", + "Failed to stream from OpenAI-compatible endpoint: %s", e, + exc_info=True, ) raise ProviderError( str(e), @@ -530,7 +531,7 @@ async def generate_response( raise except Exception as e: logger.error( - "Failed to generate response from OpenAI-compatible endpoint", + "Failed to generate response from OpenAI-compatible endpoint: %s", e, ) raise ProviderError( str(e), From f6c481c445e1bc96ac6c3d72d3cb8b8809110c12 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Wed, 29 Jul 2026 04:28:03 +0000 Subject: [PATCH 10/14] fix(telemetry): restore stripped debug logs across streaming, tools, email, embeddings MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Restore full event content, text deltas, citations, search queries, response bodies, and exception details in non-provider modules - Split bounded metadata (index, status) from sensitive content in same log call — bounded part goes to OTel, sensitive part uses otel_skip - Restore tool name in tool use block start log (bounded, no otel_skip) --- services/ai/email_service/sender.py | 5 ++++- services/ai/embeddings/batch_processor.py | 4 ++-- services/ai/streaming/generate.py | 14 +++++++++++++- services/ai/tools/searcher_client.py | 10 ++++++++-- 4 files changed, 27 insertions(+), 6 deletions(-) diff --git a/services/ai/email_service/sender.py b/services/ai/email_service/sender.py index ef224a62c..26451106d 100644 --- a/services/ai/email_service/sender.py +++ b/services/ai/email_service/sender.py @@ -94,7 +94,10 @@ async def send( data = resp.json() return SendResult(success=True, message_id=data.get("id")) - logger.error("Resend error: status=%d", resp.status_code) + logger.error( + "Resend error: status=%d body=%s", resp.status_code, resp.text, + extra={"otel_skip": True}, + ) return SendResult(success=False, error="Failed to send email via Resend") except Exception as e: logger.error("Resend send error: %s", e) diff --git a/services/ai/embeddings/batch_processor.py b/services/ai/embeddings/batch_processor.py index 886e453ec..ef97452d6 100644 --- a/services/ai/embeddings/batch_processor.py +++ b/services/ai/embeddings/batch_processor.py @@ -168,8 +168,8 @@ async def processing_loop(self): # to allow higher-priority tasks (stream requests) to run if processed_any: await asyncio.sleep(ONLINE_BATCH_DELAY) - except Exception: - logger.error("Online processing loop error", exc_info=True) + except Exception as e: + logger.error("Online processing loop error: %s", e, exc_info=True) await asyncio.sleep(10) async def _process_online_batch(self) -> bool: diff --git a/services/ai/streaming/generate.py b/services/ai/streaming/generate.py index 010e8df94..21ea959a4 100644 --- a/services/ai/streaming/generate.py +++ b/services/ai/streaming/generate.py @@ -739,6 +739,10 @@ async def stream_generator( continue logger.debug("Received event (index: %s)", event_index) + logger.debug( + "Event content: %s", event, + extra={"otel_skip": True}, + ) event_index += 1 now = asyncio.get_running_loop().time() @@ -810,6 +814,10 @@ async def stream_generator( elif event.type == "content_block_start": if event.content_block.type == "text": logger.info("Text block start") + logger.debug( + "Text block content: %s", event.content_block.text, + extra={"otel_skip": True}, + ) text_block: TextBlockParam = TextBlockParam( type="text", text=event.content_block.text ) @@ -819,7 +827,7 @@ async def stream_generator( content_blocks.append(text_block) elif event.content_block.type == "tool_use": logger.info( - "Tool use block start" + "Tool use block start: %s", event.content_block.name ) tool_block: ToolUseBlockParam = ToolUseBlockParam( type="tool_use", @@ -834,6 +842,10 @@ async def stream_generator( elif event.type == "citation": logger.info("Citation received") + logger.debug( + "Citation content: %s", event.citation, + extra={"otel_skip": True}, + ) elif event.type == "message_stop": logger.info("Message stop received.") message_stop_received = True diff --git a/services/ai/tools/searcher_client.py b/services/ai/tools/searcher_client.py index f39751c88..268eb1484 100644 --- a/services/ai/tools/searcher_client.py +++ b/services/ai/tools/searcher_client.py @@ -171,7 +171,10 @@ async def search_documents(self, request: SearchRequest) -> SearchResponse: try: search_payload = request.model_dump(mode="json") - logger.info("Calling searcher service...") + logger.info( + "Calling searcher service with query: %s", request.query, + extra={"otel_skip": True}, + ) response = await self.client.post( f"{self.searcher_url}/search", json=search_payload @@ -197,7 +200,10 @@ async def search_documents(self, request: SearchRequest) -> SearchResponse: async def search_people(self, request: PeopleSearchRequest) -> PeopleSearchResponse: """Search the people directory using omni-searcher service.""" try: - logger.info("People search...") + logger.info( + "People search with query: %s", request.query, + extra={"otel_skip": True}, + ) response = await self.client.get( f"{self.searcher_url}/people/search", params={"q": request.query, "limit": request.limit}, From 88614990896e834df73ff9f294d82d7baa067510 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Wed, 29 Jul 2026 04:31:20 +0000 Subject: [PATCH 11/14] fix(telemetry): restore document_id and exception detail in embedding error log --- services/ai/embeddings/batch_processor.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/services/ai/embeddings/batch_processor.py b/services/ai/embeddings/batch_processor.py index ef97452d6..5b3f3c4fb 100644 --- a/services/ai/embeddings/batch_processor.py +++ b/services/ai/embeddings/batch_processor.py @@ -212,7 +212,8 @@ async def _process_online_batch(self) -> bool: except Exception as exc: had_failure = True logger.error( - "Failed to process document" + "Failed to process document %s: %s", item.document_id, exc, + exc_info=True, ) await self.queue_repo.mark_failed([item.id], str(exc)) self._docs_failed += 1 From 0f37d8523a51dd8eb0806a083bedaafeec8efe72 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Wed, 29 Jul 2026 07:27:18 +0000 Subject: [PATCH 12/14] fix(telemetry): restore stripped log details in web service - Restore search query in agent search request and search error logs with otel_skip - Restore userId in add-message debug log (internal ID) - Restore error object in instrumentation shutdown and OAuth unlink console.error calls --- web/src/instrumentation.mjs | 4 ++-- web/src/routes/(public)/auth/google/unlink/+server.ts | 2 +- web/src/routes/api/chat/[chatId]/messages/+server.ts | 2 +- web/src/routes/api/search/+server.ts | 7 ++++++- web/src/routes/api/v1/search/+server.ts | 2 ++ 5 files changed, 12 insertions(+), 5 deletions(-) diff --git a/web/src/instrumentation.mjs b/web/src/instrumentation.mjs index f17dba649..170acac43 100644 --- a/web/src/instrumentation.mjs +++ b/web/src/instrumentation.mjs @@ -47,8 +47,8 @@ const shutdown = async () => { try { await sdk.shutdown() console.log('Telemetry shut down successfully') - } catch { - console.error('Error shutting down telemetry') + } catch (error) { + console.error('Error shutting down telemetry', error) } } diff --git a/web/src/routes/(public)/auth/google/unlink/+server.ts b/web/src/routes/(public)/auth/google/unlink/+server.ts index 5f0685639..2a8dfcc29 100644 --- a/web/src/routes/(public)/auth/google/unlink/+server.ts +++ b/web/src/routes/(public)/auth/google/unlink/+server.ts @@ -39,7 +39,7 @@ export const POST: RequestHandler = async ({ url, cookies }) => { // Redirect back to settings with success message throw redirect(302, '/settings/integrations?success=google_unlinked') } catch (error) { - console.error('OAuth unlink error') + console.error('OAuth unlink error:', error) // Re-throw redirects if (error instanceof Response) { diff --git a/web/src/routes/api/chat/[chatId]/messages/+server.ts b/web/src/routes/api/chat/[chatId]/messages/+server.ts index 0c6a53edf..e91819c0b 100644 --- a/web/src/routes/api/chat/[chatId]/messages/+server.ts +++ b/web/src/routes/api/chat/[chatId]/messages/+server.ts @@ -129,7 +129,7 @@ export const POST: RequestHandler = async ({ params, request, locals }) => { return json({ error: 'Content or attachments are required' }, { status: 400 }) } - logger.debug('Adding message to chat', { chatId }) + logger.debug('Adding message to chat', { chatId, userId: locals.user.id }) try { // First check if chat exists diff --git a/web/src/routes/api/search/+server.ts b/web/src/routes/api/search/+server.ts index 6191ea98a..5d3f92d99 100644 --- a/web/src/routes/api/search/+server.ts +++ b/web/src/routes/api/search/+server.ts @@ -34,6 +34,8 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { logger.debug('Sending search request to searcher service', { queryLength: queryData.query.length, mode: queryData.mode, + query: queryData.query, + otel_skip: true, }) try { @@ -65,7 +67,10 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { return json(searchResults) } catch (error) { - logger.error('Error calling search service', error) + logger.error('Error calling search service', error, { + query: queryData.query, + otel_skip: true, + }) return json( { error: 'Failed to perform search', diff --git a/web/src/routes/api/v1/search/+server.ts b/web/src/routes/api/v1/search/+server.ts index 3b6c1213a..0bd8ba02a 100644 --- a/web/src/routes/api/v1/search/+server.ts +++ b/web/src/routes/api/v1/search/+server.ts @@ -96,8 +96,10 @@ export const POST: RequestHandler = async ({ request, fetch, locals }) => { } logger.debug('Agent search request', { + query: queryData.query, queryLength: queryData.query?.length ?? 0, mode: queryData.mode, + otel_skip: true, }) try { From 704e0802c8f09f4494d6fcc78100618aeb93c7e7 Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Wed, 29 Jul 2026 07:31:43 +0000 Subject: [PATCH 13/14] fix(telemetry): restore stripped debug details in Rust services - Restore search queries in fulltext/semantic/hybrid search logs with otel_skip - Restore query embedding, RAG prompt, AI answer cache info with otel_skip - Restore connector response bodies in error logs with otel_skip - Restore document IDs in indexer create/update/delete logs --- .../connector-manager/src/connector_client.rs | 47 +++++++++++++++---- services/indexer/src/lib.rs | 6 +-- services/searcher/src/handlers.rs | 18 +++++-- services/searcher/src/search.rs | 25 ++++++++-- services/searcher/src/search_repository.rs | 2 +- 5 files changed, 77 insertions(+), 21 deletions(-) diff --git a/services/connector-manager/src/connector_client.rs b/services/connector-manager/src/connector_client.rs index 1f898337b..925258fc5 100644 --- a/services/connector-manager/src/connector_client.rs +++ b/services/connector-manager/src/connector_client.rs @@ -45,7 +45,11 @@ impl ConnectorClient { if !response.status().is_success() { let status = response.status(); - error!("Failed to get manifest: status={}", status); + let body = response.text().await.unwrap_or_default(); + error!( + otel_skip = true, + "Failed to get manifest: {} - {}", status, body + ); return Err(ClientError::ConnectorError { status: status.as_u16(), message: String::new(), @@ -73,10 +77,17 @@ impl ConnectorClient { if !response.status().is_success() { let status = response.status(); + let body = response.text().await.unwrap_or_default(); if status.as_u16() == 404 && request.sync_mode == SyncType::Realtime { - debug!("Realtime sync unavailable: status={}", status); + debug!( + otel_skip = true, + "Realtime sync unavailable: {} - {}", status, body + ); } else { - error!("Failed to trigger sync: status={}", status); + error!( + otel_skip = true, + "Failed to trigger sync: {} - {}", status, body + ); } return Err(ClientError::ConnectorError { status: status.as_u16(), @@ -183,7 +194,11 @@ impl ConnectorClient { if !response.status().is_success() { let status = response.status(); - warn!("Failed to cancel sync: status={}", status); + let body = response.text().await.unwrap_or_default(); + warn!( + otel_skip = true, + "Failed to cancel sync: {} - {}", status, body + ); return Err(ClientError::ConnectorError { status: status.as_u16(), message: String::new(), @@ -207,7 +222,11 @@ impl ConnectorClient { if !response.status().is_success() { let status = response.status(); - error!("Failed to execute action: status={}", status); + let body = response.text().await.unwrap_or_default(); + error!( + otel_skip = true, + "Failed to execute action: {} - {}", status, body + ); return Err(ClientError::ConnectorError { status: status.as_u16(), message: String::new(), @@ -297,7 +316,11 @@ impl ConnectorClient { if !response.status().is_success() { let status = response.status(); - error!("Failed to read resource: status={}", status); + let body = response.text().await.unwrap_or_default(); + error!( + otel_skip = true, + "Failed to read resource: {} - {}", status, body + ); return Err(ClientError::ConnectorError { status: status.as_u16(), message: String::new(), @@ -324,7 +347,11 @@ impl ConnectorClient { if !response.status().is_success() { let status = response.status(); - error!("Failed to get prompt: status={}", status); + let body = response.text().await.unwrap_or_default(); + error!( + otel_skip = true, + "Failed to get prompt: {} - {}", status, body + ); return Err(ClientError::ConnectorError { status: status.as_u16(), message: String::new(), @@ -351,7 +378,11 @@ impl ConnectorClient { if !response.status().is_success() { let status = response.status(); - error!("Failed to get skill: status={}", status); + let body = response.text().await.unwrap_or_default(); + error!( + otel_skip = true, + "Failed to get skill: {} - {}", status, body + ); return Err(ClientError::ConnectorError { status: status.as_u16(), message: String::new(), diff --git a/services/indexer/src/lib.rs b/services/indexer/src/lib.rs index 1721c7753..576887e4d 100644 --- a/services/indexer/src/lib.rs +++ b/services/indexer/src/lib.rs @@ -159,7 +159,7 @@ async fn create_document( let repo = DocumentRepository::new(state.db_pool.pool()); let document = repo.create(doc).await?; - info!("Created document"); + info!("Created document: {}", document.id); Ok(Json(document)) } @@ -208,7 +208,7 @@ async fn update_document( match updated_doc { Some(doc) => { - info!("Updated document"); + info!("Updated document: {}", doc.id); Ok(Json(doc)) } None => Err(error::IndexerError::NotFound(format!( @@ -232,7 +232,7 @@ async fn delete_document( ))); } - info!("Deleted document"); + info!("Deleted document: {}", id); Ok(Json(json!({ "message": "Document deleted successfully", "id": id diff --git a/services/searcher/src/handlers.rs b/services/searcher/src/handlers.rs index 1b4fd0105..ccaa51692 100644 --- a/services/searcher/src/handlers.rs +++ b/services/searcher/src/handlers.rs @@ -29,7 +29,7 @@ use std::sync::Arc; use std::task::{Context, Poll}; use std::time::Instant; use tokio::sync::Mutex; -use tracing::{error, info}; +use tracing::{debug, error, info}; /// A stream wrapper that collects chunks for caching while forwarding them to the client struct CachingStream { @@ -246,7 +246,10 @@ pub async fn ai_answer( Json(mut request): Json, ) -> Result, axum::http::StatusCode> { let answer_query_len = request.query.len(); - info!("Received AI answer request"); + info!( + otel_skip = true, + "Received AI answer request: {:?}", request + ); hydrate_user_configuration(&state, &mut request) .await .map_err(|_| axum::http::StatusCode::BAD_REQUEST)?; @@ -267,7 +270,10 @@ pub async fn ai_answer( // Try to get cached AI response first if let Ok(mut conn) = state.redis_client.get_multiplexed_async_connection().await { if let Ok(cached_answer) = conn.get::<_, String>(&cache_key).await { - info!("AI answer cache hit"); + info!( + otel_skip = true, + "Cache hit for AI answer query: '{}'", request.query + ); let response = axum::response::Response::builder() .status(StatusCode::OK) .header("Content-Type", "text/plain; charset=utf-8") @@ -279,7 +285,10 @@ pub async fn ai_answer( } // Cache miss - generate fresh response - info!("AI answer cache miss"); + info!( + otel_skip = true, + "Cache miss for AI answer query: '{}'", request.query + ); // Get RAG context by running hybrid search let context = match search_engine.get_rag_context(&request).await { @@ -292,6 +301,7 @@ pub async fn ai_answer( // Build RAG prompt with context and citation instructions let prompt = search_engine.build_rag_prompt(&request.query, &context); + debug!(otel_skip = true, "RAG prompt: {}", prompt); info!("Built RAG prompt of length: {}", prompt.len()); // Stream AI response diff --git a/services/searcher/src/search.rs b/services/searcher/src/search.rs index ac1f80c81..ad49709ed 100644 --- a/services/searcher/src/search.rs +++ b/services/searcher/src/search.rs @@ -487,7 +487,10 @@ impl SearchEngine { let content_types = request.content_types.as_deref(); let attribute_filters = request.attribute_filters.as_ref(); - debug!("Running fulltext search"); + debug!( + otel_skip = true, + "Running fulltext search for {}", &request.query + ); let search_hits = repo .search( &request.query, @@ -551,7 +554,10 @@ impl SearchEngine { offset: i64, ) -> Result> { let start_time = Instant::now(); - info!("Performing semantic search"); + info!( + otel_skip = true, + "Performing semantic search for query '{}'", request.query + ); let query_embedding = self.generate_query_embedding(&request.query).await?; @@ -889,7 +895,10 @@ impl SearchEngine { } async fn generate_query_embedding(&self, query: &str) -> Result> { - debug!("Generating query embeddings"); + debug!( + otel_skip = true, + "Generating query embeddings for query '{}'", query + ); let embeddings = self .ai_client .generate_embeddings_with_options( @@ -1034,7 +1043,10 @@ impl SearchEngine { user_groups: &[String], tantivy_query: Option<&str>, ) -> Result<(Vec, i64)> { - info!("Performing hybrid search"); + info!( + otel_skip = true, + "Performing hybrid search for query '{}'", request.query + ); let start_time = Instant::now(); let doc_repo = DocumentRepository::new(self.db_pool.pool()); @@ -1298,7 +1310,10 @@ impl SearchEngine { /// Generate RAG context from search request using chunk-based approach with expanded context pub async fn get_rag_context(&self, request: &SearchRequest) -> Result> { - info!("Generating RAG context"); + info!( + otel_skip = true, + "Generating RAG context for query '{}'", request.query + ); let user_groups = if let Some(email) = request.user_email() { let group_repo = GroupRepository::new(self.db_pool.pool()); diff --git a/services/searcher/src/search_repository.rs b/services/searcher/src/search_repository.rs index 6c408873a..20eea0280 100644 --- a/services/searcher/src/search_repository.rs +++ b/services/searcher/src/search_repository.rs @@ -283,7 +283,7 @@ impl SearchDocumentRepository { offset_idx = offset_idx, min_score_ratio_idx = min_score_ratio_idx, ); - debug!("Full search query: {} params", param_idx); + debug!(otel_skip = true, "Full search query: {}", full_query); let mut query_builder = sqlx::query_as::<_, SearchHitWithTotalRow>(&full_query).bind(tantivy_query); From ad03f3fb3809c9bb3d0fb08b34920721459939cd Mon Sep 17 00:00:00 2001 From: Praveen Sampath Date: Wed, 29 Jul 2026 10:39:39 +0000 Subject: [PATCH 14/14] fix(connector-manager): handle Option source_id after master rebase --- Cargo.lock | 115 +++++++++++++++++++++ services/connector-manager/src/handlers.rs | 15 ++- 2 files changed, 125 insertions(+), 5 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 8422d22cf..24e111bf0 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1439,6 +1439,18 @@ dependencies = [ "windows-sys 0.59.0", ] +[[package]] +name = "const-hex" +version = "1.19.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "33e2a781ebdf4467d1428dc4593067825fb646f6871475098d8577421af73558" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "proptest", + "serde_core", +] + [[package]] name = "const-oid" version = "0.9.6" @@ -3169,6 +3181,19 @@ dependencies = [ "webpki-roots 1.0.7", ] +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper 1.10.1", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + [[package]] name = "hyper-tls" version = "0.5.0" @@ -4475,6 +4500,7 @@ dependencies = [ "dashmap 6.2.1", "dotenvy", "futures", + "opentelemetry", "redis", "reqwest 0.12.28", "serde", @@ -4666,6 +4692,7 @@ dependencies = [ "tower 0.4.13", "tower-http 0.5.2", "tracing", + "tracing-opentelemetry", "tracing-subscriber", "ulid", "uuid", @@ -4900,6 +4927,18 @@ dependencies = [ "tracing", ] +[[package]] +name = "opentelemetry-appender-tracing" +version = "0.32.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c0080f0dc1d7c786f467cd85a4e395fcab11ee852004f39a29a18ab7c25d837" +dependencies = [ + "opentelemetry", + "tracing", + "tracing-core", + "tracing-subscriber", +] + [[package]] name = "opentelemetry-http" version = "0.32.0" @@ -4926,6 +4965,7 @@ dependencies = [ "opentelemetry_sdk", "prost", "reqwest 0.13.4", + "serde_json", "thiserror 2.0.18", ] @@ -4935,9 +4975,14 @@ version = "0.32.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "56d658ba1faf63f7b9c492cfbe6e0ec365440a16132d3270c1065f7b33f1b638" dependencies = [ + "base64 0.22.1", + "const-hex", "opentelemetry", "opentelemetry_sdk", "prost", + "serde", + "tonic", + "tonic-prost", ] [[package]] @@ -5584,6 +5629,21 @@ dependencies = [ "windows", ] +[[package]] +name = "proptest" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b45fcc2344c680f5025fe57779faef368840d0bd1f42f216291f0dc4ace4744" +dependencies = [ + "bitflags 2.13.0", + "num-traits", + "rand 0.9.4", + "rand_chacha 0.9.0", + "rand_xorshift", + "regex-syntax", + "unarray", +] + [[package]] name = "prost" version = "0.14.4" @@ -5837,6 +5897,15 @@ dependencies = [ "rand 0.8.6", ] +[[package]] +name = "rand_xorshift" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" +dependencies = [ + "rand_core 0.9.5", +] + [[package]] name = "raw-cpuid" version = "11.6.0" @@ -6848,12 +6917,15 @@ dependencies = [ "msg_parser", "nonzero_ext", "opentelemetry", + "opentelemetry-appender-tracing", "opentelemetry-http", "opentelemetry-otlp", + "opentelemetry-proto", "opentelemetry-semantic-conventions", "opentelemetry_sdk", "pdf_oxide", "pgvector", + "prost", "quick-xml", "rand 0.8.6", "redis", @@ -8023,6 +8095,43 @@ version = "0.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5d99f8c9a7727884afe522e9bd5edbfc91a3312b36a77b5fb8926e4c31a41801" +[[package]] +name = "tonic" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" +dependencies = [ + "async-trait", + "base64 0.22.1", + "bytes", + "http 1.4.2", + "http-body 1.0.1", + "http-body-util", + "hyper 1.10.1", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "sync_wrapper 1.0.2", + "tokio", + "tokio-stream", + "tower 0.5.3", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tonic-prost" +version = "0.14.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" +dependencies = [ + "bytes", + "prost", + "tonic", +] + [[package]] name = "tower" version = "0.4.13" @@ -8277,6 +8386,12 @@ dependencies = [ "web-time", ] +[[package]] +name = "unarray" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eaea85b334db583fe3274d12b4cd1880032beab409c0d774be044d4480ab9a94" + [[package]] name = "unicase" version = "2.9.0" diff --git a/services/connector-manager/src/handlers.rs b/services/connector-manager/src/handlers.rs index eaa661489..165e54f72 100644 --- a/services/connector-manager/src/handlers.rs +++ b/services/connector-manager/src/handlers.rs @@ -370,12 +370,17 @@ pub async fn execute_action( ) -> Result { info!("Executing action"); + let source_id = request + .source_id + .as_deref() + .ok_or_else(|| ApiError::BadRequest("source_id is required".to_string()))?; + let source_repo = SourceRepository::new(state.db_pool.pool()); let source = source_repo - .find_by_id(request.source_id.clone()) + .find_by_id(source_id.to_string()) .await .map_err(|e| ApiError::Internal(e.to_string()))? - .ok_or_else(|| ApiError::NotFound(format!("Source not found: {}", request.source_id)))?; + .ok_or_else(|| ApiError::NotFound(format!("Source not found: {}", source_id)))?; // Look up the connector manifest to get connector_url and action metadata let manifests = get_registered_manifests(&state.redis_client).await; @@ -423,7 +428,7 @@ pub async fn execute_action( .map_err(|e| ApiError::Internal(e.to_string()))?; let creds = match resolve_credentials( &creds_repo, - &request.source_id, + source_id, request.user_id.as_deref(), action_admin_only, ) @@ -432,7 +437,7 @@ pub async fn execute_action( CredentialResolution::Resolved(c) => c, CredentialResolution::NeedsUserAuth { provider } => { return Ok(needs_user_auth_response( - &request.source_id, + source_id, source.source_type, provider, )?); @@ -440,7 +445,7 @@ pub async fn execute_action( CredentialResolution::NoCredentials => { return Err(ApiError::NotFound(format!( "Credentials not found for source: {}", - request.source_id + source_id ))); } };