tenant.openagents/omega

1#[cfg(feature = "predict-edits")]
2pub mod predict_edits_v3;
3#[cfg(feature = "predict-edits")]
4pub mod predict_edits_v4;
5
6use std::str::FromStr;
7use std::sync::Arc;
8
9use anyhow::Context as _;
10use serde::{Deserialize, Serialize};
11use strum::{Display, EnumIter, EnumString};
12use uuid::Uuid;
13
14/// The name of the header used to indicate which version of Zed the client is running.
15pub const ZED_VERSION_HEADER_NAME: &str = "x-zed-version";
16
17/// The name of the header used to indicate which edit prediction experiment should be used.
18pub const PREFERRED_EXPERIMENT_HEADER_NAME: &str = "x-zed-preferred-experiment";
19
20/// The name of the header used to indicate when a request failed due to an
21/// expired LLM token.
22///
23/// The client may use this as a signal to refresh the token.
24pub const EXPIRED_LLM_TOKEN_HEADER_NAME: &str = "x-zed-expired-token";
25
26/// The name of the header used to indicate when a request failed due to an outdated LLM token.
27///
28/// A token is considered "outdated" when we can't parse the claims (e.g., after adding a new required claim).
29///
30/// This is distinct from [`EXPIRED_LLM_TOKEN_HEADER_NAME`] which indicates the token's time-based validity has passed.
31/// An outdated token means the token's structure is incompatible with the current server expectations.
32pub const OUTDATED_LLM_TOKEN_HEADER_NAME: &str = "x-zed-outdated-token";
33
34/// The name of the header used to indicate the usage limit for edit predictions.
35pub const EDIT_PREDICTIONS_USAGE_LIMIT_HEADER_NAME: &str = "x-zed-edit-predictions-usage-limit";
36
37/// The name of the header used to indicate the usage amount for edit predictions.
38pub const EDIT_PREDICTIONS_USAGE_AMOUNT_HEADER_NAME: &str = "x-zed-edit-predictions-usage-amount";
39
40pub const EDIT_PREDICTIONS_RESOURCE_HEADER_VALUE: &str = "edit_predictions";
41
42/// The name of the header used to indicate the minimum required Zed version.
43///
44/// This can be used to force a Zed upgrade in order to continue communicating
45/// with the LLM service.
46pub const MINIMUM_REQUIRED_VERSION_HEADER_NAME: &str = "x-zed-minimum-required-version";
47
48/// The name of the header used by the client to indicate to the server that it supports receiving status messages.
49pub const CLIENT_SUPPORTS_STATUS_MESSAGES_HEADER_NAME: &str =
50    "x-zed-client-supports-status-messages";
51
52/// The name of the header used by the client to indicate to the server that it supports receiving a "stream_ended" request completion status.
53pub const CLIENT_SUPPORTS_STATUS_STREAM_ENDED_HEADER_NAME: &str =
54    "x-zed-client-supports-stream-ended-request-completion-status";
55
56/// The name of the header used by the server to indicate to the client that it supports sending status messages.
57pub const SERVER_SUPPORTS_STATUS_MESSAGES_HEADER_NAME: &str =
58    "x-zed-server-supports-status-messages";
59
60/// The name of the header used by the client to indicate that it supports receiving xAI models.
61pub const CLIENT_SUPPORTS_X_AI_HEADER_NAME: &str = "x-zed-client-supports-x-ai";
62
63/// The maximum number of edit predictions that can be rejected per request.
64pub const MAX_EDIT_PREDICTION_REJECTIONS_PER_REQUEST: usize = 100;
65
66#[derive(Debug, PartialEq, Clone, Copy, Serialize, Deserialize)]
67#[serde(rename_all = "snake_case")]
68pub enum UsageLimit {
69    Limited(i32),
70    Unlimited,
71}
72
73impl FromStr for UsageLimit {
74    type Err = anyhow::Error;
75
76    fn from_str(value: &str) -> Result<Self, Self::Err> {
77        match value {
78            "unlimited" => Ok(Self::Unlimited),
79            limit => limit
80                .parse::<i32>()
81                .map(Self::Limited)
82                .context("failed to parse limit"),
83        }
84    }
85}
86
87#[derive(
88    Debug, PartialEq, Eq, Hash, Clone, Copy, Serialize, Deserialize, EnumString, EnumIter, Display,
89)]
90#[serde(rename_all = "snake_case")]
91#[strum(serialize_all = "snake_case")]
92pub enum LanguageModelProvider {
93    Anthropic,
94    OpenAi,
95    Google,
96    XAi,
97}
98
99#[derive(Debug, Clone, Serialize, Deserialize)]
100pub struct PredictEditsBody {
101    #[serde(skip_serializing_if = "Option::is_none", default)]
102    pub outline: Option<String>,
103    pub input_events: String,
104    pub input_excerpt: String,
105    #[serde(skip_serializing_if = "Option::is_none", default)]
106    pub speculated_output: Option<String>,
107    /// Whether the user provided consent for sampling this interaction.
108    #[serde(default, alias = "data_collection_permission")]
109    pub can_collect_data: bool,
110    #[serde(skip_serializing_if = "Option::is_none", default)]
111    pub diagnostic_groups: Option<Vec<(String, serde_json::Value)>>,
112    /// Info about the git repository state, only present when can_collect_data is true.
113    #[serde(skip_serializing_if = "Option::is_none", default)]
114    pub git_info: Option<PredictEditsGitInfo>,
115    /// The trigger for this request.
116    #[serde(default)]
117    pub trigger: PredictEditsRequestTrigger,
118}
119
120#[derive(
121    Default, Debug, Clone, Copy, Serialize, Deserialize, strum::AsRefStr, strum::EnumString,
122)]
123#[strum(serialize_all = "snake_case")]
124pub enum PredictEditsRequestTrigger {
125    Testing,
126    Diagnostics,
127    DiagnosticNavigation,
128    Cli,
129    Explicit,
130    BufferEdit,
131    LSPCompletionAccepted,
132    PredictionAccepted,
133    PredictionPartiallyAccepted,
134    EditorCreated,
135    ProviderChanged,
136    UserInfoChanged,
137    VimModeChanged,
138    SettingsChanged,
139    #[default]
140    Other,
141}
142
143#[derive(Debug, Clone, Serialize, Deserialize)]
144pub struct PredictEditsGitInfo {
145    /// SHA of git HEAD commit at time of prediction.
146    #[serde(skip_serializing_if = "Option::is_none", default)]
147    pub head_sha: Option<String>,
148    /// URL of the remote called `origin`.
149    #[serde(skip_serializing_if = "Option::is_none", default)]
150    pub remote_origin_url: Option<String>,
151    /// URL of the remote called `upstream`.
152    #[serde(skip_serializing_if = "Option::is_none", default)]
153    pub remote_upstream_url: Option<String>,
154}
155
156#[derive(Debug, Clone, Serialize, Deserialize)]
157pub struct PredictEditsResponse {
158    pub request_id: String,
159    pub output_excerpt: String,
160}
161
162#[derive(Debug, Clone, Serialize, Deserialize)]
163pub struct AcceptEditPredictionBody {
164    pub request_id: String,
165    #[serde(default, skip_serializing_if = "Option::is_none")]
166    pub model_version: Option<String>,
167    #[serde(default, skip_serializing_if = "Option::is_none")]
168    pub e2e_latency_ms: Option<u128>,
169}
170
171#[derive(Debug, Clone, Deserialize)]
172pub struct RejectEditPredictionsBody {
173    pub rejections: Vec<EditPredictionRejection>,
174}
175
176#[derive(Debug, Clone, Serialize)]
177pub struct RejectEditPredictionsBodyRef<'a> {
178    pub rejections: &'a [EditPredictionRejection],
179}
180
181#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
182pub struct EditPredictionRejection {
183    pub request_id: String,
184    #[serde(default)]
185    pub reason: EditPredictionRejectReason,
186    pub was_shown: bool,
187    #[serde(default, skip_serializing_if = "Option::is_none")]
188    pub model_version: Option<String>,
189    #[serde(default, skip_serializing_if = "Option::is_none")]
190    pub e2e_latency_ms: Option<u128>,
191}
192
193#[derive(
194    Default, Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Hash, strum::AsRefStr,
195)]
196#[strum(serialize_all = "snake_case")]
197pub enum EditPredictionRejectReason {
198    /// New requests were triggered before this one completed
199    Canceled,
200    /// No edits returned
201    Empty,
202    /// Edits returned, but none remained after interpolation
203    InterpolatedEmpty,
204    /// Edits returned, but could not be interpolated after buffer changes
205    InterpolateFailed,
206    /// A patch was returned, but could not be applied to the buffer
207    PatchApplyFailed,
208    /// The new prediction was preferred over the current one
209    Replaced,
210    /// The current prediction was preferred over the new one
211    CurrentPreferred,
212    /// The current prediction was discarded
213    #[default]
214    Discarded,
215    /// The current prediction was explicitly rejected by the user
216    Rejected,
217}
218
219#[derive(Debug, Serialize, Deserialize)]
220pub struct CompletionBody {
221    #[serde(skip_serializing_if = "Option::is_none", default)]
222    pub thread_id: Option<String>,
223    #[serde(skip_serializing_if = "Option::is_none", default)]
224    pub prompt_id: Option<String>,
225    pub provider: LanguageModelProvider,
226    pub model: String,
227    pub provider_request: serde_json::Value,
228}
229
230#[derive(Debug, PartialEq, Clone, Serialize, Deserialize)]
231#[serde(rename_all = "snake_case")]
232pub enum CompletionRequestStatus {
233    Queued {
234        position: usize,
235    },
236    Started,
237    Failed {
238        code: String,
239        message: String,
240        request_id: Uuid,
241        /// Retry duration in seconds.
242        retry_after: Option<f64>,
243    },
244    /// The cloud sends a StreamEnded message when the stream from the LLM provider finishes.
245    StreamEnded,
246    #[serde(other)]
247    Unknown,
248}
249
250#[derive(Serialize, Deserialize)]
251#[serde(rename_all = "snake_case")]
252pub enum CompletionEvent<T> {
253    Status(CompletionRequestStatus),
254    Event(T),
255}
256
257impl<T> CompletionEvent<T> {
258    pub fn into_status(self) -> Option<CompletionRequestStatus> {
259        match self {
260            Self::Status(status) => Some(status),
261            Self::Event(_) => None,
262        }
263    }
264
265    pub fn into_event(self) -> Option<T> {
266        match self {
267            Self::Event(event) => Some(event),
268            Self::Status(_) => None,
269        }
270    }
271}
272
273#[derive(Serialize, Deserialize)]
274pub struct WebSearchBody {
275    pub query: String,
276}
277
278#[derive(Debug, Serialize, Deserialize, Clone)]
279pub struct WebSearchResponse {
280    pub results: Vec<WebSearchResult>,
281}
282
283#[derive(Debug, Serialize, Deserialize, Clone)]
284pub struct WebSearchResult {
285    pub title: String,
286    pub url: String,
287    pub text: String,
288}
289
290#[derive(Debug, PartialEq, Eq, Hash, Clone, Serialize, Deserialize)]
291pub struct LanguageModelId(pub Arc<str>);
292
293impl std::fmt::Display for LanguageModelId {
294    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
295        write!(f, "{}", self.0)
296    }
297}
298
299#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
300pub struct LanguageModel {
301    pub provider: LanguageModelProvider,
302    pub id: LanguageModelId,
303    pub display_name: String,
304    #[serde(default)]
305    pub is_latest: bool,
306    pub max_token_count: usize,
307    pub max_token_count_in_max_mode: Option<usize>,
308    pub max_output_tokens: usize,
309    pub supports_tools: bool,
310    pub supports_images: bool,
311    pub supports_thinking: bool,
312    /// Whether thinking can be turned off entirely for this model, allowing
313    /// clients to offer an "off" choice alongside `supported_effort_levels`.
314    /// Some models (e.g. Claude Fable 5) always think and cannot honor an
315    /// "off" request. Only meaningful when `supports_thinking` is `true`.
316    #[serde(default)]
317    pub supports_disabling_thinking: bool,
318    #[serde(default)]
319    pub supports_fast_mode: bool,
320    #[serde(default)]
321    pub supports_server_side_compaction: bool,
322    pub supported_effort_levels: Vec<SupportedEffortLevel>,
323    #[serde(default)]
324    pub supports_streaming_tools: bool,
325    /// Only used by OpenAI and xAI.
326    #[serde(default)]
327    pub supports_parallel_tool_calls: bool,
328    #[serde(default)]
329    pub is_disabled: bool,
330    #[serde(default, skip_serializing_if = "Option::is_none")]
331    pub disabled_reason: Option<String>,
332}
333
334#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
335pub struct SupportedEffortLevel {
336    pub name: Arc<str>,
337    pub value: Arc<str>,
338    #[serde(default, skip_serializing_if = "Option::is_none")]
339    pub is_default: Option<bool>,
340}
341
342#[derive(Debug, Serialize, Deserialize)]
343pub struct ListModelsResponse {
344    pub models: Vec<LanguageModel>,
345    pub default_model: Option<LanguageModelId>,
346    pub default_fast_model: Option<LanguageModelId>,
347    pub recommended_models: Vec<LanguageModelId>,
348}
349
350#[derive(Debug, PartialEq, Serialize, Deserialize)]
351pub struct CurrentUsage {
352    pub edit_predictions: UsageData,
353}
354
355#[derive(Debug, PartialEq, Serialize, Deserialize)]
356pub struct UsageData {
357    pub used: u32,
358    pub limit: UsageLimit,
359}
360
361#[cfg(test)]
362mod tests {
363    use super::*;
364
365    #[test]
366    fn test_usage_limit_from_str() {
367        let limit = UsageLimit::from_str("unlimited").unwrap();
368        assert!(matches!(limit, UsageLimit::Unlimited));
369
370        let limit = UsageLimit::from_str(&0.to_string()).unwrap();
371        assert!(matches!(limit, UsageLimit::Limited(0)));
372
373        let limit = UsageLimit::from_str(&50.to_string()).unwrap();
374        assert!(matches!(limit, UsageLimit::Limited(50)));
375
376        for value in ["not_a_number", "50xyz"] {
377            let limit = UsageLimit::from_str(value);
378            assert!(limit.is_err());
379        }
380    }
381}
382