agent_runtime_infra_adapter/

lib.rs

//! Agent Infra SDK adapters for Runtime ports.

mod binding;
mod context;
mod interaction;
mod model;
mod resources;
mod trace;
mod workspace;

use std::{collections::BTreeMap, sync::Arc, time::Duration};

use agent_infra_sdk::{
    BearerCredential, CallOptions, CredentialsProvider, DelegationLeaseCredentials, InfraClient,
    InfraClientBuilder, InfraClientError, StaticCredentials,
};
use async_trait::async_trait;
use runtime_ports::{
    CommitDisposition, ExecutionDelegationLease, ExecutionSessionFactory, ExecutionSessions,
    OperationContext, PortFailure, PortFailureKind, SessionScope,
};
use runtime_types::{
    CallerScope, CredentialHandle, DelegationLeaseRef, ExecutionId, OperationId, RequestAuthority,
};

pub use binding::{RuntimeBinding, StaticRuntimeBindingResolver};
use context::{
    AgentInfraContextAuditSession, AgentInfraConversationSession, NullContextAuditSession,
    NullConversationSession,
};
use interaction::{InMemoryInteractionBroker, PendingInteractionSession};
use model::AgentInfraModelSession;
use resources::BindingResourceSession;
use trace::AgentInfraTraceSession;
use workspace::AgentInfraWorkspaceSession;

#[derive(Debug, Clone)]
pub struct AgentInfraConfig {
    pub gateway_url: String,
    pub request_timeout: Duration,
    pub max_response_bytes: usize,
    pub trusted_mesh_http: bool,
    pub workload_subject: String,
    pub workload_credential: CredentialHandle,
    pub delegation_ttl: Duration,
}

impl AgentInfraConfig {
    pub fn validate(&self) -> Result<(), PortFailure> {
        if self.gateway_url.trim().is_empty() {
            return Err(PortFailure::new(
                PortFailureKind::Invalid,
                "GATEWAY_URL_EMPTY",
                "Agent Infra Gateway URL must not be empty",
            ));
        }
        if self.workload_subject.trim().is_empty()
            || !(Duration::from_secs(30)..=Duration::from_secs(7 * 24 * 60 * 60))
                .contains(&self.delegation_ttl)
        {
            return Err(PortFailure::new(
                PortFailureKind::Invalid,
                "DELEGATION_CONFIG_INVALID",
                "workload subject and delegation TTL within 30s..=7d are required",
            ));
        }
        if self.request_timeout.is_zero() || self.request_timeout > Duration::from_secs(600) {
            return Err(PortFailure::new(
                PortFailureKind::Invalid,
                "GATEWAY_TIMEOUT_INVALID",
                "Agent Infra request timeout must be within 1ns..=600s",
            ));
        }
        if !(1024..=1024 * 1024 * 1024).contains(&self.max_response_bytes) {
            return Err(PortFailure::new(
                PortFailureKind::Invalid,
                "GATEWAY_RESPONSE_LIMIT_INVALID",
                "Agent Infra response limit must be within 1 KiB..=1 GiB",
            ));
        }
        Ok(())
    }
}

pub struct AgentInfraSessionFactory {
    config: AgentInfraConfig,
    interactions: Arc<InMemoryInteractionBroker>,
    workload_credentials: Arc<dyn CredentialsProvider>,
}

impl AgentInfraSessionFactory {
    pub fn new(config: AgentInfraConfig) -> Result<Self, PortFailure> {
        config.validate()?;
        let workload_credentials = Arc::new(
            StaticCredentials::new(config.workload_credential.expose()).map_err(|_| {
                PortFailure::new(
                    PortFailureKind::Invalid,
                    "WORKLOAD_CREDENTIAL_INVALID",
                    "Agent Infra workload credential is invalid",
                )
            })?,
        );
        Ok(Self {
            config,
            interactions: Arc::new(InMemoryInteractionBroker::new(10_000)),
            workload_credentials,
        })
    }

