From 785019fe3c9de591f977302890f8b01ba838d917 Mon Sep 17 00:00:00 2001 From: hera Date: Wed, 22 Jul 2026 09:34:16 +0000 Subject: [PATCH] hera: Hermetic unit tests for h-storage-clickhouse --- server/h-storage-clickhouse/src/calls.rs | 422 ++++++++++++++++-- server/h-storage-clickhouse/src/distincts.rs | 142 ++++-- server/h-storage-clickhouse/src/exchanges.rs | 302 +++++++++++-- server/h-storage-clickhouse/src/metrics.rs | 215 ++++++++- server/h-storage-clickhouse/src/retention.rs | 33 ++ server/h-storage-clickhouse/src/rows.rs | 434 +++++++++++++++++++ server/h-storage-clickhouse/src/services.rs | 241 +++++++--- server/h-storage-clickhouse/src/sql.rs | 89 ++++ server/h-storage-clickhouse/src/turns.rs | 346 ++++++++++++--- 9 files changed, 1964 insertions(+), 260 deletions(-) diff --git a/server/h-storage-clickhouse/src/calls.rs b/server/h-storage-clickhouse/src/calls.rs index 3a9c2bbf..4423e26d 100644 --- a/server/h-storage-clickhouse/src/calls.rs +++ b/server/h-storage-clickhouse/src/calls.rs @@ -136,53 +136,9 @@ impl ClickHouseBackend { query.sort_by ))); } - let sort_order = if query.sort_order.eq_ignore_ascii_case("ASC") { - "ASC" - } else { - "DESC" - }; + let sort_order = resolve_sort_order(&query.sort_order); - // Time-range + filter WHERE. Timestamps and numeric ports are - // interpolated (values we control); user string lists go through - // `sql_in_list`'s single-quote escaping, valid in ClickHouse. - let mut where_parts = - vec![time_where("request_time", query.time_range.start_us, query.time_range.end_us)]; - if !query.filter.wire_apis.is_empty() { - where_parts.push(format!("wire_api IN ({})", sql_in_list(&query.filter.wire_apis))); - } - if !query.filter.models.is_empty() { - where_parts.push(format!("model IN ({})", sql_in_list(&query.filter.models))); - } - if !query.filter.server_ips.is_empty() { - where_parts.push(format!("server_ip IN ({})", sql_in_list(&query.filter.server_ips))); - } - if !query.status_codes.is_empty() { - where_parts.push(format!("status_code IN ({})", join_nums(&query.status_codes))); - } - if !query.finish_reasons.is_empty() { - where_parts.push(format!( - "finish_reason IN ({})", - sql_in_list(&query.finish_reasons) - )); - } - if !query.client_ips.is_empty() { - where_parts.push(format!("client_ip IN ({})", sql_in_list(&query.client_ips))); - } - if !query.server_ports.is_empty() { - where_parts.push(format!("server_port IN ({})", join_nums(&query.server_ports))); - } - if let Some(substr) = query - .request_path_contains - .as_deref() - .map(str::trim) - .filter(|s| !s.is_empty()) - { - where_parts.push(format!("request_path LIKE '%{}%'", escape_str(substr))); - } - if let Some(stream) = query.is_stream { - where_parts.push(format!("is_stream = {}", if stream { 1 } else { 0 })); - } - let where_sql = where_parts.join(" AND "); + let where_sql = spans_where_sql(query); let total = self .client @@ -376,6 +332,63 @@ fn join_nums(values: &[T]) -> String { .join(", ") } +/// Normalize a client-supplied sort direction to the SQL literal. Anything +/// other than an ASCII-case-insensitive `ASC` (including the empty / missing +/// value the API defaults to `"desc"` for) collapses to `DESC`, so a malformed +/// `sort_order` can never inject a non-keyword into the ORDER BY clause. +pub(crate) fn resolve_sort_order(sort_order: &str) -> &'static str { + if sort_order.eq_ignore_ascii_case("ASC") { + "ASC" + } else { + "DESC" + } +} + +/// Build the `query_spans` WHERE clause: a half-open `request_time` time range +/// AND-ed with every present dimension filter. Extracted as a pure fn so the +/// escaping / IN-list / LIKE assembly is unit-testable without a live server. +/// Timestamps and numeric ports are interpolated (values we control); user +/// string lists go through `sql_in_list`'s backslash-aware single-quote +/// escaping; the `LIKE` substring goes through `escape_str` (which deliberately +/// leaves `%` / `_` so substring semantics survive). +pub(crate) fn spans_where_sql(query: &SpansQuery) -> String { + let mut where_parts = + vec![time_where("request_time", query.time_range.start_us, query.time_range.end_us)]; + if !query.filter.wire_apis.is_empty() { + where_parts.push(format!("wire_api IN ({})", sql_in_list(&query.filter.wire_apis))); + } + if !query.filter.models.is_empty() { + where_parts.push(format!("model IN ({})", sql_in_list(&query.filter.models))); + } + if !query.filter.server_ips.is_empty() { + where_parts.push(format!("server_ip IN ({})", sql_in_list(&query.filter.server_ips))); + } + if !query.status_codes.is_empty() { + where_parts.push(format!("status_code IN ({})", join_nums(&query.status_codes))); + } + if !query.finish_reasons.is_empty() { + where_parts.push(format!("finish_reason IN ({})", sql_in_list(&query.finish_reasons))); + } + if !query.client_ips.is_empty() { + where_parts.push(format!("client_ip IN ({})", sql_in_list(&query.client_ips))); + } + if !query.server_ports.is_empty() { + where_parts.push(format!("server_port IN ({})", join_nums(&query.server_ports))); + } + if let Some(substr) = query + .request_path_contains + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + { + where_parts.push(format!("request_path LIKE '%{}%'", escape_str(substr))); + } + if let Some(stream) = query.is_stream { + where_parts.push(format!("is_stream = {}", if stream { 1 } else { 0 })); + } + where_parts.join(" AND ") +} + /// Build `Option` from the three nullable `process_*` columns. /// `None` exactly when `pid` is NULL (passive-tap rows). fn row_process(pid: Option, comm: Option, exe: Option) -> Option { @@ -459,3 +472,320 @@ fn call_detail(r: SpanDetailRow) -> SpanDetail { process, } } + +#[cfg(test)] +mod tests { + use super::*; + + fn list_row(response_body: Option) -> CallListRow { + CallListRow { + id: "call-1".into(), + source_id: "src-0".into(), + request_time_ms: 1_700_000_000_000, + wire_api: "openai-chat".into(), + model: "gpt-4".into(), + status_code: Some(200), + is_stream: true, + finish_reason: Some("stop".into()), + ttft_ms: Some(500.0), + e2e_latency_ms: Some(1000.0), + input_tokens: Some(100), + output_tokens: Some(50), + client_ip: "10.0.0.1".into(), + server_ip: "10.0.0.2".into(), + server_port: 8080, + request_path: "/v1/chat/completions".into(), + response_body, + is_agent_request: false, + tool_surface: None, + agent_topology: None, + tool_call_count: 0, + tool_names_json: Some(r#"["bash","grep"]"#.into()), + process_pid: None, + process_comm: None, + process_exe: None, + } + } + + #[test] + fn join_nums_renders_numeric_in_list() { + assert_eq!(join_nums(&[8080u16, 443u16]), "8080, 443"); + assert_eq!(join_nums(&[200u16]), "200"); + assert_eq!(join_nums::(&[]), ""); + } + + #[test] + fn row_process_none_when_pid_none() { + assert_eq!(row_process(None, Some("node".into()), Some("/x".into())), None); + // pid present → Some(ProcessInfo); comm defaults to "" when absent; exe pass-through. + let p = row_process(Some(7), None, None).unwrap(); + assert_eq!(p.pid, 7); + assert_eq!(p.comm, ""); + assert_eq!(p.exe, None); + let p = row_process(Some(7), Some("node".into()), Some("/usr/bin/node".into())).unwrap(); + assert_eq!(p.comm, "node"); + assert_eq!(p.exe.as_deref(), Some("/usr/bin/node")); + } + + #[test] + fn call_list_item_maps_scalars_and_ms_timestamp() { + let item = call_list_item(list_row(None)); + assert_eq!(item.id, "call-1"); + assert_eq!(item.source_id, "src-0"); + assert_eq!(item.request_time, 1_700_000_000_000); + assert_eq!(item.wire_api, "openai-chat"); + assert_eq!(item.model, "gpt-4"); + assert_eq!(item.status_code, Some(200)); + assert!(item.is_stream); + assert_eq!(item.finish_reason.as_deref(), Some("stop")); + assert_eq!(item.ttft_ms, Some(500.0)); + assert_eq!(item.e2e_latency_ms, Some(1000.0)); + assert_eq!(item.input_tokens, Some(100)); + assert_eq!(item.output_tokens, Some(50)); + assert_eq!(item.client_ip, "10.0.0.1"); + assert_eq!(item.server_ip, "10.0.0.2"); + assert_eq!(item.server_port, 8080); + assert_eq!(item.request_path, "/v1/chat/completions"); + assert!(!item.is_agent_request); + assert_eq!(item.tool_names, vec!["bash".to_string(), "grep".to_string()]); + assert_eq!(item.process, None); + } + + #[test] + fn call_list_item_tokens_estimated_follows_usage_block() { + // Body with a positive usage block → wire tokens → estimated = false. + let with_usage = list_row(Some( + r#"{"choices":[],"usage":{"prompt_tokens":100,"completion_tokens":50}}"#.into(), + )); + assert!(!call_list_item(with_usage).tokens_estimated); + + // Body without a usage block → estimated = true. + let no_usage = list_row(Some(r#"{"choices":[{"message":{"content":"hi"}}]}"#.into())); + assert!(call_list_item(no_usage).tokens_estimated); + + // No body at all → estimated = true. + assert!(call_list_item(list_row(None)).tokens_estimated); + } + + #[test] + fn call_list_item_zero_tokens_never_estimated() { + let mut r = list_row(None); + r.input_tokens = Some(0); + r.output_tokens = Some(0); + // Even with no body, zero tokens → not estimated. + assert!(!call_list_item(r).tokens_estimated); + } + + #[test] + fn call_list_item_malformed_tool_names_json_degrades_to_empty() { + let mut r = list_row(None); + r.tool_names_json = Some("not-json".into()); + assert_eq!(call_list_item(r).tool_names, Vec::::new()); + } + + fn detail_row(response_body: Option) -> SpanDetailRow { + SpanDetailRow { + id: "call-1".into(), + source_id: "src-0".into(), + request_time_ms: 1_700_000_000_000, + response_time_ms: Some(1_700_000_000_500), + complete_time_ms: Some(1_700_000_001_000), + wire_api: "openai-chat".into(), + model: "gpt-4".into(), + api_type: "chat".into(), + is_stream: true, + request_path: "/v1/chat/completions".into(), + status_code: Some(200), + finish_reason: Some("stop".into()), + input_tokens: Some(100), + output_tokens: Some(50), + total_tokens: Some(150), + ttft_ms: Some(500.0), + e2e_latency_ms: Some(1000.0), + response_id: Some("chatcmpl-x".into()), + client_ip: "10.0.0.1".into(), + client_port: 54321, + server_ip: "10.0.0.2".into(), + server_port: 8080, + request_body: Some(r#"{"model":"gpt-4"}"#.into()), + response_body, + request_headers: r#"[["content-type","application/json"]]"#.into(), + response_headers: r#"[["x-request-id","abc"]]"#.into(), + is_agent_request: false, + tool_surface: None, + agent_topology: None, + tool_call_count: 0, + tool_names_json: Some("[]".into()), + process_pid: Some(42), + process_comm: Some("node".into()), + process_exe: Some("/usr/bin/node".into()), + } + } + + #[test] + fn call_detail_wraps_headers_in_some() { + // SpanDetailRow.request_headers / response_headers are non-null String; + // SpanDetail wraps them in Some(...) to match the API shape. + let d = call_detail(detail_row(None)); + assert_eq!(d.id, "call-1"); + assert_eq!(d.response_time, Some(1_700_000_000_500)); + assert_eq!(d.complete_time, Some(1_700_000_001_000)); + assert_eq!(d.api_type, "chat"); + assert_eq!(d.total_tokens, Some(150)); + assert_eq!(d.response_id.as_deref(), Some("chatcmpl-x")); + assert_eq!(d.client_port, 54321); + assert_eq!(d.request_headers.as_deref(), Some(r#"[["content-type","application/json"]]"#)); + assert_eq!(d.response_headers.as_deref(), Some(r#"[["x-request-id","abc"]]"#)); + assert_eq!(d.request_body.as_deref(), Some(r#"{"model":"gpt-4"}"#)); + // process built from the three process_* columns. + assert_eq!( + d.process.as_ref().unwrap(), + &ProcessInfo { + pid: 42, + comm: "node".into(), + exe: Some("/usr/bin/node".into()), + } + ); + } + + #[test] + fn call_detail_tokens_estimated_follows_usage_block() { + let with_usage = detail_row(Some( + r#"{"choices":[],"usage":{"prompt_tokens":100,"completion_tokens":50}}"#.into(), + )); + assert!(!call_detail(with_usage).tokens_estimated); + let no_usage = detail_row(Some(r#"{"choices":[]}"#.into())); + assert!(call_detail(no_usage).tokens_estimated); + } + + #[test] + fn valid_sort_fields_is_whitelisted() { + // The `query_spans` sort_by is interpolated into ORDER BY, so an unknown + // field must be rejected up front — the whitelist is the gate. + for &known in VALID_SORT_FIELDS { + // every entry is a plain column name (no injection surface). + assert!(known.chars().all(|c| c.is_alphanumeric() || c == '_')); + } + assert!(VALID_SORT_FIELDS.contains(&"request_time")); + assert!(VALID_SORT_FIELDS.contains(&"ttft_ms")); + assert!(!VALID_SORT_FIELDS.contains(&"bogus")); + } + + fn spans_query() -> SpansQuery { + SpansQuery { + time_range: TimeRange { start_us: 100, end_us: 200 }, + filter: DimensionFilter::default(), + status_codes: vec![], + finish_reasons: vec![], + client_ips: vec![], + server_ports: vec![], + request_path_contains: None, + is_stream: None, + sort_by: "request_time".into(), + sort_order: "desc".into(), + page: 1, + page_size: 10, + } + } + + #[test] + fn resolve_sort_order_normalizes() { + assert_eq!(resolve_sort_order("asc"), "ASC"); + assert_eq!(resolve_sort_order("ASC"), "ASC"); + assert_eq!(resolve_sort_order("Asc"), "ASC"); + assert_eq!(resolve_sort_order("desc"), "DESC"); + assert_eq!(resolve_sort_order("DESC"), "DESC"); + // Anything non-ASC (including garbage / empty) collapses to DESC so a + // malformed value can never inject a non-keyword into ORDER BY. + assert_eq!(resolve_sort_order(""), "DESC"); + assert_eq!(resolve_sort_order("garbage; DROP"), "DESC"); + } + + #[test] + fn spans_where_sql_is_time_range_only_by_default() { + let s = spans_where_sql(&spans_query()); + assert_eq!( + s, + "request_time >= fromUnixTimestamp64Micro(100) \ + AND request_time < fromUnixTimestamp64Micro(200)" + ); + } + + #[test] + fn spans_where_sql_combines_all_filters() { + let q = SpansQuery { + filter: DimensionFilter { + wire_apis: vec!["openai-chat".into()], + models: vec!["gpt-4".into()], + server_ips: vec!["10.0.0.2".into()], + tool_surfaces: vec![], + }, + status_codes: vec![429], + finish_reasons: vec!["stop".into()], + client_ips: vec!["10.0.0.1".into()], + server_ports: vec![8080], + request_path_contains: Some("chat/completions".into()), + is_stream: Some(true), + ..spans_query() + }; + let s = spans_where_sql(&q); + // Every filter appends an AND'd predicate in declaration order. + assert!(s.contains("wire_api IN ('openai-chat')")); + assert!(s.contains("model IN ('gpt-4')")); + assert!(s.contains("server_ip IN ('10.0.0.2')")); + assert!(s.contains("status_code IN (429)")); + assert!(s.contains("finish_reason IN ('stop')")); + assert!(s.contains("client_ip IN ('10.0.0.1')")); + assert!(s.contains("server_port IN (8080)")); + assert!(s.contains("request_path LIKE '%chat/completions%'")); + assert!(s.contains("is_stream = 1")); + // Joined by AND, no trailing/leading separator. + assert!(!s.starts_with(" AND")); + assert!(!s.ends_with("AND ")); + } + + #[test] + fn spans_where_sql_like_escapes_quotes_not_wildcards() { + let q = SpansQuery { + request_path_contains: Some("a'b".into()), + ..spans_query() + }; + let s = spans_where_sql(&q); + // The embedded quote is doubled (no breakout); the % wrappers come from + // the format string, not the value. + assert!(s.contains("request_path LIKE '%a''b%'")); + } + + #[test] + fn spans_where_sql_trims_and_ignores_empty_like_substring() { + let q = SpansQuery { + request_path_contains: Some(" ".into()), + ..spans_query() + }; + // A whitespace-only substring is trimmed to empty → no LIKE predicate. + assert!(!spans_where_sql(&q).contains("LIKE")); + } + + #[test] + fn spans_where_sql_is_stream_false() { + let q = SpansQuery { + is_stream: Some(false), + ..spans_query() + }; + assert!(spans_where_sql(&q).contains("is_stream = 0")); + } + + #[test] + fn spans_where_sql_in_list_uses_backslash_escaping() { + let q = SpansQuery { + filter: DimensionFilter { + models: vec![r"gpt\4".into()], + ..Default::default() + }, + ..spans_query() + }; + // A backslash in a model value is doubled (ClickHouse-aware), keeping + // the literal closed. + assert!(spans_where_sql(&q).contains(r"model IN ('gpt\\4')")); + } +} diff --git a/server/h-storage-clickhouse/src/distincts.rs b/server/h-storage-clickhouse/src/distincts.rs index 65e9880f..71c9790b 100644 --- a/server/h-storage-clickhouse/src/distincts.rs +++ b/server/h-storage-clickhouse/src/distincts.rs @@ -79,49 +79,10 @@ impl ClickHouseBackend { ) -> Result> { // `traces` is a ReplacingMergeTree, so reads must use FINAL to see // the latest version per turn_id. - let mut where_parts = vec![time_where( - "start_time", - query.time_range.start_us, - query.time_range.end_us, - )]; - - if !query.filter.wire_apis.is_empty() { - where_parts.push(format!( - "wire_api IN ({})", - sql_in_list(&query.filter.wire_apis) - )); - } - if !query.filter.models.is_empty() { - // `models_used` is stored as a JSON string array. DuckDB uses - // `list_has_any`; the ClickHouse equivalent parses the JSON to - // `Array(String)` and tests overlap with the filter list. - where_parts.push(format!( - "hasAny(JSONExtract(coalesce(models_used, '[]'), 'Array(String)'), [{}])", - sql_in_list(&query.filter.models) - )); - } - if !query.filter.server_ips.is_empty() { - where_parts.push(format!( - "server_ip IN ({})", - sql_in_list(&query.filter.server_ips) - )); - } - if !query.include_proxy_hops { - // `metadata` is Nullable(String) holding JSON. JSONExtractString - // returns '' when the path is absent, so DuckDB's `IS NULL` maps to - // `= ''`; the role exclusion stays an explicit NOT IN list. - where_parts.push( - "(JSONExtractString(coalesce(metadata, ''), 'proxy', 'role') = '' \ - OR JSONExtractString(coalesce(metadata, ''), 'proxy', 'role') \ - NOT IN ('proxy_out', 'mirror_secondary'))" - .to_string(), - ); - } - let sql = format!( "SELECT DISTINCT agent_kind AS v FROM traces FINAL \ WHERE {} ORDER BY agent_kind", - where_parts.join(" AND ") + distinct_agent_kinds_where_sql(query) ); let rows = self .client @@ -157,3 +118,104 @@ impl ClickHouseBackend { .collect()) } } + +/// Build the `query_distinct_agent_kinds` WHERE clause: a half-open +/// `start_time` time range AND-ed with every present dimension filter plus, +/// when `include_proxy_hops` is false, a proxy-role exclusion. Extracted as a +/// pure fn so the escaping / IN-list / JSON-array / proxy-hop assembly is +/// unit-testable without a live server. `traces` is a ReplacingMergeTree, so +/// the caller wraps the result in `FROM traces FINAL`. +pub(crate) fn distinct_agent_kinds_where_sql(query: &DistinctAgentKindsQuery) -> String { + let mut where_parts = vec![time_where( + "start_time", + query.time_range.start_us, + query.time_range.end_us, + )]; + + if !query.filter.wire_apis.is_empty() { + where_parts.push(format!("wire_api IN ({})", sql_in_list(&query.filter.wire_apis))); + } + if !query.filter.models.is_empty() { + // `models_used` is stored as a JSON string array. DuckDB uses + // `list_has_any`; the ClickHouse equivalent parses the JSON to + // `Array(String)` and tests overlap with the filter list. + where_parts.push(format!( + "hasAny(JSONExtract(coalesce(models_used, '[]'), 'Array(String)'), [{}])", + sql_in_list(&query.filter.models) + )); + } + if !query.filter.server_ips.is_empty() { + where_parts.push(format!("server_ip IN ({})", sql_in_list(&query.filter.server_ips))); + } + if !query.include_proxy_hops { + // `metadata` is Nullable(String) holding JSON. JSONExtractString returns + // '' when the path is absent, so DuckDB's `IS NULL` maps to `= ''`; the + // role exclusion stays an explicit NOT IN list. + where_parts.push( + "(JSONExtractString(coalesce(metadata, ''), 'proxy', 'role') = '' \ + OR JSONExtractString(coalesce(metadata, ''), 'proxy', 'role') \ + NOT IN ('proxy_out', 'mirror_secondary'))" + .to_string(), + ); + } + where_parts.join(" AND ") +} + +#[cfg(test)] +mod tests { + use super::*; + use h_storage::query::{DimensionFilter, TimeRange}; + + fn q() -> DistinctAgentKindsQuery { + DistinctAgentKindsQuery { + time_range: TimeRange { start_us: 100, end_us: 200 }, + filter: DimensionFilter::default(), + include_proxy_hops: false, + } + } + + #[test] + fn default_excludes_proxy_hops() { + let s = distinct_agent_kinds_where_sql(&q()); + assert!(s.starts_with("start_time >= fromUnixTimestamp64Micro(100)")); + assert!(s.contains("start_time < fromUnixTimestamp64Micro(200)")); + // Proxy exclusion: role = '' OR role NOT IN (proxy_out, mirror_secondary). + assert!(s.contains("JSONExtractString(coalesce(metadata, ''), 'proxy', 'role') = ''")); + assert!(s.contains("NOT IN ('proxy_out', 'mirror_secondary')")); + } + + #[test] + fn include_proxy_hops_omits_exclusion() { + let query = DistinctAgentKindsQuery { include_proxy_hops: true, ..q() }; + let s = distinct_agent_kinds_where_sql(&query); + assert!(!s.contains("proxy_out")); + } + + #[test] + fn models_uses_hasany_json_extract() { + let query = DistinctAgentKindsQuery { + filter: DimensionFilter { + models: vec!["gpt-4".into()], + ..Default::default() + }, + ..q() + }; + assert!(distinct_agent_kinds_where_sql(&query) + .contains("hasAny(JSONExtract(coalesce(models_used, '[]'), 'Array(String)'), ['gpt-4'])")); + } + + #[test] + fn dimension_in_lists_escape_quotes() { + let query = DistinctAgentKindsQuery { + filter: DimensionFilter { + wire_apis: vec!["a'b".into()], + server_ips: vec!["10.0.0.2".into()], + ..Default::default() + }, + ..q() + }; + let s = distinct_agent_kinds_where_sql(&query); + assert!(s.contains("wire_api IN ('a''b')")); + assert!(s.contains("server_ip IN ('10.0.0.2')")); + } +} diff --git a/server/h-storage-clickhouse/src/exchanges.rs b/server/h-storage-clickhouse/src/exchanges.rs index e087716b..aff6d2b5 100644 --- a/server/h-storage-clickhouse/src/exchanges.rs +++ b/server/h-storage-clickhouse/src/exchanges.rs @@ -112,56 +112,14 @@ impl ClickHouseBackend { query.sort_by ))); } - let sort_order = if query.sort_order.eq_ignore_ascii_case("ASC") { - "ASC" - } else { - "DESC" - }; + let sort_order = crate::calls::resolve_sort_order(&query.sort_order); - // Time-range + filter WHERE. Timestamps and numeric values are - // interpolated (values we control); user string lists go through - // `sql_in_list`'s single-quote escaping, valid in ClickHouse. - let mut where_parts = vec![time_where( - "request_time", - query.time_range.start_us, - query.time_range.end_us, - )]; - if !query.server_ips.is_empty() { - where_parts.push(format!("server_ip IN ({})", sql_in_list(&query.server_ips))); - } - if !query.client_ips.is_empty() { - where_parts.push(format!("client_ip IN ({})", sql_in_list(&query.client_ips))); - } - if !query.methods.is_empty() { - where_parts.push(format!("method IN ({})", sql_in_list(&query.methods))); - } - if !query.status_codes.is_empty() { - where_parts.push(format!("status IN ({})", join_nums(&query.status_codes))); - } - if let Some(substr) = query - .uri_contains - .as_deref() - .map(str::trim) - .filter(|s| !s.is_empty()) - { - where_parts.push(format!("uri LIKE '%{}%'", escape_str(substr))); - } - if let Some(sse) = query.is_sse { - where_parts.push(format!("is_sse = {}", if sse { 1 } else { 0 })); - } - let where_sql = where_parts.join(" AND "); + let where_sql = http_exchanges_where_sql(query); // Map virtual field → column/expression for ORDER BY. `duration_ms` // and `status` get `NULLS LAST` so incomplete (duration/status=None) // rows don't dominate a descending sort, matching the DuckDB backend. - let order_expr = match query.sort_by.as_str() { - "duration_ms" => { - "(toUnixTimestamp64Micro(response_complete_time) \ - - toUnixTimestamp64Micro(request_time)) / 1000.0 NULLS LAST" - } - "status" => "status NULLS LAST", - _ => "request_time", - }; + let order_expr = http_exchanges_order_expr(&query.sort_by); let total = self .client @@ -211,6 +169,59 @@ fn join_nums(values: &[T]) -> String { .join(", ") } +/// Build the `query_http_exchanges` WHERE clause: a half-open `request_time` +/// time range AND-ed with every present filter. Extracted as a pure fn so the +/// escaping / IN-list / LIKE assembly is unit-testable without a live server. +/// Timestamps and numeric values are interpolated (values we control); user +/// string lists go through `sql_in_list`'s backslash-aware escaping; the `LIKE` +/// substring goes through `escape_str`. +pub(crate) fn http_exchanges_where_sql(query: &HttpExchangesQuery) -> String { + let mut where_parts = vec![time_where( + "request_time", + query.time_range.start_us, + query.time_range.end_us, + )]; + if !query.server_ips.is_empty() { + where_parts.push(format!("server_ip IN ({})", sql_in_list(&query.server_ips))); + } + if !query.client_ips.is_empty() { + where_parts.push(format!("client_ip IN ({})", sql_in_list(&query.client_ips))); + } + if !query.methods.is_empty() { + where_parts.push(format!("method IN ({})", sql_in_list(&query.methods))); + } + if !query.status_codes.is_empty() { + where_parts.push(format!("status IN ({})", join_nums(&query.status_codes))); + } + if let Some(substr) = query + .uri_contains + .as_deref() + .map(str::trim) + .filter(|s| !s.is_empty()) + { + where_parts.push(format!("uri LIKE '%{}%'", escape_str(substr))); + } + if let Some(sse) = query.is_sse { + where_parts.push(format!("is_sse = {}", if sse { 1 } else { 0 })); + } + where_parts.join(" AND ") +} + +/// Map the `query_http_exchanges` virtual `sort_by` field to the ORDER BY +/// expression. `duration_ms` and `status` get `NULLS LAST` so incomplete +/// (duration/status=None) rows don't dominate a descending sort, matching the +/// DuckDB backend. Extracted as a pure fn for offline testability. +pub(crate) fn http_exchanges_order_expr(sort_by: &str) -> &'static str { + match sort_by { + "duration_ms" => { + "(toUnixTimestamp64Micro(response_complete_time) \ + - toUnixTimestamp64Micro(request_time)) / 1000.0 NULLS LAST" + } + "status" => "status NULLS LAST", + _ => "request_time", + } +} + fn exchange_detail(r: ExchangeDetailRow) -> HttpExchangeDetail { HttpExchangeDetail { id: r.id, @@ -252,3 +263,204 @@ fn exchange_list_item(r: ExchangeListRow) -> HttpExchangeListItem { duration_ms: r.duration_ms, } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn join_nums_renders_numeric_in_list() { + assert_eq!(join_nums(&[200u16, 429u16]), "200, 429"); + assert_eq!(join_nums(&[200u16]), "200"); + assert_eq!(join_nums::(&[]), ""); + } + + #[test] + fn valid_sort_fields_is_whitelisted() { + // sort_by is interpolated into ORDER BY, so the whitelist is the gate + // that keeps an unknown field from reaching the engine. + assert_eq!(VALID_SORT_FIELDS, &["request_time", "status", "duration_ms"]); + assert!(!VALID_SORT_FIELDS.contains(&"bogus")); + } + + fn detail_row() -> ExchangeDetailRow { + ExchangeDetailRow { + id: "xchg-1".into(), + source_id: "src-0".into(), + client_ip: "10.0.0.1".into(), + client_port: 54321, + server_ip: "10.0.0.2".into(), + server_port: 8080, + method: "POST".into(), + uri: "/v1/chat/completions".into(), + request_headers: r#"[["content-type","application/json"]]"#.into(), + request_body: Some(r#"{"model":"gpt-4"}"#.into()), + status: Some(200), + response_headers: r#"[["x-request-id","abc"]]"#.into(), + response_body: Some(r#"{"choices":[]}"#.into()), + is_sse: false, + sse_event_count: 0, + sse_data_bytes: 0, + request_time_us: 1_700_000_000_000_000, + response_first_byte_time_us: Some(1_700_000_000_500_000), + response_complete_time_us: Some(1_700_000_001_000_000), + } + } + + #[test] + fn exchange_detail_maps_micros_timestamps_and_bodies() { + let d = exchange_detail(detail_row()); + assert_eq!(d.id, "xchg-1"); + assert_eq!(d.source_id, "src-0"); + assert_eq!(d.client_ip, "10.0.0.1"); + assert_eq!(d.client_port, 54321); + assert_eq!(d.server_ip, "10.0.0.2"); + assert_eq!(d.server_port, 8080); + assert_eq!(d.method, "POST"); + assert_eq!(d.uri, "/v1/chat/completions"); + assert_eq!(d.request_headers, r#"[["content-type","application/json"]]"#); + assert_eq!(d.request_body.as_deref(), Some(r#"{"model":"gpt-4"}"#)); + assert_eq!(d.status, Some(200)); + assert_eq!(d.response_headers, r#"[["x-request-id","abc"]]"#); + assert_eq!(d.response_body.as_deref(), Some(r#"{"choices":[]}"#)); + assert!(!d.is_sse); + assert_eq!(d.sse_event_count, 0); + assert_eq!(d.sse_data_bytes, 0); + // Microsecond timestamps pass through verbatim (DateTime64(6) ↔ i64 micros). + assert_eq!(d.request_time, 1_700_000_000_000_000); + assert_eq!(d.response_first_byte_time_us, Some(1_700_000_000_500_000)); + assert_eq!(d.response_complete_time_us, Some(1_700_000_001_000_000)); + } + + fn list_row() -> ExchangeListRow { + ExchangeListRow { + id: "xchg-1".into(), + source_id: "src-0".into(), + request_time_ms: 1_700_000_000_000, + method: "POST".into(), + uri: "/v1/chat/completions".into(), + client_ip: "10.0.0.1".into(), + server_ip: "10.0.0.2".into(), + server_port: 8080, + status: Some(200), + is_sse: false, + duration_ms: Some(1000.0), + } + } + + #[test] + fn exchange_list_item_maps_ms_request_time_and_duration() { + let it = exchange_list_item(list_row()); + assert_eq!(it.id, "xchg-1"); + assert_eq!(it.source_id, "src-0"); + assert_eq!(it.request_time, 1_700_000_000_000); + assert_eq!(it.method, "POST"); + assert_eq!(it.uri, "/v1/chat/completions"); + assert_eq!(it.client_ip, "10.0.0.1"); + assert_eq!(it.server_ip, "10.0.0.2"); + assert_eq!(it.server_port, 8080); + assert_eq!(it.status, Some(200)); + assert!(!it.is_sse); + assert_eq!(it.duration_ms, Some(1000.0)); + } + + #[test] + fn exchange_list_item_passes_none_status_and_duration() { + let mut r = list_row(); + r.status = None; + r.duration_ms = None; + let it = exchange_list_item(r); + assert_eq!(it.status, None); + assert_eq!(it.duration_ms, None); + } + + fn exchanges_query() -> HttpExchangesQuery { + HttpExchangesQuery { + time_range: TimeRange { start_us: 100, end_us: 200 }, + server_ips: vec![], + client_ips: vec![], + methods: vec![], + status_codes: vec![], + uri_contains: None, + is_sse: None, + sort_by: "request_time".into(), + sort_order: "desc".into(), + page: 1, + page_size: 10, + } + } + + #[test] + fn http_exchanges_where_sql_is_time_range_only_by_default() { + let s = http_exchanges_where_sql(&exchanges_query()); + assert_eq!( + s, + "request_time >= fromUnixTimestamp64Micro(100) \ + AND request_time < fromUnixTimestamp64Micro(200)" + ); + } + + #[test] + fn http_exchanges_where_sql_combines_all_filters() { + let q = HttpExchangesQuery { + server_ips: vec!["10.0.0.2".into()], + client_ips: vec!["10.0.0.1".into()], + methods: vec!["POST".into()], + status_codes: vec![200, 429], + uri_contains: Some("/v1/chat".into()), + is_sse: Some(true), + ..exchanges_query() + }; + let s = http_exchanges_where_sql(&q); + assert!(s.contains("server_ip IN ('10.0.0.2')")); + assert!(s.contains("client_ip IN ('10.0.0.1')")); + assert!(s.contains("method IN ('POST')")); + assert!(s.contains("status IN (200, 429)")); + assert!(s.contains("uri LIKE '%/v1/chat%'")); + assert!(s.contains("is_sse = 1")); + assert!(!s.starts_with(" AND")); + assert!(!s.ends_with("AND ")); + } + + #[test] + fn http_exchanges_where_sql_escapes_like_quote() { + let q = HttpExchangesQuery { + uri_contains: Some("a'b".into()), + ..exchanges_query() + }; + assert!(http_exchanges_where_sql(&q).contains("uri LIKE '%a''b%'")); + } + + #[test] + fn http_exchanges_where_sql_ignores_blank_like() { + let q = HttpExchangesQuery { + uri_contains: Some(" ".into()), + ..exchanges_query() + }; + assert!(!http_exchanges_where_sql(&q).contains("LIKE")); + } + + #[test] + fn http_exchanges_where_sql_is_sse_false() { + let q = HttpExchangesQuery { + is_sse: Some(false), + ..exchanges_query() + }; + assert!(http_exchanges_where_sql(&q).contains("is_sse = 0")); + } + + #[test] + fn http_exchanges_order_expr_maps_virtual_fields() { + assert_eq!(http_exchanges_order_expr("request_time"), "request_time"); + assert_eq!(http_exchanges_order_expr("status"), "status NULLS LAST"); + // duration_ms derives from the complete−request gap with NULLS LAST so + // incomplete exchanges don't dominate a descending sort. + let d = http_exchanges_order_expr("duration_ms"); + assert!(d.contains("toUnixTimestamp64Micro(response_complete_time)")); + assert!(d.contains("toUnixTimestamp64Micro(request_time)")); + assert!(d.contains("NULLS LAST")); + // An unknown sort_by falls back to request_time (the whitelist rejects + // it before this is reached, but the fallback must be safe regardless). + assert_eq!(http_exchanges_order_expr("bogus"), "request_time"); + } +} diff --git a/server/h-storage-clickhouse/src/metrics.rs b/server/h-storage-clickhouse/src/metrics.rs index 306f4a0b..f26730b3 100644 --- a/server/h-storage-clickhouse/src/metrics.rs +++ b/server/h-storage-clickhouse/src/metrics.rs @@ -103,6 +103,22 @@ const SUM_FIELDS: &[&str] = &[ const MAX_FIELDS: &[&str] = &["active_calls_max"]; +/// Valid `sort_by` fields for `query_metrics_models`. Hoisted to module scope so +/// the reject-unknown-sort path is unit-testable without a live client (the +/// value is interpolated into `ORDER BY`, so an unknown field is rejected up +/// front rather than reaching the engine). Mirrors the DuckDB whitelist. +const MODELS_VALID_SORT_FIELDS: &[&str] = &[ + "call_count", + "error_count", + "total_input_tokens", + "total_output_tokens", + "ttft_avg", + "ttft_p95", + "e2e_avg", + "e2e_p95", + "tpot_avg", +]; + fn avg_pair(f: &str) -> Option<(&'static str, &'static str)> { match f { "active_calls_avg" => Some(("active_calls_sum", "active_calls_sample_count")), @@ -363,18 +379,7 @@ impl ClickHouseBackend { &self, query: &MetricsModelsQuery, ) -> Result> { - const VALID_SORT_FIELDS: &[&str] = &[ - "call_count", - "error_count", - "total_input_tokens", - "total_output_tokens", - "ttft_avg", - "ttft_p95", - "e2e_avg", - "e2e_p95", - "tpot_avg", - ]; - if !VALID_SORT_FIELDS.contains(&query.sort_by.as_str()) { + if !MODELS_VALID_SORT_FIELDS.contains(&query.sort_by.as_str()) { return Err(AppError::Storage(format!( "invalid sort_by field: {}", query.sort_by @@ -584,3 +589,189 @@ impl ClickHouseBackend { .collect()) } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn avg_pair_maps_known_avg_fields() { + assert_eq!( + avg_pair("ttft_avg"), + Some(("ttft_sum", "ttft_count")) + ); + assert_eq!( + avg_pair("e2e_avg"), + Some(("e2e_sum", "e2e_count")) + ); + assert_eq!( + avg_pair("tpot_avg"), + Some(("tpot_sum", "tpot_count")) + ); + assert_eq!( + avg_pair("active_calls_avg"), + Some(("active_calls_sum", "active_calls_sample_count")) + ); + assert_eq!(avg_pair("input_tokens_avg"), Some(("total_input_tokens", "input_token_count"))); + assert_eq!(avg_pair("output_tokens_avg"), Some(("total_output_tokens", "output_token_count"))); + assert_eq!(avg_pair("ttft_stream_avg"), Some(("ttft_stream_sum", "ttft_stream_count"))); + assert_eq!( + avg_pair("ttft_nonstream_avg"), + Some(("ttft_nonstream_sum", "ttft_nonstream_count")) + ); + } + + #[test] + fn avg_pair_unknown_is_none() { + assert_eq!(avg_pair("call_count"), None); + assert_eq!(avg_pair("ttft_p95"), None); + assert_eq!(avg_pair("not_a_field"), None); + } + + #[test] + fn percentile_weight_routes_by_prefix() { + assert_eq!(percentile_weight("ttft_stream_p95"), "ttft_stream_count"); + assert_eq!(percentile_weight("ttft_nonstream_p99"), "ttft_nonstream_count"); + assert_eq!(percentile_weight("ttft_p50"), "ttft_count"); + assert_eq!(percentile_weight("e2e_p95"), "e2e_count"); + assert_eq!(percentile_weight("tpot_p95"), "tpot_count"); + // Non-latency field falls through to call_count. + assert_eq!(percentile_weight("call_count"), "call_count"); + assert_eq!(percentile_weight("anything"), "call_count"); + } + + #[test] + fn ch_field_expr_sum_for_count_total() { + let e = ch_field_expr("call_count"); + assert_eq!(e, "CAST(sum(call_count) AS Nullable(Float64))"); + // A SUM_FIELD that is not an avg/percentile/peak goes through the sum arm. + let e = ch_field_expr("total_input_tokens"); + assert_eq!(e, "CAST(sum(total_input_tokens) AS Nullable(Float64))"); + } + + #[test] + fn ch_field_expr_max_for_peak() { + let e = ch_field_expr("active_calls_max"); + assert_eq!(e, "CAST(max(active_calls_max) AS Nullable(Float64))"); + } + + #[test] + fn ch_field_expr_avg_is_count_guarded_ratio() { + let e = ch_field_expr("ttft_avg"); + assert_eq!( + e, + "CAST(if(sum(ttft_count) > 0, sum(ttft_sum) / sum(ttft_count), NULL) AS Nullable(Float64))" + ); + // Zero-denominator guard avoids divide-by-zero; the outer CAST keeps the + // array element type uniform. + assert!(e.contains("if(sum(ttft_count) > 0")); + } + + #[test] + fn ch_field_expr_percentile_is_count_weighted_average() { + let e = ch_field_expr("ttft_p95"); + assert_eq!( + e, + "CAST(if(sum(ttft_count) > 0, sum(ttft_p95 * ttft_count) / sum(ttft_count), NULL) \ + AS Nullable(Float64))" + ); + // The stream/non-stream variants route to their own count columns. + let es = ch_field_expr("ttft_stream_p99"); + assert!(es.contains("sum(ttft_stream_count)")); + let en = ch_field_expr("ttft_nonstream_p50"); + assert!(en.contains("sum(ttft_nonstream_count)")); + } + + #[test] + fn ch_field_expr_unknown_falls_back_to_sum() { + // Anything unrecognized reaches the final `sum({f})` fallback arm. + let e = ch_field_expr("not_a_field"); + assert_eq!(e, "CAST(sum(not_a_field) AS Nullable(Float64))"); + } + + #[test] + fn build_vals_array_empty() { + assert_eq!( + build_vals_array(&[]), + "CAST([] AS Array(Nullable(Float64))) AS vals" + ); + } + + #[test] + fn build_vals_array_projects_each_field() { + let vals = build_vals_array(&["call_count".into(), "ttft_avg".into()]); + assert!(vals.starts_with("[")); + assert!(vals.ends_with(" AS vals")); + assert!(vals.contains("CAST(sum(call_count) AS Nullable(Float64))")); + assert!(vals.contains("if(sum(ttft_count) > 0, sum(ttft_sum) / sum(ttft_count), NULL)")); + // Fields are comma-joined inside the array projection. + assert!(vals.contains(", ")); + } + + #[test] + fn ts_where_targets_timestamp_column() { + assert_eq!( + ts_where(10, 20), + "timestamp >= fromUnixTimestamp64Micro(10) \ + AND timestamp < fromUnixTimestamp64Micro(20)" + ); + } + + #[test] + fn valid_metric_fields_contains_known_set() { + // Guards the reject-unknown-field path by construction: these must be + // present so the API field names the frontend sends are accepted. + for &known in &[ + "call_count", + "ttft_avg", + "ttft_p95", + "ttft_p99", + "ttft_stream_p95", + "ttft_nonstream_p99", + "e2e_avg", + "tpot_p50", + "error_429_count", + "active_calls_avg", + ] { + assert!( + VALID_METRIC_FIELDS.contains(&known), + "VALID_METRIC_FIELDS missing {known}" + ); + } + assert!(!VALID_METRIC_FIELDS.contains(&"bogus_field")); + } + + #[test] + fn models_sort_whitelist_contains_expected() { + for &known in &[ + "call_count", + "error_count", + "total_input_tokens", + "total_output_tokens", + "ttft_avg", + "ttft_p95", + "e2e_avg", + "e2e_p95", + "tpot_avg", + ] { + assert!( + MODELS_VALID_SORT_FIELDS.contains(&known), + "models MODELS_VALID_SORT_FIELDS missing {known}" + ); + } + assert!(!MODELS_VALID_SORT_FIELDS.contains(&"bogus")); + } + + #[test] + fn sum_fields_and_max_fields_disjoint_from_avg() { + // Every SUM_FIELD must NOT be an avg field (avg_pair None) and not a + // percentile, so the dispatch order in ch_field_expr is unambiguous. + for &f in SUM_FIELDS { + assert_eq!(avg_pair(f), None, "{f} unexpectedly mapped to an avg pair"); + assert!(!f.ends_with("_p50") && !f.ends_with("_p95") && !f.ends_with("_p99")); + } + for &f in MAX_FIELDS { + assert_eq!(avg_pair(f), None); + } + } +} diff --git a/server/h-storage-clickhouse/src/retention.rs b/server/h-storage-clickhouse/src/retention.rs index 16eef9c0..d8678bc5 100644 --- a/server/h-storage-clickhouse/src/retention.rs +++ b/server/h-storage-clickhouse/src/retention.rs @@ -51,6 +51,39 @@ fn cutoff_micros(t: SystemTime) -> Result { .map_err(|_| AppError::Storage("retention cutoff out of i64 range".to_string())) } +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + #[test] + fn cutoff_micros_known_time_to_epoch_micros() { + // 1_700_000_000 s since epoch → 1.7e15 µs, a value in the realistic + // LLM-traffic range (2023-11-14 UTC). + let t = SystemTime::UNIX_EPOCH + Duration::from_secs(1_700_000_000); + assert_eq!(cutoff_micros(t).unwrap(), 1_700_000_000_000_000); + } + + #[test] + fn cutoff_micros_preserves_subsecond_micros() { + let t = SystemTime::UNIX_EPOCH + Duration::from_micros(1_234_567); + assert_eq!(cutoff_micros(t).unwrap(), 1_234_567); + } + + #[test] + fn cutoff_micros_epoch_is_zero() { + assert_eq!(cutoff_micros(SystemTime::UNIX_EPOCH).unwrap(), 0); + } + + #[test] + fn cutoff_micros_before_epoch_is_err() { + // A time strictly before the UNIX epoch errors (the predicate would + // otherwise produce a nonsensical negative micros literal). + let before = SystemTime::UNIX_EPOCH - Duration::from_secs(1); + assert!(cutoff_micros(before).is_err()); + } +} + impl ClickHouseBackend { /// Count the rows that match `predicate` on `table` (best-effort — see the /// module docs). Used to populate the `RetentionReport` since ClickHouse diff --git a/server/h-storage-clickhouse/src/rows.rs b/server/h-storage-clickhouse/src/rows.rs index f104684b..29d7d9c0 100644 --- a/server/h-storage-clickhouse/src/rows.rs +++ b/server/h-storage-clickhouse/src/rows.rs @@ -391,3 +391,437 @@ impl From for ExchangeRow { } } } + +#[cfg(test)] +mod tests { + use super::*; + use h_common::agent::{AgentTopology, ToolSurface}; + use h_common::process::ProcessInfo; + use h_llm::model::ApiType; + use h_llm::wire_apis as wa; + use h_metrics::model::{LlmFinishMetric, LlmMetric}; + use h_protocol::HttpExchange; + use h_turn::{Trace, TraceStatus}; + use std::net::IpAddr; + + /// Minimal `LlmCall` with non-default scalar fields set so every `CallRow` + /// mapping branch is exercised. Mirrors the shape used by the live IT + /// fixtures but kept local to this module. + fn sample_call(id: &str) -> LlmCall { + LlmCall { + source_id: "src-0".into(), + id: id.into(), + wire_api: wa::OPENAI_CHAT, + model: "gpt-4".into(), + api_type: ApiType::Chat, + request_time: 1_700_000_000_000_000, + response_time: Some(1_700_000_000_500_000), + complete_time: Some(1_700_000_001_000_000), + request_path: "/v1/chat/completions".into(), + is_stream: true, + request_body: Some(r#"{"model":"gpt-4"}"#.into()), + status_code: Some(200), + finish_reason: Some("stop".into()), + response_body: Some(r#"{"choices":[]}"#.into()), + input_tokens: Some(100), + output_tokens: Some(50), + total_tokens: Some(150), + cache_read_input_tokens: Some(10), + cache_creation_input_tokens: Some(20), + ttft_ms: Some(500.0), + e2e_latency_ms: Some(1000.0), + client_ip: "10.0.0.1".parse::().unwrap(), + client_port: 54321, + server_ip: "10.0.0.2".parse::().unwrap(), + server_port: 8080, + response_id: Some("chatcmpl-x".into()), + request_headers: vec![("content-type".into(), "application/json".into())], + response_headers: vec![("x-request-id".into(), "abc".into())], + is_agent_request: true, + tool_surface: Some(ToolSurface::Mcp), + agent_topology: Some(AgentTopology::Orchestrator), + tool_call_count: 3, + tool_names: vec!["bash".into(), "grep".into()], + body_bytes_dropped: 0, + process: Some(ProcessInfo::new(42, "node")), + } + } + + #[test] + fn call_row_from_llm_call_maps_scalars_and_stringified_ips() { + let row = CallRow::from(sample_call("call-1")); + assert_eq!(row.id, "call-1"); + assert_eq!(row.source_id, "src-0"); + assert_eq!(row.client_ip, "10.0.0.1"); + assert_eq!(row.client_port, 54321); + assert_eq!(row.server_ip, "10.0.0.2"); + assert_eq!(row.server_port, 8080); + assert_eq!(row.request_time, 1_700_000_000_000_000); + assert_eq!(row.response_time, Some(1_700_000_000_500_000)); + assert_eq!(row.complete_time, Some(1_700_000_001_000_000)); + assert_eq!(row.wire_api, "openai-chat"); + assert_eq!(row.model, "gpt-4"); + assert_eq!(row.api_type, "chat"); + assert!(row.is_stream); + assert_eq!(row.request_path, "/v1/chat/completions"); + assert_eq!(row.status_code, Some(200)); + assert_eq!(row.finish_reason.as_deref(), Some("stop")); + assert_eq!(row.input_tokens, Some(100)); + assert_eq!(row.output_tokens, Some(50)); + assert_eq!(row.total_tokens, Some(150)); + assert_eq!(row.cache_read_input_tokens, Some(10)); + assert_eq!(row.cache_creation_input_tokens, Some(20)); + assert_eq!(row.ttft_ms, Some(500.0)); + assert_eq!(row.e2e_latency_ms, Some(1000.0)); + assert_eq!(row.request_body.as_deref(), Some(r#"{"model":"gpt-4"}"#)); + assert_eq!(row.response_body.as_deref(), Some(r#"{"choices":[]}"#)); + assert_eq!(row.response_id.as_deref(), Some("chatcmpl-x")); + assert!(row.is_agent_request); + assert_eq!(row.tool_call_count, 3); + assert_eq!(row.body_bytes_dropped, 0); + assert_eq!(row.process_pid, Some(42)); + assert_eq!(row.process_comm.as_deref(), Some("node")); + assert_eq!(row.process_exe, None); // ProcessInfo::new leaves exe None + assert_eq!(row.kind, "llm"); + } + + #[test] + fn call_row_headers_are_json_pair_arrays() { + let row = CallRow::from(sample_call("c")); + let req: serde_json::Value = serde_json::from_str(&row.request_headers).unwrap(); + assert!(req.is_array()); + assert_eq!(req[0][0], "content-type"); + assert_eq!(req[0][1], "application/json"); + let resp: serde_json::Value = serde_json::from_str(&row.response_headers).unwrap(); + assert_eq!(resp[0][0], "x-request-id"); + assert_eq!(resp[0][1], "abc"); + } + + #[test] + fn call_row_tool_names_and_surface_and_topology_serialized() { + let row = CallRow::from(sample_call("c")); + let names: Vec = serde_json::from_str(row.tool_names_json.as_deref().unwrap()).unwrap(); + assert_eq!(names, vec!["bash".to_string(), "grep".to_string()]); + assert_eq!(row.tool_surface.as_deref(), Some("mcp")); + assert_eq!(row.agent_topology.as_deref(), Some("orchestrator")); + } + + #[test] + fn call_row_from_empty_tool_names_yields_json_array() { + let mut c = sample_call("c"); + c.tool_names = vec![]; + let row = CallRow::from(c); + assert_eq!(row.tool_names_json.as_deref(), Some("[]")); + } + + #[test] + fn call_row_passive_tap_has_no_process() { + let mut c = sample_call("c"); + c.process = None; + let row = CallRow::from(c); + assert_eq!(row.process_pid, None); + assert_eq!(row.process_comm, None); + assert_eq!(row.process_exe, None); + } + + fn sample_metric() -> LlmMetric { + LlmMetric { + timestamp_us: 1_700_000_000_000_000, + source_id: "src-0".into(), + granularity: "1m", + wire_api: "openai-chat".into(), + model: "gpt-4".into(), + server_ip: "10.0.0.2".into(), + call_count: 5, + stream_count: 3, + non_stream_count: 2, + active_calls_sum: 7, + active_calls_sample_count: 4, + active_calls_max: 9, + total_input_tokens: 100, + input_token_count: 5, + total_output_tokens: 50, + output_token_count: 5, + total_cache_read_input_tokens: 10, + total_cache_creation_input_tokens: 20, + error_count: 1, + error_4xx_count: 1, + error_429_count: 1, + error_5xx_count: 0, + ttft_sum: 2500.0, + ttft_count: 5, + ttft_p50: Some(400.0), + ttft_p95: Some(600.0), + ttft_p99: Some(900.0), + ttft_stream_sum: 2000.0, + ttft_stream_count: 4, + ttft_stream_p50: Some(400.0), + ttft_stream_p95: Some(500.0), + ttft_stream_p99: Some(550.0), + ttft_nonstream_sum: 500.0, + ttft_nonstream_count: 1, + ttft_nonstream_p50: Some(500.0), + ttft_nonstream_p95: None, + ttft_nonstream_p99: None, + e2e_sum: 5000.0, + e2e_count: 5, + e2e_p50: Some(800.0), + e2e_p95: Some(1200.0), + e2e_p99: Some(2000.0), + tpot_sum: 50.0, + tpot_count: 5, + tpot_p50: Some(10.0), + tpot_p95: Some(12.0), + tpot_p99: Some(15.0), + tool_surface: Some("cli".into()), + } + } + + #[test] + fn metric_row_from_llm_metric_is_field_for_field() { + let row = MetricRow::from(sample_metric()); + assert_eq!(row.timestamp, 1_700_000_000_000_000); + assert_eq!(row.source_id, "src-0"); + assert_eq!(row.granularity, "1m"); + assert_eq!(row.wire_api, "openai-chat"); + assert_eq!(row.model, "gpt-4"); + assert_eq!(row.server_ip, "10.0.0.2"); + assert_eq!(row.call_count, 5); + assert_eq!(row.stream_count, 3); + assert_eq!(row.non_stream_count, 2); + assert_eq!(row.active_calls_sum, 7); + assert_eq!(row.active_calls_sample_count, 4); + assert_eq!(row.active_calls_max, 9); + assert_eq!(row.total_input_tokens, 100); + assert_eq!(row.input_token_count, 5); + assert_eq!(row.total_output_tokens, 50); + assert_eq!(row.output_token_count, 5); + assert_eq!(row.total_cache_read_input_tokens, 10); + assert_eq!(row.total_cache_creation_input_tokens, 20); + assert_eq!(row.error_count, 1); + assert_eq!(row.error_4xx_count, 1); + assert_eq!(row.error_429_count, 1); + assert_eq!(row.error_5xx_count, 0); + assert_eq!(row.ttft_sum, 2500.0); + assert_eq!(row.ttft_count, 5); + assert_eq!(row.ttft_p50, Some(400.0)); + assert_eq!(row.ttft_p95, Some(600.0)); + assert_eq!(row.ttft_p99, Some(900.0)); + assert_eq!(row.ttft_stream_sum, 2000.0); + assert_eq!(row.ttft_stream_count, 4); + assert_eq!(row.e2e_sum, 5000.0); + assert_eq!(row.e2e_count, 5); + assert_eq!(row.e2e_p99, Some(2000.0)); + assert_eq!(row.tpot_sum, 50.0); + assert_eq!(row.tpot_count, 5); + assert_eq!(row.tpot_p95, Some(12.0)); + // None percentiles must pass through as None (nullable columns). + assert_eq!(row.ttft_nonstream_p95, None); + assert_eq!(row.ttft_nonstream_p99, None); + assert_eq!(row.tool_surface.as_deref(), Some("cli")); + } + + #[test] + fn finish_metric_row_from_llm_finish_metric() { + let m = LlmFinishMetric { + timestamp_us: 1_700_000_000_000_000, + source_id: "src-0".into(), + granularity: "1m".into(), + wire_api: "openai-chat".into(), + model: "gpt-4".into(), + server_ip: "10.0.0.2".into(), + finish_reason: "stop".into(), + count: 7, + }; + let row = FinishMetricRow::from(m); + assert_eq!(row.timestamp, 1_700_000_000_000_000); + assert_eq!(row.source_id, "src-0"); + assert_eq!(row.granularity, "1m"); + assert_eq!(row.wire_api, "openai-chat"); + assert_eq!(row.model, "gpt-4"); + assert_eq!(row.server_ip, "10.0.0.2"); + assert_eq!(row.finish_reason, "stop"); + assert_eq!(row.count, 7); + } + + fn sample_turn(end_time_us: i64) -> Trace { + Trace { + source_id: "src-0".into(), + turn_id: "turn-1".into(), + session_id: "sess-1".into(), + wire_api: "openai-chat".into(), + agent_kind: "claude-cli".into(), + client_ip: "10.0.0.1".parse().unwrap(), + server_ip: "10.0.0.2".parse().unwrap(), + start_time_us: end_time_us - 5_000_000, + end_time_us, + duration_ms: 5_000, + call_count: 2, + models_used: vec!["gpt-4".into()], + subagents_used: vec!["task".into()], + total_input_tokens: 100, + total_output_tokens: 50, + total_cache_read_input_tokens: 10, + total_cache_creation_input_tokens: 20, + total_cost_usd: Some(0.0123), + status: TraceStatus::Complete, + final_finish_reason: Some("stop".into()), + user_input_preview: Some("hello".into()), + user_call_id: Some("c1".into()), + final_answer_preview: Some("world".into()), + final_call_id: Some("c2".into()), + span_ids: vec!["c1".into(), "c2".into()], + metadata: serde_json::json!({"k": "v"}), + tool_surfaces: vec![ToolSurface::Mcp, ToolSurface::Cli], + tool_call_total: 4, + agent_topology: Some(AgentTopology::SubAgent), + suspicious_skills: vec![h_turn::SuspiciousSkillRollup { + tool_name: "bash".into(), + reason: "shell".into(), + }], + } + } + + #[test] + fn turn_row_from_trace_serializes_json_columns() { + let end = 1_700_000_001_000_000_i64; + let row = TurnRow::from(sample_turn(end)); + assert_eq!(row.turn_id, "turn-1"); + assert_eq!(row.source_id, "src-0"); + assert_eq!(row.session_id, "sess-1"); + assert_eq!(row.wire_api, "openai-chat"); + assert_eq!(row.agent_kind, "claude-cli"); + assert_eq!(row.client_ip, "10.0.0.1"); + assert_eq!(row.server_ip, "10.0.0.2"); + assert_eq!(row.start_time, end - 5_000_000); + assert_eq!(row.end_time, end); + assert_eq!(row.duration_ms, 5_000); + assert_eq!(row.call_count, 2); + assert_eq!(row.total_input_tokens, 100); + assert_eq!(row.total_output_tokens, 50); + assert_eq!(row.total_cache_read_input_tokens, 10); + assert_eq!(row.total_cache_creation_input_tokens, 20); + assert_eq!(row.total_cost_usd, Some(0.0123)); + assert_eq!(row.status, "complete"); + assert_eq!(row.final_finish_reason.as_deref(), Some("stop")); + assert_eq!(row.user_input_preview.as_deref(), Some("hello")); + assert_eq!(row.user_call_id.as_deref(), Some("c1")); + assert_eq!(row.final_answer_preview.as_deref(), Some("world")); + assert_eq!(row.final_call_id.as_deref(), Some("c2")); + assert_eq!(row.tool_call_total, 4); + assert_eq!(row.agent_topology.as_deref(), Some("sub_agent")); + + // JSON-encoded columns are real arrays / objects, not raw strings. + let span_ids: Vec = serde_json::from_str(&row.span_ids).unwrap(); + assert_eq!(span_ids, vec!["c1".to_string(), "c2".to_string()]); + let models: Vec = serde_json::from_str(row.models_used.as_deref().unwrap()).unwrap(); + assert_eq!(models, vec!["gpt-4".to_string()]); + let subs: Vec = + serde_json::from_str(row.subagents_used.as_deref().unwrap()).unwrap(); + assert_eq!(subs, vec!["task".to_string()]); + let md: serde_json::Value = serde_json::from_str(row.metadata.as_deref().unwrap()).unwrap(); + assert_eq!(md["k"], "v"); + let surfaces: Vec = + serde_json::from_str(row.tool_surfaces_json.as_deref().unwrap()).unwrap(); + assert_eq!(surfaces, vec!["mcp".to_string(), "cli".to_string()]); + let susp: Vec = + serde_json::from_str(row.suspicious_skills_json.as_deref().unwrap()).unwrap(); + assert_eq!(susp[0]["tool_name"], "bash"); + assert_eq!(susp[0]["reason"], "shell"); + } + + #[test] + fn turn_row_version_is_end_time_micros() { + // Initial finalize version = end_time (micros); update_trace_metadata + // re-inserts with a strictly-greater wall-clock-micros version. + let end = 1_700_000_001_000_000_i64; + let row = TurnRow::from(sample_turn(end)); + assert_eq!(row._version, end.max(0) as u64); + } + + #[test] + fn turn_row_version_clamps_negative_end_time() { + let row = TurnRow::from(sample_turn(-5)); + assert_eq!(row._version, 0); + } + + #[test] + fn exchange_row_from_http_exchange_maps_addrs_and_bodies() { + let x = sample_exchange("xchg-1", 1_700_000_000_000_000); + let row = ExchangeRow::from(x); + assert_eq!(row.id, "xchg-1"); + assert_eq!(row.source_id, "src-0"); + assert_eq!(row.client_ip, "10.0.0.1"); + assert_eq!(row.client_port, 54321); + assert_eq!(row.server_ip, "10.0.0.2"); + assert_eq!(row.server_port, 443); + assert_eq!(row.method, "POST"); + assert_eq!(row.uri, "/v1/chat/completions"); + assert_eq!(row.request_body.as_deref(), Some(r#"{"model":"gpt-4"}"#)); + assert_eq!(row.status, Some(200)); + assert_eq!(row.response_body.as_deref(), Some(r#"{"choices":[]}"#)); + assert!(!row.is_sse); + assert_eq!(row.sse_event_count, 0); + assert_eq!(row.sse_data_bytes, 0); + assert_eq!(row.request_time, 1_700_000_000_000_000); + assert_eq!(row.response_first_byte_time, Some(1_700_000_000_500_000)); + assert_eq!(row.response_complete_time, Some(1_700_000_001_000_000)); + // Headers serialized as JSON pair arrays. + let req: serde_json::Value = serde_json::from_str(&row.request_headers).unwrap(); + assert_eq!(req[0][0], "content-type"); + } + + #[test] + fn exchange_row_empty_request_body_becomes_none() { + let mut x = sample_exchange("x", 0); + // Replace the request with one whose body is empty. + let mut req = (*x.request).clone(); + req.body = bytes::Bytes::new(); + x.request = std::sync::Arc::new(req); + let row = ExchangeRow::from(x); + assert_eq!(row.request_body, None); + } + + /// Minimal paired HTTP exchange (generic IPs, placeholder ids) for the + /// `ExchangeRow::from` mapping — kept local so the test module is + /// self-contained. Mirrors the shape of the live-IT fixture. + fn sample_exchange(id: &str, request_time_us: i64) -> HttpExchange { + use bytes::Bytes; + use h_protocol::model::{HttpRequestData, HttpResponseData}; + use h_protocol::net::FlowKey; + use std::sync::Arc; + let client_ip: IpAddr = "10.0.0.1".parse().unwrap(); + let server_ip: IpAddr = "10.0.0.2".parse().unwrap(); + let request = Arc::new(HttpRequestData { + flow_key: FlowKey::new("src-0".into(), client_ip, 54321, server_ip, 443), + client_addr: (client_ip, 54321), + server_addr: (server_ip, 443), + method: "POST".into(), + uri: "/v1/chat/completions".into(), + version: 1, + headers: vec![("content-type".into(), "application/json".into())], + body: Bytes::from_static(br#"{"model":"gpt-4"}"#), + timestamp_us: request_time_us, + process: None, + }); + let response = Arc::new(HttpResponseData { + flow_key: request.flow_key.clone(), + client_addr: request.client_addr, + server_addr: request.server_addr, + status: 200, + version: 1, + headers: vec![("x-request-id".into(), "req_abc".into())], + body: Bytes::from_static(br#"{"choices":[]}"#), + first_byte_timestamp_us: request_time_us + 500_000, + complete_timestamp_us: request_time_us + 1_000_000, + process: None, + }); + HttpExchange { + id: id.to_string(), + request, + response, + sse_event_count: 0, + sse_data_bytes: 0, + } + } +} diff --git a/server/h-storage-clickhouse/src/services.rs b/server/h-storage-clickhouse/src/services.rs index d232797b..1401247d 100644 --- a/server/h-storage-clickhouse/src/services.rs +++ b/server/h-storage-clickhouse/src/services.rs @@ -145,6 +145,15 @@ struct CallEndpointRow { client_ip: String, } +/// One turn's proxy role / pair_id + its first call_id, the in-Rust join input +/// for the no-JOIN two-step topology edge builder. Hoisted to module scope so +/// the pure edge builders are unit-testable without a live server. +struct TurnInfo { + proxy_role: String, + pair_id: String, + first_call_id: String, +} + impl ClickHouseBackend { /// "Services" view — aggregate `spans` by `(server_ip, server_port)`. /// Port of the DuckDB `query_services`; see that fn + `StorageBackend:: @@ -492,11 +501,6 @@ impl ClickHouseBackend { .map_err(|e| ch_err("query_services_topology turns", e))?; // Per-turn first call_id + role/pair_id. Skip turns with no calls. - struct TurnInfo { - proxy_role: String, - pair_id: String, - first_call_id: String, - } let mut turn_infos: Vec = Vec::with_capacity(turn_rows.len()); let mut wanted_ids: HashSet = HashSet::new(); for t in turn_rows { @@ -541,53 +545,7 @@ impl ClickHouseBackend { // pair_id non-empty, and from != to (drop dup-capture self-pairs). // Counted by number of (a,b) turn pairs, then aggregated by endpoint // quad — matching DuckDB's COUNT(*) GROUP BY both endpoints. - let mut by_pair_in: HashMap<&str, Vec<(String, u16)>> = HashMap::new(); - let mut by_pair_out: HashMap<&str, Vec<(String, u16)>> = HashMap::new(); - for ti in &turn_infos { - if ti.pair_id.is_empty() { - continue; - } - let ep = match endpoint_by_id.get(&ti.first_call_id) { - Some((ip, port, _client)) => (ip.clone(), *port), - None => continue, - }; - if ti.proxy_role == "proxy_in" { - by_pair_in.entry(ti.pair_id.as_str()).or_default().push(ep); - } else if ti.proxy_role == "proxy_out" { - by_pair_out.entry(ti.pair_id.as_str()).or_default().push(ep); - } - } - // Aggregate (from_ip, from_port, to_ip, to_port) → turn_count, the - // self-join COUNT(*): for each pair_id, every proxy_in × every - // proxy_out is one pairing. - let mut proxy_counts: HashMap<(String, u16, String, u16), u64> = HashMap::new(); - for (pair_id, ins) in &by_pair_in { - if let Some(outs) = by_pair_out.get(pair_id) { - for (fi, fp) in ins { - for (ti, tp) in outs { - // Drop same-endpoint pairs (multi-interface dup capture, - // not a real proxy hop). - if fi == ti && fp == tp { - continue; - } - *proxy_counts - .entry((fi.clone(), *fp, ti.clone(), *tp)) - .or_insert(0) += 1; - } - } - } - } - let proxy_edges: Vec = proxy_counts - .into_iter() - .map(|((fi, fp, ti, tp), c)| TopologyEdge { - from_ip: fi, - from_port: fp, - to_ip: ti, - to_port: tp, - turn_count: c, - kind: "proxy".to_string(), - }) - .collect(); + let proxy_edges = build_proxy_edges(&turn_infos, &endpoint_by_id); // --- Inbound entry edges, grouped by (caller_ip, to_ip, to_port). // DuckDB excludes proxy_out turns (their inbound side is the proxy hop, @@ -730,3 +688,182 @@ impl ClickHouseBackend { Ok(ServicesTopology { nodes, edges }) } } + +/// Build the proxy-hop edges of the service topology in Rust (the no-JOIN +/// equivalent of the DuckDB `turn_endpoint` self-join). For each `pair_id`, +/// every `proxy_in` turn's endpoint × every `proxy_out` turn's endpoint is one +/// pairing, counted and aggregated by the `(from_ip, from_port, to_ip, +/// to_port)` quad — matching DuckDB's `COUNT(*) GROUP BY both endpoints`. +/// Same-endpoint pairs (multi-interface dup capture, not a real proxy hop) are +/// dropped. Turns whose `first_call_id` is not in `endpoint_by_id` are skipped +/// (the call hasn't flushed). Extracted as a pure fn for offline testability. +fn build_proxy_edges( + turn_infos: &[TurnInfo], + endpoint_by_id: &HashMap, +) -> Vec { + let mut by_pair_in: HashMap<&str, Vec<(String, u16)>> = HashMap::new(); + let mut by_pair_out: HashMap<&str, Vec<(String, u16)>> = HashMap::new(); + for ti in turn_infos { + if ti.pair_id.is_empty() { + continue; + } + let ep = match endpoint_by_id.get(&ti.first_call_id) { + Some((ip, port, _client)) => (ip.clone(), *port), + None => continue, + }; + if ti.proxy_role == "proxy_in" { + by_pair_in.entry(ti.pair_id.as_str()).or_default().push(ep); + } else if ti.proxy_role == "proxy_out" { + by_pair_out.entry(ti.pair_id.as_str()).or_default().push(ep); + } + } + // Aggregate (from_ip, from_port, to_ip, to_port) → turn_count, the + // self-join COUNT(*): for each pair_id, every proxy_in × every proxy_out is + // one pairing. + let mut proxy_counts: HashMap<(String, u16, String, u16), u64> = HashMap::new(); + for (pair_id, ins) in &by_pair_in { + if let Some(outs) = by_pair_out.get(pair_id) { + for (fi, fp) in ins { + for (ti, tp) in outs { + // Drop same-endpoint pairs (multi-interface dup capture, + // not a real proxy hop). + if fi == ti && fp == tp { + continue; + } + *proxy_counts + .entry((fi.clone(), *fp, ti.clone(), *tp)) + .or_insert(0) += 1; + } + } + } + } + proxy_counts + .into_iter() + .map(|((fi, fp, ti, tp), c)| TopologyEdge { + from_ip: fi, + from_port: fp, + to_ip: ti, + to_port: tp, + turn_count: c, + kind: "proxy".to_string(), + }) + .collect() +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn valid_sort_fields_is_whitelisted() { + // sort_by is interpolated into ORDER BY, so an unknown field must be + // rejected up front — the whitelist is the gate. Every entry is a + // plain identifier (no injection surface). + for &known in VALID_SORT_FIELDS { + assert!(known.chars().all(|c| c.is_alphanumeric() || c == '_')); + } + assert!(VALID_SORT_FIELDS.contains(&"call_count")); + assert!(VALID_SORT_FIELDS.contains(&"server_port")); + assert!(!VALID_SORT_FIELDS.contains(&"bogus")); + } + + /// Small helper: a turn whose first call resolved to `endpoint`. + fn turn(role: &str, pair_id: &str, first_call_id: &str) -> TurnInfo { + TurnInfo { + proxy_role: role.into(), + pair_id: pair_id.into(), + first_call_id: first_call_id.into(), + } + } + + /// `endpoint_by_id` maps call_id → (server_ip, server_port, client_ip). + fn endpoints(ids: &[(&str, &str, u16, &str)]) -> HashMap { + ids.iter() + .map(|(id, ip, port, client)| (id.to_string(), (ip.to_string(), *port, client.to_string()))) + .collect() + } + + #[test] + fn build_proxy_edges_pairs_in_and_out_per_pair_id() { + // Two proxy_in turns and one proxy_out turn on the same pair_id, on + // distinct endpoints → 2 proxy edges (each in × the single out). + let turns = vec![ + turn("proxy_in", "p1", "c_in1"), + turn("proxy_in", "p1", "c_in2"), + turn("proxy_out", "p1", "c_out"), + ]; + let eps = endpoints(&[("c_in1", "10.0.0.1", 8080, "10.1.0.1"), + ("c_in2", "10.0.0.1", 8080, "10.1.0.2"), + ("c_out", "10.0.0.2", 443, "10.1.0.1")]); + let edges = build_proxy_edges(&turns, &eps); + assert_eq!(edges.len(), 2); + assert!(edges.iter().all(|e| e.kind == "proxy")); + // Both edges point at the proxy_out endpoint. + assert!(edges.iter().all(|e| e.to_ip == "10.0.0.2" && e.to_port == 443)); + // turn_count is 1 each (one (in, out) pairing per in). + assert!(edges.iter().all(|e| e.turn_count == 1)); + } + + #[test] + fn build_proxy_edges_drops_same_endpoint_self_pairs() { + // proxy_in and proxy_out resolve to the SAME endpoint → a dup-capture + // self-pair, dropped (multi-interface capture, not a real hop). + let turns = vec![ + turn("proxy_in", "p1", "c_in"), + turn("proxy_out", "p1", "c_out"), + ]; + let eps = endpoints(&[("c_in", "10.0.0.1", 8080, "10.1.0.1"), + ("c_out", "10.0.0.1", 8080, "10.1.0.1")]); + assert!(build_proxy_edges(&turns, &eps).is_empty()); + } + + #[test] + fn build_proxy_edges_ignores_empty_pair_id_and_unresolved() { + // Empty pair_id → skip. first_call_id not in endpoint_by_id → skip. + let turns = vec![ + turn("proxy_in", "", "c_in"), // empty pair_id + turn("proxy_out", "p1", "c_missing"), // endpoint unresolved + ]; + let eps = endpoints(&[("c_in", "10.0.0.1", 8080, "x")]); + assert!(build_proxy_edges(&turns, &eps).is_empty()); + } + + #[test] + fn build_proxy_edges_aggregates_repeated_endpoint_pairs() { + // Two proxy_in turns on the SAME endpoint × one proxy_out → one edge + // with turn_count = 2 (the (a,b) pair count aggregates). + let turns = vec![ + turn("proxy_in", "p1", "c_in1"), + turn("proxy_in", "p1", "c_in2"), + turn("proxy_out", "p1", "c_out"), + ]; + let eps = endpoints(&[("c_in1", "10.0.0.1", 8080, "x"), + ("c_in2", "10.0.0.1", 8080, "y"), + ("c_out", "10.0.0.2", 443, "z")]); + let edges = build_proxy_edges(&turns, &eps); + assert_eq!(edges.len(), 1); + assert_eq!(edges[0].turn_count, 2); + assert_eq!(edges[0].from_ip, "10.0.0.1"); + assert_eq!(edges[0].from_port, 8080); + assert_eq!(edges[0].to_ip, "10.0.0.2"); + assert_eq!(edges[0].to_port, 443); + } + + #[test] + fn build_proxy_edges_separates_pair_ids() { + // Two distinct pair_ids never cross-pair. + let turns = vec![ + turn("proxy_in", "p1", "a_in"), + turn("proxy_out", "p1", "a_out"), + turn("proxy_in", "p2", "b_in"), + turn("proxy_out", "p2", "b_out"), + ]; + let eps = endpoints(&[("a_in", "1.1.1.1", 1, "x"), ("a_out", "2.2.2.2", 2, "x"), + ("b_in", "3.3.3.3", 3, "x"), ("b_out", "4.4.4.4", 4, "x")]); + let edges = build_proxy_edges(&turns, &eps); + assert_eq!(edges.len(), 2); + // Each pair_id produced exactly one (from,to) edge. + assert!(edges.iter().any(|e| e.from_ip == "1.1.1.1" && e.to_ip == "2.2.2.2")); + assert!(edges.iter().any(|e| e.from_ip == "3.3.3.3" && e.to_ip == "4.4.4.4")); + } +} diff --git a/server/h-storage-clickhouse/src/sql.rs b/server/h-storage-clickhouse/src/sql.rs index f54a096b..74e5db1a 100644 --- a/server/h-storage-clickhouse/src/sql.rs +++ b/server/h-storage-clickhouse/src/sql.rs @@ -34,3 +34,92 @@ pub(crate) fn time_where(col: &str, start_us: i64, end_us: i64) -> String { AND {col} < fromUnixTimestamp64Micro({end_us})" ) } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn escape_str_doubles_backslash_before_quote() { + // The ClickHouse SQL-injection payload: a trailing backslash consumes the + // closing quote (ClickHouse treats `\'` as an escaped quote). Backslash + // must be doubled *before* the quote is doubled so the literal stays closed. + assert_eq!(escape_str(r"\') OR 1=1 --"), r"\\'') OR 1=1 --"); + } + + #[test] + fn escape_str_doubles_single_quotes() { + assert_eq!(escape_str("o'brien"), "o''brien"); + assert_eq!(escape_str("a''b"), "a''''b"); + } + + #[test] + fn escape_str_leaves_like_wildcards() { + // `%` / `_` are intentionally NOT escaped so `LIKE '%x%'` keeps substring + // semantics — the call sites wrap the value in `%...%` themselves. + assert_eq!(escape_str("a%b_c"), "a%b_c"); + } + + #[test] + fn escape_str_round_trips_clean_values() { + assert_eq!(escape_str("call-1"), "call-1"); + assert_eq!(escape_str("server\\path"), "server\\\\path"); + assert_eq!(escape_str(""), ""); + } + + #[test] + fn sql_in_list_quotes_and_comma_joins() { + assert_eq!(sql_in_list(&["a".into(), "b".into()]), "'a', 'b'"); + assert_eq!(sql_in_list(&["only".into()]), "'only'"); + // Empty input yields an empty IN-list body (call sites guard emptiness + // before emitting `IN (...)`, so this is a pure-rendering property). + assert_eq!(sql_in_list(&[]), ""); + } + + #[test] + fn sql_in_list_uses_backslash_aware_escaping() { + // A value containing both a backslash and a quote must stay a single + // closed literal — the ClickHouse-aware escaping differs from the + // backend-neutral (quote-only) `sql_in_list`. + let vals = vec![r"\')".to_string()]; + // ClickHouse: backslash doubled first, then quote doubled → '\\'' wrapped in quotes. + assert_eq!(sql_in_list(&vals), r"'\\'')'"); + // Backend-neutral (DuckDB/Postgres) would NOT double the backslash — + // only the quote is doubled, so the trailing backslash stays lone and + // would (wrongly, for ClickHouse) consume the closing quote. + assert_eq!(h_storage::dialect::sql_in_list(&vals), r"'\'')'"); + assert_ne!(sql_in_list(&vals), h_storage::dialect::sql_in_list(&vals)); + } + + #[test] + fn time_where_is_half_open_with_micro_bounds() { + let s = time_where("request_time", 100, 200); + assert_eq!( + s, + "request_time >= fromUnixTimestamp64Micro(100) \ + AND request_time < fromUnixTimestamp64Micro(200)" + ); + } + + #[test] + fn time_where_interpolates_column_verbatim() { + // The column name is caller-controlled (a constant in every call site); + // it is interpolated verbatim, not escaped — so the literal column name + // appears exactly. + let s = time_where("start_time", -5, 0); + assert!(s.starts_with("start_time >= fromUnixTimestamp64Micro(-5)")); + assert!(s.contains("start_time < fromUnixTimestamp64Micro(0)")); + } + + #[test] + fn time_where_equal_bounds_is_empty_range() { + // Half-open `>= x AND < x` matches nothing — the read-path relies on + // this to express point-adjacency exclusion. + let s = time_where("timestamp", 7, 7); + assert_eq!( + s, + "timestamp >= fromUnixTimestamp64Micro(7) \ + AND timestamp < fromUnixTimestamp64Micro(7)" + ); + } +} diff --git a/server/h-storage-clickhouse/src/turns.rs b/server/h-storage-clickhouse/src/turns.rs index c5d69f3d..d216822b 100644 --- a/server/h-storage-clickhouse/src/turns.rs +++ b/server/h-storage-clickhouse/src/turns.rs @@ -76,6 +76,212 @@ fn now_micros() -> u64 { .unwrap_or(0) } +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn extract_proxy_fields_missing_metadata() { + let (role, peer, peer_ids) = extract_proxy_fields(None); + assert_eq!(role, None); + assert_eq!(peer, None); + assert_eq!(peer_ids, None); + } + + #[test] + fn extract_proxy_fields_non_json_metadata() { + let (role, peer, peer_ids) = extract_proxy_fields(Some("not-json".into())); + assert_eq!(role, None); + assert_eq!(peer, None); + assert_eq!(peer_ids, None); + } + + #[test] + fn extract_proxy_fields_json_without_proxy() { + let (role, peer, peer_ids) = + extract_proxy_fields(Some(r#"{"other":"value"}"#.into())); + assert_eq!(role, None); + assert_eq!(peer, None); + assert_eq!(peer_ids, None); + } + + #[test] + fn extract_proxy_fields_role_only() { + let (role, peer, peer_ids) = + extract_proxy_fields(Some(r#"{"proxy":{"role":"proxy_in"}}"#.into())); + assert_eq!(role.as_deref(), Some("proxy_in")); + assert_eq!(peer, None); + assert_eq!(peer_ids, None); + } + + #[test] + fn extract_proxy_fields_role_and_peer_turn_id() { + let (role, peer, peer_ids) = extract_proxy_fields(Some( + r#"{"proxy":{"role":"proxy_out","peer_turn_id":"turn-42"}}"#.into(), + )); + assert_eq!(role.as_deref(), Some("proxy_out")); + assert_eq!(peer.as_deref(), Some("turn-42")); + assert_eq!(peer_ids, None); + } + + #[test] + fn extract_proxy_fields_peer_turn_ids_array() { + let (role, _peer, peer_ids) = extract_proxy_fields(Some( + r#"{"proxy":{"role":"proxy_in","peer_turn_ids":["turn-1","turn-2"]}}"#.into(), + )); + assert_eq!(role.as_deref(), Some("proxy_in")); + assert_eq!( + peer_ids, + Some(vec!["turn-1".to_string(), "turn-2".to_string()]) + ); + } + + #[test] + fn extract_proxy_fields_handles_sweeper_patch_shape() { + // The real patch the pair sweeper writes via update_trace_metadata: + // {"proxy":{"role":...,"pair_id":...,"peer_turn_id":...,"peer_turn_ids":[...]}}. + let raw = r#"{"proxy":{"role":"proxy_in","pair_id":"pair-7","peer_turn_id":"turn-out","peer_turn_ids":["x","y"]}}"#; + let (role, peer, peer_ids) = extract_proxy_fields(Some(raw.into())); + assert_eq!(role.as_deref(), Some("proxy_in")); + assert_eq!(peer.as_deref(), Some("turn-out")); + assert_eq!(peer_ids, Some(vec!["x".to_string(), "y".to_string()])); + // pair_id is not surfaced by this helper (only role + peer_turn_id[s]). + } + + #[test] + fn extract_proxy_fields_peer_turn_ids_not_array_is_none() { + // A non-array peer_turn_ids (e.g. a string) yields None, not a crash. + let (_role, _peer, peer_ids) = extract_proxy_fields(Some( + r#"{"proxy":{"role":"proxy_in","peer_turn_ids":"oops"}}"#.into(), + )); + assert_eq!(peer_ids, None); + } + + #[test] + fn turn_row_select_lists_span_ids_as_micros() { + // The read-modify-write SELECT must surface the two DateTime64(6) cols + // as i64 micros (via toUnixTimestamp64Micro) so they deserialize into + // TurnRow's i64 fields and re-insert round-trip. This is a compile-time + // invariant of the constant; the test pins the two aliases. + assert!(TURN_ROW_SELECT.contains("toUnixTimestamp64Micro(start_time) AS start_time")); + assert!(TURN_ROW_SELECT.contains("toUnixTimestamp64Micro(end_time) AS end_time")); + // span_ids read as the raw JSON String column (no transform). + assert!(TURN_ROW_SELECT.contains("span_ids")); + // _version read back so it can be bumped on re-insert. + assert!(TURN_ROW_SELECT.contains("_version")); + } + + #[test] + fn traces_valid_sort_fields_is_whitelisted() { + for &known in TRACES_VALID_SORT_FIELDS { + assert!(known.chars().all(|c| c.is_alphanumeric() || c == '_')); + } + assert!(TRACES_VALID_SORT_FIELDS.contains(&"start_time")); + assert!(TRACES_VALID_SORT_FIELDS.contains(&"call_count")); + assert!(!TRACES_VALID_SORT_FIELDS.contains(&"bogus")); + } + + fn traces_query() -> TracesQuery { + TracesQuery { + time_range: TimeRange { start_us: 100, end_us: 200 }, + filter: DimensionFilter::default(), + client_ips: vec![], + server_ports: vec![], + statuses: vec![], + agent_kinds: vec![], + sort_by: "start_time".into(), + sort_order: "desc".into(), + page: 1, + page_size: 10, + include_proxy_hops: false, + } + } + + #[test] + fn traces_where_sql_default_hides_proxy_hops() { + // include_proxy_hops = false (default) appends the proxy exclusion. + let s = traces_where_sql(&traces_query()); + assert!(s.starts_with("start_time >= fromUnixTimestamp64Micro(100)")); + assert!(s.contains("start_time < fromUnixTimestamp64Micro(200)")); + assert!(s.contains("NOT IN ('proxy_out', 'mirror_secondary')")); + } + + #[test] + fn traces_where_sql_include_proxy_hops_omits_exclusion() { + let q = TracesQuery { + include_proxy_hops: true, + ..traces_query() + }; + assert!(!traces_where_sql(&q).contains("NOT IN ('proxy_out'")); + } + + #[test] + fn traces_where_sql_models_uses_hasany_json_extract() { + let q = TracesQuery { + filter: DimensionFilter { + models: vec!["gpt-4".into()], + ..Default::default() + }, + ..traces_query() + }; + let s = traces_where_sql(&q); + // models_used is a JSON-array String → hasAny(JSONExtract(..., 'Array(String)'), [...]). + assert!(s.contains("hasAny(JSONExtract(coalesce(models_used, '[]'), 'Array(String)'), ['gpt-4'])")); + } + + #[test] + fn traces_where_sql_server_ports_uses_in_subquery_not_join() { + // traces has no server_port → resolve the turn's first call_id against + // spans via an uncorrelated IN-subquery (NOT a JOIN). Assert the shape + // and that no literal "JOIN" keyword is introduced. + let q = TracesQuery { + server_ports: vec![8080, 443], + ..traces_query() + }; + let s = traces_where_sql(&q); + assert!(s.contains("arrayElement(JSONExtract(span_ids, 'Array(String)'), 1) IN")); + assert!(s.contains("SELECT id FROM spans WHERE server_port IN (8080, 443)")); + assert!(!s.to_lowercase().contains(" join ")); + assert!(!s.contains(" JOIN ")); + } + + #[test] + fn traces_where_sql_combines_dimension_filters() { + let q = TracesQuery { + filter: DimensionFilter { + wire_apis: vec!["openai-chat".into()], + server_ips: vec!["10.0.0.2".into()], + ..Default::default() + }, + statuses: vec!["complete".into()], + agent_kinds: vec!["claude-cli".into()], + client_ips: vec!["10.0.0.1".into()], + ..traces_query() + }; + let s = traces_where_sql(&q); + assert!(s.contains("wire_api IN ('openai-chat')")); + assert!(s.contains("server_ip IN ('10.0.0.2')")); + assert!(s.contains("status IN ('complete')")); + assert!(s.contains("agent_kind IN ('claude-cli')")); + assert!(s.contains("client_ip IN ('10.0.0.1')")); + assert!(!s.starts_with(" AND")); + assert!(!s.ends_with("AND ")); + } + + #[test] + fn traces_where_sql_escapes_user_lists() { + let q = TracesQuery { + filter: DimensionFilter { + wire_apis: vec!["a'b".into()], + ..Default::default() + }, + ..traces_query() + }; + // A quote in a wire_api value is doubled (no breakout). + assert!(traces_where_sql(&q).contains("wire_api IN ('a''b')")); + } +} + #[derive(Row, Deserialize)] struct TurnListRow { turn_id: String, @@ -159,6 +365,78 @@ struct CountRow { n: u64, } +/// Valid `sort_by` fields for `query_traces`. Hoisted to module scope so the +/// reject-unknown-sort path is unit-testable without a live client (the value +/// is interpolated into `ORDER BY`). Mirrors the DuckDB whitelist. +const TRACES_VALID_SORT_FIELDS: &[&str] = &[ + "start_time", + "end_time", + "duration_ms", + "call_count", + "total_input_tokens", + "total_output_tokens", +]; + +/// Build the `query_traces` WHERE clause: a half-open `start_time` time range +/// AND-ed with every present dimension + per-call filter. Extracted as a pure +/// fn so the escaping / IN-list / JSON-array / IN-subquery / proxy-hop +/// assembly is unit-testable without a live server. The `server_ports` filter +/// uses an uncorrelated IN-subquery (NOT a JOIN) because `traces` carries no +/// `server_port`; the proxy-hop exclusion hides sweeper-folded hops. +pub(crate) fn traces_where_sql(query: &TracesQuery) -> String { + let mut where_parts = vec![time_where( + "start_time", + query.time_range.start_us, + query.time_range.end_us, + )]; + if !query.filter.wire_apis.is_empty() { + where_parts.push(format!("wire_api IN ({})", sql_in_list(&query.filter.wire_apis))); + } + if !query.filter.models.is_empty() { + // models_used is a JSON-array String; match if any requested model is + // present (DuckDB list_has_any → ClickHouse hasAny). + where_parts.push(format!( + "hasAny(JSONExtract(coalesce(models_used, '[]'), 'Array(String)'), [{}])", + sql_in_list(&query.filter.models) + )); + } + if !query.statuses.is_empty() { + where_parts.push(format!("status IN ({})", sql_in_list(&query.statuses))); + } + if !query.agent_kinds.is_empty() { + where_parts.push(format!("agent_kind IN ({})", sql_in_list(&query.agent_kinds))); + } + if !query.client_ips.is_empty() { + where_parts.push(format!("client_ip IN ({})", sql_in_list(&query.client_ips))); + } + if !query.server_ports.is_empty() { + // traces has no server_port; resolve via the turn's first call_id + // against spans. ClickHouse can't do the DuckDB correlated EXISTS, so + // use an uncorrelated IN-subquery (still not a JOIN): the turn's first + // call_id ∈ { calls on those ports }. + let ports: Vec = query.server_ports.iter().map(|p| p.to_string()).collect(); + where_parts.push(format!( + "arrayElement(JSONExtract(span_ids, 'Array(String)'), 1) IN \ + (SELECT id FROM spans WHERE server_port IN ({}))", + ports.join(", ") + )); + } + if !query.filter.server_ips.is_empty() { + where_parts.push(format!("server_ip IN ({})", sql_in_list(&query.filter.server_ips))); + } + if !query.include_proxy_hops { + // Hide the sweeper-folded hops. JSONExtractString returns '' when + // absent, and '' NOT IN (...) is true, so direct turns + + // proxy_in/mirror_primary stay visible. + where_parts.push( + "JSONExtractString(coalesce(metadata, ''), 'proxy', 'role') \ + NOT IN ('proxy_out', 'mirror_secondary')" + .to_string(), + ); + } + where_parts.join(" AND ") +} + impl ClickHouseBackend { pub(crate) async fn write_traces(&self, turns: Vec) -> Result<()> { let rows: Vec = turns.into_iter().map(TurnRow::from).collect(); @@ -167,77 +445,15 @@ impl ClickHouseBackend { } pub(crate) async fn query_traces(&self, query: &TracesQuery) -> Result { - const VALID_SORT_FIELDS: &[&str] = &[ - "start_time", - "end_time", - "duration_ms", - "call_count", - "total_input_tokens", - "total_output_tokens", - ]; - if !VALID_SORT_FIELDS.contains(&query.sort_by.as_str()) { + if !TRACES_VALID_SORT_FIELDS.contains(&query.sort_by.as_str()) { return Err(AppError::Storage(format!( "invalid sort_by field: {}", query.sort_by ))); } - let sort_order = if query.sort_order.eq_ignore_ascii_case("ASC") { - "ASC" - } else { - "DESC" - }; + let sort_order = crate::calls::resolve_sort_order(&query.sort_order); - let mut where_parts = vec![time_where( - "start_time", - query.time_range.start_us, - query.time_range.end_us, - )]; - if !query.filter.wire_apis.is_empty() { - where_parts.push(format!("wire_api IN ({})", sql_in_list(&query.filter.wire_apis))); - } - if !query.filter.models.is_empty() { - // models_used is a JSON-array String; match if any requested model - // is present (DuckDB list_has_any → ClickHouse hasAny). - where_parts.push(format!( - "hasAny(JSONExtract(coalesce(models_used, '[]'), 'Array(String)'), [{}])", - sql_in_list(&query.filter.models) - )); - } - if !query.statuses.is_empty() { - where_parts.push(format!("status IN ({})", sql_in_list(&query.statuses))); - } - if !query.agent_kinds.is_empty() { - where_parts.push(format!("agent_kind IN ({})", sql_in_list(&query.agent_kinds))); - } - if !query.client_ips.is_empty() { - where_parts.push(format!("client_ip IN ({})", sql_in_list(&query.client_ips))); - } - if !query.server_ports.is_empty() { - // traces has no server_port; resolve via the turn's first - // call_id against spans. ClickHouse can't do the DuckDB - // correlated EXISTS, so use an uncorrelated IN-subquery (still - // not a JOIN): the turn's first call_id ∈ { calls on those ports }. - let ports: Vec = query.server_ports.iter().map(|p| p.to_string()).collect(); - where_parts.push(format!( - "arrayElement(JSONExtract(span_ids, 'Array(String)'), 1) IN \ - (SELECT id FROM spans WHERE server_port IN ({}))", - ports.join(", ") - )); - } - if !query.filter.server_ips.is_empty() { - where_parts.push(format!("server_ip IN ({})", sql_in_list(&query.filter.server_ips))); - } - if !query.include_proxy_hops { - // Hide the sweeper-folded hops. JSONExtractString returns '' when - // absent, and '' NOT IN (...) is true, so direct turns + - // proxy_in/mirror_primary stay visible. - where_parts.push( - "JSONExtractString(coalesce(metadata, ''), 'proxy', 'role') \ - NOT IN ('proxy_out', 'mirror_secondary')" - .to_string(), - ); - } - let where_sql = where_parts.join(" AND "); + let where_sql = traces_where_sql(query); let total = self .client