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(
82 ctx: &PredictCtx,
83 config: &PredictConfig,
84 agent_id: Option<&str>,
85) -> String {
86 if let Some(id) = agent_id.filter(|s| !s.is_empty()) {
89 if let Some(m) = ctx
90 .host
91 .agent_bound_model(id)
92 .await
93 .filter(|m| !m.trim().is_empty())
94 {
95 return m;
96 }
97 }
98 let configured = config.model.trim();
99 if !configured.is_empty() {
100 return configured.to_string();
101 }
102 for var in ["RYU_PREDICT_MODEL", "RYU_DEFAULT_LLM_MODEL"] {
103 if let Ok(val) = std::env::var(var) {
104 if !val.is_empty() {
105 return val;
106 }
107 }
108 }
109 ctx.host.default_model()
110}
111
112#[utoipa::path(
114 get,
115 path = "/api/predict/config",
116 tag = "Predict",
117 summary = "Get predictive-typing config",
118 responses((status = 200, description = "OK", body = serde_json::Value))
119)]
120pub async fn get_config(State(ctx): State<PredictCtx>) -> Json<PredictConfig> {
121 Json(load_config(&ctx).await)
122}
123
124#[utoipa::path(
126 put,
127 path = "/api/predict/config",
128 tag = "Predict",
129 summary = "Update predictive-typing config",
130 request_body = serde_json::Value,
131 responses((status = 200, description = "OK", body = serde_json::Value))
132)]
133pub async fn put_config(
134 State(ctx): State<PredictCtx>,
135 Json(config): Json<PredictConfig>,
136) -> axum::response::Response {
137 let raw = match serde_json::to_string(&config) {
138 Ok(s) => s,
139 Err(e) => {
140 return json_error(StatusCode::BAD_REQUEST, format!("invalid config: {e}"));
141 }
142 };
143 match ctx.host.pref_set(PREDICT_CONFIG_PREF, &raw).await {
144 Ok(()) => Json(config).into_response(),
145 Err(e) => json_error(
146 StatusCode::INTERNAL_SERVER_ERROR,
147 format!("failed to persist config: {e}"),
148 ),
149 }
150}
151
152#[derive(Debug, Deserialize)]
157pub struct CompleteBody {
158 #[serde(default)]
159 pub context: String,
160 #[serde(default)]
161 pub app: Option<String>,
162 #[serde(default)]
163 pub control: Option<String>,
164 #[serde(default, rename = "agentId")]
165 pub agent_id: Option<String>,
166}
167
168#[derive(Debug, Serialize)]
172pub struct CompleteResponse {
173 pub suggestion: String,
174 pub model: String,
175 #[serde(skip_serializing_if = "Option::is_none")]
176 pub reason: Option<String>,
177}
178
179fn refused(reason: &str) -> Json<CompleteResponse> {
180 Json(CompleteResponse {
181 suggestion: String::new(),
182 model: String::new(),
183 reason: Some(reason.to_string()),
184 })
185}
186
187#[utoipa::path(
192 post,
193 path = "/api/predict/complete",
194 tag = "Predict",
195 summary = "Predict the inline continuation for caret context",
196 request_body = serde_json::Value,
197 responses((status = 200, description = "OK", body = serde_json::Value))
198)]
199pub async fn complete(
200 State(ctx): State<PredictCtx>,
201 Json(body): Json<CompleteBody>,
202) -> axum::response::Response {
203 if !ctx.host.is_enabled() {
208 return refused("predictive typing plugin is disabled").into_response();
209 }
210
211 let config = load_config(&ctx).await;
212
213 if let Some(control) = body.control.as_deref() {
215 if crate::is_secure_control(control) {
216 return refused("secure field").into_response();
217 }
218 }
219
220 if let Some(app) = body.app.as_deref() {
222 if !crate::app_allowed(&config.app_allowlist, app) {
223 return refused("app not in allowlist").into_response();
224 }
225 }
226
227 let context = body.context.trim();
228 if context.is_empty() {
229 return refused("no context").into_response();
230 }
231
232 let agent_id = body.agent_id.as_deref().or(config.agent_id.as_deref());
233 let model = resolve_model(&ctx, &config, agent_id).await;
234 let (system, user) = crate::build_messages(context);
235
236 match ctx
237 .host
238 .call_side_model(&model, config.effort.trim(), &system, &user)
239 .await
240 {
241 Ok(text) => {
242 let suggestion = crate::clean_suggestion(&text, config.max_chars);
243 Json(CompleteResponse {
244 suggestion,
245 model,
246 reason: None,
247 })
248 .into_response()
249 }
250 Err(e) => json_error(
251 StatusCode::BAD_GATEWAY,
252 format!("prediction model unavailable: {e}"),
253 ),
254 }
255}