roboticus-api 0.11.3

//! Helpers for annotating `PipelineTrace` spans with subsystem metrics.
//!
//! Each `annotate_*` function follows the same contract:
//! - It writes into the currently open span (if any) on the provided trace.
//! - It is a no-op when no span is open (safe to call unconditionally).
//! - It uses canonical namespace constants from `pipeline_trace::ns`.

use super::super::pipeline_trace::{PipelineTrace, ns};

/// Annotate the currently open trace span with memory retrieval metrics.
///
/// Keys written (all under the `retrieval.*` namespace):
/// - `retrieval.avg_similarity`   — average cosine similarity across results
/// - `retrieval.budget_utilization` — fraction of context budget used by memory tokens
/// - `retrieval.retrieval_count`  — total memories retrieved across all tiers
/// - `retrieval.retrieval_hit`    — whether any memories were retrieved
/// - `retrieval.tier_breakdown`   — per-tier counts (working/episodic/semantic/procedural/relationship)
pub(super) fn annotate_retrieval_metrics(
    trace: &mut PipelineTrace,
    metrics: &roboticus_agent::retrieval::RetrievalMetrics,
) {
    trace.annotate_ns(
        ns::RETRIEVAL,
        "avg_similarity",
        serde_json::json!(metrics.avg_similarity),
    );
    trace.annotate_ns(
        ns::RETRIEVAL,
        "budget_utilization",
        serde_json::json!(metrics.budget_utilization),
    );
    trace.annotate_ns(
        ns::RETRIEVAL,
        "retrieval_count",
        serde_json::json!(metrics.retrieval_count),
    );
    trace.annotate_ns(
        ns::RETRIEVAL,
        "retrieval_hit",
        serde_json::json!(metrics.retrieval_hit),
    );
    trace.annotate_ns(
        ns::RETRIEVAL,
        "tier_breakdown",
        serde_json::json!({
            "working":      metrics.tiers.working,
            "episodic":     metrics.tiers.episodic,
            "semantic":     metrics.tiers.semantic,
            "procedural":   metrics.tiers.procedural,
            "relationship": metrics.tiers.relationship,
        }),
    );
}

/// Annotate the currently open trace span with tool search metrics.
///
/// Keys written (all under the `tool_search.*` namespace):
/// - `tool_search.candidates_considered` — total tools evaluated
/// - `tool_search.candidates_selected`   — tools kept after ranking + budget
/// - `tool_search.candidates_pruned`     — tools removed
/// - `tool_search.token_savings`         — tokens saved by pruning
/// - `tool_search.top_scores`            — top-10 selected tools with scores
/// - `tool_search.embedding_status`      — "ok" or "failed"
pub(super) fn annotate_tool_search(
    trace: &mut PipelineTrace,
    stats: &roboticus_agent::tool_search::ToolSearchStats,
) {
    trace.annotate_ns(
        ns::TOOL_SEARCH,
        "candidates_considered",
        serde_json::json!(stats.candidates_considered),
    );
    trace.annotate_ns(
        ns::TOOL_SEARCH,
        "candidates_selected",
        serde_json::json!(stats.candidates_selected),
    );
    trace.annotate_ns(
        ns::TOOL_SEARCH,
        "candidates_pruned",
        serde_json::json!(stats.candidates_pruned),
    );
    trace.annotate_ns(
        ns::TOOL_SEARCH,
        "token_savings",
        serde_json::json!(stats.token_savings),
    );
    if !stats.top_scores.is_empty() {
        let scores_map: serde_json::Map<String, serde_json::Value> = stats
            .top_scores
            .iter()
            .map(|(name, score)| (name.clone(), serde_json::json!(score)))
            .collect();
        trace.annotate_ns(
            ns::TOOL_SEARCH,
            "top_scores",
            serde_json::Value::Object(scores_map),
        );
    }
    trace.annotate_ns(
        ns::TOOL_SEARCH,
        "embedding_status",
        serde_json::json!(stats.embedding_status),
    );
}