    pub async fn check_ready(&self) -> Result<(), PortFailure> {
        InfraClientBuilder::new(self.config.gateway_url.clone())
            .trusted_mesh_http(self.config.trusted_mesh_http)
            .request_timeout(self.config.request_timeout)
            .build()
            .map_err(map_sdk_error)?
            .gateway()
            .ready()
            .await
            .map_err(map_sdk_error)
            .and_then(|status| {
                if status.is_ready() {
                    Ok(())
                } else {
                    Err(PortFailure::new(
                        PortFailureKind::Unavailable,
                        "GATEWAY_NOT_READY",
                        "Agent Infra Gateway is not ready",
                    ))
                }
            })
    }

    fn delegated_client(
        &self,
        delegation: &ExecutionDelegationLease,
        scope: &SessionScope,
    ) -> Result<Arc<InfraClient>, PortFailure> {
        let capabilities = delegation_capabilities(scope);
        let provider: Arc<dyn CredentialsProvider> = Arc::new(
            if self.config.trusted_mesh_http {
                DelegationLeaseCredentials::new_trusted_mesh_http(
                    self.config.gateway_url.clone(),
                    delegation.lease_ref.as_str(),
                    self.workload_credentials.clone(),
                    capabilities,
                )
            } else {
                DelegationLeaseCredentials::new(
                    self.config.gateway_url.clone(),
                    delegation.lease_ref.as_str(),
                    self.workload_credentials.clone(),
                    capabilities,
                )
            }
            .map_err(|_| {
                PortFailure::new(
                    PortFailureKind::Invalid,
                    "DELEGATION_PROVIDER_INVALID",
                    "execution delegation credential provider could not be configured",
                )
            })?,
        );
        InfraClientBuilder::new(self.config.gateway_url.clone())
            .credentials(provider)
            .request_timeout(self.config.request_timeout)
            .max_response_bytes(self.config.max_response_bytes)
            .trusted_mesh_http(self.config.trusted_mesh_http)
            .user_agent(format!("agent-runtime/{}", env!("CARGO_PKG_VERSION")))
            .build()
            .map(Arc::new)
            .map_err(map_sdk_error)
    }
}

#[async_trait]
impl ExecutionSessionFactory for AgentInfraSessionFactory {
    async fn establish_delegation(
        &self,
        operation: &OperationContext,
        authority: &RequestAuthority,
        scope: &SessionScope,
    ) -> Result<ExecutionDelegationLease, PortFailure> {
        let client = InfraClientBuilder::new(self.config.gateway_url.clone())
            .credentials(self.workload_credentials.clone())
            .trusted_mesh_http(self.config.trusted_mesh_http)
            .request_timeout(self.config.request_timeout)
            .build()
            .map_err(map_sdk_error)?;
        let caller = BearerCredential::new(authority.credential.expose()).map_err(|_| {
            PortFailure::new(
                PortFailureKind::Forbidden,
                "CALLER_CREDENTIAL_INVALID",
                "transient caller credential is invalid",
            )
        })?;
        let lease = client
            .gateway()
            .establish_execution_delegation(
                caller,
                &agent_infra_sdk::gateway_contract::EstablishExecutionDelegationRequest {
                    execution_id: scope.execution_id.to_string(),
                    workload_subject: self.config.workload_subject.clone(),
                    audience_capabilities: delegation_capabilities(scope),
                    expires_in_seconds: self
                        .config
                        .delegation_ttl
                        .as_secs()
                        .max(operation.remaining()?.as_secs().saturating_add(30))
                        .min(7 * 24 * 60 * 60),
                    idempotency_key: format!("{}:delegation", scope.execution_id),
                },
                CallOptions::default()
                    .deadline(operation.remaining()?)
                    .request_id(operation.id.to_string()),
            )
            .await
            .map_err(map_sdk_error)?;
        if lease.execution_id != scope.execution_id.as_str() {
            return Err(PortFailure::new(
                PortFailureKind::Protocol,
                "DELEGATION_EXECUTION_MISMATCH",
                "Agent Infra delegation is bound to a different execution",
            ));
        }
        Ok(ExecutionDelegationLease {
            lease_ref: DelegationLeaseRef::new(lease.lease_ref).map_err(|_| {
                PortFailure::new(
                    PortFailureKind::Protocol,
                    "DELEGATION_REFERENCE_INVALID",
                    "Agent Infra returned an invalid delegation reference",
                )
            })?,
            expires_at_seconds: lease.expires_at_seconds,
            revision: lease.revision,
        })
    }

