diff --git a/strix/orchestration/hooks.py b/strix/orchestration/hooks.py index 8576b49..90fbff6 100644 --- a/strix/orchestration/hooks.py +++ b/strix/orchestration/hooks.py @@ -176,12 +176,13 @@ class StrixOrchestrationHooks(RunHooks[Any]): agent: Any, tool: Any, ) -> None: + del agent try: ctx = context.context if not isinstance(ctx, dict): return tracer = ctx.get("tracer") - if tracer is not None and hasattr(tracer, "log_tool_start"): + if tracer is not None: tracer.log_tool_start(ctx.get("agent_id", "?"), tool.name) except Exception: logger.exception("on_tool_start failed") @@ -193,6 +194,7 @@ class StrixOrchestrationHooks(RunHooks[Any]): tool: Any, result: str, ) -> None: + del agent try: ctx = context.context if not isinstance(ctx, dict): @@ -200,7 +202,7 @@ class StrixOrchestrationHooks(RunHooks[Any]): if tool.name in ("agent_finish", "finish_scan"): ctx["agent_finish_called"] = True tracer = ctx.get("tracer") - if tracer is not None and hasattr(tracer, "log_tool_end"): + if tracer is not None: tracer.log_tool_end(ctx.get("agent_id", "?"), tool.name, result) except Exception: logger.exception("on_tool_end failed") diff --git a/strix/telemetry/tracer.py b/strix/telemetry/tracer.py index d69433e..3f46ecb 100644 --- a/strix/telemetry/tracer.py +++ b/strix/telemetry/tracer.py @@ -42,6 +42,7 @@ class Tracer: self.agents: dict[str, dict[str, Any]] = {} self.tool_executions: dict[int, dict[str, Any]] = {} self.chat_messages: list[dict[str, Any]] = [] + self._next_exec_id = 1 self.vulnerability_reports: list[dict[str, Any]] = [] self.final_scan_result: str | None = None @@ -380,6 +381,42 @@ class Tracer: except (OSError, RuntimeError): logger.exception("Failed to save scan data") + def log_tool_start(self, agent_id: str, tool_name: str) -> int: + """Record a tool invocation in flight. Returns an exec_id.""" + exec_id = self._next_exec_id + self._next_exec_id += 1 + self.tool_executions[exec_id] = { + "agent_id": agent_id, + "tool_name": tool_name, + "status": "running", + "result": None, + "timestamp": datetime.now(UTC).isoformat(), + } + return exec_id + + def log_tool_end(self, agent_id: str, tool_name: str, result: Any) -> None: + """Mark the most recent matching exec as completed.""" + for exec_id in reversed(self.tool_executions): + entry = self.tool_executions[exec_id] + if ( + entry.get("agent_id") == agent_id + and entry.get("tool_name") == tool_name + and entry.get("status") == "running" + ): + entry["status"] = "completed" + entry["result"] = result + return + # No matching start (e.g. hooks added later in life) — record as completed. + exec_id = self._next_exec_id + self._next_exec_id += 1 + self.tool_executions[exec_id] = { + "agent_id": agent_id, + "tool_name": tool_name, + "status": "completed", + "result": result, + "timestamp": datetime.now(UTC).isoformat(), + } + def get_real_tool_count(self) -> int: return sum( 1