/// Annotate the currently open trace span with MCP tool call metrics.
///
/// Keys written (all under the `mcp.*` namespace):
/// - `mcp.server`      — name of the MCP server
/// - `mcp.tool`        — name of the tool called
/// - `mcp.duration_ms` — round-trip time in milliseconds
/// - `mcp.success`     — whether the call succeeded
#[cfg_attr(not(test), allow(dead_code))]
pub(super) fn annotate_mcp_call(
    trace: &mut PipelineTrace,
    server: &str,
    tool: &str,
    duration_ms: u64,
    success: bool,
) {
    trace.annotate_ns(ns::MCP, "server", serde_json::json!(server));
    trace.annotate_ns(ns::MCP, "tool", serde_json::json!(tool));
    trace.annotate_ns(ns::MCP, "duration_ms", serde_json::json!(duration_ms));
    trace.annotate_ns(ns::MCP, "success", serde_json::json!(success));
}

#[cfg(test)]
mod tests {
    use super::super::super::pipeline_trace::{PipelineTrace, SpanOutcome};
    use super::*;
    use roboticus_agent::retrieval::{MemoryTierBreakdown, RetrievalMetrics};

    #[test]
    fn annotate_retrieval_metrics_writes_all_keys() {
        let mut trace = PipelineTrace::new("turn-1", "api");
        trace.begin_stage("retrieval");

        let metrics = RetrievalMetrics {
            retrieval_count: 5,
            retrieval_hit: true,
            avg_similarity: 0.75,
            budget_utilization: 0.42,
            tiers: MemoryTierBreakdown {
                working: 1,
                episodic: 2,
                semantic: 1,
                procedural: 0,
                relationship: 1,
            },
        };

        annotate_retrieval_metrics(&mut trace, &metrics);
        trace.end_stage(SpanOutcome::Ok);

        let span = &trace.stages[0];
        assert_eq!(
            span.annotations.get("retrieval.avg_similarity"),
            Some(&serde_json::json!(0.75))
        );
        assert_eq!(
            span.annotations.get("retrieval.retrieval_count"),
            Some(&serde_json::json!(5_usize))
        );
        assert_eq!(
            span.annotations.get("retrieval.retrieval_hit"),
            Some(&serde_json::json!(true))
        );
        let breakdown = span
            .annotations
            .get("retrieval.tier_breakdown")
            .expect("tier_breakdown missing");
        assert_eq!(breakdown["working"], serde_json::json!(1_usize));
        assert_eq!(breakdown["episodic"], serde_json::json!(2_usize));
    }

    #[test]
    fn annotate_tool_search_writes_all_keys() {
        let mut trace = PipelineTrace::new("turn-2", "api");
        trace.begin_stage("tool_selection");

        let stats = roboticus_agent::tool_search::ToolSearchStats {
            candidates_considered: 25,
            candidates_selected: 12,
            candidates_pruned: 13,
            token_savings: 1200,
            top_scores: vec![("bash".into(), 0.95), ("memory_store".into(), 0.87)],
            embedding_status: "ok".into(),
        };

        annotate_tool_search(&mut trace, &stats);
        trace.end_stage(SpanOutcome::Ok);

        let span = &trace.stages[0];
        assert_eq!(
            span.annotations.get("tool_search.candidates_considered"),
            Some(&serde_json::json!(25_usize))
        );
        assert_eq!(
            span.annotations.get("tool_search.candidates_pruned"),
            Some(&serde_json::json!(13_usize))
        );
        assert_eq!(
            span.annotations.get("tool_search.token_savings"),
            Some(&serde_json::json!(1200_usize))
        );
    }

    #[test]
    fn annotate_mcp_call_writes_all_keys() {
        let mut trace = PipelineTrace::new("turn-3", "api");
        trace.begin_stage("inference");

        annotate_mcp_call(&mut trace, "github", "create_issue", 350, true);
        trace.end_stage(SpanOutcome::Ok);

        let span = &trace.stages[0];
        assert_eq!(
            span.annotations.get("mcp.server"),
            Some(&serde_json::json!("github"))
        );
        assert_eq!(
            span.annotations.get("mcp.tool"),
            Some(&serde_json::json!("create_issue"))
        );
        assert_eq!(
            span.annotations.get("mcp.duration_ms"),
            Some(&serde_json::json!(350_u64))
        );
        assert_eq!(
            span.annotations.get("mcp.success"),
            Some(&serde_json::json!(true))
        );
    }
}