use std::{future::Future, pin::Pin};
use futures::{Stream, StreamExt as _};
use polyc_rpc_client::{
AgentDialer, AssertedConversationVisibility, DialError, EdgeCredentials, IngressIdentity,
TurnEvent, TurnIngress, user_message,
};
const CONVERSATION_VISIBILITY: AssertedConversationVisibility =
AssertedConversationVisibility::UNKNOWN;
#[derive(Debug, Clone)]
pub struct TurnRequest {
pub conversation_id: String,
pub exec_id: String,
pub source_identity: IngressIdentity,
pub text: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IngressReceiptError {
pub message: String,
pub retryable: bool,
pub content_conflict: bool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IngressReceipt {
pub dispatch_id: String,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TurnOutcome {
Completed {
text: String,
},
InputRequired {
turn_id: String,
request_id: String,
tool_name: String,
prompt: String,
resolve_token: String,
},
Failed {
message: String,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum TurnStreamEvent {
DurablyReceived,
TextDelta(String),
Outcome(TurnOutcome),
}
pub trait TurnRunner: Send + Sync {
fn receive_ingress<'a>(
&'a self,
_req: TurnRequest,
) -> Pin<Box<dyn Future<Output = Result<IngressReceipt, IngressReceiptError>> + Send + 'a>>
{
Box::pin(async {
Ok(IngressReceipt {
dispatch_id: "test-dispatch".to_owned(),
})
})
}
fn run_turn<'a>(
&'a self,
req: TurnRequest,
) -> Pin<Box<dyn Future<Output = TurnOutcome> + Send + 'a>>;
fn run_turn_streaming<'a>(
&'a self,
req: TurnRequest,
) -> Pin<Box<dyn Stream<Item = TurnStreamEvent> + Send + 'a>> {
Box::pin(async_stream::stream! {
yield TurnStreamEvent::DurablyReceived;
yield TurnStreamEvent::Outcome(self.run_turn(req).await);
})
}
}
#[must_use]
pub fn outcome_from_events(events: &[TurnEvent]) -> TurnOutcome {
if let Some(message) = events.iter().find_map(|e| match e {
TurnEvent::TurnFailed { message, .. } => Some(message.clone()),
_ => None,
}) {
return TurnOutcome::Failed { message };
}
if let Some(pending) = events.iter().find_map(|e| match e {
TurnEvent::ApprovalPending {
turn_id,
request_id,
tool_name,
title,
args_json,
reason,
resolve_token,
preview: _,
fire_dispatch: _,
} => Some((
turn_id,
request_id,
tool_name,
title,
args_json,
reason,
resolve_token,
)),
_ => None,
}) {
let (turn_id, request_id, tool_name, title, args_json, reason, resolve_token) = pending;
let label = if title.is_empty() { tool_name } else { title };
let prompt = if reason.is_empty() {
format!("Approval required to run `{label}` with arguments {args_json}")
} else {
format!("Approval required to run `{label}` with arguments {args_json} — {reason}")
};
return TurnOutcome::InputRequired {
turn_id: turn_id.clone(),
request_id: request_id.clone(),
tool_name: tool_name.clone(),
prompt,
resolve_token: resolve_token.clone(),
};
}
let text = events
.iter()
.filter_map(|e| match e {
TurnEvent::TextDelta(t) => Some(t.as_str()),
_ => None,
})
.collect::<Vec<_>>()
.join("");
TurnOutcome::Completed { text }
}
#[derive(Clone, Copy, Debug, Default)]
pub struct UnconfiguredRunner;
const UNCONFIGURED_MESSAGE: &str = "control-plane address is unset \
(POLYCHROME_AGENT_ADDR); this edge serves its Agent Card but cannot run \
message/send tasks until it is set";
impl TurnRunner for UnconfiguredRunner {
fn receive_ingress<'a>(
&'a self,
_req: TurnRequest,
) -> Pin<Box<dyn Future<Output = Result<IngressReceipt, IngressReceiptError>> + Send + 'a>>
{
Box::pin(async {
Err(IngressReceiptError {
message: UNCONFIGURED_MESSAGE.to_owned(),
retryable: false,
content_conflict: false,
})
})
}
fn run_turn<'a>(
&'a self,
_req: TurnRequest,
) -> Pin<Box<dyn Future<Output = TurnOutcome> + Send + 'a>> {
Box::pin(async {
TurnOutcome::Failed {
message: UNCONFIGURED_MESSAGE.to_owned(),
}
})
}
}
#[derive(Clone)]
pub struct AgentDialerRunner {
dialer: AgentDialer,
}
impl AgentDialerRunner {
pub fn new(addr: &str) -> Result<Self, DialError> {
Ok(Self {
dialer: AgentDialer::new(addr)?,
})
}
pub fn with_credentials(addr: &str, creds: EdgeCredentials) -> Result<Self, DialError> {
Ok(Self {
dialer: AgentDialer::with_credentials(addr, creds)?,
})
}
}
fn classify_ingress_error(err: &polyc_rpc_client::DialError) -> IngressReceiptError {
let content_conflict = err.is_already_exists();
IngressReceiptError {
retryable: err.is_retryable(),
content_conflict,
message: if content_conflict {
"source event was already received with different content".to_owned()
} else {
err.to_string()
},
}
}
impl TurnRunner for AgentDialerRunner {
fn receive_ingress<'a>(
&'a self,
req: TurnRequest,
) -> Pin<Box<dyn Future<Output = Result<IngressReceipt, IngressReceiptError>> + Send + 'a>>
{
Box::pin(async move {
self.dialer
.receive_ingress(
TurnIngress::new(
&req.conversation_id,
&req.exec_id,
req.source_identity,
crate::rpc::claimed_namespace(),
polyc_rpc_client::ConversationIdOrigin::AuthenticatedSource,
vec![user_message(&req.text)],
)
.with_conversation_visibility(CONVERSATION_VISIBILITY),
)
.await
.map(|received| IngressReceipt {
dispatch_id: received.dispatch_id().to_owned(),
})
.map_err(|err| classify_ingress_error(&err))
})
}
fn run_turn<'a>(
&'a self,
req: TurnRequest,
) -> Pin<Box<dyn Future<Output = TurnOutcome> + Send + 'a>> {
Box::pin(async move {
let mut stream = self.run_turn_streaming(req);
let mut outcome = None;
while let Some(event) = stream.next().await {
if let TurnStreamEvent::Outcome(o) = event {
outcome = Some(o);
}
}
outcome.unwrap_or_else(|| TurnOutcome::Failed {
message: "turn stream ended without a terminal outcome".to_owned(),
})
})
}
fn run_turn_streaming<'a>(
&'a self,
req: TurnRequest,
) -> Pin<Box<dyn Stream<Item = TurnStreamEvent> + Send + 'a>> {
Box::pin(async_stream::stream! {
let stream = match self
.dialer
.receive_ingress(
TurnIngress::new(
&req.conversation_id,
&req.exec_id,
req.source_identity,
crate::rpc::claimed_namespace(),
polyc_rpc_client::ConversationIdOrigin::AuthenticatedSource,
vec![user_message(&req.text)],
)
.with_conversation_visibility(CONVERSATION_VISIBILITY),
)
.await
{
Ok(received) => {
yield TurnStreamEvent::DurablyReceived;
match self.dialer.attach_ingress(&received).await {
Ok(stream) => stream,
Err(err) => {
yield TurnStreamEvent::Outcome(TurnOutcome::Failed {
message: err.to_string(),
});
return;
}
}
}
Err(err) => {
yield TurnStreamEvent::Outcome(TurnOutcome::Failed {
message: err.to_string(),
});
return;
}
};
futures::pin_mut!(stream);
let mut events = Vec::new();
loop {
match stream.next().await {
Some(Ok(event)) => {
if let TurnEvent::TextDelta(ref text) = event {
yield TurnStreamEvent::TextDelta(text.clone());
}
events.push(event);
}
Some(Err(err)) => {
yield TurnStreamEvent::Outcome(TurnOutcome::Failed {
message: err.to_string(),
});
return;
}
None => break,
}
}
yield TurnStreamEvent::Outcome(outcome_from_events(&events));
})
}
}
pub trait ApprovalResponder: Send + Sync {
fn respond<'a>(
&'a self,
turn_id: &'a str,
request_id: &'a str,
approved: bool,
reason: &'a str,
conversation_id: &'a str,
resolve_token: &'a str,
) -> Pin<Box<dyn Future<Output = Result<bool, String>> + Send + 'a>>;
}
const UNCONFIGURED_APPROVAL_MESSAGE: &str = "control-plane address is unset \
(POLYCHROME_AGENT_ADDR); this edge cannot submit approval decisions until \
it is set";
#[derive(Clone, Copy, Debug, Default)]
pub struct UnconfiguredApprovalResponder;
impl ApprovalResponder for UnconfiguredApprovalResponder {
fn respond<'a>(
&'a self,
_turn_id: &'a str,
_request_id: &'a str,
_approved: bool,
_reason: &'a str,
_conversation_id: &'a str,
_resolve_token: &'a str,
) -> Pin<Box<dyn Future<Output = Result<bool, String>> + Send + 'a>> {
Box::pin(async { Err(UNCONFIGURED_APPROVAL_MESSAGE.to_owned()) })
}
}
#[derive(Clone)]
pub struct ApprovalDialerResponder {
dialer: polyc_rpc_client::ApprovalDialer,
}
impl ApprovalDialerResponder {
pub fn new(addr: &str) -> Result<Self, DialError> {
Ok(Self {
dialer: polyc_rpc_client::ApprovalDialer::new(addr)?,
})
}
pub fn with_bearer(addr: &str, bearer: &str) -> Result<Self, DialError> {
Ok(Self {
dialer: polyc_rpc_client::ApprovalDialer::with_bearer(addr, bearer)?,
})
}
}
impl ApprovalResponder for ApprovalDialerResponder {
fn respond<'a>(
&'a self,
turn_id: &'a str,
request_id: &'a str,
approved: bool,
reason: &'a str,
conversation_id: &'a str,
resolve_token: &'a str,
) -> Pin<Box<dyn Future<Output = Result<bool, String>> + Send + 'a>> {
Box::pin(async move {
let choice = if approved {
polyc_rpc_client::ApprovalChoice::Approve
} else {
polyc_rpc_client::ApprovalChoice::Deny
};
self.dialer
.respond(
turn_id,
request_id,
choice,
reason,
conversation_id,
"",
"",
resolve_token,
None,
)
.await
.map(|outcome| outcome.persisted)
.map_err(|err| err.to_string())
})
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs)]
use super::*;
const TEST_TURN: &str = "00000000-0000-0000-0000-000000000001";
use polyc_rpc_client::TurnFailureKind;
#[test]
fn an_already_exists_refusal_is_a_content_conflict() {
let classified = classify_ingress_error(&polyc_rpc_client::DialError::Connect(
connectrpc::ConnectError::already_exists("recorded digest differs from presented"),
));
assert!(
classified.content_conflict,
"a digest conflict must reach the peer as a bad parameter, not an internal error"
);
}
#[test]
fn a_conflict_does_not_echo_the_digests_to_the_peer() {
let recorded = "1".repeat(64);
let presented = "2".repeat(64);
let classified = classify_ingress_error(&polyc_rpc_client::DialError::Connect(
connectrpc::ConnectError::already_exists(format!(
"command receive:whatever was recorded with digest {recorded}, not {presented}"
)),
));
assert!(
classified.content_conflict,
"precondition: this is the conflict branch"
);
assert!(
!classified.message.contains(&recorded) && !classified.message.contains(&presented),
"the digests must not reach the peer, got: {}",
classified.message
);
}
#[test]
fn an_invalid_argument_refusal_is_not_a_content_conflict() {
let classified = classify_ingress_error(&polyc_rpc_client::DialError::Connect(
connectrpc::ConnectError::invalid_argument("malformed ingress"),
));
assert!(
!classified.content_conflict,
"only ALREADY-EXISTS carries the same-identity/changed-content meaning"
);
}
#[test]
fn this_edge_states_no_conversation_audience() {
assert_eq!(
CONVERSATION_VISIBILITY,
AssertedConversationVisibility::UNKNOWN
);
}
#[test]
fn text_deltas_fold_to_completed() {
let events = vec![
TurnEvent::ToolStarted {
name: "search".to_owned(),
},
TurnEvent::TextDelta("Hello, ".to_owned()),
TurnEvent::TextDelta("world.".to_owned()),
TurnEvent::Done,
];
assert_eq!(
outcome_from_events(&events),
TurnOutcome::Completed {
text: "Hello, world.".to_owned()
}
);
}
#[test]
fn approval_pending_folds_to_input_required() {
let events = vec![
TurnEvent::TextDelta("working on it".to_owned()),
TurnEvent::ApprovalPending {
turn_id: TEST_TURN.to_owned(),
request_id: "call-1".to_owned(),
tool_name: "delete_repo".to_owned(),
title: "Delete repository".to_owned(),
args_json: r#"{"repo":"x"}"#.to_owned(),
reason: "untrusted content in context".to_owned(),
resolve_token: String::new(),
preview: None,
fire_dispatch: false,
},
TurnEvent::Done,
];
match outcome_from_events(&events) {
TurnOutcome::InputRequired {
request_id,
tool_name,
prompt,
..
} => {
assert_eq!(request_id, "call-1");
assert_eq!(tool_name, "delete_repo");
assert!(prompt.contains("Delete repository"));
assert!(prompt.contains(r#"{"repo":"x"}"#));
assert!(
prompt.contains("untrusted content in context"),
"the gate reason must surface in the input-required prompt"
);
}
other => panic!("expected InputRequired, got {other:?}"),
}
}
#[test]
fn turn_failed_folds_to_failed_not_a_silent_empty_completion() {
let events = vec![
TurnEvent::TurnFailed {
kind: TurnFailureKind::RateLimit,
message: "provider returned 429".to_owned(),
},
TurnEvent::Done,
];
assert_eq!(
outcome_from_events(&events),
TurnOutcome::Failed {
message: "provider returned 429".to_owned()
}
);
}
#[test]
fn turn_failed_wins_over_an_approval_pending_in_the_same_batch() {
let events = vec![
TurnEvent::ApprovalPending {
turn_id: TEST_TURN.to_owned(),
request_id: "call-1".to_owned(),
tool_name: "delete_repo".to_owned(),
title: String::new(),
args_json: "{}".to_owned(),
reason: String::new(),
resolve_token: String::new(),
preview: None,
fire_dispatch: false,
},
TurnEvent::TurnFailed {
kind: TurnFailureKind::Other,
message: "durable failure".to_owned(),
},
TurnEvent::Done,
];
assert_eq!(
outcome_from_events(&events),
TurnOutcome::Failed {
message: "durable failure".to_owned()
}
);
}
#[test]
fn empty_stream_is_empty_completed() {
assert_eq!(
outcome_from_events(&[]),
TurnOutcome::Completed {
text: String::new()
}
);
}
#[tokio::test]
async fn unconfigured_runner_fails_naming_the_unset_address() {
let outcome = UnconfiguredRunner
.run_turn(TurnRequest {
conversation_id: "a2a:ctx".to_owned(),
exec_id: "exec".to_owned(),
source_identity: IngressIdentity::reported("a2a:test-peer", "m1").unwrap(),
text: "hello".to_owned(),
})
.await;
match outcome {
TurnOutcome::Failed { message } => assert!(
message.contains("POLYCHROME_AGENT_ADDR"),
"message must name the unset variable: {message}"
),
other => panic!("expected Failed, got {other:?}"),
}
}
#[tokio::test]
async fn unconfigured_approval_responder_fails_naming_the_unset_address() {
let err = UnconfiguredApprovalResponder
.respond(TEST_TURN, "req-1", true, "a2a:peer", "a2a:ctx", "tok-1")
.await
.expect_err("must fail closed");
assert!(
err.contains("POLYCHROME_AGENT_ADDR"),
"message must name the unset variable: {err}"
);
}
struct FixedOutcomeRunner(TurnOutcome);
impl TurnRunner for FixedOutcomeRunner {
fn run_turn<'a>(
&'a self,
_req: TurnRequest,
) -> Pin<Box<dyn Future<Output = TurnOutcome> + Send + 'a>> {
let outcome = self.0.clone();
Box::pin(async move { outcome })
}
}
#[tokio::test]
async fn default_run_turn_streaming_yields_exactly_one_outcome() {
let runner = FixedOutcomeRunner(TurnOutcome::Completed {
text: "42".to_owned(),
});
let req = TurnRequest {
conversation_id: "a2a:ctx".to_owned(),
exec_id: "exec".to_owned(),
source_identity: IngressIdentity::reported("a2a:test-peer", "m1").unwrap(),
text: "hi".to_owned(),
};
let events: Vec<TurnStreamEvent> = runner.run_turn_streaming(req).collect().await;
assert_eq!(
events,
vec![
TurnStreamEvent::DurablyReceived,
TurnStreamEvent::Outcome(TurnOutcome::Completed {
text: "42".to_owned()
})
],
"a runner with no finer-grained stream must acknowledge durability before its outcome"
);
}
}