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

server/src/io/mod.rs

@@ -1,4 +1,3 @@
-pub mod chat;
 pub mod fs;
 pub mod search;
 pub mod shell;

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) {

server/src/llm/mod.rs

@@ -1,3 +1,4 @@
+pub mod chat;
 pub mod prompts;
 pub mod providers;
 pub mod types;

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<SessionState>,
-    store: Arc<JsonFileStore>,
-}
-
-#[derive(RustEmbed)]
-#[folder = "../frontend/dist"]
-struct EmbeddedAssets;
-
-type OpenApiResult<T> = poem::Result<T>;
-
-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<String>) -> Response {
-    let asset_path = format!("assets/{path}");
-    serve_embedded(&asset_path)
-}
-
-#[handler]
-fn embedded_file(Path(path): Path<String>) -> 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<String>,
-}
-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();
-        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<PathPayload>) -> OpenApiResult<Json<String>> {
-        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<Json<bool>> {
-        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<Json<Option<String>>> {
-        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<ModelPayload>) -> OpenApiResult<Json<bool>> {
-        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<Option<String>>,
-    ) -> OpenApiResult<Json<Vec<String>>> {
-        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<Json<bool>> {
-        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<ApiKeyPayload>,
-    ) -> OpenApiResult<Json<bool>> {
-        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<FilePathPayload>) -> OpenApiResult<Json<String>> {
-        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<WriteFilePayload>) -> OpenApiResult<Json<bool>> {
-        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<FilePathPayload>,
-    ) -> OpenApiResult<Json<Vec<fs::FileEntry>>> {
-        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<SearchPayload>,
-    ) -> OpenApiResult<Json<Vec<crate::commands::search::SearchResult>>> {
-        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<ExecShellPayload>,
-    ) -> OpenApiResult<Json<crate::commands::shell::CommandOutput>> {
-        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<Json<bool>> {
-        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<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]
-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;
-    })
-}
 
 #[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)