diff --git a/crates/cli/tests/coverage/server_tests.rs b/crates/cli/tests/coverage/server_tests.rs index 677a38e79..8cd63de0f 100644 --- a/crates/cli/tests/coverage/server_tests.rs +++ b/crates/cli/tests/coverage/server_tests.rs @@ -521,6 +521,298 @@ async fn serve_listener_hermes_api_hooks_write_atof_category_profile_and_fidelit assert_eq!(lossy_start["data"]["content"]["message_count"], json!(2)); } +#[tokio::test] +async fn serve_listener_routed_gateway_wire_formats_write_atof_category_profile_and_usage() { + let _guard = PLUGIN_TEST_LOCK.lock().await; + let _ = nemo_relay::plugin::clear_plugin_configuration(); + + async fn anthropic_messages() -> TestServer { + async fn messages(_headers: HeaderMap, _request: Request) -> impl IntoResponse { + Json(json!({ + "id": "msg_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4", + "content": [ + {"type": "text", "text": "I will search."}, + {"type": "tool_use", "id": "toolu_01", "name": "search", "input": {"query": "file"}} + ], + "stop_reason": "tool_use", + "usage": { + "input_tokens": 11, + "output_tokens": 7, + "cache_read_input_tokens": 3, + "cost": {"total": 0.0042} + } + })) + } + + let app = Router::new().route("/v1/messages", post(messages)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + TestServer { + url: format!("http://{address}"), + handle, + } + } + + async fn openai_routed() -> TestServer { + async fn chat(_headers: HeaderMap, request: Request) -> impl IntoResponse { + let path = request.uri().path().to_string(); + if path == "/v1/responses" { + Json(json!({ + "id": "resp_1", + "status": "completed", + "output": [ + { + "type": "message", + "content": [{"type": "output_text", "text": "I will check the weather."}] + }, + { + "type": "function_call", + "call_id": "call_weather_1", + "name": "get_weather", + "arguments": "{\"city\":\"SF\"}", + "status": "completed" + } + ], + "usage": { + "input_tokens": 75, + "output_tokens": 20, + "total_tokens": 95, + "input_tokens_details": {"cached_tokens": 10}, + "cost_usd": 0.005 + } + })) + } else { + Json(json!({ + "choices": [{ + "message": { + "role": "assistant", + "content": "I will inspect.", + "tool_calls": [ + { + "id": "call_read_1", + "type": "function", + "function": {"name": "read", "arguments": "{\"path\":\"api.py\"}"} + } + ] + }, + "finish_reason": "tool_calls" + }], + "usage": { + "prompt_tokens": 3, + "completion_tokens": 4, + "total_tokens": 7, + "prompt_tokens_details": {"cached_tokens": 2}, + "cost_usd": 0.001 + } + })) + } + } + + let app = Router::new() + .route("/v1/chat/completions", post(chat)) + .route("/v1/responses", post(chat)); + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + TestServer { + url: format!("http://{address}"), + handle, + } + } + + let temp = tempfile::tempdir().unwrap(); + let atof_dir = temp.path().join("atof"); + std::fs::create_dir_all(&atof_dir).unwrap(); + + let anthropic_upstream = anthropic_messages().await; + let openai_upstream = openai_routed().await; + + let mut config = test_config(); + config.anthropic_base_url = anthropic_upstream.url(); + config.openai_base_url = openai_upstream.url(); + config.plugin_config = Some(json!({ + "version": 1, + "components": [ + { + "kind": "observability", + "enabled": true, + "config": { + "version": 1, + "atof": { + "enabled": true, + "output_directory": atof_dir, + "filename": "events.jsonl", + "mode": "overwrite" + } + } + } + ] + })); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let address = listener.local_addr().unwrap(); + let url = format!("http://{address}"); + let (shutdown_tx, shutdown_rx) = oneshot::channel(); + let handle = + tokio::spawn(async move { serve_listener(listener, config, Some(shutdown_rx)).await }); + + wait_for_gateway(&url).await; + let client = test_http_client(); + + let response = client + .post(format!("{url}/v1/messages")) + .header("content-type", "application/json") + .header("x-api-key", "sk-ant-test") + .header("x-nemo-relay-session-id", "hermes-routed-atof") + .json(&json!({ + "model": "claude-sonnet-4", + "messages": [{"role": "user", "content": "Find the file."}], + "tools": [{"name": "search", "input_schema": {"type": "object"}}] + })) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let response = client + .post(format!("{url}/v1/responses")) + .header("content-type", "application/json") + .header("authorization", "Bearer test") + .header("x-nemo-relay-session-id", "hermes-routed-atof") + .json(&json!({ + "model": "gpt-4o", + "input": "Find the weather.", + "tools": [{"type": "function", "name": "get_weather"}] + })) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let response = client + .post(format!("{url}/v1/chat/completions")) + .header("content-type", "application/json") + .header("authorization", "Bearer test") + .header("x-nemo-relay-session-id", "hermes-routed-atof") + .json(&json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Inspect the files."}], + "tools": [{"type": "function", "function": {"name": "read"}}] + })) + .send() + .await + .unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + shutdown_tx.send(()).unwrap(); + handle.await.unwrap().unwrap(); + + let events = std::fs::read_to_string(temp.path().join("atof/events.jsonl")).unwrap(); + let llm_events = events + .lines() + .map(|line| serde_json::from_str::(line).unwrap()) + .filter(|event| event["category"] == "llm") + .collect::>(); + assert_eq!( + llm_events.len(), + 6, + "expected three routed LLM start/end pairs, got {llm_events:?}" + ); + + let anthropic_start = llm_events + .iter() + .find(|event| { + event["scope_category"] == "start" + && event["name"] == "anthropic.messages" + && event["metadata"]["gateway_path"] == "/v1/messages" + }) + .unwrap(); + assert_eq!( + anthropic_start["category_profile"]["model_name"], + json!("claude-sonnet-4") + ); + assert_eq!( + anthropic_start["data"]["content"]["messages"][0]["content"], + json!("Find the file.") + ); + + let anthropic_end = llm_events + .iter() + .find(|event| { + event["scope_category"] == "end" + && event["name"] == "anthropic.messages" + && event["metadata"]["gateway_path"] == "/v1/messages" + }) + .unwrap(); + assert_eq!( + anthropic_end["category_profile"]["annotated_response"]["tool_calls"][0]["id"], + json!("toolu_01") + ); + assert_eq!(anthropic_end["data"]["content"][1]["id"], json!("toolu_01")); + assert_eq!(anthropic_end["data"]["usage"]["input_tokens"], json!(11)); + assert_eq!( + anthropic_end["data"]["usage"]["cost"]["total"], + json!(0.0042) + ); + + let responses_end = llm_events + .iter() + .find(|event| { + event["scope_category"] == "end" + && event["name"] == "openai.responses" + && event["metadata"]["gateway_path"] == "/v1/responses" + }) + .unwrap(); + assert_eq!( + responses_end["category_profile"]["model_name"], + json!("gpt-4o") + ); + assert_eq!( + responses_end["category_profile"]["annotated_response"]["tool_calls"][0]["id"], + json!("call_weather_1") + ); + assert_eq!( + responses_end["data"]["output"][1]["call_id"], + json!("call_weather_1") + ); + assert_eq!( + responses_end["data"]["usage"]["input_tokens_details"]["cached_tokens"], + json!(10) + ); + assert_eq!(responses_end["data"]["usage"]["cost_usd"], json!(0.005)); + + let chat_end = llm_events + .iter() + .find(|event| { + event["scope_category"] == "end" + && event["name"] == "openai.chat_completions" + && event["metadata"]["gateway_path"] == "/v1/chat/completions" + }) + .unwrap(); + assert_eq!(chat_end["category_profile"]["model_name"], json!("gpt-4o")); + assert_eq!( + chat_end["category_profile"]["annotated_response"]["tool_calls"][0]["id"], + json!("call_read_1") + ); + assert_eq!( + chat_end["data"]["choices"][0]["message"]["tool_calls"][0]["id"], + json!("call_read_1") + ); + assert_eq!( + chat_end["data"]["usage"]["prompt_tokens_details"]["cached_tokens"], + json!(2) + ); + assert_eq!(chat_end["data"]["usage"]["cost_usd"], json!(0.001)); +} + #[tokio::test] async fn serve_listener_activates_any_registered_plugin_kind() { let _guard = PLUGIN_TEST_LOCK.lock().await; diff --git a/crates/cli/tests/coverage/session_tests.rs b/crates/cli/tests/coverage/session_tests.rs index 07c76214d..a4824a1da 100644 --- a/crates/cli/tests/coverage/session_tests.rs +++ b/crates/cli/tests/coverage/session_tests.rs @@ -2058,6 +2058,224 @@ async fn hermes_subagent_child_session_embeds_non_empty_atif_trajectory() { ); } +#[tokio::test] +async fn hermes_routed_provider_payloads_write_exact_atif_trajectory() { + let _guard = OBSERVABILITY_PLUGIN_TEST_LOCK.lock().await; + let temp = tempfile::tempdir().unwrap(); + let atif_dir = temp.path().join("atif"); + install_test_atif_plugin(&atif_dir).await; + let manager = SessionManager::new(session_test_config()); + let headers = HeaderMap::new(); + + manager + .apply_events( + &headers, + vec![NormalizedEvent::AgentStarted(SessionEvent { + session_id: "hermes-routed".into(), + agent_kind: AgentKind::Hermes, + event_name: "on_session_start".into(), + payload: json!({}), + metadata: json!({}), + })], + ) + .await + .unwrap(); + + let anthropic = manager + .start_llm( + &headers, + LlmGatewayStart { + session_id: Some("hermes-routed".into()), + provider: "anthropic.messages".into(), + model_name: Some("claude-sonnet-4".into()), + subagent_id: None, + conversation_id: None, + generation_id: None, + request_id: Some("msg-request".into()), + request: LlmRequest { + headers: Map::new(), + content: json!({ + "model": "claude-sonnet-4", + "messages": [{"role": "user", "content": "Find the file."}], + "tools": [{"name": "search", "input_schema": {"type": "object"}}] + }), + }, + streaming: false, + metadata: json!({ "gateway_path": "/v1/messages" }), + }, + ) + .await + .unwrap(); + manager + .end_llm( + anthropic, + json!({ + "id": "msg_01", + "type": "message", + "content": [ + {"type": "text", "text": "I will search."}, + {"type": "tool_use", "id": "toolu_01", "name": "search", "input": {"query": "file"}} + ], + "usage": { + "input_tokens": 11, + "output_tokens": 7, + "cache_read_input_tokens": 3, + "cost": {"total": 0.0042} + } + }), + json!({}), + ) + .await + .unwrap(); + + let responses = manager + .start_llm( + &headers, + LlmGatewayStart { + session_id: Some("hermes-routed".into()), + provider: "openai.responses".into(), + model_name: Some("gpt-4o".into()), + subagent_id: None, + conversation_id: None, + generation_id: None, + request_id: Some("resp-request".into()), + request: LlmRequest { + headers: Map::new(), + content: json!({ + "model": "gpt-4o", + "input": "Find the weather.", + "tools": [{"type": "function", "name": "get_weather"}] + }), + }, + streaming: false, + metadata: json!({ "gateway_path": "/v1/responses" }), + }, + ) + .await + .unwrap(); + manager + .end_llm( + responses, + json!({ + "id": "resp_1", + "output": [ + {"type": "message", "content": [{"type": "output_text", "text": "I will check the weather."}]}, + {"type": "function_call", "call_id": "call_weather_1", "name": "get_weather", "arguments": "{\"city\":\"SF\"}"} + ], + "usage": { + "input_tokens": 75, + "output_tokens": 20, + "total_tokens": 95, + "input_tokens_details": {"cached_tokens": 10}, + "cost_usd": 0.005 + } + }), + json!({}), + ) + .await + .unwrap(); + + let chat = manager + .start_llm( + &headers, + LlmGatewayStart { + session_id: Some("hermes-routed".into()), + provider: "openai.chat_completions".into(), + model_name: Some("gpt-4o".into()), + subagent_id: None, + conversation_id: None, + generation_id: None, + request_id: Some("chat-request".into()), + request: LlmRequest { + headers: Map::new(), + content: json!({ + "model": "gpt-4o", + "messages": [{"role": "user", "content": "Inspect the files."}], + "tools": [{"type": "function", "function": {"name": "read"}}] + }), + }, + streaming: false, + metadata: json!({ "gateway_path": "/v1/chat/completions" }), + }, + ) + .await + .unwrap(); + manager + .end_llm( + chat, + json!({ + "choices": [{ + "message": { + "role": "assistant", + "content": "I will inspect.", + "tool_calls": [{"id": "call_read_1", "function": {"name": "read", "arguments": "{\"path\":\"api.py\"}"}}] + } + }], + "usage": { + "prompt_tokens": 3, + "completion_tokens": 4, + "total_tokens": 7, + "prompt_tokens_details": {"cached_tokens": 2}, + "cost_usd": 0.001 + } + }), + json!({}), + ) + .await + .unwrap(); + + manager + .apply_events( + &headers, + vec![NormalizedEvent::AgentEnded(SessionEvent { + session_id: "hermes-routed".into(), + agent_kind: AgentKind::Hermes, + event_name: "on_session_finalize".into(), + payload: json!({}), + metadata: json!({}), + })], + ) + .await + .unwrap(); + + clear_plugin_configuration().unwrap(); + let atif = read_atif_for_session(&atif_dir, "hermes-routed"); + let steps = atif["steps"].as_array().unwrap(); + assert_eq!(steps.len(), 6); + + assert_eq!(steps[0]["message"], json!("Find the file.")); + assert_eq!(steps[1]["message"], json!("I will search.")); + assert_eq!(steps[1]["tool_calls"][0]["tool_call_id"], json!("toolu_01")); + assert_eq!(steps[1]["metrics"]["prompt_tokens"], json!(11)); + assert_eq!(steps[1]["metrics"]["cached_tokens"], json!(3)); + assert_eq!(steps[1]["metrics"]["cost_usd"], json!(0.0042)); + + assert_eq!(steps[2]["message"], json!("Find the weather.")); + assert_eq!(steps[3]["message"], json!("I will check the weather.")); + assert_eq!( + steps[3]["tool_calls"][0]["tool_call_id"], + json!("call_weather_1") + ); + assert_eq!(steps[3]["metrics"]["prompt_tokens"], json!(75)); + assert_eq!(steps[3]["metrics"]["cached_tokens"], json!(10)); + assert_eq!(steps[3]["metrics"]["cost_usd"], json!(0.005)); + + assert_eq!(steps[4]["message"], json!("Inspect the files.")); + assert_eq!(steps[5]["message"], json!("I will inspect.")); + assert_eq!( + steps[5]["tool_calls"][0]["tool_call_id"], + json!("call_read_1") + ); + assert_eq!(steps[5]["metrics"]["prompt_tokens"], json!(3)); + assert_eq!(steps[5]["metrics"]["cached_tokens"], json!(2)); + assert_eq!(steps[5]["metrics"]["cost_usd"], json!(0.001)); + + assert_eq!(atif["final_metrics"]["total_prompt_tokens"], json!(89)); + assert_eq!(atif["final_metrics"]["total_completion_tokens"], json!(31)); + assert_eq!(atif["final_metrics"]["total_cached_tokens"], json!(15)); + assert_eq!(atif["final_metrics"]["total_cost_usd"], json!(0.0102)); +} + #[tokio::test] async fn empty_hook_marks_do_not_create_empty_atif_steps() { let _guard = OBSERVABILITY_PLUGIN_TEST_LOCK.lock().await; diff --git a/crates/core/src/observability/atif.rs b/crates/core/src/observability/atif.rs index b761846ee..f73927eae 100644 --- a/crates/core/src/observability/atif.rs +++ b/crates/core/src/observability/atif.rs @@ -668,6 +668,7 @@ fn extract_metrics(output: &Json) -> Option { let completion = usage_u64(usage, &["completion_tokens", "output_tokens"]); let cached = usage_u64(usage, &["cached_tokens"]) .or_else(|| prompt_tokens_detail_u64(usage, "cached_tokens")) + .or_else(|| input_tokens_detail_u64(usage, "cached_tokens")) .or_else(|| { sum_usage_u64( usage, @@ -806,6 +807,14 @@ fn prompt_tokens_detail_u64(usage: &serde_json::Map, key: &str) -> .and_then(Json::as_u64) } +fn input_tokens_detail_u64(usage: &serde_json::Map, key: &str) -> Option { + usage + .get("input_tokens_details") + .and_then(Json::as_object) + .and_then(|details| details.get(key)) + .and_then(Json::as_u64) +} + /// Extract `reasoning_effort` from an LLM request (string or number). /// /// The request content may have `reasoning_effort` (e.g. `"high"`, `"medium"`, diff --git a/crates/core/tests/unit/atif_tests.rs b/crates/core/tests/unit/atif_tests.rs index 70b76c86a..02df626e9 100644 --- a/crates/core/tests/unit/atif_tests.rs +++ b/crates/core/tests/unit/atif_tests.rs @@ -699,6 +699,27 @@ fn test_extract_metrics_supports_provider_usage_payloads() { ); assert_eq!(openai_metrics.cost_usd, Some(0.001)); + let responses_metrics = extract_metrics(&json!({ + "usage": { + "input_tokens": 75, + "output_tokens": 20, + "total_tokens": 95, + "input_tokens_details": { + "cached_tokens": 10 + }, + "cost_usd": 0.005 + } + })) + .unwrap(); + assert_eq!(responses_metrics.prompt_tokens, Some(75)); + assert_eq!(responses_metrics.completion_tokens, Some(20)); + assert_eq!(responses_metrics.cached_tokens, Some(10)); + assert_eq!( + responses_metrics.extra.as_ref().unwrap()["total_tokens"], + json!(95) + ); + assert_eq!(responses_metrics.cost_usd, Some(0.005)); + let anthropic_metrics = extract_metrics(&json!({ "usage": { "input_tokens": 11,