Implement agent logging context manager and enhance agent execution logging

Add a context manager for binding per-agent execution fields to structlog contextvars in logging.py. Update the execute_agent function in executor.py to utilize this context manager, improving traceability and logging of agent execution events, including input and output messages.

Author: stormaref

src/sub_agent_mcp/agent/executor.py

@@ -2,10 +2,12 @@
 
 from __future__ import annotations
 
+import uuid
+
 from sub_agent_mcp.agent.builder import build_agent, get_recursion_limit
 from sub_agent_mcp.agent.errors import AgentExecutionError, AgentNotFoundError
 from sub_agent_mcp.config.schema import AgentConfig, AgentsFile
-from sub_agent_mcp.logging import get_logger
+from sub_agent_mcp.logging import agent_log_context, get_logger
 
 logger = get_logger(__name__)
 
@@ -24,49 +26,58 @@ async def execute_agent(agent: AgentConfig, prompt: str) -> str:
     from langchain_core.messages import HumanMessage
     from openai import AuthenticationError, OpenAIError
 
-    logger.info("spawn_agent_start", agent_id=agent.id)
-
-    try:
-        runnable = await build_agent(agent)
-        result = await runnable.ainvoke(
-            {"messages": [HumanMessage(content=prompt)]},
-            config={"recursion_limit": get_recursion_limit()},
-        )
-    except AuthenticationError as exc:
-        logger.error("agent_auth_failed", agent_id=agent.id)
-        raise AgentExecutionError(
-            f"Invalid credentials for agent '{agent.id}' LLM provider"
-        ) from exc
-    except OpenAIError as exc:
-        logger.error("agent_llm_error", agent_id=agent.id, error=str(exc))
-        raise AgentExecutionError(f"LLM error for agent '{agent.id}': {exc}") from exc
-    except Exception as exc:
-        if "401" in str(exc) or "authentication" in str(exc).lower():
-            logger.error("agent_auth_failed", agent_id=agent.id)
+    trace_id = str(uuid.uuid4())
+
+    with agent_log_context(
+        trace_id=trace_id,
+        model_id=agent.llm.model_id,
+        agent_id=agent.id,
+    ):
+        logger.info("spawn_agent_start")
+        logger.info("agent_input", input=prompt)
+
+        try:
+            runnable = await build_agent(agent)
+            result = await runnable.ainvoke(
+                {"messages": [HumanMessage(content=prompt)]},
+                config={"recursion_limit": get_recursion_limit()},
+            )
+        except AuthenticationError as exc:
+            logger.error("agent_auth_failed")
             raise AgentExecutionError(
                 f"Invalid credentials for agent '{agent.id}' LLM provider"
             ) from exc
-        logger.error("agent_execution_failed", agent_id=agent.id, error=str(exc))
-        raise AgentExecutionError(f"Agent execution failed for '{agent.id}': {exc}") from exc
-
-    messages = result.get("messages", [])
-    if not messages:
-        raise AgentExecutionError(f"Agent '{agent.id}' returned no messages")
-
-    final_message = messages[-1]
-    content = getattr(final_message, "content", None)
-    if content is None:
-        raise AgentExecutionError(f"Agent '{agent.id}' returned empty content")
-
-    if isinstance(content, list):
-        text_parts = [
-            part.get("text", "")
-            for part in content
-            if isinstance(part, dict) and part.get("type") == "text"
-        ]
-        response = "\n".join(part for part in text_parts if part)
-    else:
-        response = str(content)
-
-    logger.info("spawn_agent_complete", agent_id=agent.id)
-    return response
+        except OpenAIError as exc:
+            logger.error("agent_llm_error", error=str(exc))
+            raise AgentExecutionError(f"LLM error for agent '{agent.id}': {exc}") from exc
+        except Exception as exc:
+            if "401" in str(exc) or "authentication" in str(exc).lower():
+                logger.error("agent_auth_failed")
+                raise AgentExecutionError(
+                    f"Invalid credentials for agent '{agent.id}' LLM provider"
+                ) from exc
+            logger.error("agent_execution_failed", error=str(exc))
+            raise AgentExecutionError(f"Agent execution failed for '{agent.id}': {exc}") from exc
+
+        messages = result.get("messages", [])
+        if not messages:
+            raise AgentExecutionError(f"Agent '{agent.id}' returned no messages")
+
+        final_message = messages[-1]
+        content = getattr(final_message, "content", None)
+        if content is None:
+            raise AgentExecutionError(f"Agent '{agent.id}' returned empty content")
+
+        if isinstance(content, list):
+            text_parts = [
+                part.get("text", "")
+                for part in content
+                if isinstance(part, dict) and part.get("type") == "text"
+            ]
+            response = "\n".join(part for part in text_parts if part)
+        else:
+            response = str(content)
+
+        logger.info("agent_output", output=response)
+        logger.info("spawn_agent_complete")
+        return response

src/sub_agent_mcp/logging.py

@@ -4,6 +4,8 @@
 
 import logging
 import sys
+from collections.abc import Iterator
+from contextlib import contextmanager
 
 import structlog
 
