use async_trait::async_trait;
use serde::{Deserialize, Serialize};
pub mod api;
pub use api::{routes, PredictCtx};
pub const PREDICT_CONFIG_PREF: &str = "predict-config";
#[async_trait]
pub trait PredictHost: Send + Sync {
fn is_enabled(&self) -> bool;
async fn pref_get(&self, key: &str) -> Option<String>;
async fn pref_set(&self, key: &str, value: &str) -> Result<(), String>;
async fn agent_bound_model(&self, agent_id: &str) -> Option<String>;
fn default_model(&self) -> String;
async fn call_side_model(
&self,
model: &str,
effort: &str,
system: &str,
user: &str,
) -> Result<String, String>;
}
pub const DEFAULT_DEBOUNCE_MS: u64 = 400;
pub const DEFAULT_MAX_CHARS: usize = 240;
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
#[serde(rename_all = "camelCase")]
pub struct PredictConfig {
#[serde(default)]
pub model: String,
#[serde(default)]
pub effort: String,
#[serde(default, rename = "agentId", skip_serializing_if = "Option::is_none")]
pub agent_id: Option<String>,
#[serde(default, rename = "appAllowlist")]
pub app_allowlist: Vec<String>,
#[serde(default = "default_debounce", rename = "debounceMs")]
pub debounce_ms: u64,
#[serde(default = "default_max_chars", rename = "maxChars")]
pub max_chars: usize,
}
fn default_debounce() -> u64 {
DEFAULT_DEBOUNCE_MS
}
fn default_max_chars() -> usize {
DEFAULT_MAX_CHARS
}
impl Default for PredictConfig {
fn default() -> Self {
Self {
model: String::new(),
effort: String::new(),
agent_id: None,
app_allowlist: Vec::new(),
debounce_ms: DEFAULT_DEBOUNCE_MS,
max_chars: DEFAULT_MAX_CHARS,
}
}
}
impl PredictConfig {
pub fn from_pref(raw: Option<&str>) -> Self {
raw.and_then(|s| serde_json::from_str::<PredictConfig>(s).ok())
.unwrap_or_default()
}
}
const SECURE_CONTROL_MARKERS: &[&str] = &[
"password", "passwd", "secure", "pin", "otp", "cvv", "secret",
];
pub fn is_secure_control(control: &str) -> bool {
let lower = control.to_lowercase();
SECURE_CONTROL_MARKERS.iter().any(|m| lower.contains(m))
}
pub fn app_allowed(allowlist: &[String], app: &str) -> bool {
if allowlist.is_empty() {
return true;
}
let name = app
.rsplit(['\\', '/'])
.next()
.unwrap_or(app)
.trim()
.to_lowercase();
if name.is_empty() {
return false;
}
let name_stem = name.trim_end_matches(".exe");
allowlist.iter().any(|entry| {
let e = entry.trim().to_lowercase();
!e.is_empty() && e.trim_end_matches(".exe") == name_stem
})
}
pub fn build_messages(context: &str) -> (String, String) {
let system = "You are an inline autocomplete engine, like GitHub Copilot but for any text \
field. Predict the immediate continuation of the user's text from the context before their cursor. \
Rules:\n\
- Output ONLY the continuation text — never repeat the context, never explain.\n\
- Continue naturally, up to roughly the next clause or sentence.\n\
- Match the existing style, tone, and language.\n\
- Do not start a new line or block; continue in place.\n\
- If you cannot confidently continue, output exactly: 0"
.to_string();
let user = format!(
"Continue the text after the cursor. Text before the cursor:\n\"\"\"\n{context}\n\"\"\""
);
(system, user)
}
pub fn clean_suggestion(raw: &str, max_chars: usize) -> String {
let mut s = raw.trim().to_string();
if s == "0" {
return String::new();
}
if let Some(rest) = s.strip_prefix("```") {
s = rest.to_string();
if let Some(idx) = s.find('\n') {
s = s[idx + 1..].to_string();
}
if let Some(idx) = s.rfind("```") {
s = s[..idx].to_string();
}
s = s.trim().to_string();
}
for (open, close) in [('"', '"'), ('\'', '\''), ('“', '”')] {
if s.starts_with(open) && s.ends_with(close) && s.chars().count() >= 2 {
let inner: String = s.chars().skip(1).take(s.chars().count() - 2).collect();
s = inner.trim().to_string();
}
}
if let Some(idx) = s.find(['\n', '\r']) {
s = s[..idx].to_string();
}
let s = s.trim().to_string();
if s == "0" || s.is_empty() {
return String::new();
}
if s.chars().count() > max_chars {
return s.chars().take(max_chars).collect::<String>();
}
s
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn secure_controls_are_refused() {
assert!(is_secure_control("Password"));
assert!(is_secure_control("password edit"));
assert!(is_secure_control("Secure Text Field"));
assert!(is_secure_control("OTP"));
assert!(!is_secure_control("edit"));
assert!(!is_secure_control("document"));
assert!(!is_secure_control("text box"));
}
#[test]
fn empty_allowlist_permits_all() {
assert!(app_allowed(&[], "notepad.exe"));
assert!(app_allowed(&[], "C:\\x\\chrome.exe"));
}
#[test]
fn allowlist_matches_basename_case_insensitive() {
let allow = vec!["Notepad.exe".to_string(), "chrome".to_string()];
assert!(app_allowed(&allow, "notepad.exe"));
assert!(app_allowed(&allow, "C:\\Windows\\System32\\notepad.exe"));
assert!(app_allowed(&allow, "chrome.exe"));
assert!(!app_allowed(&allow, "code.exe"));
assert!(!app_allowed(&allow, ""));
}
#[test]
fn cleans_sentinel_and_empty() {
assert_eq!(clean_suggestion("0", 240), "");
assert_eq!(clean_suggestion(" ", 240), "");
assert_eq!(clean_suggestion("0\n", 240), "");
}
#[test]
fn strips_quotes_and_collapses_to_one_line() {
assert_eq!(clean_suggestion("\" world\"", 240), "world");
assert_eq!(clean_suggestion("hello\nthere", 240), "hello");
assert_eq!(clean_suggestion("```\ncode here\n```", 240), "code here");
}
#[test]
fn enforces_max_chars() {
let long = "a".repeat(500);
assert_eq!(clean_suggestion(&long, 10).chars().count(), 10);
}
#[test]
fn config_roundtrips_through_pref() {
let cfg = PredictConfig {
model: "gpt-4o-mini".to_string(),
effort: "low".to_string(),
agent_id: Some("ryu".to_string()),
app_allowlist: vec!["notepad.exe".to_string()],
debounce_ms: 250,
max_chars: 120,
};
let raw = serde_json::to_string(&cfg).unwrap();
let back = PredictConfig::from_pref(Some(&raw));
assert_eq!(cfg, back);
}
#[test]
fn missing_pref_is_default() {
let cfg = PredictConfig::from_pref(None);
assert_eq!(cfg.debounce_ms, DEFAULT_DEBOUNCE_MS);
assert_eq!(cfg.max_chars, DEFAULT_MAX_CHARS);
assert!(cfg.app_allowlist.is_empty());
}
#[test]
fn garbage_pref_falls_back_to_default() {
let cfg = PredictConfig::from_pref(Some("not json"));
assert_eq!(cfg, PredictConfig::default());
}
}