Vakyartha documentation

1use std::io::Write;
2use std::path::PathBuf;
3use std::sync::Arc;
4
5use tokio_util::sync::CancellationToken;
6
7use crate::error::LlmError;
8
9#[derive(Clone)]
10pub(crate) struct ProviderGate {
11 base_url: String,
12 api_key: String,
13 slots: usize,
14 semaphore: Arc<tokio::sync::Semaphore>,
15}
16
17pub(crate) struct ProviderPermit {
18 _permit: tokio::sync::OwnedSemaphorePermit,
19 _lease: ProviderLease,
20}
21
22impl ProviderGate {
23 pub(crate) fn new(base_url: &str, api_key: &str) -> Self {
24 let slots = provider_parallelism(base_url);
25 Self {
26 base_url: base_url.to_owned(),
27 api_key: api_key.to_owned(),
28 slots,
29 semaphore: Arc::new(tokio::sync::Semaphore::new(slots)),
30 }
31 }
32
33 pub(crate) async fn acquire(
34 &self,
35 cancel: &CancellationToken,
36 ) -> Result<ProviderPermit, LlmError> {
37 let permit = tokio::select! {
38 _ = cancel.cancelled() => return Err(LlmError::Aborted { partial: None }),
39 permit = self.semaphore.clone().acquire_owned() => permit
40 .map_err(|_| LlmError::Aborted { partial: None })?,
41 };
42 let lease =
43 ProviderLease::acquire(&self.base_url, &self.api_key, self.slots, cancel).await?;
44 Ok(ProviderPermit {
45 _permit: permit,
46 _lease: lease,
47 })
48 }
49}
50
51struct ProviderLease {
52 path: PathBuf,
53}
54
55impl ProviderLease {
56 async fn acquire(
57 base_url: &str,
58 api_key: &str,
59 slots: usize,
60 cancel: &CancellationToken,
61 ) -> Result<Self, LlmError> {
62 let root = std::env::temp_dir().join("vak-provider-gates");
63 std::fs::create_dir_all(&root)
64 .map_err(|e| LlmError::Network(format!("provider gate setup failed: {e}")))?;
65 if let Ok(entries) = std::fs::read_dir(&root) {
66 for entry in entries.flatten() {
67 let path = entry.path();
68 if path.extension().is_some_and(|ext| ext == "lock") && stale_lease(&path) {
69 let _ = std::fs::remove_file(path);
70 }
71 }
72 }
73 let hash = gate_hash(base_url, api_key);
74 loop {
75 for slot in 0..slots.max(1) {
76 let path = root.join(format!("gate-{hash:016x}-{slot}.lock"));
77 match std::fs::OpenOptions::new()
78 .write(true)
79 .create_new(true)
80 .open(&path)
81 {
82 Ok(mut file) => {
83 let _ = writeln!(file, "{}", std::process::id());
84 return Ok(Self { path });
85 }
86 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
87 if stale_lease(&path) {
88 let _ = std::fs::remove_file(&path);
89 }
90 }
91 Err(error) => {
92 return Err(LlmError::Network(format!(
93 "provider gate acquire failed: {error}"
94 )));
95 }
96 }
97 }
98 tokio::select! {
99 _ = cancel.cancelled() => return Err(LlmError::Aborted { partial: None }),
100 _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => {}
101 }
102 }
103 }
104}
105
106impl Drop for ProviderLease {
107 fn drop(&mut self) {
108 let _ = std::fs::remove_file(&self.path);
109 }
110}
111
112fn provider_parallelism(base_url: &str) -> usize {
113 if let Ok(value) = std::env::var("VAK_PROVIDER_CONCURRENCY")
114 && let Ok(value) = value.parse::<usize>()
115 {
116 return value.max(1);
117 }
118 let local = reqwest::Url::parse(base_url)
119 .ok()
120 .and_then(|url| url.host_str().map(str::to_ascii_lowercase))
121 .is_some_and(|host| {
122 host == "localhost" || host == "127.0.0.1" || host == "::1" || host == "[::1]"
123 });
124 if local { 1 } else { 8 }
125}
126
127fn gate_hash(base_url: &str, api_key: &str) -> u64 {
128 let mut hash = 0xcbf29ce484222325u64;
129 for byte in base_url
130 .as_bytes()
131 .iter()
132 .chain([0].iter())
133 .chain(api_key.as_bytes())
134 {
135 hash ^= u64::from(*byte);
136 hash = hash.wrapping_mul(0x100000001b3);
137 }
138 hash
139}
140
141pub(crate) fn route_identity(name: &str, base_url: &str, api_key: &str) -> String {
142 format!("{name}:{:016x}", gate_hash(base_url, api_key))
143}
144
145/// Stable, non-secret credential identity for frozen route contracts.
146pub fn credential_id(base_url: &str, api_key: &str) -> String {
147 format!("{:016x}", gate_hash(base_url, api_key))
148}
149
150#[cfg(unix)]
151fn stale_lease(path: &std::path::Path) -> bool {
152 let Ok(pid) = std::fs::read_to_string(path)
153 .ok()
154 .and_then(|text| text.trim().parse::<i32>().ok())
155 .ok_or(())
156 else {
157 return true;
158 };
159 std::process::Command::new("kill")
160 .args(["-0", &pid.to_string()])
161 .status()
162 .map(|status| !status.success())
163 .unwrap_or(true)
164}
165
166#[cfg(not(unix))]
167fn stale_lease(_path: &std::path::Path) -> bool {
168 false
169}
170
171#[cfg(test)]
172mod tests {
173 use super::{credential_id, gate_hash, provider_parallelism};
174
175 #[test]
176 fn local_endpoints_default_to_one_slot() {
177 assert_eq!(provider_parallelism("http://localhost:11434/v1"), 1);
178 assert_eq!(provider_parallelism("http://127.0.0.1:1234/v1"), 1);
179 assert_eq!(provider_parallelism("http://[::1]:1234/v1"), 1);
180 }
181
182 #[test]
183 fn remote_endpoints_default_to_eight_slots() {
184 assert_eq!(provider_parallelism("https://api.example.test/v1"), 8);
185 assert_eq!(provider_parallelism("https://localhost.example.test/v1"), 8);
186 assert_eq!(provider_parallelism("https://127.0.0.10/v1"), 8);
187 }
188
189 #[test]
190 fn one_provider_key_shares_capacity_across_models() {
191 assert_eq!(
192 gate_hash("https://api.example.test/v1", "key"),
193 gate_hash("https://api.example.test/v1", "key")
194 );
195 }
196
197 #[test]
198 fn credentials_define_independent_pools() {
199 assert_ne!(
200 gate_hash("https://api.example.test/v1", "key-a"),
201 gate_hash("https://api.example.test/v1", "key-b")
202 );
203 }
204
205 #[test]
206 fn credential_identity_is_stable_and_does_not_contain_secret() {
207 let secret = "provider-secret-value";
208 let identity = credential_id("https://api.example.test/v1", secret);
209 assert_eq!(
210 identity,
211 credential_id("https://api.example.test/v1", secret)
212 );
213 assert!(!identity.contains(secret));
214 assert_eq!(identity.len(), 16);
215 }
216}
217

