tenant.openagents/omega

1use anyhow::Result;
2use collections::HashMap;
3use credentials_provider::CredentialsProvider;
4use futures::{FutureExt, Stream, StreamExt, future::BoxFuture};
5use gpui::{App, AppContext, AsyncApp, Context, Entity, SharedString, Task};
6use http_client::{CustomHeaders, HttpClient};
7use language_model::{
8    ApiKeyConfiguration, ApiKeyState, AuthenticateError, EnvVar, IconOrSvg, LanguageModel,
9    LanguageModelCompletionError, LanguageModelCompletionEvent, LanguageModelId, LanguageModelName,
10    LanguageModelProvider, LanguageModelProviderId, LanguageModelProviderName,
11    LanguageModelProviderState, LanguageModelRequest, LanguageModelToolChoice,
12    LanguageModelToolResultContent, LanguageModelToolSchemaFormat, LanguageModelToolUse,
13    MessageContent, ProviderSettingsView, RateLimiter, Role, StopReason, TokenUsage, env_var,
14};
15use open_router::{
16    Model, ModelMode as OpenRouterModelMode, OPEN_ROUTER_API_URL, ResponseStreamEvent, list_models,
17};
18use settings::{OpenRouterAvailableModel as AvailableModel, Settings, SettingsStore};
19use std::pin::Pin;
20use std::sync::{Arc, LazyLock};
21use ui::IconName;
22
23use language_model::util::{fix_streamed_json, parse_tool_arguments};
24
25const PROVIDER_ID: LanguageModelProviderId = LanguageModelProviderId::new("openrouter");
26const PROVIDER_NAME: LanguageModelProviderName = LanguageModelProviderName::new("OpenRouter");
27
28const API_KEY_ENV_VAR_NAME: &str = "OPENROUTER_API_KEY";
29static API_KEY_ENV_VAR: LazyLock<EnvVar> = env_var!(API_KEY_ENV_VAR_NAME);
30pub(crate) const RESERVED_HEADER_NAMES: &[&str] = &["HTTP-Referer", "X-Title"];
31const MAX_OPEN_ROUTER_SESSION_ID_LENGTH: usize = 256;
32
33#[derive(Default, Clone, Debug, PartialEq)]
34pub struct OpenRouterSettings {
35    pub api_url: String,
36    pub available_models: Vec<AvailableModel>,
37    pub custom_headers: CustomHeaders,
38}
39
40pub struct OpenRouterLanguageModelProvider {
41    http_client: Arc<dyn HttpClient>,
42    state: Entity<State>,
43}
44
45pub struct State {
46    api_key_state: ApiKeyState,
47    credentials_provider: Arc<dyn CredentialsProvider>,
48    http_client: Arc<dyn HttpClient>,
49    available_models: Vec<open_router::Model>,
50    fetch_models_task: Option<Task<Result<(), LanguageModelCompletionError>>>,
51}
52
53impl State {
54    fn is_authenticated(&self) -> bool {
55        self.api_key_state.has_key()
56    }
57
58    fn set_api_key(&mut self, api_key: Option<String>, cx: &mut Context<Self>) -> Task<Result<()>> {
59        let credentials_provider = self.credentials_provider.clone();
60        let api_url = OpenRouterLanguageModelProvider::api_url(cx);
61        let task = self.api_key_state.store(
62            api_url,
63            api_key,
64            |this| &mut this.api_key_state,
65            credentials_provider,
66            cx,
67        );
68
69        cx.spawn(async move |this, cx| {
70            let result = task.await?;
71            this.update(cx, |this, cx| this.restart_fetch_models_task(cx))
72                .ok();
73            Ok(result)
74        })
75    }
76
77    fn authenticate(&mut self, cx: &mut Context<Self>) -> Task<Result<(), AuthenticateError>> {
78        let credentials_provider = self.credentials_provider.clone();
79        let api_url = OpenRouterLanguageModelProvider::api_url(cx);
80        let task = self.api_key_state.load_if_needed(
81            api_url,
82            |this| &mut this.api_key_state,
83            credentials_provider,
84            cx,
85        );
86
87        cx.spawn(async move |this, cx| {
88            let result = task.await;
89            this.update(cx, |this, cx| this.restart_fetch_models_task(cx))
90                .ok();
91            result
92        })
93    }
94
95    fn fetch_models(
96        &mut self,
97        cx: &mut Context<Self>,
98    ) -> Task<Result<(), LanguageModelCompletionError>> {
99        let http_client = self.http_client.clone();
100        let api_url = OpenRouterLanguageModelProvider::api_url(cx);
101        let extra_headers = OpenRouterLanguageModelProvider::settings(cx)
102            .custom_headers
103            .clone();
104        let Some(api_key) = self.api_key_state.key(&api_url) else {
105            return Task::ready(Err(LanguageModelCompletionError::NoApiKey {
106                provider: PROVIDER_NAME,
107            }));
108        };
109        cx.spawn(async move |this, cx| {
110            let models = list_models(http_client.as_ref(), &api_url, &api_key, &extra_headers)
111                .await
112                .map_err(LanguageModelCompletionError::from)?;
113
114            this.update(cx, |this, cx| {
115                this.available_models = models;
116                cx.notify();
117            })
118            .map_err(|e| LanguageModelCompletionError::Other(e))?;
119
120            Ok(())
121        })
122    }
123
124    fn restart_fetch_models_task(&mut self, cx: &mut Context<Self>) {
125        if self.is_authenticated() {
126            let task = self.fetch_models(cx);
127            self.fetch_models_task.replace(task);
128        } else {
129            self.available_models.clear();
130        }
131    }
132}
133
134impl OpenRouterLanguageModelProvider {
135    pub fn new(
136        http_client: Arc<dyn HttpClient>,
137        credentials_provider: Arc<dyn CredentialsProvider>,
138        cx: &mut App,
139    ) -> Self {
140        let state = cx.new(|cx| {
141            cx.observe_global::<SettingsStore>({
142                let mut last_settings = OpenRouterLanguageModelProvider::settings(cx).clone();
143                move |this: &mut State, cx| {
144                    let current_settings = OpenRouterLanguageModelProvider::settings(cx);
145                    let settings_changed = current_settings != &last_settings;
146                    if settings_changed {
147                        last_settings = current_settings.clone();
148                        this.authenticate(cx).detach();
149                        cx.notify();
150                    }
151                }
152            })
153            .detach();
154            State {
155                api_key_state: ApiKeyState::new(Self::api_url(cx), (*API_KEY_ENV_VAR).clone()),
156                credentials_provider,
157                http_client: http_client.clone(),
158                available_models: Vec::new(),
159                fetch_models_task: None,
160            }
161        });
162
163        Self { http_client, state }
164    }
165
166    fn settings(cx: &App) -> &OpenRouterSettings {
167        &crate::AllLanguageModelSettings::get_global(cx).open_router
168    }
169
170    fn api_url(cx: &App) -> SharedString {
171        let api_url = &Self::settings(cx).api_url;
172        if api_url.is_empty() {
173            OPEN_ROUTER_API_URL.into()
174        } else {
175            SharedString::new(api_url.as_str())
176        }
177    }
178
179    fn create_language_model(&self, model: open_router::Model) -> Arc<dyn LanguageModel> {
180        Arc::new(OpenRouterLanguageModel {
181            id: LanguageModelId::from(model.id().to_string()),
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 OpenRouterLanguageModelProvider {
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 OpenRouterLanguageModelProvider {
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::AiOpenRouter)
209    }
210
211    fn default_model(&self, _cx: &App) -> Option<Arc<dyn LanguageModel>> {
212        Some(self.create_language_model(open_router::Model::default()))
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_from_api = self.state.read(cx).available_models.clone();
221        let mut settings_models = Vec::new();
222
223        for model in &Self::settings(cx).available_models {
224            settings_models.push(open_router::Model {
225                name: model.name.clone(),
226                display_name: model.display_name.clone(),
227                max_tokens: model.max_tokens,
228                supports_tools: model.supports_tools,
229                supports_images: model.supports_images,
230                mode: model.mode.unwrap_or_default(),
231                provider: model.provider.clone(),
232            });
233        }
234
235        for settings_model in &settings_models {
236            if let Some(pos) = models_from_api
237                .iter()
238                .position(|m| m.name == settings_model.name)
239            {
240                models_from_api[pos] = settings_model.clone();
241            } else {
242                models_from_api.push(settings_model.clone());
243            }
244        }
245
246        models_from_api
247            .into_iter()
248            .map(|model| self.create_language_model(model))
249            .collect()
250    }
251
252    fn is_authenticated(&self, cx: &App) -> bool {
253        self.state.read(cx).is_authenticated()
254    }
255
256    fn authenticate(&self, cx: &mut App) -> Task<Result<(), AuthenticateError>> {
257        self.state.update(cx, |state, cx| state.authenticate(cx))
258    }
259
260    fn settings_view(&self, cx: &mut App) -> Option<ProviderSettingsView> {
261        let state = self.state.read(cx);
262        Some(ProviderSettingsView::ApiKey(ApiKeyConfiguration::new(
263            state.api_key_state.has_key(),
264            state.api_key_state.is_from_env_var(),
265            state.api_key_state.env_var_name().clone(),
266            "https://openrouter.ai/keys".into(),
267        )))
268    }
269
270    fn set_api_key(&self, api_key: Option<String>, cx: &mut App) -> Task<Result<()>> {
271        self.state
272            .update(cx, |state, cx| state.set_api_key(api_key, cx))
273    }
274}
275
276pub struct OpenRouterLanguageModel {
277    id: LanguageModelId,
278    model: open_router::Model,
279    state: Entity<State>,
280    http_client: Arc<dyn HttpClient>,
281    request_limiter: RateLimiter,
282}
283
284impl OpenRouterLanguageModel {
285    fn stream_completion(
286        &self,
287        request: open_router::Request,
288        cx: &AsyncApp,
289    ) -> BoxFuture<
290        'static,
291        Result<
292            futures::stream::BoxStream<
293                'static,
294                Result<ResponseStreamEvent, open_router::OpenRouterError>,
295            >,
296            LanguageModelCompletionError,
297        >,
298    > {
299        let http_client = self.http_client.clone();
300        let (api_key, api_url, extra_headers) = self.state.read_with(cx, |state, cx| {
301            let api_url = OpenRouterLanguageModelProvider::api_url(cx);
302            let extra_headers = OpenRouterLanguageModelProvider::settings(cx)
303                .custom_headers
304                .clone();
305            (state.api_key_state.key(&api_url), api_url, extra_headers)
306        });
307
308        async move {
309            let Some(api_key) = api_key else {
310                return Err(LanguageModelCompletionError::NoApiKey {
311                    provider: PROVIDER_NAME,
312                });
313            };
314            let request = open_router::stream_completion(
315                http_client.as_ref(),
316                &api_url,
317                &api_key,
318                request,
319                &extra_headers,
320            );
321            request.await.map_err(Into::into)
322        }
323        .boxed()
324    }
325}
326
327impl LanguageModel for OpenRouterLanguageModel {
328    fn id(&self) -> LanguageModelId {
329        self.id.clone()
330    }
331
332    fn name(&self) -> LanguageModelName {
333        LanguageModelName::from(self.model.display_name().to_string())
334    }
335
336    fn provider_id(&self) -> LanguageModelProviderId {
337        PROVIDER_ID
338    }
339
340    fn provider_name(&self) -> LanguageModelProviderName {
341        PROVIDER_NAME
342    }
343
344    fn supports_tools(&self) -> bool {
345        self.model.supports_tool_calls()
346    }
347
348    fn supports_streaming_tools(&self) -> bool {
349        true
350    }
351
352    fn supports_thinking(&self) -> bool {
353        matches!(self.model.mode, OpenRouterModelMode::Thinking { .. })
354    }
355
356    fn tool_input_format(&self) -> LanguageModelToolSchemaFormat {
357        let model_id = self.model.id().trim().to_lowercase();
358        if model_id.contains("gemini") || model_id.contains("grok") {
359            LanguageModelToolSchemaFormat::JsonSchemaSubset
360        } else {
361            LanguageModelToolSchemaFormat::JsonSchema
362        }
363    }
364
365    fn telemetry_id(&self) -> String {
366        format!("openrouter/{}", self.model.id())
367    }
368
369    fn max_token_count(&self) -> u64 {
370        self.model.max_token_count()
371    }
372
373    fn max_output_tokens(&self) -> Option<u64> {
374        self.model.max_output_tokens()
375    }
376
377    fn supports_tool_choice(&self, choice: LanguageModelToolChoice) -> bool {
378        match choice {
379            LanguageModelToolChoice::Auto => true,
380            LanguageModelToolChoice::Any => true,
381            LanguageModelToolChoice::None => true,
382        }
383    }
384
385    fn supports_images(&self) -> bool {
386        self.model.supports_images.unwrap_or(false)
387    }
388
389    fn stream_completion(
390        &self,
391        request: LanguageModelRequest,
392        cx: &AsyncApp,
393    ) -> BoxFuture<
394        'static,
395        Result<
396            futures::stream::BoxStream<
397                'static,
398                Result<LanguageModelCompletionEvent, LanguageModelCompletionError>,
399            >,
400            LanguageModelCompletionError,
401        >,
402    > {
403        let openrouter_request =
404            match into_open_router(request, &self.model, self.max_output_tokens()) {
405                Ok(request) => request,
406                Err(error) => return async move { Err(error.into()) }.boxed(),
407            };
408        let request = self.stream_completion(openrouter_request, cx);
409        let future = self.request_limiter.stream(async move {
410            let response = request.await?;
411            Ok(OpenRouterEventMapper::new().map_stream(response))
412        });
413        async move { Ok(future.await?.boxed()) }.boxed()
414    }
415}
416
417pub fn into_open_router(
418    request: LanguageModelRequest,
419    model: &Model,
420    max_output_tokens: Option<u64>,
421) -> Result<open_router::Request> {
422    if request.contains_custom_tool_input() {
423        anyhow::bail!("OpenRouter does not support custom tools");
424    }
425
426    // Anthropic models via OpenRouter don't accept reasoning_details being echoed back
427    // in requests - it's an output-only field for them. However, Gemini models require
428    // the thought signatures to be echoed back for proper reasoning chain continuity.
429    // Note: OpenRouter's model API provides an `architecture.tokenizer` field (e.g. "Claude",
430    // "Gemini") which could replace this ID prefix check, but since this is the only place
431    // we need this distinction, we're just using this less invasive check instead.
432    // If we ever have a more formal distionction between the models in the future,
433    // we should revise this to use that instead.
434    let is_anthropic_model = model.id().starts_with("anthropic/");
435    let session_id = open_router_session_id(request.thread_id);
436
437    let mut messages = Vec::new();
438    let mut any_message_wants_cache = false;
439    let mut last_cache_message_index: Option<usize> = None;
440
441    for message in request.messages {
442        let mut message_added_content = false;
443        let reasoning_details_for_message = if is_anthropic_model {
444            None
445        } else {
446            message.reasoning_details.clone()
447        };
448
449        let message_wants_cache = message.cache;
450        if message_wants_cache {
451            any_message_wants_cache = true;
452        }
453
454        for content in message.content {
455            match content {
456                MessageContent::Text(text) => {
457                    add_message_content_part(
458                        open_router::MessagePart::Text {
459                            text,
460                            cache_control: None,
461                        },
462                        message.role,
463                        &mut messages,
464                        reasoning_details_for_message.clone(),
465                    );
466                    message_added_content = true;
467                }
468                MessageContent::Thinking { .. } => {}
469                MessageContent::RedactedThinking(_) => {}
470                MessageContent::Compaction(_) => {}
471                MessageContent::Image(image) => {
472                    add_message_content_part(
473                        open_router::MessagePart::Image {
474                            image_url: image.to_base64_url(),
475                        },
476                        message.role,
477                        &mut messages,
478                        reasoning_details_for_message.clone(),
479                    );
480                    message_added_content = true;
481                }
482                MessageContent::ToolUse(tool_use) => {
483                    let input = tool_use.input.as_json().ok_or_else(|| {
484                        anyhow::anyhow!("OpenRouter does not support custom tool calls")
485                    })?;
486                    let tool_call = open_router::ToolCall {
487                        id: tool_use.id.to_string(),
488                        content: open_router::ToolCallContent::Function {
489                            function: open_router::FunctionContent {
490                                name: tool_use.name.to_string(),
491                                arguments: serde_json::to_string(input).unwrap_or_default(),
492                                thought_signature: tool_use.thought_signature.clone(),
493                            },
494                        },
495                    };
496
497                    if let Some(open_router::RequestMessage::Assistant { tool_calls, .. }) =
498                        messages.last_mut()
499                    {
500                        tool_calls.push(tool_call);
501                    } else {
502                        messages.push(open_router::RequestMessage::Assistant {
503                            content: None,
504                            tool_calls: vec![tool_call],
505                            reasoning_details: reasoning_details_for_message.clone(),
506                        });
507                    }
508                    message_added_content = true;
509                }
510                MessageContent::ToolResult(tool_result) => {
511                    let content: Vec<open_router::MessagePart> = tool_result
512                        .content
513                        .iter()
514                        .map(|part| match part {
515                            LanguageModelToolResultContent::Text(text) => {
516                                open_router::MessagePart::Text {
517                                    text: text.to_string(),
518                                    cache_control: None,
519                                }
520                            }
521                            LanguageModelToolResultContent::Image(image) => {
522                                open_router::MessagePart::Image {
523                                    image_url: image.to_base64_url(),
524                                }
525                            }
526                        })
527                        .collect();
528
529                    messages.push(open_router::RequestMessage::Tool {
530                        content: content.into(),
531                        tool_call_id: tool_result.tool_use_id.to_string(),
532                    });
533                    message_added_content = true;
534                }
535            }
536        }
537
538        if message_wants_cache && message_added_content {
539            last_cache_message_index = messages.len().checked_sub(1);
540        }
541    }
542
543    if is_anthropic_model && any_message_wants_cache {
544        // OpenRouter's top-level automatic cache_control restricts routing to
545        // Anthropic direct; explicit block breakpoints also work on Bedrock and Vertex.
546        if let Some(content) = last_cache_message_index
547            .and_then(|index| messages.get_mut(index))
548            .and_then(request_message_content_mut)
549        {
550            set_last_text_cache_control(content, cache_control(None));
551        }
552
553        if let Some(content) = messages.iter_mut().find_map(|message| match message {
554            open_router::RequestMessage::System { content } => Some(content),
555            _ => None,
556        }) {
557            set_last_text_cache_control(
558                content,
559                cache_control(Some(open_router::CacheTtl::OneHour)),
560            );
561        }
562    }
563
564    Ok(open_router::Request {
565        model: model.id().into(),
566        messages,
567        stream: true,
568        session_id,
569        stop: request.stop,
570        temperature: request.temperature.unwrap_or(0.4),
571        max_tokens: max_output_tokens,
572        parallel_tool_calls: if model.supports_parallel_tool_calls() && !request.tools.is_empty() {
573            Some(false)
574        } else {
575            None
576        },
577        usage: open_router::RequestUsage { include: true },
578        reasoning: if request.thinking_allowed
579            && let OpenRouterModelMode::Thinking { budget_tokens } = model.mode
580        {
581            Some(open_router::Reasoning {
582                effort: None,
583                max_tokens: budget_tokens,
584                exclude: Some(false),
585                enabled: Some(true),
586            })
587        } else {
588            None
589        },
590        tools: request
591            .tools
592            .into_iter()
593            .map(|tool| {
594                let input_schema = match tool.input {
595                    language_model::LanguageModelRequestToolInput::Function {
596                        input_schema,
597                        ..
598                    } => input_schema,
599                    language_model::LanguageModelRequestToolInput::Custom { .. } => {
600                        return Err(anyhow::anyhow!("OpenRouter does not support custom tools"));
601                    }
602                };
603                Ok(open_router::ToolDefinition::Function {
604                    function: open_router::FunctionDefinition {
605                        name: tool.name,
606                        description: Some(tool.description),
607                        parameters: Some(input_schema),
608                    },
609                })
610            })
611            .collect::<Result<_>>()?,
612        tool_choice: request.tool_choice.map(|choice| match choice {
613            LanguageModelToolChoice::Auto => open_router::ToolChoice::Auto,
614            LanguageModelToolChoice::Any => open_router::ToolChoice::Required,
615            LanguageModelToolChoice::None => open_router::ToolChoice::None,
616        }),
617        provider: model.provider.clone(),
618    })
619}
620
621fn open_router_session_id(thread_id: Option<String>) -> Option<String> {
622    thread_id.map(|thread_id| {
623        thread_id
624            .chars()
625            .take(MAX_OPEN_ROUTER_SESSION_ID_LENGTH)
626            .collect()
627    })
628}
629
630fn cache_control(ttl: Option<open_router::CacheTtl>) -> open_router::CacheControl {
631    open_router::CacheControl {
632        cache_type: open_router::CacheControlType::Ephemeral,
633        ttl,
634    }
635}
636
637fn request_message_content_mut(
638    message: &mut open_router::RequestMessage,
639) -> Option<&mut open_router::MessageContent> {
640    match message {
641        open_router::RequestMessage::User { content }
642        | open_router::RequestMessage::System { content }
643        | open_router::RequestMessage::Tool { content, .. } => Some(content),
644        open_router::RequestMessage::Assistant {
645            content: Some(content),
646            ..
647        } => Some(content),
648        open_router::RequestMessage::Assistant { content: None, .. } => None,
649    }
650}
651
652fn set_last_text_cache_control(
653    content: &mut open_router::MessageContent,
654    cache_control: open_router::CacheControl,
655) {
656    match content {
657        open_router::MessageContent::Plain(text) => {
658            let text = std::mem::take(text);
659            *content =
660                open_router::MessageContent::Multipart(vec![open_router::MessagePart::Text {
661                    text,
662                    cache_control: Some(cache_control),
663                }]);
664        }
665        open_router::MessageContent::Multipart(parts) => {
666            for part in parts.iter_mut().rev() {
667                if let open_router::MessagePart::Text {
668                    cache_control: target,
669                    ..
670                } = part
671                {
672                    *target = Some(cache_control);
673                    break;
674                }
675            }
676        }
677    }
678}
679
680fn add_message_content_part(
681    new_part: open_router::MessagePart,
682    role: Role,
683    messages: &mut Vec<open_router::RequestMessage>,
684    reasoning_details: Option<Arc<serde_json::Value>>,
685) {
686    match (role, messages.last_mut()) {
687        (Role::User, Some(open_router::RequestMessage::User { content }))
688        | (Role::System, Some(open_router::RequestMessage::System { content })) => {
689            content.push_part(new_part);
690        }
691        (
692            Role::Assistant,
693            Some(open_router::RequestMessage::Assistant {
694                content: Some(content),
695                ..
696            }),
697        ) => {
698            content.push_part(new_part);
699        }
700        _ => {
701            messages.push(match role {
702                Role::User => open_router::RequestMessage::User {
703                    content: open_router::MessageContent::from(vec![new_part]),
704                },
705                Role::Assistant => open_router::RequestMessage::Assistant {
706                    content: Some(open_router::MessageContent::from(vec![new_part])),
707                    tool_calls: Vec::new(),
708                    reasoning_details,
709                },
710                Role::System => open_router::RequestMessage::System {
711                    content: open_router::MessageContent::from(vec![new_part]),
712                },
713            });
714        }
715    }
716}
717
718pub struct OpenRouterEventMapper {
719    tool_calls_by_index: HashMap<usize, RawToolCall>,
720    reasoning_details: Option<serde_json::Value>,
721}
722
723impl OpenRouterEventMapper {
724    pub fn new() -> Self {
725        Self {
726            tool_calls_by_index: HashMap::default(),
727            reasoning_details: None,
728        }
729    }
730
731    pub fn map_stream(
732        mut self,
733        events: Pin<
734            Box<
735                dyn Send + Stream<Item = Result<ResponseStreamEvent, open_router::OpenRouterError>>,
736            >,
737        >,
738    ) -> impl Stream<Item = Result<LanguageModelCompletionEvent, LanguageModelCompletionError>>
739    {
740        events.flat_map(move |event| {
741            futures::stream::iter(match event {
742                Ok(event) => self.map_event(event),
743                Err(error) => vec![Err(error.into())],
744            })
745        })
746    }
747
748    pub fn map_event(
749        &mut self,
750        event: ResponseStreamEvent,
751    ) -> Vec<Result<LanguageModelCompletionEvent, LanguageModelCompletionError>> {
752        let mut events = Vec::new();
753
754        if let Some(usage) = event.usage {
755            let cache_creation_input_tokens = usage
756                .prompt_tokens_details
757                .as_ref()
758                .map_or(0, |details| details.cache_write_tokens);
759            let cache_read_input_tokens = usage
760                .prompt_tokens_details
761                .as_ref()
762                .map_or(0, |details| details.cached_tokens);
763            let input_tokens = usage.prompt_tokens.saturating_sub(
764                cache_creation_input_tokens.saturating_add(cache_read_input_tokens),
765            );
766
767            events.push(Ok(LanguageModelCompletionEvent::UsageUpdate(TokenUsage {
768                input_tokens,
769                output_tokens: usage.completion_tokens,
770                cache_creation_input_tokens,
771                cache_read_input_tokens,
772            })));
773        }
774
775        let Some(choice) = event.choices.first() else {
776            return events;
777        };
778
779        if let Some(details) = choice.delta.reasoning_details.clone() {
780            // Emit reasoning_details immediately
781            events.push(Ok(LanguageModelCompletionEvent::ReasoningDetails(
782                details.clone(),
783            )));
784            self.reasoning_details = Some(details);
785        }
786
787        if let Some(reasoning) = choice.delta.reasoning.clone() {
788            events.push(Ok(LanguageModelCompletionEvent::Thinking {
789                text: reasoning,
790                signature: None,
791            }));
792        }
793
794        if let Some(content) = choice.delta.content.clone() {
795            // OpenRouter send empty content string with the reasoning content
796            // This is a workaround for the OpenRouter API bug
797            if !content.is_empty() {
798                events.push(Ok(LanguageModelCompletionEvent::Text(content)));
799            }
800        }
801
802        if let Some(tool_calls) = choice.delta.tool_calls.as_ref() {
803            for tool_call in tool_calls {
804                let entry = self.tool_calls_by_index.entry(tool_call.index).or_default();
805
806                if let Some(tool_id) = tool_call.id.clone() {
807                    entry.id = tool_id;
808                }
809
810                if let Some(function) = tool_call.function.as_ref() {
811                    if let Some(name) = function.name.clone() {
812                        entry.name = name;
813                    }
814
815                    if let Some(arguments) = function.arguments.clone() {
816                        entry.arguments.push_str(&arguments);
817                    }
818
819                    if let Some(signature) = function.thought_signature.clone() {
820                        entry.thought_signature = Some(signature);
821                    }
822                }
823
824                if !entry.id.is_empty() && !entry.name.is_empty() {
825                    if let Ok(input) = serde_json::from_str::<serde_json::Value>(
826                        &fix_streamed_json(&entry.arguments),
827                    ) {
828                        events.push(Ok(LanguageModelCompletionEvent::ToolUse(
829                            LanguageModelToolUse {
830                                id: entry.id.clone().into(),
831                                name: entry.name.as_str().into(),
832                                is_input_complete: false,
833                                input: language_model::LanguageModelToolUseInput::Json(input),
834                                raw_input: entry.arguments.clone(),
835                                thought_signature: entry.thought_signature.clone(),
836                            },
837                        )));
838                    }
839                }
840            }
841        }
842
843        match choice.finish_reason.as_deref() {
844            Some("stop") => {
845                // Don't emit reasoning_details here - already emitted immediately when captured
846                events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)));
847            }
848            Some("tool_calls") => {
849                events.extend(self.tool_calls_by_index.drain().map(|(_, tool_call)| {
850                    match parse_tool_arguments(&tool_call.arguments) {
851                        Ok(input) => Ok(LanguageModelCompletionEvent::ToolUse(
852                            LanguageModelToolUse {
853                                id: tool_call.id.clone().into(),
854                                name: tool_call.name.as_str().into(),
855                                is_input_complete: true,
856                                input: language_model::LanguageModelToolUseInput::Json(input),
857                                raw_input: tool_call.arguments.clone(),
858                                thought_signature: tool_call.thought_signature.clone(),
859                            },
860                        )),
861                        Err(error) => Ok(LanguageModelCompletionEvent::ToolUseJsonParseError {
862                            id: tool_call.id.clone().into(),
863                            tool_name: tool_call.name.as_str().into(),
864                            raw_input: tool_call.arguments.clone().into(),
865                            json_parse_error: error.to_string(),
866                        }),
867                    }
868                }));
869
870                // Don't emit reasoning_details here - already emitted immediately when captured
871                events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::ToolUse)));
872            }
873            Some(stop_reason) => {
874                log::error!("Unexpected OpenRouter stop_reason: {stop_reason:?}",);
875                // Don't emit reasoning_details here - already emitted immediately when captured
876                events.push(Ok(LanguageModelCompletionEvent::Stop(StopReason::EndTurn)));
877            }
878            None => {}
879        }
880
881        events
882    }
883}
884
885#[derive(Default)]
886struct RawToolCall {
887    id: String,
888    name: String,
889    arguments: String,
890    thought_signature: Option<String>,
891}
892
893#[cfg(test)]
894mod tests {
895    use super::*;
896
897    use open_router::{ChoiceDelta, FunctionChunk, ResponseMessageDelta, ToolCallChunk};
898
899    #[gpui::test]
900    async fn test_reasoning_details_preservation_with_tool_calls() {
901        // This test verifies that reasoning_details are properly captured and preserved
902        // when a model uses tool calling with reasoning/thinking tokens.
903        //
904        // The key regression this prevents:
905        // - OpenRouter sends multiple reasoning_details updates during streaming
906        // - First with actual content (encrypted reasoning data)
907        // - Then with empty array on completion
908        // - We must NOT overwrite the real data with the empty array
909
910        let mut mapper = OpenRouterEventMapper::new();
911
912        // Simulate the streaming events as they come from OpenRouter/Gemini
913        let events = vec![
914            // Event 1: Initial reasoning details with text
915            ResponseStreamEvent {
916                id: Some("response_123".into()),
917                created: 1234567890,
918                model: "google/gemini-3.1-pro-preview".into(),
919                choices: vec![ChoiceDelta {
920                    index: 0,
921                    delta: ResponseMessageDelta {
922                        role: None,
923                        content: None,
924                        reasoning: None,
925                        tool_calls: None,
926                        reasoning_details: Some(serde_json::json!([
927                            {
928                                "type": "reasoning.text",
929                                "text": "Let me analyze this request...",
930                                "format": "google-gemini-v1",
931                                "index": 0
932                            }
933                        ])),
934                    },
935                    finish_reason: None,
936                }],
937                usage: None,
938            },
939            // Event 2: More reasoning details
940            ResponseStreamEvent {
941                id: Some("response_123".into()),
942                created: 1234567890,
943                model: "google/gemini-3.1-pro-preview".into(),
944                choices: vec![ChoiceDelta {
945                    index: 0,
946                    delta: ResponseMessageDelta {
947                        role: None,
948                        content: None,
949                        reasoning: None,
950                        tool_calls: None,
951                        reasoning_details: Some(serde_json::json!([
952                            {
953                                "type": "reasoning.encrypted",
954                                "data": "EtgDCtUDAdHtim9OF5jm4aeZSBAtl/randomized123",
955                                "format": "google-gemini-v1",
956                                "index": 0,
957                                "id": "tool_call_abc123"
958                            }
959                        ])),
960                    },
961                    finish_reason: None,
962                }],
963                usage: None,
964            },
965            // Event 3: Tool call starts
966            ResponseStreamEvent {
967                id: Some("response_123".into()),
968                created: 1234567890,
969                model: "google/gemini-3.1-pro-preview".into(),
970                choices: vec![ChoiceDelta {
971                    index: 0,
972                    delta: ResponseMessageDelta {
973                        role: None,
974                        content: None,
975                        reasoning: None,
976                        tool_calls: Some(vec![ToolCallChunk {
977                            index: 0,
978                            id: Some("tool_call_abc123".into()),
979                            function: Some(FunctionChunk {
980                                name: Some("list_directory".into()),
981                                arguments: Some("{\"path\":\"test\"}".into()),
982                                thought_signature: Some("sha256:test_signature_xyz789".into()),
983                            }),
984                        }]),
985                        reasoning_details: None,
986                    },
987                    finish_reason: None,
988                }],
989                usage: None,
990            },
991            // Event 4: Empty reasoning_details on tool_calls finish
992            // This is the critical event - we must not overwrite with this empty array!
993            ResponseStreamEvent {
994                id: Some("response_123".into()),
995                created: 1234567890,
996                model: "google/gemini-3.1-pro-preview".into(),
997                choices: vec![ChoiceDelta {
998                    index: 0,
999                    delta: ResponseMessageDelta {
1000                        role: None,
1001                        content: None,
1002                        reasoning: None,
1003                        tool_calls: None,
1004                        reasoning_details: Some(serde_json::json!([])),
1005                    },
1006                    finish_reason: Some("tool_calls".into()),
1007                }],
1008                usage: None,
1009            },
1010        ];
1011
1012        // Process all events
1013        let mut collected_events = Vec::new();
1014        for event in events {
1015            let mapped = mapper.map_event(event);
1016            collected_events.extend(mapped);
1017        }
1018
1019        // Verify we got the expected events
1020        let mut has_tool_use = false;
1021        let mut reasoning_details_events = Vec::new();
1022        let mut thought_signature_value = None;
1023
1024        for event_result in collected_events {
1025            match event_result {
1026                Ok(LanguageModelCompletionEvent::ToolUse(tool_use)) => {
1027                    has_tool_use = true;
1028                    assert_eq!(tool_use.id.to_string(), "tool_call_abc123");
1029                    assert_eq!(tool_use.name.as_ref(), "list_directory");
1030                    thought_signature_value = tool_use.thought_signature.clone();
1031                }
1032                Ok(LanguageModelCompletionEvent::ReasoningDetails(details)) => {
1033                    reasoning_details_events.push(details);
1034                }
1035                _ => {}
1036            }
1037        }
1038
1039        // Assertions
1040        assert!(has_tool_use, "Should have emitted ToolUse event");
1041        assert!(
1042            !reasoning_details_events.is_empty(),
1043            "Should have emitted ReasoningDetails events"
1044        );
1045
1046        // We should have received multiple reasoning_details events (text, encrypted, empty)
1047        // The agent layer is responsible for keeping only the first non-empty one
1048        assert!(
1049            reasoning_details_events.len() >= 2,
1050            "Should have multiple reasoning_details events from streaming"
1051        );
1052
1053        // Verify at least one contains the encrypted data
1054        let has_encrypted = reasoning_details_events.iter().any(|details| {
1055            if let serde_json::Value::Array(arr) = details {
1056                arr.iter().any(|item| {
1057                    item["type"] == "reasoning.encrypted"
1058                        && item["data"]
1059                            .as_str()
1060                            .map_or(false, |s| s.contains("EtgDCtUDAdHtim9OF5jm4aeZSBAtl"))
1061                })
1062            } else {
1063                false
1064            }
1065        });
1066        assert!(
1067            has_encrypted,
1068            "Should have at least one reasoning_details with encrypted data"
1069        );
1070
1071        // Verify thought_signature was captured
1072        assert!(
1073            thought_signature_value.is_some(),
1074            "Tool use should have thought_signature"
1075        );
1076        assert_eq!(
1077            thought_signature_value.unwrap(),
1078            "sha256:test_signature_xyz789"
1079        );
1080    }
1081
1082    #[gpui::test]
1083    async fn test_usage_only_chunk_with_empty_choices_does_not_error() {
1084        let mut mapper = OpenRouterEventMapper::new();
1085
1086        let events = mapper.map_event(ResponseStreamEvent {
1087            id: Some("response_123".into()),
1088            created: 1234567890,
1089            model: "google/gemini-3-flash-preview".into(),
1090            choices: Vec::new(),
1091            usage: Some(open_router::Usage {
1092                prompt_tokens: 12,
1093                completion_tokens: 7,
1094                total_tokens: 19,
1095                prompt_tokens_details: Some(open_router::PromptTokensDetails {
1096                    cached_tokens: 5,
1097                    cache_write_tokens: 3,
1098                }),
1099            }),
1100        });
1101
1102        assert_eq!(events.len(), 1);
1103        match events.into_iter().next() {
1104            Some(Ok(LanguageModelCompletionEvent::UsageUpdate(usage))) => {
1105                assert_eq!(usage.input_tokens, 4);
1106                assert_eq!(usage.output_tokens, 7);
1107                assert_eq!(usage.cache_creation_input_tokens, 3);
1108                assert_eq!(usage.cache_read_input_tokens, 5);
1109                assert_eq!(usage.total_tokens(), 19);
1110            }
1111            other => panic!("Expected usage update event, got: {other:?}"),
1112        }
1113    }
1114
1115    #[gpui::test]
1116    async fn test_session_id_uses_thread_id() {
1117        let model = open_router::Model::new(
1118            "openai/gpt-4o",
1119            Some("GPT-4o"),
1120            Some(128000),
1121            Some(true),
1122            Some(false),
1123            None,
1124            None,
1125        );
1126        let expected_session_id = "a".repeat(MAX_OPEN_ROUTER_SESSION_ID_LENGTH);
1127        let request = LanguageModelRequest {
1128            thread_id: Some(format!("{expected_session_id}extra")),
1129            messages: vec![language_model::LanguageModelRequestMessage {
1130                role: Role::User,
1131                content: vec![MessageContent::Text("Hello".to_string())],
1132                cache: false,
1133                reasoning_details: None,
1134            }],
1135            ..Default::default()
1136        };
1137
1138        let result = into_open_router(request, &model, None).unwrap();
1139
1140        assert_eq!(
1141            result.session_id.as_deref(),
1142            Some(expected_session_id.as_str())
1143        );
1144    }
1145
1146    #[gpui::test]
1147    async fn test_agent_prevents_empty_reasoning_details_overwrite() {
1148        // This test verifies that the agent layer prevents empty reasoning_details
1149        // from overwriting non-empty ones, even though the mapper emits all events.
1150
1151        // Simulate what the agent does when it receives multiple ReasoningDetails events
1152        let mut agent_reasoning_details: Option<serde_json::Value> = None;
1153
1154        let events = vec![
1155            // First event: non-empty reasoning_details
1156            serde_json::json!([
1157                {
1158                    "type": "reasoning.encrypted",
1159                    "data": "real_data_here",
1160                    "format": "google-gemini-v1"
1161                }
1162            ]),
1163            // Second event: empty array (should not overwrite)
1164            serde_json::json!([]),
1165        ];
1166
1167        for details in events {
1168            // This mimics the agent's logic: only store if we don't already have it
1169            if agent_reasoning_details.is_none() {
1170                agent_reasoning_details = Some(details);
1171            }
1172        }
1173
1174        // Verify the agent kept the first non-empty reasoning_details
1175        assert!(agent_reasoning_details.is_some());
1176        let final_details = agent_reasoning_details.unwrap();
1177        if let serde_json::Value::Array(arr) = &final_details {
1178            assert!(
1179                !arr.is_empty(),
1180                "Agent should have kept the non-empty reasoning_details"
1181            );
1182            assert_eq!(arr[0]["data"], "real_data_here");
1183        } else {
1184            panic!("Expected array");
1185        }
1186    }
1187
1188    #[gpui::test]
1189    async fn test_anthropic_model_caching_two_tier() {
1190        let model = open_router::Model::new(
1191            "anthropic/claude-sonnet-4-5",
1192            Some("Claude Sonnet"),
1193            Some(200000),
1194            Some(true),
1195            Some(false),
1196            None,
1197            None,
1198        );
1199
1200        let request = LanguageModelRequest {
1201            messages: vec![
1202                language_model::LanguageModelRequestMessage {
1203                    role: Role::System,
1204                    content: vec![MessageContent::Text("You are helpful.".to_string())],
1205                    cache: false,
1206                    reasoning_details: None,
1207                },
1208                language_model::LanguageModelRequestMessage {
1209                    role: Role::User,
1210                    content: vec![MessageContent::Text("Hello".to_string())],
1211                    cache: false,
1212                    reasoning_details: None,
1213                },
1214                language_model::LanguageModelRequestMessage {
1215                    role: Role::Assistant,
1216                    content: vec![MessageContent::Text("Hi there!".to_string())],
1217                    cache: false,
1218                    reasoning_details: None,
1219                },
1220                language_model::LanguageModelRequestMessage {
1221                    role: Role::User,
1222                    content: vec![MessageContent::Text("What is 2+2?".to_string())],
1223                    cache: true,
1224                    reasoning_details: None,
1225                },
1226            ],
1227            stop: vec![],
1228            temperature: None,
1229            tools: vec![],
1230            tool_choice: None,
1231            thinking_allowed: false,
1232            thinking_effort: None,
1233            speed: None,
1234            thread_id: None,
1235            prompt_id: None,
1236            intent: None,
1237            compact_at_tokens: None,
1238        };
1239
1240        let result = into_open_router(request, &model, None).unwrap();
1241
1242        let system_cache = result.messages.iter().find_map(|m| {
1243            if let open_router::RequestMessage::System { content } = m {
1244                if let open_router::MessageContent::Multipart(parts) = content {
1245                    parts.iter().last().and_then(|p| {
1246                        if let open_router::MessagePart::Text { cache_control, .. } = p {
1247                            *cache_control
1248                        } else {
1249                            None
1250                        }
1251                    })
1252                } else {
1253                    None
1254                }
1255            } else {
1256                None
1257            }
1258        });
1259        assert!(
1260            matches!(
1261                system_cache,
1262                Some(open_router::CacheControl {
1263                    cache_type: open_router::CacheControlType::Ephemeral,
1264                    ttl: Some(open_router::CacheTtl::OneHour),
1265                })
1266            ),
1267            "System message should have 1h cache_control, got: {system_cache:?}"
1268        );
1269
1270        let tail_cache = result.messages.last().and_then(|last_message| {
1271            if let open_router::RequestMessage::User { content } = last_message {
1272                if let open_router::MessageContent::Multipart(parts) = content {
1273                    parts.iter().last().and_then(|part| {
1274                        if let open_router::MessagePart::Text { cache_control, .. } = part {
1275                            *cache_control
1276                        } else {
1277                            None
1278                        }
1279                    })
1280                } else {
1281                    None
1282                }
1283            } else {
1284                None
1285            }
1286        });
1287        assert!(
1288            matches!(
1289                tail_cache,
1290                Some(open_router::CacheControl {
1291                    cache_type: open_router::CacheControlType::Ephemeral,
1292                    ttl: None,
1293                })
1294            ),
1295            "Last cache:true message should have 5min cache_control, got: {tail_cache:?}"
1296        );
1297
1298        for (i, message) in result.messages.iter().enumerate() {
1299            let is_system = matches!(message, open_router::RequestMessage::System { .. });
1300            let is_last = i == result.messages.len() - 1;
1301            if is_system || is_last {
1302                continue;
1303            }
1304            let parts: Option<&Vec<open_router::MessagePart>> = match message {
1305                open_router::RequestMessage::User { content }
1306                | open_router::RequestMessage::System { content }
1307                | open_router::RequestMessage::Tool { content, .. } => {
1308                    if let open_router::MessageContent::Multipart(parts) = content {
1309                        Some(parts)
1310                    } else {
1311                        None
1312                    }
1313                }
1314                open_router::RequestMessage::Assistant {
1315                    content: Some(content),
1316                    ..
1317                } => {
1318                    if let open_router::MessageContent::Multipart(parts) = content {
1319                        Some(parts)
1320                    } else {
1321                        None
1322                    }
1323                }
1324                _ => None,
1325            };
1326            if let Some(parts) = parts {
1327                for part in parts {
1328                    if let open_router::MessagePart::Text { cache_control, .. } = part {
1329                        assert!(
1330                            cache_control.is_none(),
1331                            "Message {i} should not have cache_control"
1332                        );
1333                    }
1334                }
1335            }
1336        }
1337    }
1338
1339    #[gpui::test]
1340    async fn test_anthropic_model_no_cache_when_no_cache_flag() {
1341        let model = open_router::Model::new(
1342            "anthropic/claude-sonnet-4-5",
1343            Some("Claude Sonnet"),
1344            Some(200000),
1345            Some(true),
1346            Some(false),
1347            None,
1348            None,
1349        );
1350
1351        let request = LanguageModelRequest {
1352            messages: vec![
1353                language_model::LanguageModelRequestMessage {
1354                    role: Role::System,
1355                    content: vec![MessageContent::Text("You are helpful.".to_string())],
1356                    cache: false,
1357                    reasoning_details: None,
1358                },
1359                language_model::LanguageModelRequestMessage {
1360                    role: Role::User,
1361                    content: vec![MessageContent::Text("Hello".to_string())],
1362                    cache: false,
1363                    reasoning_details: None,
1364                },
1365            ],
1366            stop: vec![],
1367            temperature: None,
1368            tools: vec![],
1369            tool_choice: None,
1370            thinking_allowed: false,
1371            thinking_effort: None,
1372            speed: None,
1373            thread_id: None,
1374            prompt_id: None,
1375            intent: None,
1376            compact_at_tokens: None,
1377        };
1378
1379        let result = into_open_router(request, &model, None).unwrap();
1380
1381        for message in &result.messages {
1382            let content = match message {
1383                open_router::RequestMessage::User { content }
1384                | open_router::RequestMessage::System { content } => Some(content),
1385                _ => None,
1386            };
1387            if let Some(content) = content {
1388                if let open_router::MessageContent::Multipart(parts) = content {
1389                    for part in parts {
1390                        if let open_router::MessagePart::Text { cache_control, .. } = part {
1391                            assert!(
1392                                cache_control.is_none(),
1393                                "No message should have cache_control when no cache:true flags"
1394                            );
1395                        }
1396                    }
1397                }
1398            }
1399        }
1400    }
1401
1402    #[gpui::test]
1403    async fn test_non_anthropic_model_no_cache_control() {
1404        let model = open_router::Model::new(
1405            "openai/gpt-4o",
1406            Some("GPT-4o"),
1407            Some(128000),
1408            Some(true),
1409            Some(false),
1410            None,
1411            None,
1412        );
1413
1414        let request = LanguageModelRequest {
1415            messages: vec![
1416                language_model::LanguageModelRequestMessage {
1417                    role: Role::System,
1418                    content: vec![MessageContent::Text("You are helpful.".to_string())],
1419                    cache: false,
1420                    reasoning_details: None,
1421                },
1422                language_model::LanguageModelRequestMessage {
1423                    role: Role::User,
1424                    content: vec![MessageContent::Text("Hello".to_string())],
1425                    cache: true,
1426                    reasoning_details: None,
1427                },
1428            ],
1429            stop: vec![],
1430            temperature: None,
1431            tools: vec![],
1432            tool_choice: None,
1433            thinking_allowed: false,
1434            thinking_effort: None,
1435            speed: None,
1436            thread_id: None,
1437            prompt_id: None,
1438            intent: None,
1439            compact_at_tokens: None,
1440        };
1441
1442        let result = into_open_router(request, &model, None).unwrap();
1443
1444        for message in &result.messages {
1445            let content = match message {
1446                open_router::RequestMessage::User { content }
1447                | open_router::RequestMessage::System { content } => Some(content),
1448                _ => None,
1449            };
1450            if let Some(content) = content {
1451                if let open_router::MessageContent::Multipart(parts) = content {
1452                    for part in parts {
1453                        if let open_router::MessagePart::Text { cache_control, .. } = part {
1454                            assert!(
1455                                cache_control.is_none(),
1456                                "Non-Anthropic model should never have cache_control"
1457                            );
1458                        }
1459                    }
1460                }
1461            }
1462        }
1463    }
1464}
1465