crabllm_core/
extension.rs

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

crabllm_core/
extension.rs
