forge_orchestration/sdk/
mod.rs

1//! # Forge SDK
2//!
3//! Lightweight SDK for workloads to interact with the Forge orchestrator.
4//!
5//! This module provides utilities for workloads running under Forge orchestration:
6//! - **Lifecycle**: Signal readiness, handle graceful shutdown
7//! - **Port Allocation**: Request and release ports dynamically
8//! - **Client**: HTTP client for Forge API communication
9//!
10//! ## Quick Start
11//!
12//! ```rust,no_run
13//! use forge_orchestration::sdk::{ready, allocate_port, graceful_shutdown};
14//!
15//! #[tokio::main]
16//! async fn main() -> forge_orchestration::Result<()> {
17//!     // Signal readiness
18//!     ready()?;
19//!
20//!     // Get a port
21//!     let port = allocate_port(8000..9000)?;
22//!     println!("Listening on port {}", port);
23//!
24//!     // Handle shutdown
25//!     graceful_shutdown();
26//!     Ok(())
27//! }
28//! ```
29
30pub mod client;
31pub mod lifecycle;
32pub mod port;
33
34pub use client::{ForgeClient, MetricsReport, start_heartbeat};
35pub use lifecycle::{graceful_shutdown, is_ready, notify_shutdown, ready, shutdown_requested, shutdown_signal};
36pub use port::{allocate_port, allocate_specific_port, allocated_ports, is_port_available, release_port, PortAllocator};
37
38use crate::error::ForgeError;
39
40/// Environment variable for Forge API endpoint
41pub const FORGE_API_ENV: &str = "FORGE_API";
42
43/// Environment variable for allocation ID
44pub const FORGE_ALLOC_ID_ENV: &str = "FORGE_ALLOC_ID";
45
46/// Environment variable for task name
47pub const FORGE_TASK_NAME_ENV: &str = "FORGE_TASK_NAME";
48
49/// Get the Forge API endpoint from environment
50pub fn forge_api_url() -> Option<String> {
51    std::env::var(FORGE_API_ENV).ok()
52}
53
54/// Get the allocation ID from environment
55pub fn alloc_id() -> Option<String> {
56    std::env::var(FORGE_ALLOC_ID_ENV).ok()
57}
58
59/// Get the task name from environment
60pub fn task_name() -> Option<String> {
61    std::env::var(FORGE_TASK_NAME_ENV).ok()
62}
63
64/// SDK-specific error type
65#[derive(Debug, thiserror::Error)]
66pub enum SdkError {
67    /// Forge API not configured
68    #[error("Forge API not configured: set {0} environment variable")]
69    NotConfigured(&'static str),
70
71    /// API communication error
72    #[error("API error: {0}")]
73    Api(String),
74
75    /// Port allocation failed
76    #[error("Port allocation failed: {0}")]
77    PortAllocation(String),
78
79    /// Lifecycle error
80    #[error("Lifecycle error: {0}")]
81    Lifecycle(String),
82
83    /// IO error
84    #[error("IO error: {0}")]
85    Io(#[from] std::io::Error),
86
87    /// Serialization error
88    #[error("Serialization error: {0}")]
89    Serialization(#[from] serde_json::Error),
90
91    /// HTTP error
92    #[error("HTTP error: {0}")]
93    Http(#[from] reqwest::Error),
94}
95
96impl SdkError {
97    /// Create an API error
98    pub fn api(msg: impl Into<String>) -> Self {
99        Self::Api(msg.into())
100    }
101
102    /// Create a port allocation error
103    pub fn port(msg: impl Into<String>) -> Self {
104        Self::PortAllocation(msg.into())
105    }
106
107    /// Create a lifecycle error
108    pub fn lifecycle(msg: impl Into<String>) -> Self {
109        Self::Lifecycle(msg.into())
110    }
111}
112
113impl From<SdkError> for ForgeError {
114    fn from(err: SdkError) -> Self {
115        ForgeError::Internal(err.to_string())
116    }
117}
118
119/// SDK Result type alias
120pub type SdkResult<T> = std::result::Result<T, SdkError>;