crabtalk_core/
extension.rs

1use crate::{
2    ApiError, BoxFuture, ChatCompletionChunk, ChatCompletionRequest, ChatCompletionResponse, Error,
3    Prefix, storage_key,
4};
5use std::time::Instant;
6
7/// Per-request metadata passed to extension hooks.
8#[derive(Clone, Debug)]
9pub struct RequestContext {
10    pub request_id: String,
11    pub model: String,
12    pub provider: String,
13    pub key_name: Option<String>,
14    pub is_stream: bool,
15    pub started_at: Instant,
16}
17
18/// Error returned by `Extension::on_request` to short-circuit the pipeline.
19/// Converted to an HTTP response in the handler.
20pub struct ExtensionError {
21    pub status: u16,
22    pub body: ApiError,
23}
24
25impl ExtensionError {
26    pub fn new(status: u16, message: impl Into<String>, kind: impl Into<String>) -> Self {
27        Self {
28            status,
29            body: ApiError::new(message, kind),
30        }
31    }
32}
33
34/// Trait for request pipeline extensions (usage tracking, logging, rate limiting, etc.).
35///
36/// Extensions are registered at startup and receive hooks at each stage of request
37/// processing. All methods have default no-op implementations except `name` and `prefix`.
38///
39/// Extensions must be `Send + Sync` for use across async handler tasks.
40/// Hook methods return `BoxFuture` for dyn-compatibility.
41pub trait Extension: Send + Sync {
42    /// Human-readable name for this extension, used in logs and diagnostics.
43    fn name(&self) -> &str;
44
45    /// Fixed 4-byte prefix that namespaces this extension's storage keys.
46    fn prefix(&self) -> Prefix;
47
48    /// Build a full storage key by prepending this extension's prefix to `suffix`.
49    fn storage_key(&self, suffix: &[u8]) -> Vec<u8> {
50        storage_key(&self.prefix(), suffix)
51    }
52
53    /// Check for a cached response before provider dispatch. Return `Some` to
54    /// skip the provider call entirely. Called for non-streaming requests only.
55    fn on_cache_lookup(
56        &self,
57        _request: &ChatCompletionRequest,
58    ) -> BoxFuture<'_, Option<ChatCompletionResponse>> {
59        Box::pin(async { None })
60    }
61
62    /// Called post-auth, pre-dispatch. Return `Err` to short-circuit the pipeline
63    /// (no provider call, no further extensions run).
64    fn on_request(&self, _ctx: &RequestContext) -> BoxFuture<'_, Result<(), ExtensionError>> {
65        Box::pin(async { Ok(()) })
66    }
67
68    /// Called after a non-streaming chat completion response arrives from the provider.
69    fn on_response(
70        &self,
71        _ctx: &RequestContext,
72        _request: &ChatCompletionRequest,
73        _response: &ChatCompletionResponse,
74    ) -> BoxFuture<'_, ()> {
75        Box::pin(async {})
76    }
77
78    /// Called once per SSE chunk during a streaming response, before serialization.
79    fn on_chunk(&self, _ctx: &RequestContext, _chunk: &ChatCompletionChunk) -> BoxFuture<'_, ()> {
80        Box::pin(async {})
81    }
82
83    /// Called when the provider returns an error.
84    fn on_error(&self, _ctx: &RequestContext, _error: &Error) -> BoxFuture<'_, ()> {
85        Box::pin(async {})
86    }
87}
crabtalk_core/extension.rs

crabtalk_core/
extension.rs
