test_r::enable!();
#[cfg(test)]
#[cfg(feature = "export_golem_agentic")]
#[test_r::sequential]
#[allow(clippy::disallowed_names)]
mod tests {
use golem_rust::agentic::{
AgentStream, AgentTypeName, Multimodal, MultimodalAdvanced, MultimodalCustom, Schema,
SnapshotData, UnstructuredBinary, UnstructuredText,
};
use golem_rust::agentic::{HttpRequest, HttpResponse, HttpRouter, Principal, create_webhook};
use golem_rust::golem_agentic::golem::agent::common::{
AgentConfigDeclaration, AgentConfigSource, AgentMode, AgentType, AgentTypeKind,
CachePolicy, DurableStreamSlotSource, Snapshotting, SnapshottingConfig,
};
use golem_rust::schema::VariantValuePayload;
use golem_rust::{
AllowedLanguages, AllowedMimeTypes, ConfigSchema, FromSchema, FromWire, IntoSchema,
IntoWire, MultimodalSchema, WireSchema,
};
use golem_rust::{ScheduledTime, SchemaType, SchemaValue, http_router};
use golem_rust::{agent_definition, agent_implementation, agentic::BaseAgent};
use golem_rust_macro::{description, endpoint, prompt, read_only};
use std::fmt::Debug;
use std::sync::atomic::{AtomicUsize, Ordering};
use test_r::test;
fn schema_value<T: IntoSchema>(value: T) -> SchemaValue {
value.to_value()
}
fn decode_schema_value<T: Schema>(value: SchemaValue) -> T {
T::from_schema_value(value, T::get_type()).unwrap()
}
fn wire_input(fields: Vec<SchemaValue>) -> golem_rust::schema::wit::wire::SchemaValueTree {
golem_rust::encode_schema_value(&SchemaValue::Record { fields }).unwrap()
}
#[derive(FromWire, IntoWire, WireSchema)]
#[schema(transparent)]
struct WireOnlyValue(u32);
#[derive(Debug, PartialEq, FromWire, IntoWire, WireSchema)]
#[schema(transparent)]
struct TransparentUnit(());
impl IntoSchema for WireOnlyValue {
fn type_id() -> golem_rust::schema::TypeId {
<u32 as IntoSchema>::type_id()
}
fn register_in(_: &mut golem_rust::schema::SchemaBuilder) -> SchemaType {
panic!("generated descriptors must not build the schema model")
}
fn to_value(&self) -> SchemaValue {
panic!("generated invocation must not encode through the model")
}
}
impl FromSchema for WireOnlyValue {
fn from_value(_: &SchemaValue) -> Result<Self, golem_rust::schema::FromSchemaError> {
panic!("generated invocation must not decode through the model")
}
}
type CallerPrincipal = Principal;
#[agent_definition]
trait DirectDispatch {
fn new(seed: WireOnlyValue, caller: CallerPrincipal) -> Self;
fn run(
&mut self,
caller: CallerPrincipal,
first: WireOnlyValue,
second: Vec<Option<Result<u32, String>>>,
) -> WireOnlyValue;
fn sync_unit(&mut self) -> TransparentUnit;
async fn async_boxed_unit(&mut self) -> Box<TransparentUnit>;
}
struct DirectDispatchImpl {
calls: usize,
}
#[agent_implementation]
impl DirectDispatch for DirectDispatchImpl {
fn new(seed: WireOnlyValue, caller: CallerPrincipal) -> Self {
assert!(matches!(caller, Principal::Anonymous));
Self {
calls: seed.0 as usize,
}
}
fn run(
&mut self,
caller: CallerPrincipal,
first: WireOnlyValue,
second: Vec<Option<Result<u32, String>>>,
) -> WireOnlyValue {
assert!(matches!(caller, Principal::Anonymous));
self.calls += 1;
WireOnlyValue(
first.0
+ second
.into_iter()
.flatten()
.map(|value| value.unwrap_or_else(|s| s.len() as u32))
.sum::<u32>(),
)
}
fn sync_unit(&mut self) -> TransparentUnit {
self.calls += 1;
TransparentUnit(())
}
async fn async_boxed_unit(&mut self) -> Box<TransparentUnit> {
self.calls += 1;
Box::new(TransparentUnit(()))
}
}
#[test]
fn generated_agent_descriptor_uses_wire_schema_and_preserves_principal_alias() {
use golem_rust::golem_agentic::golem::agent::common::{FieldSource, InputSchema};
use golem_rust::schema::wit::wire::SchemaTypeBody;
let descriptor = golem_rust::agentic::get_agent_type_by_name(&AgentTypeName(
"DirectDispatch".to_string(),
))
.unwrap();
let InputSchema::Parameters(fields) = &descriptor.constructor.input_schema;
assert_eq!(
fields
.iter()
.map(|field| field.name.as_str())
.collect::<Vec<_>>(),
["seed", "caller"]
);
assert!(matches!(
descriptor.schema.type_nodes[fields[0].schema as usize].body,
SchemaTypeBody::U32Type(None)
));
assert!(matches!(fields[1].source, FieldSource::AutoInjected(_)));
let reflected = golem_rust::agentic::get_enriched_agent_type_by_name(&AgentTypeName(
"DirectDispatch".to_string(),
))
.unwrap();
assert!(
reflected
.principal_params_in_constructor()
.contains("caller")
);
}
#[test]
async fn generated_agent_dispatch_uses_direct_codecs_and_rejects_invalid_structure_before_calling()
{
use golem_rust::schema::wit::{direct, wire};
#[derive(IntoWire)]
struct Arguments {
first: u32,
second: Vec<Option<Result<u32, String>>>,
}
let input = || {
direct::encode(&Arguments {
first: 7,
second: vec![Some(Err("fail".into())), None, Some(Ok(23))],
})
.unwrap()
};
let mut agent = DirectDispatchImpl::new(WireOnlyValue(0), Principal::Anonymous);
let output = agent
.invoke("run".to_string(), input(), Principal::Anonymous)
.await
.unwrap();
assert_eq!(direct::decode::<u32>(output.value.unwrap()).unwrap(), 34);
for mode in 0..4 {
let mut bad = input();
let wire::SchemaValueNode::RecordValue(fields) =
&mut bad.value_nodes[bad.root as usize]
else {
panic!("record input")
};
match mode {
0 => {
fields.pop();
}
1 => {
fields.push(fields[0]);
}
2 => {
fields[1] = fields[0];
}
_ => {
fields[0] = -1;
}
}
assert!(
agent
.invoke("run".to_string(), bad, Principal::Anonymous)
.await
.is_err()
);
}
assert_eq!(agent.calls, 1);
#[derive(IntoWire)]
struct NoArguments {}
let sync = agent
.invoke(
"sync_unit".to_string(),
direct::encode(&NoArguments {}).unwrap(),
Principal::Anonymous,
)
.await
.unwrap();
assert!(sync.value.is_none());
assert_eq!(
direct::decode_result_payload::<TransparentUnit>(sync.value).unwrap(),
TransparentUnit(())
);
let asynchronous = agent
.invoke(
"async_boxed_unit".to_string(),
direct::encode(&NoArguments {}).unwrap(),
Principal::Anonymous,
)
.await
.unwrap();
assert!(asynchronous.value.is_none());
assert_eq!(
direct::decode_result_payload::<Box<TransparentUnit>>(asynchronous.value).unwrap(),
Box::new(TransparentUnit(()))
);
assert_eq!(agent.calls, 3);
}
#[test]
async fn generated_agent_constructor_uses_direct_codecs_and_checks_record_structure() {
use golem_rust::agentic::with_agent_initiator;
use golem_rust::schema::wit::{direct, wire};
#[derive(IntoWire)]
struct Arguments {
seed: u32,
}
DirectDispatchImpl::__register_agent_type();
let name = AgentTypeName("DirectDispatch".into());
let input = || direct::encode(&Arguments { seed: 19 }).unwrap();
assert!(
with_agent_initiator(
|initiator| async move { initiator.initiate(input(), Principal::Anonymous).await },
&name
)
.await
.is_ok()
);
for mode in 0..4 {
let mut bad = input();
let wire::SchemaValueNode::RecordValue(fields) =
&mut bad.value_nodes[bad.root as usize]
else {
panic!("record input")
};
match mode {
0 => {
fields.clear();
}
1 => {
fields.push(fields[0]);
}
2 => {
fields[0] = -1;
}
_ => {
bad.value_nodes[bad.root as usize] = wire::SchemaValueNode::BoolValue(false);
}
}
assert!(
with_agent_initiator(
|initiator| async move { initiator.initiate(bad, Principal::Anonymous).await },
&name
)
.await
.is_err()
);
}
}
#[derive(Clone, Debug, IntoSchema, FromSchema, FromWire, IntoWire, WireSchema)]
struct Config {
model: String,
}
#[agent_definition]
trait SampleAgent: BaseAgent {
fn new(init: String) -> Self;
fn num(&self, i: String) -> u32;
}
struct SampleAgentImpl {
_id: String,
}
#[agent_implementation]
impl SampleAgent for SampleAgentImpl {
fn new(init: String) -> Self {
SampleAgentImpl { _id: init }
}
fn num(&self, _i: String) -> u32 {
42
}
}
struct AgentWithTypeParameterImpl<T> {
_request_id: T,
}
#[agent_definition]
trait AgentWithTypeParameter<T: Schema + Clone + Debug> {
fn new(init: String) -> Self;
fn num(&self, i: String) -> u32;
}
#[agent_implementation]
impl AgentWithTypeParameter<String> for AgentWithTypeParameterImpl<String> {
fn new(init: String) -> Self {
AgentWithTypeParameterImpl { _request_id: init }
}
fn num(&self, _i: String) -> u32 {
1
}
}
struct AgentWithStaticMethodsImpl;
#[agent_definition]
trait AgentWithStaticMethods {
fn new(init: UserId) -> Self;
#[allow(unused)]
fn foo(&self, param: String) -> String;
#[allow(unused)]
fn bar(&self, param: String) -> String;
fn baz(&self, param: String) -> String;
}
#[agent_implementation]
impl AgentWithStaticMethods for AgentWithStaticMethodsImpl {
fn new(_init: UserId) -> Self {
AgentWithStaticMethodsImpl
}
fn foo(&self, param: String) -> String {
param
}
fn bar(&self, param: String) -> String {
self.foo(param)
}
fn baz(&self, param: String) -> String {
param
}
}
#[agent_definition]
trait AgentWithOnlyConstructor {
fn new(init: UserId) -> Self;
}
struct AgentWithOnlyConstructorImpl;
#[agent_implementation]
impl AgentWithOnlyConstructor for AgentWithOnlyConstructorImpl {
fn new(_init: UserId) -> Self {
AgentWithOnlyConstructorImpl
}
}
#[agent_definition]
trait AgentWithOnlyStaticMethods {
fn new(init: UserId) -> Self;
#[allow(unused)]
fn foo(&self) -> String;
#[allow(unused)]
fn bar(&self, param: String) -> String;
}
struct AgentWithOnlyStaticMethodsImpl;
#[agent_implementation]
impl AgentWithOnlyStaticMethods for AgentWithOnlyStaticMethodsImpl {
fn new(_init: UserId) -> Self {
AgentWithOnlyStaticMethodsImpl
}
fn foo(&self) -> String {
self.bar("foo".to_string())
}
fn bar(&self, param: String) -> String {
param
}
}
#[agent_definition]
trait FooAgent {
fn new(init: UserId) -> Self;
#[allow(unused)]
fn foo(&self) -> String;
fn bar(&self) -> String;
}
struct FooImpl;
#[agent_implementation]
impl FooAgent for FooImpl {
fn new(_init: UserId) -> FooImpl {
FooImpl
}
fn foo(&self) -> String {
"foo".to_string()
}
fn bar(&self) -> String {
"bar".to_string()
}
}
#[agent_definition]
trait Echo {
fn new(init: UserId, llm_config: Config) -> Self;
fn echo_mut(&mut self, message: String) -> String;
fn echo(&self, message: String) -> String;
fn get_id(&self) -> String;
fn echo_result(&self, result: Result<(), ()>) -> Result<(), ()>;
fn echo_result_err(&self, result: Result<(), String>) -> Result<(), String>;
fn echo_result_ok(&self, result: Result<String, ()>) -> Result<String, ()>;
fn echo_option(&self, option: Option<String>) -> Option<String>;
fn echo_multimodal_advanced(
&self,
input: MultimodalAdvanced<TextOrImage>,
) -> MultimodalAdvanced<TextOrImage>;
fn echo_multimodal(&self, input: Multimodal) -> Multimodal;
fn echo_multimodal_custom(
&self,
input: MultimodalCustom<TextOrImage>,
) -> MultimodalCustom<TextOrImage>;
fn echo_unstructured_text(&self, input: UnstructuredText) -> UnstructuredText;
fn echo_unstructured_text_lc(
&self,
input: UnstructuredText<MyLang>,
) -> UnstructuredText<MyLang>;
fn echo_unstructured_binary(
&self,
input: UnstructuredBinary<MyMimeType>,
) -> UnstructuredBinary<MyMimeType>;
}
struct EchoImpl {
_id: UserId,
_llm_config: Config,
}
#[derive(Debug, Clone, AllowedLanguages)]
enum MyLang {
#[code("de")]
German,
En,
}
#[derive(Debug, Clone, AllowedMimeTypes)]
enum MyMimeType {
#[mime_type("text/plain")]
PlainText,
#[mime_type("image/png")]
PngImage,
}
#[test]
fn direct_unstructured_codecs_preserve_roles_restrictions_and_values() {
use golem_rust::schema::wit::{decode_graph, direct};
fn check<
T: Schema + Clone + golem_rust::FromWire + golem_rust::IntoWire + golem_rust::WireSchema,
>(
value: T,
) {
let expected_schema = T::get_type().get_schema_graph().unwrap();
assert_eq!(
decode_graph(&direct::schema::<T>()).unwrap(),
expected_schema
);
let expected = value.clone().to_schema_value().unwrap();
let decoded = direct::decode::<T>(direct::encode(&value).unwrap()).unwrap();
assert_eq!(decoded.to_schema_value().unwrap(), expected);
}
check(UnstructuredText::from_inline("Grüße", MyLang::German));
check(UnstructuredText::from_inline_any("unrestricted"));
check(UnstructuredText::<MyLang>::Url(
"https://example.com/text".to_string(),
));
check(UnstructuredBinary::from_inline(
vec![0, 255, 19],
MyMimeType::PngImage,
));
check(UnstructuredBinary::<MyMimeType>::from_url(
"https://example.com/image",
));
check(UnstructuredBinary::from_inline(
vec![3, 7],
"application/test".to_string(),
));
}
#[test]
fn direct_unstructured_codecs_reject_wrong_payloads_and_restrictions() {
use golem_rust::schema::wit::{direct, wire};
let text =
UnstructuredText::from_inline("bonjour", golem_rust::agentic::AnyLanguage::new("fr"));
assert!(
direct::decode::<UnstructuredText<MyLang>>(direct::encode(&text).unwrap()).is_err()
);
let binary = UnstructuredBinary::from_inline(vec![5], "application/json".to_string());
assert!(
direct::decode::<UnstructuredBinary<MyMimeType>>(direct::encode(&binary).unwrap())
.is_err()
);
for case in [0, 2] {
let tree = || wire::SchemaValueTree {
value_nodes: vec![
wire::SchemaValueNode::UrlValue("https://example.com".to_string()),
wire::SchemaValueNode::VariantValue(wire::VariantValuePayload {
case,
payload: Some(0),
}),
],
root: 1,
};
assert!(direct::decode::<UnstructuredText<MyLang>>(tree()).is_err());
assert!(direct::decode::<UnstructuredBinary<MyMimeType>>(tree()).is_err());
}
let without_mime = wire::SchemaValueTree {
value_nodes: vec![
wire::SchemaValueNode::BinaryValue(wire::BinaryValuePayload {
bytes: vec![41],
mime_type: None,
}),
wire::SchemaValueNode::VariantValue(wire::VariantValuePayload {
case: 0,
payload: Some(0),
}),
],
root: 1,
};
assert!(
matches!(direct::decode::<UnstructuredBinary<MyMimeType>>(without_mime).unwrap(),
UnstructuredBinary::Inline { data, mime_type: MyMimeType::PlainText } if data == [41])
);
}
#[agent_implementation]
impl Echo for EchoImpl {
fn new(id: UserId, llm_config: Config) -> Self {
EchoImpl {
_id: id,
_llm_config: llm_config,
}
}
fn echo_mut(&mut self, message: String) -> String {
format!("Echo: {}", message)
}
fn echo(&self, message: String) -> String {
message.to_string()
}
fn get_id(&self) -> String {
self.get_agent_id()
}
fn echo_result(&self, result: Result<(), ()>) -> Result<(), ()> {
result
}
fn echo_result_err(&self, result: Result<(), String>) -> Result<(), String> {
result
}
fn echo_result_ok(&self, result: Result<String, ()>) -> Result<String, ()> {
result
}
fn echo_option(&self, option: Option<String>) -> Option<String> {
option
}
fn echo_multimodal_advanced(
&self,
input: MultimodalAdvanced<TextOrImage>,
) -> MultimodalAdvanced<TextOrImage> {
input
}
fn echo_multimodal(&self, input: Multimodal) -> Multimodal {
input
}
fn echo_multimodal_custom(
&self,
input: MultimodalCustom<TextOrImage>,
) -> MultimodalCustom<TextOrImage> {
input
}
fn echo_unstructured_text(&self, input: UnstructuredText) -> UnstructuredText {
input
}
fn echo_unstructured_text_lc(
&self,
input: UnstructuredText<MyLang>,
) -> UnstructuredText<MyLang> {
input
}
fn echo_unstructured_binary(
&self,
input: UnstructuredBinary<MyMimeType>,
) -> UnstructuredBinary<MyMimeType> {
input
}
async fn load_snapshot(
_bytes: Vec<u8>,
_context: golem_rust::agentic::SnapshotRestoreContext,
) -> Result<Self, String> {
Ok(Self {
_id: UserId {
id: "restored".to_string(),
},
_llm_config: Config {
model: "restored".to_string(),
},
})
}
async fn save_snapshot(&self) -> Result<Vec<u8>, String> {
Ok(vec![])
}
}
#[agent_definition(mode = "ephemeral")]
trait EchoEphemeralExplicit {
fn new(init: UserId) -> Self;
fn echo(&self, message: String) -> String;
}
struct EchoEphemeralExplicitImpl {
_id: UserId,
}
#[agent_implementation]
impl EchoEphemeralExplicit for EchoEphemeralExplicitImpl {
fn new(id: UserId) -> Self {
EchoEphemeralExplicitImpl { _id: id }
}
fn echo(&self, message: String) -> String {
message.to_string()
}
}
#[agent_definition(ephemeral)]
trait EchoEphemeralShorthand {
fn new(init: UserId) -> Self;
fn echo(&self, message: String) -> String;
}
struct EchoEphemeralShorthandImpl {
_id: UserId,
}
#[agent_implementation]
impl EchoEphemeralShorthand for EchoEphemeralShorthandImpl {
fn new(id: UserId) -> Self {
EchoEphemeralShorthandImpl { _id: id }
}
fn echo(&self, message: String) -> String {
message.to_string()
}
}
#[agent_definition(mode = "durable")]
trait EchoDurableExplicit {
fn new(init: UserId) -> Self;
fn echo(&self, message: String) -> String;
}
struct EchoDurableExplicitImpl {
_id: UserId,
}
#[agent_implementation]
impl EchoDurableExplicit for EchoDurableExplicitImpl {
fn new(id: UserId) -> Self {
EchoDurableExplicitImpl { _id: id }
}
fn echo(&self, message: String) -> String {
message.to_string()
}
}
#[agent_definition]
trait EchoDurableDefault {
fn new(init: UserId) -> Self;
fn echo(&self, message: String) -> String;
}
struct EchoDurableDefaultImpl {
_id: UserId,
}
#[agent_implementation]
impl EchoDurableDefault for EchoDurableDefaultImpl {
fn new(id: UserId) -> Self {
EchoDurableDefaultImpl { _id: id }
}
fn echo(&self, message: String) -> String {
message.to_string()
}
}
#[agent_definition]
trait EchoAsync {
async fn new(init: UserId, llm_config: Config) -> Self;
async fn echo_mut(&mut self, message: String) -> String;
async fn echo(&self, message: String) -> String;
async fn get_id(&self) -> String;
async fn echo_result(&self, result: Result<(), ()>) -> Result<(), ()>;
async fn echo_result_err(&self, result: Result<(), String>) -> Result<(), String>;
async fn echo_result_ok(&self, result: Result<String, ()>) -> Result<String, ()>;
async fn echo_option(&self, option: Option<String>) -> Option<String>;
async fn echo_multimodal_custom(
&self,
input: MultimodalAdvanced<TextOrImage>,
) -> MultimodalAdvanced<TextOrImage>;
async fn echo_multimodal(&self, input: Multimodal) -> Multimodal;
async fn echo_unstructured_text(&self, input: UnstructuredText) -> UnstructuredText;
async fn echo_unstructured_text_lc(
&self,
input: UnstructuredText<MyLang>,
) -> UnstructuredText<MyLang>;
async fn rpc_call(&self, string: String) -> String;
fn rpc_call_trigger(&self, string: String);
fn rpc_call_schedule(&self, string: String);
}
struct EchoAsyncImpl {
id: UserId,
llm_config: Config,
}
#[agent_implementation]
impl EchoAsync for EchoAsyncImpl {
async fn new(id: UserId, llm_config: Config) -> Self {
EchoAsyncImpl { id, llm_config }
}
async fn echo_mut(&mut self, message: String) -> String {
format!("Echo: {}", message)
}
async fn echo(&self, message: String) -> String {
message.to_string()
}
async fn get_id(&self) -> String {
self.get_agent_id()
}
async fn echo_result(&self, result: Result<(), ()>) -> Result<(), ()> {
result
}
async fn echo_result_err(&self, result: Result<(), String>) -> Result<(), String> {
result
}
async fn echo_result_ok(&self, result: Result<String, ()>) -> Result<String, ()> {
result
}
async fn echo_option(&self, option: Option<String>) -> Option<String> {
option
}
async fn echo_multimodal_custom(
&self,
input: MultimodalAdvanced<TextOrImage>,
) -> MultimodalAdvanced<TextOrImage> {
input
}
async fn echo_multimodal(&self, input: Multimodal) -> Multimodal {
input
}
async fn echo_unstructured_text(&self, input: UnstructuredText) -> UnstructuredText {
input
}
async fn echo_unstructured_text_lc(
&self,
input: UnstructuredText<MyLang>,
) -> UnstructuredText<MyLang> {
input
}
async fn rpc_call(&self, string: String) -> String {
let client = EchoClient::get(self.id.clone(), self.llm_config.clone());
client.echo(string).await
}
fn rpc_call_trigger(&self, string: String) {
let client = EchoClient::get(self.id.clone(), self.llm_config.clone());
client.trigger_echo(string);
}
fn rpc_call_schedule(&self, string: String) {
let client = EchoClient::get(self.id.clone(), self.llm_config.clone());
client
.schedule_echo(
string,
ScheduledTime {
seconds: 1,
nanoseconds: 1,
},
)
.expect("schedule accepts a stream-free input");
}
}
#[agent_definition]
trait AgentWithReservedMethods {
fn new(init: String) -> Self;
fn get(&self) -> String;
async fn remote_get(&mut self) -> String;
}
struct AgentWithGetFunctionImpl {}
#[agent_implementation]
impl AgentWithReservedMethods for AgentWithGetFunctionImpl {
fn new(_init: String) -> Self {
AgentWithGetFunctionImpl {}
}
fn get(&self) -> String {
"foo".to_string()
}
async fn remote_get(&mut self) -> String {
let client = AgentWithReservedMethodsClient::get_("increment".to_string());
client.get().await
}
}
mod aliased_agent_api {
use super::agent_definition;
#[agent_definition]
pub trait CanonicallyNamedAgent {
fn new(id: String) -> Self;
fn ping(&self) -> String;
}
}
use aliased_agent_api::CanonicallyNamedAgent as RenamedAgent;
struct AliasedAgentImplementation {
id: String,
}
#[agent_implementation]
impl RenamedAgent for AliasedAgentImplementation {
fn new(id: String) -> Self {
Self { id }
}
fn ping(&self) -> String {
self.id.clone()
}
}
#[derive(IntoSchema, FromSchema, MultimodalSchema, FromWire, IntoWire, WireSchema)]
enum TextOrImage {
Text(String),
Image(Vec<u8>),
}
#[derive(golem_rust::Schema, Clone)]
struct UserId {
id: String,
}
#[allow(clippy::assertions_on_constants)]
#[test] fn test_agent_compilation() {
assert!(true);
}
#[test]
async fn test_aliased_agent_implementation_uses_canonical_agent_name() {
use golem_rust::agentic::{get_agent_type_by_name, with_agent_initiator};
AliasedAgentImplementation::__register_agent_type();
let canonical_name = AgentTypeName("CanonicallyNamedAgent".to_string());
let alias_name = AgentTypeName("RenamedAgent".to_string());
let agent = AliasedAgentImplementation::new("id".to_string());
assert!(get_agent_type_by_name(&canonical_name).is_some());
assert!(get_agent_type_by_name(&alias_name).is_none());
assert_eq!(agent.get_definition().type_name, canonical_name.0);
let initiator_registered_under_canonical_name =
with_agent_initiator(|_| async { true }, &canonical_name).await;
assert!(initiator_registered_under_canonical_name);
}
#[test]
fn multimodal_structured_values_are_ordered_list_variants() {
let value = MultimodalAdvanced::new([
TextOrImage::Image(vec![1, 2, 3]),
TextOrImage::Text("hello".to_string()),
])
.convert_to_schema_value()
.unwrap();
let SchemaValue::List { elements } = value else {
panic!("expected multimodal schema list")
};
assert!(matches!(
elements[0],
SchemaValue::Variant(VariantValuePayload { case: 1, .. })
));
assert!(matches!(
elements[1],
SchemaValue::Variant(VariantValuePayload { case: 0, .. })
));
let duplicate = MultimodalAdvanced::new([
TextOrImage::Text("first".to_string()),
TextOrImage::Text("second".to_string()),
]);
assert!(duplicate.convert_to_schema_value().is_ok());
let missing_image = MultimodalAdvanced::new([TextOrImage::Text("only text".to_string())]);
assert!(missing_image.convert_to_schema_value().is_ok());
}
#[test]
async fn base_agent_invoke_accepts_multimodal_schema_value_inside_parameter_record() {
EchoImpl::__register_agent_type();
let mut agent = EchoImpl::new(
UserId {
id: "id".to_string(),
},
Config {
model: "model".to_string(),
},
);
let input = MultimodalAdvanced::new([
TextOrImage::Image(vec![1, 2, 3]),
TextOrImage::Text("hello".to_string()),
])
.convert_to_schema_value()
.unwrap();
let output = agent
.invoke(
"echo_multimodal_advanced".to_string(),
wire_input(vec![input]),
Principal::Anonymous,
)
.await
.unwrap()
.value
.unwrap();
let SchemaValue::List { elements } = golem_rust::schema::wit::decode_value(output).unwrap()
else {
panic!("expected multimodal output schema list")
};
assert!(matches!(
elements[0],
SchemaValue::Variant(VariantValuePayload { case: 1, .. })
));
assert!(matches!(
elements[1],
SchemaValue::Variant(VariantValuePayload { case: 0, .. })
));
}
#[test]
async fn base_agent_invoke_rejects_extra_parameter_record_fields() {
EchoImpl::__register_agent_type();
let mut agent = EchoImpl::new(
UserId {
id: "id".to_string(),
},
Config {
model: "model".to_string(),
},
);
let result = agent
.invoke(
"echo".to_string(),
wire_input(vec![
SchemaValue::String("hello".to_string()),
SchemaValue::String("extra".to_string()),
]),
Principal::Anonymous,
)
.await;
assert!(result.is_err());
}
#[agent_definition]
trait StreamingAgent {
fn new() -> Self;
async fn combine(
&self,
prefix: String,
first: AgentStream<String>,
second: AgentStream<u32>,
) -> String;
fn forward(&self, stream: AgentStream<String>) -> AgentStream<String>;
}
struct StreamingAgentImpl;
#[agent_implementation]
impl StreamingAgent for StreamingAgentImpl {
fn new() -> Self {
Self
}
async fn combine(
&self,
prefix: String,
first: AgentStream<String>,
second: AgentStream<u32>,
) -> String {
let first = first.into_raw().await.unwrap().take_handle();
let second = second.into_raw().await.unwrap().take_handle();
format!("{prefix}:{first}:{second}")
}
fn forward(&self, stream: AgentStream<String>) -> AgentStream<String> {
stream
}
}
#[test]
fn agent_stream_schema_retains_element_graph() {
#[derive(IntoSchema, FromSchema)]
struct Element {
value: String,
}
let golem_rust::agentic::StructuredSchema::Default(graph) =
<AgentStream<Element> as Schema>::get_type()
else {
panic!("expected value schema")
};
assert!(matches!(
graph.root,
SchemaType::Stream { inner: Some(_), .. }
));
}
#[test]
async fn agent_stream_dispatch_rejects_missing_streams() {
StreamingAgentImpl::__register_agent_type();
let mut agent = StreamingAgentImpl::new();
let scalar = wire_input(vec![SchemaValue::String("p".to_string())]);
assert!(
agent
.invoke("combine".to_string(), scalar, Principal::Anonymous)
.await
.is_err()
);
}
#[cfg(not(target_arch = "wasm32"))]
async fn assert_forwarded_stream_identity(agent: &mut impl BaseAgent) {
use golem_rust::schema::wit::wire::{SchemaValueNode, SchemaValueStream, SchemaValueTree};
let input = SchemaValueTree {
value_nodes: vec![
SchemaValueNode::RecordValue(vec![1]),
SchemaValueNode::StreamValue(unsafe { SchemaValueStream::from_handle(59) }),
],
root: 0,
};
let output = agent
.invoke("forward".to_string(), input, Principal::Anonymous)
.await
.unwrap()
.value
.unwrap();
assert_eq!(output.value_nodes.len(), 1);
match output.value_nodes.into_iter().next().unwrap() {
SchemaValueNode::StreamValue(stream) => assert_eq!(stream.take_handle(), 59),
_ => panic!("expected forwarded stream"),
}
}
#[test]
#[cfg(not(target_arch = "wasm32"))]
async fn agent_stream_dispatch_preserves_forwarded_stream_identity() {
StreamingAgentImpl::__register_agent_type();
let mut agent = StreamingAgentImpl::new();
assert_forwarded_stream_identity(&mut agent).await;
}
type AliasedAgentStream<T> = AgentStream<T>;
#[agent_definition]
trait AliasedStreamingAgent {
fn new() -> Self;
fn forward(&self, stream: AliasedAgentStream<String>) -> AliasedAgentStream<String>;
}
struct AliasedStreamingAgentImpl;
#[agent_implementation]
impl AliasedStreamingAgent for AliasedStreamingAgentImpl {
fn new() -> Self {
Self
}
fn forward(&self, stream: AliasedAgentStream<String>) -> AliasedAgentStream<String> {
stream
}
}
#[test]
#[cfg(not(target_arch = "wasm32"))]
async fn agent_stream_alias_dispatches_through_recursive_value_tree() {
AliasedStreamingAgentImpl::__register_agent_type();
let mut agent = AliasedStreamingAgentImpl::new();
assert_forwarded_stream_identity(&mut agent).await;
}
#[agent_definition]
trait NativeStreamStateShadowAgent {
fn new() -> Self;
fn echo(&self, in_streams: String) -> String;
}
struct NativeStreamStateShadowAgentImpl;
#[agent_implementation]
impl NativeStreamStateShadowAgent for NativeStreamStateShadowAgentImpl {
fn new() -> Self {
Self
}
fn echo(&self, in_streams: String) -> String {
in_streams
}
}
#[test]
async fn method_parameter_named_in_streams_does_not_shadow_abi_state() {
NativeStreamStateShadowAgentImpl::__register_agent_type();
let mut agent = NativeStreamStateShadowAgentImpl::new();
let invocation = agent
.invoke(
"echo".to_string(),
wire_input(vec![SchemaValue::String("payload".to_string())]),
Principal::Anonymous,
)
.await
.expect("a scalar parameter name must not affect native stream validation");
assert_eq!(
golem_rust::schema::wit::direct::decode::<String>(invocation.value.unwrap()).unwrap(),
"payload"
);
}
#[agent_definition]
#[description("a descriptive agent")]
pub trait DescriptiveAgent {
#[description("new constructor")]
#[prompt("new prompt")]
fn new(name: String) -> Self;
#[description("increment description")]
#[prompt("increment prompt")]
fn increment(&mut self) -> u32;
}
struct DescriptiveAgentImpl {}
#[agent_implementation]
impl DescriptiveAgent for DescriptiveAgentImpl {
fn new(_name: String) -> Self {
DescriptiveAgentImpl {}
}
fn increment(&mut self) -> u32 {
1
}
}
#[allow(unused)]
#[derive(ConfigSchema)]
struct ConfigAgentConfigNested {
foo: String,
bar: i32,
#[config_schema(secret)]
nested_secret: golem_rust::agentic::Secret<bool>,
}
#[allow(unused)]
#[derive(ConfigSchema)]
struct ConfigAgentConfig {
url: String,
port: u32,
#[config_schema(nested)]
nested: ConfigAgentConfigNested,
#[config_schema(secret)]
api_key: golem_rust::agentic::Secret<String>,
}
#[agent_definition]
trait ConfigAgent: BaseAgent {
fn new(#[agent_config] config: golem_rust::agentic::Config<ConfigAgentConfig>) -> Self;
}
struct ConfigAgentImpl;
#[agent_implementation]
impl ConfigAgent for ConfigAgentImpl {
fn new(#[agent_config] _config: golem_rust::agentic::Config<ConfigAgentConfig>) -> Self {
Self
}
}
#[test]
fn test_agent_description_and_prompts() {
DescriptiveAgentImpl::__register_agent_type();
let agent_name = AgentTypeName("DescriptiveAgent".to_string());
let agent =
golem_rust::agentic::get_agent_type_by_name(&agent_name).expect("Agent type not found");
let increment_method = agent.methods.first().expect("increment method not found");
assert_eq!(agent.description, "a descriptive agent");
assert_eq!(agent.constructor.description, "new constructor");
assert_eq!(
agent.constructor.prompt_hint,
Some("new prompt".to_string())
);
assert_eq!(increment_method.description, "increment description");
assert_eq!(
increment_method.prompt_hint,
Some("increment prompt".to_string())
);
}
#[test]
fn test_agent_modes() {
use golem_rust::agentic::get_all_agent_types;
let agent_types = get_all_agent_types();
let find_agent = |name: &str| -> Option<AgentType> {
agent_types.iter().find(|a| a.type_name == name).cloned()
};
if let Some(agent) = find_agent("EchoDurableDefault") {
assert_eq!(
&agent.mode,
&AgentMode::Durable,
"EchoDurableDefault should be Durable"
);
}
if let Some(agent) = find_agent("EchoDurableExplicit") {
assert_eq!(
&agent.mode,
&AgentMode::Durable,
"EchoDurableExplicit should be Durable"
);
}
if let Some(agent) = find_agent("EchoEphemeralShorthand") {
assert_eq!(
&agent.mode,
&AgentMode::Ephemeral,
"EchoEphemeralShorthand should be Ephemeral"
);
}
if let Some(agent) = find_agent("EchoEphemeralExplicit") {
assert_eq!(
&agent.mode,
&AgentMode::Ephemeral,
"EchoEphemeralExplicit should be Ephemeral"
);
}
}
#[agent_definition(
mount = "/c/{customer}/{instance}",
phantom_id(path = "instance", optional = true)
)]
trait SelectedPathAgent {
fn new(customer: String) -> Self;
#[endpoint(get = "/read")]
fn read(&self) -> String;
}
struct SelectedPathAgentImpl(String);
#[agent_implementation]
impl SelectedPathAgent for SelectedPathAgentImpl {
fn new(customer: String) -> Self {
Self(customer)
}
fn read(&self) -> String {
self.0.clone()
}
}
#[agent_definition(mount = "/c/{customer}", phantom_id(query = "instance"))]
trait SelectedQueryAgent {
fn new(customer: String) -> Self;
#[endpoint(get = "/read")]
fn read(&self) -> String;
}
struct SelectedQueryAgentImpl(String);
#[agent_implementation]
impl SelectedQueryAgent for SelectedQueryAgentImpl {
fn new(customer: String) -> Self {
Self(customer)
}
fn read(&self) -> String {
self.0.clone()
}
}
#[test]
fn phantom_http_metadata_from_public_macros() {
use golem_rust::golem_agentic::golem::agent::common::{
PhantomIdBinding, PhantomIdBindingDetails,
};
SelectedPathAgentImpl::__register_agent_type();
SelectedQueryAgentImpl::__register_agent_type();
let agents = golem_rust::agentic::get_all_agent_types();
for (name, expected) in [
(
"SelectedPathAgent",
PhantomIdBinding::Path(PhantomIdBindingDetails {
name: "instance".into(),
optional: true,
}),
),
(
"SelectedQueryAgent",
PhantomIdBinding::Query(PhantomIdBindingDetails {
name: "instance".into(),
optional: false,
}),
),
] {
let agent = agents.iter().find(|a| a.type_name == name).unwrap();
let actual = agent
.http_mount
.as_ref()
.unwrap()
.phantom_id_binding
.as_ref()
.unwrap();
assert_eq!(
std::mem::discriminant(actual),
std::mem::discriminant(&expected)
);
let (PhantomIdBinding::Path(actual) | PhantomIdBinding::Query(actual)) = actual;
let (PhantomIdBinding::Path(expected) | PhantomIdBinding::Query(expected)) = expected;
assert_eq!(actual.name, expected.name);
assert_eq!(actual.optional, expected.optional);
let golem_rust::golem_agentic::golem::agent::common::InputSchema::Parameters(fields) =
&agent.constructor.input_schema;
assert_eq!(
fields.iter().map(|f| f.name.as_str()).collect::<Vec<_>>(),
vec!["customer"]
);
}
}
#[agent_definition(
mount = "/chats/{agent-type}/{foo}/{bar}",
webhook_suffix = "/{agent-type}/events/{foo}/{bar}",
auth = true,
phantom_agent = true,
cors = ["https://app.acme.com", "https://staging.acme.com"],
)]
trait ComplexHttpAgent {
fn new(foo: String, bar: String) -> Self;
#[endpoint(
get = "/path-and-header/{resource_id}",
headers("X-Request-ID" = "request_id")
)]
fn path_and_header(&self, resource_id: String, request_id: String) -> String;
#[endpoint(
get = "/greet?l={location}&n={name}",
durable_streams(
input("location", name = "messages"),
output(
"$result",
name = "results",
content_type = "application/vnd.golem.events"
),
allow_external_writes = true,
allow_stream_delete = false,
allow_invocation_delete = false,
max_concurrent_readers_per_stream = 8,
max_append_requests_per_second_per_stream = 25,
)
)]
fn greet1(&self, location: String, name: String) -> String;
#[endpoint(get = "/greet?l={location}&n={name}")]
#[endpoint(get = "/greet?lx={location}&nm={name}", cors = ["*"], auth = true, headers("X-Foo" = "location", "X-Bar" = "name"))]
fn greet2(&self, location: String, name: String) -> String;
#[endpoint(get = "/greet/{name}/{*file_path}")]
fn greet3(&self, name: String, file_path: String) -> String;
#[endpoint(get = "/")]
fn empty_path(&self);
}
struct AgentWithHttpMountImpl {}
#[agent_implementation]
impl ComplexHttpAgent for AgentWithHttpMountImpl {
fn new(_foo: String, _bar: String) -> Self {
AgentWithHttpMountImpl {}
}
fn path_and_header(&self, _resource_id: String, _request_id: String) -> String {
"foo".to_string()
}
fn greet1(&self, _location: String, _name: String) -> String {
"foo".to_string()
}
fn greet2(&self, _location: String, _name: String) -> String {
"bar".to_string()
}
fn greet3(&self, _name: String, _file_path: String) -> String {
"baz".to_string()
}
fn empty_path(&self) {}
}
#[test]
fn test_agent_with_http() {
use golem_rust::agentic::get_all_agent_types;
let agent_types = get_all_agent_types();
let agent = agent_types
.iter()
.find(|a| a.type_name == "ComplexHttpAgent")
.expect("ComplexHttpAgent not found");
assert!(
agent.http_mount.is_some(),
"HTTP mount details should be set"
);
assert_eq!(
agent.http_mount.as_ref().map(|hm| hm.phantom_agent),
Some(true),
"Agent phantom property should be set"
);
assert!(
!agent.methods.is_empty(),
"Agent should have methods defined"
);
assert!(
agent.methods.iter().all(|m| !m.http_endpoint.is_empty()),
"All methods should have HTTP endpoint details"
);
assert!(agent.methods.iter().all(|m| !m.http_endpoint.is_empty()),);
let greet1 = agent
.methods
.iter()
.find(|method| method.name == "greet1")
.unwrap();
let options = greet1.http_endpoint[0].durable_streams.as_ref().unwrap();
assert_eq!(options.slots.len(), 2);
assert!(
matches!(&options.slots[0].source, DurableStreamSlotSource::Input(slot) if slot == "location")
);
assert_eq!(options.slots[0].name.as_deref(), Some("messages"));
assert_eq!(options.slots[0].content_type, None);
assert!(
matches!(&options.slots[1].source, DurableStreamSlotSource::Output(slot) if slot == "$result")
);
assert_eq!(options.slots[1].name.as_deref(), Some("results"));
assert_eq!(
options.slots[1].content_type.as_deref(),
Some("application/vnd.golem.events")
);
assert_eq!(options.allow_external_writes, Some(true));
assert_eq!(options.allow_stream_delete, Some(false));
assert_eq!(options.allow_invocation_delete, Some(false));
let load = options.load.as_ref().unwrap();
assert_eq!(load.max_concurrent_readers_per_stream, Some(8));
assert_eq!(load.max_append_requests_per_second_per_stream, Some(25));
let greet2 = agent
.methods
.iter()
.find(|method| method.name == "greet2")
.unwrap();
assert!(
greet2
.http_endpoint
.iter()
.all(|endpoint| endpoint.durable_streams.is_none())
);
}
#[agent_definition(mount = "/chats/{agent-type}")]
trait SimpleHttpAgent {
fn new() -> Self;
#[endpoint(get = "/green/{name}")]
fn greet(&self, name: String) -> String;
}
struct SimpleHttpAgentImpl;
#[agent_implementation]
impl SimpleHttpAgent for SimpleHttpAgentImpl {
fn new() -> Self {
SimpleHttpAgentImpl
}
fn greet(&self, name: String) -> String {
format!("Hello, {}!", name)
}
}
#[agent_definition]
pub trait AgentWithPrincipalAutoInjection1 {
fn new(name: String, principal: Principal) -> Self;
fn foo(&self, name: String, principal: Principal) -> String;
}
pub struct AgentWithPrincipalAutoInjection1Impl;
#[agent_implementation]
impl AgentWithPrincipalAutoInjection1 for AgentWithPrincipalAutoInjection1Impl {
fn new(_name: String, _principal: Principal) -> Self {
Self
}
fn foo(&self, name: String, _principal: Principal) -> String {
name
}
}
#[agent_definition]
pub trait AgentWithPrincipalAutoInjection2 {
fn new(name: String, text1: u64, principal: Principal, text: String) -> Self;
fn foo(&self, name: String, num: u64, principal: Principal, text: String) -> String;
}
pub struct AgentWithPrincipalAutoInjection2Impl;
#[agent_implementation]
impl AgentWithPrincipalAutoInjection2 for AgentWithPrincipalAutoInjection2Impl {
fn new(_name: String, _text1: u64, _principal: Principal, _text: String) -> Self {
Self
}
fn foo(&self, name: String, _num: u64, _principal: Principal, _text: String) -> String {
name
}
}
#[agent_definition]
pub trait AgentWithPrincipalAutoInjection3 {
fn new(name: String, text1: u64, principal: Principal, text: Option<String>) -> Self;
fn foo(
&self,
name: String,
text1: u64,
principal: Principal,
text: Option<String>,
) -> String;
fn foo_without_principal_1(&self, name: String, num: u64) -> String;
fn foo_without_principal_2(&self, name: String, num: u64, text: String) -> String;
fn foo_without_principal_3(&self, name: String, num: u64, text: Option<String>) -> String;
fn foo_without_principal_4(
&self,
name: String,
num: u64,
text1: Option<String>,
text2: Option<String>,
) -> String;
fn foo_without_principal_5(&self, name: String, num: u64, text: Option<String>) -> String;
}
pub struct AgentWithPrincipalAutoInjection5Impl;
#[agent_implementation]
impl AgentWithPrincipalAutoInjection3 for AgentWithPrincipalAutoInjection5Impl {
fn new(_name: String, _text1: u64, _principal: Principal, _text: Option<String>) -> Self {
Self
}
fn foo(
&self,
name: String,
_text1: u64,
_principal: Principal,
_text: Option<String>,
) -> String {
name
}
fn foo_without_principal_1(&self, name: String, _num: u64) -> String {
name
}
fn foo_without_principal_2(&self, name: String, _num: u64, _text: String) -> String {
name
}
fn foo_without_principal_3(
&self,
name: String,
_num: u64,
_text: Option<String>,
) -> String {
name
}
fn foo_without_principal_4(
&self,
name: String,
_num: u64,
_text1: Option<String>,
_text2: Option<String>,
) -> String {
name
}
fn foo_without_principal_5(
&self,
name: String,
_num: u64,
_text: Option<String>,
) -> String {
name
}
}
#[agent_definition]
pub trait RemoteAgentWithPrincipal {
fn new(name: String, principal: Principal) -> Self;
async fn foo(&self, name: String) -> String;
}
pub struct RemoteAgentWithPrincipalImpl;
#[agent_implementation]
impl RemoteAgentWithPrincipal for RemoteAgentWithPrincipalImpl {
fn new(_name: String, _principal: Principal) -> Self {
Self
}
async fn foo(&self, name: String) -> String {
AgentWithPrincipalAutoInjection1Client::get(name.clone())
.foo(name.clone())
.await;
AgentWithPrincipalAutoInjection2Client::get(name.clone(), 1, "required".into())
.foo(name.clone(), 1, "required".into())
.await;
AgentWithPrincipalAutoInjection3Client::get(
name.clone(),
1,
Some("not-undefined".into()),
)
.foo_without_principal_1("name".into(), 1)
.await;
"finished".into()
}
}
#[agent_definition]
pub trait WebhookAgent {
fn new(name: String) -> Self;
async fn create_webhook_and_trigger(&self) -> String;
}
fn webhook_placeholder(_url: &str) -> String {
"webhook triggered".to_string()
}
pub struct WebhookAgentImpl;
#[agent_implementation]
impl WebhookAgent for WebhookAgentImpl {
fn new(_name: String) -> Self {
Self
}
async fn create_webhook_and_trigger(&self) -> String {
let webhook = create_webhook().expect("webhook creation should be allowed");
webhook_placeholder(webhook.url());
let request = webhook.await;
request.json().unwrap()
}
}
#[agent_definition]
trait AgentSnapshottingDefault {
fn new(init: String) -> Self;
fn echo(&self, message: String) -> String;
}
struct AgentSnapshottingDefaultImpl {
_id: String,
}
#[agent_implementation]
impl AgentSnapshottingDefault for AgentSnapshottingDefaultImpl {
fn new(id: String) -> Self {
Self { _id: id }
}
fn echo(&self, message: String) -> String {
message
}
}
#[agent_definition(snapshotting = "disabled")]
trait AgentSnapshottingDisabled {
fn new(init: String) -> Self;
fn echo(&self, message: String) -> String;
}
struct AgentSnapshottingDisabledImpl {
_id: String,
}
#[agent_implementation]
impl AgentSnapshottingDisabled for AgentSnapshottingDisabledImpl {
fn new(id: String) -> Self {
Self { _id: id }
}
fn echo(&self, message: String) -> String {
message
}
}
#[agent_definition(snapshotting = "enabled")]
trait AgentSnapshottingEnabled {
fn new(init: String) -> Self;
fn echo(&self, message: String) -> String;
}
#[derive(serde::Serialize, serde::Deserialize)]
struct AgentSnapshottingEnabledImpl {
_id: String,
}
#[agent_implementation]
impl AgentSnapshottingEnabled for AgentSnapshottingEnabledImpl {
fn new(id: String) -> Self {
Self { _id: id }
}
fn echo(&self, message: String) -> String {
message
}
}
#[agent_definition(snapshotting = "periodic(5s)")]
trait AgentSnapshottingPeriodic {
fn new(init: String) -> Self;
fn echo(&self, message: String) -> String;
}
#[derive(serde::Serialize, serde::Deserialize)]
struct AgentSnapshottingPeriodicImpl {
_id: String,
}
#[agent_implementation]
impl AgentSnapshottingPeriodic for AgentSnapshottingPeriodicImpl {
fn new(id: String) -> Self {
Self { _id: id }
}
fn echo(&self, message: String) -> String {
message
}
}
#[agent_definition(snapshotting = "every(10)")]
trait AgentSnapshottingEveryN {
fn new(init: String) -> Self;
fn echo(&self, message: String) -> String;
}
#[derive(serde::Serialize, serde::Deserialize)]
struct AgentSnapshottingEveryNImpl {
_id: String,
}
#[agent_implementation]
impl AgentSnapshottingEveryN for AgentSnapshottingEveryNImpl {
fn new(id: String) -> Self {
Self { _id: id }
}
fn echo(&self, message: String) -> String {
message
}
}
#[test]
fn test_agent_snapshotting() {
use golem_rust::agentic::get_all_agent_types;
let agent_types = get_all_agent_types();
let find_agent = |name: &str| -> Option<AgentType> {
agent_types.iter().find(|a| a.type_name == name).cloned()
};
if let Some(agent) = find_agent("AgentSnapshottingDefault") {
assert!(
matches!(agent.snapshotting, Snapshotting::Disabled),
"Default should be Disabled"
);
}
if let Some(agent) = find_agent("AgentSnapshottingDisabled") {
assert!(
matches!(agent.snapshotting, Snapshotting::Disabled),
"Explicit disabled should be Disabled"
);
}
if let Some(agent) = find_agent("AgentSnapshottingEnabled") {
assert!(
matches!(
agent.snapshotting,
Snapshotting::Enabled(SnapshottingConfig::Default)
),
"Enabled should use Default config"
);
}
if let Some(agent) = find_agent("AgentSnapshottingPeriodic") {
assert!(
matches!(
agent.snapshotting,
Snapshotting::Enabled(SnapshottingConfig::Periodic(5000000000))
),
"Periodic should have correct duration"
);
}
if let Some(agent) = find_agent("AgentSnapshottingEveryN") {
assert!(
matches!(
agent.snapshotting,
Snapshotting::Enabled(SnapshottingConfig::EveryNInvocation(10))
),
"EveryNInvocation should have correct count"
);
}
}
static SERIALIZABLE_AGENT_CONSTRUCTIONS: AtomicUsize = AtomicUsize::new(0);
static CUSTOM_SNAPSHOT_AGENT_CONSTRUCTIONS: AtomicUsize = AtomicUsize::new(0);
static CUSTOM_SNAPSHOT_AGENT_RESTORATIONS: AtomicUsize = AtomicUsize::new(0);
#[agent_definition(snapshotting = "enabled")]
trait SerializableAgent {
fn new(name: String) -> Self;
fn get_name(&self) -> String;
}
#[derive(serde::Serialize, serde::Deserialize)]
struct SerializableAgentImpl {
name: String,
}
#[agent_implementation]
impl SerializableAgent for SerializableAgentImpl {
fn new(name: String) -> Self {
SERIALIZABLE_AGENT_CONSTRUCTIONS.fetch_add(1, Ordering::SeqCst);
Self { name }
}
fn get_name(&self) -> String {
self.name.clone()
}
}
#[agent_definition]
trait NonSerializableAgent {
fn new(value: u32) -> Self;
fn get_value(&self) -> u32;
}
#[derive(Debug)]
struct NonSerializableAgentImpl {
value: u32,
}
#[agent_implementation]
impl NonSerializableAgent for NonSerializableAgentImpl {
fn new(value: u32) -> Self {
Self { value }
}
fn get_value(&self) -> u32 {
self.value
}
}
#[agent_definition]
trait CustomSnapshotAgent {
fn new(data: String) -> Self;
fn get_data(&self) -> String;
}
struct CustomSnapshotAgentImpl {
data: String,
}
#[agent_implementation]
impl CustomSnapshotAgent for CustomSnapshotAgentImpl {
fn new(data: String) -> Self {
CUSTOM_SNAPSHOT_AGENT_CONSTRUCTIONS.fetch_add(1, Ordering::SeqCst);
Self { data }
}
fn get_data(&self) -> String {
self.data.clone()
}
async fn load_snapshot(
bytes: Vec<u8>,
_context: golem_rust::agentic::SnapshotRestoreContext,
) -> Result<Self, String> {
CUSTOM_SNAPSHOT_AGENT_RESTORATIONS.fetch_add(1, Ordering::SeqCst);
Ok(Self {
data: String::from_utf8(bytes).map_err(|e| e.to_string())?,
})
}
async fn save_snapshot(&self) -> Result<Vec<u8>, String> {
Ok(self.data.as_bytes().to_vec())
}
}
#[test]
fn test_serializable_agent_auto_json_snapshot() {
use golem_rust::agentic::snapshot_auto::SaveHelper;
let agent = SerializableAgentImpl {
name: "test-agent".to_string(),
};
let helper = SaveHelper(&agent);
let result = helper.snapshot_save();
assert!(result.is_ok(), "Save should succeed for serializable agent");
let SnapshotData::Json(snapshot) = result.unwrap() else {
panic!("expected JSON");
};
let json_value: serde_json::Value = serde_json::from_slice(&snapshot).unwrap();
assert_eq!(json_value["name"], "test-agent");
let agent2 =
<SerializableAgentImpl as SerializableAgent>::__golem_auto_load_snapshot(&snapshot)
.unwrap();
assert_eq!(agent2.name, "test-agent");
}
#[test]
#[allow(clippy::await_holding_refcell_ref)]
async fn test_auto_restore_factory_does_not_call_constructor_and_can_save() {
use golem_rust::agentic::{SnapshotRestoreContext, with_agent_initiator};
SerializableAgentImpl::__register_agent_type();
let agent_type = AgentTypeName("SerializableAgent".to_string());
SERIALIZABLE_AGENT_CONSTRUCTIONS.store(0, Ordering::SeqCst);
let initialized = with_agent_initiator(
|initiator| async move {
initiator
.initiate(
wire_input(vec![SchemaValue::String("created".to_string())]),
Principal::Anonymous,
)
.await
},
&agent_type,
)
.await
.unwrap();
assert_eq!(SERIALIZABLE_AGENT_CONSTRUCTIONS.load(Ordering::SeqCst), 1);
drop(initialized);
SERIALIZABLE_AGENT_CONSTRUCTIONS.store(0, Ordering::SeqCst);
let context = SnapshotRestoreContext {
principal: Principal::Anonymous,
agent_type: agent_type.0.clone(),
parameters: wire_input(vec![SchemaValue::String("created".to_string())]),
phantom_id: None,
};
let restored = with_agent_initiator(
|initiator| async move {
initiator
.restore(
SnapshotData::Json(br#"{"name":"restored"}"#.to_vec()),
context,
)
.await
},
&agent_type,
)
.await
.unwrap();
assert_eq!(SERIALIZABLE_AGENT_CONSTRUCTIONS.load(Ordering::SeqCst), 0);
let saved = restored.agent.borrow().save_snapshot_base().await.unwrap();
let SnapshotData::Json(saved) = saved else {
panic!("expected JSON");
};
assert_eq!(
serde_json::from_slice::<serde_json::Value>(&saved).unwrap()["name"],
"restored"
);
}
#[test]
fn test_non_serializable_agent_snapshot_returns_error() {
use golem_rust::agentic::snapshot_auto::{SaveHelper, SnapshotSaveFallback};
let agent = NonSerializableAgentImpl { value: 42 };
let helper = SaveHelper(&agent);
let result = helper.snapshot_save();
assert!(
result.is_err(),
"Save should fail for non-serializable agent"
);
assert!(result.unwrap_err().contains("not implemented"));
let load_result =
<NonSerializableAgentImpl as NonSerializableAgent>::__golem_auto_load_snapshot(b"{}");
assert_eq!(load_result.unwrap_err(), "snapshotting is disabled");
}
#[test]
#[allow(clippy::await_holding_refcell_ref)]
async fn test_custom_snapshot_agent_uses_custom_methods() {
use golem_rust::agentic::{
BaseAgent, SnapshotRestoreContext, get_principal, get_resolved_agent,
with_agent_initiator,
};
let agent = CustomSnapshotAgentImpl {
data: "hello-world".to_string(),
};
let save_result = agent.save_snapshot_base().await;
assert!(save_result.is_ok());
let SnapshotData::Bytes(snapshot) = save_result.unwrap() else {
panic!("expected bytes");
};
assert_eq!(snapshot, b"hello-world");
CustomSnapshotAgentImpl::__register_agent_type();
CUSTOM_SNAPSHOT_AGENT_CONSTRUCTIONS.store(0, Ordering::SeqCst);
CUSTOM_SNAPSHOT_AGENT_RESTORATIONS.store(0, Ordering::SeqCst);
let agent_type = AgentTypeName("CustomSnapshotAgent".to_string());
let context = || SnapshotRestoreContext {
principal: Principal::Anonymous,
agent_type: agent_type.0.clone(),
parameters: wire_input(vec![SchemaValue::String("original".to_string())]),
phantom_id: None,
};
let failed_context = context();
let context = context();
let failed = with_agent_initiator(
|initiator| async move {
initiator
.restore(SnapshotData::Bytes(vec![0xff]), failed_context)
.await
},
&agent_type,
)
.await;
assert!(failed.is_err());
assert!(get_principal().is_none());
assert!(get_resolved_agent().is_none());
assert_eq!(
CUSTOM_SNAPSHOT_AGENT_CONSTRUCTIONS.load(Ordering::SeqCst),
0
);
let restored = with_agent_initiator(
|initiator| async move {
initiator
.restore(SnapshotData::Bytes(b"new-data".to_vec()), context)
.await
},
&agent_type,
)
.await
.unwrap();
assert_eq!(
CUSTOM_SNAPSHOT_AGENT_CONSTRUCTIONS.load(Ordering::SeqCst),
0
);
assert_eq!(CUSTOM_SNAPSHOT_AGENT_RESTORATIONS.load(Ordering::SeqCst), 2);
let saved = restored.agent.borrow().save_snapshot_base().await.unwrap();
let SnapshotData::Bytes(saved) = saved else {
panic!("expected bytes");
};
assert_eq!(saved, b"new-data");
}
#[test]
fn test_agent_config_entries() {
ConfigAgentImpl::__register_agent_type();
let agent_name = AgentTypeName("ConfigAgent".to_string());
let agent =
golem_rust::agentic::get_agent_type_by_name(&agent_name).expect("Agent type not found");
fn project_for_comparison(
schema: &golem_rust::schema::wit::wire::SchemaGraph,
value: AgentConfigDeclaration,
) -> (AgentConfigSource, Vec<String>, String) {
let body = &schema.type_nodes[value.value_type as usize].body;
let body = match body {
golem_rust::schema::wit::wire::SchemaTypeBody::SecretType(spec) => {
format!("Secret({:?})", schema.type_nodes[spec.inner as usize].body)
}
other => format!("{other:?}"),
};
(value.source, value.path, body)
}
assert_eq!(
agent
.config
.into_iter()
.map(|entry| project_for_comparison(&agent.schema, entry))
.collect::<Vec<_>>(),
vec![
(
AgentConfigSource::Local,
vec!["url".to_string()],
"SchemaTypeBody::StringType".to_string()
),
(
AgentConfigSource::Local,
vec!["port".to_string()],
"SchemaTypeBody::U32Type(None)".to_string()
),
(
AgentConfigSource::Local,
vec!["nested".to_string(), "foo".to_string(),],
"SchemaTypeBody::StringType".to_string()
),
(
AgentConfigSource::Local,
vec!["nested".to_string(), "bar".to_string(),],
"SchemaTypeBody::S32Type(None)".to_string()
),
(
AgentConfigSource::Secret,
vec!["nested".to_string(), "nested_secret".to_string(),],
"Secret(SchemaTypeBody::BoolType)".to_string()
),
(
AgentConfigSource::Secret,
vec!["api_key".to_string()],
"Secret(SchemaTypeBody::StringType)".to_string()
)
]
);
}
#[agent_definition]
trait ReadOnlyAgent: BaseAgent {
fn new(name: String) -> Self;
#[read_only]
fn get_value(&self) -> u32;
#[read_only(cache = "no_cache")]
fn get_value_no_cache(&self) -> u32;
#[read_only(cache = "until_write")]
fn get_value_explicit_until_write(&self) -> u32;
#[read_only(cache = "ttl", ttl = "30s")]
fn get_value_ttl(&self) -> u32;
#[read_only]
fn get_for_user(&self, _principal: Principal) -> u32;
fn increment(&mut self) -> u32;
}
struct ReadOnlyAgentImpl {}
#[agent_implementation]
impl ReadOnlyAgent for ReadOnlyAgentImpl {
fn new(_name: String) -> Self {
ReadOnlyAgentImpl {}
}
fn get_value(&self) -> u32 {
0
}
fn get_value_no_cache(&self) -> u32 {
0
}
fn get_value_explicit_until_write(&self) -> u32 {
0
}
fn get_value_ttl(&self) -> u32 {
0
}
fn get_for_user(&self, _principal: Principal) -> u32 {
0
}
fn increment(&mut self) -> u32 {
1
}
}
#[test]
fn test_read_only_marker() {
ReadOnlyAgentImpl::__register_agent_type();
let agent_name = AgentTypeName("ReadOnlyAgent".to_string());
let agent =
golem_rust::agentic::get_agent_type_by_name(&agent_name).expect("Agent type not found");
let find = |name: &str| {
agent
.methods
.iter()
.find(|m| m.name == name)
.unwrap_or_else(|| panic!("method `{}` not found", name))
};
let m_default = find("get_value");
let m_no_cache = find("get_value_no_cache");
let m_explicit_until_write = find("get_value_explicit_until_write");
let m_ttl = find("get_value_ttl");
let m_principal = find("get_for_user");
let m_increment = find("increment");
let ro_default = m_default.read_only.as_ref().expect("read_only expected");
assert!(matches!(ro_default.cache_policy, CachePolicy::UntilWrite));
assert!(!ro_default.uses_principal);
let ro_no_cache = m_no_cache.read_only.as_ref().expect("read_only expected");
assert!(matches!(ro_no_cache.cache_policy, CachePolicy::NoCache));
assert!(!ro_no_cache.uses_principal);
let ro_explicit_until_write = m_explicit_until_write
.read_only
.as_ref()
.expect("read_only expected");
assert!(matches!(
ro_explicit_until_write.cache_policy,
CachePolicy::UntilWrite
));
assert!(!ro_explicit_until_write.uses_principal);
let ro_ttl = m_ttl.read_only.as_ref().expect("read_only expected");
match ro_ttl.cache_policy {
CachePolicy::Ttl(nanos) => assert_eq!(nanos, 30_000_000_000),
_ => panic!("expected Ttl policy"),
}
assert!(!ro_ttl.uses_principal);
let ro_principal = m_principal.read_only.as_ref().expect("read_only expected");
assert!(matches!(ro_principal.cache_policy, CachePolicy::UntilWrite));
assert!(
ro_principal.uses_principal,
"uses_principal should be derived from the Principal parameter"
);
assert!(
m_increment.read_only.is_none(),
"non-read-only method should have read_only = None"
);
}
#[test]
fn test_from_schema_value_roundtrip_string() {
let original = "hello world".to_string();
let value = schema_value(original.clone());
let recovered = decode_schema_value::<String>(value);
assert_eq!(recovered, original);
}
#[test]
fn test_from_schema_value_roundtrip_u64() {
let original: u64 = 123456789;
let value = schema_value(original);
let recovered = decode_schema_value::<u64>(value);
assert_eq!(recovered, original);
}
#[test]
fn test_from_schema_value_roundtrip_bool() {
for original in [true, false] {
let value = schema_value(original);
let recovered = decode_schema_value::<bool>(value);
assert_eq!(recovered, original);
}
}
#[test]
fn test_from_schema_value_roundtrip_option() {
let some_val: Option<u32> = Some(42);
let value = schema_value(some_val);
let recovered = decode_schema_value::<Option<u32>>(value);
assert_eq!(recovered, Some(42));
let none_val: Option<u32> = None;
let value = schema_value(none_val);
let recovered = decode_schema_value::<Option<u32>>(value);
assert_eq!(recovered, None);
}
#[test]
fn test_from_schema_value_roundtrip_vec() {
let original = vec![1u32, 2, 3, 4, 5];
let value = schema_value(original.clone());
let recovered = decode_schema_value::<Vec<u32>>(value);
assert_eq!(recovered, original);
}
#[derive(IntoSchema, FromSchema, PartialEq, Debug, Clone)]
struct TestStruct {
name: String,
age: u32,
active: bool,
score: f64,
tags: Vec<String>,
}
#[test]
fn test_from_schema_value_roundtrip_struct() {
let original = TestStruct {
name: "test user".to_string(),
age: 42,
active: true,
score: 99.5,
tags: vec!["a".to_string(), "b".to_string()],
};
let value = schema_value(original.clone());
let recovered = decode_schema_value::<TestStruct>(value);
assert_eq!(recovered, original);
}
#[derive(IntoSchema, FromSchema, PartialEq, Debug, Clone)]
struct NestedInner {
x: String,
y: u32,
}
#[derive(IntoSchema, FromSchema, PartialEq, Debug, Clone)]
struct NestedOuter {
id: u64,
items: Vec<NestedInner>,
label: Option<String>,
}
#[test]
fn test_from_schema_value_roundtrip_nested_struct() {
let original = NestedOuter {
id: 999,
items: vec![
NestedInner {
x: "first".to_string(),
y: 1,
},
NestedInner {
x: "second".to_string(),
y: 2,
},
],
label: Some("nested test".to_string()),
};
let value = schema_value(original.clone());
let recovered = decode_schema_value::<NestedOuter>(value);
assert_eq!(recovered, original);
}
#[derive(IntoSchema, FromSchema, PartialEq, Debug, Clone)]
enum TestEnum {
A,
B(u32),
C { name: String, value: bool },
}
#[test]
fn test_from_schema_value_roundtrip_enum() {
for original in [
TestEnum::A,
TestEnum::B(42),
TestEnum::C {
name: "test".to_string(),
value: true,
},
] {
let value = schema_value(original.clone());
let recovered = decode_schema_value::<TestEnum>(value);
assert_eq!(recovered, original);
}
}
#[test]
fn test_from_schema_value_roundtrip_result() {
let ok_val: Result<String, u32> = Ok("success".to_string());
let value = schema_value(ok_val.clone());
let recovered = decode_schema_value::<Result<String, u32>>(value);
assert_eq!(recovered, ok_val);
let err_val: Result<String, u32> = Err(404);
let value = schema_value(err_val.clone());
let recovered = decode_schema_value::<Result<String, u32>>(value);
assert_eq!(recovered, err_val);
}
#[test]
fn output_schema_value_is_not_record_wrapped_for_single_return() {
let value = schema_value("hello".to_string());
assert_eq!(value, SchemaValue::String("hello".to_string()));
let recovered = decode_schema_value::<String>(value);
assert_eq!(recovered, "hello");
}
#[test]
fn unstructured_inline_values_use_schema_native_payloads() {
let text = UnstructuredText::from_inline("hallo", MyLang::German)
.to_schema_value()
.unwrap();
match text {
SchemaValue::Variant(VariantValuePayload { case: 0, payload }) => {
assert_eq!(
*payload.expect("inline text payload"),
SchemaValue::Text(golem_rust::schema::TextValuePayload {
text: "hallo".to_string(),
language: Some("de".to_string()),
})
);
}
other => panic!("expected inline text variant, got {other:?}"),
}
let binary = UnstructuredBinary::from_inline(vec![1, 2, 3], MyMimeType::PngImage)
.to_schema_value()
.unwrap();
match binary {
SchemaValue::Variant(VariantValuePayload { case: 0, payload }) => {
assert_eq!(
*payload.expect("inline binary payload"),
SchemaValue::Binary(golem_rust::schema::BinaryValuePayload {
bytes: vec![1, 2, 3],
mime_type: Some("image/png".to_string()),
})
);
}
other => panic!("expected inline binary variant, got {other:?}"),
}
}
#[test]
fn unstructured_url_values_decode_using_expected_schema() {
let text_value = UnstructuredText::from_url_any("https://example.com/text.txt")
.to_schema_value()
.unwrap();
let recovered = UnstructuredText::<golem_rust::agentic::AnyLanguage>::from_schema_value(
text_value,
UnstructuredText::<golem_rust::agentic::AnyLanguage>::get_type(),
)
.unwrap();
assert!(matches!(
recovered,
UnstructuredText::Url(url) if url == "https://example.com/text.txt"
));
let binary_value = UnstructuredBinary::<String>::from_url("https://example.com/image.png")
.to_schema_value()
.unwrap();
let recovered = UnstructuredBinary::<String>::from_schema_value(
binary_value,
UnstructuredBinary::<String>::get_type(),
)
.unwrap();
assert!(matches!(
recovered,
UnstructuredBinary::Url(url) if url == "https://example.com/image.png"
));
}
#[test]
fn unstructured_text_rejects_unknown_language_payload() {
let value = SchemaValue::Variant(VariantValuePayload {
case: 0,
payload: Some(Box::new(SchemaValue::Text(
golem_rust::schema::TextValuePayload {
text: "bonjour".to_string(),
language: Some("fr".to_string()),
},
))),
});
let err = UnstructuredText::<MyLang>::from_schema_value(
value,
UnstructuredText::<MyLang>::get_type(),
)
.unwrap_err();
assert!(err.contains("Language code 'fr' is not allowed"));
}
#[agent_definition]
trait RawResponseAgent {
fn new(id: String) -> Self;
fn binary(&self, input: UnstructuredBinary<MyMimeType>) -> UnstructuredBinary<MyMimeType>;
fn optional_binary(
&self,
input: Option<UnstructuredBinary<MyMimeType>>,
) -> Option<UnstructuredBinary<MyMimeType>>;
fn fallible_binary(
&self,
input: Result<UnstructuredBinary<MyMimeType>, String>,
) -> Result<UnstructuredBinary<MyMimeType>, String>;
fn fallible_optional_binary(
&self,
input: Result<Option<UnstructuredBinary<MyMimeType>>, String>,
) -> Result<Option<UnstructuredBinary<MyMimeType>>, String>;
fn text(&self, input: UnstructuredText<MyLang>) -> UnstructuredText<MyLang>;
fn optional_text(&self, input: Option<UnstructuredText>) -> Option<UnstructuredText>;
fn fallible_text(
&self,
input: Result<UnstructuredText<MyLang>, String>,
) -> Result<UnstructuredText<MyLang>, String>;
fn fallible_optional_text(
&self,
input: Result<Option<UnstructuredText>, String>,
) -> Result<Option<UnstructuredText>, String>;
}
struct RawResponseAgentImpl;
#[agent_implementation]
impl RawResponseAgent for RawResponseAgentImpl {
fn new(_id: String) -> Self {
RawResponseAgentImpl
}
fn binary(&self, input: UnstructuredBinary<MyMimeType>) -> UnstructuredBinary<MyMimeType> {
input
}
fn optional_binary(
&self,
input: Option<UnstructuredBinary<MyMimeType>>,
) -> Option<UnstructuredBinary<MyMimeType>> {
input
}
fn fallible_binary(
&self,
input: Result<UnstructuredBinary<MyMimeType>, String>,
) -> Result<UnstructuredBinary<MyMimeType>, String> {
input
}
fn fallible_optional_binary(
&self,
input: Result<Option<UnstructuredBinary<MyMimeType>>, String>,
) -> Result<Option<UnstructuredBinary<MyMimeType>>, String> {
input
}
fn text(&self, input: UnstructuredText<MyLang>) -> UnstructuredText<MyLang> {
input
}
fn optional_text(&self, input: Option<UnstructuredText>) -> Option<UnstructuredText> {
input
}
fn fallible_text(
&self,
input: Result<UnstructuredText<MyLang>, String>,
) -> Result<UnstructuredText<MyLang>, String> {
input
}
fn fallible_optional_text(
&self,
input: Result<Option<UnstructuredText>, String>,
) -> Result<Option<UnstructuredText>, String> {
input
}
}
#[allow(dead_code)]
#[golem_rust::agent_client]
trait RawResponseApi {
fn binary(&self, input: UnstructuredBinary<MyMimeType>) -> UnstructuredBinary<MyMimeType>;
fn optional_binary(
&self,
input: Option<UnstructuredBinary<MyMimeType>>,
) -> Option<UnstructuredBinary<MyMimeType>>;
fn fallible_binary(
&self,
input: Result<UnstructuredBinary<MyMimeType>, String>,
) -> Result<UnstructuredBinary<MyMimeType>, String>;
fn fallible_optional_binary(
&self,
input: Result<Option<UnstructuredBinary<MyMimeType>>, String>,
) -> Result<Option<UnstructuredBinary<MyMimeType>>, String>;
fn text(&self, input: UnstructuredText<MyLang>) -> UnstructuredText<MyLang>;
fn optional_text(&self, input: Option<UnstructuredText>) -> Option<UnstructuredText>;
fn fallible_text(
&self,
input: Result<UnstructuredText<MyLang>, String>,
) -> Result<UnstructuredText<MyLang>, String>;
fn fallible_optional_text(
&self,
input: Result<Option<UnstructuredText>, String>,
) -> Result<Option<UnstructuredText>, String>;
}
#[allow(dead_code)]
#[golem_rust::agent_client(type_name = "RawResponseAgent")]
trait CompleteRawResponseApi {
fn new(id: String) -> Self;
fn binary(&self, input: UnstructuredBinary<MyMimeType>) -> UnstructuredBinary<MyMimeType>;
fn optional_binary(
&self,
input: Option<UnstructuredBinary<MyMimeType>>,
) -> Option<UnstructuredBinary<MyMimeType>>;
fn fallible_binary(
&self,
input: Result<UnstructuredBinary<MyMimeType>, String>,
) -> Result<UnstructuredBinary<MyMimeType>, String>;
fn fallible_optional_binary(
&self,
input: Result<Option<UnstructuredBinary<MyMimeType>>, String>,
) -> Result<Option<UnstructuredBinary<MyMimeType>>, String>;
fn text(&self, input: UnstructuredText<MyLang>) -> UnstructuredText<MyLang>;
fn optional_text(&self, input: Option<UnstructuredText>) -> Option<UnstructuredText>;
fn fallible_text(
&self,
input: Result<UnstructuredText<MyLang>, String>,
) -> Result<UnstructuredText<MyLang>, String>;
fn fallible_optional_text(
&self,
input: Result<Option<UnstructuredText>, String>,
) -> Result<Option<UnstructuredText>, String>;
}
fn raw_leaf(ty: &SchemaType) -> &SchemaType {
match ty {
SchemaType::Option { inner, .. } => raw_leaf(inner),
SchemaType::Result { spec, .. } => {
raw_leaf(spec.ok.as_deref().expect("result ok payload"))
}
other => other,
}
}
fn assert_raw_schema_on_both_paths<T: IntoSchema + FromSchema + WireSchema>(
expected_leaf: &SchemaType,
) {
use golem_rust::schema::wit::{decode_graph, direct};
let name = std::any::type_name::<T>();
let model = golem_rust::schema::try_into_schema_graph::<T>().unwrap();
assert_eq!(
decode_graph(&direct::schema::<T>()).unwrap(),
model,
"{name}"
);
assert_eq!(
<T as Schema>::get_type().get_schema_graph().unwrap(),
model,
"{name}"
);
assert!(model.defs.is_empty(), "{name}");
assert_eq!(raw_leaf(&model.root), expected_leaf, "{name}");
}
fn restrictions(all: &[&str]) -> Option<Vec<String>> {
(!all.is_empty()).then(|| all.iter().map(|s| s.to_string()).collect())
}
#[test]
fn raw_unstructured_schemas_are_identical_on_model_and_wire_paths() {
use golem_rust::agentic::unstructured_text_schema_type;
use golem_rust::agentic::{AnyLanguage, unstructured_binary_schema_type};
let binary = unstructured_binary_schema_type(restrictions(
<MyMimeType as golem_rust::agentic::AllowedMimeTypes>::all(),
));
assert_eq!(
binary.metadata().role,
Some(golem_rust::schema::Role::UnstructuredBinary)
);
assert_raw_schema_on_both_paths::<UnstructuredBinary<MyMimeType>>(&binary);
assert_raw_schema_on_both_paths::<Option<UnstructuredBinary<MyMimeType>>>(&binary);
assert_raw_schema_on_both_paths::<Result<UnstructuredBinary<MyMimeType>, String>>(&binary);
assert_raw_schema_on_both_paths::<Result<Option<UnstructuredBinary<MyMimeType>>, String>>(
&binary,
);
let any_binary = unstructured_binary_schema_type(None);
assert_raw_schema_on_both_paths::<UnstructuredBinary<String>>(&any_binary);
assert_raw_schema_on_both_paths::<Result<Option<UnstructuredBinary<String>>, String>>(
&any_binary,
);
let text = unstructured_text_schema_type(restrictions(
<MyLang as golem_rust::agentic::AllowedLanguages>::all(),
));
assert_eq!(
text.metadata().role,
Some(golem_rust::schema::Role::UnstructuredText)
);
assert_raw_schema_on_both_paths::<UnstructuredText<MyLang>>(&text);
assert_raw_schema_on_both_paths::<Option<UnstructuredText<MyLang>>>(&text);
assert_raw_schema_on_both_paths::<Result<UnstructuredText<MyLang>, String>>(&text);
assert_raw_schema_on_both_paths::<Result<Option<UnstructuredText<MyLang>>, String>>(&text);
let any_text = unstructured_text_schema_type(None);
assert_raw_schema_on_both_paths::<UnstructuredText<AnyLanguage>>(&any_text);
assert_raw_schema_on_both_paths::<Option<UnstructuredText>>(&any_text);
assert_raw_schema_on_both_paths::<Result<Option<UnstructuredText>, String>>(&any_text);
}
#[test]
fn raw_unstructured_agent_definition_schemas_match_agent_client_schemas() {
use golem_rust::golem_agentic::golem::agent::common::{InputSchema, OutputSchema};
use golem_rust::schema::wit::GraphDecoder;
fn model_root<T: IntoSchema>() -> SchemaType {
golem_rust::schema::try_into_schema_graph::<T>()
.unwrap()
.root
}
let expected = [
("binary", model_root::<UnstructuredBinary<MyMimeType>>()),
(
"optional_binary",
model_root::<Option<UnstructuredBinary<MyMimeType>>>(),
),
(
"fallible_binary",
model_root::<Result<UnstructuredBinary<MyMimeType>, String>>(),
),
(
"fallible_optional_binary",
model_root::<Result<Option<UnstructuredBinary<MyMimeType>>, String>>(),
),
("text", model_root::<UnstructuredText<MyLang>>()),
("optional_text", model_root::<Option<UnstructuredText>>()),
(
"fallible_text",
model_root::<Result<UnstructuredText<MyLang>, String>>(),
),
(
"fallible_optional_text",
model_root::<Result<Option<UnstructuredText>, String>>(),
),
];
let descriptor = golem_rust::agentic::get_agent_type_by_name(&AgentTypeName(
"RawResponseAgent".to_string(),
))
.expect("RawResponseAgent is registered");
let decoder = GraphDecoder::new(&descriptor.schema).unwrap();
for (name, root) in expected {
let method = descriptor
.methods
.iter()
.find(|method| method.name == name)
.unwrap_or_else(|| panic!("missing method {name}"));
let OutputSchema::Single(output) = method.output_schema else {
panic!("{name} must have a single output");
};
assert_eq!(
decoder.decode_type_at(output).unwrap(),
root,
"{name} output"
);
let InputSchema::Parameters(fields) = &method.input_schema;
assert_eq!(fields.len(), 1, "{name} input");
assert_eq!(
decoder.decode_type_at(fields[0].schema).unwrap(),
root,
"{name} input"
);
}
RawResponseApiClient::client_definition()
.expect("method-only agent_client accepts raw unstructured types");
}
fn assert_raw_value_on_both_paths<T>(value: T)
where
T: IntoSchema + FromSchema + IntoWire + FromWire + WireSchema,
{
use golem_rust::schema::wit::direct;
let name = std::any::type_name::<T>();
let model = value.to_value();
golem_rust::SchemaRef::new(golem_rust::schema::try_into_schema_graph::<T>().unwrap())
.validate_value(&model)
.unwrap_or_else(|error| panic!("{name}: {error:?}"));
assert_eq!(
golem_rust::decode_schema_value(direct::encode(&value).unwrap()).unwrap(),
model,
"{name}"
);
assert_eq!(
direct::decode::<T>(golem_rust::encode_schema_value(&model).unwrap())
.unwrap()
.to_value(),
model,
"{name}"
);
assert_eq!(T::from_value(&model).unwrap().to_value(), model, "{name}");
assert_eq!(
direct::decode::<T>(direct::encode(&value).unwrap())
.unwrap()
.to_value(),
model,
"{name}"
);
assert_eq!(
<T as Schema>::from_schema_value(model.clone(), T::get_type())
.unwrap()
.to_value(),
model,
"{name}"
);
}
#[test]
fn raw_unstructured_values_round_trip_on_model_and_wire_paths() {
use golem_rust::agentic::AnyLanguage;
let png = || UnstructuredBinary::from_inline(vec![0, 255, 7], MyMimeType::PngImage);
let binary_url = || UnstructuredBinary::<MyMimeType>::from_url("https://example.com/b");
assert_raw_value_on_both_paths(png());
assert_raw_value_on_both_paths(binary_url());
assert_raw_value_on_both_paths(Some(png()));
assert_raw_value_on_both_paths(Some(binary_url()));
assert_raw_value_on_both_paths(None::<UnstructuredBinary<MyMimeType>>);
assert_raw_value_on_both_paths(Ok::<_, String>(png()));
assert_raw_value_on_both_paths(Ok::<_, String>(binary_url()));
assert_raw_value_on_both_paths(Err::<UnstructuredBinary<MyMimeType>, _>(
"boom".to_string(),
));
assert_raw_value_on_both_paths(Ok::<_, String>(Some(png())));
assert_raw_value_on_both_paths(Ok::<_, String>(Some(binary_url())));
assert_raw_value_on_both_paths(Ok::<_, String>(None::<UnstructuredBinary<MyMimeType>>));
assert_raw_value_on_both_paths(Err::<Option<UnstructuredBinary<MyMimeType>>, _>(
"boom".to_string(),
));
assert_raw_value_on_both_paths(Ok::<_, String>(Some(UnstructuredBinary::from_inline(
vec![1],
"application/x-test".to_string(),
))));
let german = || UnstructuredText::from_inline("Grüße", MyLang::German);
let text_url = || UnstructuredText::<MyLang>::Url("https://example.com/t".to_string());
assert_raw_value_on_both_paths(german());
assert_raw_value_on_both_paths(text_url());
assert_raw_value_on_both_paths(Some(german()));
assert_raw_value_on_both_paths(None::<UnstructuredText<MyLang>>);
assert_raw_value_on_both_paths(Ok::<_, String>(german()));
assert_raw_value_on_both_paths(Err::<UnstructuredText<MyLang>, _>("boom".to_string()));
assert_raw_value_on_both_paths(Ok::<_, String>(Some(text_url())));
assert_raw_value_on_both_paths(Ok::<_, String>(None::<UnstructuredText<MyLang>>));
assert_raw_value_on_both_paths(Ok::<_, String>(Some(UnstructuredText::from_inline_any(
"no language",
))));
assert_raw_value_on_both_paths(Some(UnstructuredText::from_inline(
"bonjour",
AnyLanguage::new("fr"),
)));
assert_raw_value_on_both_paths(Err::<Option<UnstructuredText>, _>("boom".to_string()));
}
#[test]
fn raw_unstructured_model_decoding_matches_wire_restrictions() {
use golem_rust::schema::wit::direct;
let french = Some(UnstructuredText::from_inline(
"bonjour",
golem_rust::agentic::AnyLanguage::new("fr"),
))
.to_value();
let err = Option::<UnstructuredText<MyLang>>::from_value(&french).unwrap_err();
assert!(
err.to_string()
.contains("Language code 'fr' is not allowed")
);
assert!(
direct::decode::<Option<UnstructuredText<MyLang>>>(
golem_rust::encode_schema_value(&french).unwrap()
)
.is_err()
);
let json = Ok::<_, String>(UnstructuredBinary::from_inline(
vec![5],
"application/json".to_string(),
))
.to_value();
let err = Result::<UnstructuredBinary<MyMimeType>, String>::from_value(&json).unwrap_err();
assert!(
err.to_string()
.contains("'application/json' is not allowed")
);
assert!(
direct::decode::<Result<UnstructuredBinary<MyMimeType>, String>>(
golem_rust::encode_schema_value(&json).unwrap()
)
.is_err()
);
let without_mime = SchemaValue::Variant(VariantValuePayload {
case: 0,
payload: Some(Box::new(SchemaValue::Binary(
golem_rust::schema::BinaryValuePayload {
bytes: vec![41],
mime_type: None,
},
))),
});
assert!(matches!(
UnstructuredBinary::<MyMimeType>::from_value(&without_mime).unwrap(),
UnstructuredBinary::Inline { data, mime_type: MyMimeType::PlainText } if data == [41]
));
for case in [0, 2] {
let value = SchemaValue::Variant(VariantValuePayload {
case,
payload: Some(Box::new(SchemaValue::Url {
url: "https://example.com".to_string(),
})),
});
assert!(UnstructuredText::<MyLang>::from_value(&value).is_err());
assert!(UnstructuredBinary::<MyMimeType>::from_value(&value).is_err());
}
assert!(UnstructuredBinary::<MyMimeType>::from_value(&SchemaValue::Bool(true)).is_err());
}
async fn invoke_raw_response_agent<T: IntoSchema + FromSchema>(method: &str, value: T) -> T {
let mut agent = RawResponseAgentImpl;
let input = golem_rust::encode_schema_value(&SchemaValue::Record {
fields: vec![value.to_value()],
})
.unwrap();
let output = agent
.invoke(method.to_string(), input, Principal::Anonymous)
.await
.unwrap();
let output = golem_rust::decode_schema_value(output.value.expect("raw output")).unwrap();
T::from_value(&output).unwrap()
}
#[test]
async fn raw_unstructured_agent_definition_dispatch_interoperates_with_model_encoding() {
let png = || UnstructuredBinary::from_inline(vec![9, 8, 7], MyMimeType::PngImage);
let cases: Vec<Result<Option<UnstructuredBinary<MyMimeType>>, String>> = vec![
Ok(Some(png())),
Ok(Some(UnstructuredBinary::from_url(
"https://example.com/raw",
))),
Ok(None),
Err("missing".to_string()),
];
for value in cases {
let expected = value.to_value();
let echoed = invoke_raw_response_agent("fallible_optional_binary", value).await;
assert_eq!(echoed.to_value(), expected);
}
let echoed = invoke_raw_response_agent("binary", png()).await;
assert_eq!(echoed.to_value(), png().to_value());
let echoed = invoke_raw_response_agent("optional_binary", Some(png())).await;
assert_eq!(echoed.to_value(), Some(png()).to_value());
let echoed = invoke_raw_response_agent(
"fallible_binary",
Err::<UnstructuredBinary<MyMimeType>, _>("nope".to_string()),
)
.await;
assert_eq!(echoed.unwrap_err(), "nope");
let german = || UnstructuredText::from_inline("hallo", MyLang::German);
let echoed = invoke_raw_response_agent("text", german()).await;
assert_eq!(echoed.to_value(), german().to_value());
let echoed = invoke_raw_response_agent("fallible_text", Ok::<_, String>(german())).await;
assert_eq!(echoed.to_value(), Ok::<_, String>(german()).to_value());
let echoed = invoke_raw_response_agent("optional_text", None::<UnstructuredText>).await;
assert!(echoed.is_none());
let value = Ok::<_, String>(Some(UnstructuredText::from_inline_any("plain")));
let expected = value.to_value();
let echoed = invoke_raw_response_agent("fallible_optional_text", value).await;
assert_eq!(echoed.to_value(), expected);
}
#[derive(IntoSchema, FromSchema, PartialEq, Debug, Clone)]
enum EnumCasePayloadNames {
Unit,
NamedRecord { name: String, value: bool },
MultiTuple(u32, u64),
}
#[test]
fn enum_case_payload_type_name_uses_case_name() {
let graph = golem_rust::schema::try_into_schema_graph::<EnumCasePayloadNames>()
.expect("failed to build schema graph");
let SchemaType::Ref { id, .. } = &graph.root else {
panic!("expected named enum root");
};
let def = graph.lookup(id).expect("missing enum definition");
let SchemaType::Variant { cases, .. } = &def.body else {
panic!("expected variant root node");
};
let named_record = cases
.iter()
.find(|case| case.name == "NamedRecord")
.and_then(|case| case.payload.as_ref())
.expect("missing NamedRecord case payload type");
let SchemaType::Record { fields, .. } = named_record else {
panic!("expected NamedRecord payload to be a record");
};
assert_eq!(fields[0].name, "name");
assert_eq!(fields[1].name, "value");
let multi_tuple = cases
.iter()
.find(|case| case.name == "MultiTuple")
.and_then(|case| case.payload.as_ref())
.expect("missing MultiTuple case payload type");
assert!(matches!(multi_tuple, SchemaType::Tuple { .. }));
}
#[agent_definition(mount = "/test")]
trait AllHttpMethodsAgent {
fn new() -> Self;
#[endpoint(get = "/get")]
fn get_method(&self) -> String;
#[endpoint(post = "/post")]
fn post_method(&self) -> String;
#[endpoint(put = "/put")]
fn put_method(&self) -> String;
#[endpoint(delete = "/delete")]
fn delete_method(&self) -> String;
#[endpoint(patch = "/patch")]
fn patch_method(&self) -> String;
}
struct AllHttpMethodsAgentImpl;
#[agent_implementation]
impl AllHttpMethodsAgent for AllHttpMethodsAgentImpl {
fn new() -> Self {
AllHttpMethodsAgentImpl
}
fn get_method(&self) -> String {
"GET".to_string()
}
fn post_method(&self) -> String {
"POST".to_string()
}
fn put_method(&self) -> String {
"PUT".to_string()
}
fn delete_method(&self) -> String {
"DELETE".to_string()
}
fn patch_method(&self) -> String {
"PATCH".to_string()
}
}
struct HttpRouterAgent;
#[http_router(
name = "HttpRouterAgent",
mount = "/raw",
auth = false,
cors = ["https://allowed.test"],
)]
impl HttpRouter for HttpRouterAgent {
type Config = ();
fn new(_: golem_rust::agentic::Config<Self::Config>) -> Self {
Self
}
async fn handle(&self, request: HttpRequest) -> HttpResponse {
HttpResponse {
status: 200,
headers: request.headers,
body: request.body,
}
}
}
#[test]
fn test_all_http_methods_supported() {
use golem_rust::agentic::get_all_agent_types;
let agent_types = get_all_agent_types();
let agent = agent_types
.iter()
.find(|a| a.type_name == "AllHttpMethodsAgent")
.expect("AllHttpMethodsAgent not found");
assert!(matches!(agent.kind, AgentTypeKind::Regular));
let expected_methods = vec![
("get_method", "HttpMethod::Get"),
("post_method", "HttpMethod::Post"),
("put_method", "HttpMethod::Put"),
("delete_method", "HttpMethod::Delete"),
("patch_method", "HttpMethod::Patch"),
];
for (method_name, expected_http_method) in expected_methods {
let method = agent
.methods
.iter()
.find(|m| m.name == method_name)
.unwrap_or_else(|| panic!("Method {} not found", method_name));
assert_eq!(
method.http_endpoint.len(),
1,
"Method {} should have exactly one HTTP endpoint",
method_name
);
let endpoint = &method.http_endpoint[0];
let method_str = format!("{:?}", endpoint.http_method);
assert_eq!(
method_str, expected_http_method,
"HTTP method mismatch for {}",
method_name
);
}
let router = agent_types
.iter()
.find(|agent| agent.type_name == "HttpRouterAgent")
.expect("HttpRouterAgent not found");
assert!(matches!(router.kind, AgentTypeKind::HttpRouter));
assert!(matches!(router.mode, AgentMode::Ephemeral));
assert!(matches!(router.snapshotting, Snapshotting::Disabled));
let mount = router.http_mount.as_ref().expect("HTTP mount not found");
assert_eq!(
mount.auth_details.as_ref().map(|auth| auth.required),
Some(false)
);
assert_eq!(
mount.cors_options.allowed_patterns,
vec!["https://allowed.test"]
);
let route = router
.methods
.iter()
.find(|method| method.name == "handle")
.expect("handle method not found");
assert_eq!(route.http_endpoint.len(), 1);
assert!(matches!(
route.http_endpoint[0].http_method,
golem_rust::golem_agentic::golem::agent::common::HttpMethod::Any
));
assert!(route.http_endpoint[0].path_suffix.is_empty());
}
static PARTS_CONSTRUCTIONS: AtomicUsize = AtomicUsize::new(0);
static PARTS_RESTORATIONS: AtomicUsize = AtomicUsize::new(0);
#[derive(golem_rust::serde::Serialize, golem_rust::serde::Deserialize)]
struct PartsSavedState {
revision: u32,
}
#[agent_definition(snapshotting = "enabled")]
trait MultipartIndexAgent {
fn new(revision: u32) -> Self;
fn revision(&self) -> u32;
}
struct MultipartIndexAgentImpl {
revision: u32,
index: Vec<u8>,
}
#[agent_implementation]
impl MultipartIndexAgent for MultipartIndexAgentImpl {
fn new(revision: u32) -> Self {
PARTS_CONSTRUCTIONS.fetch_add(1, Ordering::SeqCst);
Self {
revision,
index: vec![],
}
}
fn revision(&self) -> u32 {
self.revision
}
async fn save_snapshot_parts(
&self,
) -> Result<golem_rust::agentic::MultipartSnapshot, String> {
golem_rust::agentic::MultipartSnapshot::from_state(
&PartsSavedState {
revision: self.revision,
},
std::collections::BTreeMap::from([(
"index".to_string(),
golem_rust::agentic::SnapshotPart {
bytes: self.index.clone(),
content_type: "application/octet-stream".to_string(),
},
)]),
)
}
async fn load_snapshot_parts(
snapshot: golem_rust::agentic::MultipartSnapshot,
context: golem_rust::agentic::SnapshotRestoreContext,
) -> Result<Self, String> {
PARTS_RESTORATIONS.fetch_add(1, Ordering::SeqCst);
assert!(matches!(context.principal, Principal::Anonymous));
let saved: PartsSavedState = snapshot.decode_state()?;
let index = snapshot
.require_part("index", "application/octet-stream")?
.to_vec();
Ok(Self {
revision: saved.revision,
index,
})
}
}
#[test]
#[allow(clippy::await_holding_refcell_ref)]
async fn multipart_hooks_restore_fresh_nonserde_live_state_and_reject_wrong_modes() {
use golem_rust::agentic::{SnapshotRestoreContext, with_agent_initiator};
MultipartIndexAgentImpl::__register_agent_type();
PARTS_CONSTRUCTIONS.store(0, Ordering::SeqCst);
PARTS_RESTORATIONS.store(0, Ordering::SeqCst);
let bytes: Vec<u8> = (0..=255).chain([13, 10]).collect();
let instance = MultipartIndexAgentImpl {
revision: 17,
index: bytes.clone(),
};
let saved = instance.save_snapshot_base().await.unwrap();
assert!(matches!(&saved, SnapshotData::Multipart { .. }));
let agent_type = AgentTypeName("MultipartIndexAgent".to_string());
let context = || SnapshotRestoreContext {
principal: Principal::Anonymous,
agent_type: agent_type.0.clone(),
parameters: wire_input(vec![SchemaValue::U32(0)]),
phantom_id: None,
};
for wrong in [
SnapshotData::Json(b"null".to_vec()),
SnapshotData::Bytes(vec![]),
SnapshotData::Multipart {
state: b"1e400".to_vec(),
parts: Default::default(),
},
SnapshotData::Multipart {
state: br#""\uD800""#.to_vec(),
parts: Default::default(),
},
SnapshotData::Multipart {
state: b"{".to_vec(),
parts: Default::default(),
},
] {
let context = context();
let result = with_agent_initiator(
|initiator| async move { initiator.restore(wrong, context).await },
&agent_type,
)
.await;
assert!(result.is_err());
}
assert_eq!(PARTS_RESTORATIONS.load(Ordering::SeqCst), 0);
let context = context();
let restored = with_agent_initiator(
|initiator| async move { initiator.restore(saved, context).await },
&agent_type,
)
.await
.unwrap();
assert_eq!(PARTS_CONSTRUCTIONS.load(Ordering::SeqCst), 0);
assert_eq!(PARTS_RESTORATIONS.load(Ordering::SeqCst), 1);
let saved: golem_rust::agentic::MultipartSnapshot = restored
.agent
.borrow()
.save_snapshot_base()
.await
.unwrap()
.try_into()
.unwrap();
assert_eq!(saved.state, serde_json::json!({ "revision": 17 }));
assert_eq!(
saved
.require_part("index", "application/octet-stream")
.unwrap(),
bytes
);
assert_eq!(restored.agent.borrow().get_definition().methods.len(), 1);
}
}