tenant.openagents/omega

1use anyhow::Result;
2use collections::BTreeMap;
3use credentials_provider::CredentialsProvider;
4use futures::{AsyncReadExt, FutureExt, StreamExt, future::BoxFuture};
5use gpui::{App, AppContext, AsyncApp, Context, Entity, SharedString, Task};
6use http_client::{
7    AsyncBody, CustomHeaders, HttpClient, Method, Request as HttpRequest, RequestBuilderExt, http,
8};
9use language_model::{
10    ApiKeyConfiguration, ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel,
11    LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName,
12    LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
13    LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice,
14    LanguageModelToolSchemaFormat, ProviderSettingsView, RateLimiter, env_var,
15};
16use open_ai::ResponseStreamEvent;
17use serde::Deserialize;
18pub use settings::OpenAiCompatibleModelCapabilities as ModelCapabilities;
19pub use settings::VercelAiGatewayAvailableModel as AvailableModel;
20use settings::{Settings, SettingsStore};
21use std::sync::{Arc, LazyLock};
22use ui::IconName;
23
24const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("vercel_ai_gateway");
25const PROVIDER_NAME: LanguageModelProviderName =
26    LanguageModelProviderName::new("Vercel AI Gateway");
27
28const API_URL: &str = "https://ai-gateway.vercel.sh/v1";
29const API_KEY_ENV_VAR_NAME: &str = "VERCEL_AI_GATEWAY_API_KEY";
30static API_KEY_ENV_VAR: LazyLock<EnvVar> = env_var!(API_KEY_ENV_VAR_NAME);
31
32#[derive(Default, Clone, Debug, PartialEq)]
33pub struct VercelAiGatewaySettings {
34    pub api_url: String,
35    pub available_models: Vec<AvailableModel>,
36    pub custom_headers: CustomHeaders,
37}
38
39pub struct VercelAiGatewayLanguageModelProvider {
40    http_client: Arc<dyn HttpClient>,
41    state: Entity<State>,
42}
43
44pub struct State {
45    api_key_state: ApiKeyState,
46    credentials_provider: Arc<dyn CredentialsProvider>,
47    http_client: Arc<dyn HttpClient>,
48    available_models: Vec<AvailableModel>,
49    fetch_models_task: Option<Task<Result<(), LanguageModelCompletionError>>>,
50}
51
52impl State {
53    fn is_authenticated(&self) -> bool {
54        self.api_key_state.has_key()
55    }
56
57    fn set_api_key(&mut self, api_key: Option<String>, cx: &mut Context<Self>) -> Task<Result<()>> {
58        let credentials_provider = self.credentials_provider.clone();
59        let api_url = VercelAiGatewayLanguageModelProvider::api_url(cx);
60        self.api_key_state.store(
61            api_url,
62            api_key,
63            |this| &mut this.api_key_state,
64            credentials_provider,
65            cx,
66        )
67    }
68
69    fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
70        let credentials_provider = self.credentials_provider.clone();
71        let api_url = VercelAiGatewayLanguageModelProvider::api_url(cx);
72        let task = self.api_key_state.load_if_needed(
73            api_url,
74            |this| &mut this.api_key_state,
75            credentials_provider,
76            cx,
77        );
78
79        cx.spawn(async move |this, cx| {
80            let result = task.await;
81            this.update(cx, |this, cx| this.restart_fetch_models_task(cx))
82                .ok();
83            result
84        })
85    }
86
87    fn fetch_models(
88        &mut self,
89        cx: &mut Context<Self>,
90    ) -> Task<Result<(), LanguageModelCompletionError>> {
91        let http_client = self.http_client.clone();
92        let api_url = VercelAiGatewayLanguageModelProvider::api_url(cx);
93        let api_key = self.api_key_state.key(&api_url);
94        let extra_headers = VercelAiGatewayLanguageModelProvider::settings(cx)
95            .custom_headers
96            .clone();
97        cx.spawn(async move |this, cx| {
98            let models = list_models(
99                http_client.as_ref(),
100                &api_url,
101                api_key.as_deref(),
102                &extra_headers,
103            )
104            .await?;
105            this.update(cx, |this, cx| {
106                this.available_models = models;
107                cx.notify();
108            })
109            .map_err(|e| LanguageModelCompletionError::Other(e))?;
110            Ok(())
111        })
112    }
113
114    fn restart_fetch_models_task(&mut self, cx: &mut Context<Self>) {
115        if self.is_authenticated() {
116            let task = self.fetch_models(cx);
117            self.fetch_models_task.replace(task);
118        } else {
119            self.available_models = Vec::new();
120        }
121    }
122}
123
124impl VercelAiGatewayLanguageModelProvider {
125    pub fn new(
126        http_client: Arc<dyn HttpClient>,
127        credentials_provider: Arc<dyn CredentialsProvider>,
128        cx: &mut App,
129    ) -> Self {
130        let state = cx.new(|cx| {
131            cx.observe_global::<SettingsStore>({
132                let mut last_settings = VercelAiGatewayLanguageModelProvider::settings(cx).clone();
133                move |this: &mut State, cx| {
134                    let current_settings = VercelAiGatewayLanguageModelProvider::settings(cx);
135                    if current_settings != &last_settings {
136                        last_settings = current_settings.clone();
137                        this.authenticate(cx).detach();
138                        cx.notify();
139                    }
140                }
141            })
142            .detach();
143            State {
144                api_key_state: ApiKeyState::new(Self::api_url(cx), (*API_KEY_ENV_VAR).clone()),
145                credentials_provider,
146                http_client: http_client.clone(),
147                available_models: Vec::new(),
148                fetch_models_task: None,
149            }
150        });
151
152        Self { http_client, state }
153    }
154
155    fn settings(cx: &App) -> &VercelAiGatewaySettings {
156        &crate::AllLanguageModelSettings::get_global(cx).vercel_ai_gateway
157    }
158
159    fn api_url(cx: &App) -> SharedString {
160        let api_url = &Self::settings(cx).api_url;
161        if api_url.is_empty() {
162            API_URL.into()
163        } else {
164            SharedString::new(api_url.as_str())
165        }
166    }
167
168    fn default_available_model() -> AvailableModel {
169        AvailableModel {
170            name: "openai/gpt-5.3-codex".to_string(),
171            display_name: Some("GPT 5.3 Codex".to_string()),
172            max_tokens: 400_000,
173            max_output_tokens: Some(128_000),
174            max_completion_tokens: None,
175            capabilities: ModelCapabilities::default(),
176        }
177    }
178
179    fn create_language_model(&self, model: AvailableModel) -> Arc<dyn LanguageModel> {
180        Arc::new(VercelAiGatewayLanguageModel {
181            id: LanguageModelId::from(model.name.clone()),
182            model,
183            state: self.state.clone(),
184            http_client: self.http_client.clone(),
185            request_limiter: RateLimiter::new(4),
186        })
187    }
188}
189
190impl LanguageModelProviderState for VercelAiGatewayLanguageModelProvider {
191    type ObservableEntity = State;
192
193    fn observable_entity(&self) -> Option<Entity<Self::ObservableEntity>> {
194        Some(self.state.clone())
195    }
196}
197
198impl LanguageModelProvider for VercelAiGatewayLanguageModelProvider {
199    fn id(&self) -> LanguageModelProviderId {
200        PROVIDER_ID
201    }
202
203    fn name(&self) -> LanguageModelProviderName {
204        PROVIDER_NAME
205    }
206
207    fn icon(&self) -> IconOrSvg {
208        IconOrSvg::Icon(IconName::AiVercel)
209    }
210
211    fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
212        Some(self.create_language_model(Self::default_available_model()))
213    }
214
215    fn default_fast_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
216        None
217    }
218
219    fn provided_models(&self, cx: &App) -> Vec<Arc<dyn LanguageModel>> {
220        let mut models = BTreeMap::default();
221
222        let default_model = Self::default_available_model();
223        models.insert(default_model.name.clone(), default_model);
224
225        for model in self.state.read(cx).available_models.clone() {
226            models.insert(model.name.clone(), model);
227        }
228
229        for model in &Self::settings(cx).available_models {
230            models.insert(model.name.clone(), model.clone());
231        }
232
233        models
234            .into_values()
235            .map(|model| self.create_language_model(model))
236            .collect()
237    }
238
239    fn is_authenticated(&self, cx: &App) -> bool {
240        self.state.read(cx).is_authenticated()
241    }
242
243    fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
244        self.state.update(cx, |state, cx| state.authenticate(cx))
245    }
246
247    fn settings_view(&self, cx: &mut App) -> Option<ProviderSettingsView> {
248        let state = self.state.read(cx);
249        Some(ProviderSettingsView::ApiKey(ApiKeyConfiguration::new(
250            state.api_key_state.has_key(),
251            state.api_key_state.is_from_env_var(),
252            state.api_key_state.env_var_name().clone(),
253            "https://vercel.com/d?to=%2F%5Bteam%5D%2F%7E%2Fai%2Fapi-keys&title=Go+to+AI+Gateway"
254                .into(),
255        )))
256    }
257
258    fn set_api_key(&self, api_key: Option<String>, cx: &mut App) -> Task<Result<()>> {
259        self.state
260            .update(cx, |state, cx| state.set_api_key(api_key, cx))
261    }
262}
263
264pub struct VercelAiGatewayLanguageModel {
265    id: LanguageModelId,
266    model: AvailableModel,
267    state: Entity<State>,
268    http_client: Arc<dyn HttpClient>,
269    request_limiter: RateLimiter,
270}
271
272impl VercelAiGatewayLanguageModel {
273    fn stream_open_ai(
274        &self,
275        request: open_ai::Request,
276        cx: &AsyncApp,
277    ) -> BoxFuture<
278        'static,
279        Result<
280            futures::stream::BoxStream<'static, Result<ResponseStreamEvent>>,
281            LanguageModelCompletionError,
282        >,
283    > {
284        let http_client = self.http_client.clone();
285        let (api_key, api_url, extra_headers) = self.state.read_with(cx, |state, cx| {
286            let api_url = VercelAiGatewayLanguageModelProvider::api_url(cx);
287            let extra_headers = VercelAiGatewayLanguageModelProvider::settings(cx)
288                .custom_headers
289                .clone();
290            (state.api_key_state.key(&api_url), api_url, extra_headers)
291        });
292
293        let future = self.request_limiter.stream(async move {
294            let provider = PROVIDER_NAME;
295            let Some(api_key) = api_key else {
296                return Err(LanguageModelCompletionError::NoApiKey { provider });
297            };
298            let request = open_ai::stream_completion(
299                http_client.as_ref(),
300                provider.0.as_str(),
301                &api_url,
302                &api_key,
303                request,
304                &extra_headers,
305            );
306            let response = request.await.map_err(map_open_ai_error)?;
307            Ok(response)
308        });
309
310        async move { Ok(future.await?.boxed()) }.boxed()
311    }
312}
313
314fn map_open_ai_error(error: open_ai::RequestError) -> LanguageModelCompletionError {
315    match error {
316        open_ai::RequestError::HttpResponseError {
317            status_code,
318            body,
319            headers,
320            ..
321        } => {
322            let retry_after = headers
323                .get(http::header::RETRY_AFTER)
324                .and_then(|value| value.to_str().ok()?.parse::<u64>().ok())
325                .map(std::time::Duration::from_secs);
326
327            LanguageModelCompletionError::from_http_status(
328                PROVIDER_NAME,
329                status_code,
330                extract_error_message(&body),
331                retry_after,
332            )
333        }
334        open_ai::RequestError::Other(error) => LanguageModelCompletionError::Other(error),
335    }
336}
337
338fn extract_error_message(body: &str) -> String {
339    let json = match serde_json::from_str::<serde_json::Value>(body) {
340        Ok(json) => json,
341        Err(_) => return body.to_string(),
342    };
343
344    let message = json
345        .get("error")
346        .and_then(|value| {
347            value
348                .get("message")
349                .and_then(serde_json::Value::as_str)
350                .or_else(|| value.as_str())
351        })
352        .or_else(|| json.get("message").and_then(serde_json::Value::as_str))
353        .map(ToString::to_string)
354        .unwrap_or_else(|| body.to_string());
355
356    clean_error_message(&message)
357}
358
359fn clean_error_message(message: &str) -> String {
360    let lower = message.to_lowercase();
361
362    if lower.contains("vercel_oidc_token") && lower.contains("oidc token") {
363        return "Authentication failed for Vercel AI Gateway. Use a Vercel AI Gateway key (vck_...).\nCreate or manage keys in Vercel AI Gateway console.\nIf this persists, regenerate the key and update it in Vercel AI Gateway provider settings in Omega.".to_string();
364    }
365
366    if lower.contains("invalid api key") || lower.contains("invalid_api_key") {
367        return "Authentication failed for Vercel AI Gateway. Check that your Vercel AI Gateway key starts with vck_ and is active.".to_string();
368    }
369
370    message.to_string()
371}
372
373fn has_tag(tags: &[String], expected: &str) -> bool {
374    tags.iter()
375        .any(|tag| tag.trim().eq_ignore_ascii_case(expected))
376}
377
378impl LanguageModel for VercelAiGatewayLanguageModel {
379    fn id(&self) -> LanguageModelId {
380        self.id.clone()
381    }
382
383    fn name(&self) -> LanguageModelName {
384        LanguageModelName::from(
385            self.model
386                .display_name
387                .clone()
388                .unwrap_or_else(|| self.model.name.clone()),
389        )
390    }
391
392    fn provider_id(&self) -> LanguageModelProviderId {
393        PROVIDER_ID
394    }
395
396    fn provider_name(&self) -> LanguageModelProviderName {
397        PROVIDER_NAME
398    }
399
400    fn supports_tools(&self) -> bool {
401        self.model.capabilities.tools
402    }
403
404    fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
405        LanguageModelToolSchemaFormat::JsonSchemaSubset
406    }
407
408    fn supports_images(&self) -> bool {
409        self.model.capabilities.images
410    }
411
412    fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
413        match choice {
414            LanguageModelToolChoice::Auto => self.model.capabilities.tools,
415            LanguageModelToolChoice::Any => self.model.capabilities.tools,
416            LanguageModelToolChoice::None => true,
417        }
418    }
419
420    fn supports_streaming_tools(&self) -> bool {
421        true
422    }
423
424    fn supports_split_token_display(&self) -> bool {
425        true
426    }
427
428    fn telemetry_id(&self) -> String {
429        format!("vercel_ai_gateway/{}", self.model.name)
430    }
431
432    fn max_token_count(&self) -> u64 {
433        self.model.max_tokens
434    }
435
436    fn max_output_tokens(&self) -> Option<u64> {
437        self.model.max_output_tokens
438    }
439
440    fn stream_completion(
441        &self,
442        request: LanguageModelRequest,
443        cx: &AsyncApp,
444    ) -> BoxFuture<
445        'static,
446        Result<
447            futures::stream::BoxStream<
448                'static,
449                Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
450            >,
451            LanguageModelCompletionError,
452        >,
453    > {
454        let request = match crate::provider::open_ai::into_open_ai(
455            request,
456            &self.model.name,
457            self.model.capabilities.parallel_tool_calls,
458            self.model.capabilities.prompt_cache_key,
459            self.max_output_tokens(),
460            crate::provider::open_ai::ChatCompletionMaxTokensParameter::MaxCompletionTokens,
461            None,
462            false,
463        ) {
464            Ok(request) => request,
465            Err(error) => return async move { Err(error.into()) }.boxed(),
466        };
467        let completions = self.stream_open_ai(request, cx);
468        async move {
469            let mapper = crate::provider::open_ai::OpenAiEventMapper::new();
470            Ok(mapper.map_stream(completions.await?).boxed())
471        }
472        .boxed()
473    }
474}
475
476#[derive(Deserialize)]
477struct ModelsResponse {
478    data: Vec<ApiModel>,
479}
480
481#[derive(Deserialize)]
482struct ApiModel {
483    id: String,
484    name: Option<String>,
485    context_window: Option<u64>,
486    max_tokens: Option<u64>,
487    #[serde(default)]
488    r#type: Option<String>,
489    #[serde(default)]
490    supported_parameters: Vec<String>,
491    #[serde(default)]
492    tags: Vec<String>,
493    architecture: Option<ApiModelArchitecture>,
494}
495
496#[derive(Deserialize)]
497struct ApiModelArchitecture {
498    #[serde(default)]
499    input_modalities: Vec<String>,
500}
501
502async fn list_models(
503    client: &dyn HttpClient,
504    api_url: &str,
505    api_key: Option<&str>,
506    extra_headers: &CustomHeaders,
507) -> Result<Vec<AvailableModel>, LanguageModelCompletionError> {
508    let uri = format!("{api_url}/models?include_mappings=true");
509    let mut request_builder = HttpRequest::builder()
510        .method(Method::GET)
511        .uri(uri)
512        .header("Accept", "application/json");
513    if let Some(api_key) = api_key {
514        request_builder = request_builder.header("Authorization", format!("Bearer {}", api_key));
515    }
516    let request = request_builder
517        .extra_headers(extra_headers)
518        .body(AsyncBody::default())
519        .map_err(|error| LanguageModelCompletionError::BuildRequestBody {
520            provider: PROVIDER_NAME,
521            error,
522        })?;
523    let mut response =
524        client
525            .send(request)
526            .await
527            .map_err(|error| LanguageModelCompletionError::HttpSend {
528                provider: PROVIDER_NAME,
529                error,
530            })?;
531
532    let mut body = String::new();
533    response
534        .body_mut()
535        .read_to_string(&mut body)
536        .await
537        .map_err(|error| LanguageModelCompletionError::ApiReadResponseError {
538            provider: PROVIDER_NAME,
539            error,
540        })?;
541
542    if !response.status().is_success() {
543        return Err(LanguageModelCompletionError::from_http_status(
544            PROVIDER_NAME,
545            response.status(),
546            extract_error_message(&body),
547            None,
548        ));
549    }
550
551    let response: ModelsResponse = serde_json::from_str(&body).map_err(|error| {
552        LanguageModelCompletionError::DeserializeResponse {
553            provider: PROVIDER_NAME,
554            error,
555        }
556    })?;
557
558    let mut models = Vec::new();
559    for model in response.data {
560        if let Some(model_type) = model.r#type.as_deref()
561            && model_type != "language"
562        {
563            continue;
564        }
565        let supports_tools = model
566            .supported_parameters
567            .iter()
568            .any(|parameter| parameter == "tools")
569            || has_tag(&model.tags, "tool-use")
570            || has_tag(&model.tags, "tools");
571        let supports_images = model.architecture.is_some_and(|architecture| {
572            architecture
573                .input_modalities
574                .iter()
575                .any(|modality| modality == "image")
576        }) || has_tag(&model.tags, "vision")
577            || has_tag(&model.tags, "image-input");
578        let parallel_tool_calls = model
579            .supported_parameters
580            .iter()
581            .any(|parameter| parameter == "parallel_tool_calls");
582        let prompt_cache_key = model
583            .supported_parameters
584            .iter()
585            .any(|parameter| parameter == "prompt_cache_key" || parameter == "cache_control");
586        models.push(AvailableModel {
587            name: model.id.clone(),
588            display_name: model.name.or(Some(model.id)),
589            max_tokens: model.context_window.or(model.max_tokens).unwrap_or(128_000),
590            max_output_tokens: model.max_tokens,
591            max_completion_tokens: None,
592            capabilities: ModelCapabilities {
593                tools: supports_tools,
594                images: supports_images,
595                parallel_tool_calls,
596                prompt_cache_key,
597                chat_completions: true,
598                interleaved_reasoning: false,
599                max_tokens_parameter: false,
600            },
601        });
602    }
603
604    Ok(models)
605}
606