1use roder_api::catalog::{PROVIDER_KIMI_CODE, models_for_provider};
2use roder_api::inference::{
3 AgentInferenceRequest, InferenceCapabilities, InferenceEngine, InferenceEventStream,
4 InferenceProviderContext, InferenceProviderMetadata, InferenceTurnContext,
5 ProviderAuthType,
6};
7use roder_ext_openai_chat_completions::{ChatCompletionsRequestConfig, stream_chat_completions};
8
9use crate::auth::{
10 DEFAULT_MANAGED_BASE_URL, DEFAULT_OPEN_PLATFORM_BASE_URL, access_token, has_stored_tokens,
11 inference_headers, managed_base_url,
12};
13
14#[derive(Debug, Clone, Default)]
15pub struct KimiCodeConfig {
16 pub api_key: Option<String>,
17 pub base_url: Option<String>,
18}
19
20#[derive(Debug, Clone)]
21pub struct KimiCodeProviderSpec {
22 pub provider_id: &'static str,
23 pub name: &'static str,
24 pub description: &'static str,
25 pub default_managed_base_url: &'static str,
26 pub default_open_platform_base_url: &'static str,
27 pub sort_order: i32,
28 pub api_key_env: &'static str,
29 pub api_key_aliases: &'static [&'static str],
30 pub base_url_env: &'static str,
31 pub base_url_aliases: &'static [&'static str],
32}
33
34impl Default for KimiCodeProviderSpec {
35 fn default() -> Self {
36 Self {
37 provider_id: PROVIDER_KIMI_CODE,
38 name: "Kimi Code",
39 description: "Kimi Code (Moonshot AI) subscription inference (direct Kimi Code route).",
40 default_managed_base_url: DEFAULT_MANAGED_BASE_URL,
41 default_open_platform_base_url: DEFAULT_OPEN_PLATFORM_BASE_URL,
42 sort_order: 25,
43 api_key_env: "KIMI_CODE_API_KEY",
44 api_key_aliases: &["RODER_KIMI_CODE_API_KEY"],
45 base_url_env: "RODER_KIMI_CODE_BASE_URL",
46 base_url_aliases: &["KIMI_CODE_BASE_URL"],
47 }
48 }
49}
50
51enum KimiAuth {
52 ApiKey { key: String },
53 OAuth { token: String },
54}
55
56pub struct KimiCodeInferenceEngine {
57 config: KimiCodeConfig,
58 spec: KimiCodeProviderSpec,
59}
60
61impl KimiCodeInferenceEngine {
62 pub fn new(config: KimiCodeConfig, spec: KimiCodeProviderSpec) -> Self {
63 Self { config, spec }
64 }
65
66 fn configured_base_url(&self) -> Option<String> {
67 self.config
68 .base_url
69 .clone()
70 .or_else(|| std::env::var(self.spec.base_url_env).ok())
71 .or_else(|| {
72 for alias in self.spec.base_url_aliases {
73 if let Ok(v) = std::env::var(alias) {
74 return Some(v);
75 }
76 }
77 None
78 })
79 }
80
81 fn base_url_for(&self, auth: &KimiAuth) -> String {
82 self.configured_base_url().unwrap_or_else(|| match auth {
83 KimiAuth::ApiKey { .. } => self.spec.default_open_platform_base_url.to_string(),
84 KimiAuth::OAuth { .. } => managed_base_url(),
85 })
86 }
87
88 fn api_key(&self) -> Option<String> {
89 self.config
90 .api_key
91 .clone()
92 .or_else(|| std::env::var(self.spec.api_key_env).ok())
93 .or_else(|| {
94 for alias in self.spec.api_key_aliases {
95 if let Ok(v) = std::env::var(alias) {
96 return Some(v);
97 }
98 }
99 None
100 })
101 }
102
103 async fn resolve_auth(&self) -> anyhow::Result<KimiAuth> {
104 if let Some(api_key) = self.api_key() {
105 return Ok(KimiAuth::ApiKey { key: api_key });
106 }
107 if let Some(access_token) = access_token().await? {
108 return Ok(KimiAuth::OAuth { token: access_token });
109 }
110 anyhow::bail!(
111 "{} auth is missing; run `roder auth login kimi-code` or set {} / {}",
112 self.spec.name,
113 self.spec.api_key_env,
114 self.spec
115 .api_key_aliases
116 .first()
117 .copied()
118 .unwrap_or("RODER_KIMI_CODE_API_KEY")
119 )
120 }
121}
122
123#[async_trait::async_trait]
124impl InferenceEngine for KimiCodeInferenceEngine {
125 fn id(&self) -> roder_api::extension::InferenceEngineId {
126 self.spec.provider_id.to_string()
127 }
128
129 fn capabilities(&self) -> InferenceCapabilities {
130 InferenceCapabilities {
131 streaming: true,
132 tool_calls: true,
133 parallel_tool_calls: true,
134 reasoning_summaries: false,
135 structured_output: true,
136 image_input: false,
137 prompt_cache: false,
138 provider_metadata: true,
139 tool_search: false,
140 }
141 }
142
143 fn metadata(&self) -> InferenceProviderMetadata {
144 InferenceProviderMetadata {
145 name: self.spec.name.to_string(),
146 description: Some(self.spec.description.to_string()),
147 auth_type: ProviderAuthType::OAuth,
148 auth_label: Some("Kimi Code subscription or API key".to_string()),
149 auth_configured: Some(self.api_key().is_some() || has_stored_tokens()),
150 recommended: true,
151 sort_order: self.spec.sort_order,
152 }
153 }
154
155 async fn list_models(
156 &self,
157 _ctx: InferenceProviderContext<'_>,
158 ) -> anyhow::Result<Vec<roder_api::inference::ModelDescriptor>> {
159 Ok(models_for_provider(self.spec.provider_id, false))
160 }
161
162 async fn stream_turn(
163 &self,
164 _ctx: InferenceTurnContext<'_>,
165 request: AgentInferenceRequest,
166 ) -> anyhow::Result<InferenceEventStream> {
167 let auth = self.resolve_auth().await?;
168 let base_url = self.base_url_for(&auth);
169 let mut config = match &auth {
170 KimiAuth::ApiKey { key } => {
171 ChatCompletionsRequestConfig::bearer(self.spec.name.to_string(), base_url, key)
172 }
173 KimiAuth::OAuth { token } => {
174 ChatCompletionsRequestConfig::bearer(self.spec.name.to_string(), base_url, token)
175 }
176 };
177 if matches!(auth, KimiAuth::OAuth { .. }) {
178 config.headers = inference_headers()?;
179 }
180 stream_chat_completions(config, request).await
181 }
182}
183
184#[cfg(test)]
185mod tests {
186 use super::*;
187
188 #[test]
189 fn metadata_reports_oauth_with_api_key_or_stored_tokens() {
190 let engine = KimiCodeInferenceEngine::new(
191 KimiCodeConfig {
192 api_key: Some("test-key".to_string()),
193 base_url: None,
194 },
195 KimiCodeProviderSpec::default(),
196 );
197 let metadata = engine.metadata();
198 assert_eq!(metadata.auth_type, ProviderAuthType::OAuth);
199 assert_eq!(metadata.auth_configured, Some(true));
200 }
201
202 #[test]
203 fn oauth_uses_managed_base_url_by_default() {
204 let engine = KimiCodeInferenceEngine::new(KimiCodeConfig::default(), KimiCodeProviderSpec::default());
205 let base_url = engine.base_url_for(&KimiAuth::OAuth {
206 token: "token".to_string(),
207 });
208 assert_eq!(base_url, DEFAULT_MANAGED_BASE_URL);
209 }
210
211 #[test]
212 fn api_key_uses_open_platform_base_url_by_default() {
213 let engine = KimiCodeInferenceEngine::new(KimiCodeConfig::default(), KimiCodeProviderSpec::default());
214 let base_url = engine.base_url_for(&KimiAuth::ApiKey {
215 key: "key".to_string(),
216 });
217 assert_eq!(base_url, DEFAULT_OPEN_PLATFORM_BASE_URL);
218 }
219}