    async fn create(
        &self,
        _: &OperationContext,
        _: &CallerScope,
        delegation: &ExecutionDelegationLease,
        scope: &SessionScope,
    ) -> Result<ExecutionSessions, PortFailure> {
        let client = self.delegated_client(delegation, scope)?;
        let (conversation, context_audit): (
            Arc<dyn runtime_ports::ConversationSession>,
            Arc<dyn runtime_ports::ContextAuditSession>,
        ) = match &scope.conversation_id {
            Some(id) => (
                Arc::new(AgentInfraConversationSession::new(
                    client.clone(),
                    id.clone(),
                    scope.execution_id.clone(),
                    scope.resolved.agent_id.clone(),
                    scope.resolved.runtime_type.clone(),
                )),
                Arc::new(AgentInfraContextAuditSession::new(
                    client.clone(),
                    id.clone(),
                    scope.resolved.agent_id.clone(),
                    scope.resolved.runtime_type.clone(),
                )),
            ),
            None => (
                Arc::new(NullConversationSession),
                Arc::new(NullContextAuditSession),
            ),
        };
        Ok(ExecutionSessions {
            model: Arc::new(AgentInfraModelSession::new(client.clone())),
            conversation,
            context_audit,
            workspace: Arc::new(AgentInfraWorkspaceSession::new(
                client.clone(),
                scope.resolved.workspace_id.clone(),
            )),
            trace: Arc::new(AgentInfraTraceSession::new(
                client,
                scope.execution_id.clone(),
                scope.conversation_id.clone(),
                scope.resolved.agent_id.clone(),
                scope.resolved.runtime_type.clone(),
                trace_metadata(&scope.resolved),
            )),
            interaction: Arc::new(PendingInteractionSession::new(self.interactions.clone())),
            subagent: Arc::new(UnsupportedSubagentSession),
            resources: Arc::new(BindingResourceSession::new(scope.resolved.clone())),
        })
    }

    async fn submit_input(
        &self,
        _: &OperationContext,
        _: &CallerScope,
        execution_id: &ExecutionId,
        operation_id: &OperationId,
        request_id: &str,
        response: runtime_ports::InteractionResponse,
    ) -> Result<(), PortFailure> {
        self.interactions
            .submit(execution_id, operation_id, request_id, response)
            .await
    }

    async fn renew_delegation(
        &self,
        operation: &OperationContext,
        _: &CallerScope,
        delegation: &ExecutionDelegationLease,
    ) -> Result<ExecutionDelegationLease, PortFailure> {
        let client = InfraClientBuilder::new(self.config.gateway_url.clone())
            .credentials(self.workload_credentials.clone())
            .trusted_mesh_http(self.config.trusted_mesh_http)
            .build()
            .map_err(map_sdk_error)?;
        let renewed = client
            .gateway()
            .renew_execution_delegation(
                &agent_infra_sdk::gateway_contract::RenewExecutionDelegationRequest {
                    lease_ref: delegation.lease_ref.as_str().into(),
                    expected_revision: delegation.revision,
                    extend_by_seconds: self
                        .config
                        .delegation_ttl
                        .as_secs()
                        .max(operation.remaining()?.as_secs().saturating_add(30))
                        .min(7 * 24 * 60 * 60),
                    idempotency_key: format!(
                        "{}:renew:{}",
                        operation.execution_id, delegation.revision
                    ),
                },
                CallOptions::default()
                    .deadline(operation.remaining()?)
                    .request_id(operation.id.to_string()),
            )
            .await
            .map_err(map_sdk_error)?;
        Ok(ExecutionDelegationLease {
            lease_ref: DelegationLeaseRef::new(renewed.lease_ref).map_err(|_| {
                PortFailure::new(
                    PortFailureKind::Protocol,
                    "DELEGATION_REFERENCE_INVALID",
                    "Agent Infra returned an invalid delegation reference",
                )
            })?,
            expires_at_seconds: renewed.expires_at_seconds,
            revision: renewed.revision,
        })
    }

