feat: intro telemetry

Signed-off-by: Chojan Shang <chojan.shang@vesoft.com>

Author: Chojan Shang

pyproject.toml

@@ -43,6 +43,9 @@ dev = [
     "python-dotenv>=1.1.1",
 ]
 
+[project.optional-dependencies]
+logfire = ["logfire>=0.14"]
+
 [build-system]
 requires = ["hatchling"]
 build-backend = "hatchling.build"

src/acp/connection.py

@@ -25,6 +25,7 @@
     SenderFactory,
     TaskSupervisor,
 )
+from .telemetry import span_context
 
 JsonValue = Any
 MethodHandler = Callable[[str, JsonValue | None, bool], Awaitable[JsonValue | None]]
@@ -135,35 +136,41 @@ async def _process_message(self, message: dict[str, Any]) -> None:
 
     async def _run_request(self, message: dict[str, Any]) -> Any:
         payload: dict[str, Any] = {"jsonrpc": "2.0", "id": message["id"]}
-        try:
-            result = await self._handler(message["method"], message.get("params"), False)
-            if isinstance(result, BaseModel):
-                result = result.model_dump()
-            payload["result"] = result if result is not None else None
-            await self._sender.send(payload)
-            return payload.get("result")
-        except RequestError as exc:
-            payload["error"] = exc.to_error_obj()
-            await self._sender.send(payload)
-            raise
-        except ValidationError as exc:
-            err = RequestError.invalid_params({"errors": exc.errors()})
-            payload["error"] = err.to_error_obj()
-            await self._sender.send(payload)
-            raise err from None
-        except Exception as exc:
+        method = message["method"]
+        with span_context(
+            "acp.request",
+            attributes={"method": method},
+        ):
             try:
-                data = json.loads(str(exc))
-            except Exception:
-                data = {"details": str(exc)}
-            err = RequestError.internal_error(data)
-            payload["error"] = err.to_error_obj()
-            await self._sender.send(payload)
-            raise err from None
+                result = await self._handler(method, message.get("params"), False)
+                if isinstance(result, BaseModel):
+                    result = result.model_dump()
+                payload["result"] = result if result is not None else None
+                await self._sender.send(payload)
+                return payload.get("result")
+            except RequestError as exc:
+                payload["error"] = exc.to_error_obj()
+                await self._sender.send(payload)
+                raise
+            except ValidationError as exc:
+                err = RequestError.invalid_params({"errors": exc.errors()})
+                payload["error"] = err.to_error_obj()
+                await self._sender.send(payload)
+                raise err from None
+            except Exception as exc:
+                try:
+                    data = json.loads(str(exc))
+                except Exception:
+                    data = {"details": str(exc)}
+                err = RequestError.internal_error(data)
+                payload["error"] = err.to_error_obj()
+                await self._sender.send(payload)
+                raise err from None
 
     async def _run_notification(self, message: dict[str, Any]) -> None:
-        with contextlib.suppress(Exception):
-            await self._handler(message["method"], message.get("params"), True)
+        method = message["method"]
+        with span_context("acp.notification", attributes={"method": method}), contextlib.suppress(Exception):
+            await self._handler(method, message.get("params"), True)
 
     async def _handle_response(self, message: dict[str, Any]) -> None:
         request_id = message["id"]

src/acp/telemetry.py

@@ -0,0 +1,41 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from contextlib import AbstractContextManager, ExitStack, nullcontext
+from typing import Any
+
+try:
+    from logfire import span as logfire_span
+except Exception:  # pragma: no cover - logfire is optional
+    logfire_span = None  # type: ignore[assignment]
+else:  # pragma: no cover - optional
+    os.environ.setdefault("LOGFIRE_IGNORE_NO_CONFIG", "1")
+
+try:  # pragma: no cover - opentelemetry is optional
+    from opentelemetry.trace import get_tracer as otel_get_tracer
+except Exception:  # pragma: no cover - opentelemetry is optional
+    otel_get_tracer = None  # type: ignore[assignment]
+
+DEFAULT_TAGS = ["acp"]
+TRACER = otel_get_tracer(__name__) if otel_get_tracer else None
+
+
+def _start_tracer_span(name: str, *, attributes: Mapping[str, Any] | None = None) -> AbstractContextManager[Any]:
+    if TRACER is None:
+        return nullcontext()
+    attrs = dict(attributes or {})
+    return TRACER.start_as_current_span(name, attributes=attrs)
+
+
+def span_context(name: str, *, attributes: Mapping[str, Any] | None = None) -> AbstractContextManager[None]:
+    if logfire_span is None and TRACER is None:
+        return nullcontext()
+    stack = ExitStack()
+    attrs: dict[str, Any] = {"logfire.tags": DEFAULT_TAGS}
+    if attributes:
+        attrs.update(attributes)
+    if logfire_span is not None:
+        stack.enter_context(logfire_span(name, attributes=attrs))
+    stack.enter_context(_start_tracer_span(name, attributes=attributes))
+    return stack