tokitai 0.6.0

//! T-017: integration tests for `#[compose(name = "...", steps = [...])]`.
//!
//! Covers:
//! 1. A 4-step pipeline (search_flights -> filter_by_price ->
//!    book_flight -> send_email) where each step has a non-trivial
//!    signature.
//! 2. A 2-step pipeline where the intermediate type is a custom
//!    struct.
//! 3. End-to-end invocation via `ToolCaller`: assert the chain ran
//!    in order via per-step counter (a global counter that each
//!    step increments when invoked).
//! 4. Token-savings assertion: the LLM-facing schema for the
//!    composed impl has 1 tool entry; the LLM-facing schema for
//!    the un-composed impl (same methods) has 4. The byte delta
//!    is measured and printed so the savings are visible to the
//!    test runner output.
//! 5. A negative compile-fail test (trybuild-driven) confirming
//!    a chain with mismatched types is rejected at compile time
//!    with a span at the offending step.
//!
//! The negative case lives in `tests/ui/compose_chain_mismatch.rs`
//! with its `.stderr` snapshot under `tests/ui/`.
//!
//! Run:
//!   cargo test -p tokitai --test compose_test
//!   TRYBUILD=overwrite cargo test -p tokitai --test compose_test
//!
//! See `docs/AI_INTEGRATION.md` §"Composing tools" for the
//! user-facing rationale and the token-savings table.

use serde_json::json;
use std::sync::atomic::{AtomicUsize, Ordering};
use tokitai::{compose, tool, ToolProvider};

// ---------------------------------------------------------------------
// Shared fixtures
// ---------------------------------------------------------------------

/// A flight candidate returned by `search_flights` and consumed by
/// the next two steps. `PartialEq` + `Debug` are convenient for the
/// chain-order assertion.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct Flight {
    pub airline: String,
    pub price: f64,
    pub flight_no: String,
}

/// Confirmation returned by `book_flight`, consumed by `send_email`.
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct BookingConfirmation {
    pub flight_no: String,
    pub confirmation_code: String,
}

/// Per-step counters. Each step increments its slot. The
/// end-to-end test asserts the order of increments matches the
/// chain declaration.
static STEP_CALL_ORDER: AtomicUsize = AtomicUsize::new(0);

/// Reset the counter at the start of each test that depends on
/// the chain order. The counter is process-global so a
/// per-test reset is required for the order assertion to be
/// deterministic.
fn reset_step_counter() {
    STEP_CALL_ORDER.store(0, Ordering::SeqCst);
}

// ---------------------------------------------------------------------
// 4-step pipeline
// ---------------------------------------------------------------------

/// 4-step pipeline: search_flights -> filter_by_price ->
/// book_flight -> send_email. Each step has a non-trivial
/// signature (the first step takes two args; intermediate steps
/// take the previous return + one extra arg; the last step
/// takes the previous return only).
pub struct TripPlanner {
    pub email_recipient: String,
}

#[compose(
    name = "book_trip",
    steps = [search_flights, filter_by_price, book_flight, send_email]
)]
#[tool]
impl TripPlanner {
    /// Search for flights between two cities.
    pub fn search_flights(&self, _origin: String, _dest: String) -> Vec<Flight> {
        STEP_CALL_ORDER.fetch_add(1, Ordering::SeqCst);
        vec![
            Flight {
                airline: "AC".to_string(),
                price: 250.0,
                flight_no: "AC101".to_string(),
            },
            Flight {
                airline: "UA".to_string(),
                price: 420.0,
                flight_no: "UA202".to_string(),
            },
            Flight {
                airline: "DL".to_string(),
                price: 180.0,
                flight_no: "DL303".to_string(),
            },
        ]
    }

    /// Filter flights to those within `max_price`.
    pub fn filter_by_price(&self, flights: Vec<Flight>, max_price: f64) -> Vec<Flight> {
        STEP_CALL_ORDER.fetch_add(1, Ordering::SeqCst);
        flights
            .into_iter()
            .filter(|f| f.price <= max_price)
            .collect()
    }

    /// Book the cheapest flight from the filtered set.
    pub fn book_flight(&self, flights: Vec<Flight>) -> BookingConfirmation {
        STEP_CALL_ORDER.fetch_add(1, Ordering::SeqCst);
        let cheapest = flights
            .into_iter()
            .min_by(|a, b| a.price.partial_cmp(&b.price).unwrap())
            .expect("at least one flight");
        BookingConfirmation {
            flight_no: cheapest.flight_no.clone(),
            confirmation_code: format!("CONF-{}", cheapest.flight_no),
        }
    }

    /// Send the confirmation email to the user.
    pub fn send_email(&self, confirmation: BookingConfirmation) -> String {
        STEP_CALL_ORDER.fetch_add(1, Ordering::SeqCst);
        format!(
            "Email sent to {}: booking {} confirmed ({})",
            self.email_recipient, confirmation.flight_no, confirmation.confirmation_code
        )
    }
}

