crank/apps/mcp-server/tests/onboarding.rs

#[path = "integration/common.rs"]
mod common;

use std::time::Duration;

use crank_core::{
    ExecutionStage, InvocationLevel, InvocationLog, InvocationLogId, InvocationSource,
    InvocationStatus, OutcomeCertainty, PlatformApiKeyScope, ProductEventKind, Retryability,
};
use crank_registry::{
    CreateInvocationLogRequest, InvocationHistoryWriteOutcome, ListInvocationLogsQuery,
    ListProductEventsQuery, PublishRequest,
};
use serde_json::json;
use time::{OffsetDateTime, format_description::well_known::Rfc3339};

use common::{
    agent_mcp_url, build_test_app, create_platform_api_key, initialize_session, post_jsonrpc,
    publish_agent_for_operation, spawn_mcp_server, spawn_upstream_server, test_operation,
    test_registry, test_workspace_id,
};

#[tokio::test]
async fn onboarding_completes_only_after_exact_key_successful_public_tool_call() {
    let registry = test_registry().await;
    let upstream_base_url = spawn_upstream_server().await;
    let operation = test_operation(&upstream_base_url, "onboarding_first_call");
    let mut operation_draft = operation.clone();
    operation_draft.status = crank_core::OperationStatus::Draft;
    operation_draft.published_at = None;

    registry
        .create_operation(
            &test_workspace_id(),
            &operation_draft,
            Some("onboarding-test"),
        )
        .await
        .unwrap();
    let admin_test = InvocationLog {
        id: InvocationLogId::new("log_onboarding_admin_test"),
        workspace_id: test_workspace_id(),
        agent_id: None,
        platform_api_key_id: None,
        operation_id: operation.id.clone(),
        operation_version: Some(1),
        source: InvocationSource::AdminTestRun,
        level: InvocationLevel::Info,
        status: InvocationStatus::Ok,
        tool_name: operation.name.clone(),
        message: "admin test succeeded".to_owned(),
        request_id: Some("req_onboarding_admin_test".to_owned()),
        trace_id: Some("4bf92f3577b34da6a3ce929d0e0e4736".to_owned()),
        status_code: Some(200),
        duration_ms: 1,
        error_kind: None,
        execution_stage: Some(ExecutionStage::Runtime),
        execution_error_code: None,
        retryability: Some(Retryability::Never),
        outcome_certainty: Some(OutcomeCertainty::Certain),
        request_preview: json!({}),
        response_preview: json!({"ok": true}),
        created_at: OffsetDateTime::parse("2026-08-23T07:59:00Z", &Rfc3339).unwrap(),
    };
    assert_eq!(
        registry
            .create_invocation_log(CreateInvocationLogRequest { log: &admin_test })
            .await,
        InvocationHistoryWriteOutcome::Recorded
    );
    registry
        .publish_operation(PublishRequest {
            workspace_id: &test_workspace_id(),
            operation_id: &operation.id,
            version: 1,
            published_at: &OffsetDateTime::parse("2026-08-23T08:00:00Z", &Rfc3339).unwrap(),
            published_by: Some("onboarding-test"),
        })
        .await
        .unwrap();
    publish_agent_for_operation(&registry, &operation, "onboarding-agent").await;

    let selected_key_name = "onboarding-raw-key-canary";
    let selected_key_id = format!("pk_{selected_key_name}");
    let selected_key = create_platform_api_key(
        &registry,
        "onboarding-agent",
        selected_key_name,
        &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
    )
    .await;
    let other_key = create_platform_api_key(
        &registry,
        "onboarding-agent",
        "onboarding-other",
        &[PlatformApiKeyScope::Read, PlatformApiKeyScope::Write],
    )
    .await;

    let pre_call_projection = registry
        .get_onboarding_projection(&test_workspace_id())
        .await
        .unwrap();
    assert!(
        pre_call_projection
            .step(crank_core::OnboardingStepId::PublishOperation)
            .is_some_and(|step| step.completed),
        "fixture must expose a Published Operation: {pre_call_projection:#?}"
    );

    let base_url = spawn_mcp_server(build_test_app(
        registry.clone(),
        Duration::from_millis(0),
        Some("https://crank.example.com".to_owned()),
    ))
    .await;
    let client = reqwest::Client::new();
    let mcp_url = agent_mcp_url(&base_url, "onboarding-agent");

    let selected_session = initialize_session(&client, &mcp_url, &selected_key).await;
    let listed = post_jsonrpc(
        &client,
        &mcp_url,
        &selected_key,
        Some(&selected_session),
        json!({"jsonrpc":"2.0","id":2,"method":"tools/list","params":{}}),
    )
    .await;
    assert_eq!(listed["result"]["tools"][0]["name"], operation.name);
    assert!(
        exact_key_successes(&registry, &operation.id, &selected_key_id)
            .await
            .is_empty(),
        "initialize and tools/list must not complete onboarding"
    );

    let failed = post_jsonrpc(
        &client,
        &mcp_url,
        &selected_key,
        Some(&selected_session),
        json!({
            "jsonrpc":"2.0",
            "id":3,
            "method":"tools/call",
            "params":{"name":operation.name,"arguments":{}}
        }),
    )
    .await;
    assert_eq!(failed["result"]["isError"], true);
    assert!(
        exact_key_successes(&registry, &operation.id, &selected_key_id)
            .await
            .is_empty(),
        "a failed tools/call must not complete onboarding"
    );

    let other_session = initialize_session(&client, &mcp_url, &other_key).await;
    let other_success = post_jsonrpc(
        &client,
        &mcp_url,
        &other_key,
        Some(&other_session),
        json!({
            "jsonrpc":"2.0",
            "id":4,
            "method":"tools/call",
            "params":{
                "name":operation.name,
                "arguments":{"email":"other@example.com"}
            }
        }),
    )
    .await;
    assert_eq!(other_success["result"]["isError"], false);
    assert!(
        exact_key_successes(&registry, &operation.id, &selected_key_id)
            .await
            .is_empty(),
        "a successful tools/call made with another key must not complete onboarding"
    );

    let selected_success = post_jsonrpc(
        &client,
        &mcp_url,
        &selected_key,
        Some(&selected_session),
        json!({
            "jsonrpc":"2.0",
            "id":5,
            "method":"tools/call",
            "params":{
                "name":operation.name,
                "arguments":{"email":"selected@example.com"}
            }
        }),
    )
    .await;
    assert_eq!(selected_success["result"]["isError"], false);

    let exact_successes = exact_key_successes(&registry, &operation.id, &selected_key_id).await;
    assert_eq!(
        exact_successes.len(),
        1,
        "onboarding must complete from exactly one successful public tools/call made with the selected key"
    );
    assert_eq!(
        exact_successes[0]["platform_api_key_id"], selected_key_id,
        "successful invocation evidence must retain the typed key identity"
    );
    assert!(
        exact_successes[0]["request_id"]
            .as_str()
            .is_some_and(|id| !id.is_empty())
    );
    assert!(
        exact_successes[0]["trace_id"]
            .as_str()
            .is_some_and(|id| id.len() == 32)
    );

    // A direct MCP call can precede the first onboarding snapshot. The later
    // server-owned eligibility write must repair the completion event without
    // replacing the invocation timestamp that anchored the exact lineage.
    let projection = registry
        .ensure_onboarding_eligibility(&test_workspace_id(), OffsetDateTime::now_utc())
        .await
        .unwrap();
    let operation_summary = registry
        .get_operation_summary(&test_workspace_id(), &operation.id)
        .await
        .unwrap();
    assert!(
        projection.completed,
        "projection did not complete: {projection:#?}; operation: {operation_summary:#?}"
    );
    let completion_events = registry
        .list_product_events(ListProductEventsQuery {
            workspace_id: &test_workspace_id(),
            kind: Some(ProductEventKind::OnboardingCompleted),
            created_after: OffsetDateTime::UNIX_EPOCH,
            created_before: OffsetDateTime::now_utc() + time::Duration::minutes(1),
            limit: 10,
        })
        .await
        .unwrap();
    assert_eq!(completion_events.len(), 1);
    assert_eq!(
        completion_events[0].event.occurred_at,
        projection.first_call_at.unwrap()
    );

    let persisted_evidence = invocation_evidence(&registry, &operation.id).await;
    assert!(!persisted_evidence.is_empty());
    for evidence in persisted_evidence {
        let serialized = serde_json::to_string(&evidence).unwrap();
        assert!(
            !serialized.contains(&selected_key),
            "raw Bearer canary must never appear in serialized invocation evidence"
        );
        assert!(
            !serialized.contains(&other_key),
            "another raw Bearer secret must never appear in serialized invocation evidence"
        );
        for field in ["message", "request_preview", "response_preview"] {
            let persisted_field = serde_json::to_string(&evidence[field]).unwrap();
            assert!(
                !persisted_field.contains(&selected_key),
                "raw Bearer canary leaked into persisted {field}"
            );
            assert!(
                !persisted_field.contains(&other_key),
                "another raw Bearer secret leaked into persisted {field}"
            );
        }
    }
}

async fn exact_key_successes(
    registry: &crank_registry::PostgresRegistry,
    operation_id: &crank_core::OperationId,
    selected_key_id: &str,
) -> Vec<serde_json::Value> {
    invocation_evidence(registry, operation_id)
        .await
        .into_iter()
        .filter(|log| {
            log.get("platform_api_key_id")
                .and_then(serde_json::Value::as_str)
                == Some(selected_key_id)
                && log.get("status").and_then(serde_json::Value::as_str) == Some("ok")
        })
        .collect()
}

async fn invocation_evidence(
    registry: &crank_registry::PostgresRegistry,
    operation_id: &crank_core::OperationId,
) -> Vec<serde_json::Value> {
    registry
        .list_invocation_logs(ListInvocationLogsQuery {
            workspace_id: &test_workspace_id(),
            level: None,
            status: None,
            outcome_group: None,
            search_text: None,
            source: Some(InvocationSource::AgentToolCall),
            operation_id: Some(operation_id),
            agent_id: None,
            created_after: None,
            created_before: None,
            cursor_created_at: None,
            cursor_id: None,
            limit: 100,
        })
        .await
        .unwrap()
        .into_iter()
        .map(|record| serde_json::to_value(record.log).unwrap())
        .collect()
}