fix: address greptile review feedback for agent tracing PR

- Fix guardrail ui_type regression: restore isinstance(ut, str) guard to
  prevent AttributeError when ui_type is a plain string instead of an enum
- Fix disconnected context_id in A2A protocol: move trace_id initialization
  before the a2a_client branch so context_id stays connected to the trace
  chain even when a pre-created client is provided
- Remove hardcoded mock data from LogViewer: delete MOCK_AGENT_TRACE_SESSIONS
  and demo trace rendering that shipped unconditionally in the guardrails
  log viewer

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Krrish Dholakia 2026-03-03 18:14:36 -08:00
parent 6feb9babc1
commit 1f51b8cd0e
6 changed files with 20 additions and 478 deletions

View file

@ -243,13 +243,15 @@ async def asend_message(
if request is None:
raise ValueError("request is required")
# Ensure trace_id is always set for context_id propagation
trace_id = trace_id or str(uuid.uuid4())
# Create A2A client if not provided but api_base is available
if a2a_client is None:
if api_base is None:
raise ValueError(
"Either a2a_client or api_base is required for standard A2A flow"
)
trace_id = trace_id or str(uuid.uuid4())
extra_headers = {"X-LiteLLM-Trace-Id": trace_id}
if agent_id:
extra_headers["X-LiteLLM-Agent-Id"] = agent_id
@ -270,9 +272,8 @@ async def asend_message(
)
card_url = getattr(agent_card, "url", None) if agent_card else None
context_id = trace_id or str(uuid.uuid4())
if request.params.message.context_id is None:
request.params.message.context_id = context_id
request.params.message.context_id = trace_id
# Retry loop: if connection fails due to localhost URL in agent card, retry with fixed URL
a2a_response = None

View file

@ -1173,7 +1173,8 @@ def _build_field_dict(
# Check for custom UI type override
field_json_schema_extra = getattr(field, "json_schema_extra", {})
if field_json_schema_extra and "ui_type" in field_json_schema_extra:
field_type = field_json_schema_extra["ui_type"].value
ut = field_json_schema_extra["ui_type"]
field_type = ut if isinstance(ut, str) else ut.value
elif field_json_schema_extra and "type" in field_json_schema_extra:
field_type = field_json_schema_extra["type"]

View file

@ -5,4 +5,3 @@ export type { TraceViewProps } from "./TraceView";
export { SpanDetail } from "./SpanDetail";
export type { SpanDetailProps } from "./SpanDetail";
export type { Span, AgentTraceSession, SpanType, SpanStatus } from "../agentTraceTypes";
export { MOCK_AGENT_TRACE_SESSIONS } from "../agentTraceTypes";

View file

@ -1,4 +1,4 @@
import { act, render, screen } from "@testing-library/react";
import { render, screen } from "@testing-library/react";
import { QueryClient, QueryClientProvider } from "@tanstack/react-query";
import { describe, expect, it, vi } from "vitest";
import { LogViewer } from "./LogViewer";
@ -80,7 +80,7 @@ function wrapper({ children }: { children: React.ReactNode }) {
}
describe("LogViewer", () => {
it("should render agent-trace demo rows first then guardrail log entries", () => {
it("should render guardrail log entries", () => {
render(
<LogViewer
guardrailName="Test Guardrail"
@ -90,38 +90,22 @@ describe("LogViewer", () => {
{ wrapper }
);
// First 2 rows are the fixed agent-trace demo rows
expect(screen.getByText("Currency Research Agent")).toBeDefined();
expect(screen.getByText("Travel Booking Agent")).toBeDefined();
// Trace badge appears for demo rows
const traceBadges = screen.getAllByText("Trace");
expect(traceBadges.length).toBeGreaterThanOrEqual(2);
// Following rows are guardrail log entries (input snippets are unique to log rows)
// Log entries are rendered
expect(screen.getByText("First guardrail log request")).toBeDefined();
expect(screen.getByText("Second guardrail log request")).toBeDefined();
expect(screen.getByText("Third guardrail log request")).toBeDefined();
});
it("should open AgentTraceDrawer when first demo trace row is clicked", async () => {
it("should show empty state when no logs", () => {
render(
<LogViewer
guardrailName="Test Guardrail"
logs={mockLogs}
logs={[]}
logsLoading={false}
/>,
{ wrapper }
);
const currencyRow = screen.getByText("Currency Research Agent");
expect(screen.queryByTestId("agent-trace-drawer")).toBeNull();
await act(async () => {
currencyRow.closest("button")?.click();
});
expect(screen.getByTestId("agent-trace-drawer")).toBeDefined();
expect(screen.getByText(/Agent trace: Currency Research Agent/)).toBeDefined();
expect(screen.getByText(/Session: 0c4b4759-83aa/)).toBeDefined();
expect(screen.getByText("No logs to display. Adjust filters or date range.")).toBeDefined();
});
});

View file

@ -5,7 +5,7 @@ import {
WarningOutlined,
} from "@ant-design/icons";
import { useQuery } from "@tanstack/react-query";
import { Activity } from "lucide-react";
import moment from "moment";
import { Button, Spin } from "antd";
import React, { useState } from "react";
@ -13,10 +13,7 @@ import { uiSpendLogsCall } from "@/components/networking";
import { LogDetailsDrawer } from "@/components/view_logs/LogDetailsDrawer";
import type { LogEntry as ViewLogsLogEntry } from "@/components/view_logs/columns";
import { AgentTraceDrawer } from "./AgentTrace";
import {
type AgentTraceSession,
MOCK_AGENT_TRACE_SESSIONS,
} from "./agentTraceTypes";
import type { AgentTraceSession } from "./agentTraceTypes";
import type { LogEntry } from "./mockData";
const actionConfig: Record<
@ -74,7 +71,6 @@ export function LogViewer({
const [selectedTraceSession, setSelectedTraceSession] =
useState<AgentTraceSession | null>(null);
const demoTraces = MOCK_AGENT_TRACE_SESSIONS.slice(0, 2);
const filteredLogs = logs.filter(
(log) => activeFilter === "all" || log.action === activeFilter
);
@ -121,10 +117,6 @@ export function LogViewer({
setDrawerOpen(true);
};
const handleTraceSessionClick = (session: AgentTraceSession) => {
setSelectedTraceSession(session);
};
const handleCloseDrawer = () => {
setDrawerOpen(false);
setSelectedRequestId(null);
@ -141,11 +133,9 @@ export function LogViewer({
<p className="text-xs text-gray-500 mt-0.5">
{logsLoading
? "Loading…"
: demoTraces.length > 0
? `2 demo traces, then showing ${displayLogs.length} of ${total} log entries`
: logs.length > 0
? `Showing ${displayLogs.length} of ${total} entries`
: "No logs for this period. Select a guardrail and date range."}
: logs.length > 0
? `Showing ${displayLogs.length} of ${total} entries`
: "No logs for this period. Select a guardrail and date range."}
</p>
</div>
{logs.length > 0 && (
@ -186,64 +176,13 @@ export function LogViewer({
<Spin />
</div>
)}
{!logsLoading && demoTraces.length === 0 && displayLogs.length === 0 && (
{!logsLoading && displayLogs.length === 0 && (
<div className="py-12 text-center text-sm text-gray-500">
No logs to display. Adjust filters or date range.
</div>
)}
{!logsLoading && (demoTraces.length > 0 || displayLogs.length > 0) && (
{!logsLoading && displayLogs.length > 0 && (
<div className="divide-y divide-gray-100">
{demoTraces.map((session) => {
const statusDotColor =
session.status === "success"
? "bg-green-500"
: session.status === "error"
? "bg-red-500"
: session.status === "running"
? "bg-blue-500"
: "bg-amber-500";
return (
<button
key={session.id}
type="button"
onClick={() => handleTraceSessionClick(session)}
className="w-full text-left px-4 py-3 hover:bg-gray-50 transition-colors flex items-start gap-3"
>
<Activity className="w-4 h-4 mt-0.5 flex-shrink-0 text-indigo-500" />
<div className="flex-1 min-w-0">
<div className="flex items-center gap-2 mb-1 flex-wrap">
<span className="inline-flex items-center px-2 py-0.5 text-xs font-medium rounded border bg-indigo-50 text-indigo-600 border-indigo-200">
Trace
</span>
<span className="text-xs text-gray-400">
{session.relativeTime}
</span>
<span className="text-xs text-gray-400">·</span>
<span className="text-xs text-gray-500">
{session.totalSpans} spans
</span>
<span className="text-xs text-gray-400">·</span>
<span className="text-xs text-gray-500">
{session.totalDurationMs}ms
</span>
<span className="text-xs text-gray-400">·</span>
<span className="text-xs text-gray-500">
${session.totalCost.toFixed(4)}
</span>
<span
className={`w-1.5 h-1.5 rounded-full flex-shrink-0 ${statusDotColor}`}
title={session.status}
aria-hidden
/>
</div>
<p className="text-sm text-gray-800 truncate font-medium">
{session.rootAgentName}
</p>
</div>
<DownOutlined className="w-4 h-4 text-gray-400 flex-shrink-0 mt-1" />
</button>
);
})}
{displayLogs.map((log) => {
const config = actionConfig[log.action];
const ActionIcon = config.icon;

View file

@ -1,5 +1,5 @@
/**
* Types and mock data for agent trace demo in GuardrailsMonitor LogViewer.
* Types for agent trace visualization in GuardrailsMonitor LogViewer.
* Supports orchestrator → sub-agents → LLM + MCP hierarchy.
*/
@ -41,385 +41,3 @@ export interface AgentTraceSession {
status: SpanStatus;
spans: Span[];
}
/** First 2 sessions for demo: always shown as first 1–2 rows in LogViewer. */
export const MOCK_AGENT_TRACE_SESSIONS: AgentTraceSession[] = [
{
id: "0c4b4759-83aa-4f2e-b91c-d3e7f8a12345",
shortId: "0c4b4759-83aa",
rootAgentName: "Currency Research Agent",
timestamp: "2026-03-02 09:36:12 PM",
relativeTime: "2 min ago",
totalSpans: 10,
totalCost: 0.00366,
totalDurationMs: 2677,
status: "success",
spans: [
{
id: "span-orch-001",
name: "OrchestratorAgent",
type: "orchestrator",
status: "success",
startMs: 0,
durationMs: 2677,
input: JSON.stringify(
{
task: "Research current currency exchange rates and recent news for USD/EUR, USD/GBP",
context: { user_id: "u_8821", session: "0c4b4759" },
},
null,
2
),
output: JSON.stringify(
{
summary:
"Completed currency research. EUR/USD at 1.0842 (+0.23%), GBP/USD at 1.2634 (+0.11%). Key news: ECB holds rates, UK inflation data due Thursday.",
agents_used: ["CurrencyAgent", "NewsAgent"],
total_tokens: 1026,
},
null,
2
),
children: [
{
id: "span-currency-001",
name: "CurrencyAgent",
type: "agent",
status: "success",
startMs: 50,
durationMs: 2100,
input: JSON.stringify(
{
task: "Fetch current exchange rates and 7-day history for USD/EUR and USD/GBP",
tools: ["forex-api"],
},
null,
2
),
output: JSON.stringify(
{
rates: { "USD/EUR": 1.0842, "USD/GBP": 1.2634 },
trend: "bullish",
confidence: 0.87,
},
null,
2
),
children: [
{
id: "span-llm-001",
name: "gpt-4o",
type: "llm",
status: "success",
startMs: 100,
durationMs: 1346,
model: "gpt-4o",
tokens: { prompt: 198, completion: 74 },
cost: 0.00096,
input: JSON.stringify(
{
messages: [
{
role: "system",
content:
"You are a currency research agent. Use the available tools to fetch exchange rates.",
},
{
role: "user",
content:
"Get current USD/EUR and USD/GBP rates and 7-day history.",
},
],
tools: [
{ name: "get_rates", description: "Get current exchange rates" },
{ name: "get_history", description: "Get historical rate data" },
],
},
null,
2
),
output: JSON.stringify(
{
role: "assistant",
content: null,
tool_calls: [
{
id: "call_abc123",
type: "function",
function: {
name: "get_rates",
arguments: '{"pairs": ["USD/EUR", "USD/GBP"]}',
},
},
],
},
null,
2
),
children: [
{
id: "span-mcp-001",
name: "forex-api/get_rates",
type: "mcp",
status: "success",
startMs: 150,
durationMs: 234,
mcpServer: "forex-api",
mcpTool: "get_rates",
input: JSON.stringify(
{ pairs: ["USD/EUR", "USD/GBP"], timestamp: "latest" },
null,
2
),
output: JSON.stringify(
{
USD_EUR: 1.0842,
USD_GBP: 1.2634,
timestamp: "2026-03-02T21:36:12Z",
source: "ECB",
},
null,
2
),
},
{
id: "span-mcp-002",
name: "forex-api/get_history",
type: "mcp",
status: "success",
startMs: 400,
durationMs: 189,
mcpServer: "forex-api",
mcpTool: "get_history",
input: JSON.stringify(
{ pairs: ["USD/EUR", "USD/GBP"], days: 7 },
null,
2
),
output: JSON.stringify(
{
data: [
{ date: "2026-02-24", USD_EUR: 1.0791 },
{ date: "2026-02-25", USD_EUR: 1.0803 },
{ date: "2026-02-26", USD_EUR: 1.0821 },
{ date: "2026-02-27", USD_EUR: 1.0834 },
{ date: "2026-02-28", USD_EUR: 1.0842 },
],
},
null,
2
),
},
],
},
{
id: "span-llm-002",
name: "gpt-4o",
type: "llm",
status: "success",
startMs: 1500,
durationMs: 826,
model: "gpt-4o",
tokens: { prompt: 171, completion: 44 },
cost: 0.00087,
input: JSON.stringify(
{
messages: [
{
role: "system",
content:
"Analyze the fetched currency data and provide a summary.",
},
{
role: "user",
content: "Summarize the exchange rate data and trends.",
},
{
role: "tool",
content: '{"USD_EUR": 1.0842, "USD_GBP": 1.2634}',
},
],
},
null,
2
),
output: JSON.stringify(
{
role: "assistant",
content:
"EUR/USD is trading at 1.0842, up 0.23% over the past week showing a steady bullish trend. GBP/USD at 1.2634 with modest gains of 0.11%. Both pairs showing positive momentum.",
},
null,
2
),
},
],
},
{
id: "span-news-001",
name: "NewsAgent",
type: "agent",
status: "success",
startMs: 50,
durationMs: 1800,
input: JSON.stringify(
{
task: "Search for recent news about USD/EUR and USD/GBP exchange rates",
tools: ["web-search"],
},
null,
2
),
output: JSON.stringify(
{
articles: 3,
sentiment: "neutral",
key_events: ["ECB rate hold", "UK CPI data pending"],
},
null,
2
),
children: [
{
id: "span-llm-003",
name: "gpt-4o",
type: "llm",
status: "success",
startMs: 100,
durationMs: 923,
model: "gpt-4o",
tokens: { prompt: 245, completion: 96 },
cost: 0.00112,
input: JSON.stringify(
{
messages: [
{
role: "system",
content: "You are a financial news research agent.",
},
{
role: "user",
content:
"Find recent news about EUR and GBP currency movements.",
},
],
tools: [
{ name: "query", description: "Search the web" },
{ name: "fetch", description: "Fetch a URL" },
],
},
null,
2
),
output: JSON.stringify(
{
role: "assistant",
content: null,
tool_calls: [
{
id: "call_def456",
type: "function",
function: {
name: "query",
arguments:
'{"q": "EUR USD GBP exchange rate news March 2026"}',
},
},
],
},
null,
2
),
children: [
{
id: "span-mcp-003",
name: "web-search/query",
type: "mcp",
status: "success",
startMs: 150,
durationMs: 445,
mcpServer: "web-search",
mcpTool: "query",
input: JSON.stringify(
{ q: "EUR USD GBP exchange rate news March 2026", num_results: 5 },
null,
2
),
output: JSON.stringify(
{
results: [
{
title: "ECB Holds Rates, Euro Steady Against Dollar",
url: "https://reuters.com/...",
snippet:
"The European Central Bank held rates steady...",
},
{
title: "UK Inflation Data Due Thursday, GBP Cautious",
url: "https://ft.com/...",
snippet: "Markets await UK CPI data...",
},
],
},
null,
2
),
},
],
},
{
id: "span-llm-004",
name: "gpt-4o",
type: "llm",
status: "success",
startMs: 1100,
durationMs: 654,
model: "gpt-4o",
tokens: { prompt: 156, completion: 42 },
cost: 0.00071,
input: JSON.stringify(
{
messages: [
{
role: "system",
content: "Summarize the news findings into key insights.",
},
{
role: "tool",
content:
"ECB holds rates steady. UK CPI data pending Thursday.",
},
],
},
null,
2
),
output: JSON.stringify(
{
role: "assistant",
content:
"Key findings: (1) ECB held rates steady — neutral for EUR short-term. (2) UK CPI data due Thursday — potential GBP volatility. (3) Overall sentiment: cautiously neutral with slight bullish bias.",
},
null,
2
),
},
],
},
],
},
],
},
{
id: "6419f0a4-82ef-4b3c-a921-e8f2d3c45678",
shortId: "6419f0a4-82ef",
rootAgentName: "Travel Booking Agent",
timestamp: "2026-03-02 09:34:27 PM",
relativeTime: "4 min ago",
totalSpans: 4,
totalCost: 0.0031,
totalDurationMs: 3200,
status: "success",
spans: [],
},
];