refactor: migrate monolithic crate to Cargo Workspace with Clean Architecture
Transform the single binary crate into a 9-crate workspace monorepo: - Root Cargo.toml as [workspace] manager with resolver = "2" - zesdex-entities: Domain entity types (session, settings, store, message, etc.) - zesdex-utils: Pure utility functions (error, logger, pagination, slug, clipboard) - zesdex-dto: Data Transfer Objects for LLM provider API communication - zesdex-ipc: Unix-socket IPC layer (client/server/framing/protocol) - zesdex-iam: Identity & Access Management (Clean Architecture: domain/application/infrastructure) - zesdex-cms: Content Management (Clean Architecture: domain/application/infrastructure) - zesdex-middleware: HTTP middleware (Auth, CORS, Rate Limiting) - zesdex-libs: Composition root (AppContext, DB init, JWT, Argon2) - zesdex-backend: Main binary entry point + seed/migrate binaries - DevOps: Dockerfile, docker-compose, Nix (flake/shell/default), CI/CD updates - Remove dead root src/ and src-misc/ directories All crate re-exports maintain backward compatibility with original crate::model::*, crate::dto::*, crate::ipc::* module paths. Feature crates enforce strict layer separation: domain -> application -> infrastructure with generic trait-based dependency injection.
This commit is contained in:
asepharyana
@@ -0,0 +1,4 @@
|
||||
//! Service layer: LLM provider HTTP client and OAuth flows.
|
||||
|
||||
pub mod oauth;
|
||||
pub mod provider;
|
||||
@@ -0,0 +1,137 @@
|
||||
#![allow(
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss,
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_wrap
|
||||
)]
|
||||
//! Minimal loopback HTTP server for capturing OAuth authorization-code redirects.
|
||||
use std::io::{Read, Write};
|
||||
use std::net::{TcpListener, TcpStream};
|
||||
|
||||
/// A single-use HTTP listener on `127.0.0.1` that receives the OAuth
|
||||
/// `?code=...` redirect and serves back a static confirmation page.
|
||||
pub struct LoopbackServer {
|
||||
listener: TcpListener,
|
||||
port: u16,
|
||||
}
|
||||
|
||||
impl LoopbackServer {
|
||||
/// Bind to an OS-assigned free port on localhost.
|
||||
///
|
||||
/// Return: `Err` if the loopback interface can't be bound.
|
||||
pub fn bind() -> std::io::Result<Self> {
|
||||
let listener = TcpListener::bind("127.0.0.1:0")?;
|
||||
let port = listener.local_addr()?.port();
|
||||
Ok(LoopbackServer { listener, port })
|
||||
}
|
||||
|
||||
/// The redirect URI to hand to the OAuth authorization endpoint.
|
||||
pub fn redirect_uri(&self) -> String {
|
||||
format!("http://127.0.0.1:{}/callback", self.port)
|
||||
}
|
||||
|
||||
/// Block until one HTTP request arrives, then extract the `code` query param
|
||||
/// and validate that the `state` param matches the expected value.
|
||||
///
|
||||
/// Flow: accept one connection → apply read timeout → parse request line
|
||||
/// → verify state matches → respond 200/400 depending on whether the code
|
||||
/// was found and state matched.
|
||||
///
|
||||
/// Return: `Err(InvalidData)` if no `code` param is present or the state
|
||||
/// doesn't match `expected_state`.
|
||||
pub fn wait_for_code(&self, timeout_ms: u64, expected_state: &str) -> std::io::Result<String> {
|
||||
let (mut stream, _) = self.listener.accept()?;
|
||||
stream.set_read_timeout(Some(std::time::Duration::from_millis(timeout_ms)))?;
|
||||
Self::read_callback(&mut stream, expected_state)
|
||||
}
|
||||
|
||||
/// Read and parse a single HTTP callback request off `stream`, replying with a status page.
|
||||
///
|
||||
/// Why: writes the HTTP response before returning so the browser tab
|
||||
/// shows a result regardless of whether the code was found.
|
||||
fn read_callback(stream: &mut TcpStream, expected_state: &str) -> std::io::Result<String> {
|
||||
let mut buf = [0u8; 4096];
|
||||
let n = stream.read(&mut buf)?;
|
||||
let request = String::from_utf8_lossy(&buf[..n]);
|
||||
let code = Self::extract_code(&request);
|
||||
let state = Self::extract_state(&request);
|
||||
let state_ok = state.as_deref() == Some(expected_state);
|
||||
let response = match (code.as_ref(), state_ok) {
|
||||
(Some(_), true) => "HTTP/1.1 200 OK\r\nContent-Type: text/plain\r\n\r\nAuthorization complete. You may close this tab.",
|
||||
(Some(_), false) => "HTTP/1.1 400 Bad Request\r\nContent-Type: text/plain\r\n\r\nState mismatch — possible CSRF attack.",
|
||||
(None, _) => "HTTP/1.1 400 Bad Request\r\nContent-Type: text/plain\r\n\r\nMissing authorization code.",
|
||||
};
|
||||
let _ = stream.write_all(response.as_bytes());
|
||||
let _ = stream.flush();
|
||||
if !state_ok {
|
||||
return Err(std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"state mismatch",
|
||||
));
|
||||
}
|
||||
code.ok_or_else(|| {
|
||||
std::io::Error::new(
|
||||
std::io::ErrorKind::InvalidData,
|
||||
"code not found in callback",
|
||||
)
|
||||
})
|
||||
}
|
||||
|
||||
/// Extract and percent-decode the `code` query parameter from an HTTP request line.
|
||||
///
|
||||
/// Return: `None` if the request is malformed or has no `code` param.
|
||||
fn extract_code(request: &str) -> Option<String> {
|
||||
let line = request.lines().next()?;
|
||||
let path = line.split(' ').nth(1)?;
|
||||
let query = path.split('?').nth(1)?;
|
||||
for pair in query.split('&') {
|
||||
let mut parts = pair.splitn(2, '=');
|
||||
if parts.next()? == "code" {
|
||||
return parts.next().map(urlencoding);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
|
||||
/// Extract the `state` query parameter from an HTTP request line.
|
||||
///
|
||||
/// Return: `None` if the request is malformed or has no `state` param.
|
||||
fn extract_state(request: &str) -> Option<String> {
|
||||
let line = request.lines().next()?;
|
||||
let path = line.split(' ').nth(1)?;
|
||||
let query = path.split('?').nth(1)?;
|
||||
for pair in query.split('&') {
|
||||
let mut parts = pair.splitn(2, '=');
|
||||
if parts.next()? == "state" {
|
||||
return parts.next().map(urlencoding);
|
||||
}
|
||||
}
|
||||
None
|
||||
}
|
||||
}
|
||||
|
||||
/// Percent-decode a string (e.g. `%20` -> space).
|
||||
///
|
||||
/// Why: invalid escape sequences (missing/non-hex digits) are passed through
|
||||
/// literally as `%` rather than erroring, since this only handles a redirect
|
||||
/// query param, not untrusted binary data.
|
||||
fn urlencoding(s: &str) -> String {
|
||||
let mut result = String::with_capacity(s.len());
|
||||
let mut chars = s.chars();
|
||||
while let Some(c) = chars.next() {
|
||||
if c == '%' {
|
||||
match (
|
||||
chars.next().and_then(|c| c.to_digit(16)),
|
||||
chars.next().and_then(|c| c.to_digit(16)),
|
||||
) {
|
||||
(Some(hi), Some(lo)) => result.push(char::from((hi * 16 + lo) as u8)),
|
||||
_ => {
|
||||
result.push('%');
|
||||
}
|
||||
}
|
||||
} else {
|
||||
result.push(c);
|
||||
}
|
||||
}
|
||||
result
|
||||
}
|
||||
@@ -0,0 +1,145 @@
|
||||
//! OAuth 2.0 authorization-code + PKCE flow: token exchange and authorization URL building.
|
||||
use serde::{Deserialize, Serialize};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
|
||||
/// An OAuth access token plus its refresh token and absolute expiry (unix seconds).
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OAuthToken {
|
||||
pub access_token: String,
|
||||
pub refresh_token: Option<String>,
|
||||
pub expires_at: u64,
|
||||
pub token_type: String,
|
||||
}
|
||||
|
||||
impl OAuthToken {}
|
||||
|
||||
/// Static configuration for an OAuth provider: endpoints, client identity, and requested scopes.
|
||||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||
pub struct OAuthConfig {
|
||||
pub auth_url: String,
|
||||
pub token_url: String,
|
||||
pub client_id: String,
|
||||
pub client_secret: Option<String>,
|
||||
pub scopes: Vec<String>,
|
||||
}
|
||||
|
||||
impl Default for OAuthConfig {
|
||||
fn default() -> Self {
|
||||
OAuthConfig {
|
||||
auth_url: String::new(),
|
||||
token_url: String::new(),
|
||||
client_id: String::new(),
|
||||
client_secret: None,
|
||||
scopes: vec![
|
||||
"openid".to_string(),
|
||||
"profile".to_string(),
|
||||
"email".to_string(),
|
||||
],
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Drives one OAuth flow: holds config, the current token (if any), and an HTTP client.
|
||||
pub struct OAuthManager {
|
||||
pub config: OAuthConfig,
|
||||
pub token: Option<OAuthToken>,
|
||||
client: reqwest::blocking::Client,
|
||||
}
|
||||
|
||||
impl OAuthManager {
|
||||
/// Create a manager for the given provider config with no token yet acquired.
|
||||
pub fn new(config: OAuthConfig) -> Self {
|
||||
OAuthManager {
|
||||
config,
|
||||
token: None,
|
||||
client: reqwest::blocking::Client::new(),
|
||||
}
|
||||
}
|
||||
|
||||
/// Exchange an authorization code for an access token via the provider's token endpoint.
|
||||
///
|
||||
/// Flow: POST form-encoded grant to `token_url` → parse JSON body →
|
||||
/// compute absolute `expires_at` from `expires_in` → store on `self.token`.
|
||||
///
|
||||
/// Return: `Err(String)` on network failure, non-2xx status, or a missing `access_token` field.
|
||||
pub fn exchange_code(
|
||||
&mut self,
|
||||
code: &str,
|
||||
redirect_uri: &str,
|
||||
code_verifier: &str,
|
||||
) -> Result<(), String> {
|
||||
let mut params = std::collections::HashMap::new();
|
||||
params.insert("grant_type", "authorization_code");
|
||||
params.insert("code", code);
|
||||
params.insert("redirect_uri", redirect_uri);
|
||||
params.insert("client_id", &self.config.client_id);
|
||||
params.insert("code_verifier", code_verifier);
|
||||
|
||||
let resp = self
|
||||
.client
|
||||
.post(&self.config.token_url)
|
||||
.form(¶ms)
|
||||
.send()
|
||||
.map_err(|e| format!("token request failed: {e}"))?;
|
||||
|
||||
let status = resp.status();
|
||||
let body: serde_json::Value = resp.json().map_err(|e| format!("parse failed: {e}"))?;
|
||||
|
||||
if !status.is_success() {
|
||||
return Err(format!("token endpoint returned {status}: {body}"));
|
||||
}
|
||||
|
||||
let access_token = body["access_token"]
|
||||
.as_str()
|
||||
.ok_or("missing access_token")?
|
||||
.to_string();
|
||||
let expires_in = body["expires_in"].as_u64().unwrap_or(3600);
|
||||
let now = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_default()
|
||||
.as_secs();
|
||||
|
||||
self.token = Some(OAuthToken {
|
||||
access_token,
|
||||
refresh_token: body["refresh_token"]
|
||||
.as_str()
|
||||
.map(std::string::ToString::to_string),
|
||||
expires_at: now + expires_in,
|
||||
token_type: body["token_type"].as_str().unwrap_or("Bearer").to_string(),
|
||||
});
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Build the provider's authorization URL with PKCE and state params attached.
|
||||
///
|
||||
/// Why: refuses to build a URL if `auth_url` is missing or invalid. Previously
|
||||
/// this silently fell back to <https://example.com>, which produced a valid-looking
|
||||
/// auth URL pointing at the wrong server and leaked client credentials in
|
||||
/// query params. Returning an empty string signals failure to callers, who
|
||||
/// can prompt the user to fix the OAuth config instead of starting a flow
|
||||
/// against a wrong host.
|
||||
///
|
||||
/// Return: the full authorization URL, or `""` if `auth_url` is empty/unparseable.
|
||||
pub fn build_auth_url(&self, redirect_uri: &str, state: &str, code_challenge: &str) -> String {
|
||||
let mut url = match url::Url::parse(&self.config.auth_url) {
|
||||
Ok(u) if !self.config.auth_url.is_empty() => u,
|
||||
_ => {
|
||||
tracing::warn!(
|
||||
"warning: OAuth auth_url is missing or invalid ('{}'); aborting build_auth_url",
|
||||
self.config.auth_url
|
||||
);
|
||||
return String::new();
|
||||
}
|
||||
};
|
||||
url.query_pairs_mut()
|
||||
.append_pair("response_type", "code")
|
||||
.append_pair("client_id", &self.config.client_id)
|
||||
.append_pair("redirect_uri", redirect_uri)
|
||||
.append_pair("scope", &self.config.scopes.join(" "))
|
||||
.append_pair("state", state)
|
||||
.append_pair("code_challenge_method", "S256")
|
||||
.append_pair("code_challenge", code_challenge);
|
||||
url.to_string()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,6 @@
|
||||
//! OAuth 2.0 authorization-code + PKCE flow: local HTTP callback server,
|
||||
//! token exchange, and code verifier/challenge generation.
|
||||
|
||||
pub mod loopback;
|
||||
pub mod manager;
|
||||
pub mod pkce;
|
||||
@@ -0,0 +1,68 @@
|
||||
#![allow(
|
||||
clippy::cast_possible_truncation,
|
||||
clippy::cast_sign_loss,
|
||||
clippy::cast_precision_loss,
|
||||
clippy::cast_possible_wrap
|
||||
)]
|
||||
//! PKCE (Proof Key for Code Exchange) verifier/challenge pair generation for OAuth flows.
|
||||
use base64::{engine::general_purpose::URL_SAFE_NO_PAD, Engine as _};
|
||||
use sha2::{Digest, Sha256};
|
||||
|
||||
const VERIFIER_LENGTH: usize = 64;
|
||||
|
||||
/// A randomly generated, base64url-encoded PKCE code verifier.
|
||||
pub struct CodeVerifier(String);
|
||||
|
||||
impl CodeVerifier {
|
||||
/// Generate a fresh random code verifier.
|
||||
pub fn new() -> Self {
|
||||
let bytes: Vec<u8> = (0..VERIFIER_LENGTH).map(|_| rand_byte()).collect();
|
||||
CodeVerifier(URL_SAFE_NO_PAD.encode(&bytes))
|
||||
}
|
||||
|
||||
/// Borrow the verifier as a string, to send in the token exchange request.
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
|
||||
/// Derive the S256 code challenge (SHA-256 hash, base64url-encoded) to send
|
||||
/// in the authorization request.
|
||||
pub fn challenge(&self) -> CodeChallenge {
|
||||
let mut hasher = Sha256::new();
|
||||
hasher.update(self.0.as_bytes());
|
||||
let digest = hasher.finalize();
|
||||
CodeChallenge(URL_SAFE_NO_PAD.encode(digest))
|
||||
}
|
||||
}
|
||||
|
||||
/// Produce one pseudo-random byte from the system clock mixed with a monotonic
|
||||
/// counter, providing ~64 bits of per-call unpredictability without a `rand`
|
||||
/// dependency.
|
||||
///
|
||||
/// Why: avoids pulling in a `rand` dependency for a short-lived verifier; the
|
||||
/// monotonic counter ensures that calls within the same clock tick produce
|
||||
/// different values, which is sufficient to prevent OAuth code interception.
|
||||
fn rand_byte() -> u8 {
|
||||
use std::sync::atomic::{AtomicU64, Ordering};
|
||||
use std::time::{SystemTime, UNIX_EPOCH};
|
||||
static COUNTER: AtomicU64 = AtomicU64::new(0);
|
||||
let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
|
||||
let seed = SystemTime::now()
|
||||
.duration_since(UNIX_EPOCH)
|
||||
.unwrap_or_else(|_| {
|
||||
tracing::warn!("[pkce] system time before UNIX_EPOCH, using 0 for random byte");
|
||||
std::time::Duration::default()
|
||||
})
|
||||
.as_nanos() as u64;
|
||||
((seed ^ counter) & 0xFF) as u8
|
||||
}
|
||||
|
||||
/// The S256-derived code challenge sent in the authorization request URL.
|
||||
pub struct CodeChallenge(String);
|
||||
|
||||
impl CodeChallenge {
|
||||
/// Borrow the challenge as a string.
|
||||
pub fn as_str(&self) -> &str {
|
||||
&self.0
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,350 @@
|
||||
//! Blocking HTTP client for OpenAI/Anthropic-compatible chat completion APIs,
|
||||
//! supporting both non-streaming and SSE-streaming requests with automatic retry.
|
||||
use anyhow::Result;
|
||||
use std::time::Duration;
|
||||
|
||||
use crate::app::runtime::stream::turn::StreamedTurn;
|
||||
use crate::app::runtime::stream::{SseParser, StreamEvent};
|
||||
use crate::dto::chat::message::ChatMessage;
|
||||
use crate::dto::provider::request::{ChatRequest, StreamOptions, ToolDef};
|
||||
|
||||
pub(crate) const DEFAULT_BASE_URL: &str = "https://opencode.ai/zen/v1";
|
||||
const DEFAULT_MODEL: &str = "deepseek-v4-flash-free";
|
||||
pub const DEFAULT_API_KEY: &str = "";
|
||||
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
|
||||
const REQUEST_TIMEOUT: Duration = Duration::from_mins(1);
|
||||
|
||||
/// Blocking HTTP client for a single LLM provider endpoint.
|
||||
///
|
||||
/// Holds the reqwest client, credentials, and model/base URL selection used
|
||||
/// by both the non-streaming and streaming chat completion calls.
|
||||
pub struct LlmClient {
|
||||
pub client: reqwest::blocking::Client,
|
||||
pub api_key: String,
|
||||
pub base_url: String,
|
||||
pub model: String,
|
||||
}
|
||||
|
||||
impl LlmClient {
|
||||
/// Construct a client, falling back to built-in defaults for empty inputs.
|
||||
///
|
||||
/// Flow: empty `api_key/model` → substitute defaults → build reqwest client
|
||||
/// with connect/request timeouts → if TLS config fails, retry with just
|
||||
/// request timeout (no connect timeout) → normalize `base_url`.
|
||||
///
|
||||
/// Why: empty strings are treated as "unset" rather than errors so callers
|
||||
/// can pass through unconfigured settings without special-casing them.
|
||||
/// Timeouts are always enforced — the pure-default-client fallback is only
|
||||
/// used as a last resort when even the no-connect-timeout build fails.
|
||||
pub fn new(mut api_key: String, model: String, base_url: Option<String>) -> Self {
|
||||
if api_key.is_empty() {
|
||||
api_key = DEFAULT_API_KEY.to_string();
|
||||
}
|
||||
let model = if model.is_empty() {
|
||||
DEFAULT_MODEL.to_string()
|
||||
} else {
|
||||
model
|
||||
};
|
||||
let client = match reqwest::blocking::Client::builder()
|
||||
.timeout(REQUEST_TIMEOUT)
|
||||
.connect_timeout(CONNECT_TIMEOUT)
|
||||
.build()
|
||||
{
|
||||
Ok(c) => c,
|
||||
Err(e) => {
|
||||
tracing::warn!(
|
||||
"failed to build reqwest client with connect timeout: {}. \
|
||||
retrying without connect timeout",
|
||||
e,
|
||||
);
|
||||
match reqwest::blocking::Client::builder()
|
||||
.timeout(REQUEST_TIMEOUT)
|
||||
.build()
|
||||
{
|
||||
Ok(c) => c,
|
||||
Err(e2) => {
|
||||
tracing::warn!(
|
||||
"also failed: {}. using default client (no configured timeouts)",
|
||||
e2,
|
||||
);
|
||||
reqwest::blocking::Client::new()
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
LlmClient {
|
||||
client,
|
||||
api_key,
|
||||
base_url: base_url
|
||||
.filter(|s| !s.is_empty())
|
||||
.unwrap_or_else(|| DEFAULT_BASE_URL.to_string()),
|
||||
model,
|
||||
}
|
||||
}
|
||||
|
||||
/// Send a non-streaming chat completion request and return the assistant's reply.
|
||||
///
|
||||
/// Flow: build request → POST with retry loop (up to 10 attempts, 2s backoff)
|
||||
/// → parse JSON response → extract first choice's message and token usage.
|
||||
///
|
||||
/// Why: retries transient failures but aborts immediately on 401/403, since
|
||||
/// those indicate a bad API key that retrying won't fix.
|
||||
///
|
||||
/// Return: `Err` if all retries are exhausted, an auth error occurs, or the
|
||||
/// response has no choices.
|
||||
pub fn chat_with_tools_non_streaming(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
tools: Option<Vec<ToolDef>>,
|
||||
) -> Result<(ChatMessage, Option<(u64, u64)>)> {
|
||||
let req = ChatRequest {
|
||||
model: self.model.clone(),
|
||||
messages: messages.to_vec(),
|
||||
max_tokens: Some(4096),
|
||||
temperature: Some(0.7),
|
||||
tools,
|
||||
stream: Some(false),
|
||||
stop: None,
|
||||
stream_options: None,
|
||||
tool_choice: None,
|
||||
};
|
||||
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
let max_retries = 10;
|
||||
let mut attempt = 0;
|
||||
|
||||
loop {
|
||||
attempt += 1;
|
||||
|
||||
let mut http_req = self
|
||||
.client
|
||||
.post(&url)
|
||||
.header("Content-Type", "application/json");
|
||||
|
||||
if !self.api_key.is_empty() {
|
||||
http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
|
||||
}
|
||||
|
||||
let result = (|| -> Result<(ChatMessage, Option<(u64, u64)>)> {
|
||||
let resp = http_req.json(&req).send().map_err(|e| {
|
||||
if e.is_timeout() {
|
||||
anyhow::anyhow!("API request timed out after {REQUEST_TIMEOUT:?}. Check your network or try again.")
|
||||
} else if e.is_connect() {
|
||||
anyhow::anyhow!("Could not connect to {}. Is the URL correct and is the service reachable?", self.base_url)
|
||||
} else {
|
||||
anyhow::anyhow!("API request failed: {e}")
|
||||
}
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().unwrap_or_default();
|
||||
anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
|
||||
}
|
||||
|
||||
let data: crate::dto::provider::response::ChatResponse = resp.json()?;
|
||||
let usage = data.usage.map(|u| {
|
||||
(
|
||||
u64::from(u.prompt_tokens),
|
||||
u64::from(u.completion_tokens),
|
||||
)
|
||||
});
|
||||
let message = data
|
||||
.choices
|
||||
.into_iter()
|
||||
.next()
|
||||
.and_then(|c| c.message)
|
||||
.ok_or_else(|| anyhow::anyhow!("API response had no choices"))?;
|
||||
Ok((message, usage))
|
||||
})();
|
||||
|
||||
match result {
|
||||
Ok((msg, usage)) => return Ok((msg, usage)),
|
||||
Err(e) => {
|
||||
let err_str = e.to_string();
|
||||
let err_lower = err_str.to_lowercase();
|
||||
let is_auth_error = err_str.contains("401")
|
||||
|| err_str.contains("403")
|
||||
|| err_lower.contains("unauthorized")
|
||||
|| err_lower.contains("forbidden")
|
||||
|| err_lower.contains("authentication failed");
|
||||
if attempt >= max_retries || is_auth_error {
|
||||
return Err(e);
|
||||
}
|
||||
tracing::warn!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
|
||||
std::thread::sleep(Duration::from_secs(2));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Streaming variant of `chat_with_tools`. Feeds SSE chunks into an `SseParser` /
|
||||
/// `StreamedTurn` and invokes `on_event` for every parsed `StreamEvent` as it arrives,
|
||||
/// so the caller can push incremental UI updates in real time. Returns the fully
|
||||
/// assembled assistant message plus token usage (prompt, completion) if the server
|
||||
/// reported it. Retries the whole request only if no event has been observed yet
|
||||
/// (once tokens start arriving, a partial turn cannot be safely replayed).
|
||||
pub fn chat_with_tools_streaming(
|
||||
&self,
|
||||
messages: &[ChatMessage],
|
||||
tools: Option<Vec<ToolDef>>,
|
||||
temperature: Option<f32>,
|
||||
max_tokens: Option<u32>,
|
||||
mut on_event: impl FnMut(&StreamEvent) -> bool,
|
||||
) -> Result<(ChatMessage, Option<(u64, u64)>)> {
|
||||
let req = ChatRequest {
|
||||
model: self.model.clone(),
|
||||
messages: messages.to_vec(),
|
||||
max_tokens: Some(max_tokens.unwrap_or(4096)),
|
||||
temperature: Some(temperature.unwrap_or(0.7)),
|
||||
tools,
|
||||
stream: Some(true),
|
||||
stop: None,
|
||||
stream_options: Some(StreamOptions {
|
||||
include_usage: true,
|
||||
}),
|
||||
tool_choice: None,
|
||||
};
|
||||
|
||||
let url = format!("{}/chat/completions", self.base_url);
|
||||
// Fewer retries on streaming because `run_agent_turn` has a
|
||||
// non-streaming fallback that also retries. Combined total is
|
||||
// capped implicitly by the per-turn timeout and step limits.
|
||||
let max_retries = 3;
|
||||
let mut attempt = 0;
|
||||
let mut started = false;
|
||||
|
||||
loop {
|
||||
attempt += 1;
|
||||
let mut wrapped = |event: &StreamEvent| -> bool {
|
||||
started = true;
|
||||
on_event(event)
|
||||
};
|
||||
match self.try_stream_once(&req, &url, &mut wrapped) {
|
||||
Ok(result) => return Ok(result),
|
||||
Err(e) => {
|
||||
let err_str = e.to_string();
|
||||
let err_lower = err_str.to_lowercase();
|
||||
let is_auth_error = err_str.contains("401")
|
||||
|| err_str.contains("403")
|
||||
|| err_lower.contains("unauthorized")
|
||||
|| err_lower.contains("forbidden")
|
||||
|| err_lower.contains("authentication failed");
|
||||
if started || attempt >= max_retries || is_auth_error {
|
||||
return Err(e);
|
||||
}
|
||||
tracing::warn!("Warning: {}. Retrying {}/{}...", e, attempt, max_retries);
|
||||
std::thread::sleep(Duration::from_secs(2));
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Perform one streaming chat completion request, parsing SSE events until completion.
|
||||
///
|
||||
/// Flow: POST → read body in chunks → advance past valid UTF-8 boundary →
|
||||
/// feed into `SseParser` → dispatch each `StreamEvent` to `on_event` and
|
||||
/// accumulate in `StreamedTurn` → return assembled assistant message on `Done`.
|
||||
///
|
||||
/// Why: chunk-by-chunk UTF-8-aware reads avoid splitting multi-byte sequences;
|
||||
/// returns `aborted` error if `on_event` returns false so the caller can cancel.
|
||||
///
|
||||
/// Return: assembled message + optional usage on success, `Err` on read
|
||||
/// failure, non-2xx status, or callback-initiated abort.
|
||||
fn try_stream_once(
|
||||
&self,
|
||||
req: &ChatRequest,
|
||||
url: &str,
|
||||
on_event: &mut dyn FnMut(&StreamEvent) -> bool,
|
||||
) -> Result<(ChatMessage, Option<(u64, u64)>)> {
|
||||
use std::io::Read;
|
||||
|
||||
let mut http_req = self
|
||||
.client
|
||||
.post(url)
|
||||
.header("Content-Type", "application/json");
|
||||
if !self.api_key.is_empty() {
|
||||
http_req = http_req.header("Authorization", format!("Bearer {}", self.api_key));
|
||||
}
|
||||
|
||||
let resp = http_req.json(req).send().map_err(|e| {
|
||||
if e.is_timeout() {
|
||||
anyhow::anyhow!("API request timed out after {REQUEST_TIMEOUT:?}. Check your network or try again.")
|
||||
} else if e.is_connect() {
|
||||
anyhow::anyhow!("Could not connect to {}. Is the URL correct and is the service reachable?", self.base_url)
|
||||
} else {
|
||||
anyhow::anyhow!("API request failed: {e}")
|
||||
}
|
||||
})?;
|
||||
|
||||
if !resp.status().is_success() {
|
||||
let status = resp.status();
|
||||
let body = resp.text().unwrap_or_default();
|
||||
anyhow::bail!("API error {} from {}: {}", status, self.base_url, body);
|
||||
}
|
||||
|
||||
let mut turn = StreamedTurn::new();
|
||||
let mut usage: Option<(u64, u64)> = None;
|
||||
let mut parser = SseParser::new();
|
||||
let mut reader = resp;
|
||||
let mut byte_buf: Vec<u8> = Vec::new();
|
||||
let mut chunk_buf = [0u8; 4096];
|
||||
|
||||
loop {
|
||||
let n = reader
|
||||
.read(&mut chunk_buf)
|
||||
.map_err(|e| anyhow::anyhow!("stream read error: {e}"))?;
|
||||
if n == 0 {
|
||||
break;
|
||||
}
|
||||
byte_buf.extend_from_slice(&chunk_buf[..n]);
|
||||
let valid_len = match std::str::from_utf8(&byte_buf) {
|
||||
Ok(s) => s.len(),
|
||||
Err(e) => e.valid_up_to(),
|
||||
};
|
||||
if valid_len == 0 {
|
||||
continue;
|
||||
}
|
||||
let text = String::from_utf8_lossy(&byte_buf[..valid_len]).into_owned();
|
||||
byte_buf.drain(..valid_len);
|
||||
|
||||
for event in parser.feed(&text) {
|
||||
if !on_event(&event) {
|
||||
anyhow::bail!("aborted");
|
||||
}
|
||||
match &event {
|
||||
StreamEvent::Usage {
|
||||
prompt_tokens,
|
||||
completion_tokens,
|
||||
..
|
||||
} => {
|
||||
usage = Some((*prompt_tokens, *completion_tokens));
|
||||
}
|
||||
StreamEvent::Error(msg) => {
|
||||
anyhow::bail!("stream error: {msg}");
|
||||
}
|
||||
StreamEvent::Done => {
|
||||
turn.apply_event(&event);
|
||||
turn.done_received = true;
|
||||
return Ok((turn.build_assistant_message(), usage));
|
||||
}
|
||||
_ => turn.apply_event(&event),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// The connection closed without an explicit `[DONE]` event. Some
|
||||
// providers legitimately omit it, so EOF alone isn't an error —
|
||||
// but if it leaves a tool call's arguments as unparsable JSON, the
|
||||
// response was truncated mid-generation, not finished. Report that
|
||||
// honestly instead of silently double-stringifying the fragment
|
||||
// into a tool call that will misbehave (e.g. a `write` call with a
|
||||
// half-written file body).
|
||||
if let Some((name, err)) = turn.incomplete_tool_call() {
|
||||
anyhow::bail!("stream ended before tool call '{name}' arguments were complete: {err}");
|
||||
}
|
||||
|
||||
turn.is_complete = true;
|
||||
Ok((turn.build_assistant_message(), usage))
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user