wehub-resource-sync
906 lines
30 KiB
Python
906 lines
30 KiB
Python
import importlib
|
|
import json
|
|
import logging
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from claude_agent_sdk.types import (
|
|
AssistantMessage,
|
|
ResultMessage,
|
|
TextBlock,
|
|
ToolResultBlock,
|
|
ToolUseBlock,
|
|
UserMessage,
|
|
)
|
|
|
|
import mlflow
|
|
import mlflow.claude_code.tracing as tracing_module
|
|
from mlflow.claude_code.tracing import (
|
|
CLAUDE_TRACING_LEVEL,
|
|
METADATA_KEY_CLAUDE_CODE_VERSION,
|
|
find_last_user_message_index,
|
|
get_hook_response,
|
|
parse_timestamp_to_ns,
|
|
process_sdk_messages,
|
|
process_transcript,
|
|
setup_logging,
|
|
)
|
|
from mlflow.entities.span import SpanType
|
|
from mlflow.tracing.constant import SpanAttributeKey, TraceMetadataKey
|
|
|
|
# ============================================================================
|
|
# TIMESTAMP PARSING TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_parse_timestamp_to_ns_iso_string():
|
|
iso_timestamp = "2024-01-15T10:30:45.123456Z"
|
|
result = parse_timestamp_to_ns(iso_timestamp)
|
|
|
|
# Verify it returns an integer (nanoseconds)
|
|
assert isinstance(result, int)
|
|
assert result > 0
|
|
|
|
|
|
def test_parse_timestamp_to_ns_unix_seconds():
|
|
unix_timestamp = 1705312245.123456
|
|
result = parse_timestamp_to_ns(unix_timestamp)
|
|
|
|
# Should convert seconds to nanoseconds
|
|
expected = int(unix_timestamp * 1_000_000_000)
|
|
assert result == expected
|
|
|
|
|
|
def test_parse_timestamp_to_ns_large_number():
|
|
large_timestamp = 1705312245123
|
|
result = parse_timestamp_to_ns(large_timestamp)
|
|
|
|
# Function treats large numbers as seconds and converts to nanoseconds
|
|
# Just verify we get a reasonable nanosecond value
|
|
assert isinstance(result, int)
|
|
assert result > 0
|
|
|
|
|
|
# ============================================================================
|
|
# LOGGING TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_setup_logging_creates_logger(monkeypatch, tmp_path):
|
|
monkeypatch.chdir(tmp_path)
|
|
logger = setup_logging()
|
|
|
|
# Verify logger was created
|
|
assert logger is not None
|
|
assert logger.name == "mlflow.claude_code.tracing"
|
|
|
|
# Verify log directory was created
|
|
log_dir = tmp_path / ".claude" / "mlflow"
|
|
assert log_dir.exists()
|
|
assert log_dir.is_dir()
|
|
|
|
|
|
def test_custom_logging_level():
|
|
setup_logging()
|
|
|
|
assert CLAUDE_TRACING_LEVEL > logging.INFO
|
|
assert CLAUDE_TRACING_LEVEL < logging.WARNING
|
|
assert logging.getLevelName(CLAUDE_TRACING_LEVEL) == "CLAUDE_TRACING"
|
|
|
|
|
|
def test_get_logger_lazy_initialization(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> None:
|
|
monkeypatch.chdir(tmp_path)
|
|
|
|
# Force reload to reset the module state
|
|
importlib.reload(tracing_module)
|
|
|
|
log_dir = tmp_path / ".claude" / "mlflow"
|
|
|
|
# Before calling get_logger(), the log directory should NOT exist
|
|
assert not log_dir.exists()
|
|
|
|
# Call get_logger() for the first time - this should trigger initialization
|
|
logger1 = tracing_module.get_logger()
|
|
|
|
# After calling get_logger(), the log directory SHOULD exist
|
|
assert log_dir.exists()
|
|
assert log_dir.is_dir()
|
|
|
|
# Verify logger was created properly
|
|
assert logger1 is not None
|
|
assert logger1.name == "mlflow.claude_code.tracing"
|
|
|
|
# Call get_logger() again - should return the same logger instance
|
|
logger2 = tracing_module.get_logger()
|
|
assert logger2 is logger1
|
|
|
|
|
|
# ============================================================================
|
|
# HOOK RESPONSE TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_get_hook_response_success():
|
|
response = get_hook_response()
|
|
assert response == {"continue": True}
|
|
|
|
|
|
def test_get_hook_response_with_error():
|
|
response = get_hook_response(error="Test error")
|
|
assert response == {"continue": False, "stopReason": "Test error"}
|
|
|
|
|
|
def test_get_hook_response_with_additional_fields():
|
|
response = get_hook_response(custom_field="value")
|
|
assert response == {"continue": True, "custom_field": "value"}
|
|
|
|
|
|
# ============================================================================
|
|
# ASYNC TRACE LOGGING UTILITY TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_flush_trace_async_logging_calls_flush(monkeypatch):
|
|
mock_exporter = type("MockExporter", (), {"_async_queue": True})()
|
|
monkeypatch.setattr(tracing_module, "_get_trace_exporter", lambda: mock_exporter)
|
|
flushed = []
|
|
monkeypatch.setattr(mlflow, "flush_trace_async_logging", lambda: flushed.append(True))
|
|
tracing_module._flush_trace_async_logging()
|
|
assert len(flushed) == 1
|
|
|
|
|
|
def test_flush_trace_async_logging_skips_without_async_queue(monkeypatch):
|
|
mock_exporter = object() # no _async_queue attribute
|
|
monkeypatch.setattr(tracing_module, "_get_trace_exporter", lambda: mock_exporter)
|
|
flushed = []
|
|
monkeypatch.setattr(mlflow, "flush_trace_async_logging", lambda: flushed.append(True))
|
|
tracing_module._flush_trace_async_logging()
|
|
assert len(flushed) == 0
|
|
|
|
|
|
# ============================================================================
|
|
# INTEGRATION TESTS
|
|
# ============================================================================
|
|
|
|
# Sample Claude Code transcript for testing
|
|
DUMMY_TRANSCRIPT = [
|
|
{
|
|
"type": "user",
|
|
"message": {"role": "user", "content": "What is 2 + 2?"},
|
|
"timestamp": "2025-01-15T10:00:00.000Z",
|
|
"sessionId": "test-session-123",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "Let me calculate that for you."}],
|
|
},
|
|
"timestamp": "2025-01-15T10:00:01.000Z",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "tool_use",
|
|
"id": "tool_123",
|
|
"name": "Bash",
|
|
"input": {"command": "echo $((2 + 2))"},
|
|
}
|
|
],
|
|
},
|
|
"timestamp": "2025-01-15T10:00:02.000Z",
|
|
},
|
|
{
|
|
"type": "user",
|
|
"message": {
|
|
"role": "user",
|
|
"content": [{"type": "tool_result", "tool_use_id": "tool_123", "content": "4"}],
|
|
},
|
|
"timestamp": "2025-01-15T10:00:03.000Z",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "The answer is 4."}],
|
|
},
|
|
"timestamp": "2025-01-15T10:00:04.000Z",
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_transcript_file(tmp_path):
|
|
transcript_path = tmp_path / "transcript.jsonl"
|
|
with open(transcript_path, "w") as f:
|
|
for entry in DUMMY_TRANSCRIPT:
|
|
f.write(json.dumps(entry) + "\n")
|
|
return str(transcript_path)
|
|
|
|
|
|
def test_process_transript_creates_trace(mock_transcript_file):
|
|
trace = process_transcript(mock_transcript_file, "test-session-123")
|
|
|
|
# Verify trace was created
|
|
assert trace is not None
|
|
|
|
# Verify trace has spans
|
|
spans = list(trace.search_spans())
|
|
assert len(spans) > 0
|
|
|
|
# Verify root span and metadata
|
|
root_span = trace.data.spans[0]
|
|
assert root_span.name == "claude_code_conversation"
|
|
assert root_span.span_type == SpanType.AGENT
|
|
assert trace.info.trace_metadata.get("mlflow.trace.session") == "test-session-123"
|
|
|
|
|
|
def test_process_transcript_creates_spans(mock_transcript_file):
|
|
trace = process_transcript(mock_transcript_file, "test-session-123")
|
|
|
|
assert trace is not None
|
|
|
|
# Verify trace has spans
|
|
spans = list(trace.search_spans())
|
|
assert len(spans) > 0
|
|
|
|
# Find LLM and tool spans
|
|
llm_spans = [s for s in spans if s.span_type == SpanType.LLM]
|
|
tool_spans = [s for s in spans if s.span_type == SpanType.TOOL]
|
|
|
|
assert len(llm_spans) == 2
|
|
assert len(tool_spans) == 1
|
|
|
|
# Verify tool span has proper attributes
|
|
tool_span = tool_spans[0]
|
|
assert tool_span.name == "tool_Bash"
|
|
|
|
# Verify LLM spans have MESSAGE_FORMAT set to "anthropic" for Chat UI rendering
|
|
for llm_span in llm_spans:
|
|
assert llm_span.get_attribute(SpanAttributeKey.MESSAGE_FORMAT) == "anthropic"
|
|
|
|
# Verify LLM span outputs are in Anthropic response format
|
|
first_llm = llm_spans[0]
|
|
outputs = first_llm.outputs
|
|
assert outputs["type"] == "message"
|
|
assert outputs["role"] == "assistant"
|
|
assert isinstance(outputs["content"], list)
|
|
|
|
# Verify LLM span inputs contain messages in Anthropic format
|
|
inputs = first_llm.inputs
|
|
assert "messages" in inputs
|
|
messages = inputs["messages"]
|
|
assert any(m["role"] == "user" for m in messages)
|
|
|
|
|
|
def test_process_transcript_returns_none_for_nonexistent_file():
|
|
result = process_transcript("/nonexistent/path/transcript.jsonl", "test-session-123")
|
|
assert result is None
|
|
|
|
|
|
def test_process_transcript_links_trace_to_run(mock_transcript_file):
|
|
with mlflow.start_run() as run:
|
|
trace = process_transcript(mock_transcript_file, "test-session-123")
|
|
|
|
assert trace is not None
|
|
assert trace.info.trace_metadata.get(TraceMetadataKey.SOURCE_RUN) == run.info.run_id
|
|
|
|
|
|
# Sample Claude Code transcript with token usage for testing
|
|
DUMMY_TRANSCRIPT_WITH_USAGE = [
|
|
{
|
|
"type": "user",
|
|
"message": {"role": "user", "content": "Hello Claude!"},
|
|
"timestamp": "2025-01-15T10:00:00.000Z",
|
|
"sessionId": "test-session-usage",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "Hello! How can I help you today?"}],
|
|
"model": "claude-sonnet-4-20250514",
|
|
"usage": {"input_tokens": 150, "output_tokens": 25},
|
|
},
|
|
"timestamp": "2025-01-15T10:00:01.000Z",
|
|
},
|
|
]
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_transcript_file_with_usage(tmp_path):
|
|
transcript_path = tmp_path / "transcript_with_usage.jsonl"
|
|
with open(transcript_path, "w") as f:
|
|
for entry in DUMMY_TRANSCRIPT_WITH_USAGE:
|
|
f.write(json.dumps(entry) + "\n")
|
|
return str(transcript_path)
|
|
|
|
|
|
def test_process_transcript_tracks_token_usage(mock_transcript_file_with_usage):
|
|
trace = process_transcript(mock_transcript_file_with_usage, "test-session-usage")
|
|
|
|
assert trace is not None
|
|
|
|
# Find the LLM span
|
|
spans = list(trace.search_spans())
|
|
llm_spans = [s for s in spans if s.span_type == SpanType.LLM]
|
|
|
|
assert len(llm_spans) == 1
|
|
llm_span = llm_spans[0]
|
|
|
|
# Verify token usage is tracked using the standardized CHAT_USAGE attribute
|
|
token_usage = llm_span.get_attribute(SpanAttributeKey.CHAT_USAGE)
|
|
assert token_usage is not None
|
|
assert token_usage["input_tokens"] == 150
|
|
assert token_usage["output_tokens"] == 25
|
|
assert token_usage["total_tokens"] == 175
|
|
|
|
# Verify trace-level token usage aggregation works
|
|
assert trace.info.token_usage is not None
|
|
assert trace.info.token_usage["input_tokens"] == 150
|
|
assert trace.info.token_usage["output_tokens"] == 25
|
|
assert trace.info.token_usage["total_tokens"] == 175
|
|
|
|
|
|
def test_process_transcript_preserves_cache_tokens(tmp_path):
|
|
"""Verify cache_read/cache_creation fields from Anthropic usage survive on the
|
|
CHAT_USAGE span attribute so prompt-cache hit rate is observable.
|
|
"""
|
|
transcript_entries = [
|
|
{
|
|
"type": "user",
|
|
"message": {"role": "user", "content": "Cached prompt"},
|
|
"timestamp": "2025-01-15T10:00:00.000Z",
|
|
"sessionId": "cache-transcript-session",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "Answer using cache."}],
|
|
"model": "claude-sonnet-4-20250514",
|
|
"usage": {
|
|
"input_tokens": 36,
|
|
"cache_creation_input_tokens": 23554,
|
|
"cache_read_input_tokens": 139035,
|
|
"output_tokens": 3344,
|
|
},
|
|
},
|
|
"timestamp": "2025-01-15T10:00:01.000Z",
|
|
},
|
|
]
|
|
|
|
transcript_path = tmp_path / "transcript_cache.jsonl"
|
|
with open(transcript_path, "w") as f:
|
|
for entry in transcript_entries:
|
|
f.write(json.dumps(entry) + "\n")
|
|
|
|
trace = process_transcript(str(transcript_path), "cache-transcript-session")
|
|
|
|
assert trace is not None
|
|
llm_spans = [s for s in trace.search_spans() if s.span_type == SpanType.LLM]
|
|
assert len(llm_spans) == 1
|
|
|
|
# input_tokens is the non-cached input the Anthropic API reports, matching
|
|
# mlflow.anthropic.autolog. Cache fields are exposed as separate keys so
|
|
# consumers can compute cache hit rate.
|
|
token_usage = llm_spans[0].get_attribute(SpanAttributeKey.CHAT_USAGE)
|
|
assert token_usage["input_tokens"] == 36
|
|
assert token_usage["output_tokens"] == 3344
|
|
assert token_usage["total_tokens"] == 36 + 3344
|
|
assert token_usage["cache_read_input_tokens"] == 139035
|
|
assert token_usage["cache_creation_input_tokens"] == 23554
|
|
|
|
|
|
# ============================================================================
|
|
# SDK MESSAGE PROCESSING TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_process_sdk_messages_empty_list():
|
|
assert process_sdk_messages([]) is None
|
|
|
|
|
|
def test_process_sdk_messages_no_user_prompt():
|
|
messages = [
|
|
AssistantMessage(
|
|
content=[TextBlock(text="Hello!")],
|
|
model="claude-sonnet-4-20250514",
|
|
),
|
|
]
|
|
assert process_sdk_messages(messages) is None
|
|
|
|
|
|
def test_process_sdk_messages_simple_conversation():
|
|
messages = [
|
|
UserMessage(content="What is 2 + 2?"),
|
|
AssistantMessage(
|
|
content=[TextBlock(text="The answer is 4.")],
|
|
model="claude-sonnet-4-20250514",
|
|
),
|
|
ResultMessage(
|
|
subtype="success",
|
|
duration_ms=1000,
|
|
duration_api_ms=800,
|
|
is_error=False,
|
|
num_turns=1,
|
|
session_id="test-sdk-session",
|
|
usage={"input_tokens": 100, "output_tokens": 20},
|
|
),
|
|
]
|
|
|
|
trace = process_sdk_messages(messages, "test-sdk-session")
|
|
|
|
assert trace is not None
|
|
spans = list(trace.search_spans())
|
|
|
|
root_span = trace.data.spans[0]
|
|
assert root_span.name == "claude_code_conversation"
|
|
assert root_span.span_type == SpanType.AGENT
|
|
|
|
# LLM span should have conversation context as input in Anthropic format
|
|
llm_spans = [s for s in spans if s.span_type == SpanType.LLM]
|
|
assert len(llm_spans) == 1
|
|
assert llm_spans[0].name == "llm"
|
|
assert llm_spans[0].inputs["model"] == "claude-sonnet-4-20250514"
|
|
assert llm_spans[0].inputs["messages"] == [{"role": "user", "content": "What is 2 + 2?"}]
|
|
assert llm_spans[0].get_attribute(SpanAttributeKey.MESSAGE_FORMAT) == "anthropic"
|
|
|
|
# Output should be in Anthropic response format
|
|
outputs = llm_spans[0].outputs
|
|
assert outputs["type"] == "message"
|
|
assert outputs["role"] == "assistant"
|
|
assert outputs["content"] == [{"type": "text", "text": "The answer is 4."}]
|
|
|
|
# Token usage from ResultMessage should be on the root span and trace level
|
|
token_usage = root_span.get_attribute(SpanAttributeKey.CHAT_USAGE)
|
|
assert token_usage is not None
|
|
assert token_usage["input_tokens"] == 100
|
|
assert token_usage["output_tokens"] == 20
|
|
assert token_usage["total_tokens"] == 120
|
|
|
|
assert trace.info.token_usage is not None
|
|
assert trace.info.token_usage["input_tokens"] == 100
|
|
assert trace.info.token_usage["output_tokens"] == 20
|
|
assert trace.info.token_usage["total_tokens"] == 120
|
|
|
|
# Duration should reflect ResultMessage.duration_ms (1000ms = 1s)
|
|
duration_ns = root_span.end_time_ns - root_span.start_time_ns
|
|
assert abs(duration_ns - 1_000_000_000) < 1_000_000 # within 1ms tolerance
|
|
|
|
assert trace.info.trace_metadata.get("mlflow.trace.session") == "test-sdk-session"
|
|
assert trace.info.request_preview == "What is 2 + 2?"
|
|
assert trace.info.response_preview == "The answer is 4."
|
|
|
|
|
|
def test_process_sdk_messages_multiple_tools():
|
|
messages = [
|
|
UserMessage(content="Read two files"),
|
|
AssistantMessage(
|
|
content=[
|
|
ToolUseBlock(id="tool_1", name="Read", input={"path": "a.py"}),
|
|
ToolUseBlock(id="tool_2", name="Read", input={"path": "b.py"}),
|
|
],
|
|
model="claude-sonnet-4-20250514",
|
|
),
|
|
UserMessage(
|
|
content=[
|
|
ToolResultBlock(tool_use_id="tool_1", content="content of a"),
|
|
ToolResultBlock(tool_use_id="tool_2", content="content of b"),
|
|
],
|
|
tool_use_result={"tool_use_id": "tool_1"},
|
|
),
|
|
AssistantMessage(
|
|
content=[TextBlock(text="Here are the contents.")],
|
|
model="claude-sonnet-4-20250514",
|
|
),
|
|
ResultMessage(
|
|
subtype="success",
|
|
duration_ms=2000,
|
|
duration_api_ms=1500,
|
|
is_error=False,
|
|
num_turns=2,
|
|
session_id="multi-tool-session",
|
|
),
|
|
]
|
|
|
|
trace = process_sdk_messages(messages, "multi-tool-session")
|
|
|
|
assert trace is not None
|
|
spans = list(trace.search_spans())
|
|
|
|
tool_spans = [s for s in spans if s.span_type == SpanType.TOOL]
|
|
assert len(tool_spans) == 2
|
|
assert all(s.name == "tool_Read" for s in tool_spans)
|
|
tool_results = {s.outputs["result"] for s in tool_spans}
|
|
assert tool_results == {"content of a", "content of b"}
|
|
|
|
|
|
def test_process_sdk_messages_cache_tokens():
|
|
messages = [
|
|
UserMessage(content="Hello"),
|
|
AssistantMessage(
|
|
content=[TextBlock(text="Hi!")],
|
|
model="claude-sonnet-4-20250514",
|
|
),
|
|
ResultMessage(
|
|
subtype="success",
|
|
duration_ms=5000,
|
|
duration_api_ms=4000,
|
|
is_error=False,
|
|
num_turns=1,
|
|
session_id="cache-session",
|
|
usage={
|
|
"input_tokens": 36,
|
|
"cache_creation_input_tokens": 23554,
|
|
"cache_read_input_tokens": 139035,
|
|
"output_tokens": 3344,
|
|
},
|
|
),
|
|
]
|
|
|
|
trace = process_sdk_messages(messages, "cache-session")
|
|
|
|
assert trace is not None
|
|
root_span = trace.data.spans[0]
|
|
|
|
# input_tokens is the non-cached input the Anthropic API reports, matching
|
|
# mlflow.anthropic.autolog. Cache fields are exposed as separate keys so
|
|
# consumers can compute cache hit rate without scraping transcripts.
|
|
token_usage = root_span.get_attribute(SpanAttributeKey.CHAT_USAGE)
|
|
assert token_usage["input_tokens"] == 36
|
|
assert token_usage["output_tokens"] == 3344
|
|
assert token_usage["total_tokens"] == 36 + 3344
|
|
assert token_usage["cache_read_input_tokens"] == 139035
|
|
assert token_usage["cache_creation_input_tokens"] == 23554
|
|
|
|
# Trace-level aggregation should match
|
|
assert trace.info.token_usage["input_tokens"] == 36
|
|
assert trace.info.token_usage["output_tokens"] == 3344
|
|
|
|
|
|
# ============================================================================
|
|
# FIND LAST USER MESSAGE INDEX TESTS
|
|
# ============================================================================
|
|
|
|
|
|
def test_find_last_user_message_skips_skill_injection():
|
|
transcript = [
|
|
{"type": "queue-operation"},
|
|
{"type": "queue-operation"},
|
|
# Entry 2: actual user prompt
|
|
{
|
|
"type": "user",
|
|
"message": {"role": "user", "content": "Enable tracing on the agent."},
|
|
"timestamp": "2025-01-01T00:00:00Z",
|
|
},
|
|
# Entry 3: assistant thinking
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "thinking", "thinking": "Let me use the skill."}],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:01Z",
|
|
},
|
|
# Entry 4: assistant invokes Skill tool
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "tool_use",
|
|
"id": "toolu_abc123",
|
|
"name": "Skill",
|
|
"input": {"skill": "instrumenting-with-mlflow-tracing"},
|
|
}
|
|
],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:02Z",
|
|
},
|
|
# Entry 5: tool result with commandName (correctly skipped by toolUseResult check)
|
|
{
|
|
"type": "user",
|
|
"toolUseResult": {
|
|
"success": True,
|
|
"commandName": "instrumenting-with-mlflow-tracing",
|
|
},
|
|
"message": {
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": "toolu_abc123",
|
|
"content": "Launching skill: instrumenting-with-mlflow-tracing",
|
|
}
|
|
],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:03Z",
|
|
},
|
|
# Entry 6: skill content injection (BUG: not flagged as tool result)
|
|
{
|
|
"type": "user",
|
|
"message": {
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "text",
|
|
"text": (
|
|
"Base directory for this skill: /path/to/skill\n\n"
|
|
"# MLflow Tracing Guide\n\n...(full skill content)..."
|
|
),
|
|
}
|
|
],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:04Z",
|
|
},
|
|
# Entry 7: assistant continues
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "thinking", "thinking": "Now let me implement tracing."}],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:05Z",
|
|
},
|
|
# Entry 8: assistant text response
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "I've enabled tracing on the agent."}],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:06Z",
|
|
},
|
|
]
|
|
|
|
idx = find_last_user_message_index(transcript)
|
|
|
|
# Should return index 2 (actual user prompt), not 6 (skill injection)
|
|
assert idx == 2
|
|
assert transcript[idx]["message"]["content"] == "Enable tracing on the agent."
|
|
|
|
|
|
def test_find_last_user_message_index_basic():
|
|
transcript = [
|
|
{"type": "queue-operation"},
|
|
{
|
|
"type": "user",
|
|
"message": {"role": "user", "content": "First question"},
|
|
"timestamp": "2025-01-01T00:00:00Z",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "First answer"}],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:01Z",
|
|
},
|
|
{
|
|
"type": "user",
|
|
"message": {"role": "user", "content": "Second question"},
|
|
"timestamp": "2025-01-01T00:00:02Z",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "Second answer"}],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:03Z",
|
|
},
|
|
]
|
|
|
|
idx = find_last_user_message_index(transcript)
|
|
|
|
assert idx == 3
|
|
assert transcript[idx]["message"]["content"] == "Second question"
|
|
|
|
|
|
def test_find_last_user_message_skips_consecutive_skill_injections():
|
|
transcript = [
|
|
# Entry 0: actual user prompt
|
|
{
|
|
"type": "user",
|
|
"message": {"role": "user", "content": "Do the thing."},
|
|
"timestamp": "2025-01-01T00:00:00Z",
|
|
},
|
|
# Entry 1: assistant invokes first Skill
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "tool_use",
|
|
"id": "toolu_1",
|
|
"name": "Skill",
|
|
"input": {"skill": "skill-one"},
|
|
}
|
|
],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:01Z",
|
|
},
|
|
# Entry 2: first skill tool result
|
|
{
|
|
"type": "user",
|
|
"toolUseResult": {"success": True, "commandName": "skill-one"},
|
|
"message": {
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": "toolu_1",
|
|
"content": "Launching skill: skill-one",
|
|
}
|
|
],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:02Z",
|
|
},
|
|
# Entry 3: first skill content injection
|
|
{
|
|
"type": "user",
|
|
"message": {
|
|
"role": "user",
|
|
"content": [{"type": "text", "text": "Base directory: /skill-one\n# Skill One"}],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:03Z",
|
|
},
|
|
# Entry 4: assistant invokes second Skill
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [
|
|
{
|
|
"type": "tool_use",
|
|
"id": "toolu_2",
|
|
"name": "Skill",
|
|
"input": {"skill": "skill-two"},
|
|
}
|
|
],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:04Z",
|
|
},
|
|
# Entry 5: second skill tool result
|
|
{
|
|
"type": "user",
|
|
"toolUseResult": {"success": True, "commandName": "skill-two"},
|
|
"message": {
|
|
"role": "user",
|
|
"content": [
|
|
{
|
|
"type": "tool_result",
|
|
"tool_use_id": "toolu_2",
|
|
"content": "Launching skill: skill-two",
|
|
}
|
|
],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:05Z",
|
|
},
|
|
# Entry 6: second skill content injection
|
|
{
|
|
"type": "user",
|
|
"message": {
|
|
"role": "user",
|
|
"content": [{"type": "text", "text": "Base directory: /skill-two\n# Skill Two"}],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:06Z",
|
|
},
|
|
# Entry 7: assistant response
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "Done."}],
|
|
},
|
|
"timestamp": "2025-01-01T00:00:07Z",
|
|
},
|
|
]
|
|
|
|
idx = find_last_user_message_index(transcript)
|
|
|
|
# Should skip both skill injections (entries 3 and 6) and return entry 0
|
|
assert idx == 0
|
|
assert transcript[idx]["message"]["content"] == "Do the thing."
|
|
|
|
|
|
def test_process_transcript_captures_claude_code_version(tmp_path):
|
|
transcript = [
|
|
{
|
|
"type": "queue-operation",
|
|
"operation": "dequeue",
|
|
"timestamp": "2025-01-15T09:59:59.000Z",
|
|
"sessionId": "test-version-session",
|
|
},
|
|
{
|
|
"type": "user",
|
|
"version": "2.1.34",
|
|
"message": {"role": "user", "content": "Hello!"},
|
|
"timestamp": "2025-01-15T10:00:00.000Z",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"version": "2.1.34",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "Hi there!"}],
|
|
},
|
|
"timestamp": "2025-01-15T10:00:01.000Z",
|
|
},
|
|
]
|
|
|
|
transcript_path = tmp_path / "version_transcript.jsonl"
|
|
transcript_path.write_text("\n".join(json.dumps(entry) for entry in transcript) + "\n")
|
|
trace = process_transcript(str(transcript_path), "test-version-session")
|
|
|
|
assert trace is not None
|
|
assert trace.info.trace_metadata.get(METADATA_KEY_CLAUDE_CODE_VERSION) == "2.1.34"
|
|
|
|
|
|
def test_process_transcript_no_version_field(mock_transcript_file):
|
|
trace = process_transcript(mock_transcript_file, "test-session-no-version")
|
|
|
|
assert trace is not None
|
|
assert METADATA_KEY_CLAUDE_CODE_VERSION not in trace.info.trace_metadata
|
|
|
|
|
|
def test_process_transcript_includes_steer_messages(tmp_path):
|
|
transcript = [
|
|
{
|
|
"type": "user",
|
|
"message": {"role": "user", "content": "Tell me about Python."},
|
|
"timestamp": "2025-01-15T10:00:00.000Z",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "Python is a programming language."}],
|
|
},
|
|
"timestamp": "2025-01-15T10:00:01.000Z",
|
|
},
|
|
{
|
|
"type": "queue-operation",
|
|
"operation": "enqueue",
|
|
"content": "also tell me about Java",
|
|
"timestamp": "2025-01-15T10:00:02.000Z",
|
|
"sessionId": "test-steer-session",
|
|
},
|
|
{
|
|
"type": "queue-operation",
|
|
"operation": "remove",
|
|
"timestamp": "2025-01-15T10:00:03.000Z",
|
|
"sessionId": "test-steer-session",
|
|
},
|
|
{
|
|
"type": "assistant",
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": [{"type": "text", "text": "Java is also a programming language."}],
|
|
},
|
|
"timestamp": "2025-01-15T10:00:04.000Z",
|
|
},
|
|
]
|
|
|
|
transcript_path = tmp_path / "steer_transcript.jsonl"
|
|
transcript_path.write_text("\n".join(json.dumps(entry) for entry in transcript) + "\n")
|
|
trace = process_transcript(str(transcript_path), "test-steer-session")
|
|
assert trace is not None
|
|
|
|
spans = list(trace.search_spans())
|
|
llm_spans = [s for s in spans if s.span_type == SpanType.LLM]
|
|
assert len(llm_spans) == 2
|
|
|
|
# The second LLM span should include the steer message in its inputs
|
|
second_llm = llm_spans[1]
|
|
input_messages = second_llm.inputs["messages"]
|
|
steer_messages = [m for m in input_messages if m.get("content") == "also tell me about Java"]
|
|
assert len(steer_messages) == 1
|
|
assert steer_messages[0]["role"] == "user"
|