1use std::{
2 collections::{HashMap, HashSet},
3 sync::{
4 Arc, Mutex, MutexGuard,
5 atomic::{AtomicBool, Ordering},
6 },
7 time::Duration,
8};
9
10trait MutexExt<T> {
13 fn lock_or_recover(&self) -> MutexGuard<'_, T>;
14}
15
16impl<T> MutexExt<T> for Mutex<T> {
17 fn lock_or_recover(&self) -> MutexGuard<'_, T> {
18 self.lock().unwrap_or_else(|e| e.into_inner())
19 }
20}
21
22use futures_util::{SinkExt, StreamExt};
23use serde::{Deserialize, Serialize};
24use serde_json::Value;
25use tokio::{
26 sync::{Notify, mpsc, oneshot},
27 time::{Instant, interval, sleep, sleep_until},
28};
29use tokio_tungstenite::{connect_async, tungstenite::Message as WsMessage};
30use uuid::Uuid;
31
32const SDK_VERSION: &str = env!("CARGO_PKG_VERSION");
33
34use iii_helpers::http::HttpInvocationConfig;
35
36use crate::{
37 channels::{ChannelReader, ChannelWriter, StreamChannelRef},
38 error::Error,
39 protocol::{
40 ErrorBody, FUNCTION_NAMESPACE_CONFLICT, Message, RegisterFunctionMessage,
41 RegisterTriggerInput, RegisterTriggerMessage, RegisterTriggerTypeMessage, TriggerAction,
42 TriggerRequest, TriggerRequestWithMetadata, UnregisterTriggerMessage,
43 UnregisterTriggerTypeMessage, WORKER_NAMESPACE_CONFLICT,
44 },
45 triggers::{Trigger, TriggerConfig, TriggerHandler},
46 types::{
47 Channel, RemoteFunctionData, RemoteFunctionHandlerWithMetadata, RemoteTriggerTypeData,
48 },
49};
50
51use iii_helpers::observability as telemetry;
52use iii_helpers::observability::OtelConfig;
53
54const DEFAULT_TIMEOUT_MS: u64 = 30_000;
55
56#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct WorkerInfo {
59 pub id: String,
60 pub name: Option<String>,
61 pub runtime: Option<String>,
62 pub version: Option<String>,
63 pub os: Option<String>,
64 pub ip_address: Option<String>,
65 pub status: String,
66 pub connected_at_ms: u64,
67 pub function_count: usize,
68 pub functions: Vec<String>,
69 pub active_invocations: usize,
70 #[serde(default)]
71 pub isolation: Option<String>,
72 #[serde(default)]
73 pub namespace: Option<String>,
74}
75
76#[derive(Debug, Clone, Serialize, Deserialize)]
78pub struct FunctionInfo {
79 pub function_id: String,
80 pub description: Option<String>,
81 pub request_format: Option<Value>,
82 pub response_format: Option<Value>,
83 pub metadata: Option<Value>,
84 #[serde(default)]
85 pub namespace: Option<String>,
86}
87
88#[derive(Debug, Clone, Serialize, Deserialize)]
90pub struct TriggerInfo {
91 pub id: String,
92 pub trigger_type: String,
93 pub function_id: String,
94 pub config: Value,
95 pub metadata: Option<Value>,
96 #[serde(default)]
97 pub namespace: Option<String>,
98}
99
100pub struct RegisterTriggerType<H, C = Value, R = Value> {
111 id: String,
112 description: String,
113 handler: H,
114 trigger_request_format: Option<Value>,
115 call_request_format: Option<Value>,
116 _phantom: std::marker::PhantomData<(C, R)>,
117}
118
119impl<H: TriggerHandler> RegisterTriggerType<H> {
120 pub fn new(id: impl Into<String>, description: impl Into<String>, handler: H) -> Self {
121 Self {
122 id: id.into(),
123 description: description.into(),
124 handler,
125 trigger_request_format: None,
126 call_request_format: None,
127 _phantom: std::marker::PhantomData,
128 }
129 }
130}
131
132impl<H: TriggerHandler, C, R> RegisterTriggerType<H, C, R> {
133 pub fn trigger_request_format<T: schemars::JsonSchema + Serialize>(
137 self,
138 ) -> RegisterTriggerType<H, T, R> {
139 RegisterTriggerType {
140 id: self.id,
141 description: self.description,
142 handler: self.handler,
143 trigger_request_format: json_schema_for::<T>(),
144 call_request_format: self.call_request_format,
145 _phantom: std::marker::PhantomData,
146 }
147 }
148
149 pub fn call_request_format<T: schemars::JsonSchema>(self) -> RegisterTriggerType<H, C, T> {
153 RegisterTriggerType {
154 id: self.id,
155 description: self.description,
156 handler: self.handler,
157 trigger_request_format: self.trigger_request_format,
158 call_request_format: json_schema_for::<T>(),
159 _phantom: std::marker::PhantomData,
160 }
161 }
162}
163
164#[derive(Clone)]
170pub struct TriggerTypeRef<C = Value, R = Value> {
171 iii: IIIClient,
172 trigger_type_id: String,
173 _phantom: std::marker::PhantomData<(C, R)>,
174}
175
176impl<C: Serialize, R> TriggerTypeRef<C, R> {
177 pub fn register_trigger(
179 &self,
180 function_id: impl Into<String>,
181 config: C,
182 ) -> Result<Trigger, Error> {
183 self.register_trigger_with_metadata(function_id, config, None)
184 }
185
186 pub fn register_trigger_with_metadata(
193 &self,
194 function_id: impl Into<String>,
195 config: C,
196 metadata: Option<Value>,
197 ) -> Result<Trigger, Error> {
198 let mut input = RegisterTriggerInput::new(
199 self.trigger_type_id.clone(),
200 function_id,
201 serde_json::to_value(config).map_err(|e| Error::Handler(e.to_string()))?,
202 );
203 input.metadata = metadata;
207 input.namespace = self.iii.namespace();
208 self.iii.register_trigger(input)
209 }
210}
211
212impl<C, R> TriggerTypeRef<C, R>
213where
214 R: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
215{
216 pub fn register_function<O, F>(&self, id: impl Into<String>, f: F) -> FunctionRef
219 where
220 O: Serialize + schemars::JsonSchema + Send + 'static,
221 F: Fn(R) -> Result<O, Error> + Send + Sync + 'static,
222 {
223 self.iii.register_function(id, RegisterFunction::new(f))
224 }
225
226 pub fn register_function_async<O, F, Fut>(&self, id: impl Into<String>, f: F) -> FunctionRef
229 where
230 O: Serialize + schemars::JsonSchema + Send + 'static,
231 F: Fn(R) -> Fut + Send + Sync + 'static,
232 Fut: std::future::Future<Output = Result<O, Error>> + Send + 'static,
233 {
234 self.iii
235 .register_function(id, RegisterFunction::new_async(f))
236 }
237}
238
239#[derive(Debug, Clone, Serialize, Deserialize, Default)]
241pub struct TelemetryOptions {
242 #[serde(skip_serializing_if = "Option::is_none")]
244 pub language: Option<String>,
245 #[serde(skip_serializing_if = "Option::is_none")]
247 pub project_name: Option<String>,
248 #[serde(skip_serializing_if = "Option::is_none")]
250 pub framework: Option<String>,
251 #[serde(skip_serializing_if = "Option::is_none")]
253 pub amplitude_api_key: Option<String>,
254}
255
256#[derive(Debug, Clone, Serialize, Deserialize)]
258pub struct WorkerMetadata {
259 pub runtime: String,
260 pub version: String,
261 pub name: String,
265 pub os: String,
266 #[serde(skip_serializing_if = "Option::is_none")]
269 pub description: Option<String>,
270 #[serde(skip_serializing_if = "Option::is_none")]
271 pub pid: Option<u32>,
272 #[serde(skip_serializing_if = "Option::is_none")]
273 pub telemetry: Option<TelemetryOptions>,
274 #[serde(skip_serializing_if = "Option::is_none")]
275 pub isolation: Option<String>,
276 #[serde(skip_serializing_if = "Option::is_none")]
281 pub namespace: Option<String>,
282}
283
284impl Default for WorkerMetadata {
285 fn default() -> Self {
286 let hostname = hostname::get()
287 .map(|h| h.to_string_lossy().to_string())
288 .unwrap_or_else(|_| "unknown".to_string());
289 let pid = std::process::id();
290 let os_info = format!(
291 "{} {} ({})",
292 std::env::consts::OS,
293 std::env::consts::ARCH,
294 std::env::consts::FAMILY
295 );
296
297 let language = std::env::var("LANG")
298 .or_else(|_| std::env::var("LC_ALL"))
299 .ok()
300 .filter(|s| !s.is_empty())
301 .map(|s| s.split('.').next().unwrap_or(&s).to_string());
302
303 let project_name = detect_project_name(None);
304
305 Self {
306 runtime: "rust".to_string(),
307 version: SDK_VERSION.to_string(),
308 name: format!("{}:{}", hostname, pid),
309 os: os_info,
310 description: None,
311 pid: Some(pid),
312 telemetry: Some(TelemetryOptions {
313 language,
314 project_name,
315 ..Default::default()
316 }),
317 isolation: std::env::var("III_ISOLATION")
318 .ok()
319 .filter(|s| !s.is_empty()),
320 namespace: None,
321 }
322 }
323}
324
325#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
327pub enum WorkerIdentityMode {
328 #[default]
331 Managed,
332 Explicit,
336}
337
338fn worker_name_from_env() -> Option<String> {
339 std::env::var("III_WORKER_NAME")
340 .ok()
341 .filter(|name| !name.is_empty())
342}
343
344fn apply_managed_identity_from_env(metadata: &mut WorkerMetadata) {
345 if let Some(managed_name) = worker_name_from_env() {
346 metadata.name = managed_name;
347 }
348 if metadata.namespace.is_none() {
349 metadata.namespace = std::env::var("III_NAMESPACE")
350 .ok()
351 .filter(|namespace| !namespace.trim().is_empty());
352 }
353}
354
355pub(crate) fn reject_blank_namespace(declared: &str, source: &str) {
387 if declared.trim().is_empty() {
388 panic!(
389 "namespace is empty: `{source}` was set to {declared:?}. \
390 Give it a name, or leave it unset to register in `default`."
391 );
392 }
393}
394
395pub(crate) fn call_namespace(
407 explicit: Option<String>,
408 worker: Option<String>,
409 source: &str,
410) -> Result<Option<String>, Error> {
411 match explicit {
412 Some(declared) if declared.trim().is_empty() => Err(Error::Handler(format!(
413 "namespace is empty: `{source}` was set to {declared:?}. Give it a name, or leave \
414 it unset to stay in this worker's namespace."
415 ))),
416 Some(declared) => Ok(Some(declared)),
417 None => Ok(worker),
418 }
419}
420
421pub(crate) fn invocation_namespace(
426 explicit: Option<String>,
427 worker: Option<String>,
428 function_id: &str,
429 source: &str,
430) -> Result<Option<String>, Error> {
431 let is_implicit = explicit.is_none();
432 let namespace = call_namespace(explicit, worker, source)?;
433
434 if is_implicit && function_id.starts_with("engine::") {
435 Ok(Some("default".to_string()))
436 } else {
437 Ok(namespace)
438 }
439}
440
441pub(crate) fn resolve_namespace(explicit: Option<String>) -> Option<String> {
442 if let Some(declared) = explicit {
443 reject_blank_namespace(&declared, "InitOptions.namespace");
444 return Some(declared);
445 }
446
447 std::env::var("III_NAMESPACE")
453 .ok()
454 .filter(|managed| !managed.trim().is_empty())
455}
456
457pub(crate) fn detect_project_name(cwd: Option<std::path::PathBuf>) -> Option<String> {
465 let cwd = cwd.or_else(|| std::env::current_dir().ok())?;
466
467 let manifest = cwd.join("Cargo.toml");
468 if let Ok(content) = std::fs::read_to_string(&manifest) {
469 if let Some(name) = parse_cargo_package_name(&content) {
470 return Some(name);
471 }
472 }
473
474 cwd.file_name()
475 .and_then(|n| n.to_str())
476 .map(|s| s.trim().to_string())
477 .filter(|s| !s.is_empty())
478}
479
480fn parse_cargo_package_name(content: &str) -> Option<String> {
483 let mut in_package = false;
484 for line in content.lines() {
485 let trimmed = line.trim();
486 if let Some(stripped) = trimmed.strip_prefix('[') {
487 in_package = stripped.trim_end_matches(']').trim() == "package";
488 continue;
489 }
490 if !in_package {
491 continue;
492 }
493 let Some(rest) = trimmed.strip_prefix("name") else {
494 continue;
495 };
496 let rest = rest.trim_start();
497 let Some(rest) = rest.strip_prefix('=') else {
498 continue;
499 };
500 let rest = rest.trim().strip_prefix('"')?;
501 let end = rest.find('"')?;
502 let name = rest[..end].trim();
503 if !name.is_empty() {
504 return Some(name.to_string());
505 }
506 }
507 None
508}
509
510#[allow(clippy::large_enum_variant)]
511enum Outbound {
512 Message(Message),
513 Shutdown,
514}
515
516type PendingInvocation = oneshot::Sender<Result<Value, Error>>;
517
518type WsTx = futures_util::stream::SplitSink<
520 tokio_tungstenite::WebSocketStream<tokio_tungstenite::MaybeTlsStream<tokio::net::TcpStream>>,
521 WsMessage,
522>;
523
524fn inject_trace_headers() -> (Option<String>, Option<String>) {
526 use iii_helpers::observability as context;
527 (context::inject_traceparent(), context::inject_baggage())
528}
529
530#[derive(Debug, Clone, Copy, PartialEq, Eq)]
532pub enum IIIConnectionState {
533 Disconnected,
534 Connecting,
535 Connected,
536 Reconnecting,
537 Failed,
538}
539
540#[derive(Clone)]
541pub struct FunctionRef {
542 pub id: String,
543 unregister_fn: Arc<dyn Fn() + Send + Sync>,
544}
545
546impl FunctionRef {
547 pub fn unregister(&self) {
548 (self.unregister_fn)();
549 }
550}
551
552fn json_schema_for<T: schemars::JsonSchema>() -> Option<Value> {
553 serde_json::to_value(
554 schemars::r#gen::SchemaSettings::draft07()
555 .into_generator()
556 .into_root_schema_for::<T>(),
557 )
558 .ok()
559}
560
561#[doc(hidden)]
564pub trait IntoSyncHandler<Marker>: Send + Sync + 'static {
565 fn into_handler(self) -> RemoteFunctionHandlerWithMetadata;
566 fn request_format() -> Option<Value> {
567 None
568 }
569 fn response_format() -> Option<Value> {
570 None
571 }
572}
573
574impl<F, T, R> IntoSyncHandler<(T, R)> for F
582where
583 F: Fn(T) -> Result<R, Error> + Send + Sync + 'static,
584 T: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
585 R: serde::Serialize + schemars::JsonSchema + Send + 'static,
586{
587 fn into_handler(self) -> RemoteFunctionHandlerWithMetadata {
588 Arc::new(move |input: Value, _metadata: Option<Value>| {
589 let output = serde_json::from_value::<T>(input)
590 .map_err(|e| Error::Serde(e.to_string()))
591 .and_then(&self)
592 .and_then(|val| {
593 serde_json::to_value(&val).map_err(|e| Error::Serde(e.to_string()))
594 });
595 Box::pin(async move { output })
596 })
597 }
598
599 fn request_format() -> Option<Value> {
600 json_schema_for::<T>()
601 }
602
603 fn response_format() -> Option<Value> {
604 json_schema_for::<R>()
605 }
606}
607
608impl<F, T, R> IntoSyncHandler<(T, Option<Value>, R)> for F
611where
612 F: Fn(T, Option<Value>) -> Result<R, Error> + Send + Sync + 'static,
613 T: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
614 R: serde::Serialize + schemars::JsonSchema + Send + 'static,
615{
616 fn into_handler(self) -> RemoteFunctionHandlerWithMetadata {
617 Arc::new(move |input: Value, metadata: Option<Value>| {
618 let output = serde_json::from_value::<T>(input)
619 .map_err(|e| Error::Serde(e.to_string()))
620 .and_then(|arg| self(arg, metadata))
621 .and_then(|val| {
622 serde_json::to_value(&val).map_err(|e| Error::Serde(e.to_string()))
623 });
624 Box::pin(async move { output })
625 })
626 }
627
628 fn request_format() -> Option<Value> {
629 json_schema_for::<T>()
630 }
631
632 fn response_format() -> Option<Value> {
633 json_schema_for::<R>()
634 }
635}
636
637#[doc(hidden)]
644pub trait IntoAsyncHandler<Marker>: Send + Sync + 'static {
645 fn into_handler(self) -> RemoteFunctionHandlerWithMetadata;
646 fn request_format() -> Option<Value> {
647 None
648 }
649 fn response_format() -> Option<Value> {
650 None
651 }
652}
653
654fn async_handler<F, T, Fut, R>(f: F) -> RemoteFunctionHandlerWithMetadata
658where
659 F: Fn(T) -> Fut + Send + Sync + 'static,
660 T: serde::de::DeserializeOwned + Send + 'static,
661 Fut: std::future::Future<Output = Result<R, Error>> + Send + 'static,
662 R: serde::Serialize + Send + 'static,
663{
664 Arc::new(
665 move |input: Value,
666 _metadata: Option<Value>|
667 -> std::pin::Pin<
668 Box<dyn std::future::Future<Output = Result<Value, Error>> + Send>,
669 > {
670 match serde_json::from_value::<T>(input) {
671 Ok(arg) => {
672 let fut = f(arg);
673 Box::pin(async move {
674 fut.await.and_then(|val| {
675 serde_json::to_value(&val).map_err(|e| Error::Serde(e.to_string()))
676 })
677 })
678 }
679 Err(e) => {
680 let err = Error::Serde(e.to_string());
681 Box::pin(async move { Err(err) })
682 }
683 }
684 },
685 )
686}
687
688fn async_handler_with_metadata<F, T, Fut, R>(f: F) -> RemoteFunctionHandlerWithMetadata
691where
692 F: Fn(T, Option<Value>) -> Fut + Send + Sync + 'static,
693 T: serde::de::DeserializeOwned + Send + 'static,
694 Fut: std::future::Future<Output = Result<R, Error>> + Send + 'static,
695 R: serde::Serialize + Send + 'static,
696{
697 Arc::new(
698 move |input: Value,
699 metadata: Option<Value>|
700 -> std::pin::Pin<
701 Box<dyn std::future::Future<Output = Result<Value, Error>> + Send>,
702 > {
703 match serde_json::from_value::<T>(input) {
704 Ok(arg) => {
705 let fut = f(arg, metadata);
706 Box::pin(async move {
707 fut.await.and_then(|val| {
708 serde_json::to_value(&val).map_err(|e| Error::Serde(e.to_string()))
709 })
710 })
711 }
712 Err(e) => {
713 let err = Error::Serde(e.to_string());
714 Box::pin(async move { Err(err) })
715 }
716 }
717 },
718 )
719}
720
721impl<F, T, Fut, R> IntoAsyncHandler<(T, Fut, R)> for F
727where
728 F: Fn(T) -> Fut + Send + Sync + 'static,
729 T: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
730 Fut: std::future::Future<Output = Result<R, Error>> + Send + 'static,
731 R: serde::Serialize + schemars::JsonSchema + Send + 'static,
732{
733 fn into_handler(self) -> RemoteFunctionHandlerWithMetadata {
734 async_handler(self)
735 }
736
737 fn request_format() -> Option<Value> {
738 json_schema_for::<T>()
739 }
740
741 fn response_format() -> Option<Value> {
742 json_schema_for::<R>()
743 }
744}
745
746impl<F, T, Fut, R> IntoAsyncHandler<(T, Option<Value>, Fut, R)> for F
749where
750 F: Fn(T, Option<Value>) -> Fut + Send + Sync + 'static,
751 T: serde::de::DeserializeOwned + schemars::JsonSchema + Send + 'static,
752 Fut: std::future::Future<Output = Result<R, Error>> + Send + 'static,
753 R: serde::Serialize + schemars::JsonSchema + Send + 'static,
754{
755 fn into_handler(self) -> RemoteFunctionHandlerWithMetadata {
756 async_handler_with_metadata(self)
757 }
758
759 fn request_format() -> Option<Value> {
760 json_schema_for::<T>()
761 }
762
763 fn response_format() -> Option<Value> {
764 json_schema_for::<R>()
765 }
766}
767
768fn empty_message() -> RegisterFunctionMessage {
773 RegisterFunctionMessage {
774 id: String::new(),
775 description: None,
776 request_format: None,
777 response_format: None,
778 metadata: None,
779 invocation: None,
780 }
781}
782
783pub struct RegisterFunction {
804 message: RegisterFunctionMessage,
805 handler: Option<RemoteFunctionHandlerWithMetadata>,
806}
807
808impl RegisterFunction {
809 pub fn new<F, M>(f: F) -> Self
814 where
815 F: IntoSyncHandler<M>,
816 {
817 let mut message = empty_message();
818 message.request_format = F::request_format();
819 message.response_format = F::response_format();
820 Self {
821 message,
822 handler: Some(f.into_handler()),
823 }
824 }
825
826 pub fn new_async<F, M>(f: F) -> Self
831 where
832 F: IntoAsyncHandler<M>,
833 {
834 let mut message = empty_message();
835 message.request_format = F::request_format();
836 message.response_format = F::response_format();
837 Self {
838 message,
839 handler: Some(f.into_handler()),
840 }
841 }
842
843 pub fn http(config: HttpInvocationConfig) -> Self {
846 let mut message = empty_message();
847 message.invocation = Some(config);
848 Self {
849 message,
850 handler: None,
851 }
852 }
853
854 pub fn description(mut self, desc: impl Into<String>) -> Self {
856 self.message.description = Some(desc.into());
857 self
858 }
859
860 pub fn metadata(mut self, meta: Value) -> Self {
862 self.message.metadata = Some(meta);
863 self
864 }
865
866 pub fn request_format(mut self, schema: Value) -> Self {
868 self.message.request_format = Some(schema);
869 self
870 }
871
872 pub fn response_format(mut self, schema: Value) -> Self {
874 self.message.response_format = Some(schema);
875 self
876 }
877
878 pub(crate) fn into_parts(
879 self,
880 ) -> (
881 RegisterFunctionMessage,
882 Option<RemoteFunctionHandlerWithMetadata>,
883 ) {
884 (self.message, self.handler)
885 }
886}
887
888#[derive(Clone, Copy, Debug)]
894struct ConnTimings {
895 connect_timeout: Duration,
898 ping_interval: Duration,
900 idle_timeout: Duration,
904 retry_delay: Duration,
906}
907
908impl Default for ConnTimings {
909 fn default() -> Self {
910 Self {
911 connect_timeout: Duration::from_secs(10),
912 ping_interval: Duration::from_secs(20),
913 idle_timeout: Duration::from_secs(60),
914 retry_delay: Duration::from_secs(2),
915 }
916 }
917}
918
919struct IIIInner {
920 address: String,
921 outbound: mpsc::UnboundedSender<Outbound>,
922 receiver: Mutex<Option<mpsc::UnboundedReceiver<Outbound>>>,
923 running: AtomicBool,
924 started: AtomicBool,
925 pending: Mutex<HashMap<Uuid, PendingInvocation>>,
926 functions: Mutex<HashMap<String, RemoteFunctionData>>,
927 trigger_types: Mutex<HashMap<String, RemoteTriggerTypeData>>,
928 triggers: Mutex<HashMap<String, RegisterTriggerMessage>>,
929 worker_metadata: Mutex<Option<WorkerMetadata>>,
930 connection_state: Mutex<IIIConnectionState>,
931 fatal_error: Mutex<Option<Error>>,
933 connection_thread: Mutex<Option<std::thread::JoinHandle<()>>>,
934 headers: Mutex<Option<HashMap<String, String>>>,
935 otel_config: Mutex<Option<OtelConfig>>,
936 timings: Mutex<ConnTimings>,
937 worker_id: Mutex<Option<String>>,
941 reattach_token: Mutex<Option<String>>,
945 identity_mode: WorkerIdentityMode,
946 registration_changed: Notify,
947}
948
949#[derive(Clone)]
953pub struct IIIClient {
954 inner: Arc<IIIInner>,
955}
956
957impl IIIClient {
958 pub fn new(address: &str) -> Self {
960 Self::with_metadata(address, WorkerMetadata::default())
961 }
962
963 pub fn with_metadata(address: &str, metadata: WorkerMetadata) -> Self {
968 Self::with_identity(address, metadata, WorkerIdentityMode::Managed)
969 }
970
971 pub(crate) fn with_identity(
973 address: &str,
974 mut metadata: WorkerMetadata,
975 identity_mode: WorkerIdentityMode,
976 ) -> Self {
977 if identity_mode == WorkerIdentityMode::Managed {
978 apply_managed_identity_from_env(&mut metadata);
979 }
980 let (tx, rx) = mpsc::unbounded_channel();
981 let inner = IIIInner {
982 address: address.into(),
983 outbound: tx,
984 receiver: Mutex::new(Some(rx)),
985 running: AtomicBool::new(false),
986 started: AtomicBool::new(false),
987 pending: Mutex::new(HashMap::new()),
988 functions: Mutex::new(HashMap::new()),
989 trigger_types: Mutex::new(HashMap::new()),
990 triggers: Mutex::new(HashMap::new()),
991 worker_metadata: Mutex::new(Some(metadata)),
992 connection_state: Mutex::new(IIIConnectionState::Disconnected),
993 fatal_error: Mutex::new(None),
994 connection_thread: Mutex::new(None),
995 headers: Mutex::new(None),
996 otel_config: Mutex::new(None),
997 timings: Mutex::new(ConnTimings::default()),
998 worker_id: Mutex::new(None),
999 reattach_token: Mutex::new(None),
1000 identity_mode,
1001 registration_changed: Notify::new(),
1002 };
1003 Self {
1004 inner: Arc::new(inner),
1005 }
1006 }
1007
1008 pub fn address(&self) -> &str {
1010 &self.inner.address
1011 }
1012
1013 pub fn set_metadata(&self, mut metadata: WorkerMetadata) {
1018 if self.inner.identity_mode == WorkerIdentityMode::Managed {
1019 apply_managed_identity_from_env(&mut metadata);
1020 }
1021 *self.inner.worker_metadata.lock_or_recover() = Some(metadata);
1022 }
1023
1024 pub fn set_namespace(&self, namespace: impl Into<String>) {
1035 let namespace = namespace.into();
1036 reject_blank_namespace(&namespace, "IIIClient::set_namespace");
1037 if let Some(md) = self.inner.worker_metadata.lock_or_recover().as_mut() {
1038 md.namespace = Some(namespace);
1039 }
1040 }
1041
1042 pub fn namespace(&self) -> Option<String> {
1045 self.inner
1046 .worker_metadata
1047 .lock_or_recover()
1048 .as_ref()
1049 .and_then(|md| md.namespace.clone())
1050 }
1051
1052 pub fn fatal_error(&self) -> Option<Error> {
1056 self.inner.fatal_error.lock_or_recover().clone()
1057 }
1058
1059 pub async fn wait_until_registered(&self, timeout: Duration) -> Result<(), Error> {
1066 if !self.inner.started.load(Ordering::SeqCst) {
1067 return Err(Error::NotConnected);
1068 }
1069
1070 let deadline = Instant::now() + timeout;
1071 loop {
1072 let notified = self.inner.registration_changed.notified();
1073 tokio::pin!(notified);
1074 notified.as_mut().enable();
1075
1076 if let Some(err) = self.fatal_error() {
1077 return Err(err);
1078 }
1079 if self.inner.worker_id.lock_or_recover().is_some() {
1080 return Ok(());
1081 }
1082 if !self.inner.running.load(Ordering::SeqCst) {
1083 return Err(Error::NotConnected);
1084 }
1085
1086 if tokio::time::timeout_at(deadline, notified).await.is_err() {
1087 return Err(Error::Timeout);
1088 }
1089 }
1090 }
1091
1092 pub fn set_headers(&self, headers: HashMap<String, String>) {
1094 *self.inner.headers.lock_or_recover() = Some(headers);
1095 }
1096
1097 pub fn set_otel_config(&self, config: OtelConfig) {
1099 *self.inner.otel_config.lock_or_recover() = Some(config);
1100 }
1101
1102 pub(crate) fn connect(&self) {
1103 if self.inner.started.swap(true, Ordering::SeqCst) {
1104 return;
1105 }
1106
1107 let receiver = self.inner.receiver.lock_or_recover().take();
1108 let Some(rx) = receiver else { return };
1109
1110 self.inner.running.store(true, Ordering::SeqCst);
1111
1112 let iii = self.clone();
1113
1114 let otel_config = {
1115 let mut config = self
1116 .inner
1117 .otel_config
1118 .lock_or_recover()
1119 .take()
1120 .unwrap_or_default();
1121 if config.engine_ws_url.is_none() {
1122 config.engine_ws_url = Some(self.inner.address.clone());
1123 }
1124 config
1125 };
1126
1127 let handle = std::thread::Builder::new()
1133 .name("iii-connection".into())
1134 .spawn(move || {
1135 let rt = tokio::runtime::Builder::new_current_thread()
1136 .enable_all()
1137 .build()
1138 .expect("failed to create iii connection runtime");
1139
1140 rt.block_on(async move {
1141 let otel_active = telemetry::init_otel(otel_config).await;
1142
1143 iii.run_connection(rx).await;
1144
1145 if otel_active {
1146 telemetry::shutdown_otel().await;
1147 }
1148 });
1149 })
1150 .expect("failed to spawn iii connection thread");
1151
1152 *self.inner.connection_thread.lock_or_recover() = Some(handle);
1153 }
1154
1155 pub fn shutdown(&self) {
1168 self.inner.running.store(false, Ordering::SeqCst);
1169 self.inner.registration_changed.notify_waiters();
1170 let _ = self.inner.outbound.send(Outbound::Shutdown);
1171 self.set_connection_state(IIIConnectionState::Disconnected);
1172
1173 if let Some(handle) = self.inner.connection_thread.lock_or_recover().take() {
1174 let _ = handle.join();
1175 }
1176 }
1177
1178 pub async fn shutdown_async(&self) {
1199 self.inner.running.store(false, Ordering::SeqCst);
1200 self.inner.registration_changed.notify_waiters();
1201 let _ = self.inner.outbound.send(Outbound::Shutdown);
1202 self.set_connection_state(IIIConnectionState::Disconnected);
1203 }
1204
1205 fn register_function_inner(
1206 &self,
1207 message: RegisterFunctionMessage,
1208 handler: Option<RemoteFunctionHandlerWithMetadata>,
1209 ) -> FunctionRef {
1210 let id = message.id.clone();
1211 if id.trim().is_empty() {
1212 panic!("id is required");
1213 }
1214 let data = RemoteFunctionData {
1215 message: message.clone(),
1216 handler,
1217 };
1218 let mut funcs = self.inner.functions.lock_or_recover();
1219 match funcs.entry(id.clone()) {
1220 std::collections::hash_map::Entry::Occupied(_) => {
1221 panic!("function id '{}' already registered", id);
1222 }
1223 std::collections::hash_map::Entry::Vacant(entry) => {
1224 entry.insert(data);
1225 }
1226 }
1227 drop(funcs);
1228 let _ = self.send_message(message.to_message());
1229
1230 let iii = self.clone();
1231 let unregister_id = id.clone();
1232 let unregister_fn = Arc::new(move || {
1233 let _ = iii.inner.functions.lock_or_recover().remove(&unregister_id);
1234 let _ = iii.send_message(Message::UnregisterFunction {
1235 id: unregister_id.clone(),
1236 });
1237 });
1238
1239 FunctionRef { id, unregister_fn }
1240 }
1241
1242 pub fn register_function(
1319 &self,
1320 id: impl Into<String>,
1321 registration: RegisterFunction,
1322 ) -> FunctionRef {
1323 let (mut message, handler) = registration.into_parts();
1324 message.id = id.into();
1325 self.register_function_inner(message, handler)
1326 }
1327
1328 pub fn register_trigger_type<H, C, R>(
1359 &self,
1360 trigger_type: RegisterTriggerType<H, C, R>,
1361 ) -> TriggerTypeRef<C, R>
1362 where
1363 H: TriggerHandler + 'static,
1364 {
1365 let message = RegisterTriggerTypeMessage {
1366 id: trigger_type.id,
1367 description: trigger_type.description,
1368 trigger_request_format: trigger_type.trigger_request_format,
1369 call_request_format: trigger_type.call_request_format,
1370 namespace: None,
1374 };
1375
1376 let trigger_type_id = message.id.clone();
1377
1378 self.inner.trigger_types.lock_or_recover().insert(
1379 message.id.clone(),
1380 RemoteTriggerTypeData {
1381 message: message.clone(),
1382 handler: Arc::new(trigger_type.handler),
1383 },
1384 );
1385
1386 let _ = self.send_message(message.to_message());
1387
1388 TriggerTypeRef {
1389 iii: self.clone(),
1390 trigger_type_id,
1391 _phantom: std::marker::PhantomData,
1392 }
1393 }
1394
1395 pub fn unregister_trigger_type(&self, id: impl Into<String>) {
1404 let id = id.into();
1405 self.inner.trigger_types.lock_or_recover().remove(&id);
1406 let msg = UnregisterTriggerTypeMessage { id };
1407 let _ = self.send_message(msg.to_message());
1408 }
1409
1410 pub fn register_trigger(&self, input: RegisterTriggerInput) -> Result<Trigger, Error> {
1432 let id = Uuid::new_v4().to_string();
1433 let message = RegisterTriggerMessage {
1434 id: id.clone(),
1435 trigger_type: input.trigger_type,
1436 function_id: input.function_id,
1437 config: input.config,
1438 metadata: input.metadata,
1439 namespace: call_namespace(
1446 input.namespace,
1447 self.namespace(),
1448 "RegisterTriggerInput.namespace",
1449 )?,
1450 trigger_namespace: input.trigger_namespace,
1458 };
1459
1460 self.inner
1461 .triggers
1462 .lock_or_recover()
1463 .insert(message.id.clone(), message.clone());
1464 let _ = self.send_message(message.to_message());
1465
1466 let iii = self.clone();
1467 let trigger_type = message.trigger_type.clone();
1468 let unregister_id = message.id.clone();
1469 let unregister_fn = Arc::new(move || {
1470 let _ = iii.inner.triggers.lock_or_recover().remove(&unregister_id);
1471 let msg = UnregisterTriggerMessage {
1472 id: unregister_id.clone(),
1473 trigger_type: trigger_type.clone(),
1474 };
1475 let _ = iii.send_message(msg.to_message());
1476 });
1477
1478 Ok(Trigger::new(unregister_fn))
1479 }
1480
1481 pub async fn trigger(
1535 &self,
1536 request: impl Into<TriggerRequestWithMetadata>,
1537 ) -> Result<Value, Error> {
1538 let request = request.into();
1539 let req = request.request;
1540 let metadata = request.metadata;
1541 let namespace = invocation_namespace(
1544 request.namespace,
1545 self.namespace(),
1546 &req.function_id,
1547 "TriggerRequest.namespace",
1548 )?;
1549 let (tp, bg) = inject_trace_headers();
1550
1551 if matches!(req.action, Some(TriggerAction::Void)) {
1553 self.send_message(Message::InvokeFunction {
1554 invocation_id: None,
1555 function_id: req.function_id,
1556 data: req.payload,
1557 traceparent: tp,
1558 baggage: bg,
1559 action: req.action,
1560 metadata,
1561 namespace,
1562 })?;
1563 return Ok(Value::Null);
1564 }
1565
1566 let timeout = Duration::from_millis(req.timeout_ms.unwrap_or(DEFAULT_TIMEOUT_MS));
1568 let invocation_id = Uuid::new_v4();
1569 let (tx, rx) = oneshot::channel();
1570
1571 self.inner
1572 .pending
1573 .lock_or_recover()
1574 .insert(invocation_id, tx);
1575
1576 self.send_message(Message::InvokeFunction {
1577 invocation_id: Some(invocation_id),
1578 function_id: req.function_id,
1579 data: req.payload,
1580 traceparent: tp,
1581 baggage: bg,
1582 action: req.action,
1583 metadata,
1584 namespace,
1585 })?;
1586
1587 match tokio::time::timeout(timeout, rx).await {
1588 Ok(Ok(result)) => result,
1589 Ok(Err(_)) => Err(Error::NotConnected),
1590 Err(_) => {
1591 self.inner.pending.lock_or_recover().remove(&invocation_id);
1592 Err(Error::Timeout)
1593 }
1594 }
1595 }
1596
1597 pub fn get_connection_state(&self) -> IIIConnectionState {
1609 *self.inner.connection_state.lock_or_recover()
1610 }
1611
1612 fn set_connection_state(&self, state: IIIConnectionState) {
1613 let mut current = self.inner.connection_state.lock_or_recover();
1614 if *current == state {
1615 return;
1616 }
1617 *current = state;
1618 }
1619
1620 fn register_worker_metadata(&self) {
1622 if let Some(mut metadata) = self.inner.worker_metadata.lock_or_recover().clone() {
1623 let fw = metadata
1624 .telemetry
1625 .as_ref()
1626 .and_then(|t| t.framework.as_deref())
1627 .unwrap_or("");
1628 if fw.is_empty() {
1629 let telem = metadata.telemetry.get_or_insert_with(Default::default);
1630 telem.framework = Some("iii-rust".to_string());
1631 }
1632 if let Ok(value) = serde_json::to_value(metadata) {
1633 let _ = self.send_message(Message::InvokeFunction {
1634 invocation_id: None,
1635 function_id: "engine::workers::register".to_string(),
1636 data: value,
1637 traceparent: None,
1638 baggage: None,
1639 action: Some(TriggerAction::Void),
1640 metadata: None,
1641 namespace: None,
1642 });
1643 }
1644 }
1645 }
1646
1647 fn send_message(&self, message: Message) -> Result<(), Error> {
1648 if !self.inner.running.load(Ordering::SeqCst) {
1649 return Ok(());
1650 }
1651
1652 self.inner
1653 .outbound
1654 .send(Outbound::Message(message))
1655 .map_err(|_| Error::NotConnected)
1656 }
1657
1658 async fn run_connection(&self, mut rx: mpsc::UnboundedReceiver<Outbound>) {
1659 let mut queue: Vec<Message> = Vec::new();
1660 let mut has_connected_before = false;
1661
1662 while self.inner.running.load(Ordering::SeqCst) {
1663 let t = *self.inner.timings.lock_or_recover();
1664 self.set_connection_state(if has_connected_before {
1665 IIIConnectionState::Reconnecting
1666 } else {
1667 IIIConnectionState::Connecting
1668 });
1669
1670 let custom_headers = self.inner.headers.lock_or_recover().clone();
1671
1672 let connect_result = tokio::time::timeout(t.connect_timeout, async {
1677 if let Some(ref h) = custom_headers {
1678 use tokio_tungstenite::tungstenite::client::IntoClientRequest;
1679 use tokio_tungstenite::tungstenite::http;
1680 let mut request = self
1681 .inner
1682 .address
1683 .as_str()
1684 .into_client_request()
1685 .expect("valid ws request");
1686 for (k, v) in h {
1687 if let (Ok(name), Ok(val)) = (
1688 http::header::HeaderName::from_bytes(k.as_bytes()),
1689 http::header::HeaderValue::from_str(v),
1690 ) {
1691 request.headers_mut().insert(name, val);
1692 }
1693 }
1694 connect_async(request).await
1695 } else {
1696 connect_async(&self.inner.address).await
1697 }
1698 })
1699 .await;
1700
1701 match connect_result {
1702 Ok(Ok((stream, _))) => {
1703 tracing::info!(address = %self.inner.address, "iii connected");
1704 has_connected_before = true;
1705 self.set_connection_state(IIIConnectionState::Connected);
1706 let (mut ws_tx, mut ws_rx) = stream.split();
1707
1708 let previous_worker_id = self.inner.worker_id.lock_or_recover().clone();
1716 if let Some(previous_worker_id) = previous_worker_id {
1717 let reattach_token = self.inner.reattach_token.lock_or_recover().clone();
1718 if let Err(err) = self
1719 .send_ws(
1720 &mut ws_tx,
1721 &Message::Reattach {
1722 previous_worker_id,
1723 reattach_token,
1724 },
1725 )
1726 .await
1727 {
1728 tracing::warn!(error = %err, "failed to send reattach; reconnecting");
1729 sleep(t.retry_delay).await;
1730 continue;
1731 }
1732 }
1733
1734 queue.extend(self.collect_registrations());
1735 Self::dedupe_registrations(&mut queue);
1736
1737 let snapshot_ids: HashSet<String> =
1744 queue.iter().filter_map(Self::registration_key).collect();
1745
1746 if let Err(err) = self.flush_queue(&mut ws_tx, &mut queue).await {
1747 tracing::warn!(error = %err, "failed to flush queue");
1748 sleep(t.retry_delay).await;
1749 continue;
1750 }
1751
1752 let shutdown =
1757 Self::drain_pre_connect_duplicates(&mut rx, &mut queue, &snapshot_ids);
1758 if shutdown {
1759 self.inner.running.store(false, Ordering::SeqCst);
1760 return;
1761 }
1762
1763 if !queue.is_empty() {
1764 if let Err(err) = self.flush_queue(&mut ws_tx, &mut queue).await {
1765 tracing::warn!(
1766 error = %err,
1767 "failed to flush post-drain queue"
1768 );
1769 sleep(t.retry_delay).await;
1770 continue;
1771 }
1772 }
1773
1774 self.register_worker_metadata();
1776
1777 let mut should_reconnect = false;
1778
1779 let mut ping = interval(t.ping_interval);
1785 let mut last_rx = Instant::now();
1786
1787 while self.inner.running.load(Ordering::SeqCst) && !should_reconnect {
1788 tokio::select! {
1789 outgoing = rx.recv() => {
1790 match outgoing {
1791 Some(Outbound::Message(message)) => {
1792 if let Err(err) = self.send_ws(&mut ws_tx, &message).await {
1793 tracing::warn!(error = %err, "send failed; reconnecting");
1794 queue.push(message);
1795 should_reconnect = true;
1796 }
1797 }
1798 Some(Outbound::Shutdown) => {
1799 self.inner.running.store(false, Ordering::SeqCst);
1800 return;
1801 }
1802 None => {
1803 self.inner.running.store(false, Ordering::SeqCst);
1804 return;
1805 }
1806 }
1807 }
1808 incoming = ws_rx.next() => {
1809 match incoming {
1810 Some(Ok(frame)) => {
1811 last_rx = Instant::now();
1812 if let Err(err) = self.handle_frame(frame) {
1813 tracing::warn!(error = %err, "failed to handle frame");
1814 }
1815 }
1816 Some(Err(err)) => {
1817 tracing::warn!(error = %err, "websocket receive error");
1818 should_reconnect = true;
1819 }
1820 None => {
1821 should_reconnect = true;
1822 }
1823 }
1824 }
1825 _ = sleep_until(last_rx + t.idle_timeout) => {
1826 tracing::warn!(
1827 idle_timeout = ?t.idle_timeout,
1828 "no frames from engine within idle window; forcing reconnect"
1829 );
1830 should_reconnect = true;
1831 }
1832 _ = ping.tick() => {
1833 let ping_send = ws_tx.send(WsMessage::Ping(Default::default()));
1837 match tokio::time::timeout(t.idle_timeout, ping_send).await {
1838 Ok(Ok(())) => {}
1839 Ok(Err(err)) => {
1840 tracing::warn!(error = %err, "keepalive ping failed; reconnecting");
1841 should_reconnect = true;
1842 }
1843 Err(_) => {
1844 tracing::warn!("keepalive ping timed out; reconnecting");
1845 should_reconnect = true;
1846 }
1847 }
1848 }
1849 }
1850 }
1851 }
1852 Ok(Err(err)) => {
1853 tracing::warn!(error = %err, "failed to connect; retrying");
1854 }
1855 Err(_) => {
1856 tracing::warn!(
1857 connect_timeout = ?t.connect_timeout,
1858 "connect attempt timed out; retrying"
1859 );
1860 }
1861 }
1862
1863 if self.inner.running.load(Ordering::SeqCst) {
1864 sleep(t.retry_delay).await;
1865 }
1866 }
1867 }
1868
1869 fn collect_registrations(&self) -> Vec<Message> {
1870 let mut messages = Vec::new();
1871
1872 for trigger_type in self.inner.trigger_types.lock_or_recover().values() {
1873 messages.push(trigger_type.message.to_message());
1874 }
1875
1876 for function in self.inner.functions.lock_or_recover().values() {
1877 messages.push(function.message.to_message());
1878 }
1879
1880 for trigger in self.inner.triggers.lock_or_recover().values() {
1881 messages.push(trigger.to_message());
1882 }
1883
1884 messages
1885 }
1886
1887 fn registration_key(message: &Message) -> Option<String> {
1894 match message {
1895 Message::RegisterTriggerType { id, .. } => Some(format!("trigger_type:{id}")),
1896 Message::RegisterTrigger { id, .. } => Some(format!("trigger:{id}")),
1897 Message::RegisterFunction { id, .. } => Some(format!("function:{id}")),
1898 _ => None,
1899 }
1900 }
1901
1902 fn drain_pre_connect_duplicates(
1910 rx: &mut mpsc::UnboundedReceiver<Outbound>,
1911 queue: &mut Vec<Message>,
1912 snapshot_ids: &HashSet<String>,
1913 ) -> bool {
1914 loop {
1915 match rx.try_recv() {
1916 Ok(Outbound::Message(msg)) => {
1917 let is_dup = Self::registration_key(&msg)
1918 .map(|k| snapshot_ids.contains(&k))
1919 .unwrap_or(false);
1920 if is_dup {
1921 continue;
1922 }
1923 queue.push(msg);
1924 }
1925 Ok(Outbound::Shutdown) => return true,
1926 Err(_) => return false,
1927 }
1928 }
1929 }
1930
1931 fn dedupe_registrations(queue: &mut Vec<Message>) {
1932 let mut seen = HashSet::new();
1933 let mut deduped_rev = Vec::with_capacity(queue.len());
1934
1935 for message in queue.iter().rev() {
1936 match Self::registration_key(message) {
1937 Some(key) => {
1938 if seen.insert(key) {
1939 deduped_rev.push(message.clone());
1940 }
1941 }
1942 None => {
1943 deduped_rev.push(message.clone());
1944 }
1945 }
1946 }
1947
1948 deduped_rev.reverse();
1949 *queue = deduped_rev;
1950 }
1951
1952 async fn flush_queue(&self, ws_tx: &mut WsTx, queue: &mut Vec<Message>) -> Result<(), Error> {
1953 let mut drained = Vec::new();
1954 std::mem::swap(queue, &mut drained);
1955
1956 let mut iter = drained.into_iter();
1957 while let Some(message) = iter.next() {
1958 if let Err(err) = self.send_ws(ws_tx, &message).await {
1959 queue.push(message);
1960 queue.extend(iter);
1961 return Err(err);
1962 }
1963 }
1964
1965 Ok(())
1966 }
1967
1968 async fn send_ws(&self, ws_tx: &mut WsTx, message: &Message) -> Result<(), Error> {
1969 let payload = serde_json::to_string(message)?;
1970 let t = *self.inner.timings.lock_or_recover();
1976 match tokio::time::timeout(t.idle_timeout, ws_tx.send(WsMessage::Text(payload.into())))
1977 .await
1978 {
1979 Ok(Ok(())) => Ok(()),
1980 Ok(Err(err)) => Err(err.into()),
1981 Err(_) => {
1982 tracing::warn!("websocket send timed out; treating link as dead");
1983 Err(Error::Timeout)
1984 }
1985 }
1986 }
1987
1988 fn handle_frame(&self, frame: WsMessage) -> Result<(), Error> {
1989 match frame {
1990 WsMessage::Text(text) => self.handle_message(&text),
1991 WsMessage::Binary(bytes) => {
1992 let text = String::from_utf8_lossy(&bytes).to_string();
1993 self.handle_message(&text)
1994 }
1995 _ => Ok(()),
1996 }
1997 }
1998
1999 fn handle_message(&self, payload: &str) -> Result<(), Error> {
2000 let message: Message = serde_json::from_str(payload)?;
2001
2002 match message {
2003 Message::InvocationResult {
2004 invocation_id,
2005 result,
2006 error,
2007 ..
2008 } => {
2009 self.handle_invocation_result(invocation_id, result, error);
2010 }
2011 Message::InvokeFunction {
2012 invocation_id,
2013 function_id,
2014 data,
2015 traceparent,
2016 baggage,
2017 action: _,
2018 metadata,
2019 namespace: _,
2020 } => {
2021 self.handle_invoke_function(
2022 invocation_id,
2023 function_id,
2024 data,
2025 traceparent,
2026 baggage,
2027 metadata,
2028 );
2029 }
2030 Message::RegisterTrigger {
2031 id,
2032 trigger_type,
2033 function_id,
2034 config,
2035 metadata,
2036 namespace,
2040 trigger_namespace: _,
2044 } => {
2045 self.handle_register_trigger(
2046 id,
2047 trigger_type,
2048 function_id,
2049 config,
2050 metadata,
2051 namespace,
2052 );
2053 }
2054 Message::UnregisterTrigger { id, trigger_type } => {
2055 self.handle_unregister_trigger(id, trigger_type);
2056 }
2057 Message::Ping => {
2058 let _ = self.send_message(Message::Pong);
2059 }
2060 Message::WorkerRegistered {
2061 worker_id,
2062 reattach_token,
2063 } => {
2064 tracing::debug!(worker_id = %worker_id, "Worker registered");
2065 *self.inner.worker_id.lock_or_recover() = Some(worker_id);
2066 *self.inner.reattach_token.lock_or_recover() = reattach_token;
2067 self.inner.registration_changed.notify_waiters();
2068 }
2069 Message::RegistrationRejected {
2070 code,
2071 namespace,
2072 worker_name,
2073 function_id,
2074 owner_worker_id,
2075 } => {
2076 self.handle_registration_rejected(
2077 code,
2078 namespace,
2079 worker_name,
2080 function_id,
2081 owner_worker_id,
2082 );
2083 }
2084 Message::TriggerRegistrationResult {
2085 id,
2086 trigger_type,
2087 function_id: _,
2088 error: Some(err),
2089 } => {
2090 tracing::error!(
2091 trigger_id = %id,
2092 trigger_type = %trigger_type,
2093 code = %err.code,
2094 "[iii] Trigger registration failed for {:?}: {}",
2095 id,
2096 err.message
2097 );
2098 }
2099 _ => {}
2100 }
2101
2102 Ok(())
2103 }
2104
2105 fn handle_registration_rejected(
2114 &self,
2115 code: String,
2116 namespace: String,
2117 worker_name: Option<String>,
2118 function_id: Option<String>,
2119 owner_worker_id: String,
2120 ) {
2121 match code.as_str() {
2122 FUNCTION_NAMESPACE_CONFLICT => {
2123 tracing::warn!(
2124 code = %code,
2125 namespace = %namespace,
2126 function_id = %function_id.as_deref().unwrap_or("<unknown>"),
2127 owner_worker_id = %owner_worker_id,
2128 "function registration rejected: another worker in this namespace already \
2129 owns this function id; the worker keeps serving its other functions"
2130 );
2131 }
2132 WORKER_NAMESPACE_CONFLICT => {
2133 self.fail_registration_fatal(
2134 code,
2135 namespace,
2136 worker_name,
2137 function_id,
2138 owner_worker_id,
2139 );
2140 }
2141 _ => {
2142 tracing::error!(
2143 code = %code,
2144 "registration rejected with an unknown code; treating as fatal"
2145 );
2146 self.fail_registration_fatal(
2147 code,
2148 namespace,
2149 worker_name,
2150 function_id,
2151 owner_worker_id,
2152 );
2153 }
2154 }
2155 }
2156
2157 fn fail_registration_fatal(
2161 &self,
2162 code: String,
2163 namespace: String,
2164 worker_name: Option<String>,
2165 function_id: Option<String>,
2166 owner_worker_id: String,
2167 ) {
2168 let err = Error::RegistrationRejected {
2169 code,
2170 namespace,
2171 worker_name,
2172 function_id,
2173 owner_worker_id,
2174 };
2175 tracing::error!(error = %err, "worker registration rejected; not reconnecting");
2176 let drained: Vec<PendingInvocation> = self
2181 .inner
2182 .pending
2183 .lock_or_recover()
2184 .drain()
2185 .map(|(_, sender)| sender)
2186 .collect();
2187 for sender in drained {
2188 let _ = sender.send(Err(err.clone()));
2189 }
2190 *self.inner.fatal_error.lock_or_recover() = Some(err);
2191 self.set_connection_state(IIIConnectionState::Failed);
2192 self.inner.running.store(false, Ordering::SeqCst);
2193 self.inner.registration_changed.notify_waiters();
2194 }
2195
2196 fn handle_invocation_result(
2197 &self,
2198 invocation_id: Uuid,
2199 result: Option<Value>,
2200 error: Option<ErrorBody>,
2201 ) {
2202 let sender = self.inner.pending.lock_or_recover().remove(&invocation_id);
2203 if let Some(sender) = sender {
2204 let result = match error {
2205 Some(error) => Err(Error::Remote {
2206 code: error.code,
2207 message: error.message,
2208 stacktrace: error.stacktrace,
2209 }),
2210 None => Ok(result.unwrap_or(Value::Null)),
2211 };
2212 let _ = sender.send(result);
2213 }
2214 }
2215
2216 fn handle_invoke_function(
2217 &self,
2218 invocation_id: Option<Uuid>,
2219 function_id: String,
2220 data: Value,
2221 traceparent: Option<String>,
2222 baggage: Option<String>,
2223 metadata: Option<Value>,
2224 ) {
2225 tracing::debug!(function_id = %function_id, traceparent = ?traceparent, baggage = ?baggage, "Invoking function");
2226
2227 let func_data = self
2228 .inner
2229 .functions
2230 .lock_or_recover()
2231 .get(&function_id)
2232 .cloned();
2233 let handler = func_data.as_ref().and_then(|d| d.handler.clone());
2234
2235 let Some(handler) = handler else {
2236 let (code, message) = match &func_data {
2237 Some(_) => (
2238 "function_not_invokable".to_string(),
2239 "Function is HTTP-invoked and cannot be invoked locally".to_string(),
2240 ),
2241 None => (
2242 "function_not_found".to_string(),
2243 "Function not found".to_string(),
2244 ),
2245 };
2246 tracing::warn!(function_id = %function_id, "Invocation: {}", message);
2247
2248 if let Some(invocation_id) = invocation_id {
2249 let (resp_tp, resp_bg) = inject_trace_headers();
2250
2251 let error = ErrorBody {
2252 code,
2253 message,
2254 stacktrace: None,
2255 };
2256 let result = self.send_message(Message::InvocationResult {
2257 invocation_id,
2258 function_id,
2259 result: None,
2260 error: Some(error),
2261 traceparent: resp_tp,
2262 baggage: resp_bg,
2263 });
2264
2265 if let Err(err) = result {
2266 tracing::warn!(error = %err, "error sending invocation result");
2267 }
2268 }
2269 return;
2270 };
2271
2272 let iii = self.clone();
2273
2274 tokio::spawn(async move {
2275 let otel_cx = {
2281 use iii_helpers::observability::extract_context;
2282 use iii_helpers::observability::opentelemetry::trace::{
2283 SpanKind, TraceContextExt, Tracer,
2284 };
2285
2286 let parent_cx = extract_context(traceparent.as_deref(), baggage.as_deref());
2287 let tracer =
2288 iii_helpers::observability::opentelemetry::global::tracer("iii-rust-sdk");
2289 let span = tracer
2296 .span_builder(format!("execute {}", function_id))
2297 .with_kind(SpanKind::Internal)
2298 .start_with_context(&tracer, &parent_cx);
2299 parent_cx.with_span(span)
2300 };
2301
2302 let trace_payloads = !std::env::var("III_DISABLE_TRACE_PAYLOADS")
2303 .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
2304 .unwrap_or(false);
2305
2306 let payload_max_bytes = iii_helpers::observability::resolve_max_bytes_from_env();
2307
2308 if trace_payloads {
2309 use iii_helpers::observability::opentelemetry::KeyValue;
2310 use iii_helpers::observability::opentelemetry::trace::TraceContextExt;
2311 use iii_helpers::observability::redact_and_truncate;
2312 let span = otel_cx.span();
2313 if span.span_context().is_valid() {
2314 let (input_json, truncated) = redact_and_truncate(&data, payload_max_bytes);
2315 span.add_event(
2316 "iii.invocation.input",
2317 vec![
2318 KeyValue::new("iii.payload.json", input_json),
2319 KeyValue::new("iii.payload.truncated", truncated),
2320 ],
2321 );
2322 }
2323 }
2324
2325 let result = {
2326 use iii_helpers::observability::opentelemetry::trace::FutureExt as OtelFutureExt;
2327 handler(data, metadata).with_context(otel_cx.clone()).await
2328 };
2329
2330 if trace_payloads {
2331 use iii_helpers::observability::opentelemetry::KeyValue;
2332 use iii_helpers::observability::opentelemetry::trace::TraceContextExt;
2333 use iii_helpers::observability::redact_and_truncate;
2334 let span = otel_cx.span();
2335 if span.span_context().is_valid() {
2336 let (output_json, truncated, ok) = match &result {
2337 Ok(value) => {
2338 let (j, t) = redact_and_truncate(value, payload_max_bytes);
2339 (j, t, true)
2340 }
2341 Err(err) => {
2342 let payload = serde_json::json!({ "error": err.to_string() });
2343 let (j, t) = redact_and_truncate(&payload, payload_max_bytes);
2344 (j, t, false)
2345 }
2346 };
2347 span.add_event(
2348 "iii.invocation.output",
2349 vec![
2350 KeyValue::new("iii.payload.json", output_json),
2351 KeyValue::new("iii.payload.truncated", truncated),
2352 KeyValue::new("iii.payload.ok", ok),
2353 ],
2354 );
2355 }
2356 }
2357
2358 let mut error_stacktrace: Option<String> = None;
2360 {
2361 use iii_helpers::observability::opentelemetry::KeyValue;
2362 use iii_helpers::observability::opentelemetry::trace::{Status, TraceContextExt};
2363 let span = otel_cx.span();
2364 match &result {
2365 Ok(_) => span.set_status(Status::Ok),
2366 Err(err) => {
2367 let (exc_type, exc_message, stacktrace) = match err {
2368 Error::Remote {
2369 code,
2370 message,
2371 stacktrace,
2372 } => (
2373 code.clone(),
2374 message.clone(),
2375 stacktrace.clone().unwrap_or_else(|| {
2376 std::backtrace::Backtrace::force_capture().to_string()
2377 }),
2378 ),
2379 other => (
2380 "InvocationError".to_string(),
2381 other.to_string(),
2382 std::backtrace::Backtrace::force_capture().to_string(),
2383 ),
2384 };
2385 span.set_status(Status::error(exc_message.clone()));
2386 span.add_event(
2387 "exception",
2388 vec![
2389 KeyValue::new("exception.type", exc_type),
2390 KeyValue::new("exception.message", exc_message),
2391 KeyValue::new("exception.stacktrace", stacktrace.clone()),
2392 ],
2393 );
2394 error_stacktrace = Some(stacktrace);
2398 }
2399 }
2400 }
2401
2402 if let Some(invocation_id) = invocation_id {
2403 let (resp_tp, resp_bg) = {
2407 let _guard = otel_cx.attach();
2408 inject_trace_headers()
2409 };
2410
2411 let message = match result {
2412 Ok(value) => Message::InvocationResult {
2413 invocation_id,
2414 function_id,
2415 result: Some(value),
2416 error: None,
2417 traceparent: resp_tp,
2418 baggage: resp_bg,
2419 },
2420 Err(err) => {
2421 let error_body = match err {
2422 Error::Remote {
2423 code,
2424 message,
2425 stacktrace,
2426 } => ErrorBody {
2427 code,
2428 message,
2429 stacktrace,
2434 },
2435 other => ErrorBody {
2436 code: "invocation_failed".to_string(),
2437 message: other.to_string(),
2438 stacktrace: error_stacktrace.or_else(|| {
2439 Some(std::backtrace::Backtrace::force_capture().to_string())
2440 }),
2441 },
2442 };
2443 Message::InvocationResult {
2444 invocation_id,
2445 function_id,
2446 result: None,
2447 error: Some(error_body),
2448 traceparent: resp_tp,
2449 baggage: resp_bg,
2450 }
2451 }
2452 };
2453
2454 let _ = iii.send_message(message);
2455 } else if let Err(err) = result {
2456 tracing::warn!(error = %err, "error handling async invocation");
2457 }
2458 });
2459 }
2460
2461 fn handle_register_trigger(
2462 &self,
2463 id: String,
2464 trigger_type: String,
2465 function_id: String,
2466 config: Value,
2467 metadata: Option<Value>,
2468 namespace: Option<String>,
2469 ) {
2470 let handler = self
2471 .inner
2472 .trigger_types
2473 .lock_or_recover()
2474 .get(&trigger_type)
2475 .map(|data| data.handler.clone());
2476
2477 let iii = self.clone();
2478
2479 tokio::spawn(async move {
2480 let message = if let Some(handler) = handler {
2481 let config = TriggerConfig {
2482 id: id.clone(),
2483 function_id: function_id.clone(),
2484 config,
2485 metadata,
2486 namespace,
2487 };
2488
2489 match handler.register_trigger(config).await {
2490 Ok(()) => Message::TriggerRegistrationResult {
2491 id,
2492 trigger_type,
2493 function_id,
2494 error: None,
2495 },
2496 Err(err) => Message::TriggerRegistrationResult {
2497 id,
2498 trigger_type,
2499 function_id,
2500 error: Some(ErrorBody {
2501 code: "trigger_registration_failed".to_string(),
2502 message: err.to_string(),
2503 stacktrace: None,
2504 }),
2505 },
2506 }
2507 } else {
2508 Message::TriggerRegistrationResult {
2509 id,
2510 trigger_type,
2511 function_id,
2512 error: Some(ErrorBody {
2513 code: "trigger_type_not_found".to_string(),
2514 message: "Trigger type not found".to_string(),
2515 stacktrace: None,
2516 }),
2517 }
2518 };
2519
2520 let _ = iii.send_message(message);
2521 });
2522 }
2523
2524 fn handle_unregister_trigger(&self, id: String, trigger_type: String) {
2525 let handler = self
2526 .inner
2527 .trigger_types
2528 .lock_or_recover()
2529 .get(&trigger_type)
2530 .map(|data| data.handler.clone());
2531
2532 let Some(handler) = handler else {
2533 return;
2534 };
2535
2536 tokio::spawn(async move {
2537 let config = TriggerConfig {
2538 id: id.clone(),
2539 function_id: String::new(),
2540 config: Value::Null,
2541 metadata: None,
2542 namespace: None,
2544 };
2545
2546 if let Err(err) = handler.unregister_trigger(config).await {
2547 tracing::warn!(trigger_id = %id, error = %err, "Error unregistering trigger");
2548 }
2549 });
2550 }
2551}
2552
2553pub(crate) async fn internal_create_channel(
2560 iii: &IIIClient,
2561 buffer_size: Option<usize>,
2562) -> Result<Channel, Error> {
2563 let result = iii
2564 .trigger(TriggerRequest {
2565 function_id: "engine::channels::create".to_string(),
2566 payload: serde_json::json!({ "buffer_size": buffer_size }),
2567 action: None,
2568 timeout_ms: None,
2569 })
2570 .await?;
2571
2572 let writer_ref: StreamChannelRef = serde_json::from_value(
2573 result
2574 .get("writer")
2575 .cloned()
2576 .ok_or_else(|| Error::Serde("missing 'writer' in channel response".into()))?,
2577 )
2578 .map_err(|e| Error::Serde(e.to_string()))?;
2579
2580 let reader_ref: StreamChannelRef = serde_json::from_value(
2581 result
2582 .get("reader")
2583 .cloned()
2584 .ok_or_else(|| Error::Serde("missing 'reader' in channel response".into()))?,
2585 )
2586 .map_err(|e| Error::Serde(e.to_string()))?;
2587
2588 Ok(Channel {
2589 writer: ChannelWriter::new(&iii.inner.address, &writer_ref),
2590 reader: ChannelReader::new(&iii.inner.address, &reader_ref),
2591 writer_ref,
2592 reader_ref,
2593 })
2594}
2595
2596#[cfg(test)]
2597mod tests {
2598 use std::collections::HashMap;
2599 use std::ffi::OsString;
2600
2601 use serde_json::json;
2602
2603 use iii_helpers::http::{HttpInvocationConfig, HttpMethod};
2604
2605 use super::*;
2606 use crate::trigger::{TriggerConfig, TriggerHandler};
2607 use crate::{InitOptions, protocol::RegisterTriggerInput, register_worker};
2608
2609 use std::sync::atomic::AtomicUsize;
2610
2611 static ENV_MUTEX: Mutex<()> = Mutex::new(());
2612
2613 struct ScopedEnvVar {
2614 key: &'static str,
2615 previous: Option<OsString>,
2616 _lock: MutexGuard<'static, ()>,
2617 }
2618
2619 impl ScopedEnvVar {
2620 fn new(key: &'static str) -> Self {
2621 let lock = ENV_MUTEX.lock_or_recover();
2622 let previous = std::env::var_os(key);
2623 Self {
2624 key,
2625 previous,
2626 _lock: lock,
2627 }
2628 }
2629
2630 fn set(&self, value: &str) {
2631 unsafe {
2634 std::env::set_var(self.key, value);
2635 }
2636 }
2637
2638 fn remove(&self) {
2639 unsafe {
2642 std::env::remove_var(self.key);
2643 }
2644 }
2645 }
2646
2647 impl Drop for ScopedEnvVar {
2648 fn drop(&mut self) {
2649 unsafe {
2652 match self.previous.take() {
2653 Some(value) => std::env::set_var(self.key, value),
2654 None => std::env::remove_var(self.key),
2655 }
2656 }
2657 }
2658 }
2659
2660 async fn spawn_stalled_listener() -> (std::net::SocketAddr, Arc<AtomicUsize>) {
2664 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2665 let addr = listener.local_addr().unwrap();
2666 let accepted = Arc::new(AtomicUsize::new(0));
2667 let counter = accepted.clone();
2668 tokio::spawn(async move {
2669 let mut held = Vec::new();
2670 while let Ok((sock, _)) = listener.accept().await {
2671 counter.fetch_add(1, Ordering::SeqCst);
2672 held.push(sock);
2673 }
2674 });
2675 (addr, accepted)
2676 }
2677
2678 async fn spawn_ws_server(silent: bool) -> (std::net::SocketAddr, Arc<AtomicUsize>) {
2683 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2684 let addr = listener.local_addr().unwrap();
2685 let handshakes = Arc::new(AtomicUsize::new(0));
2686 let counter = handshakes.clone();
2687 tokio::spawn(async move {
2688 while let Ok((sock, _)) = listener.accept().await {
2689 let counter = counter.clone();
2690 tokio::spawn(async move {
2691 let Ok(mut ws) = tokio_tungstenite::accept_async(sock).await else {
2692 return;
2693 };
2694 counter.fetch_add(1, Ordering::SeqCst);
2695 if silent {
2696 std::future::pending::<()>().await;
2697 }
2698 while let Some(frame) = ws.next().await {
2699 if frame.is_err() {
2700 break;
2701 }
2702 }
2703 });
2704 }
2705 });
2706 (addr, handshakes)
2707 }
2708
2709 fn client_with_timings(addr: std::net::SocketAddr, timings: ConnTimings) -> IIIClient {
2710 let iii = IIIClient::new(&format!("ws://{addr}"));
2711 *iii.inner.otel_config.lock_or_recover() = Some(OtelConfig {
2714 enabled: Some(false),
2715 ..Default::default()
2716 });
2717 *iii.inner.timings.lock_or_recover() = timings;
2718 iii.connect();
2719 iii
2720 }
2721
2722 async fn wait_for_count(counter: &AtomicUsize, want: usize, within: Duration) -> usize {
2723 let deadline = Instant::now() + within;
2724 loop {
2725 let n = counter.load(Ordering::SeqCst);
2726 if n >= want || Instant::now() >= deadline {
2727 return n;
2728 }
2729 sleep(Duration::from_millis(25)).await;
2730 }
2731 }
2732
2733 #[tokio::test]
2734 async fn connect_timeout_abandons_stalled_connect_and_retries() {
2735 let (addr, accepted) = spawn_stalled_listener().await;
2736 let iii = client_with_timings(
2737 addr,
2738 ConnTimings {
2739 connect_timeout: Duration::from_millis(150),
2740 retry_delay: Duration::from_millis(100),
2741 ..ConnTimings::default()
2742 },
2743 );
2744
2745 let n = wait_for_count(&accepted, 2, Duration::from_secs(5)).await;
2748 iii.shutdown_async().await;
2749 assert!(
2750 n >= 2,
2751 "stalled connect must be abandoned and retried, saw {n} attempts"
2752 );
2753 }
2754
2755 #[tokio::test]
2756 async fn idle_timeout_reconnects_when_engine_goes_silent() {
2757 let (addr, handshakes) = spawn_ws_server(true).await;
2758 let iii = client_with_timings(
2759 addr,
2760 ConnTimings {
2761 connect_timeout: Duration::from_secs(5),
2762 ping_interval: Duration::from_secs(30),
2765 idle_timeout: Duration::from_millis(250),
2766 retry_delay: Duration::from_millis(100),
2767 },
2768 );
2769
2770 let n = wait_for_count(&handshakes, 2, Duration::from_secs(5)).await;
2773 iii.shutdown_async().await;
2774 assert!(
2775 n >= 2,
2776 "silent link must force a reconnect, saw {n} handshakes"
2777 );
2778 }
2779
2780 #[tokio::test]
2781 async fn keepalive_pings_hold_healthy_connection_open() {
2782 let (addr, handshakes) = spawn_ws_server(false).await;
2783 let iii = client_with_timings(
2784 addr,
2785 ConnTimings {
2786 connect_timeout: Duration::from_secs(5),
2787 ping_interval: Duration::from_millis(100),
2788 idle_timeout: Duration::from_secs(1),
2789 retry_delay: Duration::from_millis(100),
2790 },
2791 );
2792
2793 sleep(Duration::from_millis(1500)).await;
2796 let n = handshakes.load(Ordering::SeqCst);
2797 iii.shutdown_async().await;
2798 assert_eq!(n, 1, "healthy pinged connection must not be torn down");
2799 }
2800
2801 struct NoopTriggerHandler;
2802
2803 #[async_trait::async_trait]
2804 impl TriggerHandler for NoopTriggerHandler {
2805 async fn register_trigger(&self, _: TriggerConfig) -> Result<(), Error> {
2806 Ok(())
2807 }
2808
2809 async fn unregister_trigger(&self, _: TriggerConfig) -> Result<(), Error> {
2810 Ok(())
2811 }
2812 }
2813
2814 #[tokio::test]
2815 async fn register_trigger_unregister_removes_entry() {
2816 let iii = register_worker("ws://localhost:1234", InitOptions::default());
2817 let trigger = iii
2818 .register_trigger(RegisterTriggerInput::new(
2819 "demo",
2820 "functions.echo",
2821 json!({ "foo": "bar" }),
2822 ))
2823 .unwrap();
2824
2825 assert_eq!(iii.inner.triggers.lock().unwrap().len(), 1);
2826
2827 trigger.unregister();
2828
2829 assert_eq!(iii.inner.triggers.lock().unwrap().len(), 0);
2830 }
2831
2832 #[tokio::test]
2833 async fn typed_trigger_ref_register_trigger_keeps_legacy_two_arg_shape() {
2834 let iii = register_worker("ws://localhost:1234", InitOptions::default());
2835 let trigger_type = iii.register_trigger_type(
2836 RegisterTriggerType::new("typed-trigger", "typed trigger", NoopTriggerHandler)
2837 .trigger_request_format::<Value>(),
2838 );
2839
2840 let trigger = trigger_type
2841 .register_trigger("functions.echo", json!({ "foo": "bar" }))
2842 .expect("register trigger");
2843
2844 assert_eq!(iii.inner.triggers.lock().unwrap().len(), 1);
2845 trigger.unregister();
2846 }
2847
2848 #[tokio::test]
2849 async fn typed_trigger_ref_register_trigger_with_metadata_stores_metadata() {
2850 let iii = register_worker("ws://localhost:1234", InitOptions::default());
2851 let trigger_type = iii.register_trigger_type(
2852 RegisterTriggerType::new("typed-trigger-meta", "typed trigger", NoopTriggerHandler)
2853 .trigger_request_format::<Value>(),
2854 );
2855
2856 let metadata = json!({ "owner": "billing-team" });
2857 let trigger = trigger_type
2858 .register_trigger_with_metadata(
2859 "functions.echo",
2860 json!({ "foo": "bar" }),
2861 Some(metadata.clone()),
2862 )
2863 .expect("register trigger");
2864
2865 {
2866 let triggers = iii.inner.triggers.lock().unwrap();
2867 let stored = triggers.values().next().expect("stored trigger");
2868 assert_eq!(stored.metadata, Some(metadata));
2869 }
2870 trigger.unregister();
2871 }
2872
2873 #[test]
2874 fn remote_function_handler_alias_remains_single_arg() {
2875 let legacy: crate::types::RemoteFunctionHandler =
2879 std::sync::Arc::new(|input| Box::pin(async move { Ok(input) }));
2880 let with_meta: crate::types::RemoteFunctionHandlerWithMetadata =
2881 std::sync::Arc::new(|input, _metadata| Box::pin(async move { Ok(input) }));
2882 drop((legacy, with_meta));
2883 }
2884
2885 #[tokio::test]
2886 async fn register_function_with_http_config_stores_and_unregister_removes() {
2887 let iii = register_worker("ws://localhost:1234", InitOptions::default());
2888 let config = HttpInvocationConfig {
2889 url: "https://example.com/invoke".to_string(),
2890 method: HttpMethod::Post,
2891 timeout_ms: Some(30000),
2892 headers: HashMap::new(),
2893 auth: None,
2894 };
2895
2896 let func_ref = iii.register_function("external::my_lambda", RegisterFunction::http(config));
2897
2898 assert_eq!(func_ref.id, "external::my_lambda");
2899 assert_eq!(iii.inner.functions.lock().unwrap().len(), 1);
2900
2901 func_ref.unregister();
2902
2903 assert_eq!(iii.inner.functions.lock().unwrap().len(), 0);
2904 }
2905
2906 #[tokio::test]
2907 #[should_panic(expected = "id is required")]
2908 async fn register_function_rejects_empty_id() {
2909 let iii = register_worker("ws://localhost:1234", InitOptions::default());
2910 let config = HttpInvocationConfig {
2911 url: "https://example.com/invoke".to_string(),
2912 method: HttpMethod::Post,
2913 timeout_ms: None,
2914 headers: HashMap::new(),
2915 auth: None,
2916 };
2917
2918 iii.register_function("", RegisterFunction::http(config));
2919 }
2920
2921 #[tokio::test]
2922 async fn register_function_takes_id_then_builder() {
2923 let iii = register_worker("ws://localhost:1234", InitOptions::default());
2924 let func_ref = iii.register_function(
2925 "test::reshaped::ordering",
2926 RegisterFunction::new_async(|input: Value| async move { Ok(input) })
2927 .description("reshaped")
2928 .metadata(json!({"owner": "sdk"})),
2929 );
2930 assert_eq!(func_ref.id, "test::reshaped::ordering");
2931
2932 let funcs = iii.inner.functions.lock().unwrap();
2933 let stored = funcs.get("test::reshaped::ordering").expect("stored");
2934 assert_eq!(stored.message.id, "test::reshaped::ordering");
2935 assert_eq!(stored.message.description.as_deref(), Some("reshaped"));
2936 assert_eq!(stored.message.metadata, Some(json!({"owner": "sdk"})));
2937 assert!(stored.handler.is_some());
2938 }
2939
2940 #[tokio::test]
2941 async fn register_function_metadata_builder_is_optional() {
2942 let clean = RegisterFunction::new_async(|input: Value| async move { Ok(input) });
2943 assert!(
2944 clean.message.metadata.is_none(),
2945 "metadata should be absent unless the builder method is used"
2946 );
2947
2948 let metadata = json!({"owner": "billing-team", "priority": "high"});
2949 let with_metadata = RegisterFunction::new_async(|input: Value| async move { Ok(input) })
2950 .metadata(metadata.clone());
2951 assert_eq!(with_metadata.message.metadata, Some(metadata));
2952 }
2953
2954 #[tokio::test]
2955 async fn register_function_http_variant_has_no_handler() {
2956 let iii = register_worker("ws://localhost:1234", InitOptions::default());
2957 let config = HttpInvocationConfig {
2958 url: "https://example.com/invoke".to_string(),
2959 method: HttpMethod::Post,
2960 timeout_ms: Some(30_000),
2961 headers: HashMap::new(),
2962 auth: None,
2963 };
2964
2965 let func_ref = iii.register_function("external::reshaped", RegisterFunction::http(config));
2966
2967 assert_eq!(func_ref.id, "external::reshaped");
2968 let funcs = iii.inner.functions.lock().unwrap();
2969 let stored = funcs.get("external::reshaped").expect("stored");
2970 assert!(
2971 stored.handler.is_none(),
2972 "handler should be None for HTTP invocation"
2973 );
2974 assert!(
2975 stored.message.invocation.is_some(),
2976 "invocation should be set"
2977 );
2978 }
2979
2980 #[tokio::test]
2981 async fn register_function_new_async_extracts_schemas() {
2982 #[derive(serde::Deserialize, schemars::JsonSchema)]
2983 struct In {
2984 name: String,
2985 }
2986 #[derive(serde::Serialize, schemars::JsonSchema)]
2987 struct Out {
2988 message: String,
2989 }
2990 async fn greet(input: In) -> Result<Out, Error> {
2991 Ok(Out {
2992 message: format!("Hello, {}!", input.name),
2993 })
2994 }
2995
2996 let reg = RegisterFunction::new_async(greet);
2997 assert!(reg.message.request_format.is_some());
2998 assert!(reg.message.response_format.is_some());
2999 assert_eq!(reg.message.request_format.as_ref().unwrap()["title"], "In");
3000 assert_eq!(
3001 reg.message.response_format.as_ref().unwrap()["title"],
3002 "Out"
3003 );
3004 }
3005
3006 #[tokio::test]
3007 async fn register_function_request_format_setter_overrides_auto_extraction() {
3008 #[derive(serde::Deserialize, schemars::JsonSchema)]
3009 struct In {
3010 name: String,
3011 }
3012 async fn handler(input: In) -> Result<String, Error> {
3013 Ok(input.name)
3014 }
3015
3016 let custom = json!({"custom": true});
3017 let reg = RegisterFunction::new_async(handler).request_format(custom.clone());
3018 assert_eq!(reg.message.request_format.as_ref().unwrap(), &custom);
3019 }
3020
3021 #[tokio::test]
3022 async fn register_function_untyped_runs_handler() {
3023 let iii = register_worker("ws://localhost:1234", InitOptions::default());
3024 let _func_ref = iii.register_function(
3025 "test::untyped",
3026 RegisterFunction::new_async(|input: Value| async move { Ok(json!({ "echo": input })) }),
3027 );
3028 let handler = {
3029 let funcs = iii.inner.functions.lock().unwrap();
3030 let stored = funcs.get("test::untyped").expect("stored");
3031 stored.handler.as_ref().expect("has handler").clone()
3032 };
3033 let out = handler(json!({"name": "world"}), None).await.unwrap();
3034 assert_eq!(out, json!({"echo": {"name": "world"}}));
3035 }
3036
3037 #[tokio::test]
3038 async fn new_async_delivers_metadata_as_second_arg() {
3039 let iii = register_worker("ws://localhost:1234", InitOptions::default());
3040 iii.register_function(
3041 "test::with_meta",
3042 RegisterFunction::new_async(|input: Value, metadata: Option<Value>| async move {
3043 Ok(json!({ "input": input, "metadata": metadata }))
3046 }),
3047 );
3048 let handler = {
3049 let funcs = iii.inner.functions.lock().unwrap();
3050 funcs
3051 .get("test::with_meta")
3052 .expect("stored")
3053 .handler
3054 .as_ref()
3055 .expect("has handler")
3056 .clone()
3057 };
3058
3059 let out = handler(json!({ "x": 1 }), Some(json!({ "session_id": "s" })))
3061 .await
3062 .unwrap();
3063 assert_eq!(out["input"], json!({ "x": 1 }));
3064 assert_eq!(out["metadata"], json!({ "session_id": "s" }));
3065
3066 let out_none = handler(json!({ "x": 2 }), None).await.unwrap();
3068 assert_eq!(out_none["metadata"], Value::Null);
3069 }
3070
3071 #[tokio::test]
3072 async fn trigger_request_metadata_is_sent_by_trigger() {
3073 let iii = IIIClient::new("ws://localhost:1234");
3074 iii.inner
3075 .running
3076 .store(true, std::sync::atomic::Ordering::SeqCst);
3077
3078 let _ = iii
3079 .trigger(
3080 TriggerRequest {
3081 function_id: "svc::work".to_string(),
3082 payload: json!({ "x": 1 }),
3083 action: Some(TriggerAction::Void),
3084 timeout_ms: None,
3085 }
3086 .metadata(json!({ "tenant": "acme" })),
3087 )
3088 .await
3089 .expect("void trigger should enqueue");
3090
3091 let mut rx = iii.inner.receiver.lock().unwrap().take().expect("receiver");
3092 let sent = rx.try_recv().expect("sent invoke");
3093 match sent {
3094 Outbound::Message(Message::InvokeFunction {
3095 function_id,
3096 data,
3097 metadata,
3098 action,
3099 ..
3100 }) => {
3101 assert_eq!(function_id, "svc::work");
3102 assert_eq!(data, json!({ "x": 1 }));
3103 assert_eq!(metadata, Some(json!({ "tenant": "acme" })));
3104 assert!(matches!(action, Some(TriggerAction::Void)));
3105 }
3106 _ => panic!("expected InvokeFunction"),
3107 }
3108 }
3109
3110 #[tokio::test]
3111 async fn trigger_request_without_metadata_keeps_legacy_struct_literal_shape() {
3112 let iii = IIIClient::new("ws://localhost:1234");
3113 iii.inner
3114 .running
3115 .store(true, std::sync::atomic::Ordering::SeqCst);
3116
3117 let _ = iii
3118 .trigger(TriggerRequest {
3119 function_id: "svc::work".to_string(),
3120 payload: json!({ "x": 1 }),
3121 action: Some(TriggerAction::Void),
3122 timeout_ms: None,
3123 })
3124 .await
3125 .expect("void trigger should enqueue");
3126
3127 let mut rx = iii.inner.receiver.lock().unwrap().take().expect("receiver");
3128 let sent = rx.try_recv().expect("sent invoke");
3129 match sent {
3130 Outbound::Message(Message::InvokeFunction {
3131 function_id,
3132 metadata,
3133 ..
3134 }) => {
3135 assert_eq!(function_id, "svc::work");
3136 assert!(metadata.is_none());
3137 }
3138 _ => panic!("expected invoke"),
3139 }
3140 }
3141
3142 #[tokio::test]
3143 async fn invoke_function_times_out_and_clears_pending() {
3144 let iii = register_worker("ws://localhost:1234", InitOptions::default());
3145 let result = iii
3146 .trigger(TriggerRequest {
3147 function_id: "functions.echo".to_string(),
3148 payload: json!({ "a": 1 }),
3149 action: None,
3150 timeout_ms: Some(10),
3151 })
3152 .await;
3153
3154 assert!(matches!(result, Err(Error::Timeout)));
3155 assert!(iii.inner.pending.lock().unwrap().is_empty());
3156 }
3157
3158 #[tokio::test]
3159 async fn fatal_registration_rejection_fails_pending_invocations() {
3160 let iii = register_worker("ws://localhost:1234", InitOptions::default());
3161 let id = uuid::Uuid::new_v4();
3162 let (tx, rx) = oneshot::channel();
3163 iii.inner.pending.lock().unwrap().insert(id, tx);
3164
3165 iii.fail_registration_fatal(
3166 "WORKER_NAMESPACE_CONFLICT".to_string(),
3167 "orders".to_string(),
3168 Some("state".to_string()),
3169 None,
3170 "owner-1".to_string(),
3171 );
3172
3173 match rx.await {
3176 Ok(Err(Error::RegistrationRejected { code, .. })) => {
3177 assert_eq!(code, "WORKER_NAMESPACE_CONFLICT");
3178 }
3179 other => panic!("expected RegistrationRejected, got {other:?}"),
3180 }
3181 assert!(iii.inner.pending.lock().unwrap().is_empty());
3182 }
3183
3184 #[test]
3185 fn worker_metadata_default_reads_iii_isolation_env_var() {
3186 let env = ScopedEnvVar::new("III_ISOLATION");
3187
3188 env.remove();
3189 assert!(WorkerMetadata::default().isolation.is_none());
3190
3191 env.set("docker");
3192 assert_eq!(
3193 WorkerMetadata::default().isolation.as_deref(),
3194 Some("docker")
3195 );
3196 }
3197
3198 #[test]
3199 fn worker_name_resolution_prefers_iii_worker_name_env_var() {
3200 let env = ScopedEnvVar::new("III_WORKER_NAME");
3201
3202 env.remove();
3203 let fallback = WorkerMetadata::default().name;
3204 assert!(
3205 fallback.ends_with(&format!(":{}", std::process::id())),
3206 "expected hostname:pid fallback, got {fallback}"
3207 );
3208
3209 env.set("");
3210 let metadata = WorkerMetadata {
3211 name: "explicit-worker".to_string(),
3212 ..WorkerMetadata::default()
3213 };
3214 let iii = IIIClient::with_metadata("ws://127.0.0.1:0", metadata);
3215 let explicit_name = iii
3216 .inner
3217 .worker_metadata
3218 .lock_or_recover()
3219 .as_ref()
3220 .unwrap()
3221 .name
3222 .clone();
3223 assert_eq!(explicit_name, "explicit-worker");
3224
3225 env.set("managed-worker");
3226 assert_ne!(WorkerMetadata::default().name, "managed-worker");
3227
3228 let metadata = WorkerMetadata {
3229 name: "explicit-worker".to_string(),
3230 ..WorkerMetadata::default()
3231 };
3232 let iii = IIIClient::with_metadata("ws://127.0.0.1:0", metadata);
3233 let resolved_name = iii
3234 .inner
3235 .worker_metadata
3236 .lock_or_recover()
3237 .as_ref()
3238 .unwrap()
3239 .name
3240 .clone();
3241 assert_eq!(resolved_name, "managed-worker");
3242
3243 let replacement = WorkerMetadata {
3244 name: "replacement-worker".to_string(),
3245 ..WorkerMetadata::default()
3246 };
3247 iii.set_metadata(replacement);
3248 let replaced_name = iii
3249 .inner
3250 .worker_metadata
3251 .lock_or_recover()
3252 .as_ref()
3253 .unwrap()
3254 .name
3255 .clone();
3256 assert_eq!(replaced_name, "managed-worker");
3257 }
3258
3259 #[test]
3260 fn explicit_identity_preserves_metadata_name() {
3261 let env = ScopedEnvVar::new("III_WORKER_NAME");
3262 env.set("managed-worker");
3263
3264 let metadata = WorkerMetadata {
3265 name: "remote-bridge".to_string(),
3266 ..WorkerMetadata::default()
3267 };
3268 let iii =
3269 IIIClient::with_identity("ws://127.0.0.1:0", metadata, WorkerIdentityMode::Explicit);
3270 let resolved_name = iii
3271 .inner
3272 .worker_metadata
3273 .lock_or_recover()
3274 .as_ref()
3275 .unwrap()
3276 .name
3277 .clone();
3278
3279 assert_eq!(resolved_name, "remote-bridge");
3280 }
3281
3282 #[test]
3283 fn explicit_identity_preserves_metadata_namespace() {
3284 let env = ScopedEnvVar::new("III_NAMESPACE");
3285 env.set("managed-namespace");
3286
3287 let metadata = WorkerMetadata {
3288 namespace: Some("remote-namespace".to_string()),
3289 ..WorkerMetadata::default()
3290 };
3291 let iii =
3292 IIIClient::with_identity("ws://127.0.0.1:0", metadata, WorkerIdentityMode::Explicit);
3293
3294 assert_eq!(iii.namespace().as_deref(), Some("remote-namespace"));
3295 }
3296
3297 #[test]
3298 fn managed_identity_ignores_whitespace_only_namespace() {
3299 let env = ScopedEnvVar::new("III_NAMESPACE");
3300 env.set(" \t ");
3301
3302 let iii = IIIClient::with_metadata("ws://127.0.0.1:0", WorkerMetadata::default());
3303
3304 assert!(iii.namespace().is_none());
3305 }
3306
3307 #[test]
3308 fn parse_cargo_package_name_extracts_name_field() {
3309 let toml = "[package]\nname = \"my-crate\"\nversion = \"1.0.0\"\n";
3310 assert_eq!(parse_cargo_package_name(toml), Some("my-crate".to_string()));
3311 }
3312
3313 #[test]
3314 fn parse_cargo_package_name_ignores_other_tables() {
3315 let toml = "[dependencies]\nname = \"not-the-package\"\n[package]\nname = \"the-pkg\"\n";
3316 assert_eq!(parse_cargo_package_name(toml), Some("the-pkg".to_string()));
3317 }
3318
3319 #[test]
3320 fn parse_cargo_package_name_returns_none_when_missing() {
3321 let toml = "[package]\nversion = \"1.0.0\"\n";
3322 assert_eq!(parse_cargo_package_name(toml), None);
3323 }
3324
3325 #[test]
3326 fn parse_cargo_package_name_returns_none_when_blank() {
3327 let toml = "[package]\nname = \"\"\n";
3328 assert_eq!(parse_cargo_package_name(toml), None);
3329 }
3330
3331 #[test]
3332 fn detect_project_name_reads_cargo_toml_in_cwd() {
3333 let tmp = std::env::temp_dir().join(format!("iii-rust-detect-{}", std::process::id()));
3334 std::fs::create_dir_all(&tmp).unwrap();
3335 std::fs::write(
3336 tmp.join("Cargo.toml"),
3337 "[package]\nname = \"detected-crate\"\n",
3338 )
3339 .unwrap();
3340
3341 assert_eq!(
3342 detect_project_name(Some(tmp.clone())),
3343 Some("detected-crate".to_string())
3344 );
3345
3346 std::fs::remove_dir_all(&tmp).ok();
3347 }
3348
3349 #[test]
3350 fn detect_project_name_falls_back_to_dir_basename_without_cargo_toml() {
3351 let tmp = std::env::temp_dir().join(format!("iii-rust-fallback-{}", std::process::id()));
3352 std::fs::create_dir_all(&tmp).unwrap();
3353
3354 let basename = tmp.file_name().unwrap().to_str().unwrap().to_string();
3355 assert_eq!(detect_project_name(Some(tmp.clone())), Some(basename));
3356
3357 std::fs::remove_dir_all(&tmp).ok();
3358 }
3359
3360 #[test]
3361 fn detect_project_name_falls_back_to_dir_basename_when_cargo_toml_lacks_name() {
3362 let tmp =
3363 std::env::temp_dir().join(format!("iii-rust-fallback-noname-{}", std::process::id()));
3364 std::fs::create_dir_all(&tmp).unwrap();
3365 std::fs::write(tmp.join("Cargo.toml"), "[package]\nversion = \"1.0.0\"\n").unwrap();
3366
3367 let basename = tmp.file_name().unwrap().to_str().unwrap().to_string();
3368 assert_eq!(detect_project_name(Some(tmp.clone())), Some(basename));
3369
3370 std::fs::remove_dir_all(&tmp).ok();
3371 }
3372
3373 fn make_register_function(id: &str) -> Message {
3374 Message::RegisterFunction {
3375 id: id.to_string(),
3376 description: None,
3377 request_format: None,
3378 response_format: None,
3379 metadata: None,
3380 invocation: None,
3381 }
3382 }
3383
3384 fn make_register_trigger(id: &str) -> Message {
3385 Message::RegisterTrigger {
3386 id: id.to_string(),
3387 trigger_type: "demo".to_string(),
3388 function_id: "fn".to_string(),
3389 config: json!({}),
3390 metadata: None,
3391 namespace: None,
3392 trigger_namespace: None,
3393 }
3394 }
3395
3396 fn make_register_trigger_type(id: &str) -> Message {
3397 Message::RegisterTriggerType {
3398 id: id.to_string(),
3399 description: "tt".to_string(),
3400 trigger_request_format: None,
3401 call_request_format: None,
3402 namespace: None,
3403 }
3404 }
3405
3406 fn make_invoke(function_id: &str) -> Message {
3407 Message::InvokeFunction {
3408 invocation_id: None,
3409 function_id: function_id.to_string(),
3410 data: json!({}),
3411 traceparent: None,
3412 baggage: None,
3413 action: None,
3414 metadata: None,
3415 namespace: None,
3416 }
3417 }
3418
3419 #[test]
3420 fn registration_key_returns_typed_keys_for_register_messages() {
3421 assert_eq!(
3422 IIIClient::registration_key(&make_register_function("greet")),
3423 Some("function:greet".to_string())
3424 );
3425 assert_eq!(
3426 IIIClient::registration_key(&make_register_trigger("t1")),
3427 Some("trigger:t1".to_string())
3428 );
3429 assert_eq!(
3430 IIIClient::registration_key(&make_register_trigger_type("tt1")),
3431 Some("trigger_type:tt1".to_string())
3432 );
3433 }
3434
3435 #[test]
3436 fn registration_key_returns_none_for_non_register_messages() {
3437 assert_eq!(IIIClient::registration_key(&make_invoke("f")), None);
3438 assert_eq!(IIIClient::registration_key(&Message::Ping), None);
3439 assert_eq!(IIIClient::registration_key(&Message::Pong), None);
3440 assert_eq!(
3441 IIIClient::registration_key(&Message::WorkerRegistered {
3442 worker_id: "w".to_string(),
3443 reattach_token: None
3444 }),
3445 None
3446 );
3447 }
3448
3449 #[tokio::test]
3450 async fn drain_pre_connect_duplicates_drops_only_known_register_ids() {
3451 let (tx, mut rx) = mpsc::unbounded_channel::<Outbound>();
3452
3453 tx.send(Outbound::Message(make_register_function("dup-fn")))
3454 .unwrap();
3455 tx.send(Outbound::Message(make_invoke("some::fn"))).unwrap();
3456 tx.send(Outbound::Message(make_register_function("new-fn")))
3457 .unwrap();
3458 tx.send(Outbound::Message(Message::Pong)).unwrap();
3459 tx.send(Outbound::Message(make_register_trigger("dup-trig")))
3460 .unwrap();
3461 tx.send(Outbound::Message(make_register_trigger("new-trig")))
3462 .unwrap();
3463
3464 let snapshot_ids: HashSet<String> = [
3465 "function:dup-fn".to_string(),
3466 "trigger:dup-trig".to_string(),
3467 ]
3468 .into_iter()
3469 .collect();
3470
3471 let mut queue: Vec<Message> = Vec::new();
3472 let shutdown = IIIClient::drain_pre_connect_duplicates(&mut rx, &mut queue, &snapshot_ids);
3473
3474 assert!(!shutdown);
3475 let kept_keys: Vec<Option<String>> =
3476 queue.iter().map(IIIClient::registration_key).collect();
3477 assert_eq!(
3478 kept_keys,
3479 vec![
3480 None,
3481 Some("function:new-fn".to_string()),
3482 None,
3483 Some("trigger:new-trig".to_string()),
3484 ],
3485 "kept queue mismatch: {queue:#?}"
3486 );
3487 }
3488
3489 #[tokio::test]
3490 async fn drain_pre_connect_duplicates_signals_shutdown() {
3491 let (tx, mut rx) = mpsc::unbounded_channel::<Outbound>();
3492
3493 tx.send(Outbound::Message(make_register_function("a")))
3494 .unwrap();
3495 tx.send(Outbound::Shutdown).unwrap();
3496 tx.send(Outbound::Message(make_register_function("b")))
3497 .unwrap();
3498
3499 let snapshot_ids: HashSet<String> = ["function:a".to_string()].into_iter().collect();
3500 let mut queue: Vec<Message> = Vec::new();
3501 let shutdown = IIIClient::drain_pre_connect_duplicates(&mut rx, &mut queue, &snapshot_ids);
3502
3503 assert!(shutdown, "expected shutdown signal to be reported");
3504 assert!(
3505 queue.is_empty(),
3506 "queue must be empty when shutdown short-circuits the drain: {queue:#?}"
3507 );
3508 }
3509
3510 #[tokio::test]
3511 async fn drain_pre_connect_duplicates_returns_false_on_empty_channel() {
3512 let (_tx, mut rx) = mpsc::unbounded_channel::<Outbound>();
3513 let snapshot_ids: HashSet<String> = HashSet::new();
3514 let mut queue: Vec<Message> = Vec::new();
3515
3516 let shutdown = IIIClient::drain_pre_connect_duplicates(&mut rx, &mut queue, &snapshot_ids);
3517
3518 assert!(!shutdown);
3519 assert!(queue.is_empty());
3520 }
3521
3522 #[tokio::test]
3523 #[tracing_test::traced_test]
3524 async fn trigger_registration_result_error_is_logged() {
3525 let iii = register_worker("ws://localhost:1234", InitOptions::default());
3526 let payload = serde_json::json!({
3527 "type": "triggerregistrationresult",
3528 "id": "trig-1",
3529 "trigger_type": "http",
3530 "function_id": "fn-1",
3531 "error": {
3532 "code": "trigger_type_not_found",
3533 "message": "Trigger type \"http\" not found — worker http is missing. Run: iii trigger -n <compose-daemon-namespace> compose::add worker=http",
3534 },
3535 })
3536 .to_string();
3537
3538 iii.handle_message(&payload).unwrap();
3539
3540 assert!(logs_contain("<compose-daemon-namespace>"));
3541 assert!(logs_contain("compose::add worker=http"));
3542 assert!(logs_contain("trig-1"));
3543 }
3544
3545 #[tokio::test]
3546 #[tracing_test::traced_test]
3547 async fn trigger_registration_result_success_does_not_log_error() {
3548 let iii = register_worker("ws://localhost:1234", InitOptions::default());
3549 let payload = serde_json::json!({
3550 "type": "triggerregistrationresult",
3551 "id": "trig-2",
3552 "trigger_type": "http",
3553 "function_id": "fn-2",
3554 })
3555 .to_string();
3556
3557 iii.handle_message(&payload).unwrap();
3558
3559 assert!(!logs_contain("Trigger registration failed"));
3560 }
3561
3562 #[test]
3563 fn namespace_resolution_reads_env_and_prefers_explicit_option() {
3564 let env = ScopedEnvVar::new("III_NAMESPACE");
3565
3566 env.remove();
3567 assert!(WorkerMetadata::default().namespace.is_none());
3569 assert!(resolve_namespace(None).is_none());
3570 assert_eq!(
3572 resolve_namespace(Some("payments".into())).as_deref(),
3573 Some("payments")
3574 );
3575
3576 env.set("orders");
3577 assert!(WorkerMetadata::default().namespace.is_none());
3580 assert_eq!(resolve_namespace(None).as_deref(), Some("orders"));
3582 assert_eq!(
3584 resolve_namespace(Some("payments".into())).as_deref(),
3585 Some("payments")
3586 );
3587 }
3588
3589 #[tokio::test]
3590 async fn trigger_request_namespace_is_sent_by_trigger() {
3591 let iii = IIIClient::new("ws://localhost:1234");
3592 iii.inner
3593 .running
3594 .store(true, std::sync::atomic::Ordering::SeqCst);
3595
3596 let _ = iii
3597 .trigger(
3598 TriggerRequest {
3599 function_id: "svc::work".to_string(),
3600 payload: json!({ "x": 1 }),
3601 action: Some(TriggerAction::Void),
3602 timeout_ms: None,
3603 }
3604 .namespace("orders"),
3605 )
3606 .await
3607 .expect("void trigger should enqueue");
3608
3609 let mut rx = iii.inner.receiver.lock().unwrap().take().expect("receiver");
3610 let sent = rx.try_recv().expect("sent invoke");
3611 match sent {
3612 Outbound::Message(msg @ Message::InvokeFunction { .. }) => {
3613 let Message::InvokeFunction { namespace, .. } = &msg else {
3614 unreachable!()
3615 };
3616 assert_eq!(namespace.as_deref(), Some("orders"));
3617 let wire = serde_json::to_string(&msg).unwrap();
3619 assert!(wire.contains(r#""namespace":"orders""#), "wire: {wire}");
3620 }
3621 _ => panic!("expected InvokeFunction"),
3622 }
3623 }
3624
3625 #[tokio::test]
3626 async fn trigger_request_without_namespace_omits_it_on_the_wire() {
3627 let iii = IIIClient::new("ws://localhost:1234");
3628 iii.inner
3629 .running
3630 .store(true, std::sync::atomic::Ordering::SeqCst);
3631
3632 let _ = iii
3633 .trigger(TriggerRequest {
3634 function_id: "svc::work".to_string(),
3635 payload: json!({ "x": 1 }),
3636 action: Some(TriggerAction::Void),
3637 timeout_ms: None,
3638 })
3639 .await
3640 .expect("void trigger should enqueue");
3641
3642 let mut rx = iii.inner.receiver.lock().unwrap().take().expect("receiver");
3643 let sent = rx.try_recv().expect("sent invoke");
3644 match sent {
3645 Outbound::Message(msg @ Message::InvokeFunction { .. }) => {
3646 let Message::InvokeFunction { namespace, .. } = &msg else {
3647 unreachable!()
3648 };
3649 assert!(namespace.is_none());
3650 let wire = serde_json::to_string(&msg).unwrap();
3651 assert!(!wire.contains("namespace"), "wire: {wire}");
3652 }
3653 _ => panic!("expected InvokeFunction"),
3654 }
3655 }
3656
3657 #[test]
3658 fn worker_name_conflict_is_fatal_and_stops_worker() {
3659 let iii = IIIClient::new("ws://localhost:1234");
3660 iii.inner
3661 .running
3662 .store(true, std::sync::atomic::Ordering::SeqCst);
3663
3664 let payload = json!({
3665 "type": "registrationrejected",
3666 "code": "WORKER_NAMESPACE_CONFLICT",
3667 "namespace": "orders",
3668 "worker_name": "checkout",
3669 "owner_worker_id": "worker-abc",
3670 })
3671 .to_string();
3672
3673 iii.handle_message(&payload).unwrap();
3674
3675 assert!(
3677 !iii.inner.running.load(std::sync::atomic::Ordering::SeqCst),
3678 "worker must stop; running flag should be cleared"
3679 );
3680 assert_eq!(iii.get_connection_state(), IIIConnectionState::Failed);
3681
3682 let err = iii.fatal_error().expect("fatal error must surface");
3683 match err {
3684 Error::RegistrationRejected {
3685 code,
3686 namespace,
3687 worker_name,
3688 function_id,
3689 owner_worker_id,
3690 } => {
3691 assert_eq!(code, "WORKER_NAMESPACE_CONFLICT");
3692 assert_eq!(namespace, "orders");
3693 assert_eq!(worker_name.as_deref(), Some("checkout"));
3694 assert!(function_id.is_none());
3695 assert_eq!(owner_worker_id, "worker-abc");
3696 }
3697 other => panic!("expected RegistrationRejected, got {other:?}"),
3698 }
3699 }
3700
3701 #[tokio::test]
3702 async fn wait_until_registered_completes_after_engine_accepts_worker() {
3703 let iii = IIIClient::new("ws://localhost:1234");
3704 iii.inner.started.store(true, Ordering::SeqCst);
3705 iii.inner.running.store(true, Ordering::SeqCst);
3706 let waiter = {
3707 let iii = iii.clone();
3708 tokio::spawn(async move { iii.wait_until_registered(Duration::from_secs(1)).await })
3709 };
3710
3711 tokio::task::yield_now().await;
3712 let payload = json!({
3713 "type": "workerregistered",
3714 "worker_id": "worker-accepted",
3715 "reattach_token": "secret",
3716 })
3717 .to_string();
3718 iii.handle_message(&payload).unwrap();
3719
3720 assert!(waiter.await.unwrap().is_ok());
3721 }
3722
3723 #[tokio::test]
3724 async fn wait_until_registered_returns_registration_rejection() {
3725 let iii = IIIClient::new("ws://localhost:1234");
3726 iii.inner.started.store(true, Ordering::SeqCst);
3727 iii.inner.running.store(true, Ordering::SeqCst);
3728 let payload = json!({
3729 "type": "registrationrejected",
3730 "code": "WORKER_NAMESPACE_CONFLICT",
3731 "namespace": "orders",
3732 "worker_name": "checkout",
3733 "owner_worker_id": "worker-abc",
3734 })
3735 .to_string();
3736 iii.handle_message(&payload).unwrap();
3737
3738 let result = iii.wait_until_registered(Duration::from_secs(1)).await;
3739 assert!(matches!(result, Err(Error::RegistrationRejected { .. })));
3740 }
3741
3742 #[test]
3743 fn function_conflict_is_not_fatal_and_keeps_serving() {
3744 let iii = IIIClient::new("ws://localhost:1234");
3745 iii.inner
3746 .running
3747 .store(true, std::sync::atomic::Ordering::SeqCst);
3748 iii.set_connection_state(IIIConnectionState::Connected);
3750
3751 let payload = json!({
3754 "type": "registrationrejected",
3755 "code": "FUNCTION_NAMESPACE_CONFLICT",
3756 "namespace": "orders",
3757 "function_id": "orders::charge",
3758 "owner_worker_id": "worker-abc",
3759 })
3760 .to_string();
3761
3762 iii.handle_message(&payload).unwrap();
3763
3764 assert!(
3767 iii.inner.running.load(std::sync::atomic::Ordering::SeqCst),
3768 "worker must keep running after a function-id conflict"
3769 );
3770 assert_eq!(iii.get_connection_state(), IIIConnectionState::Connected);
3771 assert!(
3772 iii.fatal_error().is_none(),
3773 "a function-id conflict must not be fatal"
3774 );
3775 }
3776
3777 #[test]
3778 fn function_info_captures_namespace_and_tolerates_absence() {
3779 let with_ns: FunctionInfo = serde_json::from_value(json!({
3781 "function_id": "svc::work",
3782 "description": null,
3783 "request_format": null,
3784 "response_format": null,
3785 "metadata": null,
3786 "namespace": "orders",
3787 }))
3788 .unwrap();
3789 assert_eq!(with_ns.namespace.as_deref(), Some("orders"));
3790
3791 let without_ns: FunctionInfo = serde_json::from_value(json!({
3793 "function_id": "svc::work",
3794 "description": null,
3795 "request_format": null,
3796 "response_format": null,
3797 "metadata": null,
3798 }))
3799 .unwrap();
3800 assert!(without_ns.namespace.is_none());
3801 }
3802}