@@ -32,3 +34,17 @@ def setup_logging(level: str = "INFO") -> None:
 
 def get_logger(name: str) -> structlog.stdlib.BoundLogger:
     return structlog.get_logger(name)
+
+
+@contextmanager
+def agent_log_context(*, trace_id: str, model_id: str, agent_id: str) -> Iterator[None]:
+    """Bind per-agent execution fields to structlog contextvars."""
+    structlog.contextvars.bind_contextvars(
+        trace_id=trace_id,
+        model_id=model_id,
+        agent_id=agent_id,
+    )
+    try:
+        yield
+    finally:
+        structlog.contextvars.clear_contextvars()

tests/test_logging.py

@@ -0,0 +1,116 @@
+"""Tests for agent execution logging."""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Iterator
+from contextlib import contextmanager
+from types import SimpleNamespace
+from unittest.mock import AsyncMock, patch
+
+import pytest
+import structlog
+from structlog.testing import LogCapture
+
+import sub_agent_mcp.agent.executor as executor_module
+from sub_agent_mcp.agent.errors import AgentExecutionError
+from sub_agent_mcp.agent.executor import spawn_agent
+from sub_agent_mcp.config.schema import AgentsFile
+from sub_agent_mcp.logging import get_logger
+
+
+@contextmanager
+def capture_agent_logs() -> Iterator[list[dict[str, object]]]:
+    """Capture structlog events including contextvars-bound fields."""
+    cap = LogCapture()
+    structlog.configure(
+        processors=[
+            structlog.contextvars.merge_contextvars,
+            structlog.processors.add_log_level,
+            cap,
+        ],
+        wrapper_class=structlog.make_filtering_bound_logger(logging.INFO),
+        logger_factory=structlog.stdlib.LoggerFactory(),
+        cache_logger_on_first_use=False,
+    )
+    executor_module.logger = get_logger(executor_module.__name__)
+    try:
+        yield cap.entries
+    finally:
+        structlog.reset_defaults()
+
+
+def _events_by_name(logs: list[dict[str, object]]) -> dict[str, list[dict[str, object]]]:
+    grouped: dict[str, list[dict[str, object]]] = {}
+    for entry in logs:
+        event = entry["event"]
+        assert isinstance(event, str)
+        grouped.setdefault(event, []).append(entry)
+    return grouped
+
+
+@pytest.mark.asyncio
+async def test_spawn_agent_success_logs_trace_model_and_io(agents_config: AgentsFile) -> None:
+    mock_runnable = AsyncMock()
+    mock_runnable.ainvoke.return_value = {
+        "messages": [SimpleNamespace(content="Research complete.")]
+    }
+    prompt = "Find info on MCP"
+
+    with (
+        patch("sub_agent_mcp.agent.executor.build_agent", return_value=mock_runnable),
+        capture_agent_logs() as logs,
+    ):
+        response = await spawn_agent(agents_config, "researcher", prompt)
+
+    assert response == "Research complete."
+
+    events = _events_by_name(logs)
+    trace_ids = {entry["trace_id"] for entry in logs if "trace_id" in entry}
+    assert len(trace_ids) == 1
+
+    trace_id = trace_ids.pop()
+    start_logs = events["spawn_agent_start"]
+    assert start_logs[0]["model_id"] == "gpt-4.1-mini"
+    assert start_logs[0]["agent_id"] == "researcher"
+    assert start_logs[0]["trace_id"] == trace_id
+
+    input_logs = events["agent_input"]
+    assert input_logs[0]["input"] == prompt
+    assert input_logs[0]["trace_id"] == trace_id
+
+    output_logs = events["agent_output"]
+    assert output_logs[0]["output"] == "Research complete."
+    assert output_logs[0]["trace_id"] == trace_id
+
+    complete_logs = events["spawn_agent_complete"]
+    assert complete_logs[0]["trace_id"] == trace_id
+
+
+@pytest.mark.asyncio
+async def test_spawn_agent_auth_error_logs_input_with_trace_id(agents_config: AgentsFile) -> None:
+    mock_runnable = AsyncMock()
+    mock_runnable.ainvoke.side_effect = Exception("401 Unauthorized")
+    prompt = "Hello"
+
+    with (
+        patch("sub_agent_mcp.agent.executor.build_agent", return_value=mock_runnable),
+        capture_agent_logs() as logs,
+        pytest.raises(AgentExecutionError, match="Invalid credentials"),
+    ):
+        await spawn_agent(agents_config, "researcher", prompt)
+
+    events = _events_by_name(logs)
+    trace_ids = {entry["trace_id"] for entry in logs if "trace_id" in entry}
+    assert len(trace_ids) == 1
+    trace_id = trace_ids.pop()
+
+    input_logs = events["agent_input"]
+    assert input_logs[0]["input"] == prompt
+    assert input_logs[0]["trace_id"] == trace_id
+
+    error_logs = events["agent_auth_failed"]
+    assert error_logs[0]["trace_id"] == trace_id
+    assert error_logs[0]["model_id"] == "gpt-4.1-mini"
+
+    assert "agent_output" not in events