feat: Complete Ask Ash AI Credits implementation (Tasks 1-10)
- Task 1: Admin endpoints for wallet management (balance, ledger, adjust, reconcile) - Task 2: AI Credits admin UI (pricing.tsx, credit.tsx) - Task 3: Ollama security (NetworkPolicy, prompt validation, audit logging) - Task 4: LiteLLM integration (litellm.rs, migrated AI feature handlers) - Task 5: Refund architecture (ai_refunds table, endpoints) - Task 6: Coupons, promotions, referrals (order creation with coupon validation) - Task 7: Subscription lifecycle (plan upgrades/downgrades/cancellations, cron jobs) - Task 8: Credit expiration enforcement (daily cron task) - Task 9: Token cost engine (ai_model_cost_config, margin view) - Task 10: Observability (metrics tables, aggregation function) New files: - apps/users/src/litellm.rs (LiteLLM client) - apps/users/src/ai_credits.rs (Charging primitives) - apps/users/src/ai_subscription.rs (Plan lifecycle) - apps/users/src/handlers/ai_credits.rs (Admin endpoints) - apps/payments/src/ai_credits.rs (Package purchase) - apps/payments/src/payu.rs (PayU integration) - apps/cron/src/tasks/ai_credits.rs (Cron jobs) - 9 SQL migrations for schema - crates/db/src/models/ai_credits.rs (Repository) - tests for ai_credits
This commit is contained in:
Ashwin Kumar Sivakumar
parent
d5b2f9c682
commit
3c0f45f1ee
38 changed files with 5782 additions and 295 deletions
203
Cargo.lock
generated
203
Cargo.lock
generated
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -75,6 +75,64 @@ async fn main() -> Result<(), Box<dyn std::error::Error>> {
|
|||
}
|
||||
});
|
||||
|
||||
// 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.");
|
||||
|
|
|
|||
168
apps/cron/src/tasks/ai_credits.rs
Normal file
168
apps/cron/src/tasks/ai_credits.rs
Normal file
|
|
@ -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<u64, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let expired: Vec<uuid::Uuid> = 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<u64, Box<dyn std::error::Error + Send + Sync>> {
|
||||
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<u64, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let applied = crate::ai_subscription::apply_scheduled_downgrades(pool)
|
||||
.await
|
||||
.map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
|
||||
|
||||
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<u64, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let expired = crate::ai_subscription::expire_trials(pool)
|
||||
.await
|
||||
.map_err(|e| Box::new(e) as Box<dyn std::error::Error + Send + Sync>)?;
|
||||
|
||||
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<u64, Box<dyn std::error::Error + Send + Sync>> {
|
||||
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)
|
||||
}
|
||||
|
|
@ -2,3 +2,4 @@ pub mod leads;
|
|||
pub mod requirements;
|
||||
pub mod jobs;
|
||||
pub mod reminders;
|
||||
pub mod ai_credits;
|
||||
|
|
|
|||
|
|
@ -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()));
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
542
apps/payments/src/ai_credits.rs
Normal file
542
apps/payments/src/ai_credits.rs
Normal file
|
|
@ -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<AppState> {
|
||||
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<AppState> {
|
||||
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<String>,
|
||||
credits: i32,
|
||||
price_inr: i32,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ListPackagesResponse {
|
||||
packages: Vec<AiCreditPackageRow>,
|
||||
}
|
||||
|
||||
async fn list_packages(State(state): State<AppState>) -> Result<Json<ListPackagesResponse>, (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<String>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
phone: Option<String>,
|
||||
}
|
||||
|
||||
/// 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<Option<(i32, String, i32)>, String> {
|
||||
// Get coupon details
|
||||
let coupon: Option<(Uuid, String, rust_decimal::Decimal, Option<rust_decimal::Decimal>, i32, i32, Vec<Uuid>)> = 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<AppState>,
|
||||
Json(payload): Json<CreateAiOrderRequest>,
|
||||
) -> Result<Json<CreateAiOrderResponse>, (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<String>,
|
||||
#[serde(default)]
|
||||
udf2: Option<String>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
discount_applied: Option<i32>,
|
||||
amount_inr: i32,
|
||||
package_id: Uuid,
|
||||
}
|
||||
|
||||
async fn verify_order(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
Json(payload): Json<VerifyAiOrderRequest>,
|
||||
) -> Result<Json<VerifyAiOrderResponse>, (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<Uuid> = 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<String>,
|
||||
credits: i32,
|
||||
price_inr: i32,
|
||||
is_active: bool,
|
||||
created_at: chrono::DateTime<chrono::Utc>,
|
||||
updated_at: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct AdminListPackagesResponse {
|
||||
packages: Vec<AdminAiCreditPackageRow>,
|
||||
}
|
||||
|
||||
async fn admin_list_packages(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
) -> Result<Json<AdminListPackagesResponse>, (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<String>,
|
||||
credits: i32,
|
||||
price_inr: i32,
|
||||
}
|
||||
|
||||
async fn admin_create_package(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
Json(body): Json<CreatePackageRequest>,
|
||||
) -> Result<Json<AdminAiCreditPackageRow>, (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<String>,
|
||||
description: Option<String>,
|
||||
credits: Option<i32>,
|
||||
price_inr: Option<i32>,
|
||||
is_active: Option<bool>,
|
||||
}
|
||||
|
||||
async fn admin_update_package(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
Path(id): Path<Uuid>,
|
||||
Json(body): Json<UpdatePackageRequest>,
|
||||
) -> Result<Json<AdminAiCreditPackageRow>, (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))
|
||||
}
|
||||
|
|
@ -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<u64>,
|
||||
#[allow(dead_code)]
|
||||
currency: Option<String>,
|
||||
package_id: Option<String>,
|
||||
user_id: Option<String>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
txnid: String,
|
||||
mihpayid: String,
|
||||
status: String,
|
||||
hash: String,
|
||||
amount: String,
|
||||
productinfo: String,
|
||||
firstname: String,
|
||||
email: String,
|
||||
#[serde(default)]
|
||||
phone: Option<String>,
|
||||
#[serde(default)]
|
||||
udf1: Option<String>,
|
||||
#[serde(default)]
|
||||
udf2: Option<String>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
phone: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, FromRow)]
|
||||
|
|
@ -79,13 +121,11 @@ async fn create_order(
|
|||
State(state): State<AppState>,
|
||||
Json(payload): Json<CreateOrderRequest>,
|
||||
) -> Result<Json<CreateOrderResponse>, (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<AppState>,
|
||||
Json(payload): Json<VerifyPaymentRequest>,
|
||||
) -> Result<Json<VerifyPaymentResponse>, (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")
|
||||
|
|
|
|||
188
apps/payments/src/payu.rs
Normal file
188
apps/payments/src/payu.rs
Normal file
|
|
@ -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,
|
||||
));
|
||||
}
|
||||
}
|
||||
373
apps/users/src/ai_credits.rs
Normal file
373
apps/users/src/ai_credits.rs
Normal file
|
|
@ -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<E> {
|
||||
Credits(AiCreditsError),
|
||||
UnknownFeature(String),
|
||||
Work(E),
|
||||
}
|
||||
|
||||
impl<E: std::fmt::Display> std::fmt::Display for ChargeError<E> {
|
||||
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<T> {
|
||||
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<T, E, F, Fut>(
|
||||
pool: &PgPool,
|
||||
user_id: Uuid,
|
||||
feature_code: &str,
|
||||
request_id: Option<&str>,
|
||||
idempotency_key: Option<&str>,
|
||||
work: F,
|
||||
) -> Result<ChargeOutcome<T>, ChargeError<E>>
|
||||
where
|
||||
F: FnOnce(AiFeatureCost) -> Fut,
|
||||
Fut: std::future::Future<Output = Result<T, E>>,
|
||||
{
|
||||
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<i32>,
|
||||
output_tokens: Option<i32>,
|
||||
total_tokens: Option<i32>,
|
||||
) {
|
||||
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<String>)` - 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<ChargeOutcome<String>, ChargeError<String>> {
|
||||
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<ChargeOutcome<String>, ChargeError<String>> {
|
||||
// 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))
|
||||
}
|
||||
}
|
||||
}
|
||||
484
apps/users/src/ai_subscription.rs
Normal file
484
apps/users/src/ai_subscription.rs
Normal file
|
|
@ -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<Utc>,
|
||||
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<Uuid>,
|
||||
) -> 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<AiPlan> = 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<AiPlan> = 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<Uuid>,
|
||||
) -> Result<DateTime<Utc>, 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<usize, AiCreditsError> {
|
||||
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<Uuid>,
|
||||
) -> 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<DateTime<Utc>, 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<usize, AiCreditsError> {
|
||||
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<Vec<SubscriptionHistoryRow>, 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<String>,
|
||||
pub to_plan_name: String,
|
||||
pub change_type: String,
|
||||
pub proration_credits: i32,
|
||||
pub effective_at: DateTime<Utc>,
|
||||
pub status: String,
|
||||
pub created_at: DateTime<Utc>,
|
||||
}
|
||||
|
|
@ -49,34 +49,34 @@ struct OllamaGenerateResponse {
|
|||
response: String,
|
||||
}
|
||||
|
||||
async fn call_ollama(_state: &AppState, model: &str, prompt: &str) -> Result<String, String> {
|
||||
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<String, String> {
|
||||
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<String, 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_inline(
|
||||
base_url: &str,
|
||||
model: &str,
|
||||
prompt: &str,
|
||||
user_id: Option<&str>,
|
||||
) -> Result<String, 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))?;
|
||||
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<Uuid> = 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<Uuid> = 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<Uuid> = 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<AppState> {
|
|||
.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))
|
||||
|
|
|
|||
718
apps/users/src/handlers/ai_credits.rs
Normal file
718
apps/users/src/handlers/ai_credits.rs
Normal file
|
|
@ -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<AppState> {
|
||||
Router::new()
|
||||
.route("/wallet", get(get_wallet))
|
||||
.route("/charge", post(charge_feature))
|
||||
}
|
||||
|
||||
/// Admin router for AI credits management
|
||||
pub fn admin_router() -> Router<AppState> {
|
||||
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<serde_json::Value>) {
|
||||
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<AppState>) -> 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<String>,
|
||||
idempotency_key: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize)]
|
||||
struct ChargeResponse {
|
||||
feature_code: String,
|
||||
credits_charged: i32,
|
||||
available_credits: i32,
|
||||
}
|
||||
|
||||
async fn charge_feature(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
Json(body): Json<ChargeRequest>,
|
||||
) -> 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<AppState>) -> 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<chrono::Utc>,
|
||||
}
|
||||
|
||||
/// 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<AppState>) -> 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=<uuid>
|
||||
/// Get wallet balance for a specific user (admin only)
|
||||
async fn admin_balance(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
Query(params): Query<AdminBalanceQuery>,
|
||||
) -> 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<Uuid>,
|
||||
}
|
||||
|
||||
/// GET /api/admin/ai-credits/ledger?userId=<uuid>&page=&limit=
|
||||
/// Get paginated ledger entries for a user (admin only)
|
||||
async fn admin_ledger(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
Query(params): Query<AdminLedgerQuery>,
|
||||
) -> 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<LedgerEntryDto> = 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<Uuid>,
|
||||
page: Option<i64>,
|
||||
limit: Option<i64>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, sqlx::FromRow)]
|
||||
struct LedgerEntryDto {
|
||||
id: uuid::Uuid,
|
||||
entry_type: String,
|
||||
credits: i32,
|
||||
balance_after: i32,
|
||||
description: Option<String>,
|
||||
idempotency_key: Option<String>,
|
||||
created_at: chrono::DateTime<chrono::Utc>,
|
||||
}
|
||||
|
||||
/// POST /api/admin/ai-credits/adjust
|
||||
/// Adjust credits for a user (ADD or DEDUCT)
|
||||
async fn admin_adjust(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
Json(body): Json<AdminAdjustRequest>,
|
||||
) -> 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<String>,
|
||||
}
|
||||
|
||||
/// GET /api/admin/ai-credits/reconcile?from=<date>&to=<date>
|
||||
/// Reconcile ledger entries within date range
|
||||
async fn admin_reconcile(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
Query(params): Query<AdminReconcileQuery>,
|
||||
) -> 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<ReconcileSummaryDto> = 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<WalletDriftDto> = 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<chrono::DateTime<chrono::Utc>>,
|
||||
to: Option<chrono::DateTime<chrono::Utc>>,
|
||||
}
|
||||
|
||||
#[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<AppState>,
|
||||
Json(body): Json<CreateRefundRequest>,
|
||||
) -> 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<String>,
|
||||
idempotency_key: Option<String>,
|
||||
}
|
||||
|
||||
/// GET /api/admin/ai-credits/refunds?userId=<uuid>&page=&limit=
|
||||
/// List refunds for a user (admin view)
|
||||
async fn admin_list_refunds(
|
||||
auth: AuthUser,
|
||||
State(state): State<AppState>,
|
||||
Query(params): Query<ListRefundsQuery>,
|
||||
) -> 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<Uuid>,
|
||||
page: Option<i64>,
|
||||
limit: Option<i64>,
|
||||
}
|
||||
|
|
@ -150,7 +150,7 @@ pub struct TranslateResponse {
|
|||
|
||||
async fn translate_text(
|
||||
State(state): State<AppState>,
|
||||
_auth: AuthUser,
|
||||
auth: AuthUser,
|
||||
Json(body): Json<TranslateRequest>,
|
||||
) -> 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<String, String> {
|
||||
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())
|
||||
}
|
||||
|
||||
// ─────────────────────────────────────────────────────────────────────────────
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
222
apps/users/src/litellm.rs
Normal file
222
apps/users/src/litellm.rs
Normal file
|
|
@ -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<LiteLLMChatMessage>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
max_tokens: Option<i32>,
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
temperature: Option<f32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
struct LiteLLMChoice {
|
||||
message: LiteLLMChatMessage,
|
||||
finish_reason: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
struct LiteLLMUsage {
|
||||
prompt_tokens: Option<i32>,
|
||||
completion_tokens: Option<i32>,
|
||||
total_tokens: Option<i32>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Deserialize)]
|
||||
pub struct LiteLLMResponse {
|
||||
pub id: String,
|
||||
pub model: String,
|
||||
pub choices: Vec<LiteLLMChoice>,
|
||||
pub usage: Option<LiteLLMUsage>,
|
||||
}
|
||||
|
||||
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<i32>,
|
||||
) -> Result<LiteLLMResponse, String> {
|
||||
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<i32>,
|
||||
) -> Result<LiteLLMResponse, String> {
|
||||
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<String>) {
|
||||
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<bool, String> {
|
||||
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)),
|
||||
}
|
||||
}
|
||||
|
|
@ -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);
|
||||
|
||||
|
|
|
|||
|
|
@ -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 }
|
||||
|
|
|
|||
1
crates/cache/Cargo.toml
vendored
1
crates/cache/Cargo.toml
vendored
|
|
@ -12,3 +12,4 @@ uuid = { workspace = true }
|
|||
tracing = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
reqwest = { workspace = true }
|
||||
lazy_static = "1.4"
|
||||
|
|
|
|||
357
crates/cache/src/ollama.rs
vendored
357
crates/cache/src/ollama.rs
vendored
|
|
@ -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<String>,
|
||||
/// User ID for rate limiting tracking
|
||||
pub user_id: Option<String>,
|
||||
/// Client IP for rate limiting
|
||||
pub client_ip: Option<String>,
|
||||
/// 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<String, OllamaError> {
|
||||
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<std::sync::Mutex<std::collections::HashMap<String, Vec<std::time::Instant>>>>,
|
||||
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<std::sync::atomic::AtomicU32>,
|
||||
threshold: u32,
|
||||
reset_after: Duration,
|
||||
last_failure: std::sync::Arc<std::sync::Mutex<Option<std::time::Instant>>>,
|
||||
}
|
||||
|
||||
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<String>) -> 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<String>) -> Result<GenerateResponse, OllamaError> {
|
||||
// 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<GenerateResponse, OllamaError> {
|
||||
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::<GenerateResponse>()
|
||||
|
|
@ -149,6 +413,30 @@ impl OllamaClient {
|
|||
Ok(result)
|
||||
}
|
||||
|
||||
/// Health check endpoint
|
||||
pub async fn health_check(&self) -> Result<bool, OllamaError> {
|
||||
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<String, OllamaError> {
|
||||
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());
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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" }
|
||||
|
|
|
|||
|
|
@ -11,3 +11,7 @@ tracing = { workspace = true }
|
|||
uuid = { workspace = true }
|
||||
chrono = { workspace = true }
|
||||
anyhow = { workspace = true }
|
||||
thiserror = { workspace = true }
|
||||
|
||||
[dev-dependencies]
|
||||
tokio = { workspace = true }
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
167
crates/db/migrations/20260703210000_ai_credits_wallet.up.sql
Normal file
167
crates/db/migrations/20260703210000_ai_credits_wallet.up.sql
Normal file
|
|
@ -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;
|
||||
|
|
@ -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;
|
||||
|
|
@ -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;
|
||||
|
|
@ -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;
|
||||
|
|
@ -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;
|
||||
|
|
@ -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;
|
||||
32
crates/db/migrations/20260706200000_ai_refunds.up.sql
Normal file
32
crates/db/migrations/20260706200000_ai_refunds.up.sql
Normal file
|
|
@ -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;
|
||||
160
crates/db/migrations/20260706300000_ai_coupons_promotions.up.sql
Normal file
160
crates/db/migrations/20260706300000_ai_coupons_promotions.up.sql
Normal file
|
|
@ -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;
|
||||
|
|
@ -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;
|
||||
119
crates/db/migrations/20260706500000_ai_token_cost_engine.up.sql
Normal file
119
crates/db/migrations/20260706500000_ai_token_cost_engine.up.sql
Normal file
|
|
@ -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;
|
||||
166
crates/db/migrations/20260706600000_ai_observability.up.sql
Normal file
166
crates/db/migrations/20260706600000_ai_observability.up.sql
Normal file
|
|
@ -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;
|
||||
888
crates/db/src/models/ai_credits.rs
Normal file
888
crates/db/src/models/ai_credits.rs
Normal file
|
|
@ -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<String>,
|
||||
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<DateTime<Utc>>,
|
||||
pub billing_cycle: String,
|
||||
pub auto_renew: bool,
|
||||
pub current_period_start: DateTime<Utc>,
|
||||
pub current_period_end: DateTime<Utc>,
|
||||
pub status: String,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub updated_at: DateTime<Utc>,
|
||||
}
|
||||
|
||||
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<i32>,
|
||||
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<i32>,
|
||||
pub max_output_tokens: Option<i32>,
|
||||
pub min_plan_code: Option<String>,
|
||||
pub timeout_ms: i32,
|
||||
pub fallback_model: Option<String>,
|
||||
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<AiWallet, sqlx::Error> {
|
||||
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<Option<AiWallet>, 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<Option<AiPlan>, 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<Uuid>,
|
||||
) -> Result<bool, sqlx::Error> {
|
||||
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<AiWallet, AiCreditsError> {
|
||||
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<Option<AiFeatureCost>, 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<AiReservationHold, AiCreditsError> {
|
||||
// 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<i32>,
|
||||
) -> Result<bool, AiCreditsError> {
|
||||
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<bool, AiCreditsError> {
|
||||
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<Uuid>,
|
||||
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<Vec<RefundRow>, 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<String>,
|
||||
pub admin_actor_id: Option<Uuid>,
|
||||
pub created_at: DateTime<Utc>,
|
||||
pub resolved_at: Option<DateTime<Utc>>,
|
||||
pub original_entry_type: String,
|
||||
pub original_credits: i32,
|
||||
}
|
||||
|
|
@ -27,4 +27,5 @@ pub mod designation;
|
|||
pub mod verification;
|
||||
pub mod user_role_profile;
|
||||
pub mod tracecoin_wallet;
|
||||
pub mod ai_credits;
|
||||
|
||||
|
|
|
|||
206
crates/db/tests/ai_credits.rs
Normal file
206
crates/db/tests/ai_credits.rs
Normal file
|
|
@ -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)));
|
||||
}
|
||||
67
crates/db/tests/ai_credits_reaper.rs
Normal file
67
crates/db/tests/ai_credits_reaper.rs
Normal file
|
|
@ -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<Uuid> = 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");
|
||||
}
|
||||
Loading…
Add table
Reference in a new issue