    async fn revoke_delegation(
        &self,
        operation: &OperationContext,
        _: &CallerScope,
        delegation: &ExecutionDelegationLease,
    ) -> Result<(), PortFailure> {
        let client = InfraClientBuilder::new(self.config.gateway_url.clone())
            .credentials(self.workload_credentials.clone())
            .trusted_mesh_http(self.config.trusted_mesh_http)
            .build()
            .map_err(map_sdk_error)?;
        client
            .gateway()
            .revoke_execution_delegation(
                &agent_infra_sdk::gateway_contract::RevokeExecutionDelegationRequest {
                    lease_ref: delegation.lease_ref.as_str().into(),
                    idempotency_key: format!("{}:revoke", operation.execution_id),
                },
                CallOptions::default()
                    .deadline(operation.remaining()?)
                    .request_id(operation.id.to_string()),
            )
            .await
            .map_err(map_sdk_error)
    }
}

fn delegation_capabilities(scope: &SessionScope) -> BTreeMap<String, Vec<String>> {
    let mut capabilities = BTreeMap::from([
        (
            "agent-model-infra".into(),
            vec![format!("model-profile:{}:invoke", scope.resolved.model)],
        ),
        (
            "workspace".into(),
            vec![
                format!("workspace:{}:read", scope.resolved.workspace_id),
                format!("workspace:{}:write", scope.resolved.workspace_id),
                format!("workspace:{}:execute", scope.resolved.workspace_id),
            ],
        ),
        (
            "agent-trace-infra".into(),
            vec![format!("trace:{}:write", scope.execution_id)],
        ),
    ]);
    if let Some(conversation_id) = &scope.conversation_id {
        capabilities.insert(
            "agent-context-infra".into(),
            vec![
                format!("conversation:{conversation_id}:read"),
                format!("conversation:{conversation_id}:append"),
                format!("context-audit:{conversation_id}:write"),
            ],
        );
    }
    capabilities
}

/// Trace correlation is intentionally allow-listed. Request instructions may
/// contain user data or secrets and therefore never become searchable Trace
/// metadata. Only task-prefixed metadata already bounded by the HTTP contract
/// is projected alongside stable resource identities.
fn trace_metadata(resolved: &runtime_ports::ResolvedExecutionContext) -> serde_json::Value {
    let task = resolved
        .metadata
        .iter()
        .filter_map(|(key, value)| {
            key.strip_prefix("task.")
                .map(|key| (key.to_string(), serde_json::Value::String(value.clone())))
        })
        .collect::<serde_json::Map<_, _>>();
    serde_json::json!({
        "workspaceId": resolved.workspace_id.as_str(),
        "model": resolved.model,
        "definitionId": resolved.definition_id,
        "definitionVersion": resolved.definition_version,
        "task": task,
    })
}

struct UnsupportedSubagentSession;
#[async_trait]
impl runtime_ports::SubagentSession for UnsupportedSubagentSession {
    async fn execute(
        &self,
        _: &OperationContext,
        _: runtime_ports::SubagentRequest,
    ) -> Result<runtime_ports::SubagentOutcome, PortFailure> {
        Err(PortFailure::new(
            PortFailureKind::Unavailable,
            "SUBAGENT_LINK_NOT_CONFIGURED",
            "subagent execution requires the Runtime Link child-execution adapter",
        ))
    }
}