1use std::io::Write;
2use std::path::PathBuf;
3use std::sync::Arc;
4
5use tokio_util::sync::CancellationToken;
6
7use crate::error::LlmError;
8
9#[derive(Clone)]
10pub(crate) struct ProviderGate {
11 base_url: String,
12 api_key: String,
13 slots: usize,
14 semaphore: Arc<tokio::sync::Semaphore>,
15}
16
17pub(crate) struct ProviderPermit {
18 _permit: tokio::sync::OwnedSemaphorePermit,
19 _lease: ProviderLease,
20}
21
22impl ProviderGate {
23 pub(crate) fn new(base_url: &str, api_key: &str) -> Self {
24 let slots = provider_parallelism(base_url);
25 Self {
26 base_url: base_url.to_owned(),
27 api_key: api_key.to_owned(),
28 slots,
29 semaphore: Arc::new(tokio::sync::Semaphore::new(slots)),
30 }
31 }
32
33 pub(crate) async fn acquire(
34 &self,
35 cancel: &CancellationToken,
36 ) -> Result<ProviderPermit, LlmError> {
37 let permit = tokio::select! {
38 _ = cancel.cancelled() => return Err(LlmError::Aborted { partial: None }),
39 permit = self.semaphore.clone().acquire_owned() => permit
40 .map_err(|_| LlmError::Aborted { partial: None })?,
41 };
42 let lease =
43 ProviderLease::acquire(&self.base_url, &self.api_key, self.slots, cancel).await?;
44 Ok(ProviderPermit {
45 _permit: permit,
46 _lease: lease,
47 })
48 }
49}
50
51struct ProviderLease {
52 path: PathBuf,
53}
54
55impl ProviderLease {
56 async fn acquire(
57 base_url: &str,
58 api_key: &str,
59 slots: usize,
60 cancel: &CancellationToken,
61 ) -> Result<Self, LlmError> {
62 let root = std::env::temp_dir().join("vak-provider-gates");
63 std::fs::create_dir_all(&root)
64 .map_err(|e| LlmError::Network(format!("provider gate setup failed: {e}")))?;
65 if let Ok(entries) = std::fs::read_dir(&root) {
66 for entry in entries.flatten() {
67 let path = entry.path();
68 if path.extension().is_some_and(|ext| ext == "lock") && stale_lease(&path) {
69 let _ = std::fs::remove_file(path);
70 }
71 }
72 }
73 let hash = gate_hash(base_url, api_key);
74 loop {
75 for slot in 0..slots.max(1) {
76 let path = root.join(format!("gate-{hash:016x}-{slot}.lock"));
77 match std::fs::OpenOptions::new()
78 .write(true)
79 .create_new(true)
80 .open(&path)
81 {
82 Ok(mut file) => {
83 let _ = writeln!(file, "{}", std::process::id());
84 return Ok(Self { path });
85 }
86 Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
87 if stale_lease(&path) {
88 let _ = std::fs::remove_file(&path);
89 }
90 }
91 Err(error) => {
92 return Err(LlmError::Network(format!(
93 "provider gate acquire failed: {error}"
94 )));
95 }
96 }
97 }
98 tokio::select! {
99 _ = cancel.cancelled() => return Err(LlmError::Aborted { partial: None }),
100 _ = tokio::time::sleep(std::time::Duration::from_millis(50)) => {}
101 }
102 }
103 }
104}
105
106impl Drop for ProviderLease {
107 fn drop(&mut self) {
108 let _ = std::fs::remove_file(&self.path);
109 }
110}
111
112fn provider_parallelism(base_url: &str) -> usize {
113 if let Ok(value) = std::env::var("VAK_PROVIDER_CONCURRENCY")
114 && let Ok(value) = value.parse::<usize>()
115 {
116 return value.max(1);
117 }
118 let local = reqwest::Url::parse(base_url)
119 .ok()
120 .and_then(|url| url.host_str().map(str::to_ascii_lowercase))
121 .is_some_and(|host| {
122 host == "localhost" || host == "127.0.0.1" || host == "::1" || host == "[::1]"
123 });
124 if local { 1 } else { 8 }
125}
126
127fn gate_hash(base_url: &str, api_key: &str) -> u64 {
128 let mut hash = 0xcbf29ce484222325u64;
129 for byte in base_url
130 .as_bytes()
131 .iter()
132 .chain([0].iter())
133 .chain(api_key.as_bytes())
134 {
135 hash ^= u64::from(*byte);
136 hash = hash.wrapping_mul(0x100000001b3);
137 }
138 hash
139}
140
141pub(crate) fn route_identity(name: &str, base_url: &str, api_key: &str) -> String {
142 format!("{name}:{:016x}", gate_hash(base_url, api_key))
143}
144
145/// Stable, non-secret credential identity for frozen route contracts.
146pub fn credential_id(base_url: &str, api_key: &str) -> String {
147 format!("{:016x}", gate_hash(base_url, api_key))
148}
149
150#[cfg(unix)]
151fn stale_lease(path: &std::path::Path) -> bool {
152 let Ok(pid) = std::fs::read_to_string(path)
153 .ok()
154 .and_then(|text| text.trim().parse::<i32>().ok())
155 .ok_or(())
156 else {
157 return true;
158 };
159 std::process::Command::new("kill")
160 .args(["-0", &pid.to_string()])
161 .status()
162 .map(|status| !status.success())
163 .unwrap_or(true)
164}
165
166#[cfg(not(unix))]
167fn stale_lease(_path: &std::path::Path) -> bool {
168 false
169}
170
171#[cfg(test)]
172mod tests {
173 use super::{credential_id, gate_hash, provider_parallelism};
174
175 #[test]
176 fn local_endpoints_default_to_one_slot() {
177 assert_eq!(provider_parallelism("http://localhost:11434/v1"), 1);
178 assert_eq!(provider_parallelism("http://127.0.0.1:1234/v1"), 1);
179 assert_eq!(provider_parallelism("http://[::1]:1234/v1"), 1);
180 }
181
182 #[test]
183 fn remote_endpoints_default_to_eight_slots() {
184 assert_eq!(provider_parallelism("https://api.example.test/v1"), 8);
185 assert_eq!(provider_parallelism("https://localhost.example.test/v1"), 8);
186 assert_eq!(provider_parallelism("https://127.0.0.10/v1"), 8);
187 }
188
189 #[test]
190 fn one_provider_key_shares_capacity_across_models() {
191 assert_eq!(
192 gate_hash("https://api.example.test/v1", "key"),
193 gate_hash("https://api.example.test/v1", "key")
194 );
195 }
196
197 #[test]
198 fn credentials_define_independent_pools() {
199 assert_ne!(
200 gate_hash("https://api.example.test/v1", "key-a"),
201 gate_hash("https://api.example.test/v1", "key-b")
202 );
203 }
204
205 #[test]
206 fn credential_identity_is_stable_and_does_not_contain_secret() {
207 let secret = "provider-secret-value";
208 let identity = credential_id("https://api.example.test/v1", secret);
209 assert_eq!(
210 identity,
211 credential_id("https://api.example.test/v1", secret)
212 );
213 assert!(!identity.contains(secret));
214 assert_eq!(identity.len(), 16);
215 }
216}
217

Indexing the workspace…