Skip to main content

sim_lib_openai_server/
server.rs

1use std::sync::{Arc, Mutex};
2
3use sim_citizen_derive::non_citizen;
4use sim_kernel::{Cx, Expr, Object, ObjectCompat, Ref, Result, Symbol};
5
6pub use crate::objects::{GatewayRequest, GatewayResponse, GatewayResponseValue};
7use crate::routes::{
8    admin::{
9        ADMIN_CACHE_STATS_PATH, ADMIN_CAPABILITY_REPORT_PATH, ADMIN_EVENTS_PATH,
10        ADMIN_MODEL_HEALTH_PATH, ADMIN_RUN_RETRIEVAL_PREFIX, ADMIN_RUN_RETRIEVAL_ROUTE,
11        ADMIN_RUNS_PATH, ADMIN_STORAGE_STATS_PATH, handle_admin_cache_stats,
12        handle_admin_capability_report, handle_admin_events, handle_admin_model_health,
13        handle_admin_run_get, handle_admin_runs, handle_admin_storage_stats,
14    },
15    audio::{
16        AUDIO_SPEECH_PATH, AUDIO_TRANSCRIPTIONS_PATH, handle_audio_speech,
17        handle_audio_transcriptions,
18    },
19    batches::{
20        BATCH_CANCEL_ROUTE, BATCH_RETRIEVAL_PREFIX, BATCH_RETRIEVAL_ROUTE, BATCHES_PATH,
21        handle_batch_cancel, handle_batch_retrieval, handle_batches,
22    },
23    chat_completions::{CHAT_COMPLETIONS_PATH, handle_chat_completions},
24    embeddings::{EMBEDDINGS_PATH, handle_embeddings},
25    files::{
26        FILE_RETRIEVAL_PREFIX, FILE_RETRIEVAL_ROUTE, FILES_PATH, handle_file_retrieval,
27        handle_files,
28    },
29    health::{HEALTH_PATH, handle_health},
30    images::{IMAGES_GENERATIONS_PATH, handle_image_generations},
31    models::{MODELS_PATH, handle_models},
32    replay::{
33        RESPONSE_EVENTS_ROUTE, RESPONSE_EVENTS_SUFFIX, RESPONSE_SIM_ROUTE, RESPONSE_SIM_SUFFIX,
34        SIM_FORK_PATH, SIM_REPLAY_PATH, handle_response_events, handle_response_sim,
35        handle_sim_fork, handle_sim_replay,
36    },
37    responses::{
38        RESPONSE_RETRIEVAL_PREFIX, RESPONSE_RETRIEVAL_ROUTE, RESPONSES_PATH,
39        handle_response_retrieval, handle_responses,
40    },
41    threads::{
42        THREAD_MESSAGES_ROUTE, THREAD_RETRIEVAL_PREFIX, THREAD_RETRIEVAL_ROUTE, THREADS_PATH,
43        handle_thread_get, handle_thread_post, handle_threads,
44    },
45    vector_stores::{
46        VECTOR_STORE_SEARCH_PREFIX, VECTOR_STORE_SEARCH_ROUTE, VECTOR_STORE_SEARCH_SUFFIX,
47        VECTOR_STORES_PATH, handle_vector_store_search, handle_vector_stores,
48    },
49};
50use crate::{
51    runtime::{
52        OpenAiFederation, OpenAiKeyTable, OpenAiPlanCache, OpenAiRunnerRegistry,
53        global_openai_key_table,
54    },
55    storage::MemoryGatewayStore,
56};
57
58type RouteHandler = fn(&GatewayRequest, &GatewayRouteState) -> GatewayResponse;
59
60/// A single registered route: an HTTP method, a path matcher, and a handler.
61#[derive(Clone, Copy)]
62pub struct GatewayRoute {
63    method: &'static str,
64    path: RoutePath,
65    handler: RouteHandler,
66}
67
68impl GatewayRoute {
69    /// Returns a route matching `path` exactly for the given method and handler.
70    pub fn new(method: &'static str, path: &'static str, handler: RouteHandler) -> Self {
71        Self {
72            method,
73            path: RoutePath::Exact(path),
74            handler,
75        }
76    }
77
78    /// Returns a route matching any path that starts with `prefix`.
79    ///
80    /// `display_path` is the canonical path reported by [`GatewayRoute::path`].
81    pub fn prefix(
82        method: &'static str,
83        display_path: &'static str,
84        prefix: &'static str,
85        handler: RouteHandler,
86    ) -> Self {
87        Self {
88            method,
89            path: RoutePath::Prefix {
90                display_path,
91                prefix,
92            },
93            handler,
94        }
95    }
96
97    /// Returns a route matching paths that start with `prefix` and end with `suffix`.
98    ///
99    /// `display_path` is the canonical path reported by [`GatewayRoute::path`].
100    pub fn prefix_suffix(
101        method: &'static str,
102        display_path: &'static str,
103        prefix: &'static str,
104        suffix: &'static str,
105        handler: RouteHandler,
106    ) -> Self {
107        Self {
108            method,
109            path: RoutePath::PrefixSuffix {
110                display_path,
111                prefix,
112                suffix,
113            },
114            handler,
115        }
116    }
117
118    /// Returns the canonical display path for this route.
119    pub fn path(&self) -> &'static str {
120        self.path.display()
121    }
122
123    fn matches_path(&self, path: &str) -> bool {
124        self.path.matches(path)
125    }
126}
127
128#[derive(Clone, Copy)]
129enum RoutePath {
130    Exact(&'static str),
131    Prefix {
132        display_path: &'static str,
133        prefix: &'static str,
134    },
135    PrefixSuffix {
136        display_path: &'static str,
137        prefix: &'static str,
138        suffix: &'static str,
139    },
140}
141
142impl RoutePath {
143    fn display(self) -> &'static str {
144        match self {
145            Self::Exact(path) => path,
146            Self::Prefix { display_path, .. } => display_path,
147            Self::PrefixSuffix { display_path, .. } => display_path,
148        }
149    }
150
151    fn matches(self, path: &str) -> bool {
152        match self {
153            Self::Exact(route_path) => route_path == path,
154            Self::Prefix { prefix, .. } => path.starts_with(prefix),
155            Self::PrefixSuffix { prefix, suffix, .. } => {
156                path.starts_with(prefix) && path.ends_with(suffix)
157            }
158        }
159    }
160}
161
162/// Shared, cloneable state threaded into every route handler.
163///
164/// Bundles the run store, plan cache, model runner registry, federation
165/// registry, and API key table; clones share the same underlying store and
166/// cache through interior mutability.
167#[derive(Clone)]
168pub struct GatewayRouteState {
169    store: Arc<Mutex<MemoryGatewayStore>>,
170    cache: Arc<Mutex<OpenAiPlanCache>>,
171    runners: OpenAiRunnerRegistry,
172    federation: OpenAiFederation,
173    keys: OpenAiKeyTable,
174}
175
176impl GatewayRouteState {
177    /// Returns route state wrapping `store` with a fresh cache, runners, and federation.
178    pub fn new(store: MemoryGatewayStore) -> Self {
179        Self {
180            store: Arc::new(Mutex::new(store)),
181            cache: Arc::new(Mutex::new(OpenAiPlanCache::new())),
182            runners: OpenAiRunnerRegistry::new(),
183            federation: OpenAiFederation::new(),
184            keys: global_openai_key_table().clone(),
185        }
186    }
187
188    /// Returns route state wrapping `store` and a preseeded plan `cache`.
189    pub fn with_cache(store: MemoryGatewayStore, cache: OpenAiPlanCache) -> Self {
190        Self {
191            store: Arc::new(Mutex::new(store)),
192            cache: Arc::new(Mutex::new(cache)),
193            runners: OpenAiRunnerRegistry::new(),
194            federation: OpenAiFederation::new(),
195            keys: global_openai_key_table().clone(),
196        }
197    }
198
199    /// Returns the state with its model runner registry replaced by `runners`.
200    pub fn with_runners(mut self, runners: OpenAiRunnerRegistry) -> Self {
201        self.runners = runners;
202        self
203    }
204
205    /// Returns the state with its federation registry replaced by `federation`.
206    pub fn with_federation(mut self, federation: OpenAiFederation) -> Self {
207        self.federation = federation;
208        self
209    }
210
211    /// Returns the state with its API key table replaced by `keys`.
212    pub fn with_keys(mut self, keys: OpenAiKeyTable) -> Self {
213        self.keys = keys;
214        self
215    }
216
217    /// Returns route state backed by a fresh in-memory store.
218    pub fn memory() -> Self {
219        Self::new(MemoryGatewayStore::new())
220    }
221
222    /// Returns the shared run store.
223    pub fn store(&self) -> &Arc<Mutex<MemoryGatewayStore>> {
224        &self.store
225    }
226
227    /// Returns the shared plan cache.
228    pub fn cache(&self) -> &Arc<Mutex<OpenAiPlanCache>> {
229        &self.cache
230    }
231
232    /// Returns the model runner registry.
233    pub fn runners(&self) -> &OpenAiRunnerRegistry {
234        &self.runners
235    }
236
237    /// Returns the federation registry.
238    pub fn federation(&self) -> &OpenAiFederation {
239        &self.federation
240    }
241
242    /// Returns the API key table.
243    pub fn keys(&self) -> &OpenAiKeyTable {
244        &self.keys
245    }
246}
247
248/// A route table paired with the state handlers dispatch against.
249#[derive(Clone)]
250pub struct GatewayRoutes {
251    routes: Vec<GatewayRoute>,
252    state: GatewayRouteState,
253}
254
255impl GatewayRoutes {
256    /// Returns a route table over `routes` with fresh in-memory state.
257    pub fn new(routes: Vec<GatewayRoute>) -> Self {
258        Self::with_state(routes, GatewayRouteState::memory())
259    }
260
261    /// Returns a route table over `routes` sharing the supplied `state`.
262    pub fn with_state(routes: Vec<GatewayRoute>, state: GatewayRouteState) -> Self {
263        Self { routes, state }
264    }
265
266    /// Dispatches `request` to the matching route, returning its response.
267    ///
268    /// Returns 405 when the path exists under a different method and 404 when no
269    /// route matches the path at all.
270    pub fn handle(&self, request: &GatewayRequest) -> GatewayResponse {
271        let mut path_exists = false;
272        for route in &self.routes {
273            if !route.matches_path(request.path()) {
274                continue;
275            }
276            path_exists = true;
277            if route.method == request.method() {
278                return (route.handler)(request, &self.state);
279            }
280        }
281        if path_exists {
282            GatewayResponse::text(405, "method not allowed")
283        } else {
284            GatewayResponse::text(404, "not found")
285        }
286    }
287
288    /// Returns the display paths of all registered routes.
289    pub fn paths(&self) -> Vec<&'static str> {
290        self.routes.iter().map(GatewayRoute::path).collect()
291    }
292
293    /// Returns the shared route state.
294    pub fn state(&self) -> &GatewayRouteState {
295        &self.state
296    }
297}
298
299/// Returns the full gateway route table backed by fresh in-memory state.
300pub fn configure_routes() -> GatewayRoutes {
301    configure_routes_with_state(GatewayRouteState::memory())
302}
303
304/// Returns the full gateway route table sharing the supplied `state`.
305///
306/// Registers every OpenAI-compatible endpoint (chat, responses, embeddings,
307/// files, batches, audio, images, vector stores, threads) plus the SIM replay
308/// and admin routes.
309pub fn configure_routes_with_state(state: GatewayRouteState) -> GatewayRoutes {
310    GatewayRoutes::with_state(
311        vec![
312            GatewayRoute::new("GET", HEALTH_PATH, handle_health),
313            GatewayRoute::new("GET", MODELS_PATH, handle_models),
314            GatewayRoute::new("POST", RESPONSES_PATH, handle_responses),
315            GatewayRoute::new("POST", SIM_REPLAY_PATH, handle_sim_replay),
316            GatewayRoute::new("POST", SIM_FORK_PATH, handle_sim_fork),
317            GatewayRoute::new("POST", CHAT_COMPLETIONS_PATH, handle_chat_completions),
318            GatewayRoute::new("POST", EMBEDDINGS_PATH, handle_embeddings),
319            GatewayRoute::new("POST", FILES_PATH, handle_files),
320            GatewayRoute::new("POST", BATCHES_PATH, handle_batches),
321            GatewayRoute::new(
322                "POST",
323                AUDIO_TRANSCRIPTIONS_PATH,
324                handle_audio_transcriptions,
325            ),
326            GatewayRoute::new("POST", AUDIO_SPEECH_PATH, handle_audio_speech),
327            GatewayRoute::new("POST", IMAGES_GENERATIONS_PATH, handle_image_generations),
328            GatewayRoute::new("POST", VECTOR_STORES_PATH, handle_vector_stores),
329            GatewayRoute::prefix_suffix(
330                "POST",
331                VECTOR_STORE_SEARCH_ROUTE,
332                VECTOR_STORE_SEARCH_PREFIX,
333                VECTOR_STORE_SEARCH_SUFFIX,
334                handle_vector_store_search,
335            ),
336            GatewayRoute::prefix(
337                "GET",
338                BATCH_RETRIEVAL_ROUTE,
339                BATCH_RETRIEVAL_PREFIX,
340                handle_batch_retrieval,
341            ),
342            GatewayRoute::prefix(
343                "POST",
344                BATCH_CANCEL_ROUTE,
345                BATCH_RETRIEVAL_PREFIX,
346                handle_batch_cancel,
347            ),
348            GatewayRoute::prefix(
349                "GET",
350                FILE_RETRIEVAL_ROUTE,
351                FILE_RETRIEVAL_PREFIX,
352                handle_file_retrieval,
353            ),
354            GatewayRoute::new("POST", THREADS_PATH, handle_threads),
355            GatewayRoute::prefix(
356                "POST",
357                THREAD_MESSAGES_ROUTE,
358                THREAD_RETRIEVAL_PREFIX,
359                handle_thread_post,
360            ),
361            GatewayRoute::prefix(
362                "GET",
363                THREAD_RETRIEVAL_ROUTE,
364                THREAD_RETRIEVAL_PREFIX,
365                handle_thread_get,
366            ),
367            GatewayRoute::prefix_suffix(
368                "GET",
369                RESPONSE_EVENTS_ROUTE,
370                RESPONSE_RETRIEVAL_PREFIX,
371                RESPONSE_EVENTS_SUFFIX,
372                handle_response_events,
373            ),
374            GatewayRoute::prefix_suffix(
375                "GET",
376                RESPONSE_SIM_ROUTE,
377                RESPONSE_RETRIEVAL_PREFIX,
378                RESPONSE_SIM_SUFFIX,
379                handle_response_sim,
380            ),
381            GatewayRoute::prefix(
382                "GET",
383                RESPONSE_RETRIEVAL_ROUTE,
384                RESPONSE_RETRIEVAL_PREFIX,
385                handle_response_retrieval,
386            ),
387            GatewayRoute::new("GET", ADMIN_RUNS_PATH, handle_admin_runs),
388            GatewayRoute::prefix(
389                "GET",
390                ADMIN_RUN_RETRIEVAL_ROUTE,
391                ADMIN_RUN_RETRIEVAL_PREFIX,
392                handle_admin_run_get,
393            ),
394            GatewayRoute::new("GET", ADMIN_EVENTS_PATH, handle_admin_events),
395            GatewayRoute::new("GET", ADMIN_STORAGE_STATS_PATH, handle_admin_storage_stats),
396            GatewayRoute::new("GET", ADMIN_MODEL_HEALTH_PATH, handle_admin_model_health),
397            GatewayRoute::new("GET", ADMIN_CACHE_STATS_PATH, handle_admin_cache_stats),
398            GatewayRoute::new(
399                "GET",
400                ADMIN_CAPABILITY_REPORT_PATH,
401                handle_admin_capability_report,
402            ),
403        ],
404        state,
405    )
406}
407
408/// A [`GatewayRoutes`] table wrapped as a runtime [`Object`].
409#[derive(Clone)]
410#[non_citizen(
411    reason = "live OpenAI gateway route table handle; route requests use openai/GatewayRequest descriptor",
412    kind = "handle"
413)]
414pub struct GatewayRoutesValue {
415    routes: GatewayRoutes,
416}
417
418impl GatewayRoutesValue {
419    /// Wraps a route table as a runtime object value.
420    pub fn new(routes: GatewayRoutes) -> Self {
421        Self { routes }
422    }
423
424    /// Returns the wrapped route table.
425    pub fn routes(&self) -> &GatewayRoutes {
426        &self.routes
427    }
428}
429
430impl Object for GatewayRoutesValue {
431    fn display(&self, _cx: &mut Cx) -> Result<String> {
432        Ok("#<openai-gateway-routes>".to_owned())
433    }
434
435    fn as_any(&self) -> &dyn std::any::Any {
436        self
437    }
438}
439
440impl ObjectCompat for GatewayRoutesValue {
441    fn as_expr(&self, _cx: &mut Cx) -> Result<Expr> {
442        Ok(Expr::Map(vec![(
443            Expr::Symbol(Symbol::new("routes")),
444            Expr::List(
445                self.routes
446                    .paths()
447                    .into_iter()
448                    .map(|path| Expr::String(path.to_owned()))
449                    .collect(),
450            ),
451        )]))
452    }
453}
454
455/// Returns the [`Ref`] naming the gateway health endpoint (`openai-gateway/health`).
456pub fn health_ref() -> Ref {
457    Ref::Symbol(Symbol::qualified("openai-gateway", "health"))
458}