Merge pull request #21 from aliyun/fix_a2a_telementery_report

feat: wire telemetry lifecycle into a2a/acp servers and propagate exc…

Author: guima-why

src/iac_code/a2a/events.py

@@ -31,6 +31,7 @@
 )
 
 _METADATA_MAX_CHARS = 4000
+_ERROR_TEXT_MAX_CHARS = 1000
 _METADATA_MAX_DEPTH = 32
 logger = logging.getLogger(__name__)
 A2APermissionResolver: TypeAlias = Callable[[PermissionRequestEvent], "bool | Awaitable[bool]"]
@@ -298,18 +299,19 @@ async def publish_stream_event(
         return None
 
     if isinstance(event, ErrorEvent):
+        if event.is_retryable:
+            text = "A temporary error occurred. Please retry."
+            state = TaskState.TASK_STATE_INPUT_REQUIRED
+        else:
+            raw = event.error or "Unknown error"
+            text = raw[:_ERROR_TEXT_MAX_CHARS]
+            state = TaskState.TASK_STATE_FAILED
         await _enqueue_status(
             event_queue,
             task_id=task_id,
             context_id=context_id,
-            state=TaskState.TASK_STATE_INPUT_REQUIRED if event.is_retryable else TaskState.TASK_STATE_FAILED,
-            message=_agent_text_message(
-                task_id=task_id,
-                context_id=context_id,
-                text="A temporary error occurred. Please retry."
-                if event.is_retryable
-                else "An internal error occurred.",
-            ),
+            state=state,
+            message=_agent_text_message(task_id=task_id, context_id=context_id, text=text),
         )
         return None
 

src/iac_code/a2a/executor.py

@@ -29,6 +29,16 @@
 
 logger = logging.getLogger(__name__)
 _CONTEXT_LOCK_ACQUIRE_TIMEOUT_SECONDS = 1
+_ERROR_TEXT_MAX_CHARS = 1000
+
+
+def _format_exception(exc: BaseException) -> str:
+    message = str(exc)
+    if not message:
+        return type(exc).__name__
+    return f"{type(exc).__name__}: {message[:_ERROR_TEXT_MAX_CHARS]}"
+
+
 A2APermissionResolver: TypeAlias = Callable[[Any], "bool | Awaitable[bool]"]
 
 
@@ -274,6 +284,17 @@ def runtime_factory(session_id: str) -> Any:
                 ctx.touch()
                 task.touch()
                 self._task_store.mirror_context(ctx)
+                # Force-flush telemetry between tasks. The a2a server may run in
+                # an ephemeral sandbox that's destroyed immediately after the
+                # response is delivered, before the natural batch interval or
+                # process-exit graceful_shutdown can run. Synchronous flush is
+                # offloaded to a worker thread so the event loop is not blocked.
+                from iac_code.services.telemetry import flush_telemetry
+
+                try:
+                    await asyncio.to_thread(flush_telemetry)
+                except Exception:
+                    logger.debug("flush_telemetry after task failed", exc_info=True)
         finally:
             lock.release()
 
@@ -338,7 +359,7 @@ def _sanitize_error(self, exc: Exception) -> str:
         if status == 401:
             return "Authentication required. Please configure your API credentials."
         logger.exception("Unhandled A2A executor error")
-        return "An internal error occurred."
+        return _format_exception(exc)
 
     async def _publish_status(
         self,

src/iac_code/acp/session.py

@@ -351,6 +351,17 @@ async def _run() -> None:
             duration_ms = (time.monotonic() - prompt_start) * 1000
             if self._metrics is not None:
                 self._metrics.record_prompt(duration_ms)
+            # Force-flush telemetry between prompts. The acp server may run in
+            # an ephemeral sandbox that's destroyed immediately after the
+            # response is delivered, before the natural batch interval or
+            # process-exit graceful_shutdown can run. Synchronous flush is
+            # offloaded to a worker thread so the event loop is not blocked.
+            from iac_code.services.telemetry import flush_telemetry
+
+            try:
+                await asyncio.to_thread(flush_telemetry)
+            except Exception:
+                logger.debug("flush_telemetry after prompt failed", exc_info=True)
 
         self.touch()
 

src/iac_code/cli/main.py

@@ -311,14 +311,91 @@ def acp(
     debug: bool = typer.Option(False, "--debug", "-d", help=_("Enable debug logging")),
 ) -> None:
     """Run iac-code as an ACP server."""
-    if transport == "http":
-        from iac_code.acp import acp_main_http
+    import atexit
+    import signal as _signal_mod
+    import time
+
+    from iac_code.services.telemetry import add_metric, bootstrap_telemetry, graceful_shutdown, log_event
+    from iac_code.services.telemetry.names import Events, Metrics
+
+    telemetry_session_id = f"acp-server-{uuid.uuid4()}"
+    bootstrap_telemetry(session_id=telemetry_session_id)
+    log_event(
+        Events.SESSION_STARTED,
+        {
+            "mode": "acp-server",
+            "transport": transport,
+        },
+    )
+    add_metric(Metrics.SESSION_COUNT, 1, {})
 
-        acp_main_http(host=host, port=port, debug=debug)
-    else:
-        from iac_code.acp import acp_main
+    started = time.monotonic()
+    exit_reason = "normal"
+    _finalized = [False]
+
+    def _finalize_telemetry(reason_override: str | None = None) -> None:
+        if _finalized[0]:
+            return
+        _finalized[0] = True
+        final_reason = reason_override or exit_reason
+        try:
+            log_event(
+                Events.SESSION_EXITED,
+                {
+                    "mode": "acp-server",
+                    "reason": final_reason,
+                    "duration_s": int(time.monotonic() - started),
+                },
+            )
+        finally:
+            graceful_shutdown()
+
+    atexit.register(_finalize_telemetry, "atexit")
+
+    def _telemetry_excepthook(exc_type, exc_value, traceback_obj):
+        try:
+            log_event(
+                Events.EXCEPTION_UNCAUGHT,
+                {
+                    "error_name": exc_type.__name__,
+                    "location": "acp",
+                },
+            )
+            _finalize_telemetry(f"exception:{exc_type.__name__}")
+        finally:
+            sys.__excepthook__(exc_type, exc_value, traceback_obj)
+
+    sys.excepthook = _telemetry_excepthook
+
+    _prev_sigterm = _signal_mod.getsignal(_signal_mod.SIGTERM)
+    _prev_sigint = _signal_mod.getsignal(_signal_mod.SIGINT)
+
+    def _telemetry_signal_handler(signum, frame):
+        _finalize_telemetry(f"signal:{signum}")
+        prev = _prev_sigterm if signum == _signal_mod.SIGTERM else _prev_sigint
+        if callable(prev):
+            prev(signum, frame)  # ty: ignore[call-top-callable]
+            return
+        _signal_mod.signal(signum, _signal_mod.SIG_DFL)
+        os.kill(os.getpid(), signum)
+
+    _signal_mod.signal(_signal_mod.SIGTERM, _telemetry_signal_handler)
+    _signal_mod.signal(_signal_mod.SIGINT, _telemetry_signal_handler)
+
+    try:
+        if transport == "http":
+            from iac_code.acp import acp_main_http
 
-        acp_main(debug=debug)
+            acp_main_http(host=host, port=port, debug=debug)
+        else:
+            from iac_code.acp import acp_main
+
+            acp_main(debug=debug)
+    except Exception:
+        exit_reason = "error"
+        raise
+    finally:
+        _finalize_telemetry()
 
 
 def _load_a2a_config(path: str) -> dict[str, Any]:
@@ -506,6 +583,86 @@ def a2a(
             err=True,
         )
         raise typer.Exit(1) from exc
+
+    import atexit
+    import signal as _signal_mod
+    import time
+
+    from iac_code.services.telemetry import add_metric, bootstrap_telemetry, graceful_shutdown, log_event
+    from iac_code.services.telemetry.names import Events, Metrics
+
+    telemetry_session_id = f"a2a-server-{uuid.uuid4()}"
+    bootstrap_telemetry(session_id=telemetry_session_id)
+    log_event(
+        Events.SESSION_STARTED,
+        {
+            "mode": "a2a-server",
+            "transport": transport,
+        },
+    )
+    add_metric(Metrics.SESSION_COUNT, 1, {})
+
+    started = time.monotonic()
+    exit_reason = "normal"
+    _finalized = [False]
+
+    def _finalize_telemetry(reason_override: str | None = None) -> None:
+        if _finalized[0]:
+            return
+        _finalized[0] = True
+        final_reason = reason_override or exit_reason
+        try:
+            log_event(
+                Events.SESSION_EXITED,
+                {
+                    "mode": "a2a-server",
+                    "reason": final_reason,
+                    "duration_s": int(time.monotonic() - started),
+                },
+            )
+        finally:
+            graceful_shutdown()
+
+    # atexit fires on normal interpreter exit, including after uvicorn returns
+    # from a graceful shutdown — covers the path where the finally below also
+    # ran (idempotent via the flag).
+    atexit.register(_finalize_telemetry, "atexit")
+
+    def _telemetry_excepthook(exc_type, exc_value, traceback_obj):
+        try:
+            log_event(
+                Events.EXCEPTION_UNCAUGHT,
+                {
+                    "error_name": exc_type.__name__,
+                    "location": "a2a",
+                },
+            )
+            _finalize_telemetry(f"exception:{exc_type.__name__}")
+        finally:
+            sys.__excepthook__(exc_type, exc_value, traceback_obj)
+
+    sys.excepthook = _telemetry_excepthook
+
+    # Install our own SIGTERM/SIGINT handlers BEFORE uvicorn does its own. In
+    # the common path uvicorn replaces these and triggers graceful shutdown
+    # itself, after which the finally below runs. These handlers only fire if
+    # a signal arrives before uvicorn installs its own (or if uvicorn never
+    # got a chance to) — the ARMS sandbox SIGTERM-then-SIGKILL pattern.
+    _prev_sigterm = _signal_mod.getsignal(_signal_mod.SIGTERM)
+    _prev_sigint = _signal_mod.getsignal(_signal_mod.SIGINT)
+
+    def _telemetry_signal_handler(signum, frame):
+        _finalize_telemetry(f"signal:{signum}")
+        prev = _prev_sigterm if signum == _signal_mod.SIGTERM else _prev_sigint
+        if callable(prev):
+            prev(signum, frame)  # ty: ignore[call-top-callable]
+            return
+        _signal_mod.signal(signum, _signal_mod.SIG_DFL)
+        os.kill(os.getpid(), signum)
+
+    _signal_mod.signal(_signal_mod.SIGTERM, _telemetry_signal_handler)
+    _signal_mod.signal(_signal_mod.SIGINT, _telemetry_signal_handler)
+
     try:
         if transport == "unix" and not socket_path:
             raise RuntimeError("socket-path is required in --config for --transport unix.")
@@ -547,8 +704,11 @@ def a2a(
             auto_approve_permissions=auto_approve_permissions,
         )
     except RuntimeError as exc:
+        exit_reason = "error"
         typer.echo(str(exc), err=True)
         raise typer.Exit(1) from exc
+    finally:
+        _finalize_telemetry()
 
 
 @a2a_client_app.command(name="call", help=_("Send a prompt to an A2A JSON-RPC endpoint."))

src/iac_code/providers/manager.py

@@ -302,7 +302,7 @@ async def stream(
                     response = await self._complete_with_retry(messages, system, tools, max_tokens)
                 except Exception as e:
                     self._emit_failure_telemetry(provider_name, sanitized_model, started, e)
-                    yield ErrorEvent(error=str(e), is_retryable=False)
+                    yield ErrorEvent(error=f"{type(e).__name__}: {str(e)[:1000]}", is_retryable=False)
                     return
                 span.set_attribute(GenAiAttr.RESPONSE_ID, response.message_id)
                 self._set_llm_response_span_attrs_from_response(span, response, self._model)
@@ -419,9 +419,9 @@ async def operation():
             except Exception as e:
                 status = getattr(e, "status_code", None) or getattr(e, "status", None)
                 if status and status in {408, 409, 429, 500, 502, 503, 529}:
-                    raise RetryableError(str(e), status_code=status) from e
+                    raise RetryableError(f"{type(e).__name__}: {e}", status_code=status) from e
                 if isinstance(e, (ConnectionError, TimeoutError, OSError)):
-                    raise RetryableError(str(e)) from e
+                    raise RetryableError(f"{type(e).__name__}: {e}") from e
                 raise
 
         try:

src/iac_code/services/telemetry/__init__.py

@@ -12,6 +12,7 @@
     "start_span",
     "bootstrap_telemetry",
     "graceful_shutdown",
+    "flush_telemetry",
     "get_client",
     "set_client",
     "get_session_id",
@@ -58,6 +59,16 @@ def graceful_shutdown() -> None:
     get_client().shutdown()
 
 
+def flush_telemetry() -> None:
+    """Force-flush pending telemetry without closing providers.
+
+    Safe to call repeatedly between units of work (e.g. per-task in a2a/acp
+    servers). Synchronous and bounded by the client's flush timeout — async
+    callers should wrap with ``asyncio.to_thread`` to avoid blocking the loop.
+    """
+    get_client().flush()
+
+
 def get_session_id() -> str:
     return get_client().get_session_id()
 

src/iac_code/services/telemetry/client.py

@@ -171,6 +171,22 @@ def bootstrap(self) -> None:
         # Retry previously failed batches
         self._retry_failed_batches()
 
+    def flush(self, timeout_ms: int = _FLUSH_TIMEOUT_MS) -> None:
+        """Force-flush providers without closing them; never raise.
+
+        Use this between units of work (e.g. per-task in a2a/acp servers) to
+        push pending batches before the runtime can be killed. Unlike
+        ``shutdown()``, the providers stay usable for subsequent work.
+        """
+        for provider, label in (
+            (self._meter_provider, "MeterProvider"),
+            (self._logger_provider, "LoggerProvider"),
+            (self._tracer_provider, "TracerProvider"),
+        ):
+            if provider is None:
+                continue
+            self._safe_force_flush(provider, label, timeout_ms)
+
     def shutdown(self) -> None:
         """Force-flush providers with bounded timeout; never raise."""
         for provider, label in (
@@ -315,13 +331,18 @@ def _retry_failed_batches(self) -> None:
                 log.warning("Failed to retry batch %s: %s", path, e)
 
     @staticmethod
-    def _safe_flush(provider: object, label: str) -> None:
+    def _safe_force_flush(provider: object, label: str, timeout_ms: int) -> None:
         flush = getattr(provider, "force_flush", None)
-        if flush is not None:
-            try:
-                flush(_FLUSH_TIMEOUT_MS)
-            except Exception as e:
-                log.warning("Flush %s failed: %s", label, e)
+        if flush is None:
+            return
+        try:
+            flush(timeout_ms)
+        except Exception as e:
+            log.warning("Flush %s failed: %s", label, e)
+
+    @classmethod
+    def _safe_flush(cls, provider: object, label: str) -> None:
+        cls._safe_force_flush(provider, label, _FLUSH_TIMEOUT_MS)
         shutdown = getattr(provider, "shutdown", None)
         if shutdown is not None:
             try: