systemprompt_runtime/
storage.rs

1//! File-storage composition: the profile's `storage:` section resolved to a
2//! [`FileStorage`] backend, with the credential source the GCS backend needs.
3//!
4//! The storage crate stays free of the security crate, so the key-file token
5//! source ([`ServiceAccountTokens`]) lives here and wraps
6//! `systemprompt_security::google::access_token`.
7//!
8//! Copyright (c) systemprompt.io — Business Source License 1.1.
9//! See <https://systemprompt.io> for licensing details.
10
11use std::sync::Arc;
12use std::time::Duration;
13
14use async_trait::async_trait;
15use systemprompt_config::paths::AppPaths;
16use systemprompt_identifiers::{InstanceId, SecretName};
17use systemprompt_manifest::Secrets;
18use systemprompt_manifest::profile::{GcsCredentials, StorageBackend, StorageConfig};
19use systemprompt_security::google::{ServiceAccountKey, access_token};
20use systemprompt_storage::{
21    DEFAULT_METADATA_TOKEN_URL, FileStorageBackend, GcsError, GcsParams, GcsTokenSource,
22    MetadataServerTokens,
23};
24use systemprompt_traits::FileStorage;
25
26use crate::error::{RuntimeError, RuntimeResult};
27
28const GCS_REQUEST_TIMEOUT: Duration = Duration::from_secs(60);
29
30/// Cloud Storage tokens minted from a service-account key held in secrets.
31#[derive(Debug)]
32pub struct ServiceAccountTokens {
33    cache_key: String,
34    key: ServiceAccountKey,
35}
36
37impl ServiceAccountTokens {
38    #[must_use]
39    pub fn new(secret: &SecretName, key: ServiceAccountKey) -> Self {
40        Self {
41            cache_key: format!("storage-gcs:{secret}"),
42            key,
43        }
44    }
45}
46
47#[async_trait]
48impl GcsTokenSource for ServiceAccountTokens {
49    async fn bearer(&self) -> Result<String, GcsError> {
50        access_token(&self.cache_key, &self.key)
51            .await
52            .map_err(|e| GcsError::Token(Box::new(e)))
53    }
54}
55
56pub(crate) async fn init_file_storage(
57    storage: &StorageConfig,
58    app_paths: &AppPaths,
59    instance_id: &InstanceId,
60    secrets: &Secrets,
61) -> RuntimeResult<Arc<dyn FileStorage>> {
62    match storage.backend {
63        StorageBackend::Local => init_local(storage, app_paths, instance_id).await,
64        StorageBackend::Gcs => init_gcs(storage, secrets),
65    }
66}
67
68async fn init_local(
69    storage: &StorageConfig,
70    app_paths: &AppPaths,
71    instance_id: &InstanceId,
72) -> RuntimeResult<Arc<dyn FileStorage>> {
73    let root = app_paths.storage().root();
74    let report = systemprompt_storage::probe_shared_mount(root, instance_id)
75        .await
76        .map_err(|source| RuntimeError::StorageProbe {
77            path: root.to_path_buf(),
78            source,
79        })?;
80    if !report.write_read_ok {
81        return Err(RuntimeError::StorageReadBack {
82            path: root.to_path_buf(),
83        });
84    }
85    match (storage.shared, report.has_siblings()) {
86        (true, false) => tracing::warn!(
87            root = %root.display(),
88            "storage.shared is true but no other replica has marked this root; \
89             it may be a per-node disk"
90        ),
91        (false, true) => tracing::warn!(
92            root = %root.display(),
93            instances = ?report.instances,
94            "storage.shared is false but other replicas have marked this root; \
95             set storage.shared: true if it is a shared mount"
96        ),
97        _ => {},
98    }
99    Ok(systemprompt_storage::build_file_storage(
100        FileStorageBackend::Local {
101            root: root.to_path_buf(),
102        },
103    ))
104}
105
106fn init_gcs(storage: &StorageConfig, secrets: &Secrets) -> RuntimeResult<Arc<dyn FileStorage>> {
107    let http = reqwest::Client::builder()
108        .timeout(GCS_REQUEST_TIMEOUT)
109        .build()
110        .map_err(RuntimeError::StorageHttp)?;
111    let tokens: Arc<dyn GcsTokenSource> = match storage.credentials.clone().unwrap_or_default() {
112        GcsCredentials::WorkloadIdentity => {
113            let endpoint = url::Url::parse(DEFAULT_METADATA_TOKEN_URL)
114                .map_err(RuntimeError::StorageEndpoint)?;
115            Arc::new(MetadataServerTokens::new(endpoint, http.clone()))
116        },
117        GcsCredentials::Secret(name) => {
118            let raw = secrets
119                .get(name.as_str())
120                .ok_or_else(|| RuntimeError::StorageCredentialMissing { name: name.clone() })?;
121            let key = serde_json::from_str::<ServiceAccountKey>(raw).map_err(|source| {
122                RuntimeError::StorageCredential {
123                    name: name.clone(),
124                    source,
125                }
126            })?;
127            Arc::new(ServiceAccountTokens::new(&name, key))
128        },
129    };
130    let bucket = storage
131        .bucket
132        .clone()
133        .filter(|bucket| !bucket.is_empty())
134        .ok_or(RuntimeError::StorageBucketMissing)?;
135    let params = GcsParams::new(bucket, storage.prefix.clone(), storage.public_read)
136        .map_err(RuntimeError::StorageEndpoint)?;
137    tracing::info!(
138        bucket = %params.bucket,
139        prefix = params.prefix.as_deref().unwrap_or(""),
140        "file storage: Cloud Storage backend"
141    );
142    Ok(systemprompt_storage::build_file_storage(
143        FileStorageBackend::Gcs {
144            params: Box::new(params),
145            tokens,
146            http,
147        },
148    ))
149}
systemprompt_runtime/storage.rs

systemprompt_runtime/
storage.rs
