diff --git a/Cargo.lock b/Cargo.lock index c13f0ad..3bba41e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -599,6 +599,12 @@ dependencies = [ "fastrand", ] +[[package]] +name = "base16ct" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" + [[package]] name = "base64" version = "0.22.1" @@ -689,6 +695,7 @@ dependencies = [ name = "cache" version = "0.1.0" dependencies = [ + "lazy_static", "redis", "reqwest", "serde", @@ -958,6 +965,18 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "crypto-bigint" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "subtle", + "zeroize", +] + [[package]] name = "crypto-common" version = "0.1.7" @@ -986,6 +1005,33 @@ dependencies = [ "cmov", ] +[[package]] +name = "curve25519-dalek" +version = "4.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "curve25519-dalek-derive", + "digest 0.10.7", + "fiat-crypto", + "rustc_version", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + [[package]] name = "customers" version = "0.1.0" @@ -1014,6 +1060,8 @@ dependencies = [ "serde", "serde_json", "sqlx", + "thiserror", + "tokio", "tracing", "uuid", ] @@ -1116,6 +1164,44 @@ version = "1.0.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" +[[package]] +name = "ecdsa" +version = "0.16.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" +dependencies = [ + "der", + "digest 0.10.7", + "elliptic-curve", + "rfc6979", + "signature", + "spki", +] + +[[package]] +name = "ed25519" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +dependencies = [ + "pkcs8", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +dependencies = [ + "curve25519-dalek", + "ed25519", + "serde", + "sha2 0.10.9", + "subtle", + "zeroize", +] + [[package]] name = "either" version = "1.16.0" @@ -1125,6 +1211,27 @@ dependencies = [ "serde", ] +[[package]] +name = "elliptic-curve" +version = "0.13.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47" +dependencies = [ + "base16ct", + "crypto-bigint", + "digest 0.10.7", + "ff", + "generic-array", + "group", + "hkdf", + "pem-rfc7468", + "pkcs8", + "rand_core 0.6.4", + "sec1", + "subtle", + "zeroize", +] + [[package]] name = "email" version = "0.1.0" @@ -1228,6 +1335,22 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" +[[package]] +name = "ff" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393" +dependencies = [ + "rand_core 0.6.4", + "subtle", +] + +[[package]] +name = "fiat-crypto" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" + [[package]] name = "find-msvc-tools" version = "0.1.9" @@ -1434,6 +1557,7 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", + "zeroize", ] [[package]] @@ -1495,6 +1619,17 @@ dependencies = [ "uuid", ] +[[package]] +name = "group" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" +dependencies = [ + "ff", + "rand_core 0.6.4", + "subtle", +] + [[package]] name = "h2" version = "0.3.27" @@ -2048,11 +2183,18 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "eba32bfb4ffdeaca3e34431072faf01745c9b26d25504aa7a6cf5684334fc4fc" dependencies = [ "base64", + "ed25519-dalek", "getrandom 0.2.17", + "hmac 0.12.1", "js-sys", + "p256", + "p384", "pem", + "rand 0.8.6", + "rsa", "serde", "serde_json", + "sha2 0.10.9", "signature", "simple_asn1", "zeroize", @@ -2451,6 +2593,30 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" +[[package]] +name = "p256" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2 0.10.9", +] + +[[package]] +name = "p384" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe42f1670a52a47d448f14b6a5c61dd78fce51856e68edaa38f7ae3a46b8d6b6" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2 0.10.9", +] + [[package]] name = "parking" version = "2.2.1" @@ -2499,9 +2665,13 @@ dependencies = [ "axum", "chrono", "contracts", + "db", + "hex", + "rand 0.8.6", "reqwest", "serde", "serde_json", + "sha2 0.10.9", "sqlx", "tokio", "tracing", @@ -2632,6 +2802,15 @@ dependencies = [ "syn", ] +[[package]] +name = "primeorder" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6" +dependencies = [ + "elliptic-curve", +] + [[package]] name = "proc-macro2" version = "1.0.106" @@ -2909,6 +3088,16 @@ dependencies = [ "webpki-roots 1.0.7", ] +[[package]] +name = "rfc6979" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" +dependencies = [ + "hmac 0.12.1", + "subtle", +] + [[package]] name = "ring" version = "0.17.14" @@ -3080,6 +3269,20 @@ dependencies = [ "untrusted", ] +[[package]] +name = "sec1" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc" +dependencies = [ + "base16ct", + "der", + "generic-array", + "pkcs8", + "subtle", + "zeroize", +] + [[package]] name = "security-framework" version = "3.7.0" diff --git a/apps/cron/src/main.rs b/apps/cron/src/main.rs index 752bd0a..edd15ee 100644 --- a/apps/cron/src/main.rs +++ b/apps/cron/src/main.rs @@ -75,6 +75,64 @@ async fn main() -> Result<(), Box> { } }); + // Spawn AI credits reservation-hold reaper (every 2 minutes -- holds + // expire after 5 minutes by default, so this catches abandoned holds + // promptly without hammering the wallet table). + let p_ai_reap = pool.clone(); + tokio::spawn(async move { + let mut interval = time::interval(Duration::from_secs(2 * 60)); + loop { + interval.tick().await; + if let Err(e) = tasks::ai_credits::sweep_expired_reservation_holds(&p_ai_reap).await { + tracing::error!("AI Credit Reservation Reaper Failed: {}", e); + } + } + }); + + // Spawn AI credits daily-usage-counter reset (hourly -- idempotent + // no-op for any wallet whose counters are already current; exists so + // admin-facing "today's usage" is accurate even for untouched wallets, + // not because charge-path correctness depends on it). + let p_ai_reset = pool.clone(); + tokio::spawn(async move { + let mut interval = time::interval(Duration::from_secs(60 * 60)); + loop { + interval.tick().await; + if let Err(e) = tasks::ai_credits::reset_stale_daily_ai_usage(&p_ai_reset).await { + tracing::error!("AI Credit Daily Usage Reset Failed: {}", e); + } + } + }); + + // Spawn AI credits subscription lifecycle tasks (hourly) + // Task 7: Apply scheduled downgrades and expire trials + let p_ai_sub = pool.clone(); + tokio::spawn(async move { + let mut interval = time::interval(Duration::from_secs(60 * 60)); + loop { + interval.tick().await; + if let Err(e) = tasks::ai_credits::apply_scheduled_downgrades(&p_ai_sub).await { + tracing::error!("AI Credit Scheduled Downgrades Failed: {}", e); + } + if let Err(e) = tasks::ai_credits::expire_trials(&p_ai_sub).await { + tracing::error!("AI Credit Trial Expiry Failed: {}", e); + } + } + }); + + // Spawn AI credits purchased credit expiration (daily) + // Task 8: Expire purchased credits past their expiration date + let p_ai_expiry = pool.clone(); + tokio::spawn(async move { + let mut interval = time::interval(Duration::from_secs(24 * 60 * 60)); + loop { + interval.tick().await; + if let Err(e) = tasks::ai_credits::expire_purchased_credits(&p_ai_expiry).await { + tracing::error!("AI Credit Purchased Credits Expiry Failed: {}", e); + } + } + }); + // Keep main thread alive tokio::signal::ctrl_c().await?; tracing::info!("Shutting down cron engine."); diff --git a/apps/cron/src/tasks/ai_credits.rs b/apps/cron/src/tasks/ai_credits.rs new file mode 100644 index 0000000..001b999 --- /dev/null +++ b/apps/cron/src/tasks/ai_credits.rs @@ -0,0 +1,168 @@ +//! Ask Ash AI credits background jobs (Phase 2 of +//! docs/ASK_ASH_BILLING_ARCHITECTURE.md in nxtgauge-ai-assistant). +//! +//! Phase 1 shipped the reserve/capture/release wallet primitives +//! (crates/db/src/models/ai_credits.rs) with an explicit note that +//! abandoned reservation holds have no reaper yet -- this is that reaper. + +use sqlx::PgPool; + +/// Release any `held` reservation past its `expires_at` back to the +/// owning wallet's available balance. A hold is only left in this state +/// by a crash or bug between `try_reserve_credits` and the matching +/// `try_capture_reservation`/`try_release_reservation` call (normal +/// success/failure paths always resolve the hold immediately via +/// `apps/users/src/ai_credits.rs::charge_ai_feature`) -- this job is the +/// backstop for that abnormal case, not the primary release path. +pub async fn sweep_expired_reservation_holds( + pool: &PgPool, +) -> Result> { + let expired: Vec = sqlx::query_scalar( + "SELECT id FROM ai_reservation_holds WHERE status = 'held' AND expires_at < NOW()", + ) + .fetch_all(pool) + .await?; + + let mut released = 0u64; + for hold_id in expired { + match db::models::ai_credits::AiCreditsRepository::try_release_reservation(pool, hold_id).await { + Ok(true) => released += 1, + Ok(false) => {} // already resolved by the normal path between our SELECT and this call + Err(e) => tracing::error!("Failed to release expired AI credit hold {}: {}", hold_id, e), + } + } + + if released > 0 { + tracing::info!("Released {} expired AI credit reservation hold(s).", released); + } + Ok(released) +} + +/// Roll `daily_actions_used`/`daily_credits_used` back to zero for any +/// wallet whose `daily_usage_date` has fallen behind today. Per-wallet +/// counters also self-heal inline on next use +/// (`AiCreditsRepository::try_reserve_credits` resets a stale wallet the +/// moment it's touched), so this batch job exists only so admin-facing +/// "today's usage" figures are accurate even for wallets nobody has +/// touched yet today -- it is not required for charge-path correctness. +pub async fn reset_stale_daily_ai_usage( + pool: &PgPool, +) -> Result> { + let result = sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET daily_actions_used = 0, daily_credits_used = 0, daily_usage_date = CURRENT_DATE + WHERE daily_usage_date < CURRENT_DATE + "#, + ) + .execute(pool) + .await?; + + let rows = result.rows_affected(); + if rows > 0 { + tracing::info!("Reset stale daily AI usage counters for {} wallet(s).", rows); + } + Ok(rows) +} + +/// Apply scheduled plan downgrades that have reached their effective date +/// (Task 7 - Subscription Lifecycle) +pub async fn apply_scheduled_downgrades( + pool: &PgPool, +) -> Result> { + let applied = crate::ai_subscription::apply_scheduled_downgrades(pool) + .await + .map_err(|e| Box::new(e) as Box)?; + + if applied > 0 { + tracing::info!("Applied {} scheduled plan downgrades.", applied); + } + Ok(applied as u64) +} + +/// Expire trials that have reached their end date +/// (Task 7 - Subscription Lifecycle) +pub async fn expire_trials( + pool: &PgPool, +) -> Result> { + let expired = crate::ai_subscription::expire_trials(pool) + .await + .map_err(|e| Box::new(e) as Box)?; + + if expired > 0 { + tracing::info!("Expired {} trials.", expired); + } + Ok(expired as u64) +} + +/// Expire purchased credits that have passed their expiration date +/// (Task 8 - Credit Expiration Enforcement) +pub async fn expire_purchased_credits( + pool: &PgPool, +) -> Result> { + let expired: Vec<(uuid::Uuid, i32, i32)> = sqlx::query_as( + r#" + SELECT id, user_id, + (purchased_credits_total - purchased_credits_used) as unused_credits + FROM user_ai_subscriptions + WHERE purchased_credits_expire_at IS NOT NULL + AND purchased_credits_expire_at < NOW() + AND purchased_credits_total > purchased_credits_used + "# + ) + .fetch_all(pool) + .await?; + + let mut expired_count = 0u64; + + for (wallet_id, user_id, unused_credits) in expired { + if unused_credits <= 0 { + continue; + } + + let mut tx = pool.begin().await?; + + // Mark credits as used (expired) + sqlx::query( + "UPDATE user_ai_subscriptions SET purchased_credits_used = purchased_credits_total, updated_at = NOW() WHERE id = $1" + ) + .bind(wallet_id) + .execute(&mut *tx) + .await?; + + // Create ledger entry for expiration + sqlx::query( + r#" + INSERT INTO ai_credit_ledger + (wallet_id, entry_type, credits, balance_after, reference_type, actor_type, metadata) + SELECT + $1, + 'expiration', + -$2, + (bonus_credits_total - bonus_credits_used + monthly_credits_total - monthly_credits_used + 0 - reserved_credits - locked_credits), + 'expiration', + 'system', + jsonb_build_object('reason', 'purchased_credits_expired', 'expired_credits', $2) + FROM user_ai_subscriptions WHERE id = $1 + "# + ) + .bind(wallet_id) + .bind(unused_credits) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + expired_count += 1; + + tracing::info!( + "Expired {} purchased credits for user {}", + unused_credits, user_id + ); + } + + if expired_count > 0 { + tracing::info!("Expired purchased credits for {} wallet(s).", expired_count); + } + + Ok(expired_count) +} diff --git a/apps/cron/src/tasks/mod.rs b/apps/cron/src/tasks/mod.rs index a48bd41..42bf577 100644 --- a/apps/cron/src/tasks/mod.rs +++ b/apps/cron/src/tasks/mod.rs @@ -2,3 +2,4 @@ pub mod leads; pub mod requirements; pub mod jobs; pub mod reminders; +pub mod ai_credits; diff --git a/apps/gateway/src/main.rs b/apps/gateway/src/main.rs index 54cf951..03e2adc 100644 --- a/apps/gateway/src/main.rs +++ b/apps/gateway/src/main.rs @@ -200,10 +200,29 @@ impl Services { else if path.starts_with("/api/credits") { Some(self.payments_url.clone()) } + // ── Ask Ash AI Credits — split across two services (see docs/ASK_ASH_BILLING_ARCHITECTURE.md). + // Wallet balance + feature-charge live in the users service; package + // catalog + PayU order/verify live in the payments service. Both + // sub-paths start with "/api/ai", so they must be matched BEFORE the + // generic AI-chat rule below, or everything here silently falls + // through to the users service and 404s. + else if path.starts_with("/api/ai-credits/wallet") || path.starts_with("/api/ai-credits/charge") { + Some(self.users_url.clone()) + } + else if path.starts_with("/api/ai-credits") { + Some(self.payments_url.clone()) + } // ── AI Chat (routes to users service, which calls Ollama directly) ─── else if path.starts_with("/api/ai") { Some(self.users_url.clone()) } + // Admin: AI credit package catalog lives alongside the public + // catalog + PayU order/verify in the payments service. All other + // /api/admin/ai-credits/* (balance/ledger/adjust/reconcile) falls + // through to the generic /api/admin/ -> users rule below. + else if path.starts_with("/api/admin/ai-credits/packages") { + Some(self.payments_url.clone()) + } // Admin runtime config management defaults to users service else if path.starts_with("/api/admin/runtime-configs") { Some(self.users_url.clone()) @@ -343,3 +362,67 @@ async fn proxy_handler( } } } + +#[cfg(test)] +mod tests { + use super::*; + + fn test_services() -> Services { + Services { + users_url: "http://users".to_string(), + companies_url: "http://companies".to_string(), + jobs_url: "http://jobs".to_string(), + leads_url: "http://leads".to_string(), + job_seekers_url: "http://job-seekers".to_string(), + customers_url: "http://customers".to_string(), + photographers_url: "http://photographers".to_string(), + makeup_artists_url: "http://makeup-artists".to_string(), + tutors_url: "http://tutors".to_string(), + developers_url: "http://developers".to_string(), + video_editors_url: "http://video-editors".to_string(), + graphic_designers_url: "http://graphic-designers".to_string(), + social_media_managers_url: "http://social-media-managers".to_string(), + fitness_trainers_url: "http://fitness-trainers".to_string(), + catering_services_url: "http://catering-services".to_string(), + ugc_content_creators_url: "http://ugc-content-creators".to_string(), + payments_url: "http://payments".to_string(), + employees_url: "http://employees".to_string(), + client: reqwest::Client::new(), + } + } + + /// Regression test: /api/ai-credits/* must split correctly across + /// users (wallet/charge) and payments (packages/order/verify), and + /// must NOT fall through to the generic /api/ai -> users rule, which + /// it would silently do if the ai-credits-specific checks were ever + /// removed or reordered after the generic AI-chat rule. + #[test] + fn ai_credits_routes_split_correctly_between_users_and_payments() { + let services = test_services(); + + assert_eq!(services.resolve_upstream("/api/ai-credits/wallet"), Some("http://users".to_string())); + assert_eq!(services.resolve_upstream("/api/ai-credits/charge"), Some("http://users".to_string())); + + assert_eq!(services.resolve_upstream("/api/ai-credits"), Some("http://payments".to_string())); + assert_eq!(services.resolve_upstream("/api/ai-credits/order"), Some("http://payments".to_string())); + assert_eq!(services.resolve_upstream("/api/ai-credits/verify"), Some("http://payments".to_string())); + + // Plain /api/ai/* (chat, generate-job-field, etc.) is unaffected. + assert_eq!(services.resolve_upstream("/api/ai/generate-job-field"), Some("http://users".to_string())); + } + + #[test] + fn admin_ai_credits_routes_split_correctly() { + let services = test_services(); + + assert_eq!(services.resolve_upstream("/api/admin/ai-credits/packages"), Some("http://payments".to_string())); + assert_eq!(services.resolve_upstream("/api/admin/ai-credits/packages/abc-123"), Some("http://payments".to_string())); + + // Everything else under /api/admin/ai-credits/* falls through to + // the generic /api/admin/ -> users rule. + assert_eq!(services.resolve_upstream("/api/admin/ai-credits/balance"), Some("http://users".to_string())); + assert_eq!(services.resolve_upstream("/api/admin/ai-credits/ledger"), Some("http://users".to_string())); + assert_eq!(services.resolve_upstream("/api/admin/ai-credits/adjust"), Some("http://users".to_string())); + assert_eq!(services.resolve_upstream("/api/admin/ai-credits/reconcile"), Some("http://users".to_string())); + } +} diff --git a/apps/payments/Cargo.toml b/apps/payments/Cargo.toml index 86cf510..0ad1e6b 100644 --- a/apps/payments/Cargo.toml +++ b/apps/payments/Cargo.toml @@ -13,6 +13,10 @@ tracing-subscriber.workspace = true reqwest = { version = "0.12", features = ["json", "stream"] } anyhow.workspace = true contracts = { path = "../../crates/contracts" } +db = { path = "../../crates/db" } sqlx.workspace = true uuid.workspace = true chrono.workspace = true +sha2 = "0.10" +hex = "0.4" +rand = "0.8" diff --git a/apps/payments/src/ai_credits.rs b/apps/payments/src/ai_credits.rs new file mode 100644 index 0000000..362f638 --- /dev/null +++ b/apps/payments/src/ai_credits.rs @@ -0,0 +1,542 @@ +//! AI credit package purchase flow via PayU, matching the already-shipped +//! frontend contract in +//! nxtgauge-frontend-solid/src/components/dashboard/CreditsPage.tsx +//! (GET /api/ai-credits, POST /api/ai-credits/order, POST /api/ai-credits/verify). +//! +//! Deliberately separate from the TraceCoins purchase flow in main.rs -- +//! its own table (`ai_credit_orders`), its own wallet +//! (`db::models::ai_credits::AiCreditsRepository`), same PayU integration +//! (crate::payu) reused for the hash logic. + +use crate::payu; +use crate::AppState; +use axum::{ + extract::{Path, State}, + http::StatusCode, + routing::{get, patch, post}, + Json, Router, +}; +use contracts::auth_middleware::{require_admin, AuthUser}; +use db::models::ai_credits::AiCreditsRepository; +use serde::{Deserialize, Serialize}; +use sqlx::FromRow; +use uuid::Uuid; + +pub fn router() -> Router { + Router::new() + .route("/", get(list_packages)) + .route("/order", post(create_order)) + .route("/verify", post(verify_order)) +} + +/// Admin CRUD for the AI credit package catalog -- mounted at +/// /api/admin/ai-credits/packages by the gateway (see apps/gateway's +/// resolve_upstream). Mirrors the shape of the existing TraceCoins +/// package admin form in nxtgauge-admin-solid's pricing.tsx, but against +/// `ai_credit_packages`, not `pricing_packages`. +pub fn admin_router() -> Router { + Router::new() + .route("/", get(admin_list_packages).post(admin_create_package)) + .route("/{id}", patch(admin_update_package)) +} + +#[derive(Debug, FromRow, Serialize)] +struct AiCreditPackageRow { + id: Uuid, + name: String, + description: Option, + credits: i32, + price_inr: i32, +} + +#[derive(Serialize)] +struct ListPackagesResponse { + packages: Vec, +} + +async fn list_packages(State(state): State) -> Result, (StatusCode, String)> { + let packages = sqlx::query_as::<_, AiCreditPackageRow>( + "SELECT id, name, description, credits, price_inr FROM ai_credit_packages WHERE is_active = TRUE ORDER BY price_inr", + ) + .fetch_all(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + Ok(Json(ListPackagesResponse { packages })) +} + +#[derive(Debug, Deserialize)] +struct CreateAiOrderRequest { + package_id: String, + coupon_code: Option, +} + +#[derive(Debug, Serialize)] +struct CreateAiOrderResponse { + key: String, + txnid: String, + amount: String, + productinfo: String, + firstname: String, + email: String, + phone: String, + surl: String, + furl: String, + hash: String, + payu_base_url: String, + udf1: String, + udf2: String, + order_id: String, + currency: String, +} + +#[derive(Debug, FromRow)] +struct AiCreditPackagePriceRow { + name: String, + credits: i32, + price_inr: i32, +} + +#[derive(Debug, FromRow)] +struct UserContactRow { + email: String, + full_name: Option, + phone: Option, +} + +/// Validate and apply a coupon code +async fn validate_coupon( + pool: &PgPool, + user_id: Uuid, + coupon_code: &str, + package_id: Uuid, + original_price: i32, +) -> Result, String> { + // Get coupon details + let coupon: Option<(Uuid, String, rust_decimal::Decimal, Option, i32, i32, Vec)> = sqlx::query_as( + r#" + SELECT + id, discount_type, discount_value, max_discount_amount, + max_redemptions_per_user, redemptions_used, applicable_package_ids + FROM ai_coupons + WHERE code = $1 + AND is_active = TRUE + AND valid_from <= NOW() + AND (valid_until IS NULL OR valid_until > NOW()) + "# + ) + .bind(coupon_code) + .fetch_optional(pool) + .await + .map_err(|e| format!("DB error: {e}"))?; + + let Some((coupon_id, discount_type, discount_value, max_discount_amount, max_per_user, total_redemptions, applicable_packages)) = coupon else { + return Ok(None); // Coupon not found or invalid + }; + + // Check if coupon has remaining redemptions + if total_redemptions >= max_redemptions_per_user { + return Err("Coupon redemption limit reached".to_string()); + } + + // Check if user already redeemed this coupon + let user_redemptions: i64 = sqlx::query_scalar( + "SELECT COUNT(*) FROM ai_coupon_redemptions WHERE coupon_id = $1 AND user_id = $2" + ) + .bind(coupon_id) + .bind(user_id) + .fetch_one(pool) + .await + .map_err(|e| format!("DB error: {e}"))?; + + if user_redemptions >= max_per_user as i64 { + return Err("You have already used this coupon".to_string()); + } + + // Check if package is applicable + if !applicable_packages.is_empty() && !applicable_packages.contains(&package_id) { + return Err("Coupon not applicable to this package".to_string()); + } + + // Calculate discount + let original_price_decimal = rust_decimal::Decimal::from(original_price) / rust_decimal::Decimal::from(100); // Convert paise to rupees + let discount = if discount_type == "percentage" { + original_price_decimal * (discount_value / rust_decimal::Decimal::from(100)) + } else { + discount_value + }; + + // Apply max discount limit + let final_discount = if let Some(max) = max_discount_amount { + discount.min(max) + } else { + discount + }; + + // Convert back to paise + let discount_paise = (final_discount * rust_decimal::Decimal::from(100)).to_i32().unwrap_or(0); + let discounted_price = original_price - discount_paise; + + Ok(Some((discount_paise, coupon_id.to_string(), discounted_price))) +} + +async fn create_order( + auth: AuthUser, + State(state): State, + Json(payload): Json, +) -> Result, (StatusCode, String)> { + let package_id = Uuid::parse_str(&payload.package_id) + .map_err(|_| (StatusCode::BAD_REQUEST, "Invalid package id".to_string()))?; + + let package = sqlx::query_as::<_, AiCreditPackagePriceRow>( + "SELECT name, credits, price_inr FROM ai_credit_packages WHERE id = $1 AND is_active = TRUE", + ) + .bind(package_id) + .fetch_optional(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))? + .ok_or((StatusCode::BAD_REQUEST, "Invalid or inactive AI credit package".to_string()))?; + + let contact = sqlx::query_as::<_, UserContactRow>("SELECT email, full_name, phone FROM users WHERE id = $1") + .bind(auth.user_id) + .fetch_optional(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))? + .ok_or((StatusCode::NOT_FOUND, "User not found".to_string()))?; + + let firstname = contact + .full_name + .as_deref() + .and_then(|n| n.split_whitespace().next()) + .unwrap_or("Customer") + .to_string(); + let phone = contact.phone.unwrap_or_default(); + + // Validate coupon if provided + let (final_price, discount_applied, coupon_id) = if let Some(code) = &payload.coupon_code { + match validate_coupon(&state.pool, auth.user_id, code, package_id, package.price_inr).await { + Ok(Some((discount, coupon_id_str, discounted_price))) => { + (discounted_price, discount, Some(coupon_id_str)) + } + Ok(None) => return Err((StatusCode::BAD_REQUEST, "Invalid coupon code".to_string())), + Err(e) => return Err((StatusCode::BAD_REQUEST, e)), + } + } else { + (package.price_inr, 0, None) + }; + + let txnid = payu::generate_txnid(); + let amount_str = payu::paise_to_rupee_string(final_price); + let productinfo = package.name.clone(); + let udf1 = package_id.to_string(); + let udf2 = coupon_id.unwrap_or_default(); // Store coupon ID in UDF2 + + let hash = payu::request_hash( + &state.payu, + &txnid, + &amount_str, + &productinfo, + &firstname, + &contact.email, + &udf1, + &udf2, + ); + + sqlx::query( + r#" + INSERT INTO ai_credit_orders (user_id, package_id, txnid, amount_inr, credits, status, coupon_code, discount_applied) + VALUES ($1, $2, $3, $4, $5, 'PENDING', $6, $7) + "#, + ) + .bind(auth.user_id) + .bind(package_id) + .bind(&txnid) + .bind(final_price) + .bind(package.credits) + .bind(payload.coupon_code.as_ref()) + .bind(discount_applied) + .execute(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + Ok(Json(CreateAiOrderResponse { + key: state.payu.merchant_key.clone(), + txnid: txnid.clone(), + amount: amount_str, + productinfo, + firstname, + email: contact.email, + phone, + surl: state.payu.surl.clone(), + furl: state.payu.furl.clone(), + hash, + payu_base_url: state.payu.base_url.clone(), + udf1, + udf2, + order_id: txnid, + currency: "INR".to_string(), + })) +} + +#[derive(Debug, Deserialize)] +struct VerifyAiOrderRequest { + txnid: String, + mihpayid: String, + status: String, + hash: String, + amount: String, + productinfo: String, + firstname: String, + email: String, + #[serde(default)] + udf1: Option, + #[serde(default)] + udf2: Option, +} + +#[derive(Debug, Serialize)] +struct VerifyAiOrderResponse { + verified: bool, + credits_added: i32, +} + +#[derive(Debug, FromRow)] +struct AiCreditOrderRow { + id: Uuid, + user_id: Uuid, + credits: i32, + coupon_code: Option, + discount_applied: Option, + amount_inr: i32, + package_id: Uuid, +} + +async fn verify_order( + auth: AuthUser, + State(state): State, + Json(payload): Json, +) -> Result, (StatusCode, String)> { + if !payu::verify_response_hash( + &state.payu, + &payload.status, + &payload.txnid, + &payload.amount, + &payload.productinfo, + &payload.firstname, + &payload.email, + payload.udf1.as_deref().unwrap_or(""), + payload.udf2.as_deref().unwrap_or(""), + &payload.hash, + ) { + return Err((StatusCode::BAD_REQUEST, "Payment hash verification failed".to_string())); + } + + if !payload.status.eq_ignore_ascii_case("success") { + return Err((StatusCode::BAD_REQUEST, "Payment was not successful".to_string())); + } + + let order = sqlx::query_as::<_, AiCreditOrderRow>( + "SELECT id, user_id, credits, coupon_code, discount_applied, amount_inr, package_id FROM ai_credit_orders WHERE txnid = $1 AND status = 'PENDING'", + ) + .bind(&payload.txnid) + .fetch_optional(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))? + .ok_or((StatusCode::NOT_FOUND, "Order not found or already processed".to_string()))?; + + if order.user_id != auth.user_id { + return Err((StatusCode::FORBIDDEN, "Order does not belong to user".to_string())); + } + + // Idempotency key ties this credit grant to the specific order -- + // AiCreditsRepository::add_purchased_credits no-ops on a retried call + // with the same key rather than double-crediting. + let idempotency_key = format!("ai-credit-order:{}", order.id); + AiCreditsRepository::ensure_wallet(&state.pool, order.user_id) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + AiCreditsRepository::add_purchased_credits( + &state.pool, + order.user_id, + order.credits, + &idempotency_key, + "ai_credit_order", + Some(order.id), + ) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + // Track coupon redemption if coupon was used + if let Some(coupon_code) = order.coupon_code { + let coupon_id: Option = sqlx::query_scalar("SELECT id FROM ai_coupons WHERE code = $1") + .bind(&coupon_code) + .fetch_optional(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + if let Some(coupon_id) = coupon_id { + // Record redemption + sqlx::query( + r#" + INSERT INTO ai_coupon_redemptions + (coupon_id, user_id, order_id, credits_purchased, discount_applied, final_amount_paid) + VALUES ($1, $2, $3, $4, $5, $6) + ON CONFLICT (coupon_id, user_id) DO UPDATE SET + order_id = EXCLUDED.order_id, + credits_purchased = EXCLUDED.credits_purchased, + discount_applied = EXCLUDED.discount_applied, + final_amount_paid = EXCLUDED.final_amount_paid, + redeemed_at = NOW() + "# + ) + .bind(coupon_id) + .bind(order.user_id) + .bind(order.id) + .bind(order.credits) + .bind(order.discount_applied.unwrap_or(0)) + .bind(order.amount_inr) + .execute(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + // Update coupon redemption count + sqlx::query("UPDATE ai_coupons SET redemptions_used = redemptions_used + 1 WHERE id = $1") + .bind(coupon_id) + .execute(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + } + } + + sqlx::query("UPDATE ai_credit_orders SET status = 'SUCCESS', payu_payment_id = $1, verified_at = NOW() WHERE id = $2") + .bind(&payload.mihpayid) + .bind(order.id) + .execute(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + Ok(Json(VerifyAiOrderResponse { + verified: true, + credits_added: order.credits, + })) +} + +// ── Admin package management ──────────────────────────────────────────── + +#[derive(Debug, FromRow, Serialize)] +struct AdminAiCreditPackageRow { + id: Uuid, + name: String, + description: Option, + credits: i32, + price_inr: i32, + is_active: bool, + created_at: chrono::DateTime, + updated_at: chrono::DateTime, +} + +#[derive(Serialize)] +struct AdminListPackagesResponse { + packages: Vec, +} + +async fn admin_list_packages( + auth: AuthUser, + State(state): State, +) -> Result, (StatusCode, String)> { + require_admin(&auth).map_err(|_| (StatusCode::FORBIDDEN, "Admin access required".to_string()))?; + + let packages = sqlx::query_as::<_, AdminAiCreditPackageRow>( + "SELECT id, name, description, credits, price_inr, is_active, created_at, updated_at FROM ai_credit_packages ORDER BY price_inr", + ) + .fetch_all(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + Ok(Json(AdminListPackagesResponse { packages })) +} + +#[derive(Debug, Deserialize)] +struct CreatePackageRequest { + name: String, + description: Option, + credits: i32, + price_inr: i32, +} + +async fn admin_create_package( + auth: AuthUser, + State(state): State, + Json(body): Json, +) -> Result, (StatusCode, String)> { + require_admin(&auth).map_err(|_| (StatusCode::FORBIDDEN, "Admin access required".to_string()))?; + + if body.credits <= 0 || body.price_inr <= 0 { + return Err((StatusCode::BAD_REQUEST, "credits and price_inr must be positive".to_string())); + } + + let package = sqlx::query_as::<_, AdminAiCreditPackageRow>( + r#" + INSERT INTO ai_credit_packages (name, description, credits, price_inr) + VALUES ($1, $2, $3, $4) + RETURNING id, name, description, credits, price_inr, is_active, created_at, updated_at + "#, + ) + .bind(&body.name) + .bind(&body.description) + .bind(body.credits) + .bind(body.price_inr) + .fetch_one(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; + + Ok(Json(package)) +} + +#[derive(Debug, Deserialize)] +struct UpdatePackageRequest { + name: Option, + description: Option, + credits: Option, + price_inr: Option, + is_active: Option, +} + +async fn admin_update_package( + auth: AuthUser, + State(state): State, + Path(id): Path, + Json(body): Json, +) -> Result, (StatusCode, String)> { + require_admin(&auth).map_err(|_| (StatusCode::FORBIDDEN, "Admin access required".to_string()))?; + + if body.credits.is_some_and(|c| c <= 0) || body.price_inr.is_some_and(|p| p <= 0) { + return Err((StatusCode::BAD_REQUEST, "credits and price_inr must be positive".to_string())); + } + + let package = sqlx::query_as::<_, AdminAiCreditPackageRow>( + r#" + UPDATE ai_credit_packages + SET name = COALESCE($1, name), + description = COALESCE($2, description), + credits = COALESCE($3, credits), + price_inr = COALESCE($4, price_inr), + is_active = COALESCE($5, is_active), + updated_at = NOW() + WHERE id = $6 + RETURNING id, name, description, credits, price_inr, is_active, created_at, updated_at + "#, + ) + .bind(&body.name) + .bind(&body.description) + .bind(body.credits) + .bind(body.price_inr) + .bind(body.is_active) + .bind(id) + .fetch_optional(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))? + .ok_or((StatusCode::NOT_FOUND, "Package not found".to_string()))?; + + Ok(Json(package)) +} diff --git a/apps/payments/src/main.rs b/apps/payments/src/main.rs index 7e5c7c8..71c625f 100644 --- a/apps/payments/src/main.rs +++ b/apps/payments/src/main.rs @@ -12,36 +12,69 @@ use uuid::Uuid; use sqlx::postgres::PgPool; use sqlx::FromRow; +pub mod ai_credits; pub mod packages; +pub mod payu; #[derive(Clone)] pub struct AppState { beeceptor_url: String, client: reqwest::Client, pool: PgPool, + payu: payu::PayuConfig, } #[derive(Debug, Serialize, Deserialize)] struct CreateOrderRequest { - amount: u64, + // Client-supplied amount/currency are accepted for backward + // compatibility but NOT trusted for pricing -- the real charge is + // always derived server-side from pricing_packages.price_inr (see + // create_order). Trusting a client-supplied amount for a real payment + // gateway would let a client pay any amount it likes for a package. + #[allow(dead_code)] + amount: Option, + #[allow(dead_code)] currency: Option, package_id: Option, - user_id: Option, } #[derive(Debug, Serialize)] struct CreateOrderResponse { + key: String, + txnid: String, + amount: String, + productinfo: String, + firstname: String, + email: String, + phone: String, + surl: String, + furl: String, + hash: String, + payu_base_url: String, + udf1: String, + udf2: String, + // Kept for callers still reading the pre-PayU response shape. order_id: String, - amount: u64, currency: String, status: String, } #[derive(Debug, Serialize, Deserialize)] struct VerifyPaymentRequest { - order_id: String, - payment_id: String, - signature: Option, + txnid: String, + mihpayid: String, + status: String, + hash: String, + amount: String, + productinfo: String, + firstname: String, + email: String, + #[serde(default)] + phone: Option, + #[serde(default)] + udf1: Option, + #[serde(default)] + udf2: Option, } #[derive(Debug, Serialize)] @@ -63,6 +96,15 @@ struct PaymentStatusResponse { #[derive(Debug, FromRow)] struct PricingPackageRow { tracecoins_amount: i32, + price_inr: i32, + name: String, +} + +#[derive(Debug, FromRow)] +struct UserContactRow { + email: String, + full_name: Option, + phone: Option, } #[derive(Debug, FromRow)] @@ -79,13 +121,11 @@ async fn create_order( State(state): State, Json(payload): Json, ) -> Result, (StatusCode, String)> { - tracing::info!("Creating payment order: amount={}", payload.amount); - let package_id_str = payload.package_id.as_ref().ok_or((StatusCode::BAD_REQUEST, "package_id is required".to_string()))?; let package_id = Uuid::parse_str(package_id_str).map_err(|_| (StatusCode::BAD_REQUEST, "Invalid package id".to_string()))?; let package = sqlx::query_as::<_, PricingPackageRow>( - "SELECT tracecoins_amount FROM pricing_packages WHERE id = $1 AND is_active = true", + "SELECT tracecoins_amount, price_inr, name FROM pricing_packages WHERE id = $1 AND is_active = true", ) .bind(package_id) .fetch_optional(&state.pool) @@ -95,61 +135,75 @@ async fn create_order( let package = package.ok_or((StatusCode::BAD_REQUEST, "Invalid or inactive package".to_string()))?; let tracecoins_credited = package.tracecoins_amount; - let resp = state - .client - .post(&state.beeceptor_url) - .header("Content-Type", "application/json") - .json(&serde_json::json!({ - "amount": payload.amount, - "currency": payload.currency.as_deref().unwrap_or("INR"), - "package_id": package_id_str, - "user_id": auth.user_id.to_string(), - })) - .send() - .await - .map_err(|e| (StatusCode::BAD_GATEWAY, format!("Beeceptor error: {}", e)))?; + tracing::info!("Creating PayU order for package {} (₹{})", package_id, payu::paise_to_rupee_string(package.price_inr)); - let status = resp.status(); - let body: serde_json::Value = resp - .json() - .await - .map_err(|e| (StatusCode::BAD_GATEWAY, format!("Parse error: {}", e)))?; + let contact = sqlx::query_as::<_, UserContactRow>( + "SELECT email, full_name, phone FROM users WHERE id = $1", + ) + .bind(auth.user_id) + .fetch_optional(&state.pool) + .await + .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))? + .ok_or((StatusCode::NOT_FOUND, "User not found".to_string()))?; - if !status.is_success() { - return Err(( - StatusCode::BAD_REQUEST, - body.get("message") - .and_then(|m| m.as_str()) - .unwrap_or("Order creation failed") - .to_string(), - )); - } - - let order_id = body - .get("order_id") - .and_then(|v| v.as_str()) - .unwrap_or("mock_order_123") + let firstname = contact + .full_name + .as_deref() + .and_then(|n| n.split_whitespace().next()) + .unwrap_or("Customer") .to_string(); + let phone = contact.phone.unwrap_or_default(); + + let txnid = payu::generate_txnid(); + // Price is derived from pricing_packages.price_inr (server-side truth), + // never from the client-supplied `payload.amount` -- see CreateOrderRequest. + let amount_str = payu::paise_to_rupee_string(package.price_inr); + let productinfo = package.name.clone(); + let udf1 = package_id_str.clone(); + let udf2 = String::new(); + + let hash = payu::request_hash( + &state.payu, + &txnid, + &amount_str, + &productinfo, + &firstname, + &contact.email, + &udf1, + &udf2, + ); sqlx::query( r#" - INSERT INTO payments (user_id, package_id, razorpay_order_id, amount, tracecoins_credited, status) + INSERT INTO payments (user_id, package_id, razorpay_order_id, amount_inr, tracecoins_credited, status) VALUES ($1, $2, $3, $4, $5, 'PENDING') "#, ) .bind(auth.user_id) .bind(package_id) - .bind(&order_id) - .bind(payload.amount as i64) + .bind(&txnid) + .bind(package.price_inr) .bind(tracecoins_credited) .execute(&state.pool) .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; Ok(Json(CreateOrderResponse { - order_id, - amount: payload.amount, - currency: payload.currency.unwrap_or("INR".to_string()), + key: state.payu.merchant_key.clone(), + txnid: txnid.clone(), + amount: amount_str, + productinfo, + firstname, + email: contact.email, + phone, + surl: state.payu.surl.clone(), + furl: state.payu.furl.clone(), + hash, + payu_base_url: state.payu.base_url.clone(), + udf1, + udf2, + order_id: txnid, + currency: "INR".to_string(), status: "created".to_string(), })) } @@ -159,32 +213,25 @@ async fn verify_payment( State(state): State, Json(payload): Json, ) -> Result, (StatusCode, String)> { - tracing::info!("Verifying payment: order_id={}", payload.order_id); + tracing::info!("Verifying PayU payment: txnid={}", payload.txnid); - let verify_url = format!("{}/verify", state.beeceptor_url.trim_end_matches('/')); - let resp = state - .client - .post(&verify_url) - .header("Content-Type", "application/json") - .json(&payload) - .send() - .await - .map_err(|e| (StatusCode::BAD_GATEWAY, format!("Beeceptor error: {}", e)))?; + if !payu::verify_response_hash( + &state.payu, + &payload.status, + &payload.txnid, + &payload.amount, + &payload.productinfo, + &payload.firstname, + &payload.email, + payload.udf1.as_deref().unwrap_or(""), + payload.udf2.as_deref().unwrap_or(""), + &payload.hash, + ) { + return Err((StatusCode::BAD_REQUEST, "Payment hash verification failed".to_string())); + } - let status = resp.status(); - let body: serde_json::Value = resp - .json() - .await - .map_err(|e| (StatusCode::BAD_GATEWAY, format!("Parse error: {}", e)))?; - - if !status.is_success() { - return Err(( - StatusCode::BAD_REQUEST, - body.get("message") - .and_then(|m| m.as_str()) - .unwrap_or("Verification failed") - .to_string(), - )); + if !payload.status.eq_ignore_ascii_case("success") { + return Err((StatusCode::BAD_REQUEST, "Payment was not successful".to_string())); } let payment = sqlx::query_as::<_, PaymentRow>( @@ -194,7 +241,7 @@ async fn verify_payment( WHERE razorpay_order_id = $1 AND status = 'PENDING' "#, ) - .bind(&payload.order_id) + .bind(&payload.txnid) .fetch_optional(&state.pool) .await .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, format!("DB error: {e}")))?; @@ -219,7 +266,7 @@ async fn verify_payment( WHERE id = $2 "#, ) - .bind(&payload.payment_id) + .bind(&payload.mihpayid) .bind(payment.id) .execute(&state.pool) .await @@ -277,7 +324,7 @@ async fn verify_payment( Ok(Json(VerifyPaymentResponse { verified: true, - payment_id: payload.payment_id, + payment_id: payload.mihpayid, status: "success".to_string(), message: "Payment verified successfully".to_string(), })) @@ -354,6 +401,7 @@ async fn main() { beeceptor_url, client: reqwest::Client::new(), pool, + payu: payu::PayuConfig::from_env(), }; let app = Router::new() @@ -361,6 +409,8 @@ async fn main() { .route("/api/payments/verify", post(verify_payment)) .route("/api/payments/{id}/status", get(get_payment_status)) .nest("/api/packages", packages::router()) + .nest("/api/ai-credits", ai_credits::router()) + .nest("/api/admin/ai-credits/packages", ai_credits::admin_router()) .with_state(state); let port: u16 = std::env::var("PORT") diff --git a/apps/payments/src/payu.rs b/apps/payments/src/payu.rs new file mode 100644 index 0000000..f1972e6 --- /dev/null +++ b/apps/payments/src/payu.rs @@ -0,0 +1,188 @@ +//! Real PayU integration, replacing the mocked Beeceptor calls in +//! `create_order`/`verify_payment`. Matches the checkout contract the +//! frontend (`nxtgauge-frontend-solid/src/lib/payu.ts`) already expects: +//! a form-redirect to `{payu_base_url}/_payment` with a request hash, and +//! a reverse-hash-verified response on callback. +//! +//! Hash formulas are PayU's standard v1 scheme (documented publicly by +//! PayU): request hash covers key/txnid/amount/productinfo/firstname/ +//! email/udf1-10/salt; response verification reverses that order and +//! substitutes `status` for the trailing empty slots. + +use sha2::{Digest, Sha512}; + +#[derive(Clone)] +pub struct PayuConfig { + pub merchant_key: String, + pub merchant_salt: String, + pub base_url: String, + pub surl: String, + pub furl: String, +} + +impl PayuConfig { + pub fn from_env() -> Self { + Self { + merchant_key: std::env::var("PAYU_MERCHANT_KEY").unwrap_or_default(), + merchant_salt: std::env::var("PAYU_MERCHANT_SALT").unwrap_or_default(), + base_url: std::env::var("PAYU_BASE_URL") + .unwrap_or_else(|_| "https://test.payu.in".to_string()), + surl: std::env::var("PAYU_SURL") + .unwrap_or_else(|_| "/api/payments/payu/callback".to_string()), + furl: std::env::var("PAYU_FURL") + .unwrap_or_else(|_| "/api/payments/payu/callback".to_string()), + } + } +} + +fn sha512_hex(input: &str) -> String { + let mut hasher = Sha512::new(); + hasher.update(input.as_bytes()); + hex::encode(hasher.finalize()) +} + +#[allow(clippy::too_many_arguments)] +pub fn request_hash( + config: &PayuConfig, + txnid: &str, + amount: &str, + productinfo: &str, + firstname: &str, + email: &str, + udf1: &str, + udf2: &str, +) -> String { + let seq = [ + config.merchant_key.as_str(), + txnid, + amount, + productinfo, + firstname, + email, + udf1, + udf2, + "", + "", + "", + "", + "", + "", + "", + "", + config.merchant_salt.as_str(), + ]; + sha512_hex(&seq.join("|")) +} + +/// Verify a PayU response hash. `status` is PayU's returned status string +/// (e.g. "success"/"failure") -- it is part of the hashed sequence, not +/// just metadata, per PayU's reverse-hash formula. +#[allow(clippy::too_many_arguments)] +pub fn verify_response_hash( + config: &PayuConfig, + status: &str, + txnid: &str, + amount: &str, + productinfo: &str, + firstname: &str, + email: &str, + udf1: &str, + udf2: &str, + received_hash: &str, +) -> bool { + let seq = [ + config.merchant_salt.as_str(), + status, + "", + "", + "", + "", + "", + "", + "", + "", + udf2, + udf1, + email, + firstname, + productinfo, + amount, + txnid, + config.merchant_key.as_str(), + ]; + let expected = sha512_hex(&seq.join("|")); + expected.eq_ignore_ascii_case(received_hash) +} + +pub fn generate_txnid() -> String { + uuid::Uuid::new_v4().to_string().replace('-', "") +} + +/// Format an integer-paise amount as the decimal-rupee string PayU expects +/// in its hash and form fields (e.g. 9900 paise -> "99.00"). +pub fn paise_to_rupee_string(paise: i32) -> String { + format!("{:.2}", paise as f64 / 100.0) +} + +#[cfg(test)] +mod tests { + use super::*; + + fn test_config() -> PayuConfig { + PayuConfig { + merchant_key: "testkey".to_string(), + merchant_salt: "testsalt".to_string(), + base_url: "https://test.payu.in".to_string(), + surl: "/s".to_string(), + furl: "/f".to_string(), + } + } + + #[test] + fn request_hash_is_deterministic_and_order_sensitive() { + let config = test_config(); + let h1 = request_hash(&config, "txn1", "99.00", "Starter AI Credits", "John", "john@test.com", "u1", "u2"); + let h2 = request_hash(&config, "txn1", "99.00", "Starter AI Credits", "John", "john@test.com", "u1", "u2"); + assert_eq!(h1, h2); + + let h3 = request_hash(&config, "txn2", "99.00", "Starter AI Credits", "John", "john@test.com", "u1", "u2"); + assert_ne!(h1, h3, "different txnid must produce a different hash"); + } + + #[test] + fn response_hash_round_trips_with_request_hash_shape() { + let config = test_config(); + // Build a response hash the way PayU would for a success callback, + // then confirm verify_response_hash accepts it. + let seq = [ + config.merchant_salt.as_str(), + "success", + "", + "", + "", + "", + "", + "", + "", + "", + "u2", + "u1", + "john@test.com", + "John", + "Starter AI Credits", + "99.00", + "txn1", + config.merchant_key.as_str(), + ]; + let hash = sha512_hex(&seq.join("|")); + + assert!(verify_response_hash( + &config, "success", "txn1", "99.00", "Starter AI Credits", "John", "john@test.com", "u1", "u2", &hash, + )); + + // Tampering with the amount must invalidate the hash. + assert!(!verify_response_hash( + &config, "success", "txn1", "999.00", "Starter AI Credits", "John", "john@test.com", "u1", "u2", &hash, + )); + } +} diff --git a/apps/users/src/ai_credits.rs b/apps/users/src/ai_credits.rs new file mode 100644 index 0000000..e8ea5a4 --- /dev/null +++ b/apps/users/src/ai_credits.rs @@ -0,0 +1,373 @@ +//! Reusable "charge an AI feature" primitive for handler code. +//! +//! Wraps `db::models::ai_credits::AiCreditsRepository`'s reserve/capture/ +//! release calls around an arbitrary unit of work (in practice, a LiteLLM +//! call): reserve credits up front, run the work, capture on success or +//! release on failure/panic-unwind-safe-error. No handler should call +//! `try_reserve_credits`/`try_capture_reservation`/`try_release_reservation` +//! directly -- always go through `charge_ai_feature` so a hold can never be +//! left unresolved by a forgotten capture/release call at a call site. + +use axum::http::StatusCode; +use db::models::ai_credits::{AiCreditsError, AiCreditsRepository, AiFeatureCost}; +use sqlx::PgPool; +use uuid::Uuid; + +/// Maximum prompt length in characters (approx 4 chars per token as per architecture doc) +/// Used for input validation before reserving credits +pub const APPROX_CHARS_PER_TOKEN: i32 = 4; + +/// Validate prompt/response lengths against feature limits before making LLM calls +/// Returns None if valid, or Some((status, error_json)) if validation fails +pub fn validate_prompt_length(prompt: &str, feature: &AiFeatureCost) -> Option<(StatusCode, serde_json::Value)> { + // Check input length if max_input_tokens is configured + if let Some(max_tokens) = feature.max_input_tokens { + let max_chars = max_tokens * APPROX_CHARS_PER_TOKEN; + if prompt.len() as i32 > max_chars { + return Some(( + StatusCode::BAD_REQUEST, + serde_json::json!({ + "error": format!( + "Prompt too long: {} characters exceeds maximum of {} (approx {} tokens)", + prompt.len(), + max_chars, + max_tokens + ), + "code": "PROMPT_TOO_LONG" + }), + )); + } + } + None +} + +/// Truncate text for preview/logging (first N chars) +pub fn truncate_for_preview(text: &str, max_chars: usize) -> String { + if text.len() <= max_chars { + text.to_string() + } else { + format!("{}...", &text[..max_chars]) + } +} + +#[derive(Debug)] +pub enum ChargeError { + Credits(AiCreditsError), + UnknownFeature(String), + Work(E), +} + +impl std::fmt::Display for ChargeError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ChargeError::Credits(e) => write!(f, "AI credits error: {e}"), + ChargeError::UnknownFeature(f2) => write!(f, "Unknown AI feature '{f2}'"), + ChargeError::Work(e) => write!(f, "{e}"), + } + } +} + +pub struct ChargeOutcome { + pub result: T, + pub feature: AiFeatureCost, + pub credits_charged: i32, +} + +/// Reserve credits for `feature_code`, run `work`, and settle the +/// reservation based on whether `work` succeeded. `work` receives the +/// resolved feature cost (so it knows which model/timeout to use) and +/// returns `Ok(T)` to capture the charge or `Err(E)` to release it +/// uncharged. +pub async fn charge_ai_feature( + pool: &PgPool, + user_id: Uuid, + feature_code: &str, + request_id: Option<&str>, + idempotency_key: Option<&str>, + work: F, +) -> Result, ChargeError> +where + F: FnOnce(AiFeatureCost) -> Fut, + Fut: std::future::Future>, +{ + let feature = AiCreditsRepository::get_feature_cost(pool, feature_code) + .await + .map_err(|e| ChargeError::Credits(AiCreditsError::Db(e)))? + .ok_or_else(|| ChargeError::UnknownFeature(feature_code.to_string()))?; + + let hold = AiCreditsRepository::try_reserve_credits( + pool, + user_id, + feature_code, + feature.credit_cost, + request_id, + idempotency_key, + ) + .await + .map_err(ChargeError::Credits)?; + + match work(feature.clone()).await { + Ok(result) => { + AiCreditsRepository::try_capture_reservation(pool, hold.id, None) + .await + .map_err(ChargeError::Credits)?; + log_usage(pool, user_id, &feature, hold.credits_held, "success", request_id, None, None, None, None, None).await; + Ok(ChargeOutcome { + result, + feature, + credits_charged: hold.credits_held, + }) + } + Err(e) => { + // Best-effort release -- if this fails, the hold's expires_at + // (Section 4.2 of the architecture doc) is the backstop until + // a reaper job exists to sweep abandoned holds. + let _ = AiCreditsRepository::try_release_reservation(pool, hold.id).await; + log_usage(pool, user_id, &feature, 0, "error", request_id, None, None, None, None, None).await; + Err(ChargeError::Work(e)) + } + } +} + +/// Best-effort usage log write -- failures here must never fail the +/// caller's request (the credit charge/release already happened; a +/// missing analytics row is not worth surfacing as an error to the user). +#[allow(clippy::too_many_arguments)] +async fn log_usage( + pool: &PgPool, + user_id: Uuid, + feature: &AiFeatureCost, + credits_charged: i32, + status: &str, + request_id: Option<&str>, + prompt_preview: Option<&str>, + response_preview: Option<&str>, + input_tokens: Option, + output_tokens: Option, + total_tokens: Option, +) { + let result = sqlx::query( + r#" + INSERT INTO ai_usage_logs + (user_id, feature_code, model_alias, credits_charged, status, request_id, + prompt_preview, response_preview, input_tokens, output_tokens, total_tokens) + VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) + "#, + ) + .bind(user_id) + .bind(&feature.feature_code) + .bind(&feature.default_model) + .bind(credits_charged) + .bind(status) + .bind(request_id) + .bind(prompt_preview) + .bind(response_preview) + .bind(input_tokens) + .bind(output_tokens) + .bind(total_tokens) + .execute(pool) + .await; + + if let Err(e) = result { + tracing::error!("Failed to write ai_usage_logs row: {e}"); + } +} + +/// Charge an AI feature using LiteLLM with full audit logging and token capture +/// +/// This is the preferred way to charge AI features - it uses LiteLLM for model routing +/// and captures token counts for cost tracking. +/// +/// # Arguments +/// * `pool` - Database connection pool +/// * `user_id` - The user being charged +/// * `feature_code` - The AI feature being used (e.g., "jd_generate", "cover_letter_generate") +/// * `request_id` - Optional request tracking ID +/// * `idempotency_key` - Optional idempotency key +/// * `prompt` - The prompt being sent to the LLM +/// * `litellm_response` - The LiteLLM response containing generated text and token counts +/// +/// # Returns +/// * `Ok(ChargeOutcome)` - The generated text from the LLM +/// * `Err(ChargeError)` - Error if credits insufficient or DB failure +pub async fn charge_ai_feature_litellm( + pool: &PgPool, + user_id: Uuid, + feature_code: &str, + request_id: Option<&str>, + idempotency_key: Option<&str>, + prompt: &str, + litellm_response: &crate::litellm::LiteLLMResponse, +) -> Result, ChargeError> { + let feature = AiCreditsRepository::get_feature_cost(pool, feature_code) + .await + .map_err(|e| ChargeError::Credits(AiCreditsError::Db(e)))? + .ok_or_else(|| ChargeError::UnknownFeature(feature_code.to_string()))?; + + // Validate prompt length before reserving + if let Some((_, error_json)) = validate_prompt_length(prompt, &feature) { + return Err(ChargeError::Work(format!( + "Validation failed: {}", + error_json["error"] + ))); + } + + let prompt_preview = truncate_for_preview(prompt, 200); + let response_text = litellm_response.generated_text(); + let response_preview = truncate_for_preview(&response_text, 200); + + let hold = AiCreditsRepository::try_reserve_credits( + pool, + user_id, + feature_code, + feature.credit_cost, + request_id, + idempotency_key, + ) + .await + .map_err(ChargeError::Credits)?; + + // Capture the charge since we have a successful LiteLLM response + AiCreditsRepository::try_capture_reservation(pool, hold.id, None) + .await + .map_err(ChargeError::Credits)?; + + // Log with token counts if available + let input_tokens = litellm_response.usage.as_ref().and_then(|u| u.prompt_tokens); + let output_tokens = litellm_response.usage.as_ref().and_then(|u| u.completion_tokens); + let total_tokens = litellm_response.usage.as_ref().and_then(|u| u.total_tokens); + + log_usage( + pool, + user_id, + &feature, + hold.credits_held, + "success", + request_id, + Some(&prompt_preview), + Some(&response_preview), + input_tokens, + output_tokens, + total_tokens, + ) + .await; + + Ok(ChargeOutcome { + result: response_text, + feature, + credits_charged: hold.credits_held, + }) +} + +/// Convenience function to call LiteLLM and charge credits in one operation +/// +/// This is the recommended pattern for handlers: +/// ```rust +/// let response = call_litellm_and_charge( +/// &state.pool, +/// auth.user_id, +/// "jd_generate", +/// None, +/// None, +/// &prompt, +/// ).await?; +/// ``` +pub async fn call_litellm_and_charge( + pool: &PgPool, + user_id: Uuid, + feature_code: &str, + request_id: Option<&str>, + idempotency_key: Option<&str>, + prompt: &str, +) -> Result, ChargeError> { + // First reserve credits + let feature = AiCreditsRepository::get_feature_cost(pool, feature_code) + .await + .map_err(|e| ChargeError::Credits(AiCreditsError::Db(e)))? + .ok_or_else(|| ChargeError::UnknownFeature(feature_code.to_string()))?; + + // Validate prompt length before reserving + if let Some((_, error_json)) = validate_prompt_length(prompt, &feature) { + return Err(ChargeError::Work(format!( + "Validation failed: {}", + error_json["error"] + ))); + } + + let prompt_preview = truncate_for_preview(prompt, 200); + + let hold = AiCreditsRepository::try_reserve_credits( + pool, + user_id, + feature_code, + feature.credit_cost, + request_id, + idempotency_key, + ) + .await + .map_err(ChargeError::Credits)?; + + // Get LiteLLM config and call + let (base_url, default_model, api_key) = crate::litellm::get_litellm_config(); + let model_alias = feature.default_model.clone(); + + // Use feature's max_output_tokens if configured + let max_tokens = feature.max_output_tokens; + + match crate::litellm::call_litellm(&base_url, &model_alias, prompt, api_key.as_deref(), max_tokens).await { + Ok(litellm_response) => { + // Capture the charge + AiCreditsRepository::try_capture_reservation(pool, hold.id, None) + .await + .map_err(ChargeError::Credits)?; + + let response_text = litellm_response.generated_text(); + let response_preview = truncate_for_preview(&response_text, 200); + + // Log with token counts + let input_tokens = litellm_response.usage.as_ref().and_then(|u| u.prompt_tokens); + let output_tokens = litellm_response.usage.as_ref().and_then(|u| u.completion_tokens); + let total_tokens = litellm_response.usage.as_ref().and_then(|u| u.total_tokens); + + log_usage( + pool, + user_id, + &feature, + hold.credits_held, + "success", + request_id, + Some(&prompt_preview), + Some(&response_preview), + input_tokens, + output_tokens, + total_tokens, + ) + .await; + + Ok(ChargeOutcome { + result: response_text, + feature, + credits_charged: hold.credits_held, + }) + } + Err(e) => { + // Release the hold on error + let _ = AiCreditsRepository::try_release_reservation(pool, hold.id).await; + log_usage( + pool, + user_id, + &feature, + 0, + "error", + request_id, + Some(&prompt_preview), + None, + None, + None, + None, + ) + .await; + Err(ChargeError::Work(e)) + } + } +} diff --git a/apps/users/src/ai_subscription.rs b/apps/users/src/ai_subscription.rs new file mode 100644 index 0000000..e0be2f8 --- /dev/null +++ b/apps/users/src/ai_subscription.rs @@ -0,0 +1,484 @@ +//! Subscription lifecycle management for AI credits +//! +//! Handles plan upgrades, downgrades, cancellations, and renewals + +use db::models::ai_credits::{AiCreditsError, AiCreditsRepository, AiPlan}; +use sqlx::PgPool; +use uuid::Uuid; +use chrono::{DateTime, Utc, Duration}; + +/// Calculate prorated credits for a mid-period upgrade +/// +/// Formula: remaining_days / total_days * (new_monthly - old_monthly) +fn calculate_proration( + current_period_end: DateTime, + old_monthly_credits: i32, + new_monthly_credits: i32, +) -> i32 { + if new_monthly_credits <= old_monthly_credits { + return 0; + } + + let now = Utc::now(); + let total_days = (current_period_end - now).num_days(); + let period_length = 30; // Monthly period assumption + + if total_days <= 0 { + return 0; + } + + let remaining_ratio = total_days as f64 / period_length as f64; + let credit_diff = (new_monthly_credits - old_monthly_credits) as f64; + + (remaining_ratio * credit_diff).ceil() as i32 +} + +/// Upgrade a user's plan immediately +/// +/// This applies the new plan immediately and grants prorated bonus credits. +pub async fn upgrade_plan( + pool: &PgPool, + user_id: Uuid, + new_plan_id: Uuid, + actor_id: Option, +) -> Result<(), AiCreditsError> { + let mut tx = pool.begin().await?; + + // Get current wallet + let wallet = sqlx::query_as::<_, db::models::ai_credits::AiWallet>( + "SELECT * FROM user_ai_subscriptions WHERE user_id = $1 FOR UPDATE" + ) + .bind(user_id) + .fetch_one(&mut *tx) + .await?; + + // Get new plan details + let new_plan: Option = sqlx::query_as::<_, AiPlan>( + "SELECT id, code, name, monthly_credits, daily_action_limit, daily_credit_limit, allowed_models, allowed_features, is_active FROM ai_plans WHERE id = $1 AND is_active = TRUE" + ) + .bind(new_plan_id) + .fetch_optional(&mut *tx) + .await?; + + let Some(new_plan) = new_plan else { + tx.rollback().await?; + return Err(AiCreditsError::UnknownFeature("Plan not found".to_string())); + }; + + // Get old plan + let old_plan: Option = sqlx::query_as::<_, AiPlan>( + "SELECT id, code, name, monthly_credits, daily_action_limit, daily_credit_limit, allowed_models, allowed_features, is_active FROM ai_plans WHERE id = $1" + ) + .bind(wallet.plan_id) + .fetch_optional(&mut *tx) + .await?; + + // Calculate proration + let proration_credits = if let Some(ref old_plan) = old_plan { + calculate_proration( + wallet.current_period_end, + old_plan.monthly_credits, + new_plan.monthly_credits, + ) + } else { + 0 + }; + + // Update plan + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET plan_id = $1, + monthly_credits_total = $2, + bonus_credits_total = bonus_credits_total + $3, + downgrade_scheduled_to = NULL, + updated_at = NOW() + WHERE id = $4 + "# + ) + .bind(new_plan_id) + .bind(new_plan.monthly_credits) + .bind(proration_credits) + .bind(wallet.id) + .execute(&mut *tx) + .await?; + + // Record in history + sqlx::query( + r#" + INSERT INTO ai_subscription_history + (user_id, from_plan_id, to_plan_id, change_type, proration_credits, + proration_days_remaining, effective_at, created_by, status) + VALUES ($1, $2, $3, 'upgrade', $4, $5, NOW(), $6, 'completed') + "# + ) + .bind(user_id) + .bind(wallet.plan_id) + .bind(new_plan_id) + .bind(proration_credits) + .bind((wallet.current_period_end - Utc::now()).num_days() as i32) + .bind(actor_id) + .execute(&mut *tx) + .await?; + + // Create ledger entry for proration bonus + if proration_credits > 0 { + let balance_after = wallet.available_credits() + proration_credits; + sqlx::query( + r#" + INSERT INTO ai_credit_ledger + (wallet_id, entry_type, credits, balance_after, reference_type, actor_type, metadata) + VALUES ($1, 'subscription_grant', $2, $3, 'plan_upgrade', 'system', + jsonb_build_object('reason', 'upgrade_proration', 'plan_code', $4)) + "# + ) + .bind(wallet.id) + .bind(proration_credits) + .bind(balance_after) + .bind(&new_plan.code) + .execute(&mut *tx) + .await?; + } + + tx.commit().await?; + Ok(()) +} + +/// Schedule a downgrade for the end of the current period +/// +/// The downgrade doesn't happen immediately - it will be applied +/// when the current period ends. +pub async fn schedule_downgrade( + pool: &PgPool, + user_id: Uuid, + downgrade_to_plan_id: Uuid, + actor_id: Option, +) -> Result, AiCreditsError> { + let wallet = AiCreditsRepository::get_wallet(pool, user_id) + .await? + .ok_or(AiCreditsError::ReservationNotFound)?; + + // Get effective date (current period end) + let effective_at = wallet.current_period_end; + + // Update wallet to mark scheduled downgrade + sqlx::query( + "UPDATE user_ai_subscriptions SET downgrade_scheduled_to = $1, updated_at = NOW() WHERE id = $2" + ) + .bind(downgrade_to_plan_id) + .bind(wallet.id) + .execute(pool) + .await?; + + // Record scheduled downgrade + sqlx::query( + r#" + INSERT INTO ai_subscription_history + (user_id, from_plan_id, to_plan_id, change_type, effective_at, created_by, status) + VALUES ($1, $2, $3, 'downgrade', $4, $5, 'scheduled') + "# + ) + .bind(user_id) + .bind(wallet.plan_id) + .bind(downgrade_to_plan_id) + .bind(effective_at) + .bind(actor_id) + .execute(pool) + .await?; + + Ok(effective_at) +} + +/// Apply scheduled downgrades that have reached their effective date +/// +/// This should be called periodically (e.g., by a cron job) to apply +/// downgrades that were scheduled for the current period end. +pub async fn apply_scheduled_downgrades(pool: &PgPool) -> Result { + let now = Utc::now(); + + // Find all wallets with scheduled downgrades that should be applied + let to_downgrade: Vec<(Uuid, Uuid, Uuid)> = sqlx::query_as( + r#" + SELECT w.id as wallet_id, w.user_id, w.downgrade_scheduled_to as new_plan_id + FROM user_ai_subscriptions w + WHERE w.downgrade_scheduled_to IS NOT NULL + AND w.current_period_end <= $1 + "# + ) + .bind(now) + .fetch_all(pool) + .await + .map_err(AiCreditsError::Db)?; + + let mut applied_count = 0; + + for (wallet_id, user_id, new_plan_id) in to_downgrade { + let mut tx = pool.begin().await?; + + // Get new plan credits + let new_monthly_credits: i32 = sqlx::query_scalar( + "SELECT monthly_credits FROM ai_plans WHERE id = $1" + ) + .bind(new_plan_id) + .fetch_one(&mut *tx) + .await?; + + // Update wallet + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET plan_id = downgrade_scheduled_to, + monthly_credits_total = $1, + monthly_credits_used = 0, -- Reset usage for new period + downgrade_scheduled_to = NULL, + current_period_start = NOW(), + current_period_end = NOW() + INTERVAL '30 days', + updated_at = NOW() + WHERE id = $2 + "# + ) + .bind(new_monthly_credits) + .bind(wallet_id) + .execute(&mut *tx) + .await?; + + // Update history record + sqlx::query( + r#" + UPDATE ai_subscription_history + SET status = 'completed' + WHERE user_id = $1 AND change_type = 'downgrade' AND status = 'scheduled' + "# + ) + .bind(user_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + applied_count += 1; + } + + Ok(applied_count) +} + +/// Cancel a subscription (move to free plan) +pub async fn cancel_subscription( + pool: &PgPool, + user_id: Uuid, + actor_id: Option, +) -> Result<(), AiCreditsError> { + let wallet = AiCreditsRepository::get_wallet(pool, user_id) + .await? + .ok_or(AiCreditsError::ReservationNotFound)?; + + // Get free plan + let free_plan: Option<(Uuid, i32)> = sqlx::query_as( + "SELECT id, monthly_credits FROM ai_plans WHERE code = 'free' AND is_active = TRUE" + ) + .fetch_optional(pool) + .await?; + + let Some((free_plan_id, free_monthly)) = free_plan else { + return Err(AiCreditsError::UnknownFeature("Free plan not found".to_string())); + }; + + // Update to free plan immediately + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET plan_id = $1, + monthly_credits_total = $2, + monthly_credits_used = 0, + downgrade_scheduled_to = NULL, + status = 'cancelled', + updated_at = NOW() + WHERE user_id = $3 + "# + ) + .bind(free_plan_id) + .bind(free_monthly) + .bind(user_id) + .execute(pool) + .await?; + + // Record cancellation + sqlx::query( + r#" + INSERT INTO ai_subscription_history + (user_id, from_plan_id, to_plan_id, change_type, effective_at, created_by, status) + VALUES ($1, $2, $3, 'cancel', NOW(), $4, 'completed') + "# + ) + .bind(user_id) + .bind(wallet.plan_id) + .bind(free_plan_id) + .bind(actor_id) + .execute(pool) + .await?; + + Ok(()) +} + +/// Start a trial for a new user +/// +/// Sets up a trial period with the specified plan. +pub async fn start_trial( + pool: &PgPool, + user_id: Uuid, + trial_plan_id: Uuid, + trial_days: i32, +) -> Result, AiCreditsError> { + let trial_ends_at = Utc::now() + Duration::days(trial_days as i64); + + let plan: Option<(i32,)> = sqlx::query_as( + "SELECT monthly_credits FROM ai_plans WHERE id = $1" + ) + .bind(trial_plan_id) + .fetch_optional(pool) + .await?; + + let Some((monthly_credits,)) = plan else { + return Err(AiCreditsError::UnknownFeature("Trial plan not found".to_string())); + }; + + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET plan_id = $1, + monthly_credits_total = $2, + is_trial = TRUE, + trial_days = $3, + trial_ends_at = $4, + current_period_end = $4, + updated_at = NOW() + WHERE user_id = $5 + "# + ) + .bind(trial_plan_id) + .bind(monthly_credits) + .bind(trial_days) + .bind(trial_ends_at) + .bind(user_id) + .execute(pool) + .await?; + + Ok(trial_ends_at) +} + +/// Check and expire trials that have ended +/// +/// Moves users whose trials have expired to the free plan. +/// Should be called periodically by a cron job. +pub async fn expire_trials(pool: &PgPool) -> Result { + let now = Utc::now(); + + // Get free plan + let free_plan: Option<(Uuid, i32)> = sqlx::query_as( + "SELECT id, monthly_credits FROM ai_plans WHERE code = 'free' AND is_active = TRUE" + ) + .fetch_optional(pool) + .await?; + + let Some((free_plan_id, free_monthly)) = free_plan else { + return Err(AiCreditsError::UnknownFeature("Free plan not found".to_string())); + }; + + // Find expired trials + let expired: Vec<(Uuid, Uuid)> = sqlx::query_as( + r#" + SELECT user_id, plan_id + FROM user_ai_subscriptions + WHERE is_trial = TRUE AND trial_ends_at <= $1 + "# + ) + .bind(now) + .fetch_all(pool) + .await + .map_err(AiCreditsError::Db)?; + + let mut expired_count = 0; + + for (user_id, old_plan_id) in expired { + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET plan_id = $1, + monthly_credits_total = $2, + is_trial = FALSE, + trial_days = NULL, + trial_ends_at = NULL, + current_period_end = NOW() + INTERVAL '30 days', + updated_at = NOW() + WHERE user_id = $3 + "# + ) + .bind(free_plan_id) + .bind(free_monthly) + .bind(user_id) + .execute(pool) + .await?; + + // Record trial expiration + sqlx::query( + r#" + INSERT INTO ai_subscription_history + (user_id, from_plan_id, to_plan_id, change_type, effective_at, status) + VALUES ($1, $2, $3, 'trial_expired', NOW(), 'completed') + "# + ) + .bind(user_id) + .bind(old_plan_id) + .bind(free_plan_id) + .execute(pool) + .await?; + + expired_count += 1; + } + + Ok(expired_count) +} + +/// Get subscription history for a user +pub async fn get_subscription_history( + pool: &PgPool, + user_id: Uuid, + limit: i64, +) -> Result, sqlx::Error> { + sqlx::query_as( + r#" + SELECT + h.id, + h.user_id, + fp.name as from_plan_name, + tp.name as to_plan_name, + h.change_type, + h.proration_credits, + h.effective_at, + h.status, + h.created_at + FROM ai_subscription_history h + LEFT JOIN ai_plans fp ON h.from_plan_id = fp.id + JOIN ai_plans tp ON h.to_plan_id = tp.id + WHERE h.user_id = $1 + ORDER BY h.created_at DESC + LIMIT $2 + "# + ) + .bind(user_id) + .bind(limit) + .fetch_all(pool) + .await +} + +#[derive(Debug, sqlx::FromRow)] +pub struct SubscriptionHistoryRow { + pub id: Uuid, + pub user_id: Uuid, + pub from_plan_name: Option, + pub to_plan_name: String, + pub change_type: String, + pub proration_credits: i32, + pub effective_at: DateTime, + pub status: String, + pub created_at: DateTime, +} diff --git a/apps/users/src/handlers/ai.rs b/apps/users/src/handlers/ai.rs index 7e859b4..c88beeb 100644 --- a/apps/users/src/handlers/ai.rs +++ b/apps/users/src/handlers/ai.rs @@ -49,34 +49,34 @@ struct OllamaGenerateResponse { response: String, } -async fn call_ollama(_state: &AppState, model: &str, prompt: &str) -> Result { - let base_url = std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://ollama.nxtgauge-ai.svc.cluster.local:11434".to_string()); - let url = format!("{}/api/generate", base_url); - - let req = OllamaGenerateRequest { - model: model.to_string(), - prompt: prompt.to_string(), - stream: false, +async fn call_ollama( + _state: &AppState, + model: &str, + prompt: &str, + user_id: Option<&str>, +) -> Result { + let base_url = std::env::var("OLLAMA_BASE_URL") + .unwrap_or_else(|_| "http://ollama.nxtgauge-ai.svc.cluster.local:11434".to_string()); + + // Build secure client with rate limiting + let security = cache::ollama::OllamaSecurityConfig { + api_key: std::env::var("OLLAMA_API_KEY").ok(), + user_id: user_id.map(|s| s.to_string()), + client_ip: None, + rate_limit_per_minute: std::env::var("OLLAMA_RATE_LIMIT_PER_MINUTE") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(60), }; - - let client = reqwest::Client::new(); - let response = client - .post(&url) - .json(&req) - .send() - .await - .map_err(|e| format!("ollama request failed: {}", e))?; - - if !response.status().is_success() { - return Err(format!("ollama returned status: {}", response.status())); + + let client = cache::ollama::OllamaClient::with_url(&base_url) + .with_model(model) + .with_security(security); + + match client.generate(prompt).await { + Ok(response) => Ok(response.response), + Err(e) => Err(format!("Ollama error: {}", e)), } - - let result: OllamaGenerateResponse = response - .json() - .await - .map_err(|e| format!("failed to parse ollama response: {}", e))?; - - Ok(result.response) } // ── Phase 1: Strict keyword fast-path for intent classification ──────────────── @@ -279,7 +279,7 @@ fn llm_guard_check(message: &str) -> Option<(StatusCode, serde_json::Value)> { None } -async fn classify_intent(message: &str, ollama_base: &str, model: &str) -> (String, f32) { +async fn classify_intent(message: &str, ollama_base: &str, model: &str, user_id: Option<&str>) -> (String, f32) { let prompt = format!( "Classify this user message into one intent category. Categories: \ ticket_creation, form_filling, help_search, job_description_generation, \ @@ -288,7 +288,7 @@ async fn classify_intent(message: &str, ollama_base: &str, model: &str) -> (Stri message ); - match call_ollama_inline(ollama_base, model, &prompt).await { + match call_ollama_inline(ollama_base, model, &prompt, user_id).await { Ok(response) => { let intent = response.trim().to_lowercase(); let confidence = if intent.is_empty() { 0.5 } else { 0.85 }; @@ -318,32 +318,31 @@ fn is_internal_admin(auth: &AuthUser) -> bool { || auth.claims.roles.contains(&"SUPER_ADMIN".to_string()) } -async fn call_ollama_inline(base_url: &str, model: &str, prompt: &str) -> Result { - let url = format!("{}/api/generate", base_url); - let req = OllamaGenerateRequest { - model: model.to_string(), - prompt: prompt.to_string(), - stream: false, +async fn call_ollama_inline( + base_url: &str, + model: &str, + prompt: &str, + user_id: Option<&str>, +) -> Result { + // Build secure client with rate limiting + let security = cache::ollama::OllamaSecurityConfig { + api_key: std::env::var("OLLAMA_API_KEY").ok(), + user_id: user_id.map(|s| s.to_string()), + client_ip: None, + rate_limit_per_minute: std::env::var("OLLAMA_RATE_LIMIT_PER_MINUTE") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(60), }; - let client = reqwest::Client::new(); - let response = client - .post(&url) - .json(&req) - .send() - .await - .map_err(|e| format!("ollama request failed: {}", e))?; + let client = cache::ollama::OllamaClient::with_url(base_url) + .with_model(model) + .with_security(security); - if !response.status().is_success() { - return Err(format!("ollama returned status: {}", response.status())); + match client.generate(prompt).await { + Ok(response) => Ok(response.response), + Err(e) => Err(format!("Ollama error: {}", e)), } - - let result: OllamaGenerateResponse = response - .json() - .await - .map_err(|e| format!("failed to parse ollama response: {}", e))?; - - Ok(result.response) } async fn ai_chat_message( @@ -364,7 +363,7 @@ async fn ai_chat_message( // ── Phase 1: Strict keyword fast-path (skips Ollama when unambiguous) ───── let (intent, confidence) = match classify_strict_keywords(&body.message) { Some((kw_intent, kw_conf)) => (kw_intent.to_string(), kw_conf), - None => classify_intent(&body.message, &ollama_base, &model).await, + None => classify_intent(&body.message, &ollama_base, &model, body.user_id.as_deref()).await, }; let response_text = match intent.as_str() { @@ -425,7 +424,7 @@ async fn ai_chat_message( Job Description:", body.message ); - match call_ollama(&state, &model, &jd_prompt).await { + match call_ollama(&state, &model, &jd_prompt, body.user_id.as_deref()).await { Ok(r) => r, Err(e) => { tracing::error!("Ollama JD generation error: {}", e); @@ -438,7 +437,7 @@ async fn ai_chat_message( Ask for: subject, description of issue, category, priority if not provided. \ Summarize the ticket in a structured way."; let full_prompt = format!("{}\n\nUser: {}\nAssistant:", system_prompt, body.message); - match call_ollama(&state, &model, &full_prompt).await { + match call_ollama(&state, &model, &full_prompt, Some(&auth.user_id.to_string())).await { Ok(r) => r, Err(e) => { tracing::error!("Ollama error: {}", e); @@ -450,7 +449,7 @@ async fn ai_chat_message( let system_prompt = "You are a form filling assistant. Help users fill out forms by extracting relevant information \ from their message. Extract key:value pairs when possible."; let full_prompt = format!("{}\n\nUser: {}\nAssistant:", system_prompt, body.message); - match call_ollama(&state, &model, &full_prompt).await { + match call_ollama(&state, &model, &full_prompt, Some(&auth.user_id.to_string())).await { Ok(r) => r, Err(e) => { tracing::error!("Ollama error: {}", e); @@ -473,7 +472,7 @@ async fn ai_chat_message( let system_prompt = "You are a helpful AI assistant for Nxtgauge platform. Provide clear, concise responses. \ If the user needs support, guide them to create a ticket."; let full_prompt = format!("{}\n\nUser: {}\nAssistant:", system_prompt, body.message); - match call_ollama(&state, &model, &full_prompt).await { + match call_ollama(&state, &model, &full_prompt, Some(&auth.user_id.to_string())).await { Ok(r) => r, Err(e) => { tracing::error!("Ollama error: {}", e); @@ -618,7 +617,7 @@ async fn ai_extract_form( form_type, body.message ); - let response_text = match call_ollama_inline(&ollama_base, &model, &prompt).await { + let response_text = match call_ollama_inline(&ollama_base, &model, &prompt, body.user_id.as_deref()).await { Ok(r) => r, Err(e) => { tracing::error!("Ollama form extraction error: {}", e); @@ -806,34 +805,10 @@ async fn ai_generate_job_field( .ok() .flatten(); - let Some(company_id) = company else { + if company.is_none() { return (StatusCode::FORBIDDEN, Json(serde_json::json!({ "error": "No company profile found" }))).into_response(); }; - let (has_pack, daily_limit) = { - let profile_id: Option = sqlx::query_scalar( - "SELECT id FROM user_role_profiles WHERE user_id = $1 AND role_key = 'COMPANY'" - ) - .bind(auth.user_id) - .fetch_optional(&state.pool) - .await - .ok() - .flatten(); - - match profile_id { - Some(pid) => has_active_ai_pack(&state.pool, pid, "COMPANY").await, - None => (false, BASE_AI_LIMIT), - } - }; - - let mut redis = state.redis.clone(); - let (used, limit) = match check_and_increment_usage(&state.pool, &mut redis, company_id, true, daily_limit).await { - Ok((u, l)) => (u, l), - Err(msg) => { - return (StatusCode::TOO_MANY_REQUESTS, Json(serde_json::json!({ "error": msg }))).into_response(); - } - }; - let ollama_base = std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://ollama.nxtgauge-ai.svc.cluster.local:11434".to_string()); let model = std::env::var("OLLAMA_CHAT_MODEL").unwrap_or_else(|_| "gemma3:270m".to_string()); @@ -861,23 +836,67 @@ async fn ai_generate_job_field( } }; - let generated = match call_ollama_inline(&ollama_base, &model, &field_prompt).await { - Ok(r) => r.trim().to_string(), - Err(e) => { - tracing::error!("Ollama job field generation error: {}", e); - return (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Generation failed" }))).into_response(); - } - }; + // Charged against the Ask Ash AI credits wallet using LiteLLM + // for model routing (Task 4 - LiteLLM Integration) + let outcome = crate::ai_credits::call_litellm_and_charge( + &state.pool, + auth.user_id, + "jd_generate", + None, + None, + &field_prompt, + ) + .await; - ( - StatusCode::OK, - Json(GenerateFieldResponse { - generated_text: generated, - remaining_today: limit - used, - daily_limit: limit, - has_ai_pack: has_pack, - }), - ).into_response() + match outcome { + Ok(outcome) => { + let wallet = db::models::ai_credits::AiCreditsRepository::get_wallet(&state.pool, auth.user_id) + .await + .ok() + .flatten(); + let plan = match &wallet { + Some(w) => db::models::ai_credits::AiCreditsRepository::get_plan(&state.pool, w.plan_id).await.ok().flatten(), + None => None, + }; + ( + StatusCode::OK, + Json(GenerateFieldResponse { + generated_text: outcome.result.trim().to_string(), + remaining_today: wallet.as_ref().map(|w| w.available_credits()).unwrap_or(0), + daily_limit: plan.map(|p| p.daily_action_limit).unwrap_or(0), + has_ai_pack: wallet.as_ref().map(|w| w.purchased_credits_total > 0).unwrap_or(false), + }), + ) + .into_response() + } + Err(crate::ai_credits::ChargeError::Credits( + db::models::ai_credits::AiCreditsError::InsufficientCredits, + )) => ( + StatusCode::PAYMENT_REQUIRED, + Json(serde_json::json!({ "error": "Insufficient AI credits. Buy more from the Credits page." })), + ) + .into_response(), + Err(crate::ai_credits::ChargeError::Credits( + db::models::ai_credits::AiCreditsError::DailyActionLimitReached + | db::models::ai_credits::AiCreditsError::DailyCreditLimitReached, + )) => ( + StatusCode::TOO_MANY_REQUESTS, + Json(serde_json::json!({ "error": "Daily AI generation limit reached. Upgrade to AI Pack for more." })), + ) + .into_response(), + Err(crate::ai_credits::ChargeError::Credits(e)) => { + tracing::error!("AI credits error in ai_generate_job_field: {e}"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Internal error" }))).into_response() + } + Err(crate::ai_credits::ChargeError::UnknownFeature(f)) => { + tracing::error!("Unknown AI feature code '{f}' in ai_generate_job_field"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Internal error" }))).into_response() + } + Err(crate::ai_credits::ChargeError::Work(e)) => { + tracing::error!("Ollama job field generation error: {e}"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Generation failed" }))).into_response() + } + } } // ── Cover Letter Generation (Job Seekers) ────────────────────────────────────── @@ -925,32 +944,7 @@ async fn ai_generate_cover_letter( return (StatusCode::NOT_FOUND, Json(serde_json::json!({ "error": "Job not found" }))).into_response(); }; - let (has_pack, daily_limit) = { - let profile_id: Option = sqlx::query_scalar( - "SELECT id FROM user_role_profiles WHERE user_id = $1 AND role_key = 'JOB_SEEKER'" - ) - .bind(auth.user_id) - .fetch_optional(&state.pool) - .await - .ok() - .flatten(); - - match profile_id { - Some(pid) => has_active_ai_pack(&state.pool, pid, "JOB_SEEKER").await, - None => (false, BASE_AI_LIMIT), - } - }; - - let mut redis = state.redis.clone(); - let (used, limit) = match check_and_increment_usage(&state.pool, &mut redis, seeker_id, false, daily_limit).await { - Ok((u, l)) => (u, l), - Err(msg) => { - return (StatusCode::TOO_MANY_REQUESTS, Json(serde_json::json!({ "error": msg }))).into_response(); - } - }; - - let ollama_base = std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://ollama.nxtgauge-ai.svc.cluster.local:11434".to_string()); - let model = std::env::var("OLLAMA_CHAT_MODEL").unwrap_or_else(|_| "gemma3:270m".to_string()); + let _ = seeker_id; // no longer used for a per-profile daily counter -- see charge_ai_feature below let notes = body.additional_notes.as_deref().unwrap_or(""); let skills_str = skills.join(", "); @@ -974,23 +968,67 @@ async fn ai_generate_cover_letter( full_name, experience, summary.as_deref().unwrap_or("N/A"), skills_str, notes, job_title, job_desc, location ); - let generated = match call_ollama_inline(&ollama_base, &model, &prompt).await { - Ok(r) => r.trim().to_string(), - Err(e) => { - tracing::error!("Ollama cover letter generation error: {}", e); - return (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Generation failed" }))).into_response(); - } - }; + // Charged against the Ask Ash AI credits wallet using LiteLLM + // for model routing (Task 4 - LiteLLM Integration) + let outcome = crate::ai_credits::call_litellm_and_charge( + &state.pool, + auth.user_id, + "cover_letter_generate", + None, + None, + &prompt, + ) + .await; - ( - StatusCode::OK, - Json(GenerateFieldResponse { - generated_text: generated, - remaining_today: limit - used, - daily_limit: limit, - has_ai_pack: has_pack, - }), - ).into_response() + match outcome { + Ok(outcome) => { + let wallet = db::models::ai_credits::AiCreditsRepository::get_wallet(&state.pool, auth.user_id) + .await + .ok() + .flatten(); + let plan = match &wallet { + Some(w) => db::models::ai_credits::AiCreditsRepository::get_plan(&state.pool, w.plan_id).await.ok().flatten(), + None => None, + }; + ( + StatusCode::OK, + Json(GenerateFieldResponse { + generated_text: outcome.result.trim().to_string(), + remaining_today: wallet.as_ref().map(|w| w.available_credits()).unwrap_or(0), + daily_limit: plan.map(|p| p.daily_action_limit).unwrap_or(0), + has_ai_pack: wallet.as_ref().map(|w| w.purchased_credits_total > 0).unwrap_or(false), + }), + ) + .into_response() + } + Err(crate::ai_credits::ChargeError::Credits( + db::models::ai_credits::AiCreditsError::InsufficientCredits, + )) => ( + StatusCode::PAYMENT_REQUIRED, + Json(serde_json::json!({ "error": "Insufficient AI credits. Buy more from the Credits page." })), + ) + .into_response(), + Err(crate::ai_credits::ChargeError::Credits( + db::models::ai_credits::AiCreditsError::DailyActionLimitReached + | db::models::ai_credits::AiCreditsError::DailyCreditLimitReached, + )) => ( + StatusCode::TOO_MANY_REQUESTS, + Json(serde_json::json!({ "error": "Daily AI generation limit reached. Upgrade to AI Pack for more." })), + ) + .into_response(), + Err(crate::ai_credits::ChargeError::Credits(e)) => { + tracing::error!("AI credits error in ai_generate_cover_letter: {e}"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Internal error" }))).into_response() + } + Err(crate::ai_credits::ChargeError::UnknownFeature(f)) => { + tracing::error!("Unknown AI feature code '{f}' in ai_generate_cover_letter"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Internal error" }))).into_response() + } + Err(crate::ai_credits::ChargeError::Work(e)) => { + tracing::error!("Ollama cover letter generation error: {e}"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Generation failed" }))).into_response() + } + } } // ── Tailor Resume (Job Seekers) ───────────────────────────────────────────────── @@ -1020,7 +1058,7 @@ async fn ai_tailor_resume( .ok() .and_then(|r| r); - let Some((seeker_id, full_name, summary, experience, skills)) = seeker else { + let Some((_seeker_id, full_name, summary, experience, skills)) = seeker else { return (StatusCode::FORBIDDEN, Json(serde_json::json!({ "error": "No job seeker profile found" }))).into_response(); }; @@ -1038,33 +1076,6 @@ async fn ai_tailor_resume( return (StatusCode::NOT_FOUND, Json(serde_json::json!({ "error": "Job not found" }))).into_response(); }; - let (has_pack, daily_limit) = { - let profile_id: Option = sqlx::query_scalar( - "SELECT id FROM user_role_profiles WHERE user_id = $1 AND role_key = 'JOB_SEEKER'" - ) - .bind(auth.user_id) - .fetch_optional(&state.pool) - .await - .ok() - .flatten(); - - match profile_id { - Some(pid) => has_active_ai_pack(&state.pool, pid, "JOB_SEEKER").await, - None => (false, BASE_AI_LIMIT), - } - }; - - let mut redis = state.redis.clone(); - let (used, limit) = match check_and_increment_usage(&state.pool, &mut redis, seeker_id, false, daily_limit).await { - Ok((u, l)) => (u, l), - Err(msg) => { - return (StatusCode::TOO_MANY_REQUESTS, Json(serde_json::json!({ "error": msg }))).into_response(); - } - }; - - let ollama_base = std::env::var("OLLAMA_BASE_URL").unwrap_or_else(|_| "http://ollama.nxtgauge-ai.svc.cluster.local:11434".to_string()); - let model = std::env::var("OLLAMA_CHAT_MODEL").unwrap_or_else(|_| "gemma3:270m".to_string()); - let existing_resume = body.resume_text.as_deref().unwrap_or("Not provided"); let skills_str = skills.join(", "); @@ -1086,23 +1097,66 @@ async fn ai_tailor_resume( full_name, experience, summary.as_deref().unwrap_or("N/A"), skills_str, existing_resume, job_title, job_desc ); - let generated = match call_ollama_inline(&ollama_base, &model, &prompt).await { - Ok(r) => r.trim().to_string(), - Err(e) => { - tracing::error!("Ollama resume tailoring error: {}", e); - return (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Generation failed" }))).into_response(); - } - }; + // Migrated to AI credits wallet using LiteLLM (Task 4) + let outcome = crate::ai_credits::call_litellm_and_charge( + &state.pool, + auth.user_id, + "resume_improve", + None, + None, + &prompt, + ) + .await; - ( - StatusCode::OK, - Json(GenerateFieldResponse { - generated_text: generated, - remaining_today: limit - used, - daily_limit: limit, - has_ai_pack: has_pack, - }), - ).into_response() + match outcome { + Ok(outcome) => { + let wallet = db::models::ai_credits::AiCreditsRepository::get_wallet(&state.pool, auth.user_id) + .await + .ok() + .flatten(); + let plan = match &wallet { + Some(w) => db::models::ai_credits::AiCreditsRepository::get_plan(&state.pool, w.plan_id).await.ok().flatten(), + None => None, + }; + ( + StatusCode::OK, + Json(GenerateFieldResponse { + generated_text: outcome.result.trim().to_string(), + remaining_today: wallet.as_ref().map(|w| w.available_credits()).unwrap_or(0), + daily_limit: plan.map(|p| p.daily_action_limit).unwrap_or(0), + has_ai_pack: wallet.as_ref().map(|w| w.purchased_credits_total > 0).unwrap_or(false), + }), + ) + .into_response() + } + Err(crate::ai_credits::ChargeError::Credits( + db::models::ai_credits::AiCreditsError::InsufficientCredits, + )) => ( + StatusCode::PAYMENT_REQUIRED, + Json(serde_json::json!({ "error": "Insufficient AI credits. Buy more from the Credits page." })), + ) + .into_response(), + Err(crate::ai_credits::ChargeError::Credits( + db::models::ai_credits::AiCreditsError::DailyActionLimitReached + | db::models::ai_credits::AiCreditsError::DailyCreditLimitReached, + )) => ( + StatusCode::TOO_MANY_REQUESTS, + Json(serde_json::json!({ "error": "Daily AI generation limit reached. Upgrade to AI Pack for more." })), + ) + .into_response(), + Err(crate::ai_credits::ChargeError::Credits(e)) => { + tracing::error!("AI credits error in ai_tailor_resume: {e}"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Internal error" }))).into_response() + } + Err(crate::ai_credits::ChargeError::UnknownFeature(f)) => { + tracing::error!("Unknown AI feature code '{f}' in ai_tailor_resume"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Internal error" }))).into_response() + } + Err(crate::ai_credits::ChargeError::Work(e)) => { + tracing::error!("Ollama resume tailoring error: {e}"); + (StatusCode::INTERNAL_SERVER_ERROR, Json(serde_json::json!({ "error": "Generation failed" }))).into_response() + } + } } // ── Auto Apply (Job Seekers) ─────────────────────────────────────────────────── @@ -1235,7 +1289,7 @@ async fn ai_auto_apply( full_name, experience, skills_str, summary.as_deref().unwrap_or(""), job_title, job_desc ); - let cover_letter = match call_ollama_inline(&ollama_base, &model, &cover_prompt).await { + let cover_letter = match call_ollama_inline(&ollama_base, &model, &cover_prompt, Some(&auth.user_id.to_string())).await { Ok(r) => r.trim().to_string(), Err(_) => "I am excited to apply for this position.".to_string(), }; @@ -3445,6 +3499,9 @@ pub fn ai_router() -> Router { .route("/auto-apply", post(ai_auto_apply)) .route("/auto-respond-to-lead", post(ai_auto_respond_to_lead)) .route("/usage", get(ai_usage_status)) + // ── Ask Ash AI Credits: usage summary/logs for CreditsPage.tsx ──── + .route("/usage/summary", get(crate::handlers::ai_credits::ai_usage_summary)) + .route("/usage/logs", get(crate::handlers::ai_credits::ai_usage_logs_handler)) // ── Phase 3: streaming, feedback, usage, GDPR clear ─────────────── .route("/chat/stream", post(phase3_chat_stream)) .route("/feedback", post(phase3::ai_feedback)) diff --git a/apps/users/src/handlers/ai_credits.rs b/apps/users/src/handlers/ai_credits.rs new file mode 100644 index 0000000..0a324a8 --- /dev/null +++ b/apps/users/src/handlers/ai_credits.rs @@ -0,0 +1,718 @@ +//! HTTP surface for the Ask Ash AI credits wallet (Phase 1). +//! +//! `/api/ai-credits/wallet` is real and safe to use from the frontend now. +//! `/api/ai-credits/charge` is a Phase-1 verification endpoint: it proves +//! the reserve -> work -> capture/release flow end-to-end over a real HTTP +//! request, but the "work" it charges for is a stub, not a real LiteLLM +//! call -- there is no LiteLLM client in this service yet (see Section 8 of +//! docs/ASK_ASH_BILLING_ARCHITECTURE.md in nxtgauge-ai-assistant). Wiring +//! real AI features to `charge_ai_feature` is Phase 2 work. + +use crate::ai_credits::{charge_ai_feature, ChargeError}; +use crate::AppState; +use axum::{ + extract::{Query, State}, + http::StatusCode, + response::IntoResponse, + routing::{get, post}, + Json, Router, +}; +use contracts::auth_middleware::{AuthUser, require_admin}; +use db::models::ai_credits::{AiCreditsError, AiCreditsRepository}; +use serde::{Deserialize, Serialize}; +use uuid::Uuid; + +pub fn router() -> Router { + Router::new() + .route("/wallet", get(get_wallet)) + .route("/charge", post(charge_feature)) +} + +/// Admin router for AI credits management +pub fn admin_router() -> Router { + Router::new() + .route("/balance", get(admin_balance)) + .route("/ledger", get(admin_ledger)) + .route("/adjust", post(admin_adjust)) + .route("/reconcile", get(admin_reconcile)) + .route("/refunds", post(admin_create_refund).get(admin_list_refunds)) +} + +#[derive(Serialize)] +struct WalletDto { + available_credits: i32, + monthly_credits_total: i32, + monthly_credits_used: i32, + purchased_credits_total: i32, + purchased_credits_used: i32, + bonus_credits_total: i32, + bonus_credits_used: i32, + reserved_credits: i32, + locked_credits: i32, + daily_actions_used: i32, + daily_credits_used: i32, +} + +fn credits_error_response(err: &AiCreditsError) -> (StatusCode, Json) { + let (status, code) = match err { + AiCreditsError::InsufficientCredits => (StatusCode::PAYMENT_REQUIRED, "INSUFFICIENT_AI_CREDITS"), + AiCreditsError::DailyActionLimitReached => (StatusCode::TOO_MANY_REQUESTS, "AI_DAILY_ACTION_LIMIT"), + AiCreditsError::DailyCreditLimitReached => (StatusCode::TOO_MANY_REQUESTS, "AI_DAILY_CREDIT_LIMIT"), + AiCreditsError::UnknownFeature(_) => (StatusCode::BAD_REQUEST, "UNKNOWN_AI_FEATURE"), + AiCreditsError::ReservationNotFound => (StatusCode::NOT_FOUND, "RESERVATION_NOT_FOUND"), + AiCreditsError::InvalidAmount(_) => (StatusCode::BAD_REQUEST, "INVALID_AMOUNT"), + AiCreditsError::Db(e) => { + tracing::error!("ai credits db error: {e}"); + (StatusCode::INTERNAL_SERVER_ERROR, "INTERNAL_ERROR") + } + }; + (status, Json(serde_json::json!({ "error": err.to_string(), "code": code }))) +} + +async fn get_wallet(auth: AuthUser, State(state): State) -> impl IntoResponse { + match AiCreditsRepository::ensure_wallet(&state.pool, auth.user_id).await { + Ok(wallet) => { + let dto = WalletDto { + available_credits: wallet.available_credits(), + monthly_credits_total: wallet.monthly_credits_total, + monthly_credits_used: wallet.monthly_credits_used, + purchased_credits_total: wallet.purchased_credits_total, + purchased_credits_used: wallet.purchased_credits_used, + bonus_credits_total: wallet.bonus_credits_total, + bonus_credits_used: wallet.bonus_credits_used, + reserved_credits: wallet.reserved_credits, + locked_credits: wallet.locked_credits, + daily_actions_used: wallet.daily_actions_used, + daily_credits_used: wallet.daily_credits_used, + }; + (StatusCode::OK, Json(serde_json::to_value(dto).unwrap())).into_response() + } + Err(e) => credits_error_response(&AiCreditsError::Db(e)).into_response(), + } +} + +#[derive(Deserialize)] +struct ChargeRequest { + feature_code: String, + request_id: Option, + idempotency_key: Option, +} + +#[derive(Serialize)] +struct ChargeResponse { + feature_code: String, + credits_charged: i32, + available_credits: i32, +} + +async fn charge_feature( + auth: AuthUser, + State(state): State, + Json(body): Json, +) -> impl IntoResponse { + let outcome = charge_ai_feature::<_, std::convert::Infallible, _, _>( + &state.pool, + auth.user_id, + &body.feature_code, + body.request_id.as_deref(), + body.idempotency_key.as_deref(), + |_feature| async move { Ok(()) }, + ) + .await; + + match outcome { + Ok(outcome) => { + let wallet = AiCreditsRepository::get_wallet(&state.pool, auth.user_id) + .await + .ok() + .flatten(); + let response = ChargeResponse { + feature_code: outcome.feature.feature_code, + credits_charged: outcome.credits_charged, + available_credits: wallet.map(|w| w.available_credits()).unwrap_or_default(), + }; + (StatusCode::OK, Json(serde_json::to_value(response).unwrap())).into_response() + } + Err(ChargeError::Credits(e)) => credits_error_response(&e).into_response(), + Err(ChargeError::UnknownFeature(f)) => credits_error_response(&AiCreditsError::UnknownFeature(f)).into_response(), + Err(ChargeError::Work(never)) => match never {}, + } +} + +#[derive(Serialize)] +struct PlanDetailsDto { + remaining_credits: i32, + remaining_daily_actions: i32, + daily_action_limit: i32, + plan_code: String, + plan_name: String, + monthly_credits_total: i32, + monthly_credits_used: i32, + purchased_credits_total: i32, + purchased_credits_used: i32, +} + +/// GET /api/ai/usage/summary -- powers CreditsPage.tsx's "AI Usage" tab +/// overview cards (`AiUsageSummary` type in that file). +pub async fn ai_usage_summary(auth: AuthUser, State(state): State) -> impl IntoResponse { + let wallet = match AiCreditsRepository::ensure_wallet(&state.pool, auth.user_id).await { + Ok(w) => w, + Err(e) => return credits_error_response(&AiCreditsError::Db(e)).into_response(), + }; + let plan = match AiCreditsRepository::get_plan(&state.pool, wallet.plan_id).await { + Ok(Some(p)) => p, + Ok(None) => return credits_error_response(&AiCreditsError::Db(sqlx::Error::RowNotFound)).into_response(), + Err(e) => return credits_error_response(&AiCreditsError::Db(e)).into_response(), + }; + + let dto = PlanDetailsDto { + remaining_credits: wallet.available_credits(), + remaining_daily_actions: (plan.daily_action_limit - wallet.daily_actions_used).max(0), + daily_action_limit: plan.daily_action_limit, + plan_code: plan.code, + plan_name: plan.name, + monthly_credits_total: wallet.monthly_credits_total, + monthly_credits_used: wallet.monthly_credits_used, + purchased_credits_total: wallet.purchased_credits_total, + purchased_credits_used: wallet.purchased_credits_used, + }; + (StatusCode::OK, Json(serde_json::json!({ "plan_details": dto }))).into_response() +} + +#[derive(Serialize, sqlx::FromRow)] +struct UsageLogDto { + id: uuid::Uuid, + feature_code: String, + model_alias: String, + credits_charged: i32, + status: String, + created_at: chrono::DateTime, +} + +/// GET /api/ai/usage/logs -- powers CreditsPage.tsx's "AI Usage" tab +/// history list (`AiUsageLog` type in that file). +pub async fn ai_usage_logs_handler(auth: AuthUser, State(state): State) -> impl IntoResponse { + let logs = sqlx::query_as::<_, UsageLogDto>( + r#" + SELECT id, feature_code, model_alias, credits_charged, status, created_at + FROM ai_usage_logs + WHERE user_id = $1 + ORDER BY created_at DESC + LIMIT 100 + "#, + ) + .bind(auth.user_id) + .fetch_all(&state.pool) + .await; + + match logs { + Ok(logs) => (StatusCode::OK, Json(serde_json::json!({ "logs": logs }))).into_response(), + Err(e) => credits_error_response(&AiCreditsError::Db(e)).into_response(), + } +} + +// ============================================================================ +// ADMIN ENDPOINTS +// ============================================================================ + + + +/// GET /api/admin/ai-credits/balance?userId= +/// Get wallet balance for a specific user (admin only) +async fn admin_balance( + auth: AuthUser, + State(state): State, + Query(params): Query, +) -> impl IntoResponse { + if let Err(e) = require_admin(&auth) { + return e.into_response(); + } + + let user_id = match params.user_id { + Some(id) => id, + None => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ "error": "userId parameter required", "code": "MISSING_USER_ID" })) + ).into_response(); + } + }; + + match AiCreditsRepository::ensure_wallet(&state.pool, user_id).await { + Ok(wallet) => { + let dto = WalletDto { + available_credits: wallet.available_credits(), + monthly_credits_total: wallet.monthly_credits_total, + monthly_credits_used: wallet.monthly_credits_used, + purchased_credits_total: wallet.purchased_credits_total, + purchased_credits_used: wallet.purchased_credits_used, + bonus_credits_total: wallet.bonus_credits_total, + bonus_credits_used: wallet.bonus_credits_used, + reserved_credits: wallet.reserved_credits, + locked_credits: wallet.locked_credits, + daily_actions_used: wallet.daily_actions_used, + daily_credits_used: wallet.daily_credits_used, + }; + (StatusCode::OK, Json(serde_json::to_value(dto).unwrap())).into_response() + } + Err(e) => credits_error_response(&AiCreditsError::Db(e)).into_response(), + } +} + +#[derive(Deserialize)] +struct AdminBalanceQuery { + user_id: Option, +} + +/// GET /api/admin/ai-credits/ledger?userId=&page=&limit= +/// Get paginated ledger entries for a user (admin only) +async fn admin_ledger( + auth: AuthUser, + State(state): State, + Query(params): Query, +) -> impl IntoResponse { + if let Err(e) = require_admin(&auth) { + return e.into_response(); + } + + let user_id = match params.user_id { + Some(id) => id, + None => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ "error": "userId parameter required", "code": "MISSING_USER_ID" })) + ).into_response(); + } + }; + + let page = params.page.unwrap_or(1).max(1); + let limit = params.limit.unwrap_or(20).clamp(1, 100); + let offset = (page - 1) * limit; + + // Get total count + let total: i64 = match sqlx::query_scalar( + r#" + SELECT COUNT(*) + FROM ai_credit_ledger l + JOIN user_ai_subscriptions w ON l.wallet_id = w.id + WHERE w.user_id = $1 + "# + ) + .bind(user_id) + .fetch_one(&state.pool) + .await + { + Ok(count) => count, + Err(e) => return credits_error_response(&AiCreditsError::Db(e)).into_response(), + }; + + // Get ledger entries + let entries: Vec = match sqlx::query_as( + r#" + SELECT + l.id, + l.entry_type, + l.credits, + l.balance_after, + l.description, + l.idempotency_key, + l.created_at + FROM ai_credit_ledger l + JOIN user_ai_subscriptions w ON l.wallet_id = w.id + WHERE w.user_id = $1 + ORDER BY l.created_at DESC + LIMIT $2 OFFSET $3 + "# + ) + .bind(user_id) + .bind(limit) + .bind(offset) + .fetch_all(&state.pool) + .await + { + Ok(entries) => entries, + Err(e) => return credits_error_response(&AiCreditsError::Db(e)).into_response(), + }; + + ( + StatusCode::OK, + Json(serde_json::json!({ + "data": entries, + "total": total, + "page": page, + "limit": limit + })) + ).into_response() +} + +#[derive(Deserialize)] +struct AdminLedgerQuery { + user_id: Option, + page: Option, + limit: Option, +} + +#[derive(Serialize, sqlx::FromRow)] +struct LedgerEntryDto { + id: uuid::Uuid, + entry_type: String, + credits: i32, + balance_after: i32, + description: Option, + idempotency_key: Option, + created_at: chrono::DateTime, +} + +/// POST /api/admin/ai-credits/adjust +/// Adjust credits for a user (ADD or DEDUCT) +async fn admin_adjust( + auth: AuthUser, + State(state): State, + Json(body): Json, +) -> impl IntoResponse { + if let Err(e) = require_admin(&auth) { + return e.into_response(); + } + + // Validate reason is not empty + if body.reason.trim().is_empty() { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "reason is required and cannot be empty", + "code": "INVALID_REASON" + })) + ).into_response(); + } + + let is_add = body.type_ == "ADD"; + let idempotency_key = body.idempotency_key + .unwrap_or_else(|| format!("admin-adjust:{}:{}", auth.user_id, uuid::Uuid::new_v4())); + + match AiCreditsRepository::admin_adjust_credits( + &state.pool, + body.user_id, + body.amount, + is_add, + &body.reason, + auth.user_id, + Some(&idempotency_key), + ).await + { + Ok(wallet) => { + let dto = WalletDto { + available_credits: wallet.available_credits(), + monthly_credits_total: wallet.monthly_credits_total, + monthly_credits_used: wallet.monthly_credits_used, + purchased_credits_total: wallet.purchased_credits_total, + purchased_credits_used: wallet.purchased_credits_used, + bonus_credits_total: wallet.bonus_credits_total, + bonus_credits_used: wallet.bonus_credits_used, + reserved_credits: wallet.reserved_credits, + locked_credits: wallet.locked_credits, + daily_actions_used: wallet.daily_actions_used, + daily_credits_used: wallet.daily_credits_used, + }; + (StatusCode::OK, Json(serde_json::json!({ + "success": true, + "wallet": dto, + "operation": if is_add { "ADD" } else { "DEDUCT" }, + "amount": body.amount + }))).into_response() + } + Err(AiCreditsError::InsufficientCredits) => ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "Insufficient credits for deduction", + "code": "INSUFFICIENT_CREDITS" + })) + ).into_response(), + Err(e) => credits_error_response(&e).into_response(), + } +} + +#[derive(Deserialize)] +struct AdminAdjustRequest { + user_id: Uuid, + amount: i32, + #[serde(rename = "type")] + type_: String, + reason: String, + idempotency_key: Option, +} + +/// GET /api/admin/ai-credits/reconcile?from=&to= +/// Reconcile ledger entries within date range +async fn admin_reconcile( + auth: AuthUser, + State(state): State, + Query(params): Query, +) -> impl IntoResponse { + if let Err(e) = require_admin(&auth) { + return e.into_response(); + } + + let from_date = match params.from { + Some(d) => d, + None => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ "error": "from date parameter required", "code": "MISSING_FROM_DATE" })) + ).into_response(); + } + }; + + let to_date = match params.to { + Some(d) => d, + None => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ "error": "to date parameter required", "code": "MISSING_TO_DATE" })) + ).into_response(); + } + }; + + // Sum credits by entry type within date range + let summary: Vec = match sqlx::query_as( + r#" + SELECT + entry_type, + COUNT(*) as count, + SUM(credits) as total_credits + FROM ai_credit_ledger + WHERE created_at >= $1 AND created_at <= $2 + GROUP BY entry_type + ORDER BY entry_type + "# + ) + .bind(from_date) + .bind(to_date) + .fetch_all(&state.pool) + .await + { + Ok(rows) => rows, + Err(e) => return credits_error_response(&AiCreditsError::Db(e)).into_response(), + }; + + // Check for drift in sample of wallets + let drift_check: Vec = match sqlx::query_as( + r#" + SELECT + w.user_id, + u.email as user_email, + w.id as wallet_id, + COALESCE(SUM(l.credits), 0) as ledger_sum, + ( + w.bonus_credits_total - w.bonus_credits_used + + w.monthly_credits_total - w.monthly_credits_used + + w.purchased_credits_total - w.purchased_credits_used - + w.reserved_credits - w.locked_credits + ) as computed_available + FROM user_ai_subscriptions w + JOIN users u ON w.user_id = u.id + LEFT JOIN ai_credit_ledger l ON l.wallet_id = w.id + WHERE w.updated_at >= $1 AND w.updated_at <= $2 + GROUP BY w.id, w.user_id, u.email, w.bonus_credits_total, w.bonus_credits_used, + w.monthly_credits_total, w.monthly_credits_used, + w.purchased_credits_total, w.purchased_credits_used, + w.reserved_credits, w.locked_credits + LIMIT 100 + "# + ) + .bind(from_date) + .bind(to_date) + .fetch_all(&state.pool) + .await + { + Ok(rows) => rows, + Err(e) => return credits_error_response(&AiCreditsError::Db(e)).into_response(), + }; + + let drift_issues: Vec<_> = drift_check + .into_iter() + .filter(|w| w.ledger_sum != w.computed_available) + .map(|w| serde_json::json!({ + "user_id": w.user_id, + "user_email": w.user_email, + "ledger_sum": w.ledger_sum, + "computed_available": w.computed_available, + "drift": w.ledger_sum - w.computed_available + })) + .collect(); + + ( + StatusCode::OK, + Json(serde_json::json!({ + "from": from_date, + "to": to_date, + "summary": summary, + "drift_check": { + "wallets_checked": drift_issues.len(), + "issues_found": drift_issues.len(), + "issues": drift_issues + } + })) + ).into_response() +} + +#[derive(Deserialize)] +struct AdminReconcileQuery { + from: Option>, + to: Option>, +} + +#[derive(Serialize, sqlx::FromRow)] +struct ReconcileSummaryDto { + entry_type: String, + count: i64, + total_credits: i64, +} + +#[derive(Serialize, sqlx::FromRow)] +struct WalletDriftDto { + user_id: Uuid, + user_email: String, + #[sqlx(rename = "wallet_id")] + _wallet_id: Uuid, + ledger_sum: i32, + computed_available: i32, +} + +// ============================================================================ +// REFUND ENDPOINTS (Task 5) +// ============================================================================ + +/// POST /api/admin/ai-credits/refunds +/// Process a manual refund for a completed charge +async fn admin_create_refund( + auth: AuthUser, + State(state): State, + Json(body): Json, +) -> impl IntoResponse { + if let Err(e) = require_admin(&auth) { + return e.into_response(); + } + + // Validate reason + let valid_reasons = ["provider_failure", "timeout", "validation_failure", "manual", "user_dispute"]; + if !valid_reasons.contains(&body.reason.as_str()) { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ + "error": "Invalid reason", + "code": "INVALID_REASON", + "valid_reasons": valid_reasons + })) + ).into_response(); + } + + // Generate idempotency key if not provided + let idempotency_key = body.idempotency_key + .unwrap_or_else(|| format!("admin-refund:{}: {}", auth.user_id, uuid::Uuid::new_v4())); + + match AiCreditsRepository::refund_credits( + &state.pool, + body.user_id, + body.ledger_debit_entry_id, + body.credits_to_refund, + &body.reason, + "admin", + Some(auth.user_id), + Some(&idempotency_key), + body.notes.as_deref(), + ).await { + Ok((refund_id, wallet)) => { + let dto = WalletDto { + available_credits: wallet.available_credits(), + monthly_credits_total: wallet.monthly_credits_total, + monthly_credits_used: wallet.monthly_credits_used, + purchased_credits_total: wallet.purchased_credits_total, + purchased_credits_used: wallet.purchased_credits_used, + bonus_credits_total: wallet.bonus_credits_total, + bonus_credits_used: wallet.bonus_credits_used, + reserved_credits: wallet.reserved_credits, + locked_credits: wallet.locked_credits, + daily_actions_used: wallet.daily_actions_used, + daily_credits_used: wallet.daily_credits_used, + }; + ( + StatusCode::OK, + Json(serde_json::json!({ + "success": true, + "refund_id": refund_id, + "wallet": dto, + "credits_refunded": body.credits_to_refund, + "reason": body.reason + })) + ).into_response() + } + Err(AiCreditsError::InvalidAmount(msg)) => ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ "error": msg, "code": "INVALID_AMOUNT" })) + ).into_response(), + Err(AiCreditsError::ReservationNotFound) => ( + StatusCode::NOT_FOUND, + Json(serde_json::json!({ + "error": "Original ledger entry not found", + "code": "LEDGER_ENTRY_NOT_FOUND" + })) + ).into_response(), + Err(e) => credits_error_response(&e).into_response(), + } +} + +#[derive(Deserialize)] +struct CreateRefundRequest { + user_id: Uuid, + ledger_debit_entry_id: Uuid, + credits_to_refund: i32, + reason: String, + notes: Option, + idempotency_key: Option, +} + +/// GET /api/admin/ai-credits/refunds?userId=&page=&limit= +/// List refunds for a user (admin view) +async fn admin_list_refunds( + auth: AuthUser, + State(state): State, + Query(params): Query, +) -> impl IntoResponse { + if let Err(e) = require_admin(&auth) { + return e.into_response(); + } + + let user_id = match params.user_id { + Some(id) => id, + None => { + return ( + StatusCode::BAD_REQUEST, + Json(serde_json::json!({ "error": "userId parameter required", "code": "MISSING_USER_ID" })) + ).into_response(); + } + }; + + let page = params.page.unwrap_or(1).max(1); + let limit = params.limit.unwrap_or(20).clamp(1, 100); + + match AiCreditsRepository::list_refunds(&state.pool, user_id, page, limit).await { + Ok(refunds) => { + ( + StatusCode::OK, + Json(serde_json::json!({ + "data": refunds, + "page": page, + "limit": limit + })) + ).into_response() + } + Err(e) => credits_error_response(&AiCreditsError::Db(e)).into_response(), + } +} + +#[derive(Deserialize)] +struct ListRefundsQuery { + user_id: Option, + page: Option, + limit: Option, +} diff --git a/apps/users/src/handlers/ai_phase4.rs b/apps/users/src/handlers/ai_phase4.rs index d2e0446..5ddf20e 100644 --- a/apps/users/src/handlers/ai_phase4.rs +++ b/apps/users/src/handlers/ai_phase4.rs @@ -150,7 +150,7 @@ pub struct TranslateResponse { async fn translate_text( State(state): State, - _auth: AuthUser, + auth: AuthUser, Json(body): Json, ) -> impl IntoResponse { // If source==target, skip the round-trip. @@ -178,7 +178,7 @@ async fn translate_text( body.target_lang, body.text ); - match call_ollama_generate(&ollama_base, &model, &prompt).await { + match call_ollama_generate(&ollama_base, &model, &prompt, Some(&auth.user_id.to_string())).await { Ok(translated) => { let _ = state; // silence unused warning; state is plumbed for future Json(TranslateResponse { @@ -855,7 +855,7 @@ pub async fn maybe_summarise_history( transcript ); - let summary = match call_ollama_generate(&ollama_base, &model, &prompt).await { + let summary = match call_ollama_generate(&ollama_base, &model, &prompt, Some(&user_id.to_string())).await { Ok(s) => s, Err(_) => { // No Ollama: fall back to a truncation summary. Better than 5xx. @@ -1269,30 +1269,27 @@ async fn call_ollama_generate( base_url: &str, model: &str, prompt: &str, + user_id: Option<&str>, ) -> Result { - let url = format!("{}/api/generate", base_url); - let client = reqwest::Client::new(); - let resp = client - .post(&url) - .json(&serde_json::json!({ - "model": model, - "prompt": prompt, - "stream": false, - })) - .send() - .await - .map_err(|e| format!("ollama request failed: {}", e))?; - if !resp.status().is_success() { - return Err(format!("ollama status: {}", resp.status())); + // Build secure client with rate limiting and security hardening + let security = cache::ollama::OllamaSecurityConfig { + api_key: std::env::var("OLLAMA_API_KEY").ok(), + user_id: user_id.map(|s| s.to_string()), + client_ip: None, + rate_limit_per_minute: std::env::var("OLLAMA_RATE_LIMIT_PER_MINUTE") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(60), + }; + + let client = cache::ollama::OllamaClient::with_url(base_url) + .with_model(model) + .with_security(security); + + match client.generate(prompt).await { + Ok(response) => Ok(response.response), + Err(e) => Err(format!("Ollama error: {}", e)), } - let v: JsonValue = resp - .json() - .await - .map_err(|e| format!("parse ollama response: {}", e))?; - Ok(v.get("response") - .and_then(|s| s.as_str()) - .unwrap_or("") - .to_string()) } // ───────────────────────────────────────────────────────────────────────────── diff --git a/apps/users/src/handlers/mod.rs b/apps/users/src/handlers/mod.rs index 694c0f6..b9e603b 100644 --- a/apps/users/src/handlers/mod.rs +++ b/apps/users/src/handlers/mod.rs @@ -4,6 +4,7 @@ pub mod activity_logs; pub mod approvals; pub mod auth; pub mod ai; +pub mod ai_credits; pub mod ai_phase4; pub mod ai_prompts; pub mod config; diff --git a/apps/users/src/litellm.rs b/apps/users/src/litellm.rs new file mode 100644 index 0000000..9d12c44 --- /dev/null +++ b/apps/users/src/litellm.rs @@ -0,0 +1,222 @@ +//! LiteLLM client for the Ask Ash AI credits system. +//! +//! Thin HTTP wrapper around LiteLLM's `/chat/completions` endpoint. +//! This is the preferred way to call LLMs - it provides model routing, +//! fallbacks, and unified interface while billing remains in nxtgauge-backend-rust. + +use serde::{Deserialize, Serialize}; + +#[derive(Debug, Clone, Serialize)] +struct LiteLLMChatMessage { + role: String, + content: String, +} + +#[derive(Debug, Clone, Serialize)] +struct LiteLLMRequest { + model: String, + messages: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + max_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + temperature: Option, +} + +#[derive(Debug, Clone, Deserialize)] +struct LiteLLMChoice { + message: LiteLLMChatMessage, + finish_reason: Option, +} + +#[derive(Debug, Clone, Deserialize)] +struct LiteLLMUsage { + prompt_tokens: Option, + completion_tokens: Option, + total_tokens: Option, +} + +#[derive(Debug, Clone, Deserialize)] +pub struct LiteLLMResponse { + pub id: String, + pub model: String, + pub choices: Vec, + pub usage: Option, +} + +impl LiteLLMResponse { + /// Extract the generated text from the response + pub fn generated_text(&self) -> String { + self.choices + .first() + .map(|c| c.message.content.clone()) + .unwrap_or_default() + } + + /// Get token counts if available + pub fn token_counts(&self) -> Option<(i32, i32, i32)> { + self.usage.as_ref().map(|u| ( + u.prompt_tokens.unwrap_or(0), + u.completion_tokens.unwrap_or(0), + u.total_tokens.unwrap_or(0), + )) + } +} + +/// Call LiteLLM with a single prompt (simplified interface) +/// +/// # Arguments +/// * `base_url` - LiteLLM base URL (e.g., "http://litellm.nxtgauge-ai.svc.cluster.local:4000") +/// * `model_alias` - LiteLLM model alias (e.g., "askash-fast", "askash-main") +/// * `prompt` - The user prompt +/// * `api_key` - Optional API key for authentication +/// * `max_tokens` - Optional max output tokens +/// +/// # Returns +/// * `Ok(LiteLLMResponse)` - The successful response with generated text and metadata +/// * `Err(String)` - Error message if the request failed +pub async fn call_litellm( + base_url: &str, + model_alias: &str, + prompt: &str, + api_key: Option<&str>, + max_tokens: Option, +) -> Result { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(60)) + .build() + .map_err(|e| format!("Failed to build HTTP client: {}", e))?; + + let url = format!("{}/v1/chat/completions", base_url.trim_end_matches('/')); + + let request = LiteLLMRequest { + model: model_alias.to_string(), + messages: vec![ + LiteLLMChatMessage { + role: "user".to_string(), + content: prompt.to_string(), + }, + ], + max_tokens, + temperature: Some(0.7), + }; + + let mut req_builder = client.post(&url) + .header("Content-Type", "application/json"); + + // Add authorization if API key is provided + if let Some(key) = api_key { + req_builder = req_builder.header("Authorization", format!("Bearer {}", key)); + } + + let response = req_builder + .json(&request) + .send() + .await + .map_err(|e| format!("LiteLLM request failed: {}", e))?; + + let status = response.status(); + if !status.is_success() { + let error_body = response.text().await.unwrap_or_else(|_| "Unknown error".to_string()); + return Err(format!("LiteLLM error ({}): {}", status, error_body)); + } + + let litellm_response: LiteLLMResponse = response + .json() + .await + .map_err(|e| format!("Failed to parse LiteLLM response: {}", e))?; + + Ok(litellm_response) +} + +/// Call LiteLLM with a system prompt and user message +pub async fn call_litellm_with_system( + base_url: &str, + model_alias: &str, + system_prompt: &str, + user_message: &str, + api_key: Option<&str>, + max_tokens: Option, +) -> Result { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(60)) + .build() + .map_err(|e| format!("Failed to build HTTP client: {}", e))?; + + let url = format!("{}/v1/chat/completions", base_url.trim_end_matches('/')); + + let request = LiteLLMRequest { + model: model_alias.to_string(), + messages: vec![ + LiteLLMChatMessage { + role: "system".to_string(), + content: system_prompt.to_string(), + }, + LiteLLMChatMessage { + role: "user".to_string(), + content: user_message.to_string(), + }, + ], + max_tokens, + temperature: Some(0.7), + }; + + let mut req_builder = client.post(&url) + .header("Content-Type", "application/json"); + + if let Some(key) = api_key { + req_builder = req_builder.header("Authorization", format!("Bearer {}", key)); + } + + let response = req_builder + .json(&request) + .send() + .await + .map_err(|e| format!("LiteLLM request failed: {}", e))?; + + let status = response.status(); + if !status.is_success() { + let error_body = response.text().await.unwrap_or_else(|_| "Unknown error".to_string()); + return Err(format!("LiteLLM error ({}): {}", status, error_body)); + } + + let litellm_response: LiteLLMResponse = response + .json() + .await + .map_err(|e| format!("Failed to parse LiteLLM response: {}", e))?; + + Ok(litellm_response) +} + +/// Get LiteLLM configuration from environment +pub fn get_litellm_config() -> (String, String, Option) { + let base_url = std::env::var("LITELLM_BASE_URL") + .unwrap_or_else(|_| "http://litellm.nxtgauge-ai.svc.cluster.local:4000".to_string()); + + let model_alias = std::env::var("LITELLM_MODEL") + .unwrap_or_else(|_| "askash-fast".to_string()); + + let api_key = std::env::var("LITELLM_API_KEY").ok(); + + (base_url, model_alias, api_key) +} + +/// Health check for LiteLLM +pub async fn check_litellm_health(base_url: &str, api_key: Option<&str>) -> Result { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(5)) + .build() + .map_err(|e| format!("Failed to build HTTP client: {}", e))?; + + let url = format!("{}/health", base_url.trim_end_matches('/')); + + let mut req_builder = client.get(&url); + + if let Some(key) = api_key { + req_builder = req_builder.header("Authorization", format!("Bearer {}", key)); + } + + match req_builder.send().await { + Ok(response) => Ok(response.status().is_success()), + Err(e) => Err(format!("Health check failed: {}", e)), + } +} diff --git a/apps/users/src/main.rs b/apps/users/src/main.rs index 9479527..acdceb7 100644 --- a/apps/users/src/main.rs +++ b/apps/users/src/main.rs @@ -1,4 +1,7 @@ +mod ai_credits; +mod ai_subscription; mod handlers; +mod litellm; mod mail; use axum::{routing::get, Router}; @@ -111,6 +114,9 @@ async fn main() { .nest("/api/admin/email", handlers::admin_email::router()) // ── AI Assistant ────────────────────────────────────────────────── .nest("/api/ai", handlers::ai::ai_router()) + // ── Ask Ash AI Credits (Phase 1) ──────────────────────────────────── + .nest("/api/ai-credits", handlers::ai_credits::router()) + .nest("/api/admin/ai-credits", handlers::ai_credits::admin_router()) .route("/health", get(|| async { "Users OK" })) .with_state(state); diff --git a/crates/auth/Cargo.toml b/crates/auth/Cargo.toml index 166016b..4ae55b3 100644 --- a/crates/auth/Cargo.toml +++ b/crates/auth/Cargo.toml @@ -4,7 +4,7 @@ version = "0.1.0" edition = "2021" [dependencies] -jsonwebtoken = "10.3" +jsonwebtoken = { version = "10.3", features = ["rust_crypto"] } argon2 = "0.5" rand_core = { version = "0.6", features = ["std"] } serde = { workspace = true } diff --git a/crates/cache/Cargo.toml b/crates/cache/Cargo.toml index df2395f..9d83080 100644 --- a/crates/cache/Cargo.toml +++ b/crates/cache/Cargo.toml @@ -12,3 +12,4 @@ uuid = { workspace = true } tracing = { workspace = true } thiserror = { workspace = true } reqwest = { workspace = true } +lazy_static = "1.4" diff --git a/crates/cache/src/ollama.rs b/crates/cache/src/ollama.rs index 8d842a9..6116d52 100644 --- a/crates/cache/src/ollama.rs +++ b/crates/cache/src/ollama.rs @@ -1,6 +1,10 @@ -//! Ollama client for AI-powered text generation +//! Ollama client for AI-powered text generation with security hardening. //! -//! Used for generating job descriptions, resume analysis, and other AI features +//! Security features: +//! - API key authentication via X-API-Key header +//! - Rate limiting per user/IP +//! - Request timeouts and circuit breaker +//! - Input validation and sanitization use reqwest::{Client, Error as ReqwestError}; use serde::{Deserialize, Serialize}; @@ -9,12 +13,39 @@ use std::time::Duration; const OLLAMA_URL: &str = "http://nxtgauge-ai-assistant:11434"; const DEFAULT_MODEL: &str = "gemma3:270m"; const REQUEST_TIMEOUT: Duration = Duration::from_secs(120); +const MAX_PROMPT_LENGTH: usize = 100_000; // 100KB limit +const MAX_RETRIES: u32 = 3; + +/// Security configuration for Ollama client +#[derive(Debug, Clone)] +pub struct OllamaSecurityConfig { + /// API key for authentication (required) + pub api_key: Option, + /// User ID for rate limiting tracking + pub user_id: Option, + /// Client IP for rate limiting + pub client_ip: Option, + /// Maximum requests per minute per user + pub rate_limit_per_minute: u32, +} + +impl Default for OllamaSecurityConfig { + fn default() -> Self { + Self { + api_key: std::env::var("OLLAMA_API_KEY").ok(), + user_id: None, + client_ip: None, + rate_limit_per_minute: 60, // Default: 60 requests/minute + } + } +} #[derive(Debug, Clone)] pub struct OllamaClient { http_client: Client, base_url: String, model: String, + security: OllamaSecurityConfig, } #[derive(Debug, Serialize)] @@ -69,6 +100,154 @@ pub enum OllamaError { #[error("Model not found: {0}")] ModelNotFound(String), + + #[error("API key required but not provided")] + MissingApiKey, + + #[error("Invalid API key")] + InvalidApiKey, + + #[error("Rate limit exceeded: {0}")] + RateLimitExceeded(String), + + #[error("Prompt too long: {0} chars (max {1})")] + PromptTooLong(usize, usize), + + #[error("Circuit breaker open: too many failures")] + CircuitBreakerOpen, + + #[error("Max retries exceeded")] + MaxRetriesExceeded, +} + +/// Validates and sanitizes user input +fn validate_prompt(prompt: &str) -> Result { + if prompt.len() > MAX_PROMPT_LENGTH { + return Err(OllamaError::PromptTooLong(prompt.len(), MAX_PROMPT_LENGTH)); + } + + // Basic sanitization - remove null bytes and control characters + let sanitized: String = prompt + .chars() + .filter(|c| !c.is_control() || *c == '\n' || *c == '\t') + .collect(); + + Ok(sanitized) +} + +/// Rate limiter using a simple token bucket approach +#[derive(Debug)] +pub struct RateLimiter { + requests: std::sync::Arc>>>, + max_requests: u32, + window_duration: Duration, +} + +impl RateLimiter { + pub fn new(max_requests: u32, window_minutes: u64) -> Self { + Self { + requests: std::sync::Arc::new(std::sync::Mutex::new(std::collections::HashMap::new())), + max_requests, + window_duration: Duration::from_secs(window_minutes * 60), + } + } + + /// Check if the request is allowed + pub fn check_rate_limit(&self, key: &str) -> Result<(), OllamaError> { + let now = std::time::Instant::now(); + let window_start = now - self.window_duration; + + let mut requests = self.requests.lock().map_err(|_| { + OllamaError::ApiError("Rate limiter lock failed".to_string()) + })?; + + // Get or create entry for this key + let user_requests = requests.entry(key.to_string()).or_default(); + + // Remove old requests outside the window + user_requests.retain(|&time| time > window_start); + + // Check if under limit + if user_requests.len() >= self.max_requests as usize { + return Err(OllamaError::RateLimitExceeded( + format!("{} requests per {} minutes exceeded", self.max_requests, self.window_duration.as_secs() / 60) + )); + } + + // Add this request + user_requests.push(now); + Ok(()) + } +} + +/// Circuit breaker for handling cascading failures +#[derive(Debug)] +pub struct CircuitBreaker { + failures: std::sync::Arc, + threshold: u32, + reset_after: Duration, + last_failure: std::sync::Arc>>, +} + +impl CircuitBreaker { + pub fn new(threshold: u32, reset_after_secs: u64) -> Self { + Self { + failures: std::sync::Arc::new(std::sync::atomic::AtomicU32::new(0)), + threshold, + reset_after: Duration::from_secs(reset_after_secs), + last_failure: std::sync::Arc::new(std::sync::Mutex::new(None)), + } + } + + pub fn can_execute(&self) -> Result<(), OllamaError> { + let failures = self.failures.load(std::sync::atomic::Ordering::Relaxed); + + if failures >= self.threshold { + // Check if enough time has passed to reset + if let Ok(last) = self.last_failure.lock() { + if let Some(last_failure) = *last { + if std::time::Instant::now().duration_since(last_failure) > self.reset_after { + // Reset circuit breaker + self.failures.store(0, std::sync::atomic::Ordering::Relaxed); + return Ok(()); + } + } + } + return Err(OllamaError::CircuitBreakerOpen); + } + + Ok(()) + } + + pub fn record_success(&self) { + self.failures.store(0, std::sync::atomic::Ordering::Relaxed); + } + + pub fn record_failure(&self) { + self.failures.fetch_add(1, std::sync::atomic::Ordering::Relaxed); + if let Ok(mut last) = self.last_failure.lock() { + *last = Some(std::time::Instant::now()); + } + } +} + +/// Global rate limiter instance (shared across all clients) +lazy_static::lazy_static! { + static ref GLOBAL_RATE_LIMITER: RateLimiter = RateLimiter::new( + std::env::var("OLLAMA_RATE_LIMIT_PER_MINUTE") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(60), + 1 // 1 minute window + ); + + static ref CIRCUIT_BREAKER: CircuitBreaker = CircuitBreaker::new( + std::env::var("OLLAMA_CIRCUIT_THRESHOLD") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(5), + 60 // Reset after 60 seconds + ); } impl OllamaClient { @@ -82,6 +261,7 @@ impl OllamaClient { http_client, base_url: OLLAMA_URL.to_string(), model: DEFAULT_MODEL.to_string(), + security: OllamaSecurityConfig::default(), } } @@ -95,9 +275,15 @@ impl OllamaClient { http_client, base_url: base_url.into(), model: DEFAULT_MODEL.to_string(), + security: OllamaSecurityConfig::default(), } } + pub fn with_security(mut self, security: OllamaSecurityConfig) -> Self { + self.security = security; + self + } + pub fn with_model(mut self, model: impl Into) -> Self { self.model = model.into(); self @@ -107,20 +293,97 @@ impl OllamaClient { &self.model } + fn get_rate_limit_key(&self) -> String { + self.security.user_id.clone() + .or_else(|| self.security.client_ip.clone()) + .unwrap_or_else(|| "anonymous".to_string()) + } + + fn build_request_headers(&self) -> reqwest::header::HeaderMap { + let mut headers = reqwest::header::HeaderMap::new(); + + // Add API key header if configured + if let Some(api_key) = &self.security.api_key { + if let Ok(val) = reqwest::header::HeaderValue::from_str(api_key) { + headers.insert("X-API-Key", val); + } + } + + // Add user tracking headers for audit + if let Some(user_id) = &self.security.user_id { + if let Ok(val) = reqwest::header::HeaderValue::from_str(user_id) { + headers.insert("X-User-Id", val); + } + } + + headers + } + /// Generate text using the configured model and prompt pub async fn generate(&self, prompt: impl Into) -> Result { + // 1. Check circuit breaker + CIRCUIT_BREAKER.can_execute()?; + + // 2. Validate and sanitize input + let sanitized_prompt = validate_prompt(&prompt.into())?; + + // 3. Check rate limit + let rate_key = self.get_rate_limit_key(); + GLOBAL_RATE_LIMITER.check_rate_limit(&rate_key)?; + + // 4. Build request let request = GenerateRequest { model: self.model.clone(), - prompt: prompt.into(), + prompt: sanitized_prompt, stream: false, options: None, }; let url = format!("{}/api/generate", self.base_url); - let response = self.http_client - .post(&url) - .json(&request) + // 5. Execute with retry logic + let mut last_error = None; + for attempt in 0..MAX_RETRIES { + let result = self.execute_request(&url, &request).await; + + match result { + Ok(response) => { + CIRCUIT_BREAKER.record_success(); + return Ok(response); + } + Err(e) => { + last_error = Some(e); + // Don't retry on client errors (4xx) + if let OllamaError::ApiError(_) = &last_error.as_ref().unwrap() { + break; + } + // Exponential backoff + if attempt < MAX_RETRIES - 1 { + tokio::time::sleep(Duration::from_millis(100 * 2_u64.pow(attempt))).await; + } + } + } + } + + // All retries failed + CIRCUIT_BREAKER.record_failure(); + Err(last_error.unwrap_or(OllamaError::MaxRetriesExceeded)) + } + + async fn execute_request( + &self, + url: &str, + request: &GenerateRequest, + ) -> Result { + let mut request_builder = self.http_client + .post(url) + .json(request); + + // Add security headers + let headers = self.build_request_headers(); + request_builder = request_builder.headers(headers); + + let response = request_builder .send() .await .map_err(|e| { @@ -135,11 +398,12 @@ impl OllamaClient { let status = response.status(); let error_text = response.text().await.unwrap_or_else(|_| "Unknown error".to_string()); - if status.as_u16() == 404 { - return Err(OllamaError::ModelNotFound(self.model.clone())); + match status.as_u16() { + 401 => return Err(OllamaError::InvalidApiKey), + 404 => return Err(OllamaError::ModelNotFound(self.model.clone())), + 429 => return Err(OllamaError::RateLimitExceeded("Server rate limit".to_string())), + _ => return Err(OllamaError::ApiError(format!("{}: {}", status, error_text))), } - - return Err(OllamaError::ApiError(format!("{}: {}", status, error_text))); } let result = response.json::() @@ -149,6 +413,30 @@ impl OllamaClient { Ok(result) } + /// Health check endpoint + pub async fn health_check(&self) -> Result { + let url = format!("{}/api/tags", self.base_url); + + let mut request_builder = self.http_client.get(&url); + + // Add security headers + let headers = self.build_request_headers(); + request_builder = request_builder.headers(headers); + + let response = request_builder + .send() + .await + .map_err(|e| { + if e.is_timeout() { + OllamaError::Timeout + } else { + OllamaError::RequestFailed(e) + } + })?; + + Ok(response.status().is_success()) + } + /// Generate a job description based on a prompt pub async fn generate_job_description(&self, prompt: &str) -> Result { let enhanced_prompt = format!( @@ -227,4 +515,53 @@ mod tests { let client = OllamaClient::with_url("http://custom:11434"); assert_eq!(client.get_model(), DEFAULT_MODEL); } + + #[test] + fn test_validate_prompt_too_long() { + let long_prompt = "a".repeat(MAX_PROMPT_LENGTH + 1); + let result = validate_prompt(&long_prompt); + assert!(matches!(result, Err(OllamaError::PromptTooLong(_, _)))); + } + + #[test] + fn test_validate_prompt_ok() { + let prompt = "Hello, world!"; + let result = validate_prompt(prompt); + assert!(result.is_ok()); + assert_eq!(result.unwrap(), prompt); + } + + #[test] + fn test_rate_limiter() { + let limiter = RateLimiter::new(2, 1); // 2 requests per minute + + // First two should succeed + assert!(limiter.check_rate_limit("user1").is_ok()); + assert!(limiter.check_rate_limit("user1").is_ok()); + + // Third should fail + assert!(limiter.check_rate_limit("user1").is_err()); + + // Different user should succeed + assert!(limiter.check_rate_limit("user2").is_ok()); + } + + #[test] + fn test_circuit_breaker() { + let breaker = CircuitBreaker::new(2, 60); + + // Initially should allow + assert!(breaker.can_execute().is_ok()); + + // Record failures + breaker.record_failure(); + breaker.record_failure(); + + // Should now be open + assert!(breaker.can_execute().is_err()); + + // Success should reset + breaker.record_success(); + assert!(breaker.can_execute().is_ok()); + } } diff --git a/crates/contracts/Cargo.toml b/crates/contracts/Cargo.toml index cf53366..8b0cf29 100644 --- a/crates/contracts/Cargo.toml +++ b/crates/contracts/Cargo.toml @@ -13,7 +13,7 @@ chrono = { workspace = true } anyhow = { workspace = true } sqlx = { workspace = true } async-trait = { workspace = true } -jsonwebtoken = "10.3" +jsonwebtoken = { version = "10.3", features = ["rust_crypto"] } db = { path = "../db" } cache = { path = "../cache" } storage = { path = "../storage" } diff --git a/crates/db/Cargo.toml b/crates/db/Cargo.toml index 2e7ce60..0533f10 100644 --- a/crates/db/Cargo.toml +++ b/crates/db/Cargo.toml @@ -11,3 +11,7 @@ tracing = { workspace = true } uuid = { workspace = true } chrono = { workspace = true } anyhow = { workspace = true } +thiserror = { workspace = true } + +[dev-dependencies] +tokio = { workspace = true } diff --git a/crates/db/migrations/20260703210000_ai_credits_wallet.down.sql b/crates/db/migrations/20260703210000_ai_credits_wallet.down.sql new file mode 100644 index 0000000..df8e43e --- /dev/null +++ b/crates/db/migrations/20260703210000_ai_credits_wallet.down.sql @@ -0,0 +1,20 @@ +BEGIN; + +DROP INDEX IF EXISTS idx_ai_reservation_holds_status_expires; +DROP INDEX IF EXISTS idx_ai_reservation_holds_wallet_id; +DROP INDEX IF EXISTS idx_ai_credit_ledger_reference; +DROP INDEX IF EXISTS idx_ai_credit_ledger_wallet_id; +DROP INDEX IF EXISTS idx_ai_usage_logs_created_at; +DROP INDEX IF EXISTS idx_ai_usage_logs_feature_code; +DROP INDEX IF EXISTS idx_ai_usage_logs_user_id; +DROP INDEX IF EXISTS idx_user_ai_subscriptions_plan_id; +DROP INDEX IF EXISTS idx_user_ai_subscriptions_user_id; + +DROP TABLE IF EXISTS ai_reservation_holds; +DROP TABLE IF EXISTS ai_credit_ledger; +DROP TABLE IF EXISTS ai_usage_logs; +DROP TABLE IF EXISTS ai_feature_costs; +DROP TABLE IF EXISTS user_ai_subscriptions; +DROP TABLE IF EXISTS ai_plans; + +COMMIT; diff --git a/crates/db/migrations/20260703210000_ai_credits_wallet.up.sql b/crates/db/migrations/20260703210000_ai_credits_wallet.up.sql new file mode 100644 index 0000000..5231e9a --- /dev/null +++ b/crates/db/migrations/20260703210000_ai_credits_wallet.up.sql @@ -0,0 +1,167 @@ +-- Ask Ash AI Credits: wallet, ledger, plans, feature pricing, usage logs. +-- Phase 1 of docs/ASK_ASH_BILLING_ARCHITECTURE.md (nxtgauge-ai-assistant repo). +-- +-- This is a net-new schema. It does not extend or replace the existing +-- company_ai_usage / job_seeker_ai_usage daily-count quota tables +-- (20260425000000_ai_usage.up.sql) -- those stay in place until a later +-- migration phase consolidates users onto this system (see Section 19 of +-- the architecture doc). It also does not touch tracecoin_wallets / +-- tracecoin_ledger, which remain a completely separate currency. + +BEGIN; + +CREATE TABLE ai_plans ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + code VARCHAR(50) UNIQUE NOT NULL, + name VARCHAR(100) NOT NULL, + monthly_credits INT NOT NULL, + daily_action_limit INT NOT NULL, + daily_credit_limit INT, + allowed_models JSONB NOT NULL DEFAULT '[]', + allowed_features JSONB NOT NULL DEFAULT '[]', + is_trial BOOLEAN NOT NULL DEFAULT FALSE, + trial_days INT, + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- The AI credits wallet. One row per user, mirroring tracecoin_wallets' +-- balance/reserved shape (crates/db/src/models/tracecoin_wallet.rs) but +-- split into pools (monthly/purchased/bonus) so expiry and consumption +-- order can differ per pool -- see Section 4 of the architecture doc. +CREATE TABLE user_ai_subscriptions ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID UNIQUE NOT NULL REFERENCES users(id) ON DELETE CASCADE, + plan_id UUID NOT NULL REFERENCES ai_plans(id), + role_code VARCHAR(50), + + monthly_credits_total INT NOT NULL DEFAULT 0, + monthly_credits_used INT NOT NULL DEFAULT 0, + purchased_credits_total INT NOT NULL DEFAULT 0, + purchased_credits_used INT NOT NULL DEFAULT 0, + bonus_credits_total INT NOT NULL DEFAULT 0, + bonus_credits_used INT NOT NULL DEFAULT 0, + + reserved_credits INT NOT NULL DEFAULT 0, + locked_credits INT NOT NULL DEFAULT 0, + + lifetime_purchased_credits INT NOT NULL DEFAULT 0, + lifetime_used_credits INT NOT NULL DEFAULT 0, + + daily_actions_used INT NOT NULL DEFAULT 0, + daily_credits_used INT NOT NULL DEFAULT 0, + daily_usage_date DATE NOT NULL DEFAULT CURRENT_DATE, + + purchased_credits_expire_at TIMESTAMPTZ, + + billing_cycle VARCHAR(20) NOT NULL DEFAULT 'monthly', + auto_renew BOOLEAN NOT NULL DEFAULT TRUE, + current_period_start TIMESTAMPTZ NOT NULL DEFAULT NOW(), + current_period_end TIMESTAMPTZ NOT NULL DEFAULT (NOW() + INTERVAL '30 days'), + status VARCHAR(30) NOT NULL DEFAULT 'active', + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + CONSTRAINT chk_ai_wallet_nonnegative CHECK ( + monthly_credits_used >= 0 AND monthly_credits_used <= monthly_credits_total + AND purchased_credits_used >= 0 AND purchased_credits_used <= purchased_credits_total + AND bonus_credits_used >= 0 AND bonus_credits_used <= bonus_credits_total + AND reserved_credits >= 0 + AND locked_credits >= 0 + ) +); + +CREATE TABLE ai_feature_costs ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + feature_code VARCHAR(100) UNIQUE NOT NULL, + display_name VARCHAR(150) NOT NULL, + default_model VARCHAR(100) NOT NULL, + credit_cost INT NOT NULL, + max_input_tokens INT, + max_output_tokens INT, + min_plan_code VARCHAR(50), + priority VARCHAR(20) NOT NULL DEFAULT 'normal', + timeout_ms INT NOT NULL DEFAULT 15000, + rate_limit_per_minute INT, + retry_policy JSONB NOT NULL DEFAULT '{"max_retries": 1, "backoff_ms": 500}', + fallback_model VARCHAR(100), + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE TABLE ai_usage_logs ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id), + role_code VARCHAR(50), + feature_code VARCHAR(100) NOT NULL, + model_alias VARCHAR(100) NOT NULL, + credits_charged INT NOT NULL, + input_tokens INT, + output_tokens INT, + total_tokens INT, + cached_tokens INT, + provider_reported_cost NUMERIC(12, 6), + internal_cost NUMERIC(12, 6), + credit_unit_price_at_time NUMERIC(12, 6), + status VARCHAR(30) NOT NULL, + request_id VARCHAR(100), + error_message TEXT, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Immutable append-only ledger -- the source of truth for wallet balances +-- (Section 5 of the architecture doc). Application code must never UPDATE +-- or DELETE rows here; corrections are new rows referencing the row they +-- correct via reference_type='ledger_entry'/reference_id. +CREATE TABLE ai_credit_ledger ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + wallet_id UUID NOT NULL REFERENCES user_ai_subscriptions(id), + entry_type VARCHAR(40) NOT NULL, + credits INT NOT NULL, + balance_after INT NOT NULL, + idempotency_key VARCHAR(150) UNIQUE, + reference_type VARCHAR(50), + reference_id UUID, + actor_type VARCHAR(20) NOT NULL DEFAULT 'user', + actor_id UUID, + metadata JSONB, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Backs the reserve -> capture/release flow (Section 4.2). A background +-- reaper (Phase 1 follow-up, not yet implemented) releases `held` rows +-- past expires_at back to the wallet's available balance. +CREATE TABLE ai_reservation_holds ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + wallet_id UUID NOT NULL REFERENCES user_ai_subscriptions(id), + credits_held INT NOT NULL, + feature_code VARCHAR(100) NOT NULL, + request_id VARCHAR(150), + idempotency_key VARCHAR(150) UNIQUE, + status VARCHAR(20) NOT NULL DEFAULT 'held', + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + expires_at TIMESTAMPTZ NOT NULL DEFAULT (NOW() + INTERVAL '5 minutes'), + resolved_at TIMESTAMPTZ +); + +CREATE INDEX idx_user_ai_subscriptions_user_id ON user_ai_subscriptions(user_id); +CREATE INDEX idx_user_ai_subscriptions_plan_id ON user_ai_subscriptions(plan_id); +CREATE INDEX idx_ai_usage_logs_user_id ON ai_usage_logs(user_id); +CREATE INDEX idx_ai_usage_logs_feature_code ON ai_usage_logs(feature_code); +CREATE INDEX idx_ai_usage_logs_created_at ON ai_usage_logs(created_at); +CREATE INDEX idx_ai_credit_ledger_wallet_id ON ai_credit_ledger(wallet_id, created_at); +CREATE INDEX idx_ai_credit_ledger_reference ON ai_credit_ledger(reference_type, reference_id); +CREATE INDEX idx_ai_reservation_holds_wallet_id ON ai_reservation_holds(wallet_id); +CREATE INDEX idx_ai_reservation_holds_status_expires ON ai_reservation_holds(status, expires_at) WHERE status = 'held'; + +-- Seed a minimal Free plan so ensure_wallet() always has a plan to assign +-- new users to. Feature costs / higher tiers are a Phase 3 (packages/ +-- plans admin) concern, not seeded here to avoid inventing real prices in +-- a migration -- see Appendix Q4 of the architecture doc. +INSERT INTO ai_plans (code, name, monthly_credits, daily_action_limit, daily_credit_limit, allowed_models, allowed_features) +VALUES ('free', 'Free', 10, 3, 10, '["askash-fast"]', '[]'); + +COMMIT; diff --git a/crates/db/migrations/20260703220000_ai_credits_seed_features.down.sql b/crates/db/migrations/20260703220000_ai_credits_seed_features.down.sql new file mode 100644 index 0000000..2a6eca3 --- /dev/null +++ b/crates/db/migrations/20260703220000_ai_credits_seed_features.down.sql @@ -0,0 +1,8 @@ +BEGIN; + +UPDATE ai_plans SET allowed_features = '[]' WHERE code = 'free'; + +DELETE FROM ai_feature_costs +WHERE feature_code IN ('help_answer', 'jd_generate', 'cover_letter_generate', 'resume_improve'); + +COMMIT; diff --git a/crates/db/migrations/20260703220000_ai_credits_seed_features.up.sql b/crates/db/migrations/20260703220000_ai_credits_seed_features.up.sql new file mode 100644 index 0000000..3b6be68 --- /dev/null +++ b/crates/db/migrations/20260703220000_ai_credits_seed_features.up.sql @@ -0,0 +1,20 @@ +-- A first batch of real feature costs so the wallet primitives (Phase 1) +-- have something concrete to charge against once wired into handlers. +-- Model aliases match the real LiteLLM deployment +-- (nxtgauge-gitops/apps/litellm/base/configmap.yaml): askash-fast / +-- askash-main. Costs are illustrative starting points, not a pricing +-- decision -- see Appendix Q4 of the architecture doc. + +BEGIN; + +INSERT INTO ai_feature_costs (feature_code, display_name, default_model, credit_cost, timeout_ms) VALUES +('help_answer', 'Help Answer', 'askash-fast', 1, 10000), +('jd_generate', 'Job Description Generate', 'askash-main', 5, 20000), +('cover_letter_generate', 'Cover Letter Generate', 'askash-main', 5, 20000), +('resume_improve', 'Resume Improve', 'askash-main', 8, 20000); + +UPDATE ai_plans +SET allowed_features = '["help_answer", "jd_generate"]' +WHERE code = 'free'; + +COMMIT; diff --git a/crates/db/migrations/20260705110000_ai_credit_packages.down.sql b/crates/db/migrations/20260705110000_ai_credit_packages.down.sql new file mode 100644 index 0000000..5e88a11 --- /dev/null +++ b/crates/db/migrations/20260705110000_ai_credit_packages.down.sql @@ -0,0 +1,10 @@ +BEGIN; + +DROP INDEX IF EXISTS idx_ai_credit_orders_txnid; +DROP INDEX IF EXISTS idx_ai_credit_orders_user_id; +DROP TABLE IF EXISTS ai_credit_orders; + +DROP INDEX IF EXISTS idx_ai_credit_packages_active; +DROP TABLE IF EXISTS ai_credit_packages; + +COMMIT; diff --git a/crates/db/migrations/20260705110000_ai_credit_packages.up.sql b/crates/db/migrations/20260705110000_ai_credit_packages.up.sql new file mode 100644 index 0000000..9190d95 --- /dev/null +++ b/crates/db/migrations/20260705110000_ai_credit_packages.up.sql @@ -0,0 +1,48 @@ +-- AI credit purchase packages + PayU order tracking, matching the +-- already-shipped frontend contract in +-- nxtgauge-frontend-solid/src/components/dashboard/CreditsPage.tsx +-- (GET /api/ai-credits, POST /api/ai-credits/order, POST /api/ai-credits/verify). +-- price_inr is stored in paise (matches pricing_packages.price_inr's +-- existing convention in this codebase), not rupees. + +BEGIN; + +CREATE TABLE ai_credit_packages ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + name VARCHAR(150) NOT NULL, + description TEXT, + credits INT NOT NULL, + price_inr INT NOT NULL, -- paise + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_ai_credit_packages_active ON ai_credit_packages(is_active, price_inr); + +INSERT INTO ai_credit_packages (name, description, credits, price_inr) VALUES +('Starter AI Credits', '50 AI credits for casual usage', 50, 9900), +('Pro AI Credits', '200 AI credits for power users', 200, 34900), +('Business AI Credits', '750 AI credits for teams', 750, 99900), +('Enterprise AI Credits', '2500 AI credits for heavy usage', 2500, 249900); + +-- PayU order tracking for AI credit purchases, separate from the +-- TraceCoins `payments` table (crates/db/migrations/20260317190300_portfolio_payments.up.sql) +-- per the earlier decision to keep AI credits fully separate from TraceCoins. +CREATE TABLE ai_credit_orders ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id), + package_id UUID NOT NULL REFERENCES ai_credit_packages(id), + txnid VARCHAR(100) UNIQUE NOT NULL, + payu_payment_id VARCHAR(100), + amount_inr INT NOT NULL, -- paise + credits INT NOT NULL, + status VARCHAR(20) NOT NULL DEFAULT 'PENDING', -- PENDING, SUCCESS, FAILED + verified_at TIMESTAMPTZ, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +CREATE INDEX idx_ai_credit_orders_user_id ON ai_credit_orders(user_id); +CREATE INDEX idx_ai_credit_orders_txnid ON ai_credit_orders(txnid); + +COMMIT; diff --git a/crates/db/migrations/20260706100000_ai_usage_logs_preview.up.sql b/crates/db/migrations/20260706100000_ai_usage_logs_preview.up.sql new file mode 100644 index 0000000..eb65e0c --- /dev/null +++ b/crates/db/migrations/20260706100000_ai_usage_logs_preview.up.sql @@ -0,0 +1,12 @@ +-- Add prompt_preview and response_preview columns to ai_usage_logs +-- for audit logging (Task 3 - Ollama Security Hardening) + +BEGIN; + +ALTER TABLE ai_usage_logs + ADD COLUMN IF NOT EXISTS prompt_preview VARCHAR(200), + ADD COLUMN IF NOT EXISTS response_preview VARCHAR(200); + +CREATE INDEX IF NOT EXISTS idx_ai_usage_logs_prompt_preview ON ai_usage_logs(prompt_preview) WHERE prompt_preview IS NOT NULL; + +COMMIT; diff --git a/crates/db/migrations/20260706200000_ai_refunds.up.sql b/crates/db/migrations/20260706200000_ai_refunds.up.sql new file mode 100644 index 0000000..f1334e9 --- /dev/null +++ b/crates/db/migrations/20260706200000_ai_refunds.up.sql @@ -0,0 +1,32 @@ +-- Refund architecture for AI credits (Task 5) +-- Tracks refunds for failed or disputed AI credit charges + +BEGIN; + +CREATE TABLE ai_refunds ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + ledger_debit_entry_id UUID NOT NULL REFERENCES ai_credit_ledger(id), + credits_refunded INT NOT NULL CHECK (credits_refunded > 0), + reason VARCHAR(50) NOT NULL CHECK (reason IN ('provider_failure', 'timeout', 'validation_failure', 'manual', 'user_dispute')), + initiated_by VARCHAR(20) NOT NULL CHECK (initiated_by IN ('system', 'admin', 'user_dispute')), + status VARCHAR(20) NOT NULL DEFAULT 'pending' CHECK (status IN ('pending', 'approved', 'rejected', 'completed')), + notes TEXT, + admin_actor_id UUID REFERENCES users(id), + idempotency_key VARCHAR(150) UNIQUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + resolved_at TIMESTAMPTZ, + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Indexes for common queries +CREATE INDEX idx_ai_refunds_user_id ON ai_refunds(user_id); +CREATE INDEX idx_ai_refunds_status ON ai_refunds(status); +CREATE INDEX idx_ai_refunds_ledger_entry ON ai_refunds(ledger_debit_entry_id); +CREATE INDEX idx_ai_refunds_created_at ON ai_refunds(created_at DESC); + +-- Add entry_type for refunds to ai_credit_ledger (if not exists) +-- The existing entry_type column should support 'refund' values +-- No schema change needed - just ensure application uses 'refund' entry_type + +COMMIT; diff --git a/crates/db/migrations/20260706300000_ai_coupons_promotions.up.sql b/crates/db/migrations/20260706300000_ai_coupons_promotions.up.sql new file mode 100644 index 0000000..8f198df --- /dev/null +++ b/crates/db/migrations/20260706300000_ai_coupons_promotions.up.sql @@ -0,0 +1,160 @@ +-- Coupons, promotions, referral, and bonus credits (Task 6) +-- Implements coupon redemption, promotions, referral tracking + +BEGIN; + +-- Coupons table for discount codes +CREATE TABLE ai_coupons ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + code VARCHAR(100) UNIQUE NOT NULL, + name VARCHAR(200) NOT NULL, + description TEXT, + + -- Discount configuration + discount_type VARCHAR(20) NOT NULL CHECK (discount_type IN ('percentage', 'fixed_amount')), + discount_value NUMERIC(10, 2) NOT NULL CHECK (discount_value >= 0), + + -- Applicability + applicable_package_ids UUID[] DEFAULT '{}', + applicable_plan_codes TEXT[] DEFAULT '{}', + min_purchase_amount NUMERIC(12, 2), + max_discount_amount NUMERIC(12, 2), + + -- Limits + max_redemptions INT NOT NULL DEFAULT 1, + redemptions_used INT NOT NULL DEFAULT 0, + max_redemptions_per_user INT DEFAULT 1, + + -- Validity + valid_from TIMESTAMPTZ NOT NULL DEFAULT NOW(), + valid_until TIMESTAMPTZ, + + -- Status + is_active BOOLEAN NOT NULL DEFAULT TRUE, + + -- Audit + created_by UUID REFERENCES users(id), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Coupon redemptions (tracks who used which coupon) +CREATE TABLE ai_coupon_redemptions ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + coupon_id UUID NOT NULL REFERENCES ai_coupons(id) ON DELETE CASCADE, + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + order_id UUID, -- Reference to ai_credit_orders + credits_purchased INT NOT NULL, + discount_applied NUMERIC(12, 2) NOT NULL, + final_amount_paid NUMERIC(12, 2) NOT NULL, + redeemed_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(coupon_id, user_id) +); + +-- Promotions table for automatic bonuses +CREATE TABLE ai_promotions ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + name VARCHAR(200) NOT NULL, + description TEXT, + + -- Trigger conditions + trigger_type VARCHAR(50) NOT NULL CHECK (trigger_type IN ('signup', 'first_purchase', 'referral', 'milestone', 'manual')), + trigger_condition JSONB DEFAULT '{}', -- e.g., {"min_purchase": 100, "referrer_bonus": 50} + + -- Reward + bonus_credits INT NOT NULL CHECK (bonus_credits > 0), + bonus_expires_after_days INT, -- NULL means no expiry + + -- Applicability + applicable_plan_codes TEXT[] DEFAULT '{}', + applicable_user_types TEXT[] DEFAULT '{}', -- e.g., ['NEW_USER', 'REFERRER'] + + -- Limits + max_grants INT, + grants_used INT NOT NULL DEFAULT 0, + max_grants_per_user INT DEFAULT 1, + + -- Validity + valid_from TIMESTAMPTZ NOT NULL DEFAULT NOW(), + valid_until TIMESTAMPTZ, + + -- Status + is_active BOOLEAN NOT NULL DEFAULT TRUE, + + -- Audit + created_by UUID REFERENCES users(id), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Promotion grants (tracks who received which promotion) +CREATE TABLE ai_promotion_grants ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + promotion_id UUID NOT NULL REFERENCES ai_promotions(id) ON DELETE CASCADE, + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + ledger_entry_id UUID REFERENCES ai_credit_ledger(id), -- Link to the bonus credit ledger entry + bonus_credits INT NOT NULL, + expires_at TIMESTAMPTZ, + granted_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(promotion_id, user_id) +); + +-- Bonus credits ledger reference +-- This extends ai_credit_ledger with source tracking for bonus credits +CREATE TABLE ai_bonus_credits ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + source_type VARCHAR(30) NOT NULL CHECK (source_type IN ('promotion', 'coupon_bonus', 'referral', 'manual', 'signup_bonus')), + source_reference_id UUID, -- References ai_promotions, ai_coupon_redemptions, etc. + credits INT NOT NULL CHECK (credits > 0), + credits_used INT NOT NULL DEFAULT 0 CHECK (credits_used >= 0 AND credits_used <= credits), + expires_at TIMESTAMPTZ, + ledger_entry_id UUID REFERENCES ai_credit_ledger(id), + is_active BOOLEAN NOT NULL DEFAULT TRUE, + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Referral tracking +CREATE TABLE ai_referrals ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + referrer_user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + referred_user_id UUID REFERENCES users(id) ON DELETE SET NULL, + referral_code VARCHAR(100) NOT NULL UNIQUE, + + -- Status + status VARCHAR(20) NOT NULL DEFAULT 'pending' CHECK (status IN ('pending', 'converted', 'expired')), + + -- Rewards (granted when referred user completes first purchase) + referrer_bonus_credits INT, + referred_bonus_credits INT, + + -- Conversion tracking + converted_at TIMESTAMPTZ, + first_purchase_order_id UUID REFERENCES ai_credit_orders(id), + + -- Expiry + expires_at TIMESTAMPTZ NOT NULL DEFAULT (NOW() + INTERVAL '30 days'), + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Indexes for performance +CREATE INDEX idx_ai_coupons_code ON ai_coupons(code) WHERE is_active = TRUE; +CREATE INDEX idx_ai_coupons_valid ON ai_coupons(valid_from, valid_until) WHERE is_active = TRUE; +CREATE INDEX idx_ai_coupon_redemptions_user ON ai_coupon_redemptions(user_id); +CREATE INDEX idx_ai_coupon_redemptions_coupon ON ai_coupon_redemptions(coupon_id); + +CREATE INDEX idx_ai_promotions_trigger ON ai_promotions(trigger_type) WHERE is_active = TRUE; +CREATE INDEX idx_ai_promotions_valid ON ai_promotions(valid_from, valid_until) WHERE is_active = TRUE; +CREATE INDEX idx_ai_promotion_grants_user ON ai_promotion_grants(user_id); + +CREATE INDEX idx_ai_bonus_credits_user ON ai_bonus_credits(user_id) WHERE is_active = TRUE; +CREATE INDEX idx_ai_bonus_credits_expires ON ai_bonus_credits(expires_at) WHERE expires_at IS NOT NULL AND is_active = TRUE; + +CREATE INDEX idx_ai_referrals_referrer ON ai_referrals(referrer_user_id); +CREATE INDEX idx_ai_referrals_code ON ai_referrals(referral_code); +CREATE INDEX idx_ai_referrals_referred ON ai_referrals(referred_user_id) WHERE referred_user_id IS NOT NULL; + +COMMIT; diff --git a/crates/db/migrations/20260706400000_ai_subscription_lifecycle.up.sql b/crates/db/migrations/20260706400000_ai_subscription_lifecycle.up.sql new file mode 100644 index 0000000..ea558e1 --- /dev/null +++ b/crates/db/migrations/20260706400000_ai_subscription_lifecycle.up.sql @@ -0,0 +1,66 @@ +-- Subscription lifecycle management (Task 7) +-- Tracks plan changes, proration, and scheduled downgrades + +BEGIN; + +CREATE TABLE ai_subscription_history ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + user_id UUID NOT NULL REFERENCES users(id) ON DELETE CASCADE, + from_plan_id UUID REFERENCES ai_plans(id), + to_plan_id UUID NOT NULL REFERENCES ai_plans(id), + change_type VARCHAR(20) NOT NULL CHECK (change_type IN ('upgrade', 'downgrade', 'cancel', 'renew', 'trial_expired')), + + -- Proration details (for upgrades) + proration_credits INT DEFAULT 0, -- Bonus credits granted for mid-period upgrade + proration_days_remaining INT, -- Days remaining in current period + + -- Timing + effective_at TIMESTAMPTZ NOT NULL, + created_by UUID REFERENCES users(id), -- System if automated + + -- Status + status VARCHAR(20) NOT NULL DEFAULT 'completed' CHECK (status IN ('scheduled', 'completed', 'cancelled')), + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW() +); + +-- Add columns to user_ai_subscriptions for trial support (if not exists) +-- These are already in the schema from the initial migration +-- - is_trial BOOLEAN +-- - trial_days INT +-- - trial_ends_at TIMESTAMPTZ (needs to be added) + +-- Add trial tracking columns if they don't exist +DO $$ +BEGIN + IF NOT EXISTS (SELECT 1 FROM information_schema.columns + WHERE table_name = 'user_ai_subscriptions' + AND column_name = 'is_trial') THEN + ALTER TABLE user_ai_subscriptions ADD COLUMN is_trial BOOLEAN NOT NULL DEFAULT FALSE; + END IF; + + IF NOT EXISTS (SELECT 1 FROM information_schema.columns + WHERE table_name = 'user_ai_subscriptions' + AND column_name = 'trial_days') THEN + ALTER TABLE user_ai_subscriptions ADD COLUMN trial_days INT; + END IF; + + IF NOT EXISTS (SELECT 1 FROM information_schema.columns + WHERE table_name = 'user_ai_subscriptions' + AND column_name = 'trial_ends_at') THEN + ALTER TABLE user_ai_subscriptions ADD COLUMN trial_ends_at TIMESTAMPTZ; + END IF; + + IF NOT EXISTS (SELECT 1 FROM information_schema.columns + WHERE table_name = 'user_ai_subscriptions' + AND column_name = 'downgrade_scheduled_to') THEN + ALTER TABLE user_ai_subscriptions ADD COLUMN downgrade_scheduled_to UUID REFERENCES ai_plans(id); + END IF; +END $$; + +-- Indexes for subscription management +CREATE INDEX idx_ai_subscription_history_user ON ai_subscription_history(user_id); +CREATE INDEX idx_ai_subscription_history_effective ON ai_subscription_history(effective_at DESC); +CREATE INDEX idx_user_ai_subscriptions_trial ON user_ai_subscriptions(user_id) WHERE is_trial = TRUE; + +COMMIT; diff --git a/crates/db/migrations/20260706500000_ai_token_cost_engine.up.sql b/crates/db/migrations/20260706500000_ai_token_cost_engine.up.sql new file mode 100644 index 0000000..efdf1b0 --- /dev/null +++ b/crates/db/migrations/20260706500000_ai_token_cost_engine.up.sql @@ -0,0 +1,119 @@ +-- Token cost engine and provider cost tracking (Task 9) +-- Configures model costs and tracks actual usage costs for margin analysis + +BEGIN; + +-- Model cost configuration table +CREATE TABLE ai_model_cost_config ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + model_alias VARCHAR(100) NOT NULL, + + -- Cost per 1K tokens + cost_per_1k_input_tokens NUMERIC(12, 8) NOT NULL, + cost_per_1k_output_tokens NUMERIC(12, 8) NOT NULL, + + -- Cost basis: compute_amortized (self-hosted) or provider_metered (paid APIs) + cost_basis VARCHAR(30) NOT NULL CHECK (cost_basis IN ('compute_amortized', 'provider_metered')), + + -- For compute_amortized: optional metadata about calculation + cost_calculation_notes TEXT, + + -- Effective date range + effective_from TIMESTAMPTZ NOT NULL DEFAULT NOW(), + effective_until TIMESTAMPTZ, + + -- Status + is_active BOOLEAN NOT NULL DEFAULT TRUE, + + -- Audit + created_by UUID REFERENCES users(id), + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(model_alias, effective_from) +); + +-- Add cost tracking columns to ai_usage_logs if not exist +DO $$ +BEGIN + -- These columns already exist per the initial migration, but verify they're populated + IF NOT EXISTS (SELECT 1 FROM information_schema.columns + WHERE table_name = 'ai_usage_logs' + AND column_name = 'input_tokens') THEN + ALTER TABLE ai_usage_logs ADD COLUMN input_tokens INT; + END IF; + + IF NOT EXISTS (SELECT 1 FROM information_schema.columns + WHERE table_name = 'ai_usage_logs' + AND column_name = 'output_tokens') THEN + ALTER TABLE ai_usage_logs ADD COLUMN output_tokens INT; + END IF; + + IF NOT EXISTS (SELECT 1 FROM information_schema.columns + WHERE table_name = 'ai_usage_logs' + AND column_name = 'total_tokens') THEN + ALTER TABLE ai_usage_logs ADD COLUMN total_tokens INT; + END IF; +END $$; + +-- Create or replace view for margin calculation +-- This provides real-time margin analysis without storing computed values +CREATE OR REPLACE VIEW ai_usage_margin_view AS +SELECT + l.id, + l.user_id, + l.feature_code, + l.model_alias, + l.credits_charged, + l.input_tokens, + l.output_tokens, + l.total_tokens, + l.status, + l.created_at, + + -- Calculate internal cost using current cost config + CASE + WHEN l.input_tokens IS NOT NULL AND l.output_tokens IS NOT NULL THEN + (l.input_tokens::NUMERIC / 1000.0 * COALESCE(c.cost_per_1k_input_tokens, 0)) + + (l.output_tokens::NUMERIC / 1000.0 * COALESCE(c.cost_per_1k_output_tokens, 0)) + ELSE NULL + END as calculated_internal_cost, + + -- Cost basis + c.cost_basis, + + -- Margin (credits charged are in some credit unit, convert to currency for margin calc) + -- Note: This assumes credits have a monetary value; adjust conversion rate as needed + CASE + WHEN l.input_tokens IS NOT NULL AND l.output_tokens IS NOT NULL THEN + l.credits_charged - ( + (l.input_tokens::NUMERIC / 1000.0 * COALESCE(c.cost_per_1k_input_tokens, 0)) + + (l.output_tokens::NUMERIC / 1000.0 * COALESCE(c.cost_per_1k_output_tokens, 0)) + ) + ELSE NULL + END as margin_credits + +FROM ai_usage_logs l +LEFT JOIN ai_model_cost_config c ON l.model_alias = c.model_alias + AND c.is_active = TRUE + AND c.effective_from <= l.created_at + AND (c.effective_until IS NULL OR c.effective_until > l.created_at) +WHERE l.status = 'success'; + +-- Indexes +CREATE INDEX idx_ai_model_cost_config_alias ON ai_model_cost_config(model_alias) WHERE is_active = TRUE; +CREATE INDEX idx_ai_usage_logs_tokens ON ai_usage_logs(user_id, created_at DESC) + WHERE input_tokens IS NOT NULL OR output_tokens IS NOT NULL; + +-- Seed initial cost config for self-hosted Ollama models (compute amortized) +-- These are example values - adjust based on actual GPU costs and throughput +INSERT INTO ai_model_cost_config (model_alias, cost_per_1k_input_tokens, cost_per_1k_output_tokens, cost_basis, cost_calculation_notes) +VALUES + ('askash-fast', 0.0001, 0.0002, 'compute_amortized', 'qwen3:4b on shared GPU - cost from GPU node amortization'), + ('askash-main', 0.0002, 0.0004, 'compute_amortized', 'qwen3:8b on shared GPU - cost from GPU node amortization'), + ('ultra-fast', 0.00005, 0.0001, 'compute_amortized', 'gemma3:270m on shared GPU - cost from GPU node amortization'), + ('jd-generator', 0.0002, 0.0004, 'compute_amortized', 'qwen3:8b - specialized for job descriptions'), + ('profile-writer', 0.0002, 0.0004, 'compute_amortized', 'qwen3:8b - specialized for profile writing') +ON CONFLICT (model_alias, effective_from) DO NOTHING; + +COMMIT; diff --git a/crates/db/migrations/20260706600000_ai_observability.up.sql b/crates/db/migrations/20260706600000_ai_observability.up.sql new file mode 100644 index 0000000..39ce969 --- /dev/null +++ b/crates/db/migrations/20260706600000_ai_observability.up.sql @@ -0,0 +1,166 @@ +-- Observability: AI credits metrics tables (Task 10) +-- Stores aggregated metrics for dashboard and monitoring + +BEGIN; + +-- Daily aggregated metrics for AI credits usage +CREATE TABLE ai_metrics_daily ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + metric_date DATE NOT NULL, + + -- Usage metrics + total_requests INT NOT NULL DEFAULT 0, + total_credits_charged INT NOT NULL DEFAULT 0, + total_credits_refunded INT NOT NULL DEFAULT 0, + + -- Token metrics + total_input_tokens BIGINT NOT NULL DEFAULT 0, + total_output_tokens BIGINT NOT NULL DEFAULT 0, + total_tokens BIGINT NOT NULL DEFAULT 0, + + -- Financial metrics + total_revenue NUMERIC(12, 2) NOT NULL DEFAULT 0, -- From purchases + estimated_cost NUMERIC(12, 6) NOT NULL DEFAULT 0, -- Computed from tokens + estimated_margin NUMERIC(12, 2) NOT NULL DEFAULT 0, -- revenue - cost + + -- Error metrics + failed_requests INT NOT NULL DEFAULT 0, + timeout_requests INT NOT NULL DEFAULT 0, + + -- Unique users + unique_active_users INT NOT NULL DEFAULT 0, + new_wallet_creations INT NOT NULL DEFAULT 0, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(metric_date) +); + +-- Hourly metrics for more granular monitoring +CREATE TABLE ai_metrics_hourly ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + metric_hour TIMESTAMPTZ NOT NULL, -- Truncated to hour + + total_requests INT NOT NULL DEFAULT 0, + total_credits_charged INT NOT NULL DEFAULT 0, + avg_response_time_ms INT, -- Average response time in milliseconds + error_count INT NOT NULL DEFAULT 0, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(metric_hour) +); + +-- Feature usage breakdown +CREATE TABLE ai_metrics_by_feature ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + metric_date DATE NOT NULL, + feature_code VARCHAR(100) NOT NULL, + + request_count INT NOT NULL DEFAULT 0, + credits_charged INT NOT NULL DEFAULT 0, + unique_users INT NOT NULL DEFAULT 0, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(metric_date, feature_code) +); + +-- Model usage breakdown +CREATE TABLE ai_metrics_by_model ( + id UUID PRIMARY KEY DEFAULT gen_random_uuid(), + metric_date DATE NOT NULL, + model_alias VARCHAR(100) NOT NULL, + + request_count INT NOT NULL DEFAULT 0, + input_tokens BIGINT NOT NULL DEFAULT 0, + output_tokens BIGINT NOT NULL DEFAULT 0, + estimated_cost NUMERIC(12, 6) NOT NULL DEFAULT 0, + + created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(), + + UNIQUE(metric_date, model_alias) +); + +-- Indexes for metrics queries +CREATE INDEX idx_ai_metrics_daily_date ON ai_metrics_daily(metric_date DESC); +CREATE INDEX idx_ai_metrics_hourly_hour ON ai_metrics_hourly(metric_hour DESC); +CREATE INDEX idx_ai_metrics_feature ON ai_metrics_by_feature(metric_date DESC, feature_code); +CREATE INDEX idx_ai_metrics_model ON ai_metrics_by_model(metric_date DESC, model_alias); + +-- Function to aggregate daily metrics +-- This should be called by a cron job daily +CREATE OR REPLACE FUNCTION aggregate_ai_metrics_daily(target_date DATE) +RETURNS VOID AS $$ +BEGIN + -- Insert or update daily aggregate + INSERT INTO ai_metrics_daily ( + metric_date, + total_requests, + total_credits_charged, + total_credits_refunded, + total_input_tokens, + total_output_tokens, + total_tokens, + failed_requests, + unique_active_users + ) + SELECT + target_date, + COUNT(*), + SUM(CASE WHEN status = 'success' THEN credits_charged ELSE 0 END), + 0, -- Refunds handled separately + SUM(COALESCE(input_tokens, 0)), + SUM(COALESCE(output_tokens, 0)), + SUM(COALESCE(total_tokens, 0)), + COUNT(CASE WHEN status = 'error' THEN 1 END), + COUNT(DISTINCT user_id) + FROM ai_usage_logs + WHERE DATE(created_at) = target_date + ON CONFLICT (metric_date) DO UPDATE SET + total_requests = EXCLUDED.total_requests, + total_credits_charged = EXCLUDED.total_credits_charged, + total_input_tokens = EXCLUDED.total_input_tokens, + total_output_tokens = EXCLUDED.total_output_tokens, + total_tokens = EXCLUDED.total_tokens, + failed_requests = EXCLUDED.failed_requests, + unique_active_users = EXCLUDED.unique_active_users, + updated_at = NOW(); + + -- Aggregate by feature + INSERT INTO ai_metrics_by_feature (metric_date, feature_code, request_count, credits_charged, unique_users) + SELECT + target_date, + feature_code, + COUNT(*), + SUM(credits_charged), + COUNT(DISTINCT user_id) + FROM ai_usage_logs + WHERE DATE(created_at) = target_date + GROUP BY feature_code + ON CONFLICT (metric_date, feature_code) DO UPDATE SET + request_count = EXCLUDED.request_count, + credits_charged = EXCLUDED.credits_charged, + unique_users = EXCLUDED.unique_users; + + -- Aggregate by model + INSERT INTO ai_metrics_by_model (metric_date, model_alias, request_count, input_tokens, output_tokens) + SELECT + target_date, + model_alias, + COUNT(*), + SUM(COALESCE(input_tokens, 0)), + SUM(COALESCE(output_tokens, 0)) + FROM ai_usage_logs + WHERE DATE(created_at) = target_date + GROUP BY model_alias + ON CONFLICT (metric_date, model_alias) DO UPDATE SET + request_count = EXCLUDED.request_count, + input_tokens = EXCLUDED.input_tokens, + output_tokens = EXCLUDED.output_tokens; + +END; +$$ LANGUAGE plpgsql; + +COMMIT; diff --git a/crates/db/src/models/ai_credits.rs b/crates/db/src/models/ai_credits.rs new file mode 100644 index 0000000..6708aff --- /dev/null +++ b/crates/db/src/models/ai_credits.rs @@ -0,0 +1,888 @@ +//! Ask Ash AI credits wallet + ledger. +//! +//! Phase 1 of docs/ASK_ASH_BILLING_ARCHITECTURE.md (nxtgauge-ai-assistant +//! repo). Mirrors the reserve/capture/release + row-locked-transaction +//! pattern already proven by `tracecoin_wallet.rs` for TraceCoins -- see +//! that file for the pattern this one is deliberately consistent with. +//! This is a completely separate currency/schema from TraceCoins; nothing +//! here should ever read or write `tracecoin_wallets`/`tracecoin_ledger`. + +use chrono::{DateTime, NaiveDate, Utc}; +use serde::{Deserialize, Serialize}; +use sqlx::{FromRow, PgPool, Postgres, Transaction}; +use uuid::Uuid; + +#[derive(Debug, Serialize, Deserialize, FromRow, Clone)] +pub struct AiWallet { + pub id: Uuid, + pub user_id: Uuid, + pub plan_id: Uuid, + pub role_code: Option, + pub monthly_credits_total: i32, + pub monthly_credits_used: i32, + pub purchased_credits_total: i32, + pub purchased_credits_used: i32, + pub bonus_credits_total: i32, + pub bonus_credits_used: i32, + pub reserved_credits: i32, + pub locked_credits: i32, + pub lifetime_purchased_credits: i32, + pub lifetime_used_credits: i32, + pub daily_actions_used: i32, + pub daily_credits_used: i32, + pub daily_usage_date: NaiveDate, + pub purchased_credits_expire_at: Option>, + pub billing_cycle: String, + pub auto_renew: bool, + pub current_period_start: DateTime, + pub current_period_end: DateTime, + pub status: String, + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl AiWallet { + /// Credits actually spendable right now. Reserved/locked credits are + /// excluded even though they're still "owned" -- see Section 4.1 of + /// the architecture doc for the full pool breakdown. + pub fn available_credits(&self) -> i32 { + let owned = (self.monthly_credits_total - self.monthly_credits_used) + + (self.purchased_credits_total - self.purchased_credits_used) + + (self.bonus_credits_total - self.bonus_credits_used); + (owned - self.reserved_credits - self.locked_credits).max(0) + } +} + +#[derive(Debug, Serialize, Deserialize, FromRow, Clone)] +pub struct AiPlan { + pub id: Uuid, + pub code: String, + pub name: String, + pub monthly_credits: i32, + pub daily_action_limit: i32, + pub daily_credit_limit: Option, + pub allowed_models: serde_json::Value, + pub allowed_features: serde_json::Value, + pub is_active: bool, +} + +#[derive(Debug, Serialize, Deserialize, FromRow, Clone)] +pub struct AiFeatureCost { + pub id: Uuid, + pub feature_code: String, + pub display_name: String, + pub default_model: String, + pub credit_cost: i32, + pub max_input_tokens: Option, + pub max_output_tokens: Option, + pub min_plan_code: Option, + pub timeout_ms: i32, + pub fallback_model: Option, + pub is_active: bool, +} + +#[derive(Debug, Serialize, Deserialize, FromRow, Clone)] +pub struct AiReservationHold { + pub id: Uuid, + pub wallet_id: Uuid, + pub credits_held: i32, + pub feature_code: String, + pub status: String, +} + +#[derive(Debug, thiserror::Error)] +pub enum AiCreditsError { + #[error("Database error: {0}")] + Db(#[from] sqlx::Error), + #[error("Unknown AI feature '{0}'")] + UnknownFeature(String), + #[error("Insufficient AI credits")] + InsufficientCredits, + #[error("Daily AI action limit reached")] + DailyActionLimitReached, + #[error("Daily AI credit limit reached")] + DailyCreditLimitReached, + #[error("Reservation not found or already resolved")] + ReservationNotFound, + #[error("Invalid amount: {0}")] + InvalidAmount(String), +} + +pub struct AiCreditsRepository; + +impl AiCreditsRepository { + /// Get-or-create the wallet for a user, defaulting to the `free` plan. + /// Uses the same `INSERT ... ON CONFLICT DO NOTHING` + re-fetch shape + /// as `TracecoinWalletRepository::ensure_wallet` so a race between two + /// concurrent first-time callers can't surface a raw duplicate-key + /// error to either caller. + pub async fn ensure_wallet(pool: &PgPool, user_id: Uuid) -> Result { + if let Some(wallet) = Self::get_wallet(pool, user_id).await? { + return Ok(wallet); + } + + let free_plan: (Uuid, i32) = sqlx::query_as( + "SELECT id, monthly_credits FROM ai_plans WHERE code = 'free' AND is_active = TRUE", + ) + .fetch_one(pool) + .await?; + + sqlx::query( + r#" + INSERT INTO user_ai_subscriptions + (user_id, plan_id, monthly_credits_total, current_period_start, current_period_end) + VALUES ($1, $2, $3, NOW(), NOW() + INTERVAL '30 days') + ON CONFLICT (user_id) DO NOTHING + "#, + ) + .bind(user_id) + .bind(free_plan.0) + .bind(free_plan.1) + .execute(pool) + .await?; + + Self::get_wallet(pool, user_id) + .await? + .ok_or_else(|| sqlx::Error::RowNotFound) + } + + pub async fn get_wallet(pool: &PgPool, user_id: Uuid) -> Result, sqlx::Error> { + sqlx::query_as::<_, AiWallet>("SELECT * FROM user_ai_subscriptions WHERE user_id = $1") + .bind(user_id) + .fetch_optional(pool) + .await + } + + pub async fn get_plan(pool: &PgPool, plan_id: Uuid) -> Result, sqlx::Error> { + sqlx::query_as::<_, AiPlan>("SELECT id, code, name, monthly_credits, daily_action_limit, daily_credit_limit, allowed_models, allowed_features, is_active FROM ai_plans WHERE id = $1") + .bind(plan_id) + .fetch_optional(pool) + .await + } + + /// Credit purchased credits to a user's wallet (e.g. after a verified + /// payment). Row-locked and idempotency-keyed like every other wallet + /// mutation in this file -- a retried/replayed verify callback with + /// the same `idempotency_key` is a no-op, not a double-credit. + /// Returns `false` (no-op) if the idempotency key was already used. + pub async fn add_purchased_credits( + pool: &PgPool, + user_id: Uuid, + credits: i32, + idempotency_key: &str, + reference_type: &str, + reference_id: Option, + ) -> Result { + let mut tx = pool.begin().await?; + + if sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM ai_credit_ledger WHERE idempotency_key = $1") + .bind(idempotency_key) + .fetch_one(&mut *tx) + .await? + > 0 + { + tx.rollback().await?; + return Ok(false); + } + + let wallet = sqlx::query_as::<_, AiWallet>( + "SELECT * FROM user_ai_subscriptions WHERE user_id = $1 FOR UPDATE", + ) + .bind(user_id) + .fetch_one(&mut *tx) + .await?; + + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET purchased_credits_total = purchased_credits_total + $1, + lifetime_purchased_credits = lifetime_purchased_credits + $1, + updated_at = NOW() + WHERE id = $2 + "#, + ) + .bind(credits) + .bind(wallet.id) + .execute(&mut *tx) + .await?; + + let balance_after = wallet.available_credits() + credits; + sqlx::query( + r#" + INSERT INTO ai_credit_ledger + (wallet_id, entry_type, credits, balance_after, idempotency_key, reference_type, reference_id, actor_type) + VALUES ($1, 'purchase', $2, $3, $4, $5, $6, 'system') + "#, + ) + .bind(wallet.id) + .bind(credits) + .bind(balance_after) + .bind(idempotency_key) + .bind(reference_type) + .bind(reference_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + Ok(true) + } + + /// Admin-issued credit adjustment (e.g. goodwill/compensation credit, + /// or a manual correction) -- mirrors the shape of the existing + /// TraceCoins admin adjust endpoint (`POST /api/admin/credits/adjust`, + /// {user_id, amount, type: ADD|DEDUCT, reason, reference_id?}) so the + /// admin UI can reuse the same form pattern for both currencies. + /// `reason` is mandatory -- every admin-issued credit change must be + /// explained, both for audit and to guard against casual misuse. + /// ADD grants bonus credits (not "purchased" -- no money changed + /// hands); DEDUCT consumes bonus first, then monthly, then purchased, + /// same order as `try_capture_reservation`. + pub async fn admin_adjust_credits( + pool: &PgPool, + user_id: Uuid, + amount: i32, + is_add: bool, + reason: &str, + actor_id: Uuid, + idempotency_key: Option<&str>, + ) -> Result { + if amount <= 0 { + return Err(AiCreditsError::InvalidAmount("amount must be positive".to_string())); + } + + let mut tx = pool.begin().await?; + + if let Some(key) = idempotency_key { + if sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM ai_credit_ledger WHERE idempotency_key = $1") + .bind(key) + .fetch_one(&mut *tx) + .await? + > 0 + { + let wallet = sqlx::query_as::<_, AiWallet>("SELECT * FROM user_ai_subscriptions WHERE user_id = $1") + .bind(user_id) + .fetch_one(&mut *tx) + .await?; + tx.rollback().await?; + return Ok(wallet); + } + } + + let wallet = sqlx::query_as::<_, AiWallet>( + "SELECT * FROM user_ai_subscriptions WHERE user_id = $1 FOR UPDATE", + ) + .bind(user_id) + .fetch_one(&mut *tx) + .await?; + + let (entry_type, balance_after) = if is_add { + sqlx::query( + "UPDATE user_ai_subscriptions SET bonus_credits_total = bonus_credits_total + $1, updated_at = NOW() WHERE id = $2", + ) + .bind(amount) + .bind(wallet.id) + .execute(&mut *tx) + .await?; + ("admin_adjustment_credit", wallet.available_credits() + amount) + } else { + if wallet.available_credits() < amount { + tx.rollback().await?; + return Err(AiCreditsError::InsufficientCredits); + } + + let bonus_available = wallet.bonus_credits_total - wallet.bonus_credits_used; + let from_bonus = amount.min(bonus_available.max(0)); + let remaining = amount - from_bonus; + let monthly_available = wallet.monthly_credits_total - wallet.monthly_credits_used; + let from_monthly = remaining.min(monthly_available.max(0)); + let from_purchased = remaining - from_monthly; + + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET bonus_credits_used = bonus_credits_used + $1, + monthly_credits_used = monthly_credits_used + $2, + purchased_credits_used = purchased_credits_used + $3, + lifetime_used_credits = lifetime_used_credits + $4, + updated_at = NOW() + WHERE id = $5 + "#, + ) + .bind(from_bonus) + .bind(from_monthly) + .bind(from_purchased) + .bind(amount) + .bind(wallet.id) + .execute(&mut *tx) + .await?; + ("admin_adjustment_debit", wallet.available_credits() - amount) + }; + + sqlx::query( + r#" + INSERT INTO ai_credit_ledger + (wallet_id, entry_type, credits, balance_after, idempotency_key, reference_type, actor_type, actor_id, metadata) + VALUES ($1, $2, $3, $4, $5, 'admin_action', 'admin', $6, jsonb_build_object('reason', $7::text)) + "#, + ) + .bind(wallet.id) + .bind(entry_type) + .bind(if is_add { amount } else { -amount }) + .bind(balance_after) + .bind(idempotency_key) + .bind(actor_id) + .bind(reason) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + sqlx::query_as::<_, AiWallet>("SELECT * FROM user_ai_subscriptions WHERE id = $1") + .bind(wallet.id) + .fetch_one(pool) + .await + .map_err(AiCreditsError::Db) + } + + pub async fn get_feature_cost( + pool: &PgPool, + feature_code: &str, + ) -> Result, sqlx::Error> { + sqlx::query_as::<_, AiFeatureCost>( + r#" + SELECT id, feature_code, display_name, default_model, credit_cost, + max_input_tokens, max_output_tokens, min_plan_code, timeout_ms, + fallback_model, is_active + FROM ai_feature_costs + WHERE feature_code = $1 AND is_active = TRUE + "#, + ) + .bind(feature_code) + .fetch_optional(pool) + .await + } + + /// Roll `daily_actions_used`/`daily_credits_used` over to zero if the + /// wallet's stored usage date isn't today. Called inline, inside the + /// same locked transaction as a reserve, rather than depending on a + /// separate cron job -- there is no reset job in this codebase yet + /// (Section 20 Phase 2 of the architecture doc), and an inline reset + /// is strictly safer than shipping enforcement that depends on a job + /// that doesn't exist. + async fn reset_daily_counters_if_needed( + tx: &mut Transaction<'_, Postgres>, + wallet_id: Uuid, + stored_date: NaiveDate, + ) -> Result<(), sqlx::Error> { + let today = Utc::now().date_naive(); + if stored_date < today { + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET daily_actions_used = 0, daily_credits_used = 0, daily_usage_date = $2 + WHERE id = $1 + "#, + ) + .bind(wallet_id) + .bind(today) + .execute(&mut **tx) + .await?; + } + Ok(()) + } + + /// Atomically reserve `credits` against a user's wallet for a feature + /// call. Mirrors `TracecoinWalletRepository::try_reserve_tracecoins`: + /// `pool.begin()` -> `SELECT ... FOR UPDATE` -> check -> mutate -> + /// ledger insert -> commit (or rollback on any failed check). Callers + /// must follow up with `try_capture_reservation` on success or + /// `try_release_reservation` on failure -- never leave a hold + /// unresolved (a background reaper for abandoned holds is a Phase 1 + /// follow-up, not yet implemented; `expires_at` is written now so that + /// reaper has something to key off of once it exists). + #[allow(clippy::too_many_arguments)] + pub async fn try_reserve_credits( + pool: &PgPool, + user_id: Uuid, + feature_code: &str, + credits: i32, + request_id: Option<&str>, + idempotency_key: Option<&str>, + ) -> Result { + // Ensure the wallet exists before opening the locked transaction -- + // ensure_wallet does its own get-or-create round trip and would + // deadlock with itself if run inside the same FOR UPDATE tx. + Self::ensure_wallet(pool, user_id).await?; + + let mut tx = pool.begin().await?; + + // Idempotency: a retried request with the same key returns the + // existing hold instead of reserving twice. + if let Some(key) = idempotency_key { + if let Some(existing) = sqlx::query_as::<_, AiReservationHold>( + "SELECT id, wallet_id, credits_held, feature_code, status FROM ai_reservation_holds WHERE idempotency_key = $1", + ) + .bind(key) + .fetch_optional(&mut *tx) + .await? + { + tx.rollback().await?; + return Ok(existing); + } + } + + let wallet = sqlx::query_as::<_, AiWallet>( + "SELECT * FROM user_ai_subscriptions WHERE user_id = $1 FOR UPDATE", + ) + .bind(user_id) + .fetch_one(&mut *tx) + .await?; + + Self::reset_daily_counters_if_needed(&mut tx, wallet.id, wallet.daily_usage_date).await?; + + // Re-read after the possible reset so the checks below see fresh + // counters (cheap: still inside the same row lock). + let wallet = sqlx::query_as::<_, AiWallet>( + "SELECT * FROM user_ai_subscriptions WHERE id = $1 FOR UPDATE", + ) + .bind(wallet.id) + .fetch_one(&mut *tx) + .await?; + + let plan = sqlx::query_as::<_, AiPlan>( + "SELECT id, code, name, monthly_credits, daily_action_limit, daily_credit_limit, allowed_models, allowed_features, is_active FROM ai_plans WHERE id = $1", + ) + .bind(wallet.plan_id) + .fetch_one(&mut *tx) + .await?; + + if wallet.daily_actions_used >= plan.daily_action_limit { + tx.rollback().await?; + return Err(AiCreditsError::DailyActionLimitReached); + } + + if let Some(daily_credit_limit) = plan.daily_credit_limit { + if wallet.daily_credits_used + credits > daily_credit_limit { + tx.rollback().await?; + return Err(AiCreditsError::DailyCreditLimitReached); + } + } + + if wallet.available_credits() < credits { + tx.rollback().await?; + return Err(AiCreditsError::InsufficientCredits); + } + + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET reserved_credits = reserved_credits + $1, + daily_actions_used = daily_actions_used + 1, + daily_credits_used = daily_credits_used + $1, + updated_at = NOW() + WHERE id = $2 + "#, + ) + .bind(credits) + .bind(wallet.id) + .execute(&mut *tx) + .await?; + + let hold_id: Uuid = sqlx::query_scalar( + r#" + INSERT INTO ai_reservation_holds + (wallet_id, credits_held, feature_code, request_id, idempotency_key, status) + VALUES ($1, $2, $3, $4, $5, 'held') + RETURNING id + "#, + ) + .bind(wallet.id) + .bind(credits) + .bind(feature_code) + .bind(request_id) + .bind(idempotency_key) + .fetch_one(&mut *tx) + .await?; + + let balance_after = wallet.available_credits() - credits; + sqlx::query( + r#" + INSERT INTO ai_credit_ledger + (wallet_id, entry_type, credits, balance_after, reference_type, reference_id, actor_type) + VALUES ($1, 'reservation_hold', $2, $3, 'reservation', $4, 'user') + "#, + ) + .bind(wallet.id) + .bind(-credits) + .bind(balance_after) + .bind(hold_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + Ok(AiReservationHold { + id: hold_id, + wallet_id: wallet.id, + credits_held: credits, + feature_code: feature_code.to_string(), + status: "held".to_string(), + }) + } + + /// Convert a hold into an actual charge. `actual_credits` lets a + /// caller capture less than was held (e.g. a streaming response that + /// used fewer tokens than the reservation's ceiling); the difference + /// is simply never moved into a `_used` column, which makes it + /// available again with no separate "release remainder" step needed. + /// Consumption order: bonus credits first, then monthly, then + /// purchased -- see Section 4.1 of the architecture doc. + pub async fn try_capture_reservation( + pool: &PgPool, + hold_id: Uuid, + actual_credits: Option, + ) -> Result { + let mut tx = pool.begin().await?; + + let hold = sqlx::query_as::<_, AiReservationHold>( + "SELECT id, wallet_id, credits_held, feature_code, status FROM ai_reservation_holds WHERE id = $1 AND status = 'held' FOR UPDATE", + ) + .bind(hold_id) + .fetch_optional(&mut *tx) + .await?; + + let Some(hold) = hold else { + tx.rollback().await?; + return Ok(false); + }; + + let captured = actual_credits.unwrap_or(hold.credits_held).min(hold.credits_held).max(0); + + let wallet = sqlx::query_as::<_, AiWallet>( + "SELECT * FROM user_ai_subscriptions WHERE id = $1 FOR UPDATE", + ) + .bind(hold.wallet_id) + .fetch_one(&mut *tx) + .await?; + + let bonus_available = wallet.bonus_credits_total - wallet.bonus_credits_used; + let from_bonus = captured.min(bonus_available.max(0)); + let remaining = captured - from_bonus; + + let monthly_available = wallet.monthly_credits_total - wallet.monthly_credits_used; + let from_monthly = remaining.min(monthly_available.max(0)); + let from_purchased = remaining - from_monthly; + + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET reserved_credits = reserved_credits - $1, + bonus_credits_used = bonus_credits_used + $2, + monthly_credits_used = monthly_credits_used + $3, + purchased_credits_used = purchased_credits_used + $4, + lifetime_used_credits = lifetime_used_credits + $5, + updated_at = NOW() + WHERE id = $6 + "#, + ) + .bind(hold.credits_held) + .bind(from_bonus) + .bind(from_monthly) + .bind(from_purchased) + .bind(captured) + .bind(wallet.id) + .execute(&mut *tx) + .await?; + + let balance_after = wallet.available_credits() + (hold.credits_held - captured); + sqlx::query( + r#" + INSERT INTO ai_credit_ledger + (wallet_id, entry_type, credits, balance_after, reference_type, reference_id, actor_type) + VALUES ($1, 'reservation_capture', $2, $3, 'reservation', $4, 'system') + "#, + ) + .bind(wallet.id) + .bind(-captured) + .bind(balance_after) + .bind(hold.id) + .execute(&mut *tx) + .await?; + + sqlx::query("UPDATE ai_reservation_holds SET status = 'captured', resolved_at = NOW() WHERE id = $1") + .bind(hold.id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + Ok(true) + } + + /// Release a hold without charging -- the LLM call failed, timed out, + /// or a downstream validation error occurred after the reserve. + pub async fn try_release_reservation(pool: &PgPool, hold_id: Uuid) -> Result { + let mut tx = pool.begin().await?; + + let hold = sqlx::query_as::<_, AiReservationHold>( + "SELECT id, wallet_id, credits_held, feature_code, status FROM ai_reservation_holds WHERE id = $1 AND status = 'held' FOR UPDATE", + ) + .bind(hold_id) + .fetch_optional(&mut *tx) + .await?; + + let Some(hold) = hold else { + tx.rollback().await?; + return Ok(false); + }; + + let wallet = sqlx::query_as::<_, AiWallet>( + "SELECT * FROM user_ai_subscriptions WHERE id = $1 FOR UPDATE", + ) + .bind(hold.wallet_id) + .fetch_one(&mut *tx) + .await?; + + // Roll back both the reservation itself AND the daily-action/ + // daily-credit counters bumped when it was created (try_reserve_credits) + // -- otherwise a request that fails after reserving (e.g. the LLM + // call errors) still permanently burns the user's daily quota even + // though they were never charged. + sqlx::query( + r#" + UPDATE user_ai_subscriptions + SET reserved_credits = reserved_credits - $1, + daily_actions_used = GREATEST(daily_actions_used - 1, 0), + daily_credits_used = GREATEST(daily_credits_used - $1, 0), + updated_at = NOW() + WHERE id = $2 + "#, + ) + .bind(hold.credits_held) + .bind(wallet.id) + .execute(&mut *tx) + .await?; + + let balance_after = wallet.available_credits() + hold.credits_held; + sqlx::query( + r#" + INSERT INTO ai_credit_ledger + (wallet_id, entry_type, credits, balance_after, reference_type, reference_id, actor_type) + VALUES ($1, 'reservation_release', $2, $3, 'reservation', $4, 'system') + "#, + ) + .bind(wallet.id) + .bind(hold.credits_held) + .bind(balance_after) + .bind(hold.id) + .execute(&mut *tx) + .await?; + + sqlx::query("UPDATE ai_reservation_holds SET status = 'released', resolved_at = NOW() WHERE id = $1") + .bind(hold.id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + Ok(true) + } + + /// Process a refund for a completed AI credit charge. + /// + /// This creates a new ledger entry with entry_type='refund' and credits back + /// the specified amount to the user's wallet. The refund is tied to the + /// original debit ledger entry for audit purposes. + /// + /// # Arguments + /// * `pool` - Database connection pool + /// * `user_id` - The user receiving the refund + /// * `ledger_debit_entry_id` - The original debit ledger entry being refunded + /// * `credits_to_refund` - Amount of credits to refund + /// * `reason` - One of: provider_failure, timeout, validation_failure, manual, user_dispute + /// * `initiated_by` - Who initiated: system, admin, user_dispute + /// * `actor_id` - The admin/user who approved the refund (if applicable) + /// * `idempotency_key` - Optional idempotency key + /// * `notes` - Optional notes about the refund + /// + /// # Returns + /// * `Ok((refund_id, AiWallet))` - The refund record ID and updated wallet + /// * `Err(AiCreditsError)` - If the refund cannot be processed + #[allow(clippy::too_many_arguments)] + pub async fn refund_credits( + pool: &PgPool, + user_id: Uuid, + ledger_debit_entry_id: Uuid, + credits_to_refund: i32, + reason: &str, + initiated_by: &str, + actor_id: Option, + idempotency_key: Option<&str>, + notes: Option<&str>, + ) -> Result<(Uuid, AiWallet), AiCreditsError> { + if credits_to_refund <= 0 { + return Err(AiCreditsError::InvalidAmount("credits_to_refund must be positive".to_string())); + } + + let mut tx = pool.begin().await?; + + // Idempotency check + if let Some(key) = idempotency_key { + if sqlx::query_scalar::<_, i64>("SELECT COUNT(*) FROM ai_refunds WHERE idempotency_key = $1") + .bind(key) + .fetch_one(&mut *tx) + .await? > 0 + { + let refund_id: Uuid = sqlx::query_scalar("SELECT id FROM ai_refunds WHERE idempotency_key = $1") + .bind(key) + .fetch_one(&mut *tx) + .await?; + let wallet = sqlx::query_as::<_, AiWallet>("SELECT * FROM user_ai_subscriptions WHERE user_id = $1") + .bind(user_id) + .fetch_one(&mut *tx) + .await?; + tx.rollback().await?; + return Ok((refund_id, wallet)); + } + } + + // Verify the original debit entry exists and is a charge + let debit_entry: Option<(Uuid, i32, String)> = sqlx::query_as( + "SELECT id, credits, entry_type FROM ai_credit_ledger WHERE id = $1" + ) + .bind(ledger_debit_entry_id) + .fetch_optional(&mut *tx) + .await?; + + let Some((_, debit_credits, entry_type)) = debit_entry else { + tx.rollback().await?; + return Err(AiCreditsError::ReservationNotFound); + }; + + // Ensure it's actually a debit (negative credits) + if debit_credits >= 0 { + tx.rollback().await?; + return Err(AiCreditsError::InvalidAmount("ledger entry is not a debit".to_string())); + } + + // Get wallet and lock it + let wallet = sqlx::query_as::<_, AiWallet>( + "SELECT * FROM user_ai_subscriptions WHERE user_id = $1 FOR UPDATE", + ) + .bind(user_id) + .fetch_one(&mut *tx) + .await?; + + // Refund as bonus credits (goodwill/compensation) + sqlx::query( + "UPDATE user_ai_subscriptions SET bonus_credits_total = bonus_credits_total + $1, updated_at = NOW() WHERE id = $2", + ) + .bind(credits_to_refund) + .bind(wallet.id) + .execute(&mut *tx) + .await?; + + // Create refund record + let refund_id: Uuid = sqlx::query_scalar( + r#" + INSERT INTO ai_refunds + (user_id, ledger_debit_entry_id, credits_refunded, reason, initiated_by, status, notes, admin_actor_id, idempotency_key, resolved_at) + VALUES ($1, $2, $3, $4, $5, 'completed', $6, $7, $8, NOW()) + RETURNING id + "#, + ) + .bind(user_id) + .bind(ledger_debit_entry_id) + .bind(credits_to_refund) + .bind(reason) + .bind(initiated_by) + .bind(notes) + .bind(actor_id) + .bind(idempotency_key) + .fetch_one(&mut *tx) + .await?; + + // Create ledger entry for the refund + let balance_after = wallet.available_credits() + credits_to_refund; + sqlx::query( + r#" + INSERT INTO ai_credit_ledger + (wallet_id, entry_type, credits, balance_after, reference_type, reference_id, actor_type, actor_id, metadata) + VALUES ($1, 'refund', $2, $3, 'refund', $4, 'system', $5, jsonb_build_object('reason', $6, 'refund_id', $7)) + "#, + ) + .bind(wallet.id) + .bind(credits_to_refund) // Positive for credit + .bind(balance_after) + .bind(ledger_debit_entry_id) + .bind(actor_id) + .bind(reason) + .bind(refund_id) + .execute(&mut *tx) + .await?; + + tx.commit().await?; + + // Return updated wallet + let updated_wallet = sqlx::query_as::<_, AiWallet>("SELECT * FROM user_ai_subscriptions WHERE id = $1") + .bind(wallet.id) + .fetch_one(pool) + .await + .map_err(AiCreditsError::Db)?; + + Ok((refund_id, updated_wallet)) + } + + /// List refunds for a user (paginated) + pub async fn list_refunds( + pool: &PgPool, + user_id: Uuid, + page: i64, + limit: i64, + ) -> Result, sqlx::Error> { + let offset = (page - 1) * limit; + sqlx::query_as::<_, RefundRow>( + r#" + SELECT + r.id, + r.user_id, + r.ledger_debit_entry_id, + r.credits_refunded, + r.reason, + r.initiated_by, + r.status, + r.notes, + r.admin_actor_id, + r.created_at, + r.resolved_at, + l.entry_type as original_entry_type, + l.credits as original_credits + FROM ai_refunds r + JOIN ai_credit_ledger l ON r.ledger_debit_entry_id = l.id + WHERE r.user_id = $1 + ORDER BY r.created_at DESC + LIMIT $2 OFFSET $3 + "#, + ) + .bind(user_id) + .bind(limit) + .bind(offset) + .fetch_all(pool) + .await + } +} + +#[derive(Debug, Serialize, Deserialize, FromRow, Clone)] +pub struct RefundRow { + pub id: Uuid, + pub user_id: Uuid, + pub ledger_debit_entry_id: Uuid, + pub credits_refunded: i32, + pub reason: String, + pub initiated_by: String, + pub status: String, + pub notes: Option, + pub admin_actor_id: Option, + pub created_at: DateTime, + pub resolved_at: Option>, + pub original_entry_type: String, + pub original_credits: i32, +} diff --git a/crates/db/src/models/mod.rs b/crates/db/src/models/mod.rs index b8ae833..cd71894 100644 --- a/crates/db/src/models/mod.rs +++ b/crates/db/src/models/mod.rs @@ -27,4 +27,5 @@ pub mod designation; pub mod verification; pub mod user_role_profile; pub mod tracecoin_wallet; +pub mod ai_credits; diff --git a/crates/db/tests/ai_credits.rs b/crates/db/tests/ai_credits.rs new file mode 100644 index 0000000..20df4e1 --- /dev/null +++ b/crates/db/tests/ai_credits.rs @@ -0,0 +1,206 @@ +//! Integration tests for the Ask Ash AI credits wallet (Phase 1). +//! Requires a live Postgres reachable via TEST_DATABASE_URL with the +//! `20260317000000_init_config_schema`, `20260317000001_init_users_schema`, +//! and `20260703210000_ai_credits_wallet` migrations already applied. +//! Run manually: `TEST_DATABASE_URL=... cargo test -p db --test ai_credits`. + +use db::models::ai_credits::{AiCreditsError, AiCreditsRepository}; +use sqlx::PgPool; +use uuid::Uuid; + +async fn pool() -> PgPool { + let url = std::env::var("TEST_DATABASE_URL") + .expect("set TEST_DATABASE_URL to run ai_credits integration tests"); + sqlx::postgres::PgPoolOptions::new() + .max_connections(10) + .connect(&url) + .await + .expect("connect to test db") +} + +async fn make_user(pool: &PgPool) -> Uuid { + let id = Uuid::new_v4(); + sqlx::query("INSERT INTO users (id, email, password_hash) VALUES ($1, $2, 'x')") + .bind(id) + .bind(format!("{id}@test.local")) + .execute(pool) + .await + .expect("insert test user"); + id +} + +#[tokio::test] +async fn ensure_wallet_grants_free_plan_credits() { + let pool = pool().await; + let user_id = make_user(&pool).await; + + let wallet = AiCreditsRepository::ensure_wallet(&pool, user_id) + .await + .expect("ensure_wallet"); + + assert_eq!(wallet.monthly_credits_total, 10); + assert_eq!(wallet.available_credits(), 10); + + // Idempotent: calling again returns the same wallet, doesn't grant twice. + let wallet2 = AiCreditsRepository::ensure_wallet(&pool, user_id) + .await + .expect("ensure_wallet again"); + assert_eq!(wallet2.id, wallet.id); + assert_eq!(wallet2.monthly_credits_total, 10); +} + +#[tokio::test] +async fn reserve_capture_debits_wallet() { + let pool = pool().await; + let user_id = make_user(&pool).await; + AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap(); + + let hold = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 3, None, None) + .await + .expect("reserve"); + + let mid_wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap(); + assert_eq!(mid_wallet.reserved_credits, 3); + assert_eq!(mid_wallet.available_credits(), 7); + + let captured = AiCreditsRepository::try_capture_reservation(&pool, hold.id, None) + .await + .expect("capture"); + assert!(captured); + + let final_wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap(); + assert_eq!(final_wallet.reserved_credits, 0); + assert_eq!(final_wallet.monthly_credits_used, 3); + assert_eq!(final_wallet.lifetime_used_credits, 3); + assert_eq!(final_wallet.available_credits(), 7); +} + +#[tokio::test] +async fn reserve_release_refunds_wallet() { + let pool = pool().await; + let user_id = make_user(&pool).await; + AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap(); + + let hold = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 4, None, None) + .await + .expect("reserve"); + + let released = AiCreditsRepository::try_release_reservation(&pool, hold.id) + .await + .expect("release"); + assert!(released); + + let wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap(); + assert_eq!(wallet.reserved_credits, 0); + assert_eq!(wallet.monthly_credits_used, 0); + assert_eq!(wallet.available_credits(), 10); + // Regression: releasing a hold must also roll back the daily action/ + // credit counters bumped at reserve time -- a failed request (e.g. the + // LLM call errors) must not permanently burn daily quota when the user + // was never actually charged. + assert_eq!(wallet.daily_actions_used, 0, "release must roll back daily_actions_used"); + assert_eq!(wallet.daily_credits_used, 0, "release must roll back daily_credits_used"); + + // Releasing an already-released hold is a no-op, not an error. + let released_again = AiCreditsRepository::try_release_reservation(&pool, hold.id) + .await + .expect("release again"); + assert!(!released_again); +} + +#[tokio::test] +async fn reserve_beyond_balance_is_rejected() { + let pool = pool().await; + let user_id = make_user(&pool).await; + AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap(); + + let result = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 11, None, None).await; + // The seeded Free plan's daily_credit_limit (10) equals its + // monthly_credits_total (10), so a request for 11 trips the + // daily-limit guardrail before the balance check ever runs -- both + // are correct rejections of the same over-request, so accept either. + assert!(matches!( + result, + Err(AiCreditsError::InsufficientCredits) | Err(AiCreditsError::DailyCreditLimitReached) + )); + + // Balance must be untouched after a rejected reservation. + let wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap(); + assert_eq!(wallet.reserved_credits, 0); + assert_eq!(wallet.available_credits(), 10); +} + +#[tokio::test] +async fn reserve_is_idempotent_on_key() { + let pool = pool().await; + let user_id = make_user(&pool).await; + AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap(); + + let key = format!("idem-{}", Uuid::new_v4()); + let hold1 = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 5, None, Some(&key)) + .await + .expect("first reserve"); + let hold2 = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 5, None, Some(&key)) + .await + .expect("retried reserve with same key"); + + assert_eq!(hold1.id, hold2.id, "retry with same idempotency key must not create a second hold"); + + let wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap(); + assert_eq!(wallet.reserved_credits, 5, "credits must only be reserved once despite the retry"); +} + +/// The concrete race condition this design exists to prevent (Section 4.2 +/// / G4 of the architecture doc): fire N concurrent reservations against a +/// wallet that can only afford one of them, and confirm exactly one wins. +#[tokio::test] +async fn concurrent_reservations_cannot_overdraw_wallet() { + let pool = pool().await; + let user_id = make_user(&pool).await; + let wallet = AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap(); + assert_eq!(wallet.available_credits(), 10); + + // 5 concurrent requests each asking for all 10 credits. At most one + // should succeed; a broken check-then-write implementation would let + // several through and drive the balance negative. + let mut handles = Vec::new(); + for _ in 0..5 { + let pool = pool.clone(); + handles.push(tokio::spawn(async move { + AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 10, None, None).await + })); + } + + let mut successes = 0; + for h in handles { + if h.await.unwrap().is_ok() { + successes += 1; + } + } + + assert_eq!(successes, 1, "exactly one concurrent reservation for the full balance should succeed"); + + let final_wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap(); + assert_eq!(final_wallet.available_credits(), 0); + assert!( + final_wallet.monthly_credits_used <= final_wallet.monthly_credits_total, + "wallet must never be driven into overdraft" + ); +} + +#[tokio::test] +async fn daily_action_limit_enforced() { + let pool = pool().await; + let user_id = make_user(&pool).await; + AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap(); + + // Free plan's daily_action_limit is 3 (seeded in the migration). + for _ in 0..3 { + AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 1, None, None) + .await + .expect("reserve within daily action limit"); + } + + let fourth = AiCreditsRepository::try_reserve_credits(&pool, user_id, "test_feature", 1, None, None).await; + assert!(matches!(fourth, Err(AiCreditsError::DailyActionLimitReached))); +} diff --git a/crates/db/tests/ai_credits_reaper.rs b/crates/db/tests/ai_credits_reaper.rs new file mode 100644 index 0000000..44a8d05 --- /dev/null +++ b/crates/db/tests/ai_credits_reaper.rs @@ -0,0 +1,67 @@ +//! Verifies the cron reaper's core query/release logic by reproducing it +//! directly against the wallet primitives (the reaper itself lives in +//! apps/cron, which can't easily be imported as a lib from here -- this +//! test exercises the same "find expired holds, release them" contract +//! the cron task relies on). + +use db::models::ai_credits::AiCreditsRepository; +use sqlx::PgPool; +use uuid::Uuid; + +async fn pool() -> PgPool { + let url = std::env::var("TEST_DATABASE_URL") + .expect("set TEST_DATABASE_URL to run this test"); + sqlx::postgres::PgPoolOptions::new() + .max_connections(5) + .connect(&url) + .await + .expect("connect to test db") +} + +#[tokio::test] +async fn expired_hold_is_swept_and_credits_return_to_available() { + let pool = pool().await; + let user_id = Uuid::parse_str("33333333-3333-3333-3333-333333333333").unwrap(); + + AiCreditsRepository::ensure_wallet(&pool, user_id).await.unwrap(); + + let hold = AiCreditsRepository::try_reserve_credits(&pool, user_id, "help_answer", 4, None, None) + .await + .expect("reserve"); + + let wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap(); + assert_eq!(wallet.reserved_credits, 4, "credits should be reserved before the sweep"); + + // Force the hold into the past, exactly as an abandoned request would + // eventually end up (the migration default is NOW() + 5 minutes). + sqlx::query("UPDATE ai_reservation_holds SET expires_at = NOW() - INTERVAL '1 minute' WHERE id = $1") + .bind(hold.id) + .execute(&pool) + .await + .unwrap(); + + // Reproduce the cron task's sweep query + release call. + let expired: Vec = sqlx::query_scalar( + "SELECT id FROM ai_reservation_holds WHERE status = 'held' AND expires_at < NOW()", + ) + .fetch_all(&pool) + .await + .unwrap(); + assert_eq!(expired, vec![hold.id]); + + for id in expired { + let released = AiCreditsRepository::try_release_reservation(&pool, id).await.unwrap(); + assert!(released); + } + + let wallet = AiCreditsRepository::get_wallet(&pool, user_id).await.unwrap().unwrap(); + assert_eq!(wallet.reserved_credits, 0, "reaper must release the reservation back to available"); + assert_eq!(wallet.available_credits(), 10, "released credits must be fully spendable again"); + + let hold_row: (String,) = sqlx::query_as("SELECT status FROM ai_reservation_holds WHERE id = $1") + .bind(hold.id) + .fetch_one(&pool) + .await + .unwrap(); + assert_eq!(hold_row.0, "released"); +}