1use std::collections::HashMap;
13use std::path::Path;
14use std::sync::Arc;
15use std::time::{SystemTime, UNIX_EPOCH};
16
17use crate::backend::storage::error::StoreError;
18use crate::backend::storage::store::{SecretRef, Store};
19use crate::backend::wps::{Broker, WaveEvent};
20
21pub mod oauth_status {
31 pub const VALID: &str = "valid";
33 pub const EXPIRED: &str = "expired";
35 pub const NEEDS_REAUTH: &str = "needs_reauth";
37 pub const UNKNOWN: &str = "unknown";
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
48pub enum OAuthProbeStatus {
49 Valid,
50 Expired,
51 NeedsReauth,
52}
53
54impl OAuthProbeStatus {
55 pub fn as_str(self) -> &'static str {
56 match self {
57 Self::Valid => oauth_status::VALID,
58 Self::Expired => oauth_status::EXPIRED,
59 Self::NeedsReauth => oauth_status::NEEDS_REAUTH,
60 }
61 }
62}
63
64pub fn probe_oauth_status(
86 provider: &str,
87 dir: &str,
88 now_ms: i64,
89) -> Option<OAuthProbeStatus> {
90 let probe_path: std::path::PathBuf = match provider {
91 "claude" | "codex" | "openclaw" => Path::new(dir).join(".credentials.json"),
101 _ => return None,
102 };
103
104 let contents = match std::fs::read_to_string(&probe_path) {
105 Ok(s) => s,
106 Err(e) => {
107 tracing::debug!(
108 target: "identity",
109 provider,
110 path = %probe_path.display(),
111 error = %e,
112 "oauth probe: token file unreadable — status=needs_reauth"
113 );
114 return Some(OAuthProbeStatus::NeedsReauth);
115 }
116 };
117 let json: serde_json::Value = match serde_json::from_str(&contents) {
118 Ok(v) => v,
119 Err(e) => {
120 tracing::debug!(
121 target: "identity",
122 provider,
123 path = %probe_path.display(),
124 error = %e,
125 "oauth probe: token file parse failed — status=needs_reauth"
126 );
127 return Some(OAuthProbeStatus::NeedsReauth);
128 }
129 };
130
131 let expires_at_ms = json
136 .get("claudeAiOauth")
137 .and_then(|o| o.get("expiresAt"))
138 .and_then(|v| v.as_i64())
139 .or_else(|| json.get("expiresAt").and_then(|v| v.as_i64()))
140 .or_else(|| json.get("expires_at").and_then(|v| v.as_i64()));
141
142 let has_refresh = json
143 .get("claudeAiOauth")
144 .and_then(|o| o.get("refreshToken"))
145 .and_then(|v| v.as_str())
146 .map(|s| !s.is_empty())
147 .unwrap_or(false)
148 || json
149 .get("refresh_token")
150 .and_then(|v| v.as_str())
151 .map(|s| !s.is_empty())
152 .unwrap_or(false);
153
154 match expires_at_ms {
155 Some(exp) if exp <= now_ms => {
156 if has_refresh {
161 Some(OAuthProbeStatus::Expired)
162 } else {
163 Some(OAuthProbeStatus::NeedsReauth)
164 }
165 }
166 Some(_) => Some(OAuthProbeStatus::Valid),
167 None => {
168 tracing::debug!(
173 target: "identity",
174 provider,
175 path = %probe_path.display(),
176 "oauth probe: file present but no parseable expiry — status=valid (best-effort)"
177 );
178 Some(OAuthProbeStatus::Valid)
179 }
180 }
181}
182
183#[derive(Debug, thiserror::Error)]
187pub enum ResolverError {
188 #[error("account not found: {0}")]
189 AccountNotFound(String),
190
191 #[error("env var not set in srv environment: {0}")]
192 EnvVarMissing(String),
193
194 #[error("AWS Secrets Manager backend not yet supported (Phase 3)")]
195 SecretsManagerUnsupported,
196
197 #[error("PlaintextDev secrets are disabled in release builds")]
198 PlaintextDevDisabledInRelease,
199
200 #[error("OAuthConfigDir is a config-dir pointer, not a resolvable secret — routed via the oauth-class injection path, not resolve_secret")]
208 OAuthConfigDirNotASecret,
209
210 #[error("keychain error: {0}")]
214 KeychainError(String),
215
216 #[error("storage error: {0}")]
217 Storage(#[from] StoreError),
218}
219
220#[derive(Debug, Clone, PartialEq, Eq)]
224pub enum ProviderClass {
225 ApiKey { env_vars: &'static [&'static str] },
231 OAuth { config_dir_env_var: &'static str },
236}
237
238pub fn provider_class(provider: &str) -> Option<ProviderClass> {
241 match provider {
242 "github" => Some(ProviderClass::ApiKey {
246 env_vars: &["GITHUB_TOKEN", "GH_TOKEN"],
247 }),
248 "anthropic" => Some(ProviderClass::ApiKey {
249 env_vars: &["ANTHROPIC_API_KEY"],
250 }),
251 "openai" => Some(ProviderClass::ApiKey {
252 env_vars: &["OPENAI_API_KEY"],
253 }),
254 "kimi" => Some(ProviderClass::ApiKey {
255 env_vars: &["MOONSHOT_API_KEY"],
256 }),
257 "aws" => Some(ProviderClass::ApiKey {
258 env_vars: &["AWS_ACCESS_KEY_ID"],
259 }),
260 "claude" | "codex" | "openclaw" => {
271 crate::backend::providers::get_provider(provider).map(|cfg| {
272 ProviderClass::OAuth {
273 config_dir_env_var: cfg.auth_config_dir_env_var,
274 }
275 })
276 }
277 _ => None,
278 }
279}
280
281pub fn provider_env_vars(provider: &str) -> Vec<&'static str> {
286 match provider_class(provider) {
287 Some(ProviderClass::ApiKey { env_vars }) => env_vars.to_vec(),
288 _ => Vec::new(),
289 }
290}
291
292pub fn resolve_secret(secret_ref: &SecretRef) -> Result<String, ResolverError> {
308 match secret_ref {
309 SecretRef::Env { env_var } => std::env::var(env_var)
310 .map_err(|_| ResolverError::EnvVarMissing(env_var.clone())),
311 SecretRef::PlaintextDev { plaintext_dev } => {
312 #[cfg(debug_assertions)]
313 {
314 Ok(plaintext_dev.clone())
315 }
316 #[cfg(not(debug_assertions))]
317 {
318 let _ = plaintext_dev;
319 Err(ResolverError::PlaintextDevDisabledInRelease)
320 }
321 }
322 SecretRef::SecretsManager { .. } => Err(ResolverError::SecretsManagerUnsupported),
323 SecretRef::OAuthConfigDir { dir } => {
324 let _ = dir;
330 Err(ResolverError::OAuthConfigDirNotASecret)
331 }
332 SecretRef::Keychain { account, .. } => {
333 let account_id = account.strip_prefix("acct:").unwrap_or(account);
338 crate::identity::secret_store::get(account_id)
339 .map(|z| z.to_string())
340 .map_err(ResolverError::KeychainError)
341 }
342 }
343}
344
345pub fn inject_identity_env(
370 wstore: Arc<Store>,
371 block_id: &str,
372 env_vars: &mut HashMap<String, String>,
373) {
374 inject_identity_env_with_broker(wstore, None, block_id, env_vars);
375}
376
377pub async fn inject_identity_env_async(
394 wstore: Arc<Store>,
395 broker: Option<Arc<Broker>>,
396 block_id: String,
397 env_vars: HashMap<String, String>,
398) -> HashMap<String, String> {
399 let fallback = env_vars.clone();
400 match tokio::task::spawn_blocking(move || {
401 let mut env = env_vars;
402 inject_identity_env_with_broker(wstore, broker, &block_id, &mut env);
403 env
404 })
405 .await
406 {
407 Ok(merged) => merged,
408 Err(e) => {
409 tracing::warn!(target: "identity", "identity injection task join failed: {e}");
410 fallback
411 }
412 }
413}
414
415pub fn inject_identity_env_with_broker(
416 wstore: Arc<Store>,
417 broker: Option<Arc<Broker>>,
418 block_id: &str,
419 env_vars: &mut HashMap<String, String>,
420) {
421 let instance = match wstore.instance_get_active_for_block(block_id) {
423 Ok(Some(i)) => i,
424 Ok(None) => {
425 return;
427 }
428 Err(e) => {
429 tracing::warn!(target: "identity", "instance lookup failed for block {}: {}", block_id, e);
430 return;
431 }
432 };
433
434 if instance.identity_id.is_empty() || instance.identity_id == "blank" {
436 tracing::warn!(
443 target: "identity",
444 "instance {} has empty/blank identity_id — falling back to ambient creds. \
445 Legacy row or UI regression?",
446 block_id
447 );
448 return;
449 }
450
451 let bindings = match wstore.bundle_identity_bindings(&instance.identity_id) {
453 Ok(b) => b,
454 Err(e) => {
455 tracing::warn!(
456 target: "identity",
457 "bindings lookup failed for identity {}: {}",
458 instance.identity_id,
459 e,
460 );
461 return;
462 }
463 };
464
465 if bindings.is_empty() {
466 return;
468 }
469
470 for binding in &bindings {
482 let class = match provider_class(&binding.provider) {
483 Some(c) => c,
484 None => {
485 tracing::warn!(
486 target: "identity",
487 "no provider class for {} (binding for identity {}) — skipping",
488 binding.provider,
489 instance.identity_id,
490 );
491 continue;
492 }
493 };
494
495 let account = match wstore.identity_get(&binding.account_id) {
496 Ok(Some(a)) => a,
497 Ok(None) => {
498 tracing::warn!(
499 target: "identity",
500 "account {} bound to identity {} but row not found — skipping",
501 binding.account_id,
502 instance.identity_id,
503 );
504 continue;
505 }
506 Err(e) => {
507 tracing::warn!(
508 target: "identity",
509 "account lookup failed for {}: {}",
510 binding.account_id,
511 e,
512 );
513 continue;
514 }
515 };
516
517 match class {
518 ProviderClass::ApiKey { env_vars: env_keys } => {
519 let secret = match resolve_secret(&account.secret_ref) {
520 Ok(s) => s,
521 Err(e) => {
522 tracing::warn!(
523 target: "identity",
524 "secret resolution failed for account {} (provider {}): {} — skipping",
525 binding.account_id,
526 binding.provider,
527 e,
528 );
529 continue;
530 }
531 };
532 let env_key_count = env_keys.len();
533 for key in env_keys {
534 env_vars.insert(key.to_string(), secret.clone());
535 }
536 tracing::info!(
537 target: "identity",
538 "injected {} env var(s) for api-key provider {} (identity={}, account={})",
539 env_key_count,
540 binding.provider,
541 instance.identity_id,
542 binding.account_id,
543 );
544 }
545 ProviderClass::OAuth { config_dir_env_var } => {
546 let dir = match &account.secret_ref {
550 SecretRef::OAuthConfigDir { dir } => dir.clone(),
551 other => {
552 tracing::warn!(
553 target: "identity",
554 "oauth-class provider {} has non-OAuthConfigDir secret_ref \
555 ({:?}) on account {} — skipping",
556 binding.provider,
557 other,
558 binding.account_id,
559 );
560 continue;
561 }
562 };
563 env_vars.insert(config_dir_env_var.to_string(), dir.clone());
564 tracing::info!(
565 target: "identity",
566 "injected {} for oauth provider {} (identity={}, account={})",
567 config_dir_env_var,
568 binding.provider,
569 instance.identity_id,
570 binding.account_id,
571 );
572
573 let now_ms = SystemTime::now()
582 .duration_since(UNIX_EPOCH)
583 .map(|d| d.as_millis() as i64)
584 .unwrap_or(0);
585 if let Some(probed) = probe_oauth_status(&binding.provider, &dir, now_ms) {
586 let new_status = probed.as_str();
587 if account.status != new_status {
588 let mut updated = account.clone();
589 updated.status = new_status.to_string();
590 updated.updated_at = now_ms;
591 match wstore.identity_upsert(&updated) {
592 Ok(()) => {
593 tracing::info!(
594 target: "identity",
595 provider = %binding.provider,
596 account_id = %binding.account_id,
597 old_status = %account.status,
598 new_status,
599 "oauth probe: status updated"
600 );
601 if let Some(b) = broker.as_ref() {
610 b.publish(WaveEvent {
611 event: format!(
612 "identitybundlebindings:changed:{}",
613 instance.identity_id,
614 ),
615 scopes: vec![],
616 sender: String::new(),
617 persist: 0,
618 data: None,
619 });
620 }
621 }
622 Err(e) => {
623 tracing::warn!(
624 target: "identity",
625 provider = %binding.provider,
626 account_id = %binding.account_id,
627 error = %e,
628 "oauth probe: identity_upsert failed — status not persisted",
629 );
630 }
631 }
632 }
633 }
634 }
635 }
636 }
637}
638
639#[cfg(test)]
640mod tests {
641 use super::*;
642 use crate::backend::storage::store::{
643 AgentInstance, Identity, IdentityAccount, InstanceStatus, SecretRef,
644 };
645
646 fn make_store() -> Arc<Store> {
647 Arc::new(Store::open_in_memory().unwrap())
648 }
649
650 fn make_account(
651 id: &str,
652 provider: &str,
653 secret_ref: SecretRef,
654 ) -> IdentityAccount {
655 IdentityAccount {
656 id: id.to_string(),
657 name: format!("{}-{}", provider, id),
658 provider: provider.to_string(),
659 kind: "pat".to_string(),
660 display_name: String::new(),
661 secret_ref,
662 context: serde_json::json!({}),
663 status: "unknown".to_string(),
664 created_at: 0,
665 updated_at: 0,
666 }
667 }
668
669 fn insert_block_for_agent(store: &Store, block_id: &str, agent_id: &str) {
675 use crate::backend::obj::{Block, MetaMapType};
676 let mut block = Block {
677 oid: block_id.to_string(),
678 parentoref: String::new(),
679 version: 0,
680 runtimeopts: None,
681 stickers: None,
682 meta: {
683 let mut m = MetaMapType::new();
684 m.insert("view".to_string(), serde_json::json!("agent"));
685 m.insert("agentId".to_string(), serde_json::json!(agent_id));
686 m
687 },
688 subblockids: None,
689 };
690 store.insert(&mut block).unwrap();
691 }
692
693 fn make_instance(block_id: &str, identity_id: &str) -> AgentInstance {
694 AgentInstance {
695 id: format!("inst-{block_id}"),
696 definition_id: "def-1".to_string(),
697 parent_instance_id: String::new(),
698 block_id: block_id.to_string(),
699 session_id: String::new(),
700 status: InstanceStatus::Running.as_str().to_string(),
701 github_context: String::new(),
702 started_at: 0,
703 ended_at: 0,
704 created_at: 0,
705 identity_id: identity_id.to_string(),
706 memory_id: String::new(),
707 instance_name: String::new(),
708 working_directory: String::new(),
709 display_hidden: false,
710 }
711 }
712
713 #[test]
714 fn provider_env_vars_matrix() {
715 assert_eq!(provider_env_vars("github"), vec!["GITHUB_TOKEN", "GH_TOKEN"]);
716 assert_eq!(provider_env_vars("anthropic"), vec!["ANTHROPIC_API_KEY"]);
717 assert_eq!(provider_env_vars("openai"), vec!["OPENAI_API_KEY"]);
718 assert_eq!(provider_env_vars("kimi"), vec!["MOONSHOT_API_KEY"]);
719 assert_eq!(provider_env_vars("aws"), vec!["AWS_ACCESS_KEY_ID"]);
720 assert!(provider_env_vars("unknown").is_empty());
721 }
722
723 #[cfg(debug_assertions)]
731 #[test]
732 fn resolve_plaintext_dev() {
733 let s = resolve_secret(&SecretRef::PlaintextDev {
734 plaintext_dev: "ghp_test123".to_string(),
735 })
736 .unwrap();
737 assert_eq!(s, "ghp_test123");
738 }
739
740 #[test]
741 fn resolve_env_var_missing() {
742 let res = resolve_secret(&SecretRef::Env {
743 env_var: "AGENTMUX_TEST_NEVER_SET_X9Q".to_string(),
744 });
745 assert!(matches!(res, Err(ResolverError::EnvVarMissing(_))));
746 }
747
748 #[test]
749 fn resolve_secrets_manager_unsupported() {
750 let res = resolve_secret(&SecretRef::SecretsManager {
751 sm_path: "ignored".to_string(),
752 sm_json_path: None,
753 });
754 assert!(matches!(res, Err(ResolverError::SecretsManagerUnsupported)));
755 }
756
757 #[test]
758 fn provider_class_oauth_providers() {
759 assert_eq!(
766 provider_class("claude"),
767 Some(ProviderClass::OAuth { config_dir_env_var: "CLAUDE_CONFIG_DIR" }),
768 );
769 assert_eq!(
770 provider_class("codex"),
771 Some(ProviderClass::OAuth { config_dir_env_var: "CODEX_HOME" }),
772 );
773 assert_eq!(
774 provider_class("openclaw"),
775 Some(ProviderClass::OAuth { config_dir_env_var: "OPENCLAW_HOME" }),
776 );
777 }
778
779 #[cfg(debug_assertions)]
780 #[test]
781 fn inject_oauth_class_sets_config_dir_env_var() {
782 let store = make_store();
783
784 let mut def = crate::backend::storage::store::AgentDefinition {
785 id: "def-1".to_string(),
786 slug: String::new(),
787 name: "T".to_string(),
788 icon: "✦".to_string(),
789 provider: "claude".to_string(),
790 description: String::new(),
791 working_directory: String::new(),
792 shell: String::new(),
793 provider_flags: String::new(),
794 auto_start: 0,
795 restart_on_crash: 0,
796 idle_timeout_minutes: 0,
797 created_at: 0,
798 agent_type: String::new(),
799 environment: String::new(),
800 agent_bus_id: String::new(),
801 is_seeded: 0,
802 accounts: String::new(),
803 parent_id: String::new(),
804 branch_label: String::new(),
805 updated_at: 0,
806 user_hidden: 0,
807 container_image: String::new(),
808 container_volumes: "[]".to_string(),
809 container_name: String::new(),
810 };
811 store.agent_def_insert(&mut def).unwrap();
812
813 let identity = Identity {
814 id: "id-oauth".to_string(),
815 name: "OAuth".to_string(),
816 description: String::new(),
817 is_blank: false,
818 created_at: 0,
819 updated_at: 0,
820 };
821 store.bundle_identity_upsert(&identity).unwrap();
822
823 let claude = make_account(
824 "acct-claude",
825 "claude",
826 SecretRef::OAuthConfigDir {
827 dir: "/var/agentmux/identities/id-oauth/claude".to_string(),
828 },
829 );
830 store.identity_upsert(&claude).unwrap();
831 store
832 .bundle_identity_bind("id-oauth", "claude", "acct-claude")
833 .unwrap();
834
835 insert_block_for_agent(&store, "block-oauth", "def-1");
836 let inst = make_instance("block-oauth", "id-oauth");
837 store.instance_create(&inst).unwrap();
838
839 let mut env: HashMap<String, String> = HashMap::new();
840 inject_identity_env(store, "block-oauth", &mut env);
841
842 assert_eq!(
844 env.get("CLAUDE_CONFIG_DIR").map(String::as_str),
845 Some("/var/agentmux/identities/id-oauth/claude"),
846 );
847 assert!(env.get("ANTHROPIC_API_KEY").is_none());
850 }
851
852 #[cfg(debug_assertions)]
853 #[test]
854 fn inject_oauth_class_skips_account_with_non_oauth_secret_ref() {
855 let store = make_store();
860
861 let mut def = crate::backend::storage::store::AgentDefinition {
862 id: "def-1".to_string(),
863 slug: String::new(),
864 name: "T".to_string(),
865 icon: "✦".to_string(),
866 provider: "claude".to_string(),
867 description: String::new(),
868 working_directory: String::new(),
869 shell: String::new(),
870 provider_flags: String::new(),
871 auto_start: 0,
872 restart_on_crash: 0,
873 idle_timeout_minutes: 0,
874 created_at: 0,
875 agent_type: String::new(),
876 environment: String::new(),
877 agent_bus_id: String::new(),
878 is_seeded: 0,
879 accounts: String::new(),
880 parent_id: String::new(),
881 branch_label: String::new(),
882 updated_at: 0,
883 user_hidden: 0,
884 container_image: String::new(),
885 container_volumes: "[]".to_string(),
886 container_name: String::new(),
887 };
888 store.agent_def_insert(&mut def).unwrap();
889
890 let identity = Identity {
891 id: "id-bad".to_string(),
892 name: "Bad".to_string(),
893 description: String::new(),
894 is_blank: false,
895 created_at: 0,
896 updated_at: 0,
897 };
898 store.bundle_identity_upsert(&identity).unwrap();
899
900 let bad = make_account(
901 "acct-bad",
902 "claude",
903 SecretRef::Env {
904 env_var: "CLAUDE_TOKEN_NOT_A_DIR".to_string(),
905 },
906 );
907 store.identity_upsert(&bad).unwrap();
908 store
909 .bundle_identity_bind("id-bad", "claude", "acct-bad")
910 .unwrap();
911
912 insert_block_for_agent(&store, "block-bad", "def-1");
913 let inst = make_instance("block-bad", "id-bad");
914 store.instance_create(&inst).unwrap();
915
916 let mut env: HashMap<String, String> = HashMap::new();
917 inject_identity_env(store, "block-bad", &mut env);
918
919 assert!(env.get("CLAUDE_CONFIG_DIR").is_none());
921 assert!(env.is_empty());
922 }
923
924 #[test]
925 fn resolve_oauth_config_dir_is_not_a_secret() {
926 let res = resolve_secret(&SecretRef::OAuthConfigDir {
934 dir: "/path/to/bundle/claude".to_string(),
935 });
936 assert!(matches!(res, Err(ResolverError::OAuthConfigDirNotASecret)));
937 }
938
939 #[test]
940 fn inject_no_instance_does_nothing() {
941 let store = make_store();
942 let mut env: HashMap<String, String> = HashMap::new();
943 inject_identity_env(store, "block-no-instance", &mut env);
944 assert!(env.is_empty());
945 }
946
947 #[test]
948 fn inject_blank_identity_does_nothing() {
949 let store = make_store();
950 let mut def = crate::backend::storage::store::AgentDefinition {
952 id: "def-1".to_string(),
953 slug: String::new(),
954 name: "T".to_string(),
955 icon: "✦".to_string(),
956 provider: "claude".to_string(),
957 description: String::new(),
958 working_directory: String::new(),
959 shell: String::new(),
960 provider_flags: String::new(),
961 auto_start: 0,
962 restart_on_crash: 0,
963 idle_timeout_minutes: 0,
964 created_at: 0,
965 agent_type: String::new(),
966 environment: String::new(),
967 agent_bus_id: String::new(),
968 is_seeded: 0,
969 accounts: String::new(),
970 parent_id: String::new(),
971 branch_label: String::new(),
972 updated_at: 0,
973 user_hidden: 0,
974 container_image: String::new(),
975 container_volumes: "[]".to_string(),
976 container_name: String::new(),
977 };
978 store.agent_def_insert(&mut def).unwrap();
979
980 insert_block_for_agent(&store, "block-blank", "def-1");
981 let mut inst = make_instance("block-blank", "blank");
982 store.instance_create(&inst).unwrap();
983 let _ = inst; let mut env: HashMap<String, String> = HashMap::new();
986 inject_identity_env(store, "block-blank", &mut env);
987 assert!(env.is_empty());
988 }
989
990 #[cfg(debug_assertions)]
991 #[test]
992 fn inject_full_round_trip_plaintext_dev() {
993 let store = make_store();
994
995 let mut def = crate::backend::storage::store::AgentDefinition {
997 id: "def-1".to_string(),
998 slug: String::new(),
999 name: "T".to_string(),
1000 icon: "✦".to_string(),
1001 provider: "claude".to_string(),
1002 description: String::new(),
1003 working_directory: String::new(),
1004 shell: String::new(),
1005 provider_flags: String::new(),
1006 auto_start: 0,
1007 restart_on_crash: 0,
1008 idle_timeout_minutes: 0,
1009 created_at: 0,
1010 agent_type: String::new(),
1011 environment: String::new(),
1012 agent_bus_id: String::new(),
1013 is_seeded: 0,
1014 accounts: String::new(),
1015 parent_id: String::new(),
1016 branch_label: String::new(),
1017 updated_at: 0,
1018 user_hidden: 0,
1019 container_image: String::new(),
1020 container_volumes: "[]".to_string(),
1021 container_name: String::new(),
1022 };
1023 store.agent_def_insert(&mut def).unwrap();
1024
1025 let identity = Identity {
1027 id: "id-work".to_string(),
1028 name: "Work".to_string(),
1029 description: String::new(),
1030 is_blank: false,
1031 created_at: 0,
1032 updated_at: 0,
1033 };
1034 store.bundle_identity_upsert(&identity).unwrap();
1035
1036 let github = make_account(
1038 "acct-gh",
1039 "github",
1040 SecretRef::PlaintextDev {
1041 plaintext_dev: "ghp_round_trip".to_string(),
1042 },
1043 );
1044 store.identity_upsert(&github).unwrap();
1045 store
1046 .bundle_identity_bind("id-work", "github", "acct-gh")
1047 .unwrap();
1048
1049 let anthropic = make_account(
1051 "acct-anth",
1052 "anthropic",
1053 SecretRef::PlaintextDev {
1054 plaintext_dev: "sk-ant-round_trip".to_string(),
1055 },
1056 );
1057 store.identity_upsert(&anthropic).unwrap();
1058 store
1059 .bundle_identity_bind("id-work", "anthropic", "acct-anth")
1060 .unwrap();
1061
1062 insert_block_for_agent(&store, "block-1", "def-1");
1064 let inst = make_instance("block-1", "id-work");
1065 store.instance_create(&inst).unwrap();
1066
1067 let mut env: HashMap<String, String> = HashMap::new();
1068 inject_identity_env(store, "block-1", &mut env);
1069
1070 assert_eq!(env.get("GITHUB_TOKEN").map(String::as_str), Some("ghp_round_trip"));
1072 assert_eq!(env.get("GH_TOKEN").map(String::as_str), Some("ghp_round_trip"));
1073 assert_eq!(
1075 env.get("ANTHROPIC_API_KEY").map(String::as_str),
1076 Some("sk-ant-round_trip"),
1077 );
1078 }
1079
1080 #[cfg(debug_assertions)]
1081 #[test]
1082 fn inject_partial_success_skips_failed_bindings() {
1083 let store = make_store();
1084
1085 let mut def = crate::backend::storage::store::AgentDefinition {
1086 id: "def-1".to_string(),
1087 slug: String::new(),
1088 name: "T".to_string(),
1089 icon: "✦".to_string(),
1090 provider: "claude".to_string(),
1091 description: String::new(),
1092 working_directory: String::new(),
1093 shell: String::new(),
1094 provider_flags: String::new(),
1095 auto_start: 0,
1096 restart_on_crash: 0,
1097 idle_timeout_minutes: 0,
1098 created_at: 0,
1099 agent_type: String::new(),
1100 environment: String::new(),
1101 agent_bus_id: String::new(),
1102 is_seeded: 0,
1103 accounts: String::new(),
1104 parent_id: String::new(),
1105 branch_label: String::new(),
1106 updated_at: 0,
1107 user_hidden: 0,
1108 container_image: String::new(),
1109 container_volumes: "[]".to_string(),
1110 container_name: String::new(),
1111 };
1112 store.agent_def_insert(&mut def).unwrap();
1113
1114 let identity = Identity {
1115 id: "id-mixed".to_string(),
1116 name: "Mixed".to_string(),
1117 description: String::new(),
1118 is_blank: false,
1119 created_at: 0,
1120 updated_at: 0,
1121 };
1122 store.bundle_identity_upsert(&identity).unwrap();
1123
1124 let good = make_account(
1126 "acct-good",
1127 "github",
1128 SecretRef::PlaintextDev {
1129 plaintext_dev: "ghp_good".to_string(),
1130 },
1131 );
1132 store.identity_upsert(&good).unwrap();
1133 store
1134 .bundle_identity_bind("id-mixed", "github", "acct-good")
1135 .unwrap();
1136
1137 let bad = make_account(
1139 "acct-bad",
1140 "anthropic",
1141 SecretRef::Env {
1142 env_var: "AGENTMUX_TEST_DEFINITELY_NOT_SET_4242".to_string(),
1143 },
1144 );
1145 store.identity_upsert(&bad).unwrap();
1146 store
1147 .bundle_identity_bind("id-mixed", "anthropic", "acct-bad")
1148 .unwrap();
1149
1150 insert_block_for_agent(&store, "block-mixed", "def-1");
1151 let inst = make_instance("block-mixed", "id-mixed");
1152 store.instance_create(&inst).unwrap();
1153
1154 let mut env: HashMap<String, String> = HashMap::new();
1155 inject_identity_env(store, "block-mixed", &mut env);
1156
1157 assert_eq!(env.get("GITHUB_TOKEN").map(String::as_str), Some("ghp_good"));
1159 assert_eq!(env.get("GH_TOKEN").map(String::as_str), Some("ghp_good"));
1160 assert!(env.get("ANTHROPIC_API_KEY").is_none());
1162 }
1163
1164 #[cfg(debug_assertions)]
1165 #[test]
1166 fn inject_unknown_provider_is_skipped() {
1167 let store = make_store();
1168
1169 let mut def = crate::backend::storage::store::AgentDefinition {
1170 id: "def-1".to_string(),
1171 slug: String::new(),
1172 name: "T".to_string(),
1173 icon: "✦".to_string(),
1174 provider: "claude".to_string(),
1175 description: String::new(),
1176 working_directory: String::new(),
1177 shell: String::new(),
1178 provider_flags: String::new(),
1179 auto_start: 0,
1180 restart_on_crash: 0,
1181 idle_timeout_minutes: 0,
1182 created_at: 0,
1183 agent_type: String::new(),
1184 environment: String::new(),
1185 agent_bus_id: String::new(),
1186 is_seeded: 0,
1187 accounts: String::new(),
1188 parent_id: String::new(),
1189 branch_label: String::new(),
1190 updated_at: 0,
1191 user_hidden: 0,
1192 container_image: String::new(),
1193 container_volumes: "[]".to_string(),
1194 container_name: String::new(),
1195 };
1196 store.agent_def_insert(&mut def).unwrap();
1197
1198 let identity = Identity {
1199 id: "id-future".to_string(),
1200 name: "Future".to_string(),
1201 description: String::new(),
1202 is_blank: false,
1203 created_at: 0,
1204 updated_at: 0,
1205 };
1206 store.bundle_identity_upsert(&identity).unwrap();
1207
1208 let custom = make_account(
1209 "acct-custom",
1210 "custom",
1211 SecretRef::PlaintextDev {
1212 plaintext_dev: "ignored".to_string(),
1213 },
1214 );
1215 store.identity_upsert(&custom).unwrap();
1216 store
1217 .bundle_identity_bind("id-future", "custom", "acct-custom")
1218 .unwrap();
1219
1220 insert_block_for_agent(&store, "block-future", "def-1");
1221 let inst = make_instance("block-future", "id-future");
1222 store.instance_create(&inst).unwrap();
1223
1224 let mut env: HashMap<String, String> = HashMap::new();
1225 inject_identity_env(store, "block-future", &mut env);
1226 assert!(env.is_empty());
1228 }
1229
1230 fn write_claude_creds(
1237 dir: &std::path::Path,
1238 expires_ms: i64,
1239 with_refresh: bool,
1240 ) {
1241 std::fs::create_dir_all(dir).unwrap();
1242 let body = serde_json::json!({
1243 "claudeAiOauth": {
1244 "accessToken": "test-access",
1245 "refreshToken": if with_refresh { "test-refresh" } else { "" },
1246 "expiresAt": expires_ms,
1247 }
1248 });
1249 std::fs::write(
1250 dir.join(".credentials.json"),
1251 serde_json::to_string(&body).unwrap(),
1252 )
1253 .unwrap();
1254 }
1255
1256 #[test]
1257 fn probe_oauth_status_unknown_provider_returns_none() {
1258 let r = probe_oauth_status("github", "/tmp/whatever", 0);
1264 assert_eq!(r, None);
1265 }
1266
1267 #[test]
1268 fn probe_oauth_status_missing_dir_is_needs_reauth() {
1269 let r = probe_oauth_status("claude", "/definitely/does/not/exist-xyz-9q", 0);
1270 assert_eq!(r, Some(OAuthProbeStatus::NeedsReauth));
1271 }
1272
1273 #[test]
1274 fn probe_oauth_status_future_expiry_is_valid() {
1275 let tmp = tempfile::tempdir().unwrap();
1276 let now_ms = 1_700_000_000_000;
1277 write_claude_creds(tmp.path(), now_ms + 3_600_000, true);
1278 let r = probe_oauth_status("claude", tmp.path().to_str().unwrap(), now_ms);
1279 assert_eq!(r, Some(OAuthProbeStatus::Valid));
1280 }
1281
1282 #[test]
1283 fn probe_oauth_status_past_expiry_with_refresh_is_expired() {
1284 let tmp = tempfile::tempdir().unwrap();
1285 let now_ms = 1_700_000_000_000;
1286 write_claude_creds(tmp.path(), now_ms - 1, true);
1287 let r = probe_oauth_status("claude", tmp.path().to_str().unwrap(), now_ms);
1288 assert_eq!(r, Some(OAuthProbeStatus::Expired));
1289 }
1290
1291 #[test]
1292 fn probe_oauth_status_past_expiry_no_refresh_is_needs_reauth() {
1293 let tmp = tempfile::tempdir().unwrap();
1297 let now_ms = 1_700_000_000_000;
1298 write_claude_creds(tmp.path(), now_ms - 1, false);
1299 let r = probe_oauth_status("claude", tmp.path().to_str().unwrap(), now_ms);
1300 assert_eq!(r, Some(OAuthProbeStatus::NeedsReauth));
1301 }
1302
1303 #[test]
1304 fn probe_oauth_status_malformed_json_is_needs_reauth() {
1305 let tmp = tempfile::tempdir().unwrap();
1306 std::fs::write(tmp.path().join(".credentials.json"), "{ not json").unwrap();
1307 let r = probe_oauth_status("claude", tmp.path().to_str().unwrap(), 0);
1308 assert_eq!(r, Some(OAuthProbeStatus::NeedsReauth));
1309 }
1310
1311 #[test]
1312 fn probe_oauth_status_codex_unknown_shape_is_valid_best_effort() {
1313 let tmp = tempfile::tempdir().unwrap();
1319 std::fs::write(
1320 tmp.path().join(".credentials.json"),
1321 r#"{"some":"opaque-codex-blob"}"#,
1322 )
1323 .unwrap();
1324 let r = probe_oauth_status("codex", tmp.path().to_str().unwrap(), 0);
1325 assert_eq!(r, Some(OAuthProbeStatus::Valid));
1326 }
1327
1328 #[cfg(debug_assertions)]
1329 #[test]
1330 fn inject_oauth_class_probes_and_flips_status_to_needs_reauth() {
1331 let store = make_store();
1336
1337 let mut def = crate::backend::storage::store::AgentDefinition {
1338 id: "def-1".to_string(),
1339 slug: String::new(),
1340 name: "T".to_string(),
1341 icon: "✦".to_string(),
1342 provider: "claude".to_string(),
1343 description: String::new(),
1344 working_directory: String::new(),
1345 shell: String::new(),
1346 provider_flags: String::new(),
1347 auto_start: 0,
1348 restart_on_crash: 0,
1349 idle_timeout_minutes: 0,
1350 created_at: 0,
1351 agent_type: String::new(),
1352 environment: String::new(),
1353 agent_bus_id: String::new(),
1354 is_seeded: 0,
1355 accounts: String::new(),
1356 parent_id: String::new(),
1357 branch_label: String::new(),
1358 updated_at: 0,
1359 user_hidden: 0,
1360 container_image: String::new(),
1361 container_volumes: "[]".to_string(),
1362 container_name: String::new(),
1363 };
1364 store.agent_def_insert(&mut def).unwrap();
1365
1366 let identity = Identity {
1367 id: "id-probe".to_string(),
1368 name: "Probe".to_string(),
1369 description: String::new(),
1370 is_blank: false,
1371 created_at: 0,
1372 updated_at: 0,
1373 };
1374 store.bundle_identity_upsert(&identity).unwrap();
1375
1376 let tmp = tempfile::tempdir().unwrap();
1379 let bundle_dir = tmp.path().to_str().unwrap().to_string();
1380
1381 let claude = IdentityAccount {
1382 id: "acct-claude".to_string(),
1383 name: "claude-acct-claude".to_string(),
1384 provider: "claude".to_string(),
1385 kind: "oauth".to_string(),
1386 display_name: String::new(),
1387 secret_ref: SecretRef::OAuthConfigDir { dir: bundle_dir },
1388 context: serde_json::json!({}),
1389 status: oauth_status::VALID.to_string(),
1392 created_at: 0,
1393 updated_at: 0,
1394 };
1395 store.identity_upsert(&claude).unwrap();
1396 store
1397 .bundle_identity_bind("id-probe", "claude", "acct-claude")
1398 .unwrap();
1399
1400 insert_block_for_agent(&store, "block-probe", "def-1");
1401 let inst = make_instance("block-probe", "id-probe");
1402 store.instance_create(&inst).unwrap();
1403
1404 let mut env: HashMap<String, String> = HashMap::new();
1405 inject_identity_env(store.clone(), "block-probe", &mut env);
1406
1407 assert!(env.get("CLAUDE_CONFIG_DIR").is_some());
1411
1412 let after = store.identity_get("acct-claude").unwrap().unwrap();
1414 assert_eq!(after.status, oauth_status::NEEDS_REAUTH);
1415 }
1416
1417 #[cfg(debug_assertions)]
1418 #[test]
1419 fn inject_oauth_class_probe_preserves_status_when_valid() {
1420 let store = make_store();
1425
1426 let mut def = crate::backend::storage::store::AgentDefinition {
1427 id: "def-1".to_string(),
1428 slug: String::new(),
1429 name: "T".to_string(),
1430 icon: "✦".to_string(),
1431 provider: "claude".to_string(),
1432 description: String::new(),
1433 working_directory: String::new(),
1434 shell: String::new(),
1435 provider_flags: String::new(),
1436 auto_start: 0,
1437 restart_on_crash: 0,
1438 idle_timeout_minutes: 0,
1439 created_at: 0,
1440 agent_type: String::new(),
1441 environment: String::new(),
1442 agent_bus_id: String::new(),
1443 is_seeded: 0,
1444 accounts: String::new(),
1445 parent_id: String::new(),
1446 branch_label: String::new(),
1447 updated_at: 0,
1448 user_hidden: 0,
1449 container_image: String::new(),
1450 container_volumes: "[]".to_string(),
1451 container_name: String::new(),
1452 };
1453 store.agent_def_insert(&mut def).unwrap();
1454
1455 let identity = Identity {
1456 id: "id-ok".to_string(),
1457 name: "Ok".to_string(),
1458 description: String::new(),
1459 is_blank: false,
1460 created_at: 0,
1461 updated_at: 0,
1462 };
1463 store.bundle_identity_upsert(&identity).unwrap();
1464
1465 let tmp = tempfile::tempdir().unwrap();
1466 let now_ms = SystemTime::now()
1467 .duration_since(UNIX_EPOCH)
1468 .unwrap()
1469 .as_millis() as i64;
1470 write_claude_creds(tmp.path(), now_ms + 3_600_000, true);
1471
1472 let claude = IdentityAccount {
1473 id: "acct-ok".to_string(),
1474 name: "claude-acct-ok".to_string(),
1475 provider: "claude".to_string(),
1476 kind: "oauth".to_string(),
1477 display_name: String::new(),
1478 secret_ref: SecretRef::OAuthConfigDir {
1479 dir: tmp.path().to_str().unwrap().to_string(),
1480 },
1481 context: serde_json::json!({}),
1482 status: oauth_status::VALID.to_string(),
1483 created_at: 0,
1484 updated_at: 0,
1485 };
1486 store.identity_upsert(&claude).unwrap();
1487 store
1488 .bundle_identity_bind("id-ok", "claude", "acct-ok")
1489 .unwrap();
1490
1491 insert_block_for_agent(&store, "block-ok", "def-1");
1492 let inst = make_instance("block-ok", "id-ok");
1493 store.instance_create(&inst).unwrap();
1494
1495 let mut env: HashMap<String, String> = HashMap::new();
1496 inject_identity_env(store.clone(), "block-ok", &mut env);
1497
1498 let after = store.identity_get("acct-ok").unwrap().unwrap();
1499 assert_eq!(after.status, oauth_status::VALID);
1500 assert_eq!(after.updated_at, 0);
1503 }
1504}