use std::sync::{
Arc,
atomic::{AtomicUsize, Ordering},
};
use agent_infra_sdk::{
CallOptions, InfraClient,
context_contract::{
ContextDerivationKind as InfraDerivationKind, ContextSnapshotItem, ContextSnapshotItemKind,
ContextSnapshotStage as InfraSnapshotStage, ConversationRecord, ConversationScope,
CreateContextDerivationRequest, CreateContextSnapshotRequest,
ModelInvocationState as InfraInvocationState, PutModelInvocationRequest,
},
};
use async_trait::async_trait;
use runtime_ports::{
ContextAuditDerivation, ContextAuditSession, ContextAuditSnapshot, ContextAuditStage,
ContextDerivationKind, ConversationMessage, ConversationSession, ModelInvocationAudit,
ModelInvocationAuditState, ModelRole, OperationContext, PortFailure,
};
use runtime_types::{ConversationId, ExecutionId};
use serde_json::Value;
use crate::map_sdk_error;
pub struct AgentInfraConversationSession {
client: Arc<InfraClient>,
conversation_id: ConversationId,
execution_id: ExecutionId,
agent_id: String,
runtime_type: String,
sequence: AtomicUsize,
}
pub struct AgentInfraContextAuditSession {
client: Arc<InfraClient>,
conversation_id: ConversationId,
agent_id: String,
runtime_type: String,
}
impl AgentInfraContextAuditSession {
pub fn new(
client: Arc<InfraClient>,
conversation_id: ConversationId,
agent_id: String,
runtime_type: String,
) -> Self {
Self {
client,
conversation_id,
agent_id,
runtime_type,
}
}
fn conversation_scope(&self) -> ConversationScope {
let mut scope = ConversationScope::conversation(self.conversation_id.as_str());
scope.agent_id = Some(self.agent_id.clone());
scope.runtime_type = Some(self.runtime_type.clone());
scope
}
}
#[async_trait]
impl ContextAuditSession for AgentInfraContextAuditSession {
async fn snapshot(
&self,
operation: &OperationContext,
snapshot: ContextAuditSnapshot,
) -> Result<(), PortFailure> {
let remaining = operation.remaining()?;
let message_count = snapshot.messages.len();
let item_metadata = serde_json::json!({
"requestDigest": snapshot.request_digest,
"protocolOverheadTokens": snapshot.protocol_overhead_tokens,
"reservedOutputTokens": snapshot.reserved_output_tokens,
});
let mut items = snapshot
.messages
.into_iter()
.enumerate()
.map(|(ordinal, message)| {
Ok(ContextSnapshotItem {
id: format!("{}-item-{ordinal}", snapshot.id),
ordinal: u32::try_from(ordinal).map_err(|_| context_audit_overflow())?,
kind: ContextSnapshotItemKind::Message,
role: Some(format!("{:?}", message.role).to_ascii_lowercase()),
content: serde_json::to_value(message).map_err(|error| {
PortFailure::new(
runtime_ports::PortFailureKind::Protocol,
"CONTEXT_AUDIT_SERIALIZATION_FAILED",
error.to_string(),
)
})?,
source_record_id: None,
knowledge_evidence: None,
metadata: item_metadata.clone(),
})
})
.collect::<Result<Vec<_>, PortFailure>>()?;
items.extend(
snapshot
.tools
.into_iter()
.enumerate()
.map(|(index, tool)| {
let ordinal = message_count.saturating_add(index);
Ok(ContextSnapshotItem {
id: format!("{}-tool-{index}", snapshot.id),
ordinal: u32::try_from(ordinal).map_err(|_| context_audit_overflow())?,
kind: ContextSnapshotItemKind::ContextInsight,
role: None,
content: serde_json::to_value(tool).map_err(|error| {
PortFailure::new(
runtime_ports::PortFailureKind::Protocol,
"CONTEXT_AUDIT_SERIALIZATION_FAILED",
error.to_string(),
)
})?,
source_record_id: None,
knowledge_evidence: None,
metadata: item_metadata.clone(),
})
})
.collect::<Result<Vec<_>, PortFailure>>()?,
);
self.client
.context()
.create_snapshot_with_options(
&CreateContextSnapshotRequest {
id: snapshot.id,
conversation: self.conversation_scope(),
stage: match snapshot.stage {
ContextAuditStage::Original => InfraSnapshotStage::Original,
ContextAuditStage::Assembled => InfraSnapshotStage::Assembled,
ContextAuditStage::Compressed => InfraSnapshotStage::Compressed,
},
purpose: snapshot.purpose,
token_count: snapshot.estimated_tokens,
items,
},
CallOptions::default().deadline(remaining),
)
.await
.map(|_| ())
.map_err(crate::map_sdk_error)
}
async fn derivation(
&self,
operation: &OperationContext,
derivation: ContextAuditDerivation,
) -> Result<(), PortFailure> {
let remaining = operation.remaining()?;
self.client
.context()
.create_derivation_with_options(
&CreateContextDerivationRequest {
id: derivation.id,
source_snapshot_id: derivation.source_snapshot_id,
target_snapshot_id: derivation.target_snapshot_id,
kind: match derivation.kind {
ContextDerivationKind::Assemble => InfraDerivationKind::Assemble,
ContextDerivationKind::Compress => InfraDerivationKind::Compress,
ContextDerivationKind::Trim => InfraDerivationKind::Trim,
ContextDerivationKind::Reorder => InfraDerivationKind::Reorder,
},
strategy: "runtime-context-manager".into(),
strategy_version: env!("CARGO_PKG_VERSION").into(),
input_token_count: derivation.input_tokens,
output_token_count: derivation.output_tokens,
input_item_count: u32::try_from(derivation.input_items)
.map_err(|_| context_audit_overflow())?,
output_item_count: u32::try_from(derivation.output_items)
.map_err(|_| context_audit_overflow())?,
metadata: Value::Null,
},
CallOptions::default().deadline(remaining),
)
.await
.map(|_| ())
.map_err(crate::map_sdk_error)
}
async fn model_invocation(
&self,
operation: &OperationContext,
invocation: ModelInvocationAudit,
) -> Result<(), PortFailure> {
let remaining = operation.remaining()?;
self.client
.context()
.put_model_invocation_with_options(
&PutModelInvocationRequest {
id: invocation.id,
conversation_id: self.conversation_id.to_string(),
request_snapshot_id: invocation.request_snapshot_id,
model: invocation.model,
provider: "agent-infra".into(),
state: match invocation.state {
ModelInvocationAuditState::Prepared => InfraInvocationState::Prepared,
ModelInvocationAuditState::Succeeded => InfraInvocationState::Succeeded,
ModelInvocationAuditState::Failed => InfraInvocationState::Failed,
},
response: None,
error_code: invocation.error_code,
usage: serde_json::json!({
"inputTokens": invocation.usage.input_tokens,
"outputTokens": invocation.usage.output_tokens
}),
completed_at_ms: (invocation.state != ModelInvocationAuditState::Prepared)
.then(|| chrono::Utc::now().timestamp_millis()),
},
CallOptions::default().deadline(remaining),
)
.await
.map(|_| ())
.map_err(crate::map_sdk_error)
}
}
fn context_audit_overflow() -> PortFailure {
PortFailure::new(
runtime_ports::PortFailureKind::Invalid,
"CONTEXT_AUDIT_LIMIT_EXCEEDED",
"context audit item count exceeds the Agent Infra contract",
)
}
pub struct NullContextAuditSession;
#[async_trait]
impl ContextAuditSession for NullContextAuditSession {
async fn snapshot(
&self,
_: &OperationContext,
_: ContextAuditSnapshot,
) -> Result<(), PortFailure> {
Ok(())
}
async fn derivation(
&self,
_: &OperationContext,
_: ContextAuditDerivation,
) -> Result<(), PortFailure> {
Ok(())
}
async fn model_invocation(
&self,
_: &OperationContext,
_: ModelInvocationAudit,
) -> Result<(), PortFailure> {
Ok(())
}
}
impl AgentInfraConversationSession {
pub fn new(
client: Arc<InfraClient>,
conversation_id: ConversationId,
execution_id: ExecutionId,
agent_id: String,
runtime_type: String,
) -> Self {
Self {
client,
conversation_id,
execution_id,
agent_id,
runtime_type,
sequence: AtomicUsize::new(0),
}
}
fn conversation_scope(&self) -> ConversationScope {
let mut scope = ConversationScope::conversation(self.conversation_id.as_str());
scope.agent_id = Some(self.agent_id.clone());
scope.runtime_type = Some(self.runtime_type.clone());
scope
}
}
#[async_trait]
impl ConversationSession for AgentInfraConversationSession {
async fn load_recent(
&self,
operation: &OperationContext,
limit: usize,
) -> Result<Vec<ConversationMessage>, PortFailure> {
let remaining = operation.remaining()?;
let records = self
.client
.context()
.recent_messages_with_options(
self.conversation_id.as_str(),
Some(limit.min(1000)),
&["user", "assistant"],
CallOptions::default().deadline(remaining),
)
.await
.map_err(map_sdk_error)?;
Ok(records
.into_iter()
.filter_map(|record| {
let role = match record.role.as_str() {
"user" => ModelRole::User,
"assistant" => ModelRole::Assistant,
_ => return None,
};
let text = match record.content {
Value::String(value) => value,
value => value.to_string(),
};
Some(ConversationMessage { role, text })
})
.collect())
}
async fn append(
&self,
operation: &OperationContext,
messages: &[ConversationMessage],
) -> Result<(), PortFailure> {
let remaining = operation.remaining()?;
if messages.is_empty() {
return Ok(());
}
let records = messages
.iter()
.map(|message| {
let index = self.sequence.fetch_add(1, Ordering::SeqCst);
ConversationRecord {
id: format!("{}-{index}", self.execution_id),
scope: self.conversation_scope(),
role: match message.role {
ModelRole::User => "user",
ModelRole::Assistant => "assistant",
ModelRole::System => "system",
ModelRole::Tool => "tool",
}
.into(),
content: Value::String(message.text.clone()),
name: None,
tool_call_id: None,
sequence: 0,
created_at_ms: chrono::Utc::now().timestamp_millis(),
metadata: Value::Null,
}
})
.collect();
self.client
.context()
.append_messages_with_options(records, CallOptions::default().deadline(remaining))
.await
.map(|_| ())
.map_err(map_sdk_error)
}
}
pub struct NullConversationSession;
#[async_trait]
impl ConversationSession for NullConversationSession {
async fn load_recent(
&self,
_: &OperationContext,
_: usize,
) -> Result<Vec<ConversationMessage>, PortFailure> {
Ok(Vec::new())
}
async fn append(
&self,
_: &OperationContext,
_: &[ConversationMessage],
) -> Result<(), PortFailure> {
Ok(())
}
}