use std::sync::Arc;
use std::time::Duration;
use agent_evaluate_contract::{
CANCEL_RUN_PATH, CLAIM_REVIEW_PATH, COMPARE_REPORTS_PATH, COMPLETE_REVIEW_PATH,
CancelRunRequest, CaseRunPage, ClaimReviewRequest, CompareReportsRequest,
CompleteReviewRequest, CreateEvalRunRequest, DURABLE_REPORT_PATH, EnqueueReviewRequest,
EvalOperation, EvalReport, EvalSuiteVersion, GetOperationRequest, GetReportRequest,
HumanReview, ListCaseRunsRequest, OPERATIONS_PATH, OperationState, REVIEW_QUEUE_PATH,
RUN_CASES_PATH, RUNS_PATH, RegisterSuiteRequest, ReportComparison, SUITES_PATH,
};
use async_trait::async_trait;
use reqwest::Client;
use crate::operation::{OperationHandle, OperationObservation, OperationPoller, OperationProgress};
use crate::transport::{
CallOptions, ClientOptions, HttpTransport, InfraClientError, ServiceEndpoint,
};
#[derive(Clone, Debug)]
pub struct EvaluateClient {
transport: HttpTransport,
}
impl EvaluateClient {
#[cfg(test)]
fn new(base_url: impl Into<String>) -> Self {
Self::new_with_endpoint(
Client::new(),
ServiceEndpoint::new(base_url),
ClientOptions::default(),
)
}
pub(crate) fn new_with_endpoint(
http: Client,
endpoint: ServiceEndpoint,
options: ClientOptions,
) -> Self {
let endpoint = endpoint.with_default_credential_audience(agent_evaluate_contract::AUDIENCE);
Self {
transport: HttpTransport::new_with_options(http, "evaluate", endpoint, options),
}
}
pub async fn register_suite(
&self,
request: &RegisterSuiteRequest,
) -> Result<EvalSuiteVersion, InfraClientError> {
self.transport
.post_json_with_options(
SUITES_PATH,
request,
CallOptions::default()
.idempotency_key(format!("{}:{}", request.suite_id, request.version)),
)
.await
}
pub async fn create_run(
&self,
request: &CreateEvalRunRequest,
) -> Result<OperationHandle<EvalOperation>, InfraClientError> {
let options = if request.idempotency_key.is_empty() {
CallOptions::default()
} else {
CallOptions::default().idempotency_key(&request.idempotency_key)
};
let operation: EvalOperation = self
.transport
.post_json_with_options(RUNS_PATH, request, options)
.await?;
Ok(OperationHandle::new(
operation.id.clone(),
Some(operation),
Arc::new(EvaluatePoller {
transport: self.transport.clone(),
}),
))
}
pub async fn operation(
&self,
operation_id: impl Into<String>,
) -> Result<OperationHandle<EvalOperation>, InfraClientError> {
let operation_id = operation_id.into();
let operation: EvalOperation = self
.transport
.post_json_idempotent(
OPERATIONS_PATH,
&GetOperationRequest {
operation_id: operation_id.clone(),
},
)
.await?;
Ok(OperationHandle::new(
operation_id,
Some(operation),
Arc::new(EvaluatePoller {
transport: self.transport.clone(),
}),
))
}
pub async fn list_case_runs(
&self,
request: &ListCaseRunsRequest,
) -> Result<CaseRunPage, InfraClientError> {
self.transport
.post_json_idempotent(RUN_CASES_PATH, request)
.await
}
pub async fn next_case_runs(
&self,
request: &ListCaseRunsRequest,
current: &CaseRunPage,
) -> Result<Option<CaseRunPage>, InfraClientError> {
let Some(cursor) = current.next_cursor.clone() else {
return Ok(None);
};
let mut next = request.clone();
next.cursor = Some(cursor);
self.list_case_runs(&next).await.map(Some)
}
pub async fn report(&self, report_id: &str) -> Result<EvalReport, InfraClientError> {
self.transport
.post_json_idempotent(
DURABLE_REPORT_PATH,
&GetReportRequest {
report_id: report_id.to_string(),
},
)
.await
}
pub async fn compare_reports(
&self,
request: &CompareReportsRequest,
) -> Result<ReportComparison, InfraClientError> {
self.transport
.post_json_idempotent(COMPARE_REPORTS_PATH, request)
.await
}
pub async fn enqueue_review(
&self,
request: &EnqueueReviewRequest,
idempotency_key: &str,
) -> Result<HumanReview, InfraClientError> {
self.transport
.post_json_with_options(
REVIEW_QUEUE_PATH,
request,
CallOptions::default().idempotency_key(idempotency_key),
)
.await
}
pub async fn claim_review(
&self,
request: &ClaimReviewRequest,
) -> Result<Option<HumanReview>, InfraClientError> {
self.transport.post_json(CLAIM_REVIEW_PATH, request).await
}
pub async fn complete_review(
&self,
request: &CompleteReviewRequest,
) -> Result<HumanReview, InfraClientError> {
self.transport
.post_json_with_options(
COMPLETE_REVIEW_PATH,
request,
CallOptions::default().idempotency_key(format!(
"review:{}:{}",
request.review_id, request.fencing_token
)),
)
.await
}
}
#[derive(Clone, Debug)]
struct EvaluatePoller {
transport: HttpTransport,
}
#[async_trait]
impl OperationPoller<EvalOperation> for EvaluatePoller {
async fn poll(&self, operation_id: &str) -> Result<EvalOperation, InfraClientError> {
self.transport
.post_json_idempotent(
OPERATIONS_PATH,
&GetOperationRequest {
operation_id: operation_id.to_string(),
},
)
.await
}
async fn cancel(&self, operation_id: &str) -> Result<(), InfraClientError> {
let operation = self.poll(operation_id).await?;
let _: () = self
.transport
.post_json_with_options(
CANCEL_RUN_PATH,
&CancelRunRequest {
run_id: operation.run_id,
},
CallOptions::default().idempotency_key(format!("cancel:{operation_id}")),
)
.await?;
Ok(())
}
fn observe(&self, operation: &EvalOperation) -> OperationObservation {
let progress = match operation.state {
OperationState::Succeeded => OperationProgress::Succeeded,
OperationState::Failed => OperationProgress::Failed,
OperationState::Canceled => OperationProgress::Canceled,
OperationState::Queued | OperationState::Running | OperationState::Canceling => {
OperationProgress::Pending
}
};
OperationObservation {
progress,
next_poll_after: operation.next_poll_after_ms.map(Duration::from_millis),
error_code: operation.error_code.clone(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
use wiremock::matchers::{method, path};
use wiremock::{Mock, MockServer, ResponseTemplate};
#[tokio::test]
async fn cancel_accepts_an_empty_202_representation() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path(OPERATIONS_PATH))
.respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
"id": "operation-1",
"runId": "run-1",
"state": "running",
"phase": "execute",
"attempt": 1,
"completedCases": 0,
"totalCases": 1,
"createdAtMs": 1,
"updatedAtMs": 1,
"version": 1
})))
.expect(2)
.mount(&server)
.await;
Mock::given(method("POST"))
.and(path(CANCEL_RUN_PATH))
.respond_with(ResponseTemplate::new(202))
.expect(1)
.mount(&server)
.await;
let handle = EvaluateClient::new(server.uri())
.operation("operation-1")
.await
.unwrap();
handle.cancel().await.unwrap();
}
#[tokio::test]
async fn empty_review_queue_is_none_instead_of_a_decode_error() {
let server = MockServer::start().await;
Mock::given(method("POST"))
.and(path(CLAIM_REVIEW_PATH))
.respond_with(ResponseTemplate::new(204))
.expect(1)
.mount(&server)
.await;
let review = EvaluateClient::new(server.uri())
.claim_review(&ClaimReviewRequest { lease_ms: 1_000 })
.await
.unwrap();
assert!(review.is_none());
}
}