Skip to main content

codetether_agent/provenance/
types.rs

1use super::ClaimProvenance;
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use uuid::Uuid;
5
6#[derive(Debug, Clone, Serialize, Deserialize)]
7pub struct AgentIdentity {
8    pub tenant_id: Option<String>,
9    pub agent_identity_id: Option<String>,
10    pub agent_name: String,
11    pub origin: ExecutionOrigin,
12    pub worker_id: Option<String>,
13    pub key_id: Option<String>,
14    pub github_installation_id: Option<String>,
15    pub github_app_id: Option<String>,
16}
17
18#[derive(Debug, Clone, Copy, Default, Eq, PartialEq, Serialize, Deserialize)]
19#[serde(rename_all = "snake_case")]
20pub enum ExecutionOrigin {
21    #[default]
22    LocalCli,
23    Worker,
24    Ralph,
25    Swarm,
26    Relay,
27}
28
29impl ExecutionOrigin {
30    pub fn as_str(self) -> &'static str {
31        match self {
32            Self::LocalCli => "local_cli",
33            Self::Worker => "worker",
34            Self::Ralph => "ralph",
35            Self::Swarm => "swarm",
36            Self::Relay => "relay",
37        }
38    }
39}
40
41#[derive(Debug, Clone, Serialize, Deserialize)]
42pub struct ExecutionProvenance {
43    pub provenance_id: String,
44    pub session_id: Option<String>,
45    pub task_id: Option<String>,
46    pub run_id: Option<String>,
47    pub attempt_id: Option<String>,
48    pub issued_at: DateTime<Utc>,
49    pub identity: AgentIdentity,
50}
51
52impl ExecutionProvenance {
53    pub fn for_session(session_id: &str, agent_name: &str) -> Self {
54        Self::new(Some(session_id.to_string()), agent_name, runtime_origin())
55    }
56
57    pub fn for_operation(agent_name: &str, origin: ExecutionOrigin) -> Self {
58        Self::new(None, agent_name, origin)
59    }
60
61    pub fn set_agent_name(&mut self, agent_name: &str) {
62        self.identity.agent_name = agent_name.to_string();
63    }
64
65    pub fn apply_worker_task(&mut self, worker_id: &str, task_id: &str) {
66        self.identity.origin = ExecutionOrigin::Worker;
67        self.identity.worker_id = Some(worker_id.to_string());
68        self.task_id = Some(task_id.to_string());
69    }
70
71    pub fn apply_claim(&mut self, claim: &ClaimProvenance) {
72        self.apply_worker_task(&claim.worker_id, &claim.task_id);
73        self.run_id = claim.run_id.clone().or_else(|| self.run_id.clone());
74        self.attempt_id = claim.attempt_id.clone().or_else(|| self.attempt_id.clone());
75        self.identity.tenant_id = claim
76            .tenant_id
77            .clone()
78            .or_else(|| self.identity.tenant_id.clone());
79        self.identity.agent_identity_id = claim
80            .agent_identity_id
81            .clone()
82            .or_else(|| self.identity.agent_identity_id.clone());
83    }
84
85    pub fn set_run_id(&mut self, run_id: impl Into<String>) {
86        self.run_id = Some(run_id.into());
87    }
88
89    fn new(session_id: Option<String>, agent_name: &str, origin: ExecutionOrigin) -> Self {
90        Self {
91            provenance_id: format!("ctprov_{}", Uuid::new_v4().simple()),
92            session_id,
93            task_id: std::env::var("CODETETHER_TASK_ID").ok(),
94            run_id: std::env::var("CODETETHER_RUN_ID").ok(),
95            attempt_id: std::env::var("CODETETHER_ATTEMPT_ID").ok(),
96            issued_at: Utc::now(),
97            identity: AgentIdentity {
98                tenant_id: std::env::var("CODETETHER_TENANT_ID").ok(),
99                agent_identity_id: std::env::var("CODETETHER_AGENT_IDENTITY_ID").ok(),
100                agent_name: agent_name.to_string(),
101                origin,
102                worker_id: std::env::var("CODETETHER_WORKER_ID").ok(),
103                key_id: std::env::var("CODETETHER_KEY_ID").ok(),
104                github_installation_id: std::env::var("CODETETHER_GITHUB_INSTALLATION_ID").ok(),
105                github_app_id: std::env::var("CODETETHER_GITHUB_APP_ID").ok(),
106            },
107        }
108    }
109}
110
111fn runtime_origin() -> ExecutionOrigin {
112    if std::env::var("CODETETHER_WORKER_ID").is_ok() {
113        ExecutionOrigin::Worker
114    } else {
115        ExecutionOrigin::LocalCli
116    }
117}