From 5923165fcfbe23386b6ba7be6bf2000714b41fee Mon Sep 17 00:00:00 2001 From: Dave Date: Mon, 16 Feb 2026 16:24:21 +0000 Subject: [PATCH] Refactoring the structure a bit --- server/src/http/anthropic.rs | 18 ++ server/src/http/assets.rs | 64 +++++ server/src/http/chat.rs | 8 + server/src/http/context.rs | 16 ++ server/src/http/fs.rs | 34 +++ server/src/http/health.rs | 6 + server/src/http/mod.rs | 34 +++ server/src/http/model.rs | 27 ++ server/src/http/payloads.rs | 39 +++ server/src/http/project.rs | 24 ++ server/src/http/rest.rs | 125 +++++++++ server/src/http/search.rs | 13 + server/src/http/shell.rs | 13 + server/src/http/ws.rs | 92 +++++++ server/src/{commands => io}/fs.rs | 0 server/src/{commands => io}/mod.rs | 1 - server/src/{commands => io}/search.rs | 0 server/src/{commands => io}/shell.rs | 0 server/src/{commands => llm}/chat.rs | 2 +- server/src/llm/mod.rs | 1 + server/src/main.rs | 377 +------------------------- 21 files changed, 522 insertions(+), 372 deletions(-) create mode 100644 server/src/http/anthropic.rs create mode 100644 server/src/http/assets.rs create mode 100644 server/src/http/chat.rs create mode 100644 server/src/http/context.rs create mode 100644 server/src/http/fs.rs create mode 100644 server/src/http/health.rs create mode 100644 server/src/http/mod.rs create mode 100644 server/src/http/model.rs create mode 100644 server/src/http/payloads.rs create mode 100644 server/src/http/project.rs create mode 100644 server/src/http/rest.rs create mode 100644 server/src/http/search.rs create mode 100644 server/src/http/shell.rs create mode 100644 server/src/http/ws.rs rename server/src/{commands => io}/fs.rs (100%) rename server/src/{commands => io}/mod.rs (75%) rename server/src/{commands => io}/search.rs (100%) rename server/src/{commands => io}/shell.rs (100%) rename server/src/{commands => llm}/chat.rs (99%) diff --git a/server/src/http/anthropic.rs b/server/src/http/anthropic.rs new file mode 100644 index 0000000..1fc90e1 --- /dev/null +++ b/server/src/http/anthropic.rs @@ -0,0 +1,18 @@ +use crate::http::context::{AppContext, OpenApiResult, bad_request}; +use crate::http::payloads::ApiKeyPayload; +use crate::llm; +use poem_openapi::payload::Json; + +pub async fn get_anthropic_api_key_exists(ctx: &AppContext) -> OpenApiResult> { + let exists = + llm::chat::get_anthropic_api_key_exists(ctx.store.as_ref()).map_err(bad_request)?; + Ok(Json(exists)) +} + +pub async fn set_anthropic_api_key( + payload: Json, + ctx: &AppContext, +) -> OpenApiResult> { + llm::chat::set_anthropic_api_key(ctx.store.as_ref(), payload.0.api_key).map_err(bad_request)?; + Ok(Json(true)) +} diff --git a/server/src/http/assets.rs b/server/src/http/assets.rs new file mode 100644 index 0000000..0a837a0 --- /dev/null +++ b/server/src/http/assets.rs @@ -0,0 +1,64 @@ +use poem::{ + Response, handler, + http::{StatusCode, header}, + web::Path, +}; +use rust_embed::RustEmbed; + +#[derive(RustEmbed)] +#[folder = "../frontend/dist"] +struct EmbeddedAssets; + +fn serve_embedded(path: &str) -> Response { + let normalized = if path.is_empty() { + "index.html" + } else { + path.trim_start_matches('/') + }; + + let is_asset_request = normalized.starts_with("assets/"); + let asset = if is_asset_request { + EmbeddedAssets::get(normalized) + } else { + EmbeddedAssets::get(normalized).or_else(|| { + if normalized == "index.html" { + None + } else { + EmbeddedAssets::get("index.html") + } + }) + }; + + match asset { + Some(content) => { + let body = content.data.into_owned(); + let mime = mime_guess::from_path(normalized) + .first_or_octet_stream() + .to_string(); + + Response::builder() + .status(StatusCode::OK) + .header(header::CONTENT_TYPE, mime) + .body(body) + } + None => Response::builder() + .status(StatusCode::NOT_FOUND) + .body("Not Found"), + } +} + +#[handler] +pub fn embedded_asset(Path(path): Path) -> Response { + let asset_path = format!("assets/{path}"); + serve_embedded(&asset_path) +} + +#[handler] +pub fn embedded_file(Path(path): Path) -> Response { + serve_embedded(&path) +} + +#[handler] +pub fn embedded_index() -> Response { + serve_embedded("index.html") +} diff --git a/server/src/http/chat.rs b/server/src/http/chat.rs new file mode 100644 index 0000000..e851127 --- /dev/null +++ b/server/src/http/chat.rs @@ -0,0 +1,8 @@ +use crate::http::context::{AppContext, OpenApiResult, bad_request}; +use crate::llm::chat; +use poem_openapi::payload::Json; + +pub async fn cancel_chat(ctx: &AppContext) -> OpenApiResult> { + chat::cancel_chat(&ctx.state).map_err(bad_request)?; + Ok(Json(true)) +} diff --git a/server/src/http/context.rs b/server/src/http/context.rs new file mode 100644 index 0000000..ca75df4 --- /dev/null +++ b/server/src/http/context.rs @@ -0,0 +1,16 @@ +use crate::state::SessionState; +use crate::store::JsonFileStore; +use poem::http::StatusCode; +use std::sync::Arc; + +#[derive(Clone)] +pub struct AppContext { + pub state: Arc, + pub store: Arc, +} + +pub type OpenApiResult = poem::Result; + +pub fn bad_request(message: String) -> poem::Error { + poem::Error::from_string(message, StatusCode::BAD_REQUEST) +} diff --git a/server/src/http/fs.rs b/server/src/http/fs.rs new file mode 100644 index 0000000..b26707b --- /dev/null +++ b/server/src/http/fs.rs @@ -0,0 +1,34 @@ +use crate::http::context::{AppContext, OpenApiResult, bad_request}; +use crate::http::payloads::{FilePathPayload, WriteFilePayload}; +use crate::io::fs; +use poem_openapi::payload::Json; + +pub async fn read_file( + payload: Json, + ctx: &AppContext, +) -> OpenApiResult> { + let content = fs::read_file(payload.0.path, &ctx.state) + .await + .map_err(bad_request)?; + Ok(Json(content)) +} + +pub async fn write_file( + payload: Json, + ctx: &AppContext, +) -> OpenApiResult> { + fs::write_file(payload.0.path, payload.0.content, &ctx.state) + .await + .map_err(bad_request)?; + Ok(Json(true)) +} + +pub async fn list_directory( + payload: Json, + ctx: &AppContext, +) -> OpenApiResult>> { + let entries = fs::list_directory(payload.0.path, &ctx.state) + .await + .map_err(bad_request)?; + Ok(Json(entries)) +} diff --git a/server/src/http/health.rs b/server/src/http/health.rs new file mode 100644 index 0000000..7eb3bd5 --- /dev/null +++ b/server/src/http/health.rs @@ -0,0 +1,6 @@ +use poem::handler; + +#[handler] +pub fn health() -> &'static str { + "ok" +} diff --git a/server/src/http/mod.rs b/server/src/http/mod.rs new file mode 100644 index 0000000..c16053e --- /dev/null +++ b/server/src/http/mod.rs @@ -0,0 +1,34 @@ +pub mod anthropic; +pub mod assets; +pub mod chat; +pub mod context; +pub mod fs; +pub mod health; +pub mod model; +pub mod payloads; +pub mod project; +pub mod rest; +pub mod search; +pub mod shell; +pub mod ws; + +use crate::http::context::AppContext; +use crate::http::rest::build_openapi_service; +use poem::EndpointExt; +use poem::{Route, get}; + +pub fn build_routes(ctx: AppContext) -> impl poem::Endpoint { + let ctx_arc = std::sync::Arc::new(ctx); + + let (api_service, docs_service) = build_openapi_service(ctx_arc.clone()); + + Route::new() + .nest("/api", api_service) + .nest("/docs", docs_service.swagger_ui()) + .at("/ws", get(ws::ws_handler)) + .at("/health", get(health::health)) + .at("/assets/*path", get(assets::embedded_asset)) + .at("/", get(assets::embedded_index)) + .at("/*path", get(assets::embedded_file)) + .data(ctx_arc) +} diff --git a/server/src/http/model.rs b/server/src/http/model.rs new file mode 100644 index 0000000..8fa7272 --- /dev/null +++ b/server/src/http/model.rs @@ -0,0 +1,27 @@ +use crate::http::context::{AppContext, OpenApiResult, bad_request}; +use crate::http::payloads::ModelPayload; +use crate::io::fs; +use crate::llm::chat; +use poem_openapi::{param::Query, payload::Json}; + +pub async fn get_model_preference(ctx: &AppContext) -> OpenApiResult>> { + let result = fs::get_model_preference(ctx.store.as_ref()).map_err(bad_request)?; + Ok(Json(result)) +} + +pub async fn set_model_preference( + payload: Json, + ctx: &AppContext, +) -> OpenApiResult> { + fs::set_model_preference(payload.0.model, ctx.store.as_ref()).map_err(bad_request)?; + Ok(Json(true)) +} + +pub async fn get_ollama_models( + base_url: Query>, +) -> OpenApiResult>> { + let models = chat::get_ollama_models(base_url.0) + .await + .map_err(bad_request)?; + Ok(Json(models)) +} diff --git a/server/src/http/payloads.rs b/server/src/http/payloads.rs new file mode 100644 index 0000000..dda0bc3 --- /dev/null +++ b/server/src/http/payloads.rs @@ -0,0 +1,39 @@ +use poem_openapi::Object; +use serde::Deserialize; + +#[derive(Deserialize, Object)] +pub struct PathPayload { + pub path: String, +} + +#[derive(Deserialize, Object)] +pub struct ModelPayload { + pub model: String, +} + +#[derive(Deserialize, Object)] +pub struct ApiKeyPayload { + pub api_key: String, +} + +#[derive(Deserialize, Object)] +pub struct FilePathPayload { + pub path: String, +} + +#[derive(Deserialize, Object)] +pub struct WriteFilePayload { + pub path: String, + pub content: String, +} + +#[derive(Deserialize, Object)] +pub struct SearchPayload { + pub query: String, +} + +#[derive(Deserialize, Object)] +pub struct ExecShellPayload { + pub command: String, + pub args: Vec, +} diff --git a/server/src/http/project.rs b/server/src/http/project.rs new file mode 100644 index 0000000..291c41c --- /dev/null +++ b/server/src/http/project.rs @@ -0,0 +1,24 @@ +use crate::http::context::{AppContext, OpenApiResult, bad_request}; +use crate::http::payloads::PathPayload; +use crate::io::fs; +use poem_openapi::payload::Json; + +pub async fn get_current_project(ctx: &AppContext) -> OpenApiResult>> { + let result = fs::get_current_project(&ctx.state, ctx.store.as_ref()).map_err(bad_request)?; + Ok(Json(result)) +} + +pub async fn open_project( + payload: Json, + ctx: &AppContext, +) -> OpenApiResult> { + let confirmed = fs::open_project(payload.0.path, &ctx.state, ctx.store.as_ref()) + .await + .map_err(bad_request)?; + Ok(Json(confirmed)) +} + +pub async fn close_project(ctx: &AppContext) -> OpenApiResult> { + fs::close_project(&ctx.state, ctx.store.as_ref()).map_err(bad_request)?; + Ok(Json(true)) +} diff --git a/server/src/http/rest.rs b/server/src/http/rest.rs new file mode 100644 index 0000000..c60b892 --- /dev/null +++ b/server/src/http/rest.rs @@ -0,0 +1,125 @@ +use crate::http::context::{AppContext, OpenApiResult}; +use crate::http::payloads::{ + ApiKeyPayload, ExecShellPayload, FilePathPayload, ModelPayload, PathPayload, SearchPayload, + WriteFilePayload, +}; +use crate::http::{anthropic, chat as chat_http, fs as fs_http, model, project, search, shell}; +use poem_openapi::{OpenApi, OpenApiService, param::Query, payload::Json}; +use std::sync::Arc; + +pub struct Api { + ctx: Arc, +} + +#[OpenApi] +impl Api { + #[oai(path = "/project", method = "get")] + async fn get_current_project(&self) -> OpenApiResult>> { + let ctx = self.ctx.clone(); + project::get_current_project(&ctx).await + } + + #[oai(path = "/project", method = "post")] + async fn open_project(&self, payload: Json) -> OpenApiResult> { + let ctx = self.ctx.clone(); + project::open_project(payload, &ctx).await + } + + #[oai(path = "/project", method = "delete")] + async fn close_project(&self) -> OpenApiResult> { + let ctx = self.ctx.clone(); + project::close_project(&ctx).await + } + + #[oai(path = "/model", method = "get")] + async fn get_model_preference(&self) -> OpenApiResult>> { + let ctx = self.ctx.clone(); + model::get_model_preference(&ctx).await + } + + #[oai(path = "/model", method = "post")] + async fn set_model_preference(&self, payload: Json) -> OpenApiResult> { + let ctx = self.ctx.clone(); + model::set_model_preference(payload, &ctx).await + } + + #[oai(path = "/ollama/models", method = "get")] + async fn get_ollama_models( + &self, + base_url: Query>, + ) -> OpenApiResult>> { + model::get_ollama_models(base_url).await + } + + #[oai(path = "/anthropic/key/exists", method = "get")] + async fn get_anthropic_api_key_exists(&self) -> OpenApiResult> { + let ctx = self.ctx.clone(); + anthropic::get_anthropic_api_key_exists(&ctx).await + } + + #[oai(path = "/anthropic/key", method = "post")] + async fn set_anthropic_api_key( + &self, + payload: Json, + ) -> OpenApiResult> { + let ctx = self.ctx.clone(); + anthropic::set_anthropic_api_key(payload, &ctx).await + } + + #[oai(path = "/fs/read", method = "post")] + async fn read_file(&self, payload: Json) -> OpenApiResult> { + let ctx = self.ctx.clone(); + fs_http::read_file(payload, &ctx).await + } + + #[oai(path = "/fs/write", method = "post")] + async fn write_file(&self, payload: Json) -> OpenApiResult> { + let ctx = self.ctx.clone(); + fs_http::write_file(payload, &ctx).await + } + + #[oai(path = "/fs/list", method = "post")] + async fn list_directory( + &self, + payload: Json, + ) -> OpenApiResult>> { + let ctx = self.ctx.clone(); + fs_http::list_directory(payload, &ctx).await + } + + #[oai(path = "/fs/search", method = "post")] + async fn search_files( + &self, + payload: Json, + ) -> OpenApiResult>> { + let ctx = self.ctx.clone(); + search::search_files(payload, &ctx).await + } + + #[oai(path = "/shell/exec", method = "post")] + async fn exec_shell( + &self, + payload: Json, + ) -> OpenApiResult> { + let ctx = self.ctx.clone(); + shell::exec_shell(payload, &ctx).await + } + + #[oai(path = "/chat/cancel", method = "post")] + async fn cancel_chat(&self) -> OpenApiResult> { + let ctx = self.ctx.clone(); + chat_http::cancel_chat(&ctx).await + } +} + +pub fn build_openapi_service( + ctx: Arc, +) -> (OpenApiService, OpenApiService) { + let api_service = OpenApiService::new(Api { ctx: ctx.clone() }, "Story Kit API", "1.0") + .server("http://127.0.0.1:3001/api"); + + let docs_service = OpenApiService::new(Api { ctx }, "Story Kit API", "1.0") + .server("http://127.0.0.1:3001/api"); + + (api_service, docs_service) +} diff --git a/server/src/http/search.rs b/server/src/http/search.rs new file mode 100644 index 0000000..a529e7b --- /dev/null +++ b/server/src/http/search.rs @@ -0,0 +1,13 @@ +use crate::http::context::{AppContext, OpenApiResult, bad_request}; +use crate::http::payloads::SearchPayload; +use poem_openapi::payload::Json; + +pub async fn search_files( + payload: Json, + ctx: &AppContext, +) -> OpenApiResult>> { + let results = crate::io::search::search_files(payload.0.query, &ctx.state) + .await + .map_err(bad_request)?; + Ok(Json(results)) +} diff --git a/server/src/http/shell.rs b/server/src/http/shell.rs new file mode 100644 index 0000000..6ccfb14 --- /dev/null +++ b/server/src/http/shell.rs @@ -0,0 +1,13 @@ +use crate::http::context::{AppContext, OpenApiResult, bad_request}; +use crate::http::payloads::ExecShellPayload; +use poem_openapi::payload::Json; + +pub async fn exec_shell( + payload: Json, + ctx: &AppContext, +) -> OpenApiResult> { + let output = crate::io::shell::exec_shell(payload.0.command, payload.0.args, &ctx.state) + .await + .map_err(bad_request)?; + Ok(Json(output)) +} diff --git a/server/src/http/ws.rs b/server/src/http/ws.rs new file mode 100644 index 0000000..87fb560 --- /dev/null +++ b/server/src/http/ws.rs @@ -0,0 +1,92 @@ +use crate::http::context::AppContext; +use crate::llm::chat; +use crate::llm::types::Message; +use futures::{SinkExt, StreamExt}; +use poem::handler; +use poem::web::Data; +use poem::web::websocket::{Message as WsMessage, WebSocket}; +use serde::{Deserialize, Serialize}; +use tokio::sync::mpsc; + +#[derive(Deserialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum WsRequest { + Chat { + messages: Vec, + config: chat::ProviderConfig, + }, + Cancel, +} + +#[derive(Serialize)] +#[serde(tag = "type", rename_all = "snake_case")] +enum WsResponse { + Token { content: String }, + Update { messages: Vec }, + Error { message: String }, +} + +#[handler] +pub async fn ws_handler(ws: WebSocket, ctx: Data<&AppContext>) -> impl poem::IntoResponse { + let ctx = ctx.0.clone(); + ws.on_upgrade(move |socket| async move { + let (mut sink, mut stream) = socket.split(); + let (tx, mut rx) = mpsc::unbounded_channel::(); + + let forward = tokio::spawn(async move { + while let Some(msg) = rx.recv().await { + if let Ok(text) = serde_json::to_string(&msg) + && sink.send(WsMessage::Text(text)).await.is_err() + { + break; + } + } + }); + + while let Some(Ok(msg)) = stream.next().await { + if let WsMessage::Text(text) = msg { + let parsed: Result = serde_json::from_str(&text); + match parsed { + Ok(WsRequest::Chat { messages, config }) => { + let tx_updates = tx.clone(); + let tx_tokens = tx.clone(); + let ctx_clone = ctx.clone(); + + let result = chat::chat( + messages, + config, + &ctx_clone.state, + ctx_clone.store.as_ref(), + |history| { + let _ = tx_updates.send(WsResponse::Update { + messages: history.to_vec(), + }); + }, + |token| { + let _ = tx_tokens.send(WsResponse::Token { + content: token.to_string(), + }); + }, + ) + .await; + + if let Err(err) = result { + let _ = tx.send(WsResponse::Error { message: err }); + } + } + Ok(WsRequest::Cancel) => { + let _ = chat::cancel_chat(&ctx.state); + } + Err(err) => { + let _ = tx.send(WsResponse::Error { + message: format!("Invalid request: {err}"), + }); + } + } + } + } + + drop(tx); + let _ = forward.await; + }) +} diff --git a/server/src/commands/fs.rs b/server/src/io/fs.rs similarity index 100% rename from server/src/commands/fs.rs rename to server/src/io/fs.rs diff --git a/server/src/commands/mod.rs b/server/src/io/mod.rs similarity index 75% rename from server/src/commands/mod.rs rename to server/src/io/mod.rs index 943956d..71e6bf5 100644 --- a/server/src/commands/mod.rs +++ b/server/src/io/mod.rs @@ -1,4 +1,3 @@ -pub mod chat; pub mod fs; pub mod search; pub mod shell; diff --git a/server/src/commands/search.rs b/server/src/io/search.rs similarity index 100% rename from server/src/commands/search.rs rename to server/src/io/search.rs diff --git a/server/src/commands/shell.rs b/server/src/io/shell.rs similarity index 100% rename from server/src/commands/shell.rs rename to server/src/io/shell.rs diff --git a/server/src/commands/chat.rs b/server/src/llm/chat.rs similarity index 99% rename from server/src/commands/chat.rs rename to server/src/llm/chat.rs index 228377c..6bc3799 100644 --- a/server/src/commands/chat.rs +++ b/server/src/llm/chat.rs @@ -314,7 +314,7 @@ where } async fn execute_tool(call: &ToolCall, state: &SessionState) -> String { - use crate::commands::{fs, search, shell}; + use crate::io::{fs, search, shell}; let name = call.function.name.as_str(); let args: serde_json::Value = match parse_tool_arguments(&call.function.arguments) { diff --git a/server/src/llm/mod.rs b/server/src/llm/mod.rs index fc67ad4..431529b 100644 --- a/server/src/llm/mod.rs +++ b/server/src/llm/mod.rs @@ -1,3 +1,4 @@ +pub mod chat; pub mod prompts; pub mod providers; pub mod types; diff --git a/server/src/main.rs b/server/src/main.rs index 580aca8..038cbb7 100644 --- a/server/src/main.rs +++ b/server/src/main.rs @@ -1,354 +1,17 @@ -mod commands; +mod http; +mod io; mod llm; mod state; mod store; -use crate::commands::{chat, fs}; -use crate::llm::types::Message; +use crate::http::build_routes; +use crate::http::context::AppContext; use crate::state::SessionState; use crate::store::JsonFileStore; -use futures::{SinkExt, StreamExt}; -use poem::web::websocket::{Message as WsMessage, WebSocket}; -use poem::{ - EndpointExt, Response, Route, Server, get, handler, - http::{StatusCode, header}, - listener::TcpListener, - web::{Data, Path}, -}; -use poem_openapi::{Object, OpenApi, OpenApiService, param::Query, payload::Json}; -use rust_embed::RustEmbed; -use serde::{Deserialize, Serialize}; +use poem::Server; +use poem::listener::TcpListener; use std::path::PathBuf; use std::sync::Arc; -use tokio::sync::mpsc; - -#[derive(Clone)] -struct AppContext { - state: Arc, - store: Arc, -} - -#[derive(RustEmbed)] -#[folder = "../frontend/dist"] -struct EmbeddedAssets; - -type OpenApiResult = poem::Result; - -fn bad_request(message: String) -> poem::Error { - poem::Error::from_string(message, StatusCode::BAD_REQUEST) -} - -#[handler] -fn health() -> &'static str { - "ok" -} - -fn serve_embedded(path: &str) -> Response { - let normalized = if path.is_empty() { - "index.html" - } else { - path.trim_start_matches('/') - }; - let is_asset_request = normalized.starts_with("assets/"); - let asset = if is_asset_request { - EmbeddedAssets::get(normalized) - } else { - EmbeddedAssets::get(normalized).or_else(|| { - if normalized == "index.html" { - None - } else { - EmbeddedAssets::get("index.html") - } - }) - }; - - match asset { - Some(content) => { - let body = content.data.into_owned(); - let mime = mime_guess::from_path(normalized) - .first_or_octet_stream() - .to_string(); - - Response::builder() - .status(StatusCode::OK) - .header(header::CONTENT_TYPE, mime) - .body(body) - } - None => Response::builder() - .status(StatusCode::NOT_FOUND) - .body("Not Found"), - } -} - -#[handler] -fn embedded_asset(Path(path): Path) -> Response { - let asset_path = format!("assets/{path}"); - serve_embedded(&asset_path) -} - -#[handler] -fn embedded_file(Path(path): Path) -> Response { - serve_embedded(&path) -} - -#[handler] -fn embedded_index() -> Response { - serve_embedded("index.html") -} - -#[derive(Deserialize, Object)] -struct PathPayload { - path: String, -} - -#[derive(Deserialize, Object)] -struct ModelPayload { - model: String, -} - -#[derive(Deserialize, Object)] -struct ApiKeyPayload { - api_key: String, -} - -#[derive(Deserialize, Object)] -struct FilePathPayload { - path: String, -} - -#[derive(Deserialize, Object)] -struct WriteFilePayload { - path: String, - content: String, -} - -#[derive(Deserialize, Object)] -struct SearchPayload { - query: String, -} - -#[derive(Deserialize, Object)] -struct ExecShellPayload { - command: String, - args: Vec, -} -struct Api { - ctx: Arc, -} - -#[OpenApi] -impl Api { - #[oai(path = "/project", method = "get")] - async fn get_current_project(&self) -> OpenApiResult>> { - let ctx = self.ctx.clone(); - let result = - fs::get_current_project(&ctx.state, ctx.store.as_ref()).map_err(bad_request)?; - Ok(Json(result)) - } - - #[oai(path = "/project", method = "post")] - async fn open_project(&self, payload: Json) -> OpenApiResult> { - let ctx = self.ctx.clone(); - let confirmed = fs::open_project(payload.0.path, &ctx.state, ctx.store.as_ref()) - .await - .map_err(bad_request)?; - Ok(Json(confirmed)) - } - - #[oai(path = "/project", method = "delete")] - async fn close_project(&self) -> OpenApiResult> { - let ctx = self.ctx.clone(); - fs::close_project(&ctx.state, ctx.store.as_ref()).map_err(bad_request)?; - Ok(Json(true)) - } - - #[oai(path = "/model", method = "get")] - async fn get_model_preference(&self) -> OpenApiResult>> { - let ctx = self.ctx.clone(); - let result = fs::get_model_preference(ctx.store.as_ref()).map_err(bad_request)?; - Ok(Json(result)) - } - - #[oai(path = "/model", method = "post")] - async fn set_model_preference(&self, payload: Json) -> OpenApiResult> { - let ctx = self.ctx.clone(); - fs::set_model_preference(payload.0.model, ctx.store.as_ref()).map_err(bad_request)?; - Ok(Json(true)) - } - - #[oai(path = "/ollama/models", method = "get")] - async fn get_ollama_models( - &self, - base_url: Query>, - ) -> OpenApiResult>> { - let models = chat::get_ollama_models(base_url.0) - .await - .map_err(bad_request)?; - Ok(Json(models)) - } - - #[oai(path = "/anthropic/key/exists", method = "get")] - async fn get_anthropic_api_key_exists(&self) -> OpenApiResult> { - let ctx = self.ctx.clone(); - let exists = chat::get_anthropic_api_key_exists(ctx.store.as_ref()).map_err(bad_request)?; - Ok(Json(exists)) - } - - #[oai(path = "/anthropic/key", method = "post")] - async fn set_anthropic_api_key( - &self, - payload: Json, - ) -> OpenApiResult> { - let ctx = self.ctx.clone(); - chat::set_anthropic_api_key(ctx.store.as_ref(), payload.0.api_key).map_err(bad_request)?; - Ok(Json(true)) - } - - #[oai(path = "/fs/read", method = "post")] - async fn read_file(&self, payload: Json) -> OpenApiResult> { - let ctx = self.ctx.clone(); - let content = fs::read_file(payload.0.path, &ctx.state) - .await - .map_err(bad_request)?; - Ok(Json(content)) - } - - #[oai(path = "/fs/write", method = "post")] - async fn write_file(&self, payload: Json) -> OpenApiResult> { - let ctx = self.ctx.clone(); - fs::write_file(payload.0.path, payload.0.content, &ctx.state) - .await - .map_err(bad_request)?; - Ok(Json(true)) - } - - #[oai(path = "/fs/list", method = "post")] - async fn list_directory( - &self, - payload: Json, - ) -> OpenApiResult>> { - let ctx = self.ctx.clone(); - let entries = fs::list_directory(payload.0.path, &ctx.state) - .await - .map_err(bad_request)?; - Ok(Json(entries)) - } - - #[oai(path = "/fs/search", method = "post")] - async fn search_files( - &self, - payload: Json, - ) -> OpenApiResult>> { - let ctx = self.ctx.clone(); - let results = crate::commands::search::search_files(payload.0.query, &ctx.state) - .await - .map_err(bad_request)?; - Ok(Json(results)) - } - - #[oai(path = "/shell/exec", method = "post")] - async fn exec_shell( - &self, - payload: Json, - ) -> OpenApiResult> { - let ctx = self.ctx.clone(); - let output = - crate::commands::shell::exec_shell(payload.0.command, payload.0.args, &ctx.state) - .await - .map_err(bad_request)?; - Ok(Json(output)) - } - - #[oai(path = "/chat/cancel", method = "post")] - async fn cancel_chat(&self) -> OpenApiResult> { - let ctx = self.ctx.clone(); - chat::cancel_chat(&ctx.state).map_err(bad_request)?; - Ok(Json(true)) - } -} - -#[derive(Deserialize)] -#[serde(tag = "type", rename_all = "snake_case")] -enum WsRequest { - Chat { - messages: Vec, - config: chat::ProviderConfig, - }, - Cancel, -} - -#[derive(Serialize)] -#[serde(tag = "type", rename_all = "snake_case")] -enum WsResponse { - Token { content: String }, - Update { messages: Vec }, - Error { message: String }, -} - -#[handler] -async fn ws_handler(ws: WebSocket, ctx: Data<&AppContext>) -> impl poem::IntoResponse { - let ctx = ctx.0.clone(); - ws.on_upgrade(move |socket| async move { - let (mut sink, mut stream) = socket.split(); - let (tx, mut rx) = mpsc::unbounded_channel::(); - - let forward = tokio::spawn(async move { - while let Some(msg) = rx.recv().await { - if let Ok(text) = serde_json::to_string(&msg) - && sink.send(WsMessage::Text(text)).await.is_err() - { - break; - } - } - }); - - while let Some(Ok(msg)) = stream.next().await { - if let WsMessage::Text(text) = msg { - let parsed: Result = serde_json::from_str(&text); - match parsed { - Ok(WsRequest::Chat { messages, config }) => { - let tx_updates = tx.clone(); - let tx_tokens = tx.clone(); - let ctx_clone = ctx.clone(); - - let result = chat::chat( - messages, - config, - &ctx_clone.state, - ctx_clone.store.as_ref(), - |history| { - let _ = tx_updates.send(WsResponse::Update { - messages: history.to_vec(), - }); - }, - |token| { - let _ = tx_tokens.send(WsResponse::Token { - content: token.to_string(), - }); - }, - ) - .await; - - if let Err(err) = result { - let _ = tx.send(WsResponse::Error { message: err }); - } - } - Ok(WsRequest::Cancel) => { - let _ = chat::cancel_chat(&ctx.state); - } - Err(err) => { - let _ = tx.send(WsResponse::Error { - message: format!("Invalid request: {err}"), - }); - } - } - } - } - - drop(tx); - let _ = forward.await; - }) -} #[tokio::main] async fn main() -> Result<(), std::io::Error> { @@ -361,34 +24,8 @@ async fn main() -> Result<(), std::io::Error> { state: app_state, store, }; - let ctx_arc = Arc::new(ctx.clone()); - let api_service = OpenApiService::new( - Api { - ctx: ctx_arc.clone(), - }, - "Living Spec API", - "1.0", - ) - .server("http://127.0.0.1:3001/api"); - let docs_service = OpenApiService::new( - Api { - ctx: ctx_arc.clone(), - }, - "Living Spec API", - "1.0", - ) - .server("http://127.0.0.1:3001/api"); - - let app = Route::new() - .nest("/api", api_service) - .nest("/docs", docs_service.swagger_ui()) - .at("/ws", get(ws_handler)) - .at("/health", get(health)) - .at("/assets/*path", get(embedded_asset)) - .at("/", get(embedded_index)) - .at("/*path", get(embedded_file)) - .data(ctx); + let app = build_routes(ctx); Server::new(TcpListener::bind("127.0.0.1:3001")) .run(app)