Skip to main content

systemprompt_database/repository/service/
endpoint.rs

1//! Advertised replica addresses and the lookup that finds which live replica
2//! owns a named MCP or agent process.
3//!
4//! A replica that supervises processes records the address its peers reach it
5//! on with every registry heartbeat; a replica without a local process
6//! forwards to the owner [`ServiceRepository::find_process_owner`] names. The
7//! owner must be heartbeating both its endpoint and its `services` row within
8//! [`super::INSTANCE_DEAD_AFTER_SECS`], so a crashed owner is never chosen,
9//! and a `worker` owner must also hold the process's unexpired
10//! `<module>:<name>` lease, so a running row left behind on a replica that
11//! handed the process over never wins.
12//!
13//! Copyright (c) systemprompt.io — Business Source License 1.1.
14//! See <https://systemprompt.io> for licensing details.
15
16use systemprompt_identifiers::{InstanceId, ServiceName};
17use systemprompt_manifest::profile::NodeRole;
18use systemprompt_manifest::services::ServiceStatus;
19
20use super::liveness::INSTANCE_DEAD_AFTER_SECS;
21use super::repo::ServiceRepository;
22use crate::error::DatabaseResult;
23
24/// The live replica running a named process and the address it advertises.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct ProcessOwner {
27    pub instance_id: InstanceId,
28    pub advertise_url: String,
29}
30
31impl ServiceRepository {
32    pub async fn advertise_endpoint(&self, advertise_url: &str) -> DatabaseResult<()> {
33        sqlx::query!(
34            r#"
35            INSERT INTO instance_endpoints (instance_id, advertise_url, heartbeat_at)
36            VALUES ($1, $2, CURRENT_TIMESTAMP)
37            ON CONFLICT (instance_id) DO UPDATE SET
38                advertise_url = EXCLUDED.advertise_url,
39                heartbeat_at = EXCLUDED.heartbeat_at
40            "#,
41            self.instance_id.as_str(),
42            advertise_url
43        )
44        .execute(&*self.write_pool)
45        .await?;
46        Ok(())
47    }
48
49    pub async fn find_process_owner(
50        &self,
51        service_name: &ServiceName,
52    ) -> DatabaseResult<Option<ProcessOwner>> {
53        let row = sqlx::query!(
54            r#"
55            SELECT s.instance_id AS "instance_id: InstanceId", e.advertise_url
56            FROM services s
57            JOIN instance_endpoints e ON e.instance_id = s.instance_id
58            JOIN instance_presence p ON p.instance_id = s.instance_id
59            WHERE s.name = $1
60              AND s.status = $2
61              AND s.instance_id <> $3
62              AND s.heartbeat_at
63                  >= CURRENT_TIMESTAMP - make_interval(secs => $4::double precision)
64              AND e.heartbeat_at
65                  >= CURRENT_TIMESTAMP - make_interval(secs => $4::double precision)
66              AND (
67                  p.role <> $5
68                  OR EXISTS (
69                      SELECT 1 FROM instance_leases l
70                      WHERE l.lease_key = s.module_name || ':' || s.name
71                        AND l.owner_instance_id = s.instance_id
72                        AND l.expires_at > CURRENT_TIMESTAMP
73                  )
74              )
75            ORDER BY s.heartbeat_at DESC
76            LIMIT 1
77            "#,
78            service_name.as_str(),
79            ServiceStatus::Running.as_str(),
80            self.instance_id.as_str(),
81            INSTANCE_DEAD_AFTER_SECS as f64,
82            NodeRole::Worker.as_str()
83        )
84        .fetch_optional(&*self.write_pool)
85        .await?;
86        Ok(row.map(|row| ProcessOwner {
87            instance_id: row.instance_id,
88            advertise_url: row.advertise_url,
89        }))
90    }
91
92    pub async fn delete_stale_endpoints(&self, older_than_secs: i64) -> DatabaseResult<u64> {
93        let result = sqlx::query!(
94            r#"
95            DELETE FROM instance_endpoints
96            WHERE heartbeat_at < CURRENT_TIMESTAMP - make_interval(secs => $1::double precision)
97            "#,
98            older_than_secs as f64
99        )
100        .execute(&*self.write_pool)
101        .await?;
102        Ok(result.rows_affected())
103    }
104}