agent-runtime-infra-adapter 0.1.1

use std::sync::{
    Arc,
    atomic::{AtomicI64, Ordering},
};

use agent_infra_sdk::{
    CallOptions, InfraClient,
    trace_contract::{RuntimeEventRecord, TraceScope},
};
use async_trait::async_trait;
use runtime_ports::{OperationContext, PortFailure, TraceSession};
use runtime_types::{ConversationId, ExecutionId};
use serde_json::Value;

use crate::map_sdk_error;

pub struct AgentInfraTraceSession {
    client: Arc<InfraClient>,
    execution_id: ExecutionId,
    conversation_id: Option<ConversationId>,
    agent_id: String,
    runtime_type: String,
    metadata: Value,
    sequence: AtomicI64,
}
impl AgentInfraTraceSession {
    pub fn new(
        client: Arc<InfraClient>,
        execution_id: ExecutionId,
        conversation_id: Option<ConversationId>,
        agent_id: String,
        runtime_type: String,
        metadata: Value,
    ) -> Self {
        Self {
            client,
            execution_id,
            conversation_id,
            agent_id,
            runtime_type,
            metadata,
            sequence: AtomicI64::new(0),
        }
    }
}

#[async_trait]
impl TraceSession for AgentInfraTraceSession {
    async fn append(
        &self,
        operation: &OperationContext,
        event_type: &str,
        payload: Value,
    ) -> Result<(), PortFailure> {
        let remaining = operation.remaining()?;
        let sequence = self.sequence.fetch_add(1, Ordering::SeqCst) + 1;
        let event = RuntimeEventRecord {
            id: format!("{}-{sequence}", self.execution_id),
            scope: TraceScope {
                conversation_id: self
                    .conversation_id
                    .as_ref()
                    .map(ToString::to_string)
                    .unwrap_or_else(|| self.execution_id.to_string()),
                run_id: Some(self.execution_id.to_string()),
                agent_id: Some(self.agent_id.clone()),
                runtime_type: Some(self.runtime_type.clone()),
                ..Default::default()
            },
            event_type: event_type.into(),
            payload,
            sequence,
            created_at_ms: chrono::Utc::now().timestamp_millis(),
            metadata: self.metadata.clone(),
        };
        self.client
            .trace()
            .append_events_with_options(vec![event], CallOptions::default().deadline(remaining))
            .await
            .map(|_| ())
            .map_err(map_sdk_error)
    }
}