1use std::collections::HashMap;
11use std::net::SocketAddr;
12use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering};
13use std::sync::Arc;
14use std::time::{Duration, Instant};
15
16use axum::{
17 extract::State,
18 http::{HeaderMap, StatusCode},
19 response::{IntoResponse, Response},
20 routing::{get, post},
21 Json, Router,
22};
23use reqwest::Client;
24use serde::{Deserialize, Serialize};
25use serde_json::{json, Value};
26use socket2::{Domain, Protocol, Socket, Type};
27use tokio::net::TcpListener;
28use tokio::sync::Mutex;
29
30use crate::backend_stability::{observe_backend_stability, BackendStabilityObservation};
31use crate::errors::{IicpError, Result};
32
33const DEFAULT_DIRECTORY: &str = "https://iicp.network/api";
34const HEARTBEAT_INTERVAL_SECS: u64 = 30;
35const NONCE_TTL_SECS: u64 = 300;
36
37fn filter_excluded_models(mut models: Vec<String>, excluded: &[String]) -> Vec<String> {
41 if excluded.is_empty() {
42 return models;
43 }
44 models.retain(|model| {
45 !excluded
46 .iter()
47 .any(|excluded_model| excluded_model == model)
48 });
49 models
50}
51
52fn attach_update_status(body: &mut Value) {
53 if let (Some(dst), Some(src)) = (
54 body.as_object_mut(),
55 crate::updater::auto_update_status_json().as_object(),
56 ) {
57 for (k, v) in src {
58 dst.insert(k.clone(), v.clone());
59 }
60 }
61}
62
63async fn probe_health_models_bg(
66 http: &Client,
67 backend_url: &str,
68 api_key: &Option<String>,
69 excluded_models: &[String],
70) -> Option<Vec<String>> {
71 let base = backend_url.trim_end_matches('/');
72 if base.is_empty() {
73 return None;
74 }
75 let root = base.strip_suffix("/v1").unwrap_or(base);
76 let mut rb = http
77 .get(format!("{root}/api/tags"))
78 .timeout(std::time::Duration::from_secs(2));
79 if let Some(key) = api_key {
80 if !key.is_empty() {
81 rb = rb.bearer_auth(key);
82 }
83 }
84 if let Ok(resp) = rb.send().await {
85 if resp.status().is_success() {
86 if let Ok(data) = resp.json::<Value>().await {
87 if let Some(arr) = data["models"].as_array() {
88 let mut names: Vec<String> = arr
89 .iter()
90 .filter_map(|m| m["name"].as_str().map(str::to_string))
91 .collect::<std::collections::HashSet<_>>()
92 .into_iter()
93 .collect();
94 names.sort();
95 return Some(filter_excluded_models(names, excluded_models));
96 }
97 }
98 }
99 }
100 let mut rb2 = http
101 .get(format!("{root}/v1/models"))
102 .timeout(std::time::Duration::from_secs(2));
103 if let Some(key) = api_key {
104 if !key.is_empty() {
105 rb2 = rb2.bearer_auth(key);
106 }
107 }
108 if let Ok(resp) = rb2.send().await {
109 if resp.status().is_success() {
110 if let Ok(data) = resp.json::<Value>().await {
111 if let Some(arr) = data["data"].as_array() {
112 return Some(filter_excluded_models(
113 arr.iter()
114 .filter_map(|m| m["id"].as_str().map(str::to_string))
115 .collect(),
116 excluded_models,
117 ));
118 }
119 }
120 }
121 }
122 None
123}
124
125#[derive(Debug, Clone, PartialEq, Eq)]
126struct RegistrationCredentials {
127 node_token: String,
128 node_hmac_key: Option<String>,
129}
130
131async fn reregister(
135 http: &Client,
136 url: &str,
137 payload: &serde_json::Value,
138) -> Option<RegistrationCredentials> {
139 let resp = http.post(url).json(payload).send().await.ok()?;
140 if !resp.status().is_success() {
141 return None;
142 }
143 let data = resp.json::<serde_json::Value>().await.ok()?;
144 let token = data["node_token"]
145 .as_str()
146 .or_else(|| data["token"].as_str())?;
147 Some(RegistrationCredentials {
148 node_token: token.to_string(),
149 node_hmac_key: data["node_hmac_key"]
150 .as_str()
151 .filter(|s| !s.is_empty())
152 .map(str::to_string),
153 })
154}
155
156fn update_saved_identity_credentials(
157 identity: &mut crate::identity::NodeIdentity,
158 token: &str,
159 hmac_key: Option<&str>,
160) {
161 identity.node_token = Some(token.to_string());
162 if let Some(key) = hmac_key.filter(|k| !k.is_empty()) {
163 identity.node_hmac_key = Some(key.to_string());
164 }
165}
166
167fn persist_saved_credentials(saved_node_name: Option<&str>, token: &str, hmac_key: Option<&str>) {
168 let Some(name) = saved_node_name.filter(|n| !n.trim().is_empty()) else {
169 return;
170 };
171 if token.is_empty() {
172 return;
173 }
174 match crate::identity::load_node(name) {
175 Ok(Some(mut identity)) => {
176 update_saved_identity_credentials(&mut identity, token, hmac_key);
177 if let Err(err) = crate::identity::save_node(&identity) {
178 tracing::warn!(
179 "could not persist refreshed node credentials for saved node {name}: {err}"
180 );
181 }
182 }
183 Ok(None) => tracing::warn!(
184 "saved node config {name} not found; cannot persist refreshed node credentials"
185 ),
186 Err(err) => tracing::warn!(
187 "could not load saved node config {name} to persist refreshed credentials: {err}"
188 ),
189 }
190}
191
192fn apply_runtime_credentials(
193 token_slot: &Arc<std::sync::RwLock<String>>,
194 hmac_slot: &Arc<std::sync::RwLock<String>>,
195 saved_node_name: Option<&str>,
196 credentials: RegistrationCredentials,
197) -> String {
198 let token = credentials.node_token;
199 if let Some(hmac) = credentials.node_hmac_key.filter(|s| !s.is_empty()) {
200 if let Ok(mut guard) = hmac_slot.write() {
201 *guard = hmac;
202 }
203 }
204 if let Ok(mut guard) = token_slot.write() {
205 *guard = token.clone();
206 }
207 let hmac_snapshot = hmac_slot.read().map(|g| g.clone()).unwrap_or_default();
208 persist_saved_credentials(
209 saved_node_name,
210 &token,
211 Some(hmac_snapshot.as_str()).filter(|s| !s.is_empty()),
212 );
213 token
214}
215
216fn intent_for_model(model: &str, default_intent: &str) -> String {
222 if model.to_lowercase().contains("embed") {
223 "urn:iicp:intent:llm:embedding:v1".to_string()
224 } else {
225 default_intent.to_string()
226 }
227}
228
229fn modalities_for_model(model: &str) -> Vec<&'static str> {
236 let m = model.to_lowercase();
237 let has_image = m.contains("-vl-")
238 || m.ends_with("-vl")
239 || m.contains("vision")
240 || m.contains("llava")
241 || m.contains("omni");
242 let has_audio = m.contains("audio") || m.contains("voxtral") || m.contains("omni");
243 let mut mods = vec!["text"];
244 if has_image {
245 mods.push("image");
246 }
247 if has_audio {
248 mods.push("audio");
249 }
250 mods
251}
252
253fn build_capabilities(models: &[String], default_intent: &str, max_tokens: u32) -> Vec<Value> {
262 if models.is_empty() {
263 return vec![json!({
264 "intent": default_intent, "models": [], "max_tokens": max_tokens,
265 "input_modalities": ["text"],
266 })];
267 }
268 let mut order: Vec<String> = Vec::new();
270 let mut groups: HashMap<String, (String, Vec<&'static str>, Vec<String>)> = HashMap::new();
271 for m in models {
272 let intent = intent_for_model(m, default_intent);
273 let modalities = modalities_for_model(m);
274 let key = format!("{intent}\u{0}{}", modalities.join(","));
275 let entry = groups.entry(key.clone()).or_insert_with(|| {
276 order.push(key.clone());
277 (intent.clone(), modalities.clone(), Vec::new())
278 });
279 if !entry.2.contains(m) {
280 entry.2.push(m.clone());
281 }
282 }
283 order
284 .into_iter()
285 .map(|key| {
286 let (intent, modalities, models) = groups.remove(&key).expect("key from order");
287 json!({
288 "intent": intent,
289 "models": models,
290 "max_tokens": max_tokens,
291 "input_modalities": modalities,
292 })
293 })
294 .collect()
295}
296
297#[derive(Debug, Clone)]
299pub struct NodeConfig {
300 pub node_id: String,
301 pub endpoint: String,
302 pub intent: String,
303 pub model: Option<String>,
304 pub backend: Option<String>,
307 pub region: Option<String>,
308 pub capabilities: Vec<String>,
309 pub directory_url: String,
310 pub timeout_ms: u64,
311 pub max_concurrent: usize,
313 pub tokens_per_min: u32,
315 pub max_tokens: u32,
317 pub transport_endpoint: Option<String>,
322 pub transport_method: Option<String>,
330 pub nat_type: Option<String>,
333 pub transport_metadata: Option<serde_json::Value>,
335 pub exposure_mode: Option<String>,
338 pub cip_policy: Option<std::sync::Arc<crate::cip_policy::CooperativeInferencePolicy>>,
343 pub pricing: Option<crate::pricing::PricingConfig>,
346 pub node_hmac_key: String,
350 pub availability_windows: Vec<crate::availability::Window>,
354 pub enable_idempotency: bool,
358 pub enable_mesh: bool,
361 pub relay_capable: bool,
364 pub relay_accept_port: u16,
367 pub relay_worker_endpoint: Option<String>,
370 pub relay_bind_ticket_public_key_hex: Option<String>,
372 pub relay_require_bind_ticket: bool,
374 pub log_dir: Option<std::path::PathBuf>,
377 pub operator_delegation: Option<serde_json::Value>,
383 pub operator_display_name: Option<String>,
384 pub operator_created_at: Option<String>,
385 pub operator_integrity_hash: Option<String>,
386 pub policy_manifest: Option<serde_json::Value>,
390 pub backend_url: Option<String>,
395 pub backend_api_key: Option<String>,
397 pub excluded_models: Vec<String>,
400 pub supported_receipt_profiles: Vec<String>,
403}
404
405impl NodeConfig {
406 pub fn new(
407 node_id: impl Into<String>,
408 endpoint: impl Into<String>,
409 intent: impl Into<String>,
410 ) -> Self {
411 Self {
412 node_id: node_id.into(),
413 endpoint: endpoint.into(),
414 intent: intent.into(),
415 model: None,
416 backend: None,
417 region: None,
418 capabilities: vec![],
419 directory_url: DEFAULT_DIRECTORY.into(),
420 timeout_ms: 5_000,
421 max_concurrent: 4,
422 tokens_per_min: 10_000,
423 max_tokens: 8_192,
424 transport_endpoint: None,
425 transport_method: None,
426 nat_type: None,
427 transport_metadata: None,
428 exposure_mode: None,
429 cip_policy: None,
430 pricing: None,
431 node_hmac_key: String::new(),
432 availability_windows: Vec::new(),
433 enable_idempotency: false,
434 enable_mesh: false,
435 relay_capable: false,
436 relay_accept_port: 9485,
437 relay_worker_endpoint: None,
438 relay_bind_ticket_public_key_hex: std::env::var("IICP_RELAY_BIND_TICKET_PUBLIC_KEY")
439 .ok()
440 .filter(|s| !s.is_empty()),
441 relay_require_bind_ticket: std::env::var("IICP_RELAY_REQUIRE_BIND_TICKET")
442 .ok()
443 .as_deref()
444 == Some("1"),
445 log_dir: None,
446 operator_delegation: None,
447 operator_display_name: None,
448 operator_created_at: None,
449 operator_integrity_hash: None,
450 policy_manifest: None,
451 backend_url: None,
452 backend_api_key: None,
453 excluded_models: Vec::new(),
454 supported_receipt_profiles: Vec::new(),
455 }
456 }
457}
458
459#[derive(Debug, Deserialize)]
460pub struct TaskRequest {
461 pub task_id: String,
462 pub intent: String,
463 #[serde(default)]
464 pub payload: Value,
465 #[serde(default)]
466 pub iicp_conf: Option<HashMap<String, Value>>,
467 #[serde(default)]
468 pub cx_response_encryption: Option<String>,
469 pub constraints: Option<Value>,
470 pub auth: Option<Value>,
471 pub nonce: Option<String>,
472 pub source_node_id: Option<String>,
474 #[serde(skip_deserializing)]
476 pub _trace: Option<Value>,
477}
478
479pub fn derive_native_endpoint(endpoint: &str) -> Option<String> {
485 let (scheme, rest) = if let Some(r) = endpoint.strip_prefix("http://") {
486 ("iicp", r)
487 } else {
488 let r = endpoint.strip_prefix("https://")?;
489 ("iicpsec", r)
490 };
491 let authority = rest.split('/').next().unwrap_or(rest);
492 if authority.is_empty() {
493 return None;
494 }
495 Some(format!("{scheme}://{authority}"))
496}
497
498#[derive(Debug, Serialize)]
499pub struct TaskResponse {
500 pub task_id: String,
501 pub status: String,
502 #[serde(skip_serializing_if = "Option::is_none")]
503 pub result: Option<Value>,
504 #[serde(skip_serializing_if = "Option::is_none")]
505 pub iicp_conf_resp: Option<HashMap<String, Value>>,
506 #[serde(skip_serializing_if = "Option::is_none")]
507 pub error: Option<Value>,
508 pub generated_by_ai: bool,
509}
510
511pub type TaskHandlerFn = Arc<
512 dyn Fn(
513 TaskRequest,
514 ) -> std::pin::Pin<Box<dyn std::future::Future<Output = Result<Value>> + Send>>
515 + Send
516 + Sync,
517>;
518
519struct AppState {
520 handler: TaskHandlerFn,
521 node_id: String,
522 region: String,
523 intent: String,
524 model: String,
525 models: Vec<String>,
527 active_jobs: Arc<AtomicUsize>,
528 directory_url: String,
530 node_token: Arc<std::sync::RwLock<String>>,
532 node_hmac_key: Arc<std::sync::RwLock<String>>,
534 tasks_success: Arc<AtomicUsize>,
536 tasks_failed: Arc<AtomicUsize>,
537 tasks_latency_total_ms: Arc<AtomicUsize>,
538 max_concurrent: usize,
539 availability: Arc<crate::availability::AvailabilityEvaluator>,
540 cip_policy: Arc<crate::cip_policy::CooperativeInferencePolicy>,
542 idempotency: Arc<crate::idempotency::IdempotencyGuard>,
543 enable_idempotency: bool,
544 relay_bind_ticket_public_key_hex: Option<String>,
545 relay_require_bind_ticket: bool,
546 peer_manager: Arc<crate::peer_manager::PeerManager>,
547 http: reqwest::Client,
548 nonce_cache: Arc<Mutex<HashMap<String, Instant>>>,
549 pinhole_uid: Arc<std::sync::RwLock<Option<u32>>>,
551 pinhole_lease_seconds: Arc<std::sync::RwLock<u32>>,
552 #[cfg(feature = "iicp-tcp")]
554 relay_sessions: Arc<crate::relay_session::RelaySessionRegistry>,
555 task_rate_limit: u32,
559 task_rate_buckets: Arc<std::sync::Mutex<HashMap<String, (Instant, u32)>>>,
560 cx_private_key: Option<[u8; 32]>,
562 backend_stability: Arc<std::sync::RwLock<BackendStabilityObservation>>,
563}
564
565const TASK_RATE_WINDOW: Duration = Duration::from_secs(60);
566
567fn task_rate_step(
570 buckets: &mut HashMap<String, (Instant, u32)>,
571 limit: u32,
572 origin: &str,
573 now: Instant,
574) -> bool {
575 let entry = buckets.entry(origin.to_string()).or_insert((now, 0));
576 if now.duration_since(entry.0) >= TASK_RATE_WINDOW {
577 *entry = (now, 0);
578 }
579 entry.1 += 1;
580 let allowed = entry.1 <= limit;
581 if buckets.len() > 4096 {
582 buckets.retain(|_, (start, _)| now.duration_since(*start) < TASK_RATE_WINDOW);
583 }
584 allowed
585}
586
587fn task_rate_allow(state: &AppState, origin: &str) -> bool {
588 let mut buckets = state.task_rate_buckets.lock().unwrap();
589 task_rate_step(&mut buckets, state.task_rate_limit, origin, Instant::now())
590}
591
592async fn health_endpoint(State(state): State<Arc<AppState>>) -> impl IntoResponse {
595 let active = state.active_jobs.load(Ordering::Relaxed);
596 let uid = state.pinhole_uid.read().ok().and_then(|g| *g);
597 let lease = state
598 .pinhole_lease_seconds
599 .read()
600 .map(|g| *g)
601 .unwrap_or(3600);
602 let pinhole_state = if let Some(uid) = uid {
603 json!({ "active": true, "unique_id": uid, "lease_seconds": lease })
604 } else {
605 json!({ "active": false })
606 };
607 let eff_max = state
608 .availability
609 .effective_max_concurrent(state.max_concurrent);
610 Json(json!({
611 "status": "ok",
612 "node_id": state.node_id,
613 "region": state.region,
614 "load": (active as f64 / state.max_concurrent.max(1) as f64),
615 "active_jobs": active,
616 "max_concurrent": state.max_concurrent,
617 "effective_max_concurrent": eff_max,
618 "available": active < eff_max,
619 "model": state.model,
620 "models": state.models,
621 "intent": state.intent,
622 "pinhole_state": pinhole_state,
623 "backend_stability": state.backend_stability.read().expect("poisoned").public_json(),
624 }))
625}
626
627async fn metrics_endpoint() -> Response {
630 #[cfg(feature = "metrics")]
631 {
632 use prometheus::{Encoder, TextEncoder};
633 let encoder = TextEncoder::new();
634 let mf = prometheus::gather();
635 let mut buf = Vec::new();
636 if encoder.encode(&mf, &mut buf).is_ok() {
637 return (
638 StatusCode::OK,
639 [(
640 axum::http::header::CONTENT_TYPE,
641 "text/plain; version=0.0.4",
642 )],
643 buf,
644 )
645 .into_response();
646 }
647 }
648 (
649 StatusCode::SERVICE_UNAVAILABLE,
650 "metrics feature not enabled",
651 )
652 .into_response()
653}
654
655async fn peers_endpoint(
658 State(state): State<Arc<AppState>>,
659 headers: HeaderMap,
660 body: axum::body::Bytes,
661) -> Response {
662 let sig = headers
663 .get("x-iicp-signature")
664 .and_then(|v| v.to_str().ok());
665 if !state.peer_manager.verify_exchange(&body, sig) {
666 return (
667 StatusCode::UNAUTHORIZED,
668 Json(json!({"error":{"code":"IICP-E012","message":"invalid_signature"}})),
669 )
670 .into_response();
671 }
672 if let Ok(parsed) = serde_json::from_slice::<Value>(&body) {
673 if let Some(arr) = parsed.get("known_peers").and_then(Value::as_array) {
674 let dicts: Vec<Value> = arr.iter().filter(|p| p.is_object()).cloned().collect();
675 state.peer_manager.merge_peers(&dicts);
676 }
677 }
678 let peers: Vec<Value> = state
679 .peer_manager
680 .get_peers()
681 .iter()
682 .map(|p| {
683 json!({
684 "node_id": p.node_id,
685 "endpoint": p.endpoint,
686 "region": p.region,
687 "last_seen": p.last_seen,
688 })
689 })
690 .collect();
691 Json(json!({ "peers": peers })).into_response()
692}
693
694async fn relay_endpoint(
697 State(state): State<Arc<AppState>>,
698 Json(payload): Json<Value>,
699) -> Response {
700 let target_id = payload
701 .get("target_node_id")
702 .and_then(Value::as_str)
703 .unwrap_or("");
704 let task = payload.get("task");
705 if target_id.is_empty() || task.is_none() {
706 return (
707 StatusCode::UNPROCESSABLE_ENTITY,
708 Json(
709 json!({"error":{"code":"IICP-E000","message":"target_node_id and task required"}}),
710 ),
711 )
712 .into_response();
713 }
714 let task_val = task.expect("checked above").clone();
715
716 #[cfg(feature = "iicp-tcp")]
718 if let Some(session) = state.relay_sessions.get(target_id) {
719 match session.forward_task(&task_val, 120).await {
720 Ok(result) => {
721 let task_id = task_val
722 .get("task_id")
723 .and_then(Value::as_str)
724 .unwrap_or("");
725 return Json(json!({
726 "task_id": task_id,
727 "status": "success", "result": result
729 }))
730 .into_response();
731 }
732 Err(e) => {
733 return (
734 StatusCode::BAD_GATEWAY,
735 Json(json!({"error":{"code":"IICP-E031","message":format!("relay session forward failed: {e}")}})),
736 )
737 .into_response();
738 }
739 }
740 }
741
742 let target = match state.peer_manager.relay_target(target_id) {
744 Some(t) => t,
745 None => {
746 return (
747 StatusCode::NOT_FOUND,
748 Json(json!({"error":{"code":"IICP-E030","message":"target not in peer list and not a bound relay worker"}})),
749 )
750 .into_response();
751 }
752 };
753 let url = format!("{}/v1/task", target.endpoint.trim_end_matches('/'));
754 match state
755 .http
756 .post(&url)
757 .timeout(Duration::from_secs(120))
758 .json(&task_val)
759 .send()
760 .await
761 {
762 Ok(resp) => {
763 let status = StatusCode::from_u16(resp.status().as_u16()).unwrap_or(StatusCode::OK);
764 let bytes = resp.bytes().await.unwrap_or_default();
765 (status, bytes).into_response()
766 }
767 Err(e) => (
768 StatusCode::BAD_GATEWAY,
769 Json(json!({"error":{"code":"IICP-E031","message":format!("relay failed: {e}")}})),
770 )
771 .into_response(),
772 }
773}
774
775fn relay_cors(mut resp: Response) -> Response {
782 let h = resp.headers_mut();
783 h.insert("Access-Control-Allow-Origin", "*".parse().expect("static"));
784 h.insert(
785 "Access-Control-Allow-Methods",
786 "GET, POST, OPTIONS".parse().expect("static"),
787 );
788 h.insert(
789 "Access-Control-Allow-Headers",
790 "Content-Type, Authorization".parse().expect("static"),
791 );
792 resp
793}
794
795async fn relay_cors_preflight() -> Response {
796 let mut resp = StatusCode::NO_CONTENT.into_response();
797 resp.headers_mut()
798 .insert("Access-Control-Max-Age", "86400".parse().expect("static"));
799 relay_cors(resp)
800}
801
802#[cfg(feature = "iicp-tcp")]
803fn relay_authed_session(
804 state: &AppState,
805 headers: &HeaderMap,
806) -> Option<crate::relay_session::HttpPollWorkerSession> {
807 let auth = headers
808 .get("authorization")
809 .and_then(|v| v.to_str().ok())
810 .unwrap_or("");
811 let token = auth.strip_prefix("Bearer ").unwrap_or("");
812 state.relay_sessions.get_by_token(token)
813}
814
815#[cfg(feature = "iicp-tcp")]
816async fn relay_bind_endpoint(
817 State(state): State<Arc<AppState>>,
818 Json(payload): Json<Value>,
819) -> Response {
820 let worker_id = payload
821 .get("worker_id")
822 .and_then(Value::as_str)
823 .unwrap_or("");
824 if worker_id.is_empty() {
825 return relay_cors(
826 (
827 StatusCode::UNPROCESSABLE_ENTITY,
828 Json(json!({"error":{"code":"IICP-E001","message":"worker_id is required"}})),
829 )
830 .into_response(),
831 );
832 }
833 let bind_ticket = payload
834 .get("bind_ticket")
835 .and_then(Value::as_str)
836 .unwrap_or("");
837 let ticket_public_key = state
838 .relay_bind_ticket_public_key_hex
839 .as_deref()
840 .unwrap_or("");
841 let mut ticket_claims = None;
842 if !bind_ticket.is_empty() && !ticket_public_key.is_empty() {
843 let now_s = std::time::SystemTime::now()
844 .duration_since(std::time::UNIX_EPOCH)
845 .map(|d| d.as_secs() as i64)
846 .unwrap_or(0);
847 ticket_claims = crate::relay_ticket::verify_relay_bind_ticket(
848 bind_ticket,
849 ticket_public_key,
850 worker_id,
851 &state.node_id,
852 now_s,
853 );
854 if ticket_claims.is_none() {
855 return relay_cors(
856 (
857 StatusCode::UNAUTHORIZED,
858 Json(
859 json!({"error":{"code":"IICP-E040","message":"relay bind ticket invalid"}}),
860 ),
861 )
862 .into_response(),
863 );
864 }
865 } else if state.relay_require_bind_ticket {
866 return relay_cors(
867 (
868 StatusCode::UNAUTHORIZED,
869 Json(json!({"error":{"code":"IICP-E040","message":"relay bind ticket required"}})),
870 )
871 .into_response(),
872 );
873 } else if bind_ticket.is_empty() {
874 tracing::warn!("HTTP-poll relay bind without ticket: {}", worker_id);
875 }
876
877 if let Some(existing) = state.relay_sessions.get(worker_id) {
879 if existing.is_alive() {
880 return relay_cors((
881 StatusCode::CONFLICT,
882 Json(json!({"error":{"code":"IICP-E038","message":"worker_id has an alive relay session — rebind rejected"}})),
883 ).into_response());
884 }
885 }
886 let rebind = state.relay_sessions.get(worker_id).is_some();
887 let source = format!("worker:{worker_id}");
891 if !state
892 .relay_sessions
893 .allow_bind(&source, rebind, Instant::now())
894 {
895 return relay_cors(
896 (
897 StatusCode::TOO_MANY_REQUESTS,
898 Json(json!({"error":{
899 "code":"IICP-E039",
900 "reason":"relay_bind_rate_limited",
901 "message":"relay bind rate limit exceeded"
902 }})),
903 )
904 .into_response(),
905 );
906 }
907 if state.relay_sessions.at_capacity(worker_id) {
909 return relay_cors((
910 StatusCode::SERVICE_UNAVAILABLE,
911 Json(json!({"error":{"code":"IICP-E039","message":"relay at session capacity — try another relay"}})),
912 ).into_response());
913 }
914 if let Some(claims) = &ticket_claims {
915 let now_s = std::time::SystemTime::now()
916 .duration_since(std::time::UNIX_EPOCH)
917 .map(|d| d.as_secs() as i64)
918 .unwrap_or(0);
919 if !crate::relay_ticket::consume_relay_bind_ticket(claims, now_s) {
920 return relay_cors((
921 StatusCode::CONFLICT,
922 Json(json!({"error":{"code":"IICP-E040","message":"relay bind ticket replayed"}})),
923 ).into_response());
924 }
925 }
926 let intent = payload
927 .get("intent")
928 .and_then(Value::as_str)
929 .unwrap_or("")
930 .to_string();
931 let models: Vec<String> = payload
932 .get("models")
933 .and_then(Value::as_array)
934 .map(|a| {
935 a.iter()
936 .filter_map(Value::as_str)
937 .map(|s| s.to_string())
938 .collect()
939 })
940 .unwrap_or_default();
941 let session = crate::relay_session::HttpPollWorkerSession::new(
942 worker_id.to_string(),
943 intent,
944 models.clone(),
945 );
946 let token = session.session_token.clone();
947 state.relay_sessions.bind(
948 worker_id.to_string(),
949 crate::relay_session::RelaySession::HttpPoll(session),
950 );
951 tracing::info!(
952 "HTTP-poll relay worker bound: {} (models={})",
953 worker_id,
954 models.join(",")
955 );
956 relay_cors(
957 Json(json!({
958 "session_token": token,
959 "poll_timeout_s": 25,
960 "worker_endpoint_path": format!("/v1/relay-for/{worker_id}"),
961 }))
962 .into_response(),
963 )
964}
965
966#[cfg(feature = "iicp-tcp")]
967async fn relay_pull_endpoint(State(state): State<Arc<AppState>>, headers: HeaderMap) -> Response {
968 let Some(session) = relay_authed_session(&state, &headers) else {
969 return relay_cors((
970 StatusCode::UNAUTHORIZED,
971 Json(json!({"error":{"code":"IICP-E021","message":"invalid or missing relay session token"}})),
972 ).into_response());
973 };
974 match session.next_call(Duration::from_secs(25)).await {
975 Some(call) => relay_cors(Json(call).into_response()),
976 None => relay_cors(StatusCode::NO_CONTENT.into_response()),
977 }
978}
979
980#[cfg(feature = "iicp-tcp")]
981async fn relay_result_endpoint(
982 State(state): State<Arc<AppState>>,
983 headers: HeaderMap,
984 Json(payload): Json<Value>,
985) -> Response {
986 let Some(session) = relay_authed_session(&state, &headers) else {
987 return relay_cors((
988 StatusCode::UNAUTHORIZED,
989 Json(json!({"error":{"code":"IICP-E021","message":"invalid or missing relay session token"}})),
990 ).into_response());
991 };
992 let call_id = payload.get("call_id").and_then(Value::as_str).unwrap_or("");
993 let result = payload.get("result");
994 if call_id.is_empty() || !result.map(Value::is_object).unwrap_or(false) {
995 return relay_cors(
996 (
997 StatusCode::UNPROCESSABLE_ENTITY,
998 Json(json!({"error":{"code":"IICP-E001","message":"call_id and result are required"}})),
999 )
1000 .into_response(),
1001 );
1002 }
1003 session.on_response(call_id, result.expect("checked above").clone());
1004 relay_cors(StatusCode::NO_CONTENT.into_response())
1005}
1006
1007#[cfg(feature = "iicp-tcp")]
1008async fn relay_unbind_endpoint(State(state): State<Arc<AppState>>, headers: HeaderMap) -> Response {
1009 let Some(session) = relay_authed_session(&state, &headers) else {
1010 return relay_cors((
1011 StatusCode::UNAUTHORIZED,
1012 Json(json!({"error":{"code":"IICP-E021","message":"invalid or missing relay session token"}})),
1013 ).into_response());
1014 };
1015 session.close();
1016 state.relay_sessions.unbind(&session.worker_id);
1017 tracing::info!("HTTP-poll relay worker unbound: {}", session.worker_id);
1018 relay_cors(StatusCode::NO_CONTENT.into_response())
1019}
1020
1021#[cfg(feature = "iicp-tcp")]
1027async fn relay_for_task_endpoint(
1028 State(state): State<Arc<AppState>>,
1029 axum::extract::Path(wid): axum::extract::Path<String>,
1030 Json(task): Json<Value>,
1031) -> Response {
1032 let session = match state.relay_sessions.get(&wid) {
1033 Some(s) if s.is_alive() => s,
1034 _ => {
1035 return relay_cors((
1036 StatusCode::NOT_FOUND,
1037 Json(json!({"error":{"code":"IICP-E030","message":"no alive relay session for this worker"}})),
1038 ).into_response());
1039 }
1040 };
1041 match session.forward_task(&task, 120).await {
1042 Ok(result) => {
1043 let task_id = task.get("task_id").and_then(Value::as_str).unwrap_or("");
1044 let mut body = json!({"task_id": task_id, "status": "completed"});
1048 if let (Some(obj), Some(res_obj)) = (body.as_object_mut(), result.as_object()) {
1049 for (k, v) in res_obj {
1050 obj.insert(k.clone(), v.clone());
1051 }
1052 }
1053 relay_cors(Json(body).into_response())
1054 }
1055 Err(e) => relay_cors((
1056 StatusCode::BAD_GATEWAY,
1057 Json(json!({"error":{"code":"IICP-E031","message":format!("relay session forward failed: {e}")}})),
1058 ).into_response()),
1059 }
1060}
1061
1062#[cfg(feature = "iicp-tcp")]
1063async fn relay_for_health_endpoint(
1064 State(state): State<Arc<AppState>>,
1065 axum::extract::Path(wid): axum::extract::Path<String>,
1066) -> Response {
1067 match state.relay_sessions.get(&wid) {
1068 Some(s) if s.is_alive() => relay_cors(
1069 Json(json!({
1070 "status": "ok",
1071 "node_id": wid,
1072 "via_relay": true,
1073 "models": s.models(),
1074 }))
1075 .into_response(),
1076 ),
1077 _ => relay_cors((
1078 StatusCode::NOT_FOUND,
1079 Json(json!({"error":{"code":"IICP-E030","message":"no alive relay session for this worker"}})),
1080 ).into_response()),
1081 }
1082}
1083
1084fn canonical_json_node(v: &serde_json::Value) -> String {
1089 use serde_json::Value;
1090 match v {
1091 Value::Object(map) => {
1092 let mut keys: Vec<&String> = map.keys().collect();
1093 keys.sort();
1094 let parts: Vec<String> = keys
1095 .iter()
1096 .map(|k| {
1097 format!(
1098 "{}:{}",
1099 serde_json::to_string(k).unwrap_or_default(),
1100 canonical_json_node(&map[*k])
1101 )
1102 })
1103 .collect();
1104 format!("{{{}}}", parts.join(","))
1105 }
1106 Value::Array(arr) => {
1107 format!(
1108 "[{}]",
1109 arr.iter()
1110 .map(canonical_json_node)
1111 .collect::<Vec<_>>()
1112 .join(",")
1113 )
1114 }
1115 other => serde_json::to_string(other).unwrap_or_default(),
1116 }
1117}
1118
1119#[allow(clippy::too_many_arguments)]
1127async fn post_cip_receipt(
1128 http: reqwest::Client,
1129 directory_url: String,
1130 token: String,
1131 hmac_key: String,
1132 node_id: String,
1133 task_id: String,
1134 tokens_used: u64,
1135 result: serde_json::Value,
1136 querying_node_id: Option<String>,
1137) {
1138 use hmac::{Hmac, Mac};
1139 use sha2::{Digest, Sha256};
1140 type HmacSha256 = Hmac<Sha256>;
1141
1142 if token.is_empty() || hmac_key.is_empty() {
1143 return;
1144 }
1145
1146 let result_bytes = canonical_json_node(&result).into_bytes();
1147 let response_hash = hex::encode(Sha256::digest(&result_bytes));
1148
1149 let nonce: [u8; 16] = rand::random();
1150 let nonce = hex::encode(nonce);
1151
1152 let expires_at = {
1153 use chrono::Utc;
1154 (Utc::now() + chrono::Duration::seconds(300)).to_rfc3339()
1155 };
1156
1157 let querying_node_id = querying_node_id.filter(|s| !s.is_empty());
1160 let canonical = if let Some(ref qid) = querying_node_id {
1161 format!("{task_id}:{tokens_used}:::{nonce}:{response_hash}:{qid}")
1162 } else {
1163 format!("{task_id}:{tokens_used}:::{nonce}:{response_hash}")
1164 };
1165 let amount = (tokens_used.max(1) as f64) / 1000.0;
1166
1167 let mut mac = match HmacSha256::new_from_slice(hmac_key.as_bytes()) {
1168 Ok(m) => m,
1169 Err(_) => return,
1170 };
1171 mac.update(canonical.as_bytes());
1172 let signature = hex::encode(mac.finalize().into_bytes());
1173
1174 let url = format!("{}/v1/credits/award", directory_url.trim_end_matches('/'));
1175 let mut body = serde_json::json!({
1176 "node_id": node_id,
1177 "task_id": task_id,
1178 "tokens_used": tokens_used,
1179 "amount": amount,
1180 "nonce": nonce,
1181 "expires_at": expires_at,
1182 "signature": signature,
1183 "response_hash": response_hash,
1184 "reason": "task_completion",
1185 });
1186 if let Some(qid) = querying_node_id {
1188 body["querying_node_id"] = serde_json::Value::String(qid);
1189 }
1190 let _ = http
1191 .post(&url)
1192 .header("Authorization", format!("Bearer {token}"))
1193 .json(&body)
1194 .send()
1195 .await;
1196 }
1198
1199async fn admit(state: &AppState, qos: &str) -> bool {
1204 let cap = state
1207 .availability
1208 .effective_max_concurrent(state.max_concurrent);
1209 let prev = state.active_jobs.fetch_add(1, Ordering::Relaxed);
1210 if prev < cap {
1211 return true;
1212 }
1213 state.active_jobs.fetch_sub(1, Ordering::Relaxed);
1214 if !crate::scheduler::is_queue_eligible(qos) {
1215 return false;
1216 }
1217 let deadline = Instant::now() + crate::scheduler::QUEUE_WAIT;
1218 while Instant::now() < deadline {
1219 tokio::time::sleep(Duration::from_millis(50)).await;
1220 let cap = state
1221 .availability
1222 .effective_max_concurrent(state.max_concurrent);
1223 let prev = state.active_jobs.fetch_add(1, Ordering::Relaxed);
1224 if prev < cap {
1225 return true;
1226 }
1227 state.active_jobs.fetch_sub(1, Ordering::Relaxed);
1228 }
1229 false
1230}
1231
1232async fn task_endpoint(
1233 State(state): State<Arc<AppState>>,
1234 headers: HeaderMap,
1235 Json(mut req): Json<TaskRequest>,
1236) -> Response {
1237 if state.task_rate_limit > 0 {
1240 if let Some(origin) = headers.get("origin").and_then(|v| v.to_str().ok()) {
1241 if !task_rate_allow(&state, origin) {
1242 return (
1243 StatusCode::TOO_MANY_REQUESTS,
1244 [("Retry-After", "60"), ("Content-Type", "application/json")],
1245 Json(json!({
1246 "error": { "code": "IICP-E023", "message": "per-origin task rate limit exceeded" }
1247 })),
1248 )
1249 .into_response();
1250 }
1251 }
1252 }
1253
1254 if !state.cip_policy.permits_intent(&req.intent) {
1259 return (
1260 StatusCode::FORBIDDEN,
1261 Json(json!({
1262 "error": {
1263 "code": "tool_execution_denied",
1264 "message": "Tool-execution intents are not permitted by this node's CIP policy",
1265 }
1266 })),
1267 )
1268 .into_response();
1269 }
1270
1271 let stability = state.backend_stability.read().expect("poisoned").clone();
1273 if stability.is_draining() {
1274 if let Some(retry) = stability.retry_after_s() {
1275 return (
1276 StatusCode::SERVICE_UNAVAILABLE,
1277 [
1278 ("Retry-After", retry.to_string()),
1279 ("Content-Type", "application/json".to_string()),
1280 ],
1281 Json(json!({
1282 "error": {
1283 "code": "IICP-E024",
1284 "message": "backend temporarily draining",
1285 "reason": stability.reason_class,
1286 "retry_after_ms": retry * 1000,
1287 }
1288 })),
1289 )
1290 .into_response();
1291 }
1292 }
1293
1294 let qos = req
1296 .constraints
1297 .as_ref()
1298 .and_then(|c| c.get("qos_class"))
1299 .and_then(|v| v.as_str())
1300 .unwrap_or("best_effort")
1301 .to_string();
1302 if !admit(&state, &qos).await {
1303 return (
1304 StatusCode::TOO_MANY_REQUESTS,
1305 [("Retry-After", "2"), ("Content-Type", "application/json")],
1306 Json(json!({
1307 "error": {
1308 "code": "IICP-E021",
1309 "message": "capacity_exceeded",
1310 "qos_class": qos,
1311 "retry_after_ms": 2000,
1312 }
1313 })),
1314 )
1315 .into_response();
1316 }
1317
1318 if let Some(ref nonce) = req.nonce {
1320 let mut cache = state.nonce_cache.lock().await;
1321 cache.retain(|_, inserted_at| inserted_at.elapsed().as_secs() < NONCE_TTL_SECS);
1322 if cache.contains_key(nonce) {
1323 state.active_jobs.fetch_sub(1, Ordering::Relaxed);
1324 return (
1325 StatusCode::CONFLICT,
1326 Json(json!({
1327 "error": { "code": "IICP-E011", "message": "replay_detected" }
1328 })),
1329 )
1330 .into_response();
1331 }
1332 cache.insert(nonce.clone(), Instant::now());
1333 }
1334
1335 if state.enable_idempotency && !state.idempotency.check_and_register(&req.task_id) {
1338 state.active_jobs.fetch_sub(1, Ordering::Relaxed);
1339 return (
1340 StatusCode::CONFLICT,
1341 Json(json!({
1342 "error": { "code": "IICP-E010", "message": "duplicate_task" }
1343 })),
1344 )
1345 .into_response();
1346 }
1347
1348 let mut cx_shared_secret: Option<[u8; 32]> = None;
1349 if req.payload.is_null() {
1350 if let Some(conf) = req.iicp_conf.as_ref() {
1351 match state
1352 .cx_private_key
1353 .as_ref()
1354 .ok_or_else(|| IicpError::Node("node has no CX private key".to_string()))
1355 .and_then(|private_key| {
1356 crate::confidentiality::decrypt_payload_with_context(conf, private_key)
1357 }) {
1358 Ok((payload, shared_secret)) => {
1359 req.payload = payload;
1360 cx_shared_secret = Some(shared_secret);
1361 }
1362 Err(err) => {
1363 state.active_jobs.fetch_sub(1, Ordering::Relaxed);
1364 return (
1365 StatusCode::BAD_REQUEST,
1366 Json(json!({
1367 "error": {
1368 "code": "IICP-CX-02",
1369 "message": format!("iicp_conf decrypt failed: {err}")
1370 }
1371 })),
1372 )
1373 .into_response();
1374 }
1375 }
1376 }
1377 }
1378
1379 if let Some(tp) = headers.get("traceparent").and_then(|v| v.to_str().ok()) {
1381 req._trace = Some(json!({ "traceparent": tp }));
1382 }
1383
1384 let task_id = req.task_id.clone();
1385 let cx_response_encryption_required = req.cx_response_encryption.as_deref() == Some("required");
1386 if cx_response_encryption_required && cx_shared_secret.is_none() {
1387 state.active_jobs.fetch_sub(1, Ordering::Relaxed);
1388 return (
1389 StatusCode::BAD_REQUEST,
1390 Json(json!({
1391 "error": {
1392 "code": "IICP-CX-03",
1393 "message": "encrypted response requested without an encrypted request",
1394 }
1395 })),
1396 )
1397 .into_response();
1398 }
1399 let querying_node_id = req.source_node_id.clone();
1401 let started = Instant::now();
1406 let result = {
1407 let span = tracing::info_span!(
1408 "iicp.task.execute",
1409 "iicp.task_id" = %task_id,
1410 "iicp.intent" = %req.intent,
1411 );
1412 let _guard = span.enter();
1413 (state.handler)(req).await
1414 };
1415 state.active_jobs.fetch_sub(1, Ordering::Relaxed);
1416
1417 match result {
1418 Ok(value) => {
1419 let latency_ms = started.elapsed().as_millis().min(usize::MAX as u128) as usize;
1420 state.tasks_success.fetch_add(1, Ordering::Relaxed);
1421 if latency_ms > 0 {
1422 state
1423 .tasks_latency_total_ms
1424 .fetch_add(latency_ms, Ordering::Relaxed);
1425 }
1426 let hmac_key = state.node_hmac_key.read().expect("poisoned").clone();
1429 if !hmac_key.is_empty() {
1430 let token = state.node_token.read().expect("poisoned").clone();
1431 let tokens_used: u64 = value
1435 .get("usage")
1436 .and_then(|u| u.get("total_tokens"))
1437 .and_then(|t| t.as_u64())
1438 .unwrap_or(0);
1439 tokio::spawn(post_cip_receipt(
1440 state.http.clone(),
1441 state.directory_url.clone(),
1442 token,
1443 hmac_key,
1444 state.node_id.clone(),
1445 task_id.clone(),
1446 tokens_used,
1447 value.clone(),
1448 querying_node_id,
1450 ));
1451 }
1452 let plain_response = json!({
1453 "task_id": task_id,
1454 "status": "success",
1455 "result": value,
1456 "generated_by_ai": true
1457 });
1458 let encrypted_response = if cx_response_encryption_required {
1459 cx_shared_secret
1460 .as_ref()
1461 .map(|secret| {
1462 crate::confidentiality::encrypt_response(&plain_response, secret, &task_id)
1463 })
1464 .transpose()
1465 } else {
1466 Ok(None)
1467 };
1468 let encrypted_response = match encrypted_response {
1469 Ok(value) => value,
1470 Err(err) => {
1471 return (
1472 StatusCode::INTERNAL_SERVER_ERROR,
1473 Json(TaskResponse {
1474 task_id,
1475 status: "error".into(),
1476 result: None,
1477 iicp_conf_resp: None,
1478 error: Some(json!({"message": err.to_string()})),
1479 generated_by_ai: false,
1480 }),
1481 )
1482 .into_response();
1483 }
1484 };
1485 Json(TaskResponse {
1486 task_id,
1487 status: if encrypted_response.is_some() {
1491 "encrypted".into()
1492 } else {
1493 "success".into()
1494 },
1495 result: if encrypted_response.is_some() {
1496 None
1497 } else {
1498 Some(value)
1499 },
1500 iicp_conf_resp: encrypted_response,
1501 error: None,
1502 generated_by_ai: true,
1503 })
1504 .into_response()
1505 }
1506 Err(e) => {
1507 let latency_ms = started.elapsed().as_millis().min(usize::MAX as u128) as usize;
1508 state.tasks_failed.fetch_add(1, Ordering::Relaxed);
1509 if latency_ms > 0 {
1510 state
1511 .tasks_latency_total_ms
1512 .fetch_add(latency_ms, Ordering::Relaxed);
1513 }
1514 (
1515 StatusCode::INTERNAL_SERVER_ERROR,
1516 Json(TaskResponse {
1517 task_id,
1518 status: "error".into(),
1519 result: None,
1520 iicp_conf_resp: None,
1521 error: Some(json!({ "message": e.to_string() })),
1522 generated_by_ai: false,
1523 }),
1524 )
1525 .into_response()
1526 }
1527 }
1528}
1529
1530pub struct IicpNode {
1534 cfg: NodeConfig,
1535 http: Client,
1536 runtime_hmac_key: Arc<std::sync::RwLock<String>>,
1542 runtime_token: Arc<std::sync::RwLock<String>>,
1545 #[allow(dead_code)]
1549 pinhole_uid: std::sync::RwLock<Option<u32>>,
1550 #[allow(dead_code)]
1551 pinhole_lease_seconds: std::sync::RwLock<u32>,
1552 liveness_challenge: Arc<std::sync::RwLock<Option<String>>>,
1555 registered_models: Arc<std::sync::RwLock<Vec<String>>>,
1558 endpoint_override: Arc<std::sync::RwLock<Option<String>>>,
1563 runtime_available: Arc<AtomicBool>,
1566 runtime_relay_bound: Arc<AtomicBool>,
1570 registered_endpoint: Arc<std::sync::RwLock<String>>,
1574 cx_public_key: Option<crate::types::CxPublicKey>,
1576 cx_private_key: Option<[u8; 32]>,
1577 backend_stability: Arc<std::sync::RwLock<BackendStabilityObservation>>,
1578 saved_node_name: Option<String>,
1582}
1583
1584impl IicpNode {
1585 pub fn new(cfg: NodeConfig) -> Self {
1586 let http = Client::builder()
1587 .timeout(Duration::from_millis(cfg.timeout_ms + 2_000))
1588 .use_rustls_tls()
1589 .build()
1590 .expect("failed to build HTTP client");
1591 let runtime_hmac_key = Arc::new(std::sync::RwLock::new(cfg.node_hmac_key.clone()));
1592 let (cx_public_key, cx_private_key) =
1593 match crate::confidentiality::load_or_create_node_cx_key(&cfg.node_id, &cfg.endpoint) {
1594 Ok((public_key, private_key)) => (Some(public_key), Some(private_key)),
1595 Err(err) => {
1596 eprintln!(
1597 "[iicp-node] IICP-CX provider key unavailable; node will not advertise CX: {err}"
1598 );
1599 (None, None)
1600 }
1601 };
1602 Self {
1603 cfg,
1604 http,
1605 runtime_hmac_key,
1606 runtime_token: Arc::new(std::sync::RwLock::new(String::new())),
1607 pinhole_uid: std::sync::RwLock::new(None),
1608 pinhole_lease_seconds: std::sync::RwLock::new(3600),
1609 liveness_challenge: Arc::new(std::sync::RwLock::new(None)),
1610 registered_models: Arc::new(std::sync::RwLock::new(Vec::new())),
1611 endpoint_override: Arc::new(std::sync::RwLock::new(None)),
1612 runtime_available: Arc::new(AtomicBool::new(true)),
1613 runtime_relay_bound: Arc::new(AtomicBool::new(false)),
1614 registered_endpoint: Arc::new(std::sync::RwLock::new(String::new())),
1615 cx_public_key,
1616 cx_private_key,
1617 backend_stability: Arc::new(std::sync::RwLock::new(
1618 BackendStabilityObservation::default(),
1619 )),
1620 saved_node_name: None,
1621 }
1622 }
1623
1624 pub fn set_saved_node_name(&mut self, name: impl Into<String>) {
1628 let name = name.into();
1629 self.saved_node_name = (!name.trim().is_empty()).then_some(name);
1630 }
1631
1632 fn effective_endpoint(&self) -> String {
1635 self.endpoint_override
1636 .read()
1637 .expect("poisoned")
1638 .clone()
1639 .unwrap_or_else(|| self.cfg.endpoint.clone())
1640 }
1641
1642 pub fn set_endpoint(&self, endpoint: String) {
1648 let mut g = self.endpoint_override.write().expect("poisoned");
1649 if endpoint.is_empty() {
1650 *g = None;
1651 } else {
1652 *g = Some(endpoint);
1653 }
1654 }
1655
1656 pub fn endpoint_override_handle(&self) -> Arc<std::sync::RwLock<Option<String>>> {
1657 Arc::clone(&self.endpoint_override)
1658 }
1659
1660 pub fn runtime_available_handle(&self) -> Arc<AtomicBool> {
1661 Arc::clone(&self.runtime_available)
1662 }
1663
1664 pub fn node_hmac_key(&self) -> String {
1667 self.runtime_hmac_key.read().expect("poisoned").clone()
1668 }
1669
1670 pub fn cfg(&self) -> &NodeConfig {
1674 &self.cfg
1675 }
1676
1677 pub fn registered_models(&self) -> &Arc<std::sync::RwLock<Vec<String>>> {
1679 &self.registered_models
1680 }
1681
1682 #[doc(hidden)]
1683 pub fn set_backend_stability_for_test(&self, observation: BackendStabilityObservation) {
1684 *self.backend_stability.write().expect("poisoned") = observation;
1685 }
1686
1687 fn set_backend_stability(&self, observation: BackendStabilityObservation) {
1688 let mut guard = self.backend_stability.write().expect("poisoned");
1689 if guard.is_draining() && !observation.is_draining() {
1690 return;
1691 }
1692 *guard = observation;
1693 }
1694
1695 async fn observe_backend_stability(&self) -> BackendStabilityObservation {
1696 let obs = if let Some(url) = self.cfg.backend_url.as_deref() {
1697 observe_backend_stability(
1698 &self.http,
1699 url,
1700 self.cfg.backend.as_deref(),
1701 self.cfg.model.as_deref(),
1702 self.cfg.backend_api_key.as_deref(),
1703 )
1704 .await
1705 } else {
1706 BackendStabilityObservation::default()
1707 };
1708 self.set_backend_stability(obs.clone());
1709 obs
1710 }
1711
1712 pub async fn check_model_drift_and_reregister(&self) {
1715 let registered_ep = self.registered_endpoint.read().expect("poisoned").clone();
1723 let endpoint_changed =
1724 !registered_ep.is_empty() && self.effective_endpoint() != registered_ep;
1725
1726 let live = self.probe_health_models().await.unwrap_or_default();
1728 let models_changed = if live.is_empty() {
1729 false
1730 } else {
1731 let registered = self.registered_models.read().expect("poisoned").clone();
1732 let live_set: std::collections::HashSet<_> = live.iter().cloned().collect();
1733 let reg_set: std::collections::HashSet<_> = registered.into_iter().collect();
1734 live_set != reg_set
1735 };
1736
1737 if !endpoint_changed && !models_changed {
1738 return;
1739 }
1740
1741 let mut new_payload = self.build_register_payload(); if models_changed {
1743 let new_caps = build_capabilities(&live, &self.cfg.intent, self.cfg.max_tokens);
1744 new_payload["capabilities"] = serde_json::to_value(&new_caps).unwrap_or(json!([]));
1745 }
1746 let url = format!(
1747 "{}/v1/register",
1748 self.cfg.directory_url.trim_end_matches('/')
1749 );
1750 if let Some(credentials) = reregister(&self.http, &url, &new_payload).await {
1751 if models_changed {
1752 *self.registered_models.write().expect("poisoned") = live;
1753 }
1754 *self.registered_endpoint.write().expect("poisoned") = self.effective_endpoint();
1755 apply_runtime_credentials(
1756 &self.runtime_token,
1757 &self.runtime_hmac_key,
1758 self.saved_node_name.as_deref(),
1759 credentials,
1760 );
1761 if endpoint_changed {
1762 tracing::info!(
1763 "[iicp-node] re-registered after endpoint rotation → {}",
1764 self.effective_endpoint()
1765 );
1766 }
1767 }
1768 }
1769
1770 pub fn set_relay_worker_endpoint(&mut self, endpoint: String) {
1774 self.cfg.relay_worker_endpoint = Some(endpoint);
1775 }
1776
1777 pub fn set_transport_endpoint(&mut self, endpoint: String) {
1780 self.cfg.transport_endpoint = Some(endpoint);
1781 }
1782
1783 #[cfg(feature = "nat")]
1794 pub fn apply_nat_profile(&mut self, profile: &crate::nat_detection::NatProfile) {
1795 if profile.is_reachable() {
1796 if let Some(pub_ep) = &profile.public_endpoint {
1797 self.cfg.endpoint = pub_ep.clone();
1798 }
1799 }
1800 if let Some(tep) = &profile.transport_endpoint {
1801 self.cfg.transport_endpoint = Some(tep.clone());
1802 }
1803 let tm = match profile.transport_method {
1804 crate::nat_detection::TransportMethod::Direct => Some("direct"),
1805 crate::nat_detection::TransportMethod::UpnpMapped => Some("upnp_mapped"),
1806 crate::nat_detection::TransportMethod::StunHolePunch => Some("stun_hole_punch"),
1807 crate::nat_detection::TransportMethod::TurnRelay => Some("turn_relay"),
1808 crate::nat_detection::TransportMethod::ExternalTunnel => Some("external_tunnel"),
1809 crate::nat_detection::TransportMethod::Unreachable => None,
1810 };
1811 if let Some(name) = tm {
1812 self.cfg.transport_method = Some(name.into());
1813 }
1814 if self.cfg.nat_type.is_none() {
1815 self.cfg.nat_type = Some("unknown".into());
1816 }
1817 let tail: Vec<&str> = profile
1818 .detection_log
1819 .iter()
1820 .rev()
1821 .take(1)
1822 .map(|s| s.as_str())
1823 .collect();
1824 self.cfg.transport_metadata = Some(serde_json::json!({
1825 "tier": profile.tier,
1826 "detection_log_tail": tail,
1827 }));
1828 self.cfg.exposure_mode = Some(
1831 crate::qualify::qualify_service(profile)
1832 .exposure_mode
1833 .to_string(),
1834 );
1835 if let Some(v6) = &profile.ipv6 {
1837 if v6.pinhole_active {
1838 if let Some(uid) = v6.pinhole_unique_id {
1839 if let Ok(mut slot) = self.pinhole_uid.write() {
1840 *slot = Some(uid);
1841 }
1842 }
1843 if let Some(lease) = v6.pinhole_lease_seconds {
1844 if let Ok(mut slot) = self.pinhole_lease_seconds.write() {
1845 *slot = lease;
1846 }
1847 }
1848 }
1849 }
1850 }
1851
1852 #[cfg(feature = "nat")]
1854 pub async fn revoke_pinhole(&self) -> bool {
1855 let uid = match self.pinhole_uid.write() {
1856 Ok(mut slot) => slot.take(),
1857 Err(_) => None,
1858 };
1859 match uid {
1860 Some(uid) => crate::nat_detection::delete_ipv6_pinhole(uid).await,
1861 None => false,
1862 }
1863 }
1864
1865 pub fn seed_token(&self, token: &str) {
1873 if !token.is_empty() {
1874 *self.runtime_token.write().expect("poisoned") = token.to_string();
1875 }
1876 }
1877
1878 pub async fn deregister(&self, node_token: Option<&str>) -> Result<()> {
1881 let stashed = self.runtime_token.read().expect("poisoned").clone();
1882 let token = node_token.map(str::to_string).unwrap_or(stashed);
1883 if token.is_empty() {
1884 return Err(crate::errors::IicpError::Node(
1885 "deregister() requires a node_token (none stashed — call register() first)".into(),
1886 ));
1887 }
1888 let url = format!(
1889 "{}/v1/register",
1890 self.cfg.directory_url.trim_end_matches('/')
1891 );
1892 let resp = self
1893 .http
1894 .delete(&url)
1895 .bearer_auth(&token)
1896 .json(&serde_json::json!({"node_id": self.cfg.node_id}))
1897 .send()
1898 .await?;
1899 let status = resp.status();
1900 if !status.is_success() && status.as_u16() != 404 {
1901 return Err(crate::errors::IicpError::Node(format!(
1902 "Deregister failed: {status}"
1903 )));
1904 }
1905 Ok(())
1906 }
1907
1908 #[doc(hidden)]
1919 pub fn register_payload_for_test(&self) -> Value {
1920 self.build_register_payload()
1921 }
1922
1923 fn build_register_payload(&self) -> Value {
1924 let mut models: Vec<String> = match &self.cfg.model {
1927 Some(m) => vec![m.clone()],
1928 None => Vec::new(),
1929 };
1930 for cap in &self.cfg.capabilities {
1931 if !models.contains(cap) {
1932 models.push(cap.clone());
1933 }
1934 }
1935 let region = self
1936 .cfg
1937 .region
1938 .clone()
1939 .unwrap_or_else(|| "unknown".to_string());
1940
1941 let mut payload = json!({
1942 "endpoint": self.effective_endpoint(),
1944 "region": region,
1945 "capabilities": build_capabilities(&models, &self.cfg.intent, self.cfg.max_tokens),
1949 "limits": {
1950 "max_concurrent": self.cfg.max_concurrent,
1951 "tokens_per_min": self.cfg.tokens_per_min,
1952 },
1953 });
1954 if !self.cfg.node_id.is_empty() {
1955 payload["node_id"] = json!(self.cfg.node_id);
1956 }
1957 let stashed_token = self.runtime_token.read().expect("poisoned").clone();
1961 if !stashed_token.is_empty() {
1962 payload["current_node_token"] = json!(stashed_token);
1963 }
1964 if let Some(t) = &self.cfg.transport_endpoint {
1965 payload["transport_endpoint"] = json!(t);
1966 }
1967 if let Some(m) = &self.cfg.transport_method {
1968 payload["transport_method"] = json!(m);
1969 }
1970 if let Some(n) = &self.cfg.nat_type {
1971 payload["nat_type"] = json!(n);
1972 }
1973 if let Some(md) = &self.cfg.transport_metadata {
1974 payload["transport_metadata"] = md.clone();
1975 }
1976 if let Some(e) = &self.cfg.exposure_mode {
1977 payload["exposure_mode"] = json!(e);
1978 }
1979 let receipt_profiles: Vec<&str> = self
1980 .cfg
1981 .supported_receipt_profiles
1982 .iter()
1983 .map(String::as_str)
1984 .filter(|profile| *profile == "consumer_cosignature_v1")
1985 .collect::<std::collections::HashSet<_>>()
1986 .into_iter()
1987 .collect();
1988 if !receipt_profiles.is_empty() {
1989 payload["supported_receipt_profiles"] = json!(receipt_profiles);
1990 }
1991 payload["sdk_language"] = json!("rust");
1992 payload["sdk_version"] = json!(env!("CARGO_PKG_VERSION"));
1993 attach_update_status(&mut payload);
1994 if let Some(cx_public_key) = &self.cx_public_key {
1995 payload["cx_public_key"] = json!(cx_public_key);
1996 }
1997 if self.cfg.relay_capable {
1998 payload["relay_capable"] = json!(true);
1999 payload["relay_accept_port"] = json!(self.cfg.relay_accept_port);
2000 }
2001 if let Some(b) = &self.cfg.backend {
2002 payload["backend"] = json!(b);
2003 }
2004 let policy_arc = self
2005 .cfg
2006 .cip_policy
2007 .clone()
2008 .unwrap_or_else(crate::cip_policy::get_cip_policy);
2009 if let Some(block) = policy_arc.as_register_policy_block() {
2010 payload["policy"] = block;
2011 }
2012 if let Some(pricing) = &self.cfg.pricing {
2013 let hmac_key = self.runtime_hmac_key.read().expect("poisoned").clone();
2014 payload["pricing"] = crate::pricing::build_pricing_block(pricing, &hmac_key);
2015 }
2016 if !self.cfg.node_hmac_key.is_empty() {
2017 payload["node_hmac_key"] = json!(self.cfg.node_hmac_key);
2018 }
2019 if let Some(del) = &self.cfg.operator_delegation {
2022 payload["operator_delegation"] = del.clone();
2023 if let Some(dn) = &self.cfg.operator_display_name {
2024 payload["operator_display_name"] = json!(dn);
2025 }
2026 if let Some(ca) = &self.cfg.operator_created_at {
2027 payload["operator_created_at"] = json!(ca);
2028 }
2029 if let Some(ih) = &self.cfg.operator_integrity_hash {
2030 payload["operator_integrity_hash"] = json!(ih);
2031 }
2032 }
2033 if let Some(manifest) = &self.cfg.policy_manifest {
2034 payload["policy_manifest"] = manifest.clone();
2035 }
2036 payload
2037 }
2038
2039 pub async fn register(&self) -> Result<String> {
2040 let payload = self.build_register_payload();
2041
2042 let resp = self
2043 .http
2044 .post(format!(
2045 "{}/v1/register",
2046 self.cfg.directory_url.trim_end_matches('/')
2047 ))
2048 .json(&payload)
2049 .send()
2050 .await
2051 .map_err(|e| IicpError::Node(e.to_string()))?;
2052
2053 if !resp.status().is_success() {
2054 return Err(IicpError::Node(format!(
2055 "register failed: {}",
2056 resp.status()
2057 )));
2058 }
2059 let data: Value = resp
2060 .json()
2061 .await
2062 .map_err(|e| IicpError::Node(e.to_string()))?;
2063 let token = data["node_token"]
2064 .as_str()
2065 .or_else(|| data["token"].as_str())
2066 .ok_or_else(|| IicpError::Node(format!("no node_token in response: {data}")))?;
2067 *self.runtime_token.write().expect("poisoned") = token.to_string();
2069 if self.cfg.node_hmac_key.is_empty() {
2073 if let Some(dir_key) = data["node_hmac_key"].as_str() {
2074 if !dir_key.is_empty() {
2075 let mut guard = self.runtime_hmac_key.write().expect("poisoned");
2076 *guard = dir_key.to_string();
2077 }
2078 }
2079 }
2080 let hmac_snapshot = self
2081 .runtime_hmac_key
2082 .read()
2083 .map(|g| g.clone())
2084 .unwrap_or_default();
2085 persist_saved_credentials(
2086 self.saved_node_name.as_deref(),
2087 token,
2088 Some(hmac_snapshot.as_str()).filter(|s| !s.is_empty()),
2089 );
2090 {
2092 let mut models: Vec<String> = match &self.cfg.model {
2093 Some(m) => vec![m.clone()],
2094 None => Vec::new(),
2095 };
2096 for cap in &self.cfg.capabilities {
2097 if !models.contains(cap) {
2098 models.push(cap.clone());
2099 }
2100 }
2101 *self.registered_models.write().expect("poisoned") = models;
2102 }
2103 *self.registered_endpoint.write().expect("poisoned") = self.effective_endpoint();
2106 Ok(token.to_string())
2107 }
2108
2109 async fn probe_health_models(&self) -> Option<Vec<String>> {
2113 let base = self.cfg.backend_url.as_deref()?;
2114 probe_health_models_bg(
2115 &self.http,
2116 base,
2117 &self.cfg.backend_api_key,
2118 &self.cfg.excluded_models,
2119 )
2120 .await
2121 }
2122
2123 pub async fn heartbeat(&self, node_token: &str) -> Result<()> {
2125 let public_available = self.runtime_available.load(Ordering::Relaxed);
2126 let mut body = json!({
2127 "node_id": self.cfg.node_id,
2128 "node_token": node_token,
2129 "status": if public_available { "available" } else { "recovering" },
2130 "available": public_available,
2136 "max_concurrent": crate::availability::AvailabilityEvaluator::new(
2138 self.cfg.availability_windows.clone(),
2139 )
2140 .effective_max_concurrent(self.cfg.max_concurrent),
2141 });
2142 attach_update_status(&mut body);
2143 let hmac_key = self.node_hmac_key();
2147 let stored = self.liveness_challenge.read().expect("poisoned").clone();
2148 if let Some(ch) = &stored {
2149 if !hmac_key.is_empty() {
2150 body["challenge_response"] =
2151 json!(crate::pricing::sign_body(ch.as_bytes(), &hmac_key));
2152 }
2153 }
2154
2155 if self.cfg.backend_url.is_some() {
2157 if let Some(models) = self.probe_health_models().await {
2158 body["health_models"] = json!(models);
2159 }
2160 let stability = self.observe_backend_stability().await;
2161 body["backend_stability"] = stability.public_json();
2162 }
2163
2164 let resp = self
2165 .http
2166 .post(format!(
2170 "{}/v1/heartbeat",
2171 self.cfg.directory_url.trim_end_matches('/')
2172 ))
2173 .bearer_auth(node_token)
2176 .json(&body)
2177 .send()
2178 .await
2179 .map_err(|e| IicpError::Node(e.to_string()))?;
2180
2181 if !resp.status().is_success() {
2182 return Err(IicpError::Node(format!(
2183 "heartbeat failed: {}",
2184 resp.status()
2185 )));
2186 }
2187 if let Ok(data) = resp.json::<Value>().await {
2189 if let Some(ch) = data["challenge"].as_str() {
2190 *self.liveness_challenge.write().expect("poisoned") = Some(ch.to_string());
2191 }
2192 }
2193 Ok(())
2194 }
2195
2196 pub async fn serve<F, Fut>(
2201 &self,
2202 handler: F,
2203 addr: &str,
2204 node_token: Option<String>,
2205 ) -> Result<()>
2206 where
2207 F: Fn(TaskRequest) -> Fut + Send + Sync + 'static,
2208 Fut: std::future::Future<Output = Result<Value>> + Send + 'static,
2209 {
2210 let handler: TaskHandlerFn = Arc::new(move |req| Box::pin(handler(req)));
2211 #[cfg(feature = "iicp-tcp")]
2213 let handler_for_relay = Arc::clone(&handler);
2214 #[cfg(feature = "iicp-tcp")]
2216 let handler_for_native = Arc::clone(&handler);
2217 #[cfg(feature = "iicp-tcp")]
2219 let bind_host: String = addr.split(':').next().unwrap_or("0.0.0.0").to_string();
2220 let active_jobs = Arc::new(AtomicUsize::new(0));
2221 let nonce_cache = Arc::new(Mutex::new(HashMap::new()));
2222 let shared_pinhole_uid: Arc<std::sync::RwLock<Option<u32>>> = Arc::new(
2224 std::sync::RwLock::new(self.pinhole_uid.read().ok().and_then(|g| *g)),
2225 );
2226 let shared_pinhole_lease: Arc<std::sync::RwLock<u32>> = Arc::new(std::sync::RwLock::new(
2227 self.pinhole_lease_seconds
2228 .read()
2229 .map(|g| *g)
2230 .unwrap_or(3600),
2231 ));
2232
2233 let tasks_success = Arc::new(AtomicUsize::new(0));
2234 let tasks_failed = Arc::new(AtomicUsize::new(0));
2235 let tasks_latency_total_ms = Arc::new(AtomicUsize::new(0));
2236 let mut all_models: Vec<String> = match &self.cfg.model {
2237 Some(m) => vec![m.clone()],
2238 None => Vec::new(),
2239 };
2240 for cap in &self.cfg.capabilities {
2241 if !all_models.contains(cap) {
2242 all_models.push(cap.clone());
2243 }
2244 }
2245 let state = Arc::new(AppState {
2246 handler,
2247 node_id: self.cfg.node_id.clone(),
2248 region: self.cfg.region.clone().unwrap_or_else(|| "unknown".into()),
2249 intent: self.cfg.intent.clone(),
2250 model: self.cfg.model.clone().unwrap_or_default(),
2251 models: all_models,
2252 active_jobs,
2253 directory_url: self.cfg.directory_url.clone(),
2254 node_token: Arc::clone(&self.runtime_token),
2255 node_hmac_key: Arc::clone(&self.runtime_hmac_key),
2256 tasks_success: Arc::clone(&tasks_success),
2257 tasks_failed: Arc::clone(&tasks_failed),
2258 tasks_latency_total_ms: Arc::clone(&tasks_latency_total_ms),
2259 max_concurrent: self.cfg.max_concurrent,
2260 availability: Arc::new(crate::availability::AvailabilityEvaluator::new(
2261 self.cfg.availability_windows.clone(),
2262 )),
2263 cip_policy: self
2265 .cfg
2266 .cip_policy
2267 .clone()
2268 .unwrap_or_else(crate::cip_policy::get_cip_policy),
2269 idempotency: Arc::new(crate::idempotency::IdempotencyGuard::default()),
2270 enable_idempotency: self.cfg.enable_idempotency,
2271 relay_bind_ticket_public_key_hex: self.cfg.relay_bind_ticket_public_key_hex.clone(),
2272 relay_require_bind_ticket: self.cfg.relay_require_bind_ticket,
2273 peer_manager: Arc::new(crate::peer_manager::PeerManager::with_opts(
2274 self.cfg.directory_url.clone(),
2275 self.cfg.node_hmac_key.clone(),
2276 crate::peer_manager::PeerManagerOpts {
2277 relay_capable: self.cfg.relay_capable,
2278 relay_accept_port: self.cfg.relay_accept_port,
2279 },
2280 )),
2281 http: self.http.clone(),
2282 nonce_cache,
2283 pinhole_uid: Arc::clone(&shared_pinhole_uid),
2284 pinhole_lease_seconds: Arc::clone(&shared_pinhole_lease),
2285 #[cfg(feature = "iicp-tcp")]
2286 relay_sessions: Arc::new(crate::relay_session::RelaySessionRegistry::new()),
2287 task_rate_limit: std::env::var("IICP_TASK_RATE_LIMIT")
2290 .ok()
2291 .and_then(|v| v.parse().ok())
2292 .unwrap_or(120),
2293 task_rate_buckets: Arc::new(std::sync::Mutex::new(HashMap::new())),
2294 cx_private_key: self.cx_private_key,
2295 backend_stability: Arc::clone(&self.backend_stability),
2296 });
2297
2298 let hb_availability = Arc::clone(&state.availability);
2301 if self.cfg.enable_mesh {
2303 let pm = Arc::clone(&state.peer_manager);
2304 let node_id = self.cfg.node_id.clone();
2305 let own_endpoint = self.cfg.endpoint.clone();
2306 tokio::spawn(async move {
2307 pm.start(&node_id, &own_endpoint).await;
2308 let interval = pm.gossip_interval();
2309 loop {
2310 tokio::time::sleep(interval).await;
2311 pm.gossip_round().await;
2312 }
2313 });
2314 }
2315
2316 let mut app = Router::new()
2321 .route(
2322 "/v1/task",
2323 post(task_endpoint).options(relay_cors_preflight),
2324 )
2325 .route(
2326 "/iicp/health",
2327 get(health_endpoint).options(relay_cors_preflight),
2328 )
2329 .route("/metrics", get(metrics_endpoint));
2330 if self.cfg.enable_mesh {
2331 app = app.route("/v1/peers", post(peers_endpoint));
2332 }
2333 if self.cfg.relay_capable {
2334 app = app.route("/v1/relay", post(relay_endpoint));
2335 #[cfg(feature = "iicp-tcp")]
2339 {
2340 app = app
2341 .route(
2342 "/v1/relay/bind",
2343 post(relay_bind_endpoint).options(relay_cors_preflight),
2344 )
2345 .route(
2346 "/v1/relay/pull",
2347 get(relay_pull_endpoint).options(relay_cors_preflight),
2348 )
2349 .route(
2350 "/v1/relay/result",
2351 post(relay_result_endpoint).options(relay_cors_preflight),
2352 )
2353 .route(
2354 "/v1/relay/unbind",
2355 post(relay_unbind_endpoint).options(relay_cors_preflight),
2356 )
2357 .route(
2358 "/v1/relay-for/:wid/v1/task",
2359 post(relay_for_task_endpoint).options(relay_cors_preflight),
2360 )
2361 .route(
2362 "/v1/relay-for/:wid/iicp/health",
2363 get(relay_for_health_endpoint).options(relay_cors_preflight),
2364 );
2365 }
2366 }
2367 #[cfg(feature = "iicp-tcp")]
2369 let relay_sessions_arc = Arc::clone(&state.relay_sessions);
2370 let app = app
2371 .layer(axum::middleware::map_response(
2372 |resp: Response| async move { relay_cors(resp) },
2373 ))
2374 .with_state(state);
2375
2376 let addr: SocketAddr = addr
2377 .parse()
2378 .map_err(|e| IicpError::Node(format!("invalid addr: {e}")))?;
2379
2380 let listener = if addr.is_ipv6() {
2385 let socket = Socket::new(Domain::IPV6, Type::STREAM, Some(Protocol::TCP))
2386 .map_err(|e| IicpError::Node(format!("socket create: {e}")))?;
2387 socket
2388 .set_only_v6(false)
2389 .map_err(|e| IicpError::Node(format!("set_only_v6: {e}")))?;
2390 socket
2391 .set_reuse_address(true)
2392 .map_err(|e| IicpError::Node(format!("set_reuse_address: {e}")))?;
2393 socket
2394 .bind(&addr.into())
2395 .map_err(|e| IicpError::Node(format!("bind {addr}: {e}")))?;
2396 socket
2397 .listen(1024)
2398 .map_err(|e| IicpError::Node(format!("listen: {e}")))?;
2399 let std_listener: std::net::TcpListener = socket.into();
2400 std_listener
2401 .set_nonblocking(true)
2402 .map_err(|e| IicpError::Node(e.to_string()))?;
2403 TcpListener::from_std(std_listener).map_err(|e| IicpError::Node(e.to_string()))?
2404 } else {
2405 TcpListener::bind(addr)
2406 .await
2407 .map_err(|e| IicpError::Node(e.to_string()))?
2408 };
2409
2410 tracing::info!("IICP node {} listening on {}", self.cfg.node_id, addr);
2411
2412 if let Some(token) = node_token {
2413 let node_id = self.cfg.node_id.clone();
2414 let dir = self.cfg.directory_url.clone();
2415 let http = self.http.clone();
2416 let avail = Arc::clone(&hb_availability);
2417 let max_c = self.cfg.max_concurrent;
2418 let hb_log: Option<Arc<crate::node_log::NodeLog>> =
2420 self.cfg.log_dir.as_deref().and_then(|d| {
2421 crate::node_log::NodeLog::open(d, &node_id)
2422 .map(Arc::new)
2423 .ok()
2424 });
2425 let hb_node_id = node_id.clone();
2426 let hb_tasks_success = Arc::clone(&tasks_success);
2427 let hb_tasks_failed = Arc::clone(&tasks_failed);
2428 let hb_tasks_latency_total_ms = Arc::clone(&tasks_latency_total_ms);
2429 let hb_register_payload = self.build_register_payload();
2433 let hb_token_arc = Arc::clone(&self.runtime_token);
2434 let hb_saved_node_name = self.saved_node_name.clone();
2435 let hb_register_url = format!("{}/v1/register", dir.trim_end_matches('/'));
2436 let hb_backend_url = self.cfg.backend_url.clone();
2438 let hb_backend_api_key = self.cfg.backend_api_key.clone();
2439 let hb_excluded_models = self.cfg.excluded_models.clone();
2440 let hb_backend = self.cfg.backend.clone();
2441 let hb_model = self.cfg.model.clone();
2442 let hb_backend_stability = Arc::clone(&self.backend_stability);
2443 let hb_intent = self.cfg.intent.clone();
2444 let hb_max_tokens = self.cfg.max_tokens;
2445 let hb_registered_models = Arc::clone(&self.registered_models);
2446 let hb_endpoint_override = self.endpoint_override_handle();
2449 let hb_registered_endpoint = Arc::clone(&self.registered_endpoint);
2450 let hb_liveness_challenge = Arc::clone(&self.liveness_challenge);
2451 let hb_runtime_hmac_key = Arc::clone(&self.runtime_hmac_key);
2452 let hb_runtime_available = Arc::clone(&self.runtime_available);
2453 let hb_runtime_relay_bound = Arc::clone(&self.runtime_relay_bound);
2454 let hb_recovery_grace = crate::recovery::env_grace_checks();
2455 let hb_recovery_check_every = crate::recovery::env_check_every_heartbeats();
2456 let hb_recovery_supervised = crate::recovery::supervised_recovery_enabled();
2457 tokio::spawn(async move {
2458 let mut token = token;
2459 let mut seq: u64 = 0;
2460 let mut recovery_failures: u32 = 0;
2461 loop {
2462 tokio::time::sleep(Duration::from_secs(HEARTBEAT_INTERVAL_SECS)).await;
2463 seq += 1;
2464 let ok = hb_tasks_success.swap(0, Ordering::Relaxed);
2468 let fail = hb_tasks_failed.swap(0, Ordering::Relaxed);
2469 let latency_total_ms = hb_tasks_latency_total_ms.swap(0, Ordering::Relaxed);
2470 let metrics = if ok > 0 || fail > 0 {
2471 let mut m = json!({"tasks_success": ok, "tasks_failed": fail});
2472 let total = ok + fail;
2473 if total > 0 && latency_total_ms > 0 {
2474 m["avg_latency_ms"] = json!(
2475 (latency_total_ms as f64 / total as f64 * 100.0).round() / 100.0
2476 );
2477 }
2478 m
2479 } else {
2480 json!({})
2481 };
2482 let live_models = if let Some(ref bu) = hb_backend_url {
2484 probe_health_models_bg(&http, bu, &hb_backend_api_key, &hb_excluded_models)
2485 .await
2486 } else {
2487 None
2488 };
2489 let backend_stability = if let Some(ref bu) = hb_backend_url {
2490 let obs = observe_backend_stability(
2491 &http,
2492 bu,
2493 hb_backend.as_deref(),
2494 hb_model.as_deref(),
2495 hb_backend_api_key.as_deref(),
2496 )
2497 .await;
2498 if let Ok(mut guard) = hb_backend_stability.write() {
2499 if !guard.is_draining() || obs.is_draining() {
2500 *guard = obs.clone();
2501 }
2502 }
2503 Some(obs)
2504 } else {
2505 None
2506 };
2507 let public_available = hb_runtime_available.load(Ordering::Relaxed);
2509 let mut hb_body = json!({
2510 "node_id": &node_id,
2511 "node_token": &token,
2512 "status": if public_available { "available" } else { "recovering" },
2513 "available": public_available,
2517 "max_concurrent": avail.effective_max_concurrent(max_c),
2519 "metrics": metrics,
2522 });
2523 attach_update_status(&mut hb_body);
2524 let stored_challenge = hb_liveness_challenge.read().expect("poisoned").clone();
2529 let hmac_key = hb_runtime_hmac_key.read().expect("poisoned").clone();
2530 if let Some(challenge) = stored_challenge {
2531 if !hmac_key.is_empty() {
2532 hb_body["challenge_response"] =
2533 json!(crate::pricing::sign_body(challenge.as_bytes(), &hmac_key));
2534 }
2535 }
2536 if let Some(ref hm) = live_models {
2537 hb_body["health_models"] = json!(hm);
2538 }
2539 if let Some(ref obs) = backend_stability {
2540 hb_body["backend_stability"] = obs.public_json();
2541 }
2542 match http
2543 .post(format!("{}/v1/heartbeat", dir.trim_end_matches('/')))
2546 .bearer_auth(&token)
2547 .json(&hb_body)
2548 .send()
2549 .await
2550 {
2551 Ok(resp) if resp.status().is_success() => {
2552 if let Some(ref log) = hb_log {
2553 log.write("heartbeat_ok", &hb_node_id, &format!("seq={seq}"));
2554 }
2555 let do_recovery_check =
2562 hb_recovery_check_every > 0 && seq % hb_recovery_check_every == 0;
2563 if let Ok(data) = resp.json::<Value>().await {
2564 if let Some(ch) = data["challenge"].as_str() {
2565 *hb_liveness_challenge.write().expect("poisoned") =
2566 Some(ch.to_string());
2567 }
2568 }
2569 if let Some(live) = live_models {
2571 if !live.is_empty() {
2572 let registered =
2573 hb_registered_models.read().expect("poisoned").clone();
2574 let live_set: std::collections::HashSet<_> =
2575 live.iter().cloned().collect();
2576 let reg_set: std::collections::HashSet<_> =
2577 registered.into_iter().collect();
2578 if live_set != reg_set {
2579 let mut new_payload = hb_register_payload.clone();
2580 let new_caps =
2581 build_capabilities(&live, &hb_intent, hb_max_tokens);
2582 new_payload["capabilities"] =
2583 serde_json::to_value(&new_caps).unwrap_or(json!([]));
2584 if let Some(credentials) =
2585 reregister(&http, &hb_register_url, &new_payload).await
2586 {
2587 *hb_registered_models.write().expect("poisoned") =
2588 live.clone();
2589 token = apply_runtime_credentials(
2590 &hb_token_arc,
2591 &hb_runtime_hmac_key,
2592 hb_saved_node_name.as_deref(),
2593 credentials,
2594 );
2595 tracing::info!(
2596 "seq={seq} model drift: re-registered with {} models",
2597 live.len()
2598 );
2599 }
2600 }
2601 }
2602 }
2603 let override_ep =
2610 hb_endpoint_override.read().expect("poisoned").clone();
2611 if let Some(ep) = override_ep {
2612 let registered_ep =
2613 hb_registered_endpoint.read().expect("poisoned").clone();
2614 if public_available && ep != registered_ep {
2615 let mut new_payload = hb_register_payload.clone();
2616 new_payload["endpoint"] = json!(ep);
2617 new_payload["current_node_token"] = json!(token);
2618 if let Some(credentials) =
2619 reregister(&http, &hb_register_url, &new_payload).await
2620 {
2621 *hb_registered_endpoint.write().expect("poisoned") =
2622 ep.clone();
2623 token = apply_runtime_credentials(
2624 &hb_token_arc,
2625 &hb_runtime_hmac_key,
2626 hb_saved_node_name.as_deref(),
2627 credentials,
2628 );
2629 tracing::info!(
2630 "seq={seq} endpoint rotation: re-registered → {ep}"
2631 );
2632 }
2633 }
2634 }
2635 if do_recovery_check {
2636 let route_status = crate::recovery::registry_route_status(
2637 &http,
2638 &dir,
2639 &node_id,
2640 Duration::from_secs(5),
2641 )
2642 .await;
2643 let presence = route_status.presence;
2644 let public_available_now =
2645 hb_runtime_available.load(Ordering::Relaxed);
2646 let relay_bound_now =
2647 hb_runtime_relay_bound.load(Ordering::Relaxed);
2648 let route_needs_promotion = route_status.route_needs_promotion;
2649 let effective_public_route =
2654 crate::recovery::effective_public_route_available(
2655 public_available_now,
2656 route_needs_promotion,
2657 relay_bound_now,
2658 );
2659 if !effective_public_route
2660 || matches!(
2661 presence,
2662 crate::recovery::DirectoryPresence::Absent
2663 )
2664 {
2665 recovery_failures = recovery_failures.saturating_add(1);
2666 } else if matches!(
2667 presence,
2668 crate::recovery::DirectoryPresence::Present
2669 ) {
2670 recovery_failures = 0;
2671 }
2672
2673 let backend_attention = backend_stability
2674 .as_ref()
2675 .map(|obs| obs.is_draining())
2676 .unwrap_or_else(|| {
2677 hb_backend_stability
2678 .read()
2679 .map(|obs| obs.is_draining())
2680 .unwrap_or(false)
2681 });
2682 let (recovery_state, recovery_action) = crate::recovery::classify(
2683 true,
2684 effective_public_route,
2685 presence,
2686 recovery_failures,
2687 hb_recovery_grace,
2688 backend_attention,
2689 );
2690 if !matches!(
2691 recovery_state,
2692 crate::recovery::RecoveryState::Healthy
2693 | crate::recovery::RecoveryState::Unknown
2694 ) {
2695 tracing::warn!(
2696 "seq={seq} recovery check state={recovery_state:?} action={recovery_action:?} failures={recovery_failures}/{hb_recovery_grace} route_needs_promotion={route_needs_promotion} relay_bound={relay_bound_now}"
2697 );
2698 if let Some(ref log) = hb_log {
2699 log.write(
2700 "recovery_check",
2701 &hb_node_id,
2702 &format!(
2703 "seq={seq} state={recovery_state:?} action={recovery_action:?} failures={recovery_failures}/{hb_recovery_grace} route_needs_promotion={route_needs_promotion} relay_bound={relay_bound_now}"
2704 ),
2705 );
2706 }
2707 }
2708 match recovery_action {
2709 crate::recovery::RecoveryAction::Reregister => {
2710 let override_ep =
2711 hb_endpoint_override.read().expect("poisoned").clone();
2712 let mut new_payload = hb_register_payload.clone();
2713 if let Some(ep) = override_ep.clone() {
2714 new_payload["endpoint"] = json!(ep);
2715 }
2716 new_payload["current_node_token"] = json!(token);
2717 match reregister(&http, &hb_register_url, &new_payload)
2718 .await
2719 {
2720 Some(credentials) => {
2721 if let Some(ep) = new_payload["endpoint"].as_str() {
2722 *hb_registered_endpoint
2723 .write()
2724 .expect("poisoned") = ep.to_string();
2725 }
2726 token = apply_runtime_credentials(
2727 &hb_token_arc,
2728 &hb_runtime_hmac_key,
2729 hb_saved_node_name.as_deref(),
2730 credentials,
2731 );
2732 recovery_failures = 0;
2733 if let Some(ref log) = hb_log {
2734 log.write(
2735 "reregister_ok",
2736 &hb_node_id,
2737 &format!("seq={seq} recovery_check"),
2738 );
2739 }
2740 tracing::info!(
2741 "seq={seq} recovery check: re-registered public directory entry"
2742 );
2743 }
2744 None => {
2745 tracing::warn!(
2746 "seq={seq} recovery check: re-registration failed"
2747 );
2748 if let Some(ref log) = hb_log {
2749 log.write(
2750 "reregister_fail",
2751 &hb_node_id,
2752 &format!("seq={seq} recovery_check"),
2753 );
2754 }
2755 }
2756 }
2757 }
2758 crate::recovery::RecoveryAction::RestartSelf
2759 if hb_recovery_supervised =>
2760 {
2761 eprintln!(
2762 "[iicp-node] recovery check recommends restart \
2763 (state={recovery_state:?}, failures={recovery_failures}/{hb_recovery_grace}); \
2764 exiting with code {} so the supervisor can rebuild route + registration.",
2765 crate::recovery::RECOVERY_EXIT_CODE
2766 );
2767 std::process::exit(crate::recovery::RECOVERY_EXIT_CODE);
2768 }
2769 _ => {}
2770 }
2771 }
2772 }
2773 Ok(resp) if matches!(resp.status().as_u16(), 401 | 404 | 410) => {
2779 let code = resp.status().as_u16();
2780 tracing::warn!(
2781 "heartbeat rejected ({code}) — node unknown to directory; re-registering"
2782 );
2783 match reregister(&http, &hb_register_url, &hb_register_payload).await {
2784 Some(credentials) => {
2785 token = apply_runtime_credentials(
2786 &hb_token_arc,
2787 &hb_runtime_hmac_key,
2788 hb_saved_node_name.as_deref(),
2789 credentials,
2790 );
2791 if let Some(ref log) = hb_log {
2792 log.write(
2793 "reregister_ok",
2794 &hb_node_id,
2795 &format!("seq={seq} after_status={code}"),
2796 );
2797 }
2798 }
2799 None => {
2800 tracing::warn!("re-registration failed (after status {code})");
2801 if let Some(ref log) = hb_log {
2802 log.write(
2803 "reregister_fail",
2804 &hb_node_id,
2805 &format!("seq={seq} after_status={code}"),
2806 );
2807 }
2808 }
2809 }
2810 }
2811 Ok(resp) => {
2812 if let Some(ref log) = hb_log {
2813 log.write(
2814 "heartbeat_fail",
2815 &hb_node_id,
2816 &format!("seq={seq} status={}", resp.status().as_u16()),
2817 );
2818 }
2819 }
2820 Err(e) => {
2821 tracing::warn!("heartbeat failed: {e}");
2822 if let Some(ref log) = hb_log {
2823 log.write(
2824 "heartbeat_fail",
2825 &hb_node_id,
2826 &format!("seq={seq} error={e}"),
2827 );
2828 }
2829 }
2830 }
2831 }
2832 });
2833 }
2834
2835 #[cfg(feature = "nat")]
2837 {
2838 let uid_arc = Arc::clone(&shared_pinhole_uid);
2839 let lease_arc = Arc::clone(&shared_pinhole_lease);
2840 tokio::spawn(async move {
2841 loop {
2842 let (_uid, lease) = {
2843 let u = uid_arc.read().ok().and_then(|g| *g);
2844 let l = lease_arc.read().map(|g| *g).unwrap_or(3600);
2845 (u, l)
2846 };
2847 let delay = Duration::from_secs(u64::from((lease / 2).max(60)));
2848 tokio::time::sleep(delay).await;
2849 let uid = match uid_arc.read().ok().and_then(|g| *g) {
2850 Some(u) => u,
2851 None => return,
2852 };
2853 let ok = crate::nat_detection::renew_ipv6_pinhole(uid, lease).await;
2854 if ok {
2855 tracing::debug!("UPnP IPv6 pinhole uid={uid} renewed (lease={lease}s)");
2856 } else {
2857 tracing::warn!("UPnP IPv6 pinhole uid={uid} renewal failed — will retry");
2858 }
2859 }
2860 });
2861 }
2862
2863 #[cfg(feature = "iicp-tcp")]
2865 if self.cfg.relay_capable {
2866 let relay_reg = relay_sessions_arc;
2867 let relay_host_str = bind_host.clone();
2868 let relay_port = self.cfg.relay_accept_port;
2869 let relay_http_port = addr.port();
2870 tokio::spawn(async move {
2871 let srv = Arc::new(crate::relay_session::RelayAcceptServer::with_http_port(
2872 (*relay_reg).clone(),
2873 relay_host_str,
2874 relay_port,
2875 relay_http_port,
2876 ));
2877 if let Err(e) = srv.serve().await {
2878 tracing::warn!("Relay accept server error: {e}");
2879 }
2880 });
2881 }
2882
2883 #[cfg(feature = "iicp-tcp")]
2885 if let Some(ref ep) = self.cfg.relay_worker_endpoint {
2886 let ep = ep.clone();
2887 let node_id = self.cfg.node_id.clone();
2888 let intent = self.cfg.intent.clone();
2889 let models = self.cfg.model.clone().map(|m| vec![m]).unwrap_or_default();
2890 let relay_register_payload = self.build_register_payload();
2891 let relay_register_url = format!(
2892 "{}/v1/register",
2893 self.cfg.directory_url.trim_end_matches('/')
2894 );
2895 let relay_token_arc = Arc::clone(&self.runtime_token);
2896 let relay_hmac_arc = Arc::clone(&self.runtime_hmac_key);
2897 let relay_saved_node_name = self.saved_node_name.clone();
2898 let relay_registered_endpoint = Arc::clone(&self.registered_endpoint);
2899 let relay_bound_arc = Arc::clone(&self.runtime_relay_bound);
2900 let relay_cx_private_key = self.cx_private_key;
2901 let handler_fn: crate::relay_worker_client::RelayHandlerFn = Arc::new(
2902 move |task: Value| {
2903 let h = Arc::clone(&handler_for_relay);
2904 let cx_private_key = relay_cx_private_key;
2905 Box::pin(async move {
2906 let mut req = crate::node::TaskRequest {
2907 task_id: task
2908 .get("task_id")
2909 .and_then(|v| v.as_str())
2910 .unwrap_or("")
2911 .to_string(),
2912 intent: task
2913 .get("intent")
2914 .and_then(|v| v.as_str())
2915 .unwrap_or("")
2916 .to_string(),
2917 payload: task.get("payload").cloned().unwrap_or(Value::Null),
2918 iicp_conf: task
2919 .get("iicp_conf")
2920 .and_then(|v| serde_json::from_value(v.clone()).ok()),
2921 cx_response_encryption: task
2922 .get("cx_response_encryption")
2923 .and_then(|value| value.as_str())
2924 .map(str::to_string),
2925 constraints: task.get("constraints").cloned(),
2926 auth: task.get("auth").cloned(),
2927 nonce: None,
2928 source_node_id: task
2929 .get("source_node_id")
2930 .and_then(|v| v.as_str())
2931 .map(|s| s.to_string()),
2932 _trace: None,
2933 };
2934 let mut shared_secret = None;
2935 if req.payload.is_null() {
2936 if let Some(conf) = req.iicp_conf.as_ref() {
2937 match cx_private_key
2938 .as_ref()
2939 .ok_or_else(|| {
2940 IicpError::Node("node has no CX private key".to_string())
2941 })
2942 .and_then(|key| {
2943 crate::confidentiality::decrypt_payload_with_context(
2944 conf, key,
2945 )
2946 }) {
2947 Ok((payload, secret)) => {
2948 req.payload = payload;
2949 shared_secret = Some(secret);
2950 }
2951 Err(error) => {
2952 return json!({
2953 "error": {
2954 "code": "IICP-CX-02",
2955 "message": format!("iicp_conf decrypt failed: {error}"),
2956 }
2957 })
2958 }
2959 }
2960 }
2961 }
2962 let response_required =
2963 req.cx_response_encryption.as_deref() == Some("required");
2964 let task_id = req.task_id.clone();
2965 let result = h(req)
2966 .await
2967 .unwrap_or_else(|e| json!({"error": e.to_string()}));
2968 if response_required {
2969 let Some(secret) = shared_secret else {
2970 return json!({
2971 "error": {
2972 "code": "IICP-CX-03",
2973 "message": "encrypted response requested without an encrypted request",
2974 }
2975 });
2976 };
2977 let plain_response = json!({
2978 "task_id": task_id.clone(),
2979 "status": "success",
2980 "result": result,
2981 "generated_by_ai": true,
2982 });
2983 match crate::confidentiality::encrypt_response(
2984 &plain_response,
2985 &secret,
2986 &task_id,
2987 ) {
2988 Ok(envelope) => json!({
2989 "task_id": task_id,
2990 "status": "encrypted",
2991 "iicp_conf_resp": envelope,
2992 }),
2993 Err(error) => json!({
2994 "error": {
2995 "code": "IICP-CX-04",
2996 "message": format!("response encryption failed: {error}"),
2997 }
2998 }),
2999 }
3000 } else {
3001 result
3002 }
3003 })
3004 },
3005 );
3006 let (rhost, rport) = {
3007 if let Some(pos) = ep.rfind(':') {
3008 let port = ep[pos + 1..].parse::<u16>().unwrap_or(9485);
3009 (ep[..pos].to_string(), port)
3010 } else {
3011 (ep.clone(), 9485u16)
3012 }
3013 };
3014 let http_client = self.http.clone();
3017 let on_bind_cb: crate::relay_worker_client::OnBindFn =
3018 Arc::new(move |rh: String, rp: u16, wid: String| {
3019 let http = http_client.clone();
3020 let url = relay_register_url.clone();
3021 let mut payload = relay_register_payload.clone();
3022 let token_arc = Arc::clone(&relay_token_arc);
3023 let hmac_arc = Arc::clone(&relay_hmac_arc);
3024 let saved_node_name = relay_saved_node_name.clone();
3025 let registered_endpoint = Arc::clone(&relay_registered_endpoint);
3026 let relay_bound = Arc::clone(&relay_bound_arc);
3027 Box::pin(async move {
3028 relay_bound.store(true, Ordering::Relaxed);
3032 let relay_endpoint = format!("http://{rh}:{rp}/v1/relay-for/{wid}");
3036 payload["endpoint"] = json!(relay_endpoint);
3037 payload["transport_method"] = json!("turn_relay");
3038 payload["transport_metadata"] = json!({
3039 "relay_for": wid,
3040 "relay_host": rh,
3041 "relay_port": rp,
3042 });
3043 let current_token = token_arc.read().expect("poisoned").clone();
3044 if !current_token.is_empty() {
3045 payload["current_node_token"] = json!(current_token);
3046 }
3047 match reregister(&http, &url, &payload).await {
3048 Some(credentials) => {
3049 if let Some(ep) = payload["endpoint"].as_str() {
3050 *registered_endpoint.write().expect("poisoned") =
3051 ep.to_string();
3052 }
3053 apply_runtime_credentials(
3054 &token_arc,
3055 &hmac_arc,
3056 saved_node_name.as_deref(),
3057 credentials,
3058 );
3059 tracing::info!("Relay worker bound — re-registered relay endpoint");
3060 }
3061 None => {
3062 tracing::warn!(
3063 "Relay worker bound but directory re-registration failed"
3064 );
3065 }
3066 }
3067 })
3068 });
3069 let relay_bound_on_disconnect = Arc::clone(&self.runtime_relay_bound);
3070 let on_disconnect_cb: crate::relay_worker_client::OnDisconnectFn =
3071 Arc::new(move || {
3072 let relay_bound = Arc::clone(&relay_bound_on_disconnect);
3073 Box::pin(async move {
3074 relay_bound.store(false, Ordering::Relaxed);
3075 tracing::warn!("Relay worker session ended — route is no longer bound");
3076 })
3077 });
3078 tokio::spawn(async move {
3079 let rwc = Arc::new(
3080 crate::relay_worker_client::RelayWorkerClient::new(
3081 node_id, intent, rhost, rport, handler_fn, models,
3082 )
3083 .with_on_bind(on_bind_cb)
3084 .with_on_disconnect(on_disconnect_cb),
3085 );
3086 rwc.run().await;
3087 });
3088 }
3089
3090 #[cfg(feature = "iicp-tcp")]
3099 {
3100 let native = crate::iicp_tcp::IicpTcpServer::new(&bind_host, addr.port())
3101 .with_node_id(self.cfg.node_id.clone())
3102 .with_handler(Arc::new(move |t: crate::iicp_tcp::TcpTask| {
3103 let h = Arc::clone(&handler_for_native);
3104 Box::pin(async move {
3105 let req = TaskRequest {
3106 task_id: t.task_id,
3107 intent: t.intent,
3108 payload: t.payload,
3109 iicp_conf: None,
3110 cx_response_encryption: None,
3111 constraints: None,
3112 auth: None,
3113 nonce: None,
3114 source_node_id: None,
3115 _trace: None,
3116 };
3117 h(req)
3118 .await
3119 .unwrap_or_else(|e| json!({"error": e.to_string()}))
3120 })
3121 as std::pin::Pin<Box<dyn std::future::Future<Output = Value> + Send>>
3122 }));
3123
3124 let internal = TcpListener::bind("127.0.0.1:0")
3125 .await
3126 .map_err(|e| IicpError::Node(e.to_string()))?;
3127 let internal_addr = internal
3128 .local_addr()
3129 .map_err(|e| IicpError::Node(e.to_string()))?;
3130 tokio::spawn(async move {
3131 let _ = axum::serve(internal, app).await;
3132 });
3133
3134 loop {
3135 let (stream, _peer) = match listener.accept().await {
3136 Ok(s) => s,
3137 Err(_) => continue,
3138 };
3139 let native = native.clone();
3140 tokio::spawn(async move {
3141 let mut buf = [0u8; 4];
3142 let mut got = 0usize;
3143 for _ in 0..20 {
3147 match stream.peek(&mut buf).await {
3148 Ok(n) => {
3149 got = n;
3150 if n >= 4 {
3151 break;
3152 }
3153 }
3154 Err(_) => return,
3155 }
3156 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
3157 }
3158 if got >= 4 && &buf == crate::iicp_tcp::IICP_MAGIC {
3159 let _ = native.handle_connection(stream).await;
3160 } else if let Ok(mut inner) =
3161 tokio::net::TcpStream::connect(internal_addr).await
3162 {
3163 let mut stream = stream;
3164 let _ = tokio::io::copy_bidirectional(&mut stream, &mut inner).await;
3165 }
3166 });
3167 }
3168 }
3169
3170 #[cfg(not(feature = "iicp-tcp"))]
3171 {
3172 axum::serve(listener, app)
3173 .await
3174 .map_err(|e| IicpError::Node(e.to_string()))
3175 }
3176 }
3177}
3178
3179#[cfg(test)]
3180mod task_rate_tests {
3181 use super::task_rate_step;
3182 use std::collections::HashMap;
3183 use std::time::Instant;
3184
3185 #[test]
3186 fn allows_under_limit_then_blocks() {
3187 let mut b = HashMap::new();
3188 let now = Instant::now();
3189 assert!(task_rate_step(&mut b, 3, "o-a", now));
3190 assert!(task_rate_step(&mut b, 3, "o-a", now));
3191 assert!(task_rate_step(&mut b, 3, "o-a", now));
3192 assert!(!task_rate_step(&mut b, 3, "o-a", now)); }
3194
3195 #[test]
3196 fn origins_are_independent() {
3197 let mut b = HashMap::new();
3198 let now = Instant::now();
3199 assert!(task_rate_step(&mut b, 1, "o-a", now));
3200 assert!(task_rate_step(&mut b, 1, "o-b", now)); assert!(!task_rate_step(&mut b, 1, "o-a", now)); }
3203
3204 #[test]
3205 fn window_resets() {
3206 let mut b = HashMap::new();
3207 let now = Instant::now();
3208 assert!(task_rate_step(&mut b, 1, "k", now));
3209 assert!(!task_rate_step(&mut b, 1, "k", now));
3210 let past = now
3212 .checked_sub(super::TASK_RATE_WINDOW + std::time::Duration::from_secs(1))
3213 .unwrap();
3214 b.insert("k".to_string(), (past, 1));
3215 assert!(task_rate_step(&mut b, 1, "k", now));
3216 }
3217}
3218
3219#[cfg(test)]
3220mod capability_tests {
3221 use super::{build_capabilities, filter_excluded_models};
3222
3223 const CHAT: &str = "urn:iicp:intent:llm:chat:v1";
3224 const EMBED: &str = "urn:iicp:intent:llm:embedding:v1";
3225
3226 #[test]
3227 fn excluded_backend_alias_is_not_used_as_health_or_registration_evidence() {
3228 let models = vec!["stable-model".to_string(), "mesh".to_string()];
3229 assert_eq!(
3230 filter_excluded_models(models, &["mesh".to_string()]),
3231 vec!["stable-model".to_string()]
3232 );
3233 }
3234
3235 #[test]
3238 fn chat_plus_embedding_models_advertise_two_intents() {
3239 let models = vec![
3240 "qwen2.5-coder-14b-instruct".to_string(),
3241 "text-embedding-nomic-embed-text-v1.5".to_string(),
3242 ];
3243 let caps = build_capabilities(&models, CHAT, 4096);
3244 assert_eq!(caps.len(), 2, "should advertise chat + embedding");
3245 assert_eq!(caps[0]["intent"], CHAT);
3247 assert_eq!(
3248 caps[0]["models"],
3249 serde_json::json!(["qwen2.5-coder-14b-instruct"])
3250 );
3251 assert_eq!(caps[1]["intent"], EMBED);
3252 assert_eq!(
3253 caps[1]["models"],
3254 serde_json::json!(["text-embedding-nomic-embed-text-v1.5"])
3255 );
3256 }
3257
3258 #[test]
3260 fn chat_only_yields_single_capability() {
3261 let caps = build_capabilities(&["qwen2.5:0.5b".to_string()], CHAT, 4096);
3262 assert_eq!(caps.len(), 1);
3263 assert_eq!(caps[0]["intent"], CHAT);
3264 assert_eq!(caps[0]["models"], serde_json::json!(["qwen2.5:0.5b"]));
3265 assert_eq!(caps[0]["input_modalities"], serde_json::json!(["text"]));
3266 }
3267
3268 #[test]
3271 fn vision_model_advertises_image_modality_chat_capability() {
3272 let models = vec![
3273 "qwen2.5-coder-14b".to_string(),
3274 "qwen/qwen3-vl-8b".to_string(),
3275 ];
3276 let caps = build_capabilities(&models, CHAT, 4096);
3277 assert_eq!(
3278 caps.len(),
3279 2,
3280 "text-chat and vision-chat are distinct capabilities"
3281 );
3282 assert_eq!(caps[0]["intent"], CHAT);
3283 assert_eq!(caps[0]["input_modalities"], serde_json::json!(["text"]));
3284 assert_eq!(caps[0]["models"], serde_json::json!(["qwen2.5-coder-14b"]));
3285 assert_eq!(caps[1]["intent"], CHAT);
3286 assert_eq!(
3287 caps[1]["input_modalities"],
3288 serde_json::json!(["text", "image"])
3289 );
3290 assert_eq!(caps[1]["models"], serde_json::json!(["qwen/qwen3-vl-8b"]));
3291 }
3292
3293 #[test]
3296 fn audio_model_advertises_audio_modality_chat_capability() {
3297 let models = vec!["qwen2.5:0.5b".to_string(), "qwen2-audio-7b".to_string()];
3298 let caps = build_capabilities(&models, CHAT, 4096);
3299 assert_eq!(caps.len(), 2);
3300 assert_eq!(caps[0]["input_modalities"], serde_json::json!(["text"]));
3301 assert_eq!(caps[1]["intent"], CHAT);
3302 assert_eq!(
3303 caps[1]["input_modalities"],
3304 serde_json::json!(["text", "audio"])
3305 );
3306 assert_eq!(caps[1]["models"], serde_json::json!(["qwen2-audio-7b"]));
3307 }
3308
3309 #[test]
3311 fn omni_model_advertises_image_and_audio_modalities() {
3312 let caps = build_capabilities(&["qwen2.5-omni-7b".to_string()], CHAT, 4096);
3313 assert_eq!(caps.len(), 1);
3314 assert_eq!(
3315 caps[0]["input_modalities"],
3316 serde_json::json!(["text", "image", "audio"])
3317 );
3318 }
3319
3320 #[test]
3322 fn empty_models_yields_default_intent_capability() {
3323 let caps = build_capabilities(&[], CHAT, 1024);
3324 assert_eq!(caps.len(), 1);
3325 assert_eq!(caps[0]["intent"], CHAT);
3326 assert_eq!(caps[0]["models"], serde_json::json!([]));
3327 }
3328}
3329
3330#[cfg(test)]
3331mod reregister_tests {
3332 use super::{probe_health_models_bg, reregister, update_saved_identity_credentials};
3333 use serde_json::json;
3334
3335 #[tokio::test]
3338 async fn reregister_returns_fresh_token() {
3339 let mut server = mockito::Server::new_async().await;
3340 let m = server
3341 .mock("POST", "/v1/register")
3342 .with_status(201)
3343 .with_body(json!({"node_token": "recovered-xyz"}).to_string())
3344 .create_async()
3345 .await;
3346 let http = reqwest::Client::new();
3347 let payload = json!({"endpoint": "https://x", "region": "r"});
3348 let url = format!("{}/v1/register", server.url());
3349 let tok = reregister(&http, &url, &payload).await;
3350 assert_eq!(
3351 tok,
3352 Some(super::RegistrationCredentials {
3353 node_token: "recovered-xyz".to_string(),
3354 node_hmac_key: None,
3355 })
3356 );
3357 m.assert_async().await;
3358 }
3359
3360 #[test]
3361 fn update_saved_identity_credentials_updates_token_and_hmac() {
3362 let mut node = crate::identity::NodeIdentity {
3363 node_id: "drift-rs-1".to_string(),
3364 operator_id: "op-test".to_string(),
3365 name: "drift".to_string(),
3366 backend_url: "http://mock-backend".to_string(),
3367 backend_type: "openai_compat".to_string(),
3368 model: "phi3:mini".to_string(),
3369 intent: "urn:iicp:intent:llm:chat:v1".to_string(),
3370 region: "eu-central".to_string(),
3371 directory_url: "http://mock-dir".to_string(),
3372 max_concurrent: 4,
3373 port: 9484,
3374 host: "0.0.0.0".to_string(),
3375 public_endpoint: "http://node.local:8080".to_string(),
3376 auto_detect_nat: false,
3377 external_ip_probe_url: String::new(),
3378 supported_receipt_profiles: Vec::new(),
3379 node_token: Some("old-token".to_string()),
3380 node_hmac_key: Some("old-hmac".to_string()),
3381 created_at: "2026-07-09T00:00:00Z".to_string(),
3382 };
3383
3384 update_saved_identity_credentials(&mut node, "new-token", Some("new-hmac"));
3385
3386 assert_eq!(node.node_token.as_deref(), Some("new-token"));
3387 assert_eq!(node.node_hmac_key.as_deref(), Some("new-hmac"));
3388 }
3389
3390 #[tokio::test]
3391 async fn reregister_none_on_failure() {
3392 let mut server = mockito::Server::new_async().await;
3393 let _m = server
3394 .mock("POST", "/v1/register")
3395 .with_status(500)
3396 .create_async()
3397 .await;
3398 let http = reqwest::Client::new();
3399 let url = format!("{}/v1/register", server.url());
3400 let tok = reregister(&http, &url, &json!({})).await;
3401 assert_eq!(tok, None);
3402 }
3403
3404 #[tokio::test]
3405 async fn health_inventory_excludes_meshllm_preview_alias_before_drift_logic() {
3406 let mut server = mockito::Server::new_async().await;
3407 let _models = server
3408 .mock("GET", "/v1/models")
3409 .with_status(200)
3410 .with_body(json!({"data": [{"id": "stable-model"}, {"id": "mesh"}]}).to_string())
3411 .create_async()
3412 .await;
3413
3414 let observed = probe_health_models_bg(
3415 &reqwest::Client::new(),
3416 &format!("{}/v1", server.url()),
3417 &None,
3418 &["mesh".to_string()],
3419 )
3420 .await;
3421 assert_eq!(observed, Some(vec!["stable-model".to_string()]));
3422 }
3423}
3424
3425#[cfg(test)]
3426mod operator_wiring_tests {
3427 use super::{IicpNode, NodeConfig};
3431 use crate::delegation::issue_delegation;
3432 use crate::identity::OperatorIdentity;
3433
3434 const CHAT: &str = "urn:iicp:intent:llm:chat:v1";
3435
3436 #[test]
3437 fn register_payload_carries_operator_fields_never_secret() {
3438 let op = OperatorIdentity::generate("Rebel One", "me@example.com");
3439 let node_id = "test-node-1";
3440 let token = issue_delegation(&op.signing_key().unwrap(), node_id, 3600);
3441 let mut cfg = NodeConfig::new(node_id, "http://h.test:9484", CHAT);
3442 cfg.operator_delegation = serde_json::to_value(&token).ok();
3443 cfg.operator_display_name = Some(op.display_name.clone());
3444 cfg.operator_created_at = Some(op.created_at.clone());
3445 cfg.operator_integrity_hash = Some(op.operator_integrity_hash.clone());
3446 let p = IicpNode::new(cfg).build_register_payload();
3447 assert_eq!(
3449 p["operator_delegation"]["operator_pub"],
3450 serde_json::json!(op.operator_id)
3451 );
3452 assert_eq!(p["operator_display_name"], serde_json::json!("Rebel One"));
3453 assert_eq!(
3454 p["operator_integrity_hash"],
3455 serde_json::json!(op.operator_integrity_hash)
3456 );
3457 let raw = p.to_string();
3458 assert!(
3459 !raw.contains(&op.operator_secret),
3460 "secret key must never be sent"
3461 );
3462 assert!(
3463 !raw.contains("me@example.com"),
3464 "contact/email must never be sent"
3465 );
3466 assert!(!raw.contains("operator_secret"));
3467 assert!(!raw.contains("contact"));
3468 }
3469
3470 #[test]
3471 fn register_payload_omits_operator_fields_when_unbound() {
3472 let p = IicpNode::new(NodeConfig::new("n2", "http://h.test:9484", CHAT))
3473 .build_register_payload();
3474 assert!(p.get("operator_delegation").is_none());
3475 assert!(p.get("operator_display_name").is_none());
3476 }
3477
3478 #[test]
3479 fn endpoint_override_updates_register_payload_endpoint() {
3480 let node = IicpNode::new(NodeConfig::new("n3", "https://seed.example.com", CHAT));
3481 assert_eq!(
3482 node.build_register_payload()["endpoint"],
3483 serde_json::json!("https://seed.example.com")
3484 );
3485 node.set_endpoint("https://rotated.example.net".to_string());
3486 assert_eq!(
3487 node.build_register_payload()["endpoint"],
3488 serde_json::json!("https://rotated.example.net")
3489 );
3490 node.set_endpoint(String::new());
3491 assert_eq!(
3492 node.build_register_payload()["endpoint"],
3493 serde_json::json!("https://seed.example.com")
3494 );
3495 }
3496
3497 #[test]
3501 fn token_extraction_uses_direct_usage_path() {
3502 let handler_value = serde_json::json!({
3503 "choices": [{"message": {"role": "assistant", "content": "hi"}}],
3504 "model": "qwen2.5:0.5b",
3505 "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}
3506 });
3507 let tokens: u64 = handler_value
3508 .get("usage")
3509 .and_then(|u| u.get("total_tokens"))
3510 .and_then(|t| t.as_u64())
3511 .unwrap_or(0);
3512 assert_eq!(
3513 tokens, 15,
3514 "must extract total_tokens from top-level usage key"
3515 );
3516
3517 let wrong: u64 = handler_value
3519 .get("result")
3520 .and_then(|r| r.get("usage"))
3521 .and_then(|u| u.get("total_tokens"))
3522 .and_then(|t| t.as_u64())
3523 .unwrap_or(0);
3524 assert_eq!(
3525 wrong, 0,
3526 "nested result.usage path must not exist in handler value"
3527 );
3528 }
3529
3530 #[tokio::test]
3534 async fn cip_receipt_signature_verifies() {
3535 use super::{canonical_json_node, post_cip_receipt};
3536 use hmac::{Hmac, Mac};
3537 use mockito::Server;
3538 use sha2::{Digest, Sha256};
3539 type HmacSha256 = Hmac<Sha256>;
3540
3541 let mut server = Server::new_async().await;
3542 let hmac_key = "test-hmac-key-1234567890abcdef";
3543 let task_id = "task-receipt-test-001";
3544 let node_id = "node-receipt-test";
3545 let tokens_used = 75u64;
3546
3547 let m = server
3548 .mock("POST", "/api/v1/credits/award")
3549 .with_status(200)
3550 .with_body("{}")
3551 .create_async()
3552 .await;
3553
3554 let result = serde_json::json!({"content": "hello world"});
3555 post_cip_receipt(
3556 reqwest::Client::new(),
3557 format!("{}/api", server.url()),
3558 "test-token".to_string(),
3559 hmac_key.to_string(),
3560 node_id.to_string(),
3561 task_id.to_string(),
3562 tokens_used,
3563 result.clone(),
3564 None, )
3566 .await;
3567
3568 m.assert_async().await;
3569
3570 let test_msg = format!("{task_id}:{tokens_used}:::fixed-nonce:fixed-hash");
3577 let mut mac = HmacSha256::new_from_slice(hmac_key.as_bytes()).unwrap();
3578 mac.update(test_msg.as_bytes());
3579 let expected = mac.finalize().into_bytes();
3580 assert_eq!(expected.len(), 32, "HMAC-SHA256 output must be 32 bytes");
3581
3582 let result_bytes = canonical_json_node(&result).into_bytes();
3584 let hash = hex::encode(Sha256::digest(&result_bytes));
3585 assert_eq!(
3586 hash.len(),
3587 64,
3588 "response_hash must be a 64-char hex SHA-256"
3589 );
3590 assert!(!hash.chars().any(|c| !c.is_ascii_hexdigit()), "must be hex");
3591 }
3592
3593 #[tokio::test]
3596 async fn cip_receipt_forwards_querying_node_id() {
3597 use super::post_cip_receipt;
3598 use mockito::{Matcher, Server};
3599
3600 let mut server = Server::new_async().await;
3601
3602 let m = server
3603 .mock("POST", "/api/v1/credits/award")
3604 .match_body(Matcher::PartialJson(serde_json::json!({
3605 "querying_node_id": "querying-node-abc"
3606 })))
3607 .with_status(200)
3608 .with_body("{}")
3609 .create_async()
3610 .await;
3611
3612 post_cip_receipt(
3613 reqwest::Client::new(),
3614 format!("{}/api", server.url()),
3615 "tok".to_string(),
3616 "key".to_string(),
3617 "serving-node".to_string(),
3618 "task-qni".to_string(),
3619 10u64,
3620 serde_json::json!({"content": "hi"}),
3621 Some("querying-node-abc".to_string()),
3622 )
3623 .await;
3624
3625 m.assert_async().await;
3626 }
3627
3628 #[tokio::test]
3630 async fn cip_receipt_omits_querying_node_id_when_none() {
3631 use super::post_cip_receipt;
3632 use mockito::Server;
3633
3634 let mut server = Server::new_async().await;
3635
3636 let m = server
3641 .mock("POST", "/api/v1/credits/award")
3642 .with_status(200)
3643 .with_body("{}")
3644 .create_async()
3645 .await;
3646
3647 post_cip_receipt(
3648 reqwest::Client::new(),
3649 format!("{}/api", server.url()),
3650 "tok".to_string(),
3651 "key".to_string(),
3652 "serving-node".to_string(),
3653 "task-no-qni".to_string(),
3654 10u64,
3655 serde_json::json!({"content": "hi"}),
3656 None,
3657 )
3658 .await;
3659
3660 m.assert_async().await; }
3662
3663 #[test]
3666 fn cip_receipt_canonical_includes_querying_node_id() {
3667 use hmac::{Hmac, Mac};
3668 use sha2::Sha256;
3669 type HmacSha256 = Hmac<Sha256>;
3670
3671 let hmac_key = "test-hmac-490";
3672 let task_id = "task-490";
3673 let tokens: u64 = 50;
3674 let nonce = "abc123";
3675 let response_hash = "deadbeef".repeat(8); let querying_node_id = "querying-node-xyz";
3677
3678 let extended = format!("{task_id}:{tokens}:::{nonce}:{response_hash}:{querying_node_id}");
3680 let mut mac = HmacSha256::new_from_slice(hmac_key.as_bytes()).unwrap();
3681 mac.update(extended.as_bytes());
3682 let expected = hex::encode(mac.finalize().into_bytes());
3683
3684 let short = format!("{task_id}:{tokens}:::{nonce}:{response_hash}");
3686 let mut mac2 = HmacSha256::new_from_slice(hmac_key.as_bytes()).unwrap();
3687 mac2.update(short.as_bytes());
3688 let short_sig = hex::encode(mac2.finalize().into_bytes());
3689
3690 assert_ne!(
3692 expected, short_sig,
3693 "extended canonical must produce a different signature than short canonical"
3694 );
3695
3696 assert_eq!(
3700 extended,
3701 format!("{task_id}:{tokens}:::{nonce}:{response_hash}:{querying_node_id}"),
3702 "canonical must include querying_node_id at the end"
3703 );
3704 }
3705}