crank/crates/crank-registry/tests/integration/agents_usage.rs

#![allow(dead_code, unused_imports)]

use super::common::*;

use std::collections::BTreeMap;

use crank_core::{
    AgentId, AgentOperationBinding, AgentStatus, AgentVersion, ApiKeyHeaderAuthConfig, AuthConfig,
    AuthKind, AuthProfile, ConfigExport, ExecutionConfig, ExportMode, GeneratedDraft,
    GeneratedDraftStatus, HttpMethod, InvocationLog, MembershipRole, OperationId,
    OperationSecurityLevel, OperationStatus, PlatformApiKey, PlatformApiKeyId, PlatformApiKeyScope,
    PlatformApiKeyStatus, Protocol, RestTarget, RetryPolicy, Samples, SecretId, Target,
    ToolDescription, ToolExample, User, UserId, UserSessionId, WizardState, Workspace, WorkspaceId,
};
use crank_mapping::{MappingRule, MappingSet};
use crank_schema::{Schema, SchemaKind};
use serde_json::json;
use sqlx::{Executor, PgPool, postgres::PgPoolOptions};
use time::{OffsetDateTime, format_description::well_known::Rfc3339};

use crank_registry::{
    CreateAgentRequest, CreateInvocationLogRequest, CreatePlatformApiKeyRequest,
    CreateVersionRequest, CreateWorkspaceRequest, CreateYamlImportJobRequest, DescriptorKind,
    DescriptorMetadata, OperationSampleMetadata, PlatformApiKeyRecord, PostgresRegistry,
    PublishAgentRequest, PublishRequest, RegistryError, RegistryOperation, SampleKind,
    SaveAuthProfileRequest, SaveDescriptorMetadataRequest, SaveSampleMetadataRequest,
    WorkspaceRecord, YamlImportJobCompletion, YamlImportJobId, YamlImportJobStatus,
};

fn test_workspace_id() -> WorkspaceId {
    WorkspaceId::new("ws_default")
}

fn timestamp(value: &str) -> OffsetDateTime {
    OffsetDateTime::parse(value, &Rfc3339).unwrap()
}

#[tokio::test]
async fn manages_agent_read_paths() {
    let database = TestDatabase::new().await;
    let registry = database.registry().await;
    let operation = test_operation("op_agent_01", 1, OperationStatus::Draft);
    let agent = test_agent("agent_01", AgentStatus::Draft);
    let version = test_agent_version(&agent.id, 1, AgentStatus::Draft);
    let bindings = vec![AgentOperationBinding {
        operation_id: operation.id.clone(),
        operation_version: operation.version,
        tool_name: "create_lead".to_owned(),
        tool_title: "Create lead".to_owned(),
        tool_description_override: Some("Creates CRM lead".to_owned()),
        enabled: true,
    }];

    registry
        .create_operation(&test_workspace_id(), &operation, None)
        .await
        .unwrap();
    registry
        .create_agent(CreateAgentRequest {
            agent: &agent,
            version: &version,
            bindings: &bindings,
        })
        .await
        .unwrap();

    let listed = registry.list_agents(&test_workspace_id()).await.unwrap();
    let summary = registry
        .get_agent_summary(&test_workspace_id(), &agent.id)
        .await
        .unwrap()
        .unwrap();
    let loaded_version = registry
        .get_agent_version(&test_workspace_id(), &agent.id, version.version)
        .await
        .unwrap()
        .unwrap();

    assert_eq!(listed, vec![summary.clone()]);
    assert_eq!(summary.id, agent.id);
    assert_eq!(summary.slug, agent.slug);
    assert_eq!(summary.display_name, agent.display_name);
    assert_eq!(summary.status, agent.status);
    assert_eq!(loaded_version.agent_id, agent.id);
    assert_eq!(loaded_version.version, version.version);
    assert_eq!(loaded_version.status, version.status);
    assert_eq!(loaded_version.snapshot, version);
    assert_eq!(loaded_version.bindings, bindings);

    database.cleanup().await;
}

