Refactoring the structure a bit

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<Json<bool>> {
+    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<ApiKeyPayload>,
+    ctx: &AppContext,
+) -> OpenApiResult<Json<bool>> {
+    llm::chat::set_anthropic_api_key(ctx.store.as_ref(), payload.0.api_key).map_err(bad_request)?;
+    Ok(Json(true))
+}

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<String>) -> Response {
+    let asset_path = format!("assets/{path}");
+    serve_embedded(&asset_path)
+}
+
+#[handler]
+pub fn embedded_file(Path(path): Path<String>) -> Response {
+    serve_embedded(&path)
+}
+
+#[handler]
+pub fn embedded_index() -> Response {
+    serve_embedded("index.html")
+}

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<Json<bool>> {
+    chat::cancel_chat(&ctx.state).map_err(bad_request)?;
+    Ok(Json(true))
+}

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<SessionState>,
+    pub store: Arc<JsonFileStore>,
+}
+
+pub type OpenApiResult<T> = poem::Result<T>;
+
+pub fn bad_request(message: String) -> poem::Error {
+    poem::Error::from_string(message, StatusCode::BAD_REQUEST)
+}

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<FilePathPayload>,
+    ctx: &AppContext,
+) -> OpenApiResult<Json<String>> {
+    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<WriteFilePayload>,
+    ctx: &AppContext,
+) -> OpenApiResult<Json<bool>> {
+    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<FilePathPayload>,
+    ctx: &AppContext,
+) -> OpenApiResult<Json<Vec<fs::FileEntry>>> {
+    let entries = fs::list_directory(payload.0.path, &ctx.state)
+        .await
+        .map_err(bad_request)?;
+    Ok(Json(entries))
+}

server/src/http/health.rs

@@ -0,0 +1,6 @@
+use poem::handler;
+
+#[handler]
+pub fn health() -> &'static str {
+    "ok"
+}

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)
+}

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<Json<Option<String>>> {
+    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<ModelPayload>,
+    ctx: &AppContext,
+) -> OpenApiResult<Json<bool>> {
+    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<Option<String>>,
+) -> OpenApiResult<Json<Vec<String>>> {
+    let models = chat::get_ollama_models(base_url.0)
+        .await
+        .map_err(bad_request)?;
+    Ok(Json(models))
+}

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<String>,
+}

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<Json<Option<String>>> {
+    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<PathPayload>,
+    ctx: &AppContext,
+) -> OpenApiResult<Json<String>> {
+    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<Json<bool>> {
+    fs::close_project(&ctx.state, ctx.store.as_ref()).map_err(bad_request)?;
+    Ok(Json(true))
+}

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<AppContext>,
+}
+
+#[OpenApi]
+impl Api {
+    #[oai(path = "/project", method = "get")]
+    async fn get_current_project(&self) -> OpenApiResult<Json<Option<String>>> {
+        let ctx = self.ctx.clone();
+        project::get_current_project(&ctx).await
+    }
+
+    #[oai(path = "/project", method = "post")]
+    async fn open_project(&self, payload: Json<PathPayload>) -> OpenApiResult<Json<String>> {
+        let ctx = self.ctx.clone();
+        project::open_project(payload, &ctx).await
+    }
+
+    #[oai(path = "/project", method = "delete")]
+    async fn close_project(&self) -> OpenApiResult<Json<bool>> {
+        let ctx = self.ctx.clone();
+        project::close_project(&ctx).await
+    }
+
+    #[oai(path = "/model", method = "get")]
+    async fn get_model_preference(&self) -> OpenApiResult<Json<Option<String>>> {
+        let ctx = self.ctx.clone();
+        model::get_model_preference(&ctx).await
+    }
+
+    #[oai(path = "/model", method = "post")]
+    async fn set_model_preference(&self, payload: Json<ModelPayload>) -> OpenApiResult<Json<bool>> {
+        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<Option<String>>,
+    ) -> OpenApiResult<Json<Vec<String>>> {
+        model::get_ollama_models(base_url).await
+    }
+
+    #[oai(path = "/anthropic/key/exists", method = "get")]
+    async fn get_anthropic_api_key_exists(&self) -> OpenApiResult<Json<bool>> {
+        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<ApiKeyPayload>,
+    ) -> OpenApiResult<Json<bool>> {
+        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<FilePathPayload>) -> OpenApiResult<Json<String>> {
+        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<WriteFilePayload>) -> OpenApiResult<Json<bool>> {
+        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<FilePathPayload>,
+    ) -> OpenApiResult<Json<Vec<crate::io::fs::FileEntry>>> {
+        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<SearchPayload>,
+    ) -> OpenApiResult<Json<Vec<crate::io::search::SearchResult>>> {
+        let ctx = self.ctx.clone();
+        search::search_files(payload, &ctx).await
+    }
+
+    #[oai(path = "/shell/exec", method = "post")]
+    async fn exec_shell(
+        &self,
+        payload: Json<ExecShellPayload>,
+    ) -> OpenApiResult<Json<crate::io::shell::CommandOutput>> {
+        let ctx = self.ctx.clone();
+        shell::exec_shell(payload, &ctx).await
+    }
+
+    #[oai(path = "/chat/cancel", method = "post")]
+    async fn cancel_chat(&self) -> OpenApiResult<Json<bool>> {
+        let ctx = self.ctx.clone();
+        chat_http::cancel_chat(&ctx).await
+    }
+}
+
+pub fn build_openapi_service(
+    ctx: Arc<AppContext>,
+) -> (OpenApiService<Api, ()>, OpenApiService<Api, ()>) {
+    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)
+}

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<SearchPayload>,
+    ctx: &AppContext,
+) -> OpenApiResult<Json<Vec<crate::io::search::SearchResult>>> {
+    let results = crate::io::search::search_files(payload.0.query, &ctx.state)
+        .await
+        .map_err(bad_request)?;
+    Ok(Json(results))
+}

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<ExecShellPayload>,
+    ctx: &AppContext,
+) -> OpenApiResult<Json<crate::io::shell::CommandOutput>> {
+    let output = crate::io::shell::exec_shell(payload.0.command, payload.0.args, &ctx.state)
+        .await
+        .map_err(bad_request)?;
+    Ok(Json(output))
+}

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<Message>,
+        config: chat::ProviderConfig,
+    },
+    Cancel,
+}
+
+#[derive(Serialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+enum WsResponse {
+    Token { content: String },
+    Update { messages: Vec<Message> },
+    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::<WsResponse>();
+
+        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<WsRequest, _> = 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;
+    })
+}