use buffa::Message as _;
use polyc_proto::proto::polychrome::subagent::v1::{SubagentResult, SubagentSpawn};
use crate::signing_role::{RoleTrustSet, SignatureVerdict, SubagentRole, SubagentSigner};
fn spawn_canonical_bytes(spawn: &SubagentSpawn) -> Vec<u8> {
let mut canonical = spawn.clone();
canonical.signature_hex.clear();
canonical.encode_to_vec()
}
fn result_canonical_bytes(result: &SubagentResult) -> Vec<u8> {
let mut canonical = result.clone();
canonical.signature_hex.clear();
canonical.encode_to_vec()
}
pub fn sign_subagent_spawn_into(signer: &SubagentSigner, spawn: &mut SubagentSpawn) {
spawn.signed_by = signer.public_key_bytes();
spawn.signature_hex.clear();
let signature = signer.sign_subagent_spawn(&spawn_canonical_bytes(spawn));
spawn.signature_hex = crate::hex::lower(&signature);
}
pub fn sign_subagent_result_into(signer: &SubagentSigner, result: &mut SubagentResult) {
result.signed_by = signer.public_key_bytes();
result.signature_hex.clear();
let signature = signer.sign_subagent_result(&result_canonical_bytes(result));
result.signature_hex = crate::hex::lower(&signature);
}
#[must_use]
pub fn classify_subagent_spawn(
trust: &RoleTrustSet<SubagentRole>,
spawn: &SubagentSpawn,
) -> SignatureVerdict {
let Some(signature) = crate::hex::decode(&spawn.signature_hex) else {
return SignatureVerdict::Invalid;
};
trust.classify_subagent_spawn(&spawn.signed_by, &spawn_canonical_bytes(spawn), &signature)
}
#[must_use]
pub fn classify_subagent_result(
trust: &RoleTrustSet<SubagentRole>,
result: &SubagentResult,
) -> SignatureVerdict {
let Some(signature) = crate::hex::decode(&result.signature_hex) else {
return SignatureVerdict::Invalid;
};
trust.classify_subagent_result(
&result.signed_by,
&result_canonical_bytes(result),
&signature,
)
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs)]
use super::*;
fn trust(signer: &SubagentSigner) -> RoleTrustSet<SubagentRole> {
RoleTrustSet::current(signer)
}
fn sample_spawn() -> SubagentSpawn {
SubagentSpawn {
sub_agent_id: "call-7".to_owned(),
target_agent_id: "researcher".to_owned(),
resolved_provider: "vertex".to_owned(),
resolved_model: "fable-pro".to_owned(),
task: "find prior art".to_owned(),
..Default::default()
}
}
fn sample_result() -> SubagentResult {
SubagentResult {
sub_agent_id: "call-7".to_owned(),
target_agent_id: "researcher".to_owned(),
succeeded: true,
error: String::new(),
input_tokens: 42,
output_tokens: 7,
first_party: true,
..Default::default()
}
}
#[test]
fn spawn_round_trips() {
let signer = SubagentSigner::from_seed(21);
let mut s = sample_spawn();
sign_subagent_spawn_into(&signer, &mut s);
assert!(!s.signature_hex.is_empty());
assert_eq!(s.signed_by, signer.public_key_bytes());
assert_eq!(
classify_subagent_spawn(&trust(&signer), &s),
SignatureVerdict::Verified
);
}
#[test]
fn spawn_tampered_task_fails() {
let signer = SubagentSigner::from_seed(21);
let mut s = sample_spawn();
sign_subagent_spawn_into(&signer, &mut s);
s.task = "leaked task".to_owned();
assert_eq!(
classify_subagent_spawn(&trust(&signer), &s),
SignatureVerdict::Invalid
);
}
#[test]
fn spawn_from_an_untrusted_signer_reads_untrusted() {
let signer = SubagentSigner::from_seed(21);
let other = SubagentSigner::from_seed(22);
let mut s = sample_spawn();
sign_subagent_spawn_into(&signer, &mut s);
assert_eq!(
classify_subagent_spawn(&trust(&other), &s),
SignatureVerdict::Untrusted
);
}
#[test]
fn spawn_unsigned_fails() {
let signer = SubagentSigner::from_seed(21);
let s = sample_spawn();
assert_eq!(
classify_subagent_spawn(&trust(&signer), &s),
SignatureVerdict::Invalid
);
}
#[test]
fn result_round_trips() {
let signer = SubagentSigner::from_seed(23);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
assert!(!r.signature_hex.is_empty());
assert_eq!(r.signed_by, signer.public_key_bytes());
assert_eq!(
classify_subagent_result(&trust(&signer), &r),
SignatureVerdict::Verified
);
}
#[test]
fn result_tampered_usage_fails() {
let signer = SubagentSigner::from_seed(23);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
r.input_tokens = 999_999;
assert_eq!(
classify_subagent_result(&trust(&signer), &r),
SignatureVerdict::Invalid
);
}
#[test]
fn result_tampered_succeeded_fails() {
let signer = SubagentSigner::from_seed(23);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
r.succeeded = false;
r.error = "worker turn failed: forged".to_owned();
assert_eq!(
classify_subagent_result(&trust(&signer), &r),
SignatureVerdict::Invalid
);
}
#[test]
fn tampered_first_party_fails_verification() {
let signer = SubagentSigner::from_seed(23);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
r.first_party = false;
assert_eq!(
classify_subagent_result(&trust(&signer), &r),
SignatureVerdict::Invalid
);
}
#[test]
fn result_from_an_untrusted_signer_reads_untrusted() {
let signer = SubagentSigner::from_seed(23);
let other = SubagentSigner::from_seed(24);
let mut r = sample_result();
sign_subagent_result_into(&signer, &mut r);
assert_eq!(
classify_subagent_result(&trust(&other), &r),
SignatureVerdict::Untrusted
);
}
#[test]
fn result_unsigned_fails() {
let signer = SubagentSigner::from_seed(23);
let r = sample_result();
assert_eq!(
classify_subagent_result(&trust(&signer), &r),
SignatureVerdict::Invalid
);
}
#[test]
fn garbage_signature_hex_reads_invalid() {
let pinned = trust(&SubagentSigner::from_seed(25));
let mut s = sample_spawn();
s.signature_hex = "not-hex!".to_owned();
assert_eq!(
classify_subagent_spawn(&pinned, &s),
SignatureVerdict::Invalid
);
s.signature_hex = "abcd".to_owned();
assert_eq!(
classify_subagent_spawn(&pinned, &s),
SignatureVerdict::Invalid
);
}
#[test]
fn a_spawn_signature_does_not_verify_as_a_result() {
let signer = SubagentSigner::from_seed(26);
let mut spawn = sample_spawn();
sign_subagent_spawn_into(&signer, &mut spawn);
let replayed = SubagentResult {
sub_agent_id: spawn.sub_agent_id.clone(),
target_agent_id: spawn.target_agent_id.clone(),
signed_by: spawn.signed_by.clone(),
signature_hex: spawn.signature_hex.clone(),
..Default::default()
};
assert_eq!(
classify_subagent_result(&trust(&signer), &replayed),
SignatureVerdict::Invalid
);
assert_eq!(
classify_subagent_spawn(&trust(&signer), &spawn),
SignatureVerdict::Verified
);
let mut result = sample_result();
sign_subagent_result_into(&signer, &mut result);
let replayed = SubagentSpawn {
sub_agent_id: result.sub_agent_id.clone(),
target_agent_id: result.target_agent_id.clone(),
signed_by: result.signed_by.clone(),
signature_hex: result.signature_hex.clone(),
..Default::default()
};
assert_eq!(
classify_subagent_spawn(&trust(&signer), &replayed),
SignatureVerdict::Invalid
);
}
}