File size: 5,658 Bytes
f291c90
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
"""OpenTelemetry tracing — distributed traces across webhook → queue → worker → agent → tools.

Activated only when ``DEEPAGENT_OTEL_ENABLED=1`` and an OTLP endpoint is
configured (``OTEL_EXPORTER_OTLP_ENDPOINT``, e.g. ``http://tempo:4317``).

Spans emitted:

- ``webhook.receive``    HTTP request lifecycle (FastAPI auto-instr)
- ``job.enqueue``        webhook persists a Job into Redis
- ``job.dequeue``        worker claims a Job
- ``job.process``        worker runs the agent (parent for everything below)
- ``agent.stream``       per agent.stream() invocation
- ``tool.<name>``        every tool call from the agent
- ``llm.<provider>``     every LLM call (via langchain-otel auto-instrumentation if installed)

A ``traceparent`` header is propagated through the job payload so the worker's
span chains under the webhook's request span.
"""
from __future__ import annotations

import contextlib
import logging
import os
from typing import Any, Iterator, Optional

log = logging.getLogger(__name__)

_ENABLED = False
_TRACER = None


def is_enabled() -> bool:
    return _ENABLED


def setup_tracing(service_name: str = "gh-deepagent") -> None:
    """Initialise the OTel tracer provider + OTLP exporter.

    No-op when ``DEEPAGENT_OTEL_ENABLED`` is falsy. Safe to call multiple times.
    """
    global _ENABLED, _TRACER
    if _ENABLED:
        return
    if os.getenv("DEEPAGENT_OTEL_ENABLED", "0").lower() not in ("1", "true", "yes"):
        return
    try:
        from opentelemetry import trace
        from opentelemetry.exporter.otlp.proto.grpc.trace_exporter import OTLPSpanExporter
        from opentelemetry.sdk.resources import Resource, SERVICE_NAME, SERVICE_VERSION
        from opentelemetry.sdk.trace import TracerProvider
        from opentelemetry.sdk.trace.export import BatchSpanProcessor
    except ImportError:
        log.warning("opentelemetry deps not installed; tracing disabled")
        return

    resource = Resource.create({
        SERVICE_NAME: service_name,
        SERVICE_VERSION: os.getenv("DEEPAGENT_VERSION", "0.3.0"),
    })
    provider = TracerProvider(resource=resource)
    exporter = OTLPSpanExporter()  # picks OTEL_EXPORTER_OTLP_ENDPOINT from env
    provider.add_span_processor(BatchSpanProcessor(exporter))
    trace.set_tracer_provider(provider)

    # Auto-instrument FastAPI + Redis when their instrumentors are present.
    for mod_name, instr_path in (
        ("fastapi", "opentelemetry.instrumentation.fastapi"),
        ("redis", "opentelemetry.instrumentation.redis"),
        ("requests", "opentelemetry.instrumentation.requests"),
        ("httpx", "opentelemetry.instrumentation.httpx"),
    ):
        try:
            instr_mod = __import__(instr_path, fromlist=["*"])
            instr_cls_name = f"{mod_name.capitalize()}Instrumentor"
            getattr(instr_mod, instr_cls_name)().instrument()
        except Exception:
            pass

    _TRACER = trace.get_tracer(service_name)
    _ENABLED = True

    # Inject trace IDs into structlog so logs are correlated.
    try:
        import structlog
        from opentelemetry import trace as _trace

        def add_trace_context(_logger, _method, event_dict):
            ctx = _trace.get_current_span().get_span_context()
            if ctx and ctx.is_valid:
                event_dict["trace_id"] = format(ctx.trace_id, "032x")
                event_dict["span_id"] = format(ctx.span_id, "016x")
            return event_dict

        # Patch the structlog config to add our processor in front of the renderer.
        current = structlog.get_config()
        processors = list(current["processors"])
        processors.insert(-1, add_trace_context)
        structlog.configure(processors=processors)
    except Exception:
        pass

    log.info("OpenTelemetry tracing enabled (service=%s)", service_name)


@contextlib.contextmanager
def span(name: str, **attrs: Any) -> Iterator[Optional[object]]:
    """Context manager that creates an OTel span when tracing is on; no-op otherwise.

    Always yields a span object (or None) so calling code stays identical.
    """
    if not _ENABLED or _TRACER is None:
        yield None
        return
    with _TRACER.start_as_current_span(name) as s:
        for k, v in attrs.items():
            if v is not None:
                try:
                    s.set_attribute(k, v if isinstance(v, (str, int, float, bool)) else str(v))
                except Exception:
                    pass
        yield s


def current_traceparent() -> Optional[str]:
    """Return the W3C ``traceparent`` for the active span (for cross-process propagation)."""
    if not _ENABLED:
        return None
    try:
        from opentelemetry.propagate import inject
        carrier: dict[str, str] = {}
        inject(carrier)
        return carrier.get("traceparent")
    except Exception:
        return None


@contextlib.contextmanager
def continue_from(traceparent: Optional[str], name: str, **attrs: Any) -> Iterator[Optional[object]]:
    """Continue a trace from a serialised ``traceparent`` (set by the webhook)."""
    if not _ENABLED or _TRACER is None or not traceparent:
        with span(name, **attrs) as s:
            yield s
        return
    try:
        from opentelemetry import context as otel_context
        from opentelemetry.propagate import extract
        ctx = extract({"traceparent": traceparent})
        token = otel_context.attach(ctx)
        try:
            with span(name, **attrs) as s:
                yield s
        finally:
            otel_context.detach(token)
    except Exception:
        with span(name, **attrs) as s:
            yield s