add-manual-tracing

name	add-manual-tracing
description	Adds manual MLflow trace wrapping for tool and agent spans in Level B and C agents where autolog doesn't cover everything.
argument-hint	<agent_path>
disable-model-invocation	true

Add Manual Tracing (Level B and C)

Usage: /add-manual-tracing <agent_path> Example: /add-manual-tracing agents/autogen/chat_agent

You are adding manual wrap_func_with_mlflow_trace() calls to an agent template where autolog does not fully cover all tracing layers.

This skill is only for Level B (partial autolog) and Level C (no framework autolog). If the framework is Level A, skip this skill entirely — autolog handles everything.

Input

The agent path is: $ARGUMENTS

You also need the package name, coverage level (B or C), and autolog support report. If not provided, determine the package name from pyproject.toml or src/, and the coverage level from the agent's tracing.py.

Steps

1. Read the agent's code

Read these files to understand how the agent works:

• <agent_path>/main.py — FastAPI app, _handle_chat, _handle_stream
• <agent_path>/src/<package>/agent.py — Agent class/factory, how tools are registered, the main entry point
• <agent_path>/src/<package>/tools.py — Tool function definitions
• Any other file where the agent or tools are assembled (e.g., crew.py for CrewAI)

Identify:

• The agent entry point: The function that runs the full agent loop (e.g., agent.query(), agent.run(), crew.kickoff())
• Tool registration: How tools are attached to the agent (function list, tool objects, decorator-based, etc.)
• Streaming vs non-streaming: Whether the streaming path creates the agent differently

2. Determine what needs manual wrapping

Compare the autolog report with the three tracing layers:

Layer	Needs manual wrapping if...
Agent orchestration	Autolog doesn't create a parent AGENT span for the full request
Tool execution	Autolog doesn't capture tool calls with TOOL spans
LLM calls	Autolog doesn't capture model API calls (rare — usually covered by provider autolog)

3. Add tool wrapping

Wrap each tool so its execution creates a TOOL span. The wrapping location depends on how the agent registers tools:

If tools are functions registered by name (common in Level C):

In agent.py or wherever tools are registered (e.g., from openai_responses_agent.tracing import wrap_func_with_mlflow_trace):

from <package>.tracing import wrap_func_with_mlflow_trace

# Wrap each tool function before registering
for name, func in self._tools:
    func = wrap_func_with_mlflow_trace(func, span_type="tool")
    agent.register_tool(name, func)

If tools are class instances with a _run method (common in Level B, e.g., CrewAI):

In the file where tool objects are created, e.g., crew.py (e.g., from crewai_web_search.tracing import wrap_func_with_mlflow_trace):

from <package>.tracing import wrap_func_with_mlflow_trace

tools = [MyTool()]
for tool in tools:
    tool._run = wrap_func_with_mlflow_trace(tool._run, span_type="tool", name=tool.name)

If tools are decorated functions (e.g., @tool decorator):

Wrap the underlying function after the agent is assembled:

from <package>.tracing import wrap_func_with_mlflow_trace

for tool in agent.tools:
    tool.func = wrap_func_with_mlflow_trace(tool.func, span_type="tool")

4. Add agent orchestration wrapping

Wrap the main agent entry point to create a parent AGENT span that groups all LLM calls and tool calls under one trace.

In agent.py (preferred — wrap inside the adapter/closure):

agent.query = wrap_func_with_mlflow_trace(agent.query, span_type="agent")

Or in main.py if the agent is created directly there:

# In _handle_chat:
agent = get_agent()
agent.run = wrap_func_with_mlflow_trace(agent.run, span_type="agent")
result = await agent.run(input=messages)

5. Handle the streaming path

This is critical. Read _handle_stream in main.py carefully.

If streaming uses the same agent instance as non-streaming (goes through the same closure/adapter):

• No extra work — the wrapping from step 3-4 applies to both paths.

If streaming creates a new agent instance directly (bypasses the adapter):

• You MUST duplicate the wrapping inside the streaming path.
• This is the pattern used by the Vanilla Python agent. See the run_agent() function inside _handle_stream in agents/vanilla_python/openai_responses_agent/main.py:

# Inside _handle_stream's run_agent():
def run_agent():
    adapter = get_agent()
    agent = SomeAgent(model=adapter._model_id, ...)

    # Wrap tools
    for name, func in adapter._tools:
        func = wrap_func_with_mlflow_trace(func, span_type="tool")
        agent.register_tool(name, func)

    # Wrap agent entry point
    agent.query = wrap_func_with_mlflow_trace(agent.query, span_type="agent")

    return agent.query(user_message, ...)

Why this matters: Without a parent AGENT span, mlflow.<provider>.autolog() creates a separate trace for every LLM call instead of grouping them under one trace. This results in N traces per request instead of 1.

6. Verify wrapping is conditional

Confirm that wrap_func_with_mlflow_trace() in tracing.py returns the original function unchanged when MLFLOW_TRACKING_URI is not set. This means the wrapping calls in agent code are always safe — they're no-ops when tracing is disabled.

Reference Files

• Level B tool wrapping: agents/crewai/websearch_agent/src/crewai_web_search/crew.py — ai_assistant() method
• Level C tool + agent wrapping (non-streaming): agents/vanilla_python/openai_responses_agent/src/openai_responses_agent/agent.py — _AIAgentAdapter.run() method
• Level C streaming path wrapping: agents/vanilla_python/openai_responses_agent/main.py — run_agent() inside _handle_stream()

Final Checklist

• Every tool function/method is wrapped with span_type="tool"
• The main agent entry point is wrapped with span_type="agent"
• Both streaming and non-streaming paths have wrapping applied
• Wrapping is safe when tracing is disabled (no-op via wrap_func_with_mlflow_trace)
• No duplicate wrapping — each function is wrapped exactly once per request
• The name parameter is used when wrapping tool objects (so span names are meaningful, not _run)

Self-Update

Before finishing, check whether this skill file needs updating. If any of the following are true, propose the specific changes to the user and only update this file if they approve:

• You deviated from these instructions because they were inaccurate, outdated, or insufficient — and your deviation worked.
• You encountered a new pattern not covered here (e.g., a framework that uses OpenTelemetry instead of autolog, or a new provider path).
• File paths, function names, or API patterns referenced here have changed.
• A step didn't work as described and needed a different approach.

If nothing needed changing, move on.