#[tokio::test]
async fn manages_published_agent_tool_reads() {
    let database = TestDatabase::new().await;
    let registry = database.registry().await;
    let operation = test_operation("op_agent_pub_01", 1, OperationStatus::Draft);
    let operation_v2 = test_operation("op_agent_pub_01", 2, OperationStatus::Draft);
    let agent = test_agent("agent_pub_01", AgentStatus::Draft);
    let version = test_agent_version(&agent.id, 1, AgentStatus::Draft);
    let bindings = vec![AgentOperationBinding {
        operation_id: operation.id.clone(),
        operation_version: operation_v2.version,
        tool_name: "create_lead".to_owned(),
        tool_title: "Create lead".to_owned(),
        tool_description_override: Some("Creates CRM lead".to_owned()),
        enabled: true,
    }];

    registry
        .create_operation(&test_workspace_id(), &operation, None)
        .await
        .unwrap();
    registry
        .create_version(CreateVersionRequest {
            workspace_id: &test_workspace_id(),
            snapshot: &operation_v2,
            change_note: Some("publishable"),
            created_by: Some("alice"),
        })
        .await
        .unwrap();
    registry
        .publish_operation(PublishRequest {
            workspace_id: &test_workspace_id(),
            operation_id: &operation.id,
            version: operation_v2.version,
            published_at: &timestamp("2026-03-25T12:10:00Z"),
            published_by: Some("alice"),
        })
        .await
        .unwrap();
    registry
        .create_agent(CreateAgentRequest {
            agent: &agent,
            version: &version,
            bindings: &bindings,
        })
        .await
        .unwrap();
    registry
        .publish_agent(PublishAgentRequest {
            workspace_id: &test_workspace_id(),
            agent_id: &agent.id,
            version: version.version,
            published_at: &timestamp("2026-03-25T12:11:00Z"),
            published_by: Some("alice"),
        })
        .await
        .unwrap();

    let tools = registry
        .get_published_agent_tools_by_slug("default", &agent.slug)
        .await
        .unwrap();

    assert_eq!(tools.len(), 1);
    assert_eq!(tools[0].workspace_id, test_workspace_id());
    assert_eq!(tools[0].workspace_slug, "default");
    assert_eq!(tools[0].agent_id, agent.id);
    assert_eq!(tools[0].agent_slug, agent.slug);
    assert_eq!(tools[0].tool_name, bindings[0].tool_name);
    assert_eq!(tools[0].tool_title, bindings[0].tool_title);
    assert_eq!(tools[0].tool_description, "Creates CRM lead");
    assert_eq!(tools[0].operation.id, operation_v2.id);
    assert_eq!(tools[0].operation.version, operation_v2.version);
    assert_eq!(tools[0].operation.protocol, operation_v2.protocol);
    assert!(tools[0].operation.is_published());

    database.cleanup().await;
}

#[tokio::test]
async fn manages_operation_usage_and_agent_ref_reads() {
    let database = TestDatabase::new().await;
    let registry = database.registry().await;
    let operation = test_operation("op_usage_01", 1, OperationStatus::Draft);
    let operation_v2 = test_operation("op_usage_01", 2, OperationStatus::Draft);
    let agent = test_agent("agent_usage_01", AgentStatus::Draft);
    let version = test_agent_version(&agent.id, 1, AgentStatus::Draft);
    let bindings = vec![AgentOperationBinding {
        operation_id: operation.id.clone(),
        operation_version: operation_v2.version,
        tool_name: "create_lead".to_owned(),
        tool_title: "Create lead".to_owned(),
        tool_description_override: None,
        enabled: true,
    }];

    registry
        .create_operation(&test_workspace_id(), &operation, None)
        .await
        .unwrap();
    registry
        .create_version(CreateVersionRequest {
            workspace_id: &test_workspace_id(),
            snapshot: &operation_v2,
            change_note: Some("publishable"),
            created_by: Some("alice"),
        })
        .await
        .unwrap();
    registry
        .publish_operation(PublishRequest {
            workspace_id: &test_workspace_id(),
            operation_id: &operation.id,
            version: operation_v2.version,
            published_at: &timestamp("2026-03-25T12:10:00Z"),
            published_by: Some("alice"),
        })
        .await
        .unwrap();
    registry
        .create_agent(CreateAgentRequest {
            agent: &agent,
            version: &version,
            bindings: &bindings,
        })
        .await
        .unwrap();
    registry
        .publish_agent(PublishAgentRequest {
            workspace_id: &test_workspace_id(),
            agent_id: &agent.id,
            version: version.version,
            published_at: &timestamp("2026-03-25T12:11:00Z"),
            published_by: Some("alice"),
        })
        .await
        .unwrap();
    registry
        .create_invocation_log(CreateInvocationLogRequest {
            log: &test_invocation_log(
                "log_usage_ok",
                &operation.id,
                Some(agent.id.clone()),
                crank_core::InvocationStatus::Ok,
                120,
                "2026-03-25T12:20:00Z",
            ),
        })
        .await
        .unwrap();
    registry
        .create_invocation_log(CreateInvocationLogRequest {
            log: &test_invocation_log(
                "log_usage_err",
                &operation.id,
                Some(agent.id.clone()),
                crank_core::InvocationStatus::Error,
                240,
                "2026-03-25T12:21:00Z",
            ),
        })
        .await
        .unwrap();

    let has_bindings = registry
        .has_published_agent_bindings_for_operation(&test_workspace_id(), &operation.id)
        .await
        .unwrap();
    let agent_refs = registry
        .list_operation_agent_refs(&test_workspace_id())
        .await
        .unwrap();
    let usage = registry
        .list_operation_usage_summaries(&test_workspace_id(), "2026-03-25T12:00:00Z")
        .await
        .unwrap();

    assert!(has_bindings);
    assert_eq!(agent_refs.len(), 1);
    assert_eq!(agent_refs[0].operation_id, operation.id);
    assert_eq!(agent_refs[0].agent_id, agent.id);
    assert_eq!(agent_refs[0].agent_slug, agent.slug);
    assert_eq!(agent_refs[0].display_name, agent.display_name);
    assert_eq!(usage.len(), 1);
    assert_eq!(usage[0].operation_id, operation.id);
    assert_eq!(usage[0].calls_today, 2);
    assert_eq!(usage[0].error_rate_pct, 50.0);
    assert_eq!(usage[0].avg_latency_ms, 180);

    database.cleanup().await;
}