pub(crate) fn map_sdk_error(error: InfraClientError) -> PortFailure {
    // Classify before erasing the SDK error. A response proves the server saw
    // the operation, while transport/deadline/decode failures after dispatch
    // cannot prove whether a write committed. Callers may narrow this for a
    // known read-only operation, but must never widen Unknown to NotCommitted.
    let commit = match &error {
        InfraClientError::ClientBuild(_)
        | InfraClientError::InvalidOptions { .. }
        | InfraClientError::Credential { .. }
        | InfraClientError::InvalidEndpoint { .. } => CommitDisposition::NotCommitted,
        InfraClientError::OperationTerminal { .. } => CommitDisposition::Committed,
        InfraClientError::HttpStatus { status, .. }
            if matches!(*status, 400 | 401 | 403 | 404 | 405 | 413 | 415 | 422 | 429) =>
        {
            CommitDisposition::NotCommitted
        }
        InfraClientError::HttpStatus { .. } => CommitDisposition::Unknown,
        InfraClientError::Request { .. }
        | InfraClientError::DeadlineExceeded { .. }
        | InfraClientError::Canceled { .. }
        | InfraClientError::ResponseTooLarge { .. }
        | InfraClientError::Decode { .. }
        | InfraClientError::Protocol { .. } => CommitDisposition::Unknown,
    };
    let retryable = error.retryable();
    let code = error.code().to_string();
    let kind = match code.as_str() {
        "TIMEOUT" => PortFailureKind::Timeout,
        "CANCELED" => PortFailureKind::Canceled,
        "UNAUTHENTICATED" => PortFailureKind::Forbidden,
        "FORBIDDEN" => PortFailureKind::Forbidden,
        "NOT_FOUND" => PortFailureKind::NotFound,
        "CONFLICT" => PortFailureKind::Conflict,
        "RATE_LIMITED" => PortFailureKind::RateLimited,
        "DECODE" | "PROTOCOL" => PortFailureKind::Protocol,
        _ => PortFailureKind::Unavailable,
    };
    // SDK displays may contain endpoint details from reqwest. Runtime-facing
    // failures deliberately retain only the stable class; request IDs remain
    // in SDK telemetry and must not turn into model/public payload.
    let message = match kind {
        PortFailureKind::Timeout => "Agent Infra operation exceeded its deadline",
        PortFailureKind::Canceled => "Agent Infra operation was canceled",
        PortFailureKind::Forbidden => "Agent Infra rejected the execution authority",
        PortFailureKind::NotFound => "Agent Infra resource was not found",
        PortFailureKind::Conflict => "Agent Infra operation conflicted with current state",
        PortFailureKind::RateLimited => "Agent Infra rate limit was exceeded",
        PortFailureKind::Protocol => "Agent Infra returned an invalid protocol response",
        _ => "Agent Infra dependency is unavailable",
    };
    let mut failure = PortFailure::new(kind, code, message);
    failure.retryable = retryable;
    failure.commit = commit;
    failure
}

#[cfg(test)]
mod sdk_error_tests {
    use super::*;

    #[test]
    fn trace_metadata_includes_task_correlation_but_not_instructions() {
        let resolved = runtime_ports::ResolvedExecutionContext {
            runtime_instance_id: runtime_types::RuntimeInstanceId::new("runtime-1").unwrap(),
            agent_id: "agt_test".into(),
            runtime_type: "test".into(),
            model: "model-1".into(),
            workspace_id: runtime_types::WorkspaceId::new("workspace-1").unwrap(),
            definition_id: "definition-1".into(),
            definition_version: "1".into(),
            definition_digest: None,
            metadata: BTreeMap::from([
                ("instructions".into(), "secret prompt".into()),
                ("task.traceId".into(), "trace-1".into()),
            ]),
        };
        let metadata = trace_metadata(&resolved);
        assert_eq!(metadata["workspaceId"], "workspace-1");
        assert_eq!(metadata["task"]["traceId"], "trace-1");
        assert!(!metadata.to_string().contains("secret prompt"));
    }

    #[test]
    fn preserves_commit_disposition_when_erasing_sdk_errors() {
        let before_send = map_sdk_error(InfraClientError::InvalidOptions {
            message: "bad options".into(),
        });
        assert_eq!(before_send.commit, CommitDisposition::NotCommitted);

        let server_failure = map_sdk_error(InfraClientError::HttpStatus {
            service: "model",
            status: 503,
            code: "UNAVAILABLE".into(),
            message: "unavailable".into(),
            retryable: true,
            request_id: Some("request-1".into()),
            retry_after: None,
        });
        assert_eq!(server_failure.commit, CommitDisposition::Unknown);

        let rejected = map_sdk_error(InfraClientError::HttpStatus {
            service: "model",
            status: 422,
            code: "INVALID".into(),
            message: "invalid".into(),
            retryable: false,
            request_id: Some("request-2".into()),
            retry_after: None,
        });
        assert_eq!(rejected.commit, CommitDisposition::NotCommitted);

        let invalid_response = map_sdk_error(InfraClientError::Protocol {
            service: "model",
            message: "invalid response".into(),
        });
        assert_eq!(invalid_response.commit, CommitDisposition::Unknown);
    }
}