providers_kit/provider.rs
1//! Root provider contract and opaque provider handle.
2
3use std::fmt;
4use std::future::Future;
5use std::num::NonZeroUsize;
6use std::pin::Pin;
7use std::sync::{Arc, Mutex};
8use std::task::{Context, Poll};
9
10use futures_util::Stream;
11use tokio::sync::Notify;
12use tokio::task::JoinHandle;
13
14use crate::KitError;
15use crate::batch::{BatchJob, BatchListRequest, BatchPage, BatchResult};
16use crate::batch_ownership::BatchRegistry;
17use crate::config::{ProviderConfig, ProviderKind};
18use crate::error::Operation;
19use crate::error::ValidationViolation;
20use crate::files::{FileListRequest, FileObject, FilePage, FileUploadRequest};
21use crate::generation::{GenerateEvent, GenerateEventStream, GenerationContinuation};
22use crate::models::{Model, ModelBackend, ModelId, ModelInner};
23use crate::models::{ModelListRequest, ModelPage, ModelSummary};
24use crate::providers::{
25 anthropic::AnthropicBackend,
26 gemini::GeminiBackend,
27 ollama::OllamaBackend,
28 openai::{OpenAiBackend, OpenAiBatchSubmission},
29};
30use crate::runtime::HealthStatus;
31
32/// Public implementation contract for one provider root.
33///
34/// Provider and Model clones share one lifecycle. Shutdown seals model admission,
35/// drains accepted work, and publishes one terminal result to every waiter; cancelling
36/// a waiter cannot cancel owned cleanup. Models have no independent shutdown.
37/// Caller implementations retain their concrete associated children directly.
38///
39/// # Examples
40///
41/// ```
42///
43/// # #[tokio::main(flavor = "current_thread")]
44/// # async fn main() -> Result<(), providers_kit::KitError> {
45/// use providers_kit::{
46/// ModelBackend, ModelId, Provider, ProviderBackend, ProviderConfig,
47/// config::{OllamaConfig, ProviderKind, TransportConfig},
48/// };
49/// let provider = Provider::from_config(ProviderConfig::new(
50/// ProviderKind::Ollama(OllamaConfig::local()),
51/// TransportConfig::default(),
52/// ))?;
53/// let model = provider.model(ModelId::new("exact-model")?)?;
54/// assert_eq!(model.model_id().as_str(), "exact-model");
55/// provider.shutdown().await?;
56///
57/// # Ok(())
58/// # }
59/// ```
60///
61/// # Implement a concrete local hierarchy
62///
63/// This runnable caller implementation is a deliberately limited, in-process text
64/// echo service, **not** a language model or proof of any supplied provider. It
65/// supports one text message with default settings and text-only Sessions. Optional
66/// operations retain their Unsupported defaults. Its work is synchronous and finite:
67/// there is no transport, spawned task or asynchronous resource to drain. A real
68/// network implementation must additionally own admitted work and shutdown cleanup.
69///
70/// The continuation is caller-owned, cloneable, and deliberately has no serialization
71/// contract. Finite streams yield on demand without a forwarding task. Shared closure
72/// seals the root and existing children; consuming Session close releases its local queue.
73///
74/// ```
75/// use std::{collections::VecDeque, sync::{Arc, atomic::{AtomicBool, Ordering}}};
76/// use futures_util::{stream, StreamExt};
77/// use providers_kit::{ProviderBackend, ModelBackend, SessionBackend, ModelId, KitError};
78/// use providers_kit::error::{ErrorKind, ErrorPhase, Operation};
79/// use providers_kit::generation::{ContentPart, FinishReason, GenerateEvent,
80/// GenerateEventStream, GenerateRequest, GenerateResponse, GenerationConfig, Message, Role, Usage};
81/// use providers_kit::session::{SessionConfig, SessionEvent, SessionInput};
82///
83/// #[derive(Clone, Default)]
84/// struct EchoProvider { closed: Arc<AtomicBool> }
85/// #[derive(Clone)]
86/// struct EchoContinuation(Arc<()>);
87/// struct EchoModel { id: ModelId, root: EchoProvider, replay: EchoContinuation }
88/// struct EchoSession { model: ModelId, root: EchoProvider, events: VecDeque<SessionEvent> }
89///
90/// impl EchoProvider {
91/// fn check(&self, operation: Operation) -> Result<(), KitError> {
92/// if self.closed.load(Ordering::SeqCst) {
93/// return Err(KitError::new(ErrorKind::Closed, operation, ErrorPhase::Lifecycle));
94/// }
95/// Ok(())
96/// }
97/// }
98/// impl ProviderBackend for EchoProvider {
99/// type Model = EchoModel;
100/// fn model(&self, id: ModelId) -> Result<EchoModel, KitError> {
101/// self.check(Operation::Generate)?;
102/// if id.as_str() != "local-echo" {
103/// return Err(KitError::new(ErrorKind::Unsupported, Operation::Generate,
104/// ErrorPhase::InputValidation));
105/// }
106/// Ok(EchoModel { id, root: self.clone(), replay: EchoContinuation(Arc::new(())) })
107/// }
108/// async fn shutdown(&self) -> Result<(), KitError> {
109/// // No await: sealing and this resource-free cleanup complete together.
110/// self.closed.store(true, Ordering::SeqCst);
111/// Ok(())
112/// }
113/// }
114/// impl EchoModel {
115/// fn response(&self, request: &GenerateRequest<EchoContinuation>)
116/// -> Result<GenerateResponse<EchoContinuation>, KitError>
117/// {
118/// self.root.check(Operation::Generate)?;
119/// let unsupported = || KitError::new(ErrorKind::Unsupported,
120/// Operation::Generate, ErrorPhase::InputValidation);
121/// if !request.tools().is_empty() || request.reasoning().is_some()
122/// || request.response_format().is_some()
123/// || request.generation() != &GenerationConfig::default()
124/// || request.messages().len() != 1 {
125/// return Err(unsupported());
126/// }
127/// let message = &request.messages()[0];
128/// if message.role() != Role::User { return Err(unsupported()); }
129/// let [ContentPart::Text { text }] = message.content() else {
130/// return Err(unsupported());
131/// };
132/// if text.len() > 64 * 1024 { return Err(unsupported()); }
133/// Ok(GenerateResponse {
134/// model: self.id.clone(),
135/// message: Message::assistant_text(text.clone())?
136/// .with_continuation(self.replay.clone())?,
137/// finish_reason: Some(FinishReason::Stop),
138/// citations: Vec::new(), artifacts: Vec::new(), hosted_tool_outputs: Vec::new(),
139/// usage: Usage::default(), // No token accounting is invented.
140/// })
141/// }
142/// }
143/// impl ModelBackend for EchoModel {
144/// type Continuation = EchoContinuation;
145/// type Session = EchoSession;
146/// fn model_id(&self) -> &ModelId { &self.id }
147/// async fn generate(&self, request: &GenerateRequest<EchoContinuation>)
148/// -> Result<GenerateResponse<EchoContinuation>, KitError> { self.response(request) }
149/// async fn generate_stream(&self, request: &GenerateRequest<EchoContinuation>)
150/// -> Result<GenerateEventStream<EchoContinuation>, KitError>
151/// {
152/// let response = self.response(request)?;
153/// let text: String = response.message.content().iter().filter_map(|part| match part {
154/// ContentPart::Text { text } => Some(text.as_str()), _ => None,
155/// }).collect();
156/// Ok(GenerateEventStream::new(stream::iter([
157/// Ok(GenerateEvent::Started { model: self.id.clone() }),
158/// Ok(GenerateEvent::ContentDelta { delta: text }),
159/// Ok(GenerateEvent::Completed(Box::new(response))),
160/// ])))
161/// }
162/// async fn open_session(&self, config: SessionConfig) -> Result<EchoSession, KitError> {
163/// self.root.check(Operation::OpenSession)?;
164/// if config != SessionConfig::default() {
165/// return Err(KitError::new(ErrorKind::Unsupported, Operation::OpenSession,
166/// ErrorPhase::InputValidation));
167/// }
168/// Ok(EchoSession { model: self.id.clone(), root: self.root.clone(), events: VecDeque::new() })
169/// }
170/// }
171/// impl SessionBackend for EchoSession {
172/// fn id(&self) -> &str { "local-echo-session" }
173/// fn model_id(&self) -> &ModelId { &self.model }
174/// async fn send(&mut self, input: SessionInput) -> Result<(), KitError> {
175/// self.root.check(Operation::SessionSend)?;
176/// if !self.events.is_empty() {
177/// return Err(KitError::new(ErrorKind::InvalidState, Operation::SessionSend,
178/// ErrorPhase::Lifecycle));
179/// }
180/// match input {
181/// SessionInput::Text { text } if !text.trim().is_empty() && text.len() <= 64 * 1024 => {
182/// self.events.push_back(SessionEvent::TextDelta { delta: text });
183/// self.events.push_back(SessionEvent::Completed);
184/// Ok(())
185/// }
186/// _ => Err(KitError::new(ErrorKind::Unsupported, Operation::SessionSend,
187/// ErrorPhase::InputValidation)),
188/// }
189/// }
190/// async fn recv(&mut self) -> Result<SessionEvent, KitError> {
191/// self.root.check(Operation::SessionReceive)?;
192/// self.events.pop_front().ok_or_else(|| KitError::new(ErrorKind::InvalidState,
193/// Operation::SessionReceive, ErrorPhase::Lifecycle))
194/// }
195/// async fn close(self) -> Result<(), KitError> { Ok(()) }
196/// }
197///
198/// # #[tokio::main(flavor = "current_thread")]
199/// # async fn main() -> Result<(), KitError> {
200/// let provider = EchoProvider::default();
201/// let model = provider.model(ModelId::new("local-echo")?)?;
202/// let request = GenerateRequest::new(vec![Message::user_text("hello")?])?;
203/// let response = model.generate(&request).await?;
204/// assert!(matches!(response.message.content(), [ContentPart::Text { text }] if text == "hello"));
205/// assert!(Arc::ptr_eq(&response.message.continuation().unwrap().0, &model.replay.0));
206/// let mut events = model.generate_stream(&request).await?;
207/// assert!(matches!(events.next().await.transpose()?, Some(GenerateEvent::Started { .. })));
208/// assert!(
209/// matches!(events.next().await.transpose()?, Some(GenerateEvent::ContentDelta { delta }) if delta == "hello")
210/// );
211/// assert!(matches!(events.next().await.transpose()?, Some(GenerateEvent::Completed(_))));
212/// assert!(events.next().await.is_none());
213/// let mut session = model.open_session(SessionConfig::default()).await?;
214/// session.send(SessionInput::Text { text: "hello again".into() }).await?;
215/// assert!(
216/// matches!(session.recv().await?, SessionEvent::TextDelta { delta } if delta == "hello again")
217/// );
218/// assert_eq!(session.recv().await?, SessionEvent::Completed);
219/// session.close().await?;
220/// provider.shutdown().await?;
221/// assert_eq!(model.generate(&request).await.err().unwrap().kind(), ErrorKind::Closed);
222/// provider.shutdown().await?; // Replays the same resource-free terminal result.
223/// # Ok(())
224/// # }
225/// ```
226///
227pub trait ProviderBackend: Send + Sync + 'static {
228 /// Exact model child returned by this provider.
229 type Model: ModelBackend;
230
231 /// Selects one exact provider-local model.
232 ///
233 /// # Errors
234 ///
235 /// Returns an error when root admission is closed or the selected implementation cannot construct this exact model. Selection never silently falls back to another model.
236 ///
237 /// See [`crate::ProviderBackend`] for a complete example.
238 ///
239 fn model(&self, id: ModelId) -> Result<Self::Model, KitError>;
240
241 /// Lists provider-visible models.
242 ///
243 /// This example requires a real configured service, valid credentials, and any exact
244 /// resource ids shown.
245 /// The default trait implementation returns Unsupported.
246 ///
247 /// # Errors
248 ///
249 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
250 ///
251 /// # Examples
252 ///
253 /// ```no_run
254 ///
255 /// # #[tokio::main(flavor = "current_thread")]
256 /// # async fn main() -> Result<(), providers_kit::KitError> {
257 /// use providers_kit::{
258 /// Provider, ProviderBackend, ProviderConfig,
259 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
260 /// };
261 /// let provider = Provider::from_config(ProviderConfig::new(
262 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
263 /// TransportConfig::default(),
264 /// ))?;
265 /// let result = provider.list_models(providers_kit::models::ModelListRequest::default()).await;
266 /// let cleanup = provider.shutdown().await;
267 /// if let Err(error) = &cleanup {
268 /// eprintln!("Cleanup failed: {error}");
269 /// }
270 /// let page = result?;
271 /// cleanup?;
272 /// for model in page.models {
273 /// println!("{}", model.id.as_str());
274 /// }
275 ///
276 /// # Ok(())
277 /// # }
278 /// ```
279 ///
280 fn list_models(
281 &self,
282 _request: ModelListRequest,
283 ) -> impl Future<Output = Result<ModelPage, KitError>> + Send + '_ {
284 async {
285 Err(KitError::unsupported(
286 ValidationViolation::Unsupported {
287 field: "models.list",
288 rule: "operation is not implemented",
289 },
290 "provider model discovery",
291 )
292 .at_operation(Operation::ListModels))
293 }
294 }
295
296 /// Fetches descriptive metadata for one exact provider-local model.
297 ///
298 /// This example requires a real configured service, valid credentials, and any exact
299 /// resource ids shown.
300 /// The default trait implementation returns Unsupported.
301 ///
302 /// # Errors
303 ///
304 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
305 ///
306 /// # Examples
307 ///
308 /// ```no_run
309 ///
310 /// # #[tokio::main(flavor = "current_thread")]
311 /// # async fn main() -> Result<(), providers_kit::KitError> {
312 /// use providers_kit::{
313 /// Provider, ProviderBackend, ProviderConfig,
314 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
315 /// };
316 /// let provider = Provider::from_config(ProviderConfig::new(
317 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
318 /// TransportConfig::default(),
319 /// ))?;
320 /// let result = provider.model_info(providers_kit::ModelId::new("your-exact-model-id")?).await;
321 /// let cleanup = provider.shutdown().await;
322 /// if let Err(error) = &cleanup {
323 /// eprintln!("Cleanup failed: {error}");
324 /// }
325 /// let _value = result?;
326 /// cleanup?;
327 ///
328 /// # Ok(())
329 /// # }
330 /// ```
331 ///
332 fn model_info(
333 &self,
334 _id: ModelId,
335 ) -> impl Future<Output = Result<ModelSummary, KitError>> + Send + '_ {
336 async {
337 Err(KitError::unsupported(
338 ValidationViolation::Unsupported {
339 field: "models.info",
340 rule: "operation is not implemented",
341 },
342 "provider model metadata",
343 )
344 .at_operation(Operation::ModelInfo))
345 }
346 }
347
348 /// Uploads one provider-managed file.
349 ///
350 /// This example requires a real configured service, valid credentials, and any exact
351 /// resource ids shown. Only delete or cancel resources the caller owns.
352 /// The default trait implementation returns Unsupported.
353 ///
354 /// # Errors
355 ///
356 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
357 ///
358 /// # Examples
359 ///
360 /// ```no_run
361 ///
362 /// # #[tokio::main(flavor = "current_thread")]
363 /// # async fn main() -> Result<(), providers_kit::KitError> {
364 /// use providers_kit::{
365 /// Provider, ProviderBackend, ProviderConfig,
366 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
367 /// };
368 /// let provider = Provider::from_config(ProviderConfig::new(
369 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
370 /// TransportConfig::default(),
371 /// ))?;
372 /// let result = provider
373 /// .upload_file(providers_kit::files::FileUploadRequest::new(
374 /// "notes.txt",
375 /// "text/plain",
376 /// b"notes".to_vec(),
377 /// )?)
378 /// .await;
379 /// let deletion = match &result {
380 /// Ok(file) => provider.delete_file(file.id.clone()).await,
381 /// Err(_) => Ok(()),
382 /// };
383 /// let cleanup = provider.shutdown().await;
384 /// if let Err(error) = &cleanup {
385 /// eprintln!("Cleanup failed: {error}");
386 /// }
387 /// let _value = result?;
388 /// deletion?;
389 /// cleanup?;
390 ///
391 /// # Ok(())
392 /// # }
393 /// ```
394 ///
395 fn upload_file(
396 &self,
397 _request: FileUploadRequest,
398 ) -> impl Future<Output = Result<FileObject, KitError>> + Send + '_ {
399 async {
400 Err(KitError::unsupported(
401 ValidationViolation::Unsupported {
402 field: "files.upload",
403 rule: "operation is not implemented",
404 },
405 "provider file upload",
406 )
407 .at_operation(Operation::UploadFile))
408 }
409 }
410
411 /// Lists provider-managed files.
412 ///
413 /// This example requires a real configured service, valid credentials, and any exact
414 /// resource ids shown. Only delete or cancel resources the caller owns.
415 /// The default trait implementation returns Unsupported.
416 ///
417 /// # Errors
418 ///
419 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
420 ///
421 /// # Examples
422 ///
423 /// ```no_run
424 ///
425 /// # #[tokio::main(flavor = "current_thread")]
426 /// # async fn main() -> Result<(), providers_kit::KitError> {
427 /// use providers_kit::{
428 /// Provider, ProviderBackend, ProviderConfig,
429 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
430 /// };
431 /// let provider = Provider::from_config(ProviderConfig::new(
432 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
433 /// TransportConfig::default(),
434 /// ))?;
435 /// let result = provider.list_files(providers_kit::files::FileListRequest::default()).await;
436 /// let cleanup = provider.shutdown().await;
437 /// if let Err(error) = &cleanup {
438 /// eprintln!("Cleanup failed: {error}");
439 /// }
440 /// let _value = result?;
441 /// cleanup?;
442 ///
443 /// # Ok(())
444 /// # }
445 /// ```
446 ///
447 fn list_files(
448 &self,
449 _request: FileListRequest,
450 ) -> impl Future<Output = Result<FilePage, KitError>> + Send + '_ {
451 async {
452 Err(KitError::unsupported(
453 ValidationViolation::Unsupported {
454 field: "files.list",
455 rule: "operation is not implemented",
456 },
457 "provider file listing",
458 )
459 .at_operation(Operation::ListFiles))
460 }
461 }
462
463 /// Fetches metadata for one provider-managed file.
464 ///
465 /// This example requires a real configured service, valid credentials, and any exact
466 /// resource ids shown. Only delete or cancel resources the caller owns.
467 /// The default trait implementation returns Unsupported.
468 ///
469 /// # Errors
470 ///
471 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
472 ///
473 /// # Examples
474 ///
475 /// ```no_run
476 ///
477 /// # #[tokio::main(flavor = "current_thread")]
478 /// # async fn main() -> Result<(), providers_kit::KitError> {
479 /// use providers_kit::{
480 /// Provider, ProviderBackend, ProviderConfig,
481 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
482 /// };
483 /// let provider = Provider::from_config(ProviderConfig::new(
484 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
485 /// TransportConfig::default(),
486 /// ))?;
487 /// let result = provider.file_info("owned-file-id".into()).await;
488 /// let cleanup = provider.shutdown().await;
489 /// if let Err(error) = &cleanup {
490 /// eprintln!("Cleanup failed: {error}");
491 /// }
492 /// let _value = result?;
493 /// cleanup?;
494 ///
495 /// # Ok(())
496 /// # }
497 /// ```
498 ///
499 fn file_info(
500 &self,
501 _id: String,
502 ) -> impl Future<Output = Result<FileObject, KitError>> + Send + '_ {
503 async {
504 Err(KitError::unsupported(
505 ValidationViolation::Unsupported {
506 field: "files.info",
507 rule: "operation is not implemented",
508 },
509 "provider file metadata",
510 )
511 .at_operation(Operation::FileInfo))
512 }
513 }
514
515 /// Deletes one provider-managed file.
516 ///
517 /// This example requires a real configured service, valid credentials, and any exact
518 /// resource ids shown. Only delete or cancel resources the caller owns.
519 /// The default trait implementation returns Unsupported.
520 ///
521 /// # Errors
522 ///
523 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
524 ///
525 /// # Examples
526 ///
527 /// ```no_run
528 ///
529 /// # #[tokio::main(flavor = "current_thread")]
530 /// # async fn main() -> Result<(), providers_kit::KitError> {
531 /// use providers_kit::{
532 /// Provider, ProviderBackend, ProviderConfig,
533 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
534 /// };
535 /// let provider = Provider::from_config(ProviderConfig::new(
536 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
537 /// TransportConfig::default(),
538 /// ))?;
539 /// let result = provider.delete_file("owned-file-id".into()).await;
540 /// let cleanup = provider.shutdown().await;
541 /// if let Err(error) = &cleanup {
542 /// eprintln!("Cleanup failed: {error}");
543 /// }
544 /// let _value = result?;
545 /// cleanup?;
546 ///
547 /// # Ok(())
548 /// # }
549 /// ```
550 ///
551 fn delete_file(&self, _id: String) -> impl Future<Output = Result<(), KitError>> + Send + '_ {
552 async {
553 Err(KitError::unsupported(
554 ValidationViolation::Unsupported {
555 field: "files.delete",
556 rule: "operation is not implemented",
557 },
558 "provider file deletion",
559 )
560 .at_operation(Operation::DeleteFile))
561 }
562 }
563
564 /// Downloads one provider-managed file under the caller's byte limit.
565 ///
566 /// This example requires a real configured service, valid credentials, and any exact
567 /// resource ids shown. Only delete or cancel resources the caller owns.
568 /// The default trait implementation returns Unsupported.
569 ///
570 /// # Errors
571 ///
572 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
573 ///
574 /// # Examples
575 ///
576 /// ```no_run
577 ///
578 /// # #[tokio::main(flavor = "current_thread")]
579 /// # async fn main() -> Result<(), providers_kit::KitError> {
580 /// use providers_kit::{
581 /// Provider, ProviderBackend, ProviderConfig,
582 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
583 /// };
584 /// let provider = Provider::from_config(ProviderConfig::new(
585 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
586 /// TransportConfig::default(),
587 /// ))?;
588 /// let result = provider
589 /// .download_file(
590 /// "owned-file-id".into(),
591 /// std::num::NonZeroUsize::new(1024 * 1024).unwrap(),
592 /// )
593 /// .await;
594 /// let cleanup = provider.shutdown().await;
595 /// if let Err(error) = &cleanup {
596 /// eprintln!("Cleanup failed: {error}");
597 /// }
598 /// let _value = result?;
599 /// cleanup?;
600 ///
601 /// # Ok(())
602 /// # }
603 /// ```
604 ///
605 fn download_file(
606 &self,
607 _id: String,
608 _max_bytes: NonZeroUsize,
609 ) -> impl Future<Output = Result<Vec<u8>, KitError>> + Send + '_ {
610 async {
611 Err(KitError::unsupported(
612 ValidationViolation::Unsupported {
613 field: "files.download",
614 rule: "operation is not implemented",
615 },
616 "provider file download",
617 )
618 .at_operation(Operation::DownloadFile))
619 }
620 }
621
622 /// Fetches one batch job.
623 ///
624 /// Supplied OpenAI status reads opt into the configured bounded transport
625 /// retries; every attempt shares the original absolute request deadline.
626 /// This does not retry decoding or successful-response body failures. The
627 /// lookup used by `batch_results` is excluded. See the [batch workflow](crate::batch).
628 ///
629 /// This example requires a real configured service, valid credentials, and any exact
630 /// resource ids shown. Only delete or cancel resources the caller owns.
631 /// The default trait implementation returns Unsupported.
632 ///
633 /// # Errors
634 ///
635 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
636 ///
637 /// # Examples
638 ///
639 /// ```no_run
640 ///
641 /// # #[tokio::main(flavor = "current_thread")]
642 /// # async fn main() -> Result<(), providers_kit::KitError> {
643 /// use providers_kit::{
644 /// Provider, ProviderBackend, ProviderConfig,
645 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
646 /// };
647 /// let provider = Provider::from_config(ProviderConfig::new(
648 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
649 /// TransportConfig::default(),
650 /// ))?;
651 /// let result = provider.batch_status("owned-batch-id".into()).await;
652 /// let cleanup = provider.shutdown().await;
653 /// if let Err(error) = &cleanup {
654 /// eprintln!("Cleanup failed: {error}");
655 /// }
656 /// let _value = result?;
657 /// cleanup?;
658 ///
659 /// # Ok(())
660 /// # }
661 /// ```
662 ///
663 fn batch_status(
664 &self,
665 _id: String,
666 ) -> impl Future<Output = Result<BatchJob, KitError>> + Send + '_ {
667 async {
668 Err(KitError::unsupported(
669 ValidationViolation::Unsupported {
670 field: "batch.status",
671 rule: "operation is not implemented",
672 },
673 "provider batch status",
674 )
675 .at_operation(Operation::BatchStatus))
676 }
677 }
678
679 /// Requests cancellation of one batch job.
680 ///
681 /// Acknowledgement is not terminality: the returned job may still be running.
682 /// Observe its status until terminal before cleaning up owned input files; see
683 /// the [bounded batch workflow](crate::batch).
684 ///
685 /// This example requires a real configured service, valid credentials, and any exact
686 /// resource ids shown. Only delete or cancel resources the caller owns.
687 /// The default trait implementation returns Unsupported.
688 ///
689 /// # Errors
690 ///
691 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
692 ///
693 /// # Examples
694 ///
695 /// ```no_run
696 ///
697 /// # #[tokio::main(flavor = "current_thread")]
698 /// # async fn main() -> Result<(), providers_kit::KitError> {
699 /// use providers_kit::{
700 /// Provider, ProviderBackend, ProviderConfig,
701 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
702 /// };
703 /// let provider = Provider::from_config(ProviderConfig::new(
704 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
705 /// TransportConfig::default(),
706 /// ))?;
707 /// let result = provider.cancel_batch("owned-batch-id".into()).await;
708 /// let cleanup = provider.shutdown().await;
709 /// if let Err(error) = &cleanup {
710 /// eprintln!("Cleanup failed: {error}");
711 /// }
712 /// let _value = result?;
713 /// cleanup?;
714 ///
715 /// # Ok(())
716 /// # }
717 /// ```
718 ///
719 fn cancel_batch(
720 &self,
721 _id: String,
722 ) -> impl Future<Output = Result<BatchJob, KitError>> + Send + '_ {
723 async {
724 Err(KitError::unsupported(
725 ValidationViolation::Unsupported {
726 field: "batch.cancel",
727 rule: "operation is not implemented",
728 },
729 "provider batch cancellation",
730 )
731 .at_operation(Operation::CancelBatch))
732 }
733 }
734
735 /// Downloads and normalizes batch result rows under the caller's byte limit.
736 ///
737 /// This example requires a real configured service, valid credentials, and any exact
738 /// resource ids shown. Only delete or cancel resources the caller owns.
739 /// The default trait implementation returns Unsupported.
740 /// The byte bound covers the downloaded result; each terminal row must be inspected
741 /// independently. Row failures are result data, while download/normalization failures
742 /// are outer errors.
743 ///
744 /// # Errors
745 ///
746 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
747 ///
748 /// # Examples
749 ///
750 /// ```no_run
751 ///
752 /// # #[tokio::main(flavor = "current_thread")]
753 /// # async fn main() -> Result<(), providers_kit::KitError> {
754 /// use providers_kit::{
755 /// Provider, ProviderBackend, ProviderConfig,
756 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
757 /// };
758 /// let provider = Provider::from_config(ProviderConfig::new(
759 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
760 /// TransportConfig::default(),
761 /// ))?;
762 /// let result = provider
763 /// .batch_results(
764 /// "owned-batch-id".into(),
765 /// std::num::NonZeroUsize::new(1024 * 1024).unwrap(),
766 /// )
767 /// .await;
768 /// let cleanup = provider.shutdown().await;
769 /// if let Err(error) = &cleanup {
770 /// eprintln!("Cleanup failed: {error}");
771 /// }
772 /// let _value = result?;
773 /// cleanup?;
774 ///
775 /// # Ok(())
776 /// # }
777 /// ```
778 ///
779 fn batch_results(
780 &self,
781 _id: String,
782 _max_bytes: NonZeroUsize,
783 ) -> impl Future<
784 Output = Result<Vec<BatchResult<<Self::Model as ModelBackend>::Continuation>>, KitError>,
785 > + Send
786 + '_ {
787 async {
788 Err(KitError::unsupported(
789 ValidationViolation::Unsupported {
790 field: "batch.results",
791 rule: "operation is not implemented",
792 },
793 "provider batch results",
794 )
795 .at_operation(Operation::BatchResults))
796 }
797 }
798
799 /// Lists provider batch jobs.
800 ///
801 /// This example requires a real configured service, valid credentials, and any exact
802 /// resource ids shown. Only delete or cancel resources the caller owns.
803 /// The default trait implementation returns Unsupported.
804 ///
805 /// # Errors
806 ///
807 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
808 ///
809 /// # Examples
810 ///
811 /// ```no_run
812 ///
813 /// # #[tokio::main(flavor = "current_thread")]
814 /// # async fn main() -> Result<(), providers_kit::KitError> {
815 /// use providers_kit::{
816 /// Provider, ProviderBackend, ProviderConfig,
817 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
818 /// };
819 /// let provider = Provider::from_config(ProviderConfig::new(
820 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
821 /// TransportConfig::default(),
822 /// ))?;
823 /// let result = provider.list_batches(providers_kit::batch::BatchListRequest::default()).await;
824 /// let cleanup = provider.shutdown().await;
825 /// if let Err(error) = &cleanup {
826 /// eprintln!("Cleanup failed: {error}");
827 /// }
828 /// let _value = result?;
829 /// cleanup?;
830 ///
831 /// # Ok(())
832 /// # }
833 /// ```
834 ///
835 fn list_batches(
836 &self,
837 _request: BatchListRequest,
838 ) -> impl Future<Output = Result<BatchPage, KitError>> + Send + '_ {
839 async {
840 Err(KitError::unsupported(
841 ValidationViolation::Unsupported {
842 field: "batch.list",
843 rule: "operation is not implemented",
844 },
845 "provider batch listing",
846 )
847 .at_operation(Operation::ListBatches))
848 }
849 }
850
851 /// Probes provider runtime health.
852 ///
853 /// This example requires a real configured service, valid credentials, and any exact
854 /// resource ids shown.
855 /// The default trait implementation returns Unsupported.
856 ///
857 /// # Errors
858 ///
859 /// Returns [`KitError`] classified as Unsupported when unavailable, Closed after sealing, or the owning input, transport, deadline, service-rejection, or response-validation category. Inspect `kind()` and `context()` rather than diagnostic prose.
860 ///
861 /// # Examples
862 ///
863 /// ```no_run
864 ///
865 /// # #[tokio::main(flavor = "current_thread")]
866 /// # async fn main() -> Result<(), providers_kit::KitError> {
867 /// use providers_kit::{
868 /// Provider, ProviderBackend, ProviderConfig,
869 /// config::{ApiKey, OpenAiConfig, ProviderKind, TransportConfig},
870 /// };
871 /// let provider = Provider::from_config(ProviderConfig::new(
872 /// ProviderKind::OpenAi(OpenAiConfig::new(ApiKey::new("replace-with-your-key")?)),
873 /// TransportConfig::default(),
874 /// ))?;
875 /// let result = provider.health().await;
876 /// let cleanup = provider.shutdown().await;
877 /// if let Err(error) = &cleanup {
878 /// eprintln!("Cleanup failed: {error}");
879 /// }
880 /// let _value = result?;
881 /// cleanup?;
882 ///
883 /// # Ok(())
884 /// # }
885 /// ```
886 ///
887 fn health(&self) -> impl Future<Output = Result<HealthStatus, KitError>> + Send + '_ {
888 async {
889 Err(KitError::unsupported(
890 ValidationViolation::Unsupported {
891 field: "runtime.health",
892 rule: "operation is not implemented",
893 },
894 "provider runtime health",
895 )
896 .at_operation(Operation::Health))
897 }
898 }
899
900 /// Shuts down backend-owned resources.
901 ///
902 /// Retain cleanup independently of each borrowed waiter. Repeated and concurrent
903 /// calls replay the same terminal result; last-handle drop is not an async cleanup API.
904 ///
905 /// # Errors
906 ///
907 /// Returns the shared backend cleanup failure, if any. A cancelled waiter does not reset shutdown or abandon cleanup.
908 ///
909 /// # Examples
910 ///
911 /// ```
912 ///
913 /// # #[tokio::main(flavor = "current_thread")]
914 /// # async fn main() -> Result<(), providers_kit::KitError> {
915 /// use providers_kit::{
916 /// ModelBackend, ModelId, Provider, ProviderBackend, ProviderConfig,
917 /// config::{OllamaConfig, ProviderKind, TransportConfig},
918 /// };
919 /// let provider = Provider::from_config(ProviderConfig::new(
920 /// ProviderKind::Ollama(OllamaConfig::local()),
921 /// TransportConfig::default(),
922 /// ))?;
923 /// let clone = provider.clone();
924 /// provider.shutdown().await?;
925 /// clone.shutdown().await?;
926 /// assert!(clone.model(ModelId::new("exact-model")?).is_err());
927 ///
928 /// # Ok(())
929 /// # }
930 /// ```
931 ///
932 fn shutdown(&self) -> impl Future<Output = Result<(), KitError>> + Send + '_;
933}
934
935enum ProviderInner {
936 OpenAi(OpenAiBackend),
937 Anthropic(AnthropicBackend),
938 Gemini(GeminiBackend),
939 Ollama(OllamaBackend),
940}
941
942impl ProviderBackend for ProviderInner {
943 type Model = ModelInner;
944
945 fn model(&self, id: ModelId) -> Result<ModelInner, KitError> {
946 match self {
947 Self::OpenAi(backend) => backend.model(id).map(ModelInner::OpenAi),
948 Self::Anthropic(backend) => backend.model(id).map(ModelInner::Anthropic),
949 Self::Gemini(backend) => backend.model(id).map(ModelInner::Gemini),
950 Self::Ollama(backend) => backend.model(id).map(ModelInner::Ollama),
951 }
952 }
953
954 async fn list_models(&self, request: ModelListRequest) -> Result<ModelPage, KitError> {
955 match self {
956 Self::OpenAi(backend) => backend.list_models(request).await,
957 Self::Anthropic(backend) => backend.list_models(request).await,
958 Self::Gemini(backend) => backend.list_models(request).await,
959 Self::Ollama(backend) => backend.list_models(request).await,
960 }
961 }
962
963 async fn model_info(&self, id: ModelId) -> Result<ModelSummary, KitError> {
964 match self {
965 Self::OpenAi(backend) => backend.model_info(id).await,
966 Self::Anthropic(backend) => backend.model_info(id).await,
967 Self::Gemini(backend) => backend.model_info(id).await,
968 Self::Ollama(backend) => backend.model_info(id).await,
969 }
970 }
971
972 async fn upload_file(&self, request: FileUploadRequest) -> Result<FileObject, KitError> {
973 match self {
974 Self::OpenAi(backend) => backend.upload_file(request).await,
975 Self::Anthropic(backend) => backend.upload_file(request).await,
976 Self::Gemini(backend) => backend.upload_file(request).await,
977 Self::Ollama(backend) => backend.upload_file(request).await,
978 }
979 }
980
981 async fn list_files(&self, request: FileListRequest) -> Result<FilePage, KitError> {
982 match self {
983 Self::OpenAi(backend) => backend.list_files(request).await,
984 Self::Anthropic(backend) => backend.list_files(request).await,
985 Self::Gemini(backend) => backend.list_files(request).await,
986 Self::Ollama(backend) => backend.list_files(request).await,
987 }
988 }
989
990 async fn file_info(&self, id: String) -> Result<FileObject, KitError> {
991 match self {
992 Self::OpenAi(backend) => backend.file_info(id).await,
993 Self::Anthropic(backend) => backend.file_info(id).await,
994 Self::Gemini(backend) => backend.file_info(id).await,
995 Self::Ollama(backend) => backend.file_info(id).await,
996 }
997 }
998
999 async fn delete_file(&self, id: String) -> Result<(), KitError> {
1000 match self {
1001 Self::OpenAi(backend) => backend.delete_file(id).await,
1002 Self::Anthropic(backend) => backend.delete_file(id).await,
1003 Self::Gemini(backend) => backend.delete_file(id).await,
1004 Self::Ollama(backend) => backend.delete_file(id).await,
1005 }
1006 }
1007
1008 async fn download_file(
1009 &self,
1010 id: String,
1011 max_bytes: NonZeroUsize,
1012 ) -> Result<Vec<u8>, KitError> {
1013 match self {
1014 Self::OpenAi(backend) => backend.download_file(id, max_bytes).await,
1015 Self::Anthropic(backend) => backend.download_file(id, max_bytes).await,
1016 Self::Gemini(backend) => backend.download_file(id, max_bytes).await,
1017 Self::Ollama(backend) => backend.download_file(id, max_bytes).await,
1018 }
1019 }
1020
1021 async fn batch_status(&self, id: String) -> Result<BatchJob, KitError> {
1022 match self {
1023 Self::OpenAi(backend) => backend.batch_status(id).await,
1024 Self::Anthropic(backend) => backend.batch_status(id).await,
1025 Self::Gemini(backend) => backend.batch_status(id).await,
1026 Self::Ollama(backend) => backend.batch_status(id).await,
1027 }
1028 }
1029
1030 async fn cancel_batch(&self, id: String) -> Result<BatchJob, KitError> {
1031 match self {
1032 Self::OpenAi(backend) => backend.cancel_batch(id).await,
1033 Self::Anthropic(backend) => backend.cancel_batch(id).await,
1034 Self::Gemini(backend) => backend.cancel_batch(id).await,
1035 Self::Ollama(backend) => backend.cancel_batch(id).await,
1036 }
1037 }
1038
1039 async fn batch_results(
1040 &self,
1041 id: String,
1042 max_bytes: NonZeroUsize,
1043 ) -> Result<Vec<BatchResult<GenerationContinuation>>, KitError> {
1044 match self {
1045 Self::OpenAi(backend) => backend.batch_results(id, max_bytes).await,
1046 Self::Anthropic(backend) => backend.batch_results(id, max_bytes).await,
1047 Self::Gemini(backend) => backend.batch_results(id, max_bytes).await,
1048 Self::Ollama(backend) => backend.batch_results(id, max_bytes).await,
1049 }
1050 }
1051
1052 async fn list_batches(&self, request: BatchListRequest) -> Result<BatchPage, KitError> {
1053 match self {
1054 Self::OpenAi(backend) => backend.list_batches(request).await,
1055 Self::Anthropic(backend) => backend.list_batches(request).await,
1056 Self::Gemini(backend) => backend.list_batches(request).await,
1057 Self::Ollama(backend) => backend.list_batches(request).await,
1058 }
1059 }
1060
1061 async fn health(&self) -> Result<HealthStatus, KitError> {
1062 match self {
1063 Self::OpenAi(backend) => backend.health().await,
1064 Self::Anthropic(backend) => backend.health().await,
1065 Self::Gemini(backend) => backend.health().await,
1066 Self::Ollama(backend) => backend.health().await,
1067 }
1068 }
1069
1070 async fn shutdown(&self) -> Result<(), KitError> {
1071 match self {
1072 Self::OpenAi(backend) => backend.shutdown().await,
1073 Self::Anthropic(backend) => backend.shutdown().await,
1074 Self::Gemini(backend) => backend.shutdown().await,
1075 Self::Ollama(backend) => backend.shutdown().await,
1076 }
1077 }
1078}
1079
1080/// Cloneable opaque provider handle.
1081///
1082/// Provider and Model clones share one lifecycle. Shutdown seals model admission,
1083/// drains accepted work, and publishes one terminal result to every waiter; cancelling
1084/// a waiter cannot cancel owned cleanup. Models have no independent shutdown.
1085/// Caller implementations retain their concrete associated children directly.
1086///
1087/// See [`crate::ProviderBackend`] for a complete example.
1088///
1089#[derive(Clone)]
1090pub struct Provider {
1091 inner: Arc<ProviderRuntime>,
1092}
1093
1094pub(crate) struct ProviderRuntime {
1095 backend: ProviderInner,
1096 state: Mutex<ProviderLifecycleState>,
1097 batches: Mutex<BatchRegistry>,
1098 changed: Notify,
1099}
1100
1101enum ProviderLifecycleState {
1102 Open {
1103 admitted: usize,
1104 },
1105 Closing {
1106 admitted: usize,
1107 task: Option<JoinHandle<Result<(), KitError>>>,
1108 },
1109 Closed {
1110 terminal: Result<(), KitError>,
1111 },
1112}
1113
1114pub(crate) struct ProviderAdmission {
1115 runtime: Arc<ProviderRuntime>,
1116 batch_bytes: Option<usize>,
1117}
1118
1119struct ProviderAdmittedStream {
1120 inner: GenerateEventStream,
1121 admission: Option<ProviderAdmission>,
1122}
1123
1124struct ShutdownTaskLease<'a> {
1125 runtime: &'a ProviderRuntime,
1126 task: Option<JoinHandle<Result<(), KitError>>>,
1127}
1128
1129impl ProviderRuntime {
1130 fn new(backend: ProviderInner) -> Self {
1131 Self {
1132 backend,
1133 state: Mutex::new(ProviderLifecycleState::Open { admitted: 0 }),
1134 batches: Mutex::new(BatchRegistry::new()),
1135 changed: Notify::new(),
1136 }
1137 }
1138
1139 pub(crate) fn admit(
1140 self: &Arc<Self>,
1141 operation: Operation,
1142 ) -> Result<ProviderAdmission, KitError> {
1143 let mut state = self.state();
1144 match &mut *state {
1145 ProviderLifecycleState::Open { admitted } => {
1146 *admitted = admitted.checked_add(1).ok_or_else(|| {
1147 KitError::invalid_state("provider admission is exhausted")
1148 .at_operation(operation)
1149 })?;
1150 Ok(ProviderAdmission {
1151 runtime: Arc::clone(self),
1152 batch_bytes: None,
1153 })
1154 }
1155 ProviderLifecycleState::Closing { .. } | ProviderLifecycleState::Closed { .. } => {
1156 Err(KitError::closed("provider admission closed").at_operation(operation))
1157 }
1158 }
1159 }
1160
1161 pub(crate) async fn run_owned_batch(
1162 &self,
1163 submission: OpenAiBatchSubmission,
1164 mut admission: ProviderAdmission,
1165 ) -> Result<BatchJob, KitError> {
1166 let handle = tokio::runtime::Handle::try_current().map_err(|_| {
1167 KitError::invalid_state("inline batch ownership requires a Tokio runtime")
1168 .at_operation(Operation::SubmitBatch)
1169 })?;
1170 let bytes = submission.payload_bytes();
1171 let (observer, activate) = {
1172 let mut batches = self.batches();
1173 batches.reserve(bytes)?;
1174 batches.register(submission, &handle)
1175 };
1176 // Preserve the original admission acquired before preparation. Even
1177 // shutdown sealing during rendering cannot harvest tasks until this
1178 // admission is released. Activation must follow registry lock release.
1179 admission.batch_bytes = Some(bytes);
1180 if let Err(admission) = activate.send(admission) {
1181 // The pre-spawn publication guard has already reported worker loss.
1182 // This release also returns the byte charge exactly once.
1183 drop(admission);
1184 }
1185 observer
1186 .wait()
1187 .await
1188 .map_err(|error| error.at_operation(Operation::SubmitBatch))
1189 }
1190
1191 async fn shutdown(self: &Arc<Self>) -> Result<(), KitError> {
1192 self.start_shutdown()?;
1193 self.terminal().await
1194 }
1195
1196 fn start_shutdown(self: &Arc<Self>) -> Result<(), KitError> {
1197 let mut state = self.state();
1198 let admitted = match &*state {
1199 ProviderLifecycleState::Open { admitted } => *admitted,
1200 ProviderLifecycleState::Closing { .. } | ProviderLifecycleState::Closed { .. } => {
1201 return Ok(());
1202 }
1203 };
1204 let handle = tokio::runtime::Handle::try_current()
1205 .map_err(|_| KitError::invalid_state("provider shutdown requires a Tokio runtime"))?;
1206 let runtime = Arc::clone(self);
1207 let task = handle.spawn(async move { runtime.finish_shutdown().await });
1208 *state = ProviderLifecycleState::Closing {
1209 admitted,
1210 task: Some(task),
1211 };
1212 Ok(())
1213 }
1214
1215 async fn finish_shutdown(self: Arc<Self>) -> Result<(), KitError> {
1216 self.drain().await;
1217 // Admissions cover preparation, registration, execution and cleanup.
1218 // Only after draining can no previously admitted call add another task.
1219 let mut tasks = self.batches().take_tasks();
1220 while let Some(result) = tasks.join_next().await {
1221 self.batches().joined(result);
1222 }
1223 let ownership = self.batches().terminal();
1224 let backend = self.backend.shutdown().await;
1225 match (ownership, backend) {
1226 (Ok(()), result) => result,
1227 (Err(primary), Ok(())) => Err(primary),
1228 (Err(primary), Err(cleanup)) => Err(primary.with_supplementary(cleanup)),
1229 }
1230 }
1231
1232 fn batches(&self) -> std::sync::MutexGuard<'_, BatchRegistry> {
1233 self.batches
1234 .lock()
1235 .unwrap_or_else(std::sync::PoisonError::into_inner)
1236 }
1237
1238 fn state(&self) -> std::sync::MutexGuard<'_, ProviderLifecycleState> {
1239 self.state
1240 .lock()
1241 .unwrap_or_else(std::sync::PoisonError::into_inner)
1242 }
1243
1244 fn release(&self) {
1245 let should_notify = {
1246 let mut state = self.state();
1247 match &mut *state {
1248 ProviderLifecycleState::Open { admitted } => {
1249 debug_assert_ne!(*admitted, 0);
1250 *admitted = admitted.saturating_sub(1);
1251 false
1252 }
1253 ProviderLifecycleState::Closing { admitted, .. } => {
1254 debug_assert_ne!(*admitted, 0);
1255 *admitted = admitted.saturating_sub(1);
1256 *admitted == 0
1257 }
1258 ProviderLifecycleState::Closed { .. } => {
1259 debug_assert!(false, "closed provider retained an admission");
1260 false
1261 }
1262 }
1263 };
1264 if should_notify {
1265 self.changed.notify_waiters();
1266 }
1267 }
1268
1269 async fn drain(&self) {
1270 loop {
1271 let notified = self.changed.notified();
1272 let drained = match &*self.state() {
1273 ProviderLifecycleState::Open { admitted }
1274 | ProviderLifecycleState::Closing { admitted, .. } => *admitted == 0,
1275 ProviderLifecycleState::Closed { .. } => true,
1276 };
1277 if drained {
1278 return;
1279 }
1280 notified.await;
1281 }
1282 }
1283
1284 async fn terminal(&self) -> Result<(), KitError> {
1285 loop {
1286 let notified = self.changed.notified();
1287 let mut lease = {
1288 let mut state = self.state();
1289 match &mut *state {
1290 ProviderLifecycleState::Closed { terminal } => return terminal.clone(),
1291 ProviderLifecycleState::Closing { task, .. } => {
1292 task.take().map(|task| ShutdownTaskLease::new(self, task))
1293 }
1294 ProviderLifecycleState::Open { .. } => {
1295 return Err(KitError::invalid_state(
1296 "provider shutdown lost lifecycle ownership",
1297 ));
1298 }
1299 }
1300 };
1301 let Some(lease) = lease.as_mut() else {
1302 notified.await;
1303 continue;
1304 };
1305 let terminal = lease.join().await;
1306 lease.complete(terminal.clone());
1307 return terminal;
1308 }
1309 }
1310}
1311
1312impl<'a> ShutdownTaskLease<'a> {
1313 fn new(runtime: &'a ProviderRuntime, task: JoinHandle<Result<(), KitError>>) -> Self {
1314 Self {
1315 runtime,
1316 task: Some(task),
1317 }
1318 }
1319
1320 async fn join(&mut self) -> Result<(), KitError> {
1321 let Some(task) = self.task.as_mut() else {
1322 return Err(KitError::invalid_state(
1323 "provider shutdown task lease was empty",
1324 ));
1325 };
1326 task.await.unwrap_or_else(|_| {
1327 Err(KitError::invalid_state(
1328 "provider shutdown owner could not be joined",
1329 ))
1330 })
1331 }
1332
1333 fn complete(&mut self, terminal: Result<(), KitError>) {
1334 self.task = None;
1335 let mut state = self.runtime.state();
1336 if matches!(&*state, ProviderLifecycleState::Closing { task: None, .. }) {
1337 *state = ProviderLifecycleState::Closed { terminal };
1338 drop(state);
1339 self.runtime.changed.notify_waiters();
1340 }
1341 }
1342}
1343
1344impl Drop for ShutdownTaskLease<'_> {
1345 fn drop(&mut self) {
1346 let Some(task) = self.task.take() else {
1347 return;
1348 };
1349 let mut state = self.runtime.state();
1350 if let ProviderLifecycleState::Closing { task: slot, .. } = &mut *state
1351 && slot.is_none()
1352 {
1353 *slot = Some(task);
1354 drop(state);
1355 self.runtime.changed.notify_waiters();
1356 return;
1357 }
1358 debug_assert!(false, "shutdown task lease lost lifecycle ownership");
1359 }
1360}
1361
1362impl Drop for ProviderAdmission {
1363 fn drop(&mut self) {
1364 if let Some(bytes) = self.batch_bytes {
1365 self.runtime.batches().release(bytes);
1366 }
1367 self.runtime.release();
1368 }
1369}
1370
1371impl ProviderAdmission {
1372 pub(crate) fn bind(self, stream: GenerateEventStream) -> GenerateEventStream {
1373 GenerateEventStream::new(ProviderAdmittedStream {
1374 inner: stream,
1375 admission: Some(self),
1376 })
1377 }
1378}
1379
1380impl Stream for ProviderAdmittedStream {
1381 type Item = Result<GenerateEvent, KitError>;
1382
1383 fn poll_next(mut self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Option<Self::Item>> {
1384 let polled = Pin::new(&mut self.inner).poll_next(context);
1385 if matches!(
1386 &polled,
1387 Poll::Ready(None)
1388 | Poll::Ready(Some(Err(_)))
1389 | Poll::Ready(Some(Ok(GenerateEvent::Completed(_))))
1390 ) {
1391 drop(self.admission.take());
1392 }
1393 polled.map(|item| {
1394 item.map(|result| result.map_err(|error| error.at_operation(Operation::GenerateStream)))
1395 })
1396 }
1397}
1398
1399impl Provider {
1400 /// Constructs a built-in provider from validated configuration.
1401 ///
1402 /// Construction creates clients locally and does not probe credentials or model availability.
1403 ///
1404 /// # Errors
1405 ///
1406 /// Returns a configuration or client-construction error for invalid endpoints, transport settings, or HTTP/TLS client setup. Service authentication is checked by subsequent operations.
1407 ///
1408 /// See [`crate::ProviderBackend`] for a complete example.
1409 ///
1410 pub fn from_config(config: ProviderConfig) -> Result<Self, KitError> {
1411 let (kind, transport) = config.into_parts();
1412 let backend = match kind {
1413 ProviderKind::OpenAi(config) => ProviderInner::OpenAi(
1414 OpenAiBackend::from_config(config, transport)
1415 .map_err(|error| error.at_operation(Operation::Configure))?,
1416 ),
1417 ProviderKind::Anthropic(config) => ProviderInner::Anthropic(
1418 AnthropicBackend::from_config(config, transport)
1419 .map_err(|error| error.at_operation(Operation::Configure))?,
1420 ),
1421 ProviderKind::Gemini(config) => ProviderInner::Gemini(
1422 GeminiBackend::from_config(config, transport)
1423 .map_err(|error| error.at_operation(Operation::Configure))?,
1424 ),
1425 ProviderKind::Ollama(config) => ProviderInner::Ollama(
1426 OllamaBackend::from_config(config, transport)
1427 .map_err(|error| error.at_operation(Operation::Configure))?,
1428 ),
1429 };
1430 Ok(Self {
1431 inner: Arc::new(ProviderRuntime::new(backend)),
1432 })
1433 }
1434}
1435
1436impl ProviderBackend for Provider {
1437 type Model = Model;
1438
1439 fn model(&self, id: ModelId) -> Result<Model, KitError> {
1440 let _admission = self.inner.admit(Operation::SelectModel)?;
1441 let model = self
1442 .inner
1443 .backend
1444 .model(id)
1445 .map_err(|error| error.at_operation(Operation::SelectModel))?;
1446 Ok(Model::new(model, Arc::clone(&self.inner)))
1447 }
1448
1449 async fn list_models(&self, request: ModelListRequest) -> Result<ModelPage, KitError> {
1450 let _admission = self.inner.admit(Operation::ListModels)?;
1451 self.inner
1452 .backend
1453 .list_models(request)
1454 .await
1455 .map_err(|error| error.at_operation(Operation::ListModels))
1456 }
1457
1458 async fn model_info(&self, id: ModelId) -> Result<ModelSummary, KitError> {
1459 let _admission = self.inner.admit(Operation::ModelInfo)?;
1460 self.inner
1461 .backend
1462 .model_info(id)
1463 .await
1464 .map_err(|error| error.at_operation(Operation::ModelInfo))
1465 }
1466
1467 async fn upload_file(&self, request: FileUploadRequest) -> Result<FileObject, KitError> {
1468 let _admission = self.inner.admit(Operation::UploadFile)?;
1469 self.inner
1470 .backend
1471 .upload_file(request)
1472 .await
1473 .map_err(|error| error.at_operation(Operation::UploadFile))
1474 }
1475
1476 async fn list_files(&self, request: FileListRequest) -> Result<FilePage, KitError> {
1477 let _admission = self.inner.admit(Operation::ListFiles)?;
1478 self.inner
1479 .backend
1480 .list_files(request)
1481 .await
1482 .map_err(|error| error.at_operation(Operation::ListFiles))
1483 }
1484
1485 async fn file_info(&self, id: String) -> Result<FileObject, KitError> {
1486 let _admission = self.inner.admit(Operation::FileInfo)?;
1487 self.inner
1488 .backend
1489 .file_info(id)
1490 .await
1491 .map_err(|error| error.at_operation(Operation::FileInfo))
1492 }
1493
1494 async fn delete_file(&self, id: String) -> Result<(), KitError> {
1495 let _admission = self.inner.admit(Operation::DeleteFile)?;
1496 self.inner
1497 .backend
1498 .delete_file(id)
1499 .await
1500 .map_err(|error| error.at_operation(Operation::DeleteFile))
1501 }
1502
1503 async fn download_file(
1504 &self,
1505 id: String,
1506 max_bytes: NonZeroUsize,
1507 ) -> Result<Vec<u8>, KitError> {
1508 let _admission = self.inner.admit(Operation::DownloadFile)?;
1509 self.inner
1510 .backend
1511 .download_file(id, max_bytes)
1512 .await
1513 .map_err(|error| error.at_operation(Operation::DownloadFile))
1514 }
1515
1516 async fn batch_status(&self, id: String) -> Result<BatchJob, KitError> {
1517 let _admission = self.inner.admit(Operation::BatchStatus)?;
1518 self.inner
1519 .backend
1520 .batch_status(id)
1521 .await
1522 .map_err(|error| error.at_operation(Operation::BatchStatus))
1523 }
1524
1525 async fn cancel_batch(&self, id: String) -> Result<BatchJob, KitError> {
1526 let _admission = self.inner.admit(Operation::CancelBatch)?;
1527 self.inner
1528 .backend
1529 .cancel_batch(id)
1530 .await
1531 .map_err(|error| error.at_operation(Operation::CancelBatch))
1532 }
1533
1534 async fn batch_results(
1535 &self,
1536 id: String,
1537 max_bytes: NonZeroUsize,
1538 ) -> Result<Vec<BatchResult<GenerationContinuation>>, KitError> {
1539 let _admission = self.inner.admit(Operation::BatchResults)?;
1540 self.inner
1541 .backend
1542 .batch_results(id, max_bytes)
1543 .await
1544 .map_err(|error| error.at_operation(Operation::BatchResults))
1545 }
1546
1547 async fn list_batches(&self, request: BatchListRequest) -> Result<BatchPage, KitError> {
1548 let _admission = self.inner.admit(Operation::ListBatches)?;
1549 self.inner
1550 .backend
1551 .list_batches(request)
1552 .await
1553 .map_err(|error| error.at_operation(Operation::ListBatches))
1554 }
1555
1556 async fn health(&self) -> Result<HealthStatus, KitError> {
1557 let _admission = self.inner.admit(Operation::Health)?;
1558 self.inner
1559 .backend
1560 .health()
1561 .await
1562 .map_err(|error| error.at_operation(Operation::Health))
1563 }
1564
1565 async fn shutdown(&self) -> Result<(), KitError> {
1566 self.inner
1567 .shutdown()
1568 .await
1569 .map_err(|error| error.at_operation(Operation::Shutdown))
1570 }
1571}
1572
1573impl fmt::Debug for Provider {
1574 fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
1575 formatter.debug_struct("Provider").finish_non_exhaustive()
1576 }
1577}