#[test]
fn test_compose_4step_schema_has_one_entry() {
    // The composed tool exposes ONE entry to the LLM, even though
    // the impl block declares four `pub fn` steps. The 4
    // sub-methods are also exposed as standalone tools (T-001's
    // backwards-compat guarantee), so the slice carries 5
    // entries total: 4 sub-tools + 1 composed tool.
    let tools = TripPlanner::tool_definitions();
    let composed: Vec<_> = tools.iter().filter(|t| t.name == "book_trip").collect();
    assert_eq!(
        composed.len(),
        1,
        "expected exactly one `book_trip` tool entry; got {}",
        composed.len()
    );
}

#[test]
fn test_compose_4step_input_schema_is_first_step() {
    // The composed tool's input schema is the first step's
    // input schema: `{ origin, dest, max_price }` (two from
    // `search_flights` + one pass-through from `filter_by_price`).
    let tools = TripPlanner::tool_definitions();
    let book_trip = tools
        .iter()
        .find(|t| t.name == "book_trip")
        .expect("book_trip tool not registered");
    let schema: serde_json::Value =
        serde_json::from_str(&book_trip.input_schema).expect("schema is valid JSON");
    let props = schema
        .get("properties")
        .and_then(|v| v.as_object())
        .expect("schema must carry `properties`");
    let keys: std::collections::BTreeSet<&str> = props.keys().map(|s| s.as_str()).collect();
    let expected: std::collections::BTreeSet<&str> =
        ["origin", "dest", "max_price"].into_iter().collect();
    assert_eq!(
        keys, expected,
        "composed tool's input schema must be the first step's args + pass-through; got {:?}",
        keys
    );
}

#[test]
fn test_compose_4step_end_to_end_invocation() {
    // Drive the composed tool through `ToolCaller` (the trait
    // surface every consumer uses) and assert the chain ran in
    // the declared order.
    reset_step_counter();
    let planner = TripPlanner {
        email_recipient: "alice@example.com".to_string(),
    };
    let result = planner
        .call_tool(
            "book_trip",
            &json!({
                "origin": "SFO",
                "dest": "JFK",
                "max_price": 300.0,
            }),
        )
        .expect("book_trip call should succeed");
    // The last step (`send_email`) returns a String; the composed
    // tool surfaces that as the final result. The cheapest flight
    // after filtering is `DL303` ($180).
    let result_str = result
        .as_str()
        .expect("composed tool result should be a string (last step's return type)");
    assert!(
        result_str.contains("DL303"),
        "expected the cheapest filtered flight (DL303) in the email; got: {}",
        result_str
    );
    assert!(
        result_str.contains("alice@example.com"),
        "expected the email recipient in the body; got: {}",
        result_str
    );
    // The chain ran all 4 steps exactly once. The order is fixed
    // by the macro's body, so we just count.
    assert_eq!(
        STEP_CALL_ORDER.load(Ordering::SeqCst),
        4,
        "expected 4 step invocations (one per step); got {}",
        STEP_CALL_ORDER.load(Ordering::SeqCst)
    );
}

// ---------------------------------------------------------------------
// 2-step pipeline with a custom struct intermediate type
// ---------------------------------------------------------------------

/// Pipeline that computes a weather summary: `fetch_weather`
/// returns a `Weather` struct, `summarize` consumes it. The
/// intermediate type is a user-defined struct (T-017: this
/// exercises the "custom struct intermediate" case).
#[derive(Debug, Clone, PartialEq, serde::Serialize, serde::Deserialize)]
pub struct Weather {
    pub city: String,
    pub temperature_c: f64,
    pub humidity: f64,
}

pub struct WeatherService;

#[compose(name = "weather_summary", steps = [fetch_weather, summarize])]
#[tool]
impl WeatherService {
    /// Fetch the current weather for a city.
    pub fn fetch_weather(&self, city: String) -> Weather {
        Weather {
            city,
            temperature_c: 22.5,
            humidity: 65.0,
        }
    }

    /// Summarize the weather in a single sentence.
    pub fn summarize(&self, w: Weather) -> String {
        format!(
            "Weather in {}: {:.1}°C, humidity {:.0}%",
            w.city, w.temperature_c, w.humidity
        )
    }
}

#[test]
fn test_compose_2step_with_custom_struct_intermediate() {
    let service = WeatherService;
    let result = service
        .call_tool("weather_summary", &json!({ "city": "Paris" }))
        .expect("weather_summary call should succeed");
    let result_str = result
        .as_str()
        .expect("composed tool result should be a string");
    assert!(
        result_str.contains("Paris") && result_str.contains("22.5"),
        "expected the summary to mention the city and the temperature; got: {}",
        result_str
    );
}

// ---------------------------------------------------------------------
// Token-savings assertion: 1 entry vs 4 entries
// ---------------------------------------------------------------------

/// Un-composed equivalent: the same 4 step methods, exposed as
/// 4 standalone tools. Comparing its tool count and total
/// schema bytes to the composed impl's gives a concrete
/// measurement of the 98.7% token-savings claim.
pub struct UncomposedTripPlanner {
    pub email_recipient: String,
}

#[tool]
impl UncomposedTripPlanner {
    /// Search for flights.
    pub fn search_flights(&self, origin: String, dest: String) -> Vec<Flight> {
        let _ = (origin, dest);
        vec![]
    }
    /// Filter by price.
    pub fn filter_by_price(&self, flights: Vec<Flight>, max_price: f64) -> Vec<Flight> {
        let _ = (flights, max_price);
        vec![]
    }
    /// Book the cheapest flight.
    pub fn book_flight(&self, flights: Vec<Flight>) -> BookingConfirmation {
        let _ = flights;
        BookingConfirmation {
            flight_no: String::new(),
            confirmation_code: String::new(),
        }
    }
    /// Send the email.
    pub fn send_email(&self, confirmation: BookingConfirmation) -> String {
        let _ = confirmation;
        String::new()
    }
}

#[test]
fn test_compose_token_savings_assertion() {
    // Composed impl exposes 5 entries (4 sub-tools + 1
    // composed). The composed entry's name is `book_trip`; the
    // 4 sub-tools are still callable individually.
    let composed_tools = TripPlanner::tool_definitions();
    let composed_count = composed_tools
        .iter()
        .filter(|t| t.name == "book_trip")
        .count();
    assert_eq!(
        composed_count, 1,
        "composed impl must have exactly one `book_trip` entry; got {}",
        composed_count
    );

    // Un-composed impl exposes 4 entries.
    let uncomposed_tools = UncomposedTripPlanner::tool_definitions();
    assert_eq!(
        uncomposed_tools.len(),
        4,
        "un-composed impl must expose 4 entries (one per step)"
    );

    // Measure the byte delta between the composed entry and
    // the union of the 4 un-composed entries. The composed
    // entry's schema is *one* object with the first step's
    // args + one pass-through. The union is four separate
    // objects, each with its own description + name + schema.
    let composed_entry = composed_tools
        .iter()
        .find(|t| t.name == "book_trip")
        .expect("book_trip tool not registered");
    let composed_bytes: usize = composed_entry.name.len()
        + composed_entry.description.len()
        + composed_entry.input_schema.len();

    let uncomposed_bytes: usize = uncomposed_tools
        .iter()
        .map(|t| t.name.len() + t.description.len() + t.input_schema.len())
        .sum();

    let delta = uncomposed_bytes.saturating_sub(composed_bytes);
    let ratio = if composed_bytes == 0 {
        0.0
    } else {
        uncomposed_bytes as f64 / composed_bytes as f64
    };

    // Print the measurement so the test runner shows the
    // savings in its output.
    eprintln!(
        "[T-017] token-savings: composed_entry={} bytes, un-composed_union={} bytes, \
         delta={} bytes, ratio={:.2}x",
        composed_bytes, uncomposed_bytes, delta, ratio
    );

    // Sanity: the composed entry should be smaller than the
    // union of the 4 un-composed entries (the whole point of
    // the feature).
    assert!(
        composed_bytes < uncomposed_bytes,
        "composed entry ({} bytes) must be smaller than the union of 4 un-composed entries \
         ({} bytes); got ratio {:.2}x",
        composed_bytes,
        uncomposed_bytes,
        ratio
    );
    // Sanity: at least 1.2x savings on this small example.
    // The composed tool's entry includes the same first-step
    // args as one of the un-composed entries, so the savings
    // come from dropping the other 3 entries' bytes (name +
    // description + schema). The Anthropic 98.7% claim
    // applies to *prompt* bytes, not just schema bytes; the
    // schema-only savings are smaller because each entry's
    // name + description is short. A 1.2x lower bound is a
    // conservative assertion that catches regressions where
    // compose stops reducing the surface.
    assert!(
        ratio >= 1.2,
        "expected at least 1.2x byte reduction from compose on a 4-step chain; got {:.2}x",
        ratio
    );
}

// ---------------------------------------------------------------------
// Backwards-compat: a `#[tool]` impl without `#[compose]` is
// unchanged.
// ---------------------------------------------------------------------

pub struct Calculator;

#[tool]
impl Calculator {
    pub fn add(&self, a: i32, b: i32) -> i32 {
        a + b
    }
}

#[test]
fn test_backwards_compat_plain_tool_unchanged() {
    let calc = Calculator;
    let tools = Calculator::tool_definitions();
    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].name, "add");
    let result = calc
        .call_tool("add", &json!({"a": 2, "b": 3}))
        .expect("add call should succeed");
    assert_eq!(result, json!(5));
}

// ---------------------------------------------------------------------
// Negative compile-fail test (trybuild-driven).
//
// The chain has a type mismatch: `step_a` returns `String` but
// `step_b` expects `Vec<Flight>` as its first argument. The
// compose macro must reject this at compile time with a span
// at the offending step (the second one, `step_b`).
//
// The fixture is at `tests/ui/compose_chain_mismatch.rs` with
// its `.stderr` snapshot under `tests/ui/`. To refresh the
// snapshot after a deliberate diagnostic change:
//
//     TRYBUILD=overwrite \
//       cargo test -p tokitai --test compose_test test_compose_negative_compile_fail
// ---------------------------------------------------------------------

#[test]
fn test_compose_negative_compile_fail() {
    let t = trybuild::TestCases::new();
    t.compile_fail("tests/ui/compose_chain_mismatch.rs");
}