nxtgauge-ai-assistant/docs/ASK_ASH_IMPLEMENTATION_PLAN.md
Ashwin Kumar Sivakumar 4e3180745f
All checks were successful
build-and-release / build (push) Successful in 4m23s
docs: add Ask Ash billing architecture and implementation plan docs
2026-08-15 22:32:13 +05:30

186 lines
24 KiB
Markdown
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# Ask Ash — Remaining Implementation Plan
**Status as of 2026-07-06.** This is an execution checklist, not a redesign — it assumes the architecture in `ASK_ASH_BILLING_ARCHITECTURE.md` and the code already built (Section "Done" below). Each remaining item lists exact file paths, function signatures, and DB shapes so it can be implemented without re-deriving design decisions. Once implemented, hand back for testing/bug-hunting the same way prior phases were verified (real DB, real HTTP, real edge cases — not just `cargo check`).
---
## 0. What's already done (don't redo)
- **Wallet + ledger**: `crates/db/src/models/ai_credits.rs``AiCreditsRepository` with `ensure_wallet`, `try_reserve_credits`, `try_capture_reservation`, `try_release_reservation`, `add_purchased_credits`, `admin_adjust_credits`. Row-locked, idempotent, tested (`crates/db/tests/ai_credits.rs`, `ai_credits_reaper.rs`).
- **Schema**: migrations `20260703210000_ai_credits_wallet`, `20260703220000_ai_credits_seed_features`, `20260705110000_ai_credit_packages` in `crates/db/migrations/`.
- **Charge helper**: `apps/users/src/ai_credits.rs::charge_ai_feature` — reserve → work closure → capture/release, writes `ai_usage_logs`.
- **HTTP (user-facing)**: `apps/users/src/handlers/ai_credits.rs``GET /api/ai-credits/wallet`, `POST /api/ai-credits/charge`, `GET /api/ai/usage/summary`, `GET /api/ai/usage/logs`.
- **HTTP (purchase, PayU)**: `apps/payments/src/ai_credits.rs` + `apps/payments/src/payu.rs``GET /api/ai-credits` (packages), `POST /api/ai-credits/order`, `POST /api/ai-credits/verify`. Real PayU HMAC-SHA512 hash, cross-verified against an independent Python implementation. Also fixed: TraceCoins' `/api/payments/create-order`/`verify` (was calling a Beeceptor mock, not any real gateway; also had a broken `amount` vs `amount_inr` column bug).
- **Admin package CRUD (payments)**: `apps/payments/src/ai_credits.rs::admin_router()``GET/POST /api/admin/ai-credits/packages`, `PATCH /api/admin/ai-credits/packages/{id}`. Nested in `apps/payments/src/main.rs`. **Not yet nested into the gateway routing test suite for the users-side admin endpoints below — those don't exist yet.**
- **Admin credit adjustment primitive**: `AiCreditsRepository::admin_adjust_credits(pool, user_id, amount, is_add, reason, actor_id, idempotency_key)` in `crates/db/src/models/ai_credits.rs` — ADD grants bonus credits, DEDUCT consumes bonus→monthly→purchased in order, mandatory `reason`, ledger entry types `admin_adjustment_credit`/`admin_adjustment_debit`. **HTTP handler not yet written** (see Task 2 below).
- **Gateway routing**: `apps/gateway/src/main.rs` — fixed a real bug where `/api/ai-credits/*` (order/verify/packages, hosted in `payments`) silently fell through to the generic `/api/ai``users` rule and would have 404'd in production. Added `/api/admin/ai-credits/packages``payments` routing. Both covered by `#[cfg(test)] mod tests` in the same file — **extend this test module** when adding the users-side admin routes in Task 2.
- **Cron jobs**: `apps/cron/src/tasks/ai_credits.rs``sweep_expired_reservation_holds` (every 2 min), `reset_stale_daily_ai_usage` (hourly). Verified against the real running binary.
- **JWT crypto fix**: `crates/auth/Cargo.toml` + `crates/contracts/Cargo.toml` — added `features = ["rust_crypto"]` to the `jsonwebtoken` dependency. Was causing every authenticated request to panic its worker thread.
- **Real features wired to credits**: `apps/users/src/handlers/ai.rs::ai_generate_job_field` (title/description/skills/category — all four, one handler) and `::ai_generate_cover_letter`, both migrated off the old `company_ai_usage`/`job_seeker_ai_usage` daily-counter onto `charge_ai_feature`. Fixed a real bug found while testing this: `try_release_reservation` wasn't rolling back `daily_actions_used`/`daily_credits_used`, so a failed request permanently burned daily quota even with no charge — fixed, regression-tested.
- **Frontend fixes**: `nxtgauge-frontend-solid/src/lib/payu.ts` split into `submitPayuCheckout` (takes an already-built order) + `openPayuCheckout` (creates a TraceCoins order, then delegates) — fixes a bug where `CreditsPage.tsx`'s AI-credit checkout was silently creating a second, mismatched order. `CompanyJobsPage.tsx`'s description/skills buttons fixed to call the one real `/api/ai/generate-job-field` endpoint instead of two nonexistent ones.
**Known pre-existing gap, not caused by this work, not yet fixed**: `nxtgauge-admin-solid`'s existing TraceCoins credit-adjustment page (`src/routes/admin/credit.tsx`) calls `/api/admin/credits/{adjust,ledger,reconcile,balance}`**none of these exist anywhere in the backend.** That admin feature doesn't work today, independent of AI credits. Not in scope here unless you want it added — flagging so it's a conscious decision, not a surprise.
---
## Task 1 — Finish the admin endpoints for AI credits wallet (users service)
**Files**: `apps/users/src/handlers/ai_credits.rs` (add handlers), `apps/users/src/handlers/mod.rs` (already has `ai_credits` registered), `apps/users/src/main.rs` (nest a new admin router).
Add to `apps/users/src/handlers/ai_credits.rs`:
```rust
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))
}
```
- **`GET /balance?userId=<uuid>`** — `require_admin(&auth)` gate (same pattern as `apps/payments/src/ai_credits.rs::admin_list_packages`). Parse `userId` query param, call `AiCreditsRepository::ensure_wallet`, return the same `WalletDto` shape already defined in this file (reuse it, don't redefine).
- **`GET /ledger?userId=<uuid>&page=&limit=`** — paginated query against `ai_credit_ledger WHERE wallet_id = (SELECT id FROM user_ai_subscriptions WHERE user_id = $1) ORDER BY created_at DESC`. Response shape `{ data: [...], total: i64 }` to match the existing TraceCoins ledger UI pattern in `nxtgauge-admin-solid/src/routes/admin/credit.tsx` (`GET /api/admin/credits/ledger?userId=`).
- **`POST /adjust`** — body `{ user_id: Uuid, amount: i32, type: "ADD" | "DEDUCT", reason: String, reference_id: Option<Uuid> }` (exact shape of the existing, already-defined-but-backend-missing TraceCoins contract in `credit.tsx` line ~103 — reuse the field names so the admin UI's fetch call is identical apart from the URL). Validate `reason` is non-empty (return 400 `INVALID_REASON` if blank — this is deliberately mandatory, not optional, per the architecture doc's "every admin-issued credit change must be explained" rule). Call `AiCreditsRepository::admin_adjust_credits(pool, body.user_id, body.amount, body.type == "ADD", &body.reason, auth.user_id, idempotency_key)`. Generate a server-side idempotency key if the client doesn't send one (e.g. `format!("admin-adjust:{}:{}", auth.user_id, uuid::Uuid::new_v4())`) — or better, require the frontend to send one so a double-click can't double-adjust (recommended: frontend generates a UUID per form submission attempt, resent unchanged on retry).
- **`GET /reconcile?from=<date>&to=<date>`** — sums `ai_credit_ledger.credits` grouped by `entry_type` within the date range, plus a per-wallet drift check: for a sample (or all, if the table is still small) of wallets touched in the range, recompute `SUM(credits) FROM ai_credit_ledger WHERE wallet_id = X` and compare to the wallet's derived `available_credits()` — flag any mismatch. This is the concrete "ledger is the source of truth, reconcile against it" mechanism from the architecture doc's Section 5.
**Error handling**: reuse `credits_error_response` already in this file; add any new error codes needed (e.g. a 400 for missing/invalid `userId` query param — just return `(StatusCode::BAD_REQUEST, Json(json!({"error": "..."})))` inline, no need for a new `AiCreditsError` variant for pure request-parsing issues).
**Wire it up** in `apps/users/src/main.rs`:
```rust
.nest("/api/admin/ai-credits", handlers::ai_credits::admin_router())
```
Place this near the existing `.nest("/api/ai-credits", handlers::ai_credits::router())` line.
**Gateway**: no gateway change needed — `/api/admin/ai-credits/balance` etc. don't match the `/api/admin/ai-credits/packages` rule just added, so they correctly fall through to the existing generic `/api/admin/``users_url` catch-all. **But add test cases** to `apps/gateway/src/main.rs`'s existing `#[cfg(test)] mod tests` confirming this (the pattern is already there — `admin_ai_credits_routes_split_correctly` already asserts `/api/admin/ai-credits/balance` etc. route to `users` — just make sure the real handlers now exist to back that assumption).
**Verification checklist** (mirror how every prior phase was tested — spin up a throwaway Postgres+Redis via Docker, run migrations, start the real binary, hit it with `curl` + a minted JWT):
- Admin balance for a user with no wallet yet auto-provisions Free plan (same as the user-facing endpoint).
- Ledger pagination returns newest-first, respects `page`/`limit`.
- Adjust ADD: bonus_credits_total increases, ledger row `entry_type='admin_adjustment_credit'`, `credits > 0`.
- Adjust DEDUCT: consumes bonus→monthly→purchased in that order (write a test wallet with credits in multiple pools and confirm the split), rejects with `InsufficientCredits` if amount exceeds `available_credits()`.
- Adjust with blank `reason` → 400, no DB write.
- Adjust twice with the same idempotency key → second call is a no-op (wallet unchanged, no second ledger row) — this exact behavior is already implemented in `admin_adjust_credits`, just confirm the HTTP layer passes the key through correctly.
- Reconcile flags a deliberately-introduced drift (e.g. manually `UPDATE ai_credit_ledger SET credits = credits + 999` on one row in a test DB, confirm the endpoint reports the mismatch).
---
## Task 2 — Admin UI: AI credits tabs in nxtgauge-admin-solid
**Files**: `src/routes/admin/pricing.tsx` (add a tab), `src/routes/admin/credit.tsx` (add a tab). Do not create new route files or new sidebar entries — the existing sidebar entries for "Pricing" and "Credits" (`src/components/AdminSidebar.tsx` lines ~237-248) already navigate to these files; adding a tab is enough.
### 2a. `pricing.tsx` — "AI Credit Packages" tab
Add a third tab value alongside the existing `view` signal's `"packages" | "create"` (check the exact signal name/values in the file — described as "two tabs... via local `view` signal" in the research). New tab: `"ai_packages"`.
- List view: `GET /api/admin/ai-credits/packages` (admin-scoped, shows inactive too) → render as a table, same `.data-table`/`.table-card` classes as the existing TraceCoins package table. Columns: Name, Credits, Price (₹, convert from paise: `price_inr / 100`), Active toggle, Edit.
- Create form: Name*, Description, Credits*, Price (₹)* → convert rupees to paise before `POST /api/admin/ai-credits/packages` (body `{ name, description, credits, price_inr }`).
- Inline edit / toggle active: `PATCH /api/admin/ai-credits/packages/{id}` with whichever fields changed (all optional in the body per Task 1's `UpdatePackageRequest` shape — `{ name?, description?, credits?, price_inr?, is_active? }`).
- Auth: same `authHeaders()` helper already in this file (`sessionStorage.getItem('nxtgauge_admin_access_token')`).
### 2b. `credit.tsx` — "AI Credits" tab
Add a fourth tab alongside the existing `ledger | adjust | reconcile` (`activeTab` signal). New tab: `"ai_credits"`, itself with the same three sub-views (or just reuse three separate tab entries prefixed "AI " — whichever reads more naturally in the existing tab bar; the file already copy-pastes the tab-bar markup between sections, so follow that convention).
- **Balance & Ledger**: `GET /api/admin/ai-credits/balance?userId=`, `GET /api/admin/ai-credits/ledger?userId=`. Same User ID input field as the existing TraceCoins version.
- **Reward/Deduct**: same form fields as the existing TraceCoins "Reward/Deduct" tab (User ID, Amount, Type ADD/DEDUCT, Reason required, Reference ID optional) → `POST /api/admin/ai-credits/adjust`. This is deliberately a copy of the existing form's shape, just pointed at a new endpoint — reuse the component/markup, don't redesign it.
- **Reconcile**: `GET /api/admin/ai-credits/reconcile?from=&to=`.
**Do not touch** the existing TraceCoins tabs/endpoints in either file — those stay exactly as they are (including the fact that their backend doesn't exist yet, per the known gap noted in Section 0; that's explicitly out of scope here).
**Verification**: once Task 1's backend is live, click through every button in both new tabs against a real running admin-solid dev server + real backend (not just a TypeScript compile check) — create a package, edit its price, deactivate it, look up a real user's AI wallet, adjust their credits, confirm the ledger updates.
---
## Task 3 — Ollama security hardening
**Context already gathered, don't re-research**: `nxtgauge-gitops/apps/ollama/base/``deployment.yaml` (`OLLAMA_HOST=0.0.0.0:11434`, resource limits `cpu: 1000m / memory: 1500Mi`, no auth env vars — Ollama has no built-in auth), `service.yaml` (`ClusterIP`, not internet-exposed — no `ingress.yaml` present), no `NetworkPolicy` anywhere restricting which pods can reach it. Cluster is Flux-managed (confirmed via existing `NetworkPolicy` resources in the `flux-system`/`data` namespaces) — **changes must go into the `nxtgauge-gitops` repo and be committed/reconciled through Flux, never applied directly via `kubectl apply` against the live cluster.**
Concrete work items, roughly in priority order:
1. **NetworkPolicy restricting ingress to the `ollama` Service** in `nxtgauge-gitops/apps/ollama/base/networkpolicy.yaml` (new file, add to `kustomization.yaml`'s `resources:` list). Restrict to pods with labels matching the `users`, `payments`, and `nxtgauge-ai-assistant` deployments (check their actual pod labels in `nxtgauge-gitops/apps/*/base/deployment.yaml` before writing the selector — don't guess label values). Confirm the cluster's CNI actually enforces `NetworkPolicy` (k3s ships an optional lightweight controller; this needs verifying against the live cluster, not assumed) before treating this as a real control rather than a documentation-only artifact.
2. **Application-layer prompt/response length caps**, enforced in `apps/users/src/handlers/ai.rs` before calling `call_ollama_inline` — use `ai_feature_costs.max_input_tokens`/`max_output_tokens` (already columns in the schema, currently unenforced anywhere) as a character-count approximation (e.g. `max_chars ≈ max_tokens * 4`) since a real tokenizer isn't wired up. Reject oversized requests with 400 before spending a reservation.
3. **Truncated prompt/response audit logging** — add `prompt_preview`/`response_preview` columns to `ai_usage_logs` (new migration, short varchar, e.g. first 200 chars), populate in `apps/users/src/ai_credits.rs::log_usage`. Mirrors the pattern already used (but currently dead) in the broken `apps/companies/src/handlers/ai.rs` stub — same idea, done for real this time.
4. **TLS between services and Ollama** — recommend but treat as a separate, heavier infra task (cert-manager + Ollama TLS termination or a service mesh); don't block the above items on this.
5. **Health-check + circuit breaker** for Ollama outages — already flagged as an open item in the architecture doc's Section 8; implement as a periodic `GET {OLLAMA_BASE_URL}/` poll (Ollama's root path responds 200 when healthy) feeding a simple in-memory circuit-breaker state in `apps/users`, returning a fast "AI temporarily unavailable" instead of every request timing out individually.
---
## Task 4 — Wire remaining real AI features to credits
Before writing any code, **audit which of these have a real (non-404) backend handler** — several already-shipped frontend buttons turned out to call nonexistent endpoints (found twice now: `CompanyJobsPage.tsx`'s description/skills buttons, and the entire pre-existing `admin/credits/*` page). Grep `apps/users/src/handlers/*.rs` for actual `.route(...)` registrations, don't trust frontend `fetch()` calls as evidence a backend exists.
Candidates from the original feature list: AI Chat, Help Center Assistant, Knowledge Base Search, Resume Improvement, Professional Profile Improvement, AI Suggestions, AI Form Filling, Requirement Analysis, Ticket Summaries, Admin Insights.
For each real handler found, migrate it the same way `ai_generate_job_field`/`ai_generate_cover_letter` were migrated:
1. Identify what daily-quota/rate-limit mechanism it currently uses (if any).
2. Seed a feature cost row in `ai_feature_costs` if one doesn't already exist for it (4 exist today: `help_answer`, `jd_generate`, `cover_letter_generate`, `resume_improve` — check `crates/db/migrations/20260703220000_ai_credits_seed_features.up.sql`).
3. Wrap the existing generation call (`call_ollama_inline` or equivalent) in `crate::ai_credits::charge_ai_feature(...)`, following the exact pattern in `ai_generate_job_field`.
4. Test the failure path specifically (point `OLLAMA_BASE_URL` at an unreachable host, confirm the reservation releases and daily counters roll back — this exact bug was found and fixed once already; don't assume the fix generalizes without testing each new call site).
---
## Task 5 — Refund architecture
New table (migration): `ai_refunds``id, user_id, ledger_debit_entry_id (FK ai_credit_ledger.id), credits_refunded, reason (provider_failure|timeout|validation_failure|manual), initiated_by (system|admin|user_dispute), status (pending|approved|rejected|completed), created_at, resolved_at`.
Two paths:
- **Automatic**: triggered from `charge_ai_feature`'s error path — currently a failed request just releases the *reservation* (never charged in the first place, so nothing to refund). A genuine refund is needed only if a charge was *captured* and the feature then failed asynchronously after the fact (e.g., a downstream step fails after the LLM call succeeded and credits were captured) — audit whether this can currently happen anywhere; if not, automatic refunds may not be needed yet and this table exists for the manual path only.
- **Manual (admin)**: new endpoint `POST /api/admin/ai-credits/refunds` — body `{ user_id, ledger_debit_entry_id, reason, notes? }`, looks up the original debit ledger entry, credits back via a new ledger entry `entry_type='refund'` (extend `AiCreditsRepository` with a `refund_credits` function mirroring `admin_adjust_credits`'s ADD path but tagged distinctly), writes an `ai_refunds` row, requires `require_admin`.
---
## Task 6 — Coupons, promotions, referral, admin-issued credits
New tables (one migration): `ai_coupons` (`id, code UNIQUE, discount_type, value, applicable_package_ids UUID[], applicable_plan_codes TEXT[], max_redemptions, redemptions_used, valid_from, valid_until, is_active`), `ai_coupon_redemptions` (`coupon_id, user_id, redeemed_at, UNIQUE(coupon_id, user_id)` — enforces single-use-per-user at the DB level, not just application logic), `ai_promotions` (`id, name, trigger_type, bonus_credits, applicable_plan_codes TEXT[], valid_from, valid_until, is_active`), `ai_bonus_credits` (`id, user_id, source_type, source_reference_id, credits, expires_at, ledger_entry_id FK`).
Coupon redemption at purchase time: extend `apps/payments/src/ai_credits.rs::create_order` to accept an optional `coupon_code`, validate against `ai_coupons` (active, not expired, redemption count under max, not already redeemed by this user via `ai_coupon_redemptions`), apply the discount to the PayU order amount, and only insert into `ai_coupon_redemptions` after `verify_order` confirms payment success (not at order-creation time, to avoid burning a redemption on an abandoned checkout).
---
## Task 7 — Subscription lifecycle
New table: `ai_subscription_history` (`id, user_id, from_plan_id, to_plan_id, change_type (upgrade|downgrade|cancel|renew), proration_credits, effective_at, created_by, created_at`).
Add columns to `user_ai_subscriptions` (already has `billing_cycle`, `auto_renew` from the initial migration — check before adding duplicates): confirm `status` supports `active|past_due|cancelled|expired` (currently only `'active'` is ever written).
Logic (new module, e.g. `apps/users/src/ai_subscription.rs`):
- **Upgrade**: immediate — set new `plan_id`, new `monthly_credits_total`, compute prorated bonus (`remaining_days / days_in_period × (new_monthly - old_monthly)`), grant via a new ledger entry type `subscription_grant`, write `ai_subscription_history`.
- **Downgrade**: don't change `plan_id` immediately — write `ai_subscription_history` with `effective_at = current_period_end`, and have the (currently-hourly) cron reset job also check for and apply due downgrades when it rolls a wallet into a new period.
- **Trial plans**: add `is_trial`, `trial_days` columns to `ai_plans`; on trial expiry (checked in the same cron job), auto-downgrade to `free` unless the user has since upgraded/purchased.
---
## Task 8 — Credit expiration enforcement
`user_ai_subscriptions.purchased_credits_expire_at` already exists as a column but nothing sets or checks it.
1. Set it in `apps/payments/src/ai_credits.rs::verify_order` when granting purchased credits (e.g. `NOW() + INTERVAL '12 months'` — confirm the actual policy with whoever owns pricing; the architecture doc recommends this but doesn't mandate a specific duration).
2. Add expiry-first consumption order to `try_capture_reservation`/`admin_adjust_credits`'s DEDUCT path — currently consumption order is bonus → monthly → purchased; per the architecture doc's Section 4.1, it should be *expiring-soonest-first* across whichever pools have an expiry, which today only `purchased_credits` does. Given only one pool expires, the existing order (bonus → monthly → purchased) already happens to leave purchased-with-expiry for last, which is backwards — purchased credits closest to expiring should be spent *before* they're lost. This needs revisiting once expiry is actually set.
3. New cron task in `apps/cron/src/tasks/ai_credits.rs`: sweep wallets where `purchased_credits_expire_at < NOW()` and `purchased_credits_total > purchased_credits_used`, write a ledger entry `entry_type='expiration'` for the unused remainder, zero out the expired portion.
---
## Task 9 — Token cost engine / provider cost / margin tracking
`ai_usage_logs` already has `cached_tokens`, `provider_reported_cost`, `internal_cost`, `credit_unit_price_at_time` columns (added in the Phase 1 migration) — none are populated.
1. New table `ai_model_cost_config`: `model_alias, cost_per_1k_input_tokens, cost_per_1k_output_tokens, cost_basis ('compute_amortized'|'provider_metered'), effective_from`.
2. In `apps/users/src/ai_credits.rs::charge_ai_feature`/`log_usage`, after the LLM call, compute `internal_cost = tokens_used × cost_per_1k(model, cost_basis) / 1000` using the config table, and store it. `provider_reported_cost` stays `NULL` until a real per-token-billed provider is added to LiteLLM (self-hosted Ollama has no real per-request cost signal today).
3. Margin (`credits_charged × credit_unit_price internal_cost`) is a read-time analytics computation for the admin dashboard (Task 2/reconcile-adjacent), not stored per-row.
---
## Task 10 — Observability
Per the architecture doc's explicit recommendation: **do not add Prometheus/Grafana** — this cluster doesn't run them (confirmed: only an OpenTelemetry Collector shipping to OpenObserve exists in `nxtgauge-gitops`). Add OTel metrics (`ai_ledger_writes_total`, `ai_reservation_holds_active`, `ai_charge_latency_ms`) emitted from `apps/users`/`apps/payments` via whatever OTel Rust crate is already used elsewhere in this workspace (check for one before adding a new dependency), feeding the existing collector pipeline.
---
## Suggested order
1. Task 1 + 2 (finish the admin dashboard you're already partway through) — no design decisions left, just execution.
2. Task 3 (Ollama security) — real, currently-open vulnerability.
3. Task 4 (wire remaining features) — mostly mechanical once the audit is done, directly increases how much of the system is actually load-bearing.
4. Tasks 59 in whatever order matches business priority — none block each other.
5. Task 10 last — needs the others to exist before there's anything meaningful to instrument.