1use std::sync::Arc;
21
22use axum::{
23 extract::State,
24 http::StatusCode,
25 response::IntoResponse,
26 routing::{get, post},
27 Json, Router,
28};
29use serde::{Deserialize, Serialize};
30use serde_json::json;
31
32use crate::{PredictConfig, PredictHost, PREDICT_CONFIG_PREF};
33
34#[derive(Clone)]
39pub struct PredictCtx {
40 host: Arc<dyn PredictHost>,
41}
42
43impl PredictCtx {
44 pub fn new(host: Arc<dyn PredictHost>) -> Self {
45 Self { host }
46 }
47}
48
49pub fn routes(ctx: PredictCtx) -> Router<()> {
52 Router::new()
53 .route("/config", get(get_config).put(put_config))
54 .route("/complete", post(complete))
55 .with_state(ctx)
56}
57
58pub fn openapi() -> utoipa::openapi::OpenApi {
60 <PredictApiDoc as utoipa::OpenApi>::openapi()
61}
62
63#[derive(utoipa::OpenApi)]
64#[openapi(paths(complete, get_config, put_config))]
65struct PredictApiDoc;
66
67fn json_error(status: StatusCode, msg: String) -> axum::response::Response {
69 (status, Json(json!({ "error": msg }))).into_response()
70}
71
72async fn load_config(ctx: &PredictCtx) -> PredictConfig {
74 let raw = ctx.host.pref_get(PREDICT_CONFIG_PREF).await;
75 PredictConfig::from_pref(raw.as_deref())
76}
77
78async fn resolve_model(ctx: &PredictCtx, config: &PredictConfig, agent_id: Option<&str>) -> String {
82 if let Some(id) = agent_id.filter(|s| !s.is_empty()) {
85 if let Some(m) = ctx
86 .host
87 .agent_bound_model(id)
88 .await
89 .filter(|m| !m.trim().is_empty())
90 {
91 return m;
92 }
93 }
94 let configured = config.model.trim();
95 if !configured.is_empty() {
96 return configured.to_string();
97 }
98 for var in ["RYU_PREDICT_MODEL", "RYU_DEFAULT_LLM_MODEL"] {
99 if let Ok(val) = std::env::var(var) {
100 if !val.is_empty() {
101 return val;
102 }
103 }
104 }
105 ctx.host.default_model()
106}
107
108#[utoipa::path(
110 get,
111 path = "/api/predict/config",
112 tag = "Predict",
113 summary = "Get predictive-typing config",
114 responses((status = 200, description = "OK", body = serde_json::Value))
115)]
116pub async fn get_config(State(ctx): State<PredictCtx>) -> Json<PredictConfig> {
117 Json(load_config(&ctx).await)
118}
119
120#[utoipa::path(
122 put,
123 path = "/api/predict/config",
124 tag = "Predict",
125 summary = "Update predictive-typing config",
126 request_body = serde_json::Value,
127 responses((status = 200, description = "OK", body = serde_json::Value))
128)]
129pub async fn put_config(
130 State(ctx): State<PredictCtx>,
131 Json(config): Json<PredictConfig>,
132) -> axum::response::Response {
133 let raw = match serde_json::to_string(&config) {
134 Ok(s) => s,
135 Err(e) => {
136 return json_error(StatusCode::BAD_REQUEST, format!("invalid config: {e}"));
137 }
138 };
139 match ctx.host.pref_set(PREDICT_CONFIG_PREF, &raw).await {
140 Ok(()) => Json(config).into_response(),
141 Err(e) => json_error(
142 StatusCode::INTERNAL_SERVER_ERROR,
143 format!("failed to persist config: {e}"),
144 ),
145 }
146}
147
148#[derive(Debug, Deserialize)]
153pub struct CompleteBody {
154 #[serde(default)]
155 pub context: String,
156 #[serde(default)]
157 pub app: Option<String>,
158 #[serde(default)]
159 pub control: Option<String>,
160 #[serde(default, rename = "agentId")]
161 pub agent_id: Option<String>,
162}
163
164#[derive(Debug, Serialize)]
168pub struct CompleteResponse {
169 pub suggestion: String,
170 pub model: String,
171 #[serde(skip_serializing_if = "Option::is_none")]
172 pub reason: Option<String>,
173}
174
175fn refused(reason: &str) -> Json<CompleteResponse> {
176 Json(CompleteResponse {
177 suggestion: String::new(),
178 model: String::new(),
179 reason: Some(reason.to_string()),
180 })
181}
182
183#[utoipa::path(
188 post,
189 path = "/api/predict/complete",
190 tag = "Predict",
191 summary = "Predict the inline continuation for caret context",
192 request_body = serde_json::Value,
193 responses((status = 200, description = "OK", body = serde_json::Value))
194)]
195pub async fn complete(
196 State(ctx): State<PredictCtx>,
197 Json(body): Json<CompleteBody>,
198) -> axum::response::Response {
199 if !ctx.host.is_enabled() {
204 return refused("predictive typing plugin is disabled").into_response();
205 }
206
207 let config = load_config(&ctx).await;
208
209 if let Some(control) = body.control.as_deref() {
211 if crate::is_secure_control(control) {
212 return refused("secure field").into_response();
213 }
214 }
215
216 if let Some(app) = body.app.as_deref() {
218 if !crate::app_allowed(&config.app_allowlist, app) {
219 return refused("app not in allowlist").into_response();
220 }
221 }
222
223 let context = body.context.trim();
224 if context.is_empty() {
225 return refused("no context").into_response();
226 }
227
228 let agent_id = body.agent_id.as_deref().or(config.agent_id.as_deref());
229 let model = resolve_model(&ctx, &config, agent_id).await;
230 let (system, user) = crate::build_messages(context);
231
232 match ctx
233 .host
234 .call_side_model(&model, config.effort.trim(), &system, &user)
235 .await
236 {
237 Ok(text) => {
238 let suggestion = crate::clean_suggestion(&text, config.max_chars);
239 Json(CompleteResponse {
240 suggestion,
241 model,
242 reason: None,
243 })
244 .into_response()
245 }
246 Err(e) => json_error(
247 StatusCode::BAD_GATEWAY,
248 format!("prediction model unavailable: {e}"),
249 ),
250 }
251}