UI revamp Phase 5d: Budget page, budgets saved per month

backend/ai/budget.py

@@ -55,8 +55,8 @@ def _round_up(avg: float) -> int:
 
 
 def draft_budget(txns: list[dict], today: date) -> list[dict]:
-    """One line per category: the mean over the 3 complete months before today's (only those with any spend),
-    rounded up.
+    """One line per category: the mean over the 3 complete months before today's (only those with any spend, which
+    are the ones basis.months lists), rounded up.
     Empty when fewer than 2 of those months have any spend (not enough history to suggest limits)."""
     months = []
     y, m = today.year, today.month
@@ -65,15 +65,15 @@ def draft_budget(txns: list[dict], today: date) -> list[dict]:
         months.append((y, m))
     months.reverse()
     totals = _totals(txns, set(months))
-    active = len({mo for (_, mo), v in totals.items() if v})  # months with any spend (a new user may have only 2)
-    if active < 2:
+    spent = sorted({mo for (_, mo), v in totals.items() if v})  # months with any spend (a new user may have only 2)
+    if len(spent) < 2:
         return []
     names = _display_names(txns)
     lines = []
     for key in {k for k, _ in totals}:
-        avg = sum(totals.get((key, mo), 0.0) for mo in months) / active
+        avg = sum(totals.get((key, mo), 0.0) for mo in months) / len(spent)
         if avg >= MIN_MONTHLY:
             lines.append({"category": names[key], "limit": _round_up(avg),
-                          "basis": {"months": [f"{yy}-{mm:02d}" for yy, mm in months], "average": round(avg, 2)}})
+                          "basis": {"months": [f"{yy}-{mm:02d}" for yy, mm in spent], "average": round(avg, 2)}})
     lines.sort(key=lambda l: l["limit"], reverse=True)
     return lines[:MAX_LINES]

backend/api/agent.py

@@ -158,16 +158,18 @@ def _decide(proposal_id: str, user_id: str, decision: str, deps: AgentDeps) -> d
         p.status, p.decided_at = decision, now
         applied = None
         if decision == "approved" and p.type == "set_budget":
-            from api.budget import upsert_budget
+            from api.budget import this_month, upsert_budget, valid_month
 
             # the agent can file set_budget with any payload: only a real positive limit is ever applied
             try:
-                limit = float(json.loads(p.payload or "{}")["limit"])
+                payload = json.loads(p.payload or "{}")
+                limit = float(payload["limit"])
             except (KeyError, TypeError, ValueError):
-                limit = math.nan
+                payload, limit = {}, math.nan
             if not math.isfinite(limit) or limit <= 0 or not p.target.strip():
                 raise HTTPException(422, "This budget suggestion has no valid limit; reject it instead.")
-            upsert_budget(session, user_id, p.target.strip(), limit)
+            # from the month the draft was for (a draft made last month still starts there), else this month
+            upsert_budget(session, user_id, p.target.strip(), limit, valid_month(payload.get("month")) or this_month())
             applied = {"category": p.target.strip(), "limit": limit}
         if decision == "approved" and p.type == "update_category":
             category = json.loads(p.payload or "{}").get("category")

backend/api/budget.py

@@ -1,12 +1,14 @@
 """UI revamp Phase 3: budgets. The overall monthly budget is User.monthly_budget (set in onboarding); per-category
 limits live in the budgets table. A draft (ai/budget.py) never writes budgets: it files set_budget proposals that the
-user approves one by one through /api/agent/proposals/{id}/approve (api/agent.py applies them)."""
+user approves one by one through /api/agent/proposals/{id}/approve (api/agent.py applies them).
+Phase 5d: limits are versioned by month (models.Budget); every read is for one month, every write is from one month on."""
 from __future__ import annotations
 
 import json
+import re
 from datetime import date, datetime
 
-from fastapi import APIRouter, Depends
+from fastapi import APIRouter, Depends, Query
 from pydantic import BaseModel, Field
 from sqlalchemy import select, text
 from sqlalchemy.engine import Engine
@@ -18,6 +20,30 @@
 router = APIRouter(prefix="/api/budget", tags=["budget"])
 
 _MONTH = "%b"
+MONTH_RE = r"^\d{4}-(0[1-9]|1[0-2])$"
+
+
+def this_month(today: date | None = None) -> str:
+    return (today or date.today()).strftime("%Y-%m")
+
+
+def valid_month(value) -> str | None:
+    return value if isinstance(value, str) and re.match(MONTH_RE, value) else None
+
+
+def _shift(month: str, n: int) -> str:
+    y, m = divmod(int(month[:4]) * 12 + int(month[5:]) - 1 + n, 12)
+    return f"{y}-{m + 1:02d}"
+
+
+def limits_in_force(rows, month: str) -> dict[str, tuple[str, float]]:
+    """Case-folded category -> (name, limit) for `month`: each category's latest row from that month or before.
+    A limit of 0 is a removal, so the category has none."""
+    latest: dict[str, object] = {}
+    for b in sorted(rows, key=lambda b: b.effective_from or ""):
+        if (b.effective_from or "") <= month:
+            latest[b.category.casefold()] = b
+    return {k: (b.category, b.monthly_limit) for k, b in latest.items() if b.monthly_limit > 0}
 
 
 class CategoryLimit(BaseModel):
@@ -37,31 +63,42 @@ def _txns(conn, user_id: str) -> list[dict]:
         {"u": user_id}).mappings()]
 
 
-def _state(engine: Engine, user_id: str, today: date) -> dict:
+def _spend(engine: Engine, user_id: str, month: str) -> dict[str, float]:
     from ai.budget import month_spend_by_category
-    from models import Budget, User
 
     with engine.connect() as conn:
-        spend = month_spend_by_category(_txns(conn, user_id), today.year, today.month)
+        return month_spend_by_category(_txns(conn, user_id), int(month[:4]), int(month[5:]))
+
+
+def _rows(session: Session, user_id: str) -> list:
+    from models import Budget
+
+    return session.scalars(select(Budget).where(Budget.user_id == user_id)).all()
+
+
+def _state(engine: Engine, user_id: str, month: str) -> dict:
+    from models import User
+
+    spend = _spend(engine, user_id, month)
     by_key = {k.casefold(): v for k, v in spend.items()}
     with Session(engine) as session:
         user = session.get(User, user_id)
-        rows = session.scalars(select(Budget).where(Budget.user_id == user_id)).all()
-        cats = [{"category": b.category, "limit": b.monthly_limit, "spent": round(by_key.get(b.category.casefold(), 0), 2)}
-                for b in rows]
+        cats = [{"category": name, "limit": limit, "spent": round(by_key.get(key, 0), 2)}
+                for key, (name, limit) in limits_in_force(_rows(session, user_id), month).items()]
     cats.sort(key=lambda c: c["spent"] / c["limit"], reverse=True)
-    return {"success": True, "month": today.strftime("%Y-%m"), "monthly_budget": user.monthly_budget if user else None,
+    return {"success": True, "month": month, "monthly_budget": user.monthly_budget if user else None,
             "spent": round(sum(spend.values()), 2), "categories": cats}
 
 
-def upsert_budget(session: Session, user_id: str, category: str, limit: float) -> None:
-    """One row per (user, category), matched case-insensitively so 'food' and 'Food' don't become two budgets."""
+def upsert_budget(session: Session, user_id: str, category: str, limit: float, month: str) -> None:
+    """Sets `category`'s limit from `month` on (0 removes it). One row per (user, category, month), matched
+    case-insensitively so 'food' and 'Food' don't become two budgets."""
     from models import Budget
 
-    rows = session.scalars(select(Budget).where(Budget.user_id == user_id)).all()
-    row = next((b for b in rows if b.category.casefold() == category.casefold()), None)
+    row = next((b for b in _rows(session, user_id)
+                if b.category.casefold() == category.casefold() and b.effective_from == month), None)
     if row is None:
-        session.add(Budget(user_id=user_id, category=category, monthly_limit=limit))
+        session.add(Budget(user_id=user_id, category=category, monthly_limit=limit, effective_from=month))
     else:
         row.monthly_limit, row.updated_at = limit, datetime.utcnow()
 
@@ -85,7 +122,7 @@ def create_draft(engine: Engine, user_id: str, today: date) -> list[dict]:
             months = line["basis"]["months"]
             span = f"{date.fromisoformat(months[0] + '-01'):{_MONTH}}–{date.fromisoformat(months[-1] + '-01'):{_MONTH}}"
             p = Proposal(user_id=user_id, type="set_budget", target=line["category"], risk="low",
-                         payload=json.dumps({"limit": line["limit"], "basis": line["basis"]}),
+                         payload=json.dumps({"limit": line["limit"], "basis": line["basis"], "month": this_month(today)}),
                          reason=f"Averaged {line['basis']['average']:,.0f} a month over {span}, rounded up.",
                          evidence="[]", status="pending")
             session.add(p)
@@ -97,27 +134,46 @@ def create_draft(engine: Engine, user_id: str, today: date) -> list[dict]:
 
 
 @router.get("", dependencies=[Depends(rate_limit())])
-def get_budget(claims: dict = Depends(current_user), engine: Engine = Depends(get_engine)):
-    return _state(engine, claims["sub"], date.today())
+def get_budget(month: str | None = Query(default=None, pattern=MONTH_RE), claims: dict = Depends(current_user),
+               engine: Engine = Depends(get_engine)):
+    return _state(engine, claims["sub"], month or this_month())
+
+
+@router.get("/history", dependencies=[Depends(rate_limit())])
+def get_history(months: int = Query(default=6, ge=1, le=12), claims: dict = Depends(current_user),
+                engine: Engine = Depends(get_engine)):
+    """Budget vs spent for the last `months` months, oldest first, ending with this one. A past month with no
+    category limit in force has no budget (null): spend only."""
+    from models import User
+
+    user_id, now = claims["sub"], this_month()
+    with Session(engine) as session:
+        user = session.get(User, user_id)
+        overall = user.monthly_budget if user else None
+        rows = _rows(session, user_id)
+    from ai.budget import month_spend_by_category
+
+    with engine.connect() as conn:
+        txns = _txns(conn, user_id)
+    out = []
+    for month in (_shift(now, -n) for n in range(months - 1, -1, -1)):
+        limits = limits_in_force(rows, month)
+        budget = (overall or round(sum(v for _, v in limits.values()), 2) or None) if limits or month == now else None
+        spent = sum(month_spend_by_category(txns, int(month[:4]), int(month[5:])).values())
+        out.append({"month": month, "budget": budget, "spent": round(spent, 2)})
+    return {"success": True, "months": out}
 
 
 @router.put("", dependencies=[Depends(rate_limit())])
 def put_budget(body: BudgetUpdate, claims: dict = Depends(current_user), engine: Engine = Depends(get_engine)):
-    from models import Budget
-
-    user_id = claims["sub"]
+    """Changes apply from the current month on; past months keep the limits they had."""
+    user_id, month = claims["sub"], this_month()
     with Session(engine) as session, session.begin():
         if "monthly_budget" in body.model_fields_set:
             ensure_user(session, user_id, claims.get("email")).monthly_budget = body.monthly_budget
         for c in body.categories:
-            name = c.category.strip()
-            if c.limit is None:
-                for b in session.scalars(select(Budget).where(Budget.user_id == user_id)).all():
-                    if b.category.casefold() == name.casefold():
-                        session.delete(b)
-            else:
-                upsert_budget(session, user_id, name, c.limit)
-    return _state(engine, user_id, date.today())
+            upsert_budget(session, user_id, c.category.strip(), c.limit or 0, month)
+    return _state(engine, user_id, month)
 
 
 @router.post("/draft", dependencies=[Depends(rate_limit("10 per minute"))])

backend/models/__init__.py

@@ -302,17 +302,20 @@ class Proposal(db.Model):
 
 class Budget(db.Model):
     """A monthly spending limit for one category (UI revamp Phase 3). The overall budget is User.monthly_budget.
-    Rows are written only by approving a set_budget proposal or by PUT /api/budget."""
+    Rows are written only by approving a set_budget proposal or by PUT /api/budget.
+    Versioned by month (Phase 5d): a month's limit is the latest row with effective_from <= that month, so changing
+    a budget never rewrites past months. A monthly_limit of 0 removes the category from that month on."""
 
     __tablename__ = "budgets"
 
     id = db.Column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
     user_id = db.Column(String(36), nullable=False, index=True)
     category = db.Column(db.String, nullable=False)
     monthly_limit = db.Column(db.Float, nullable=False)
+    effective_from = db.Column(db.String(7))  # "YYYY-MM"; always set by the writers (older rows: migrated)
     updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
 
-    __table_args__ = (db.UniqueConstraint("user_id", "category", name="u_budget_user_category"),)
+    __table_args__ = (db.UniqueConstraint("user_id", "category", "effective_from", name="u_budget_user_category_month"),)
 
 class RecurringState(db.Model):
     """What the user said about a recurring charge on the Bills page (UI revamp Phase 5c): "cancelled" keeps it listed

backend/models/database.py

@@ -1,4 +1,5 @@
 import logging
+from datetime import datetime
 
 from flask import Flask
 from flask_sqlalchemy import SQLAlchemy
@@ -34,6 +35,7 @@ def _migrate_transaction_unique_index(app: Flask) -> None:
 
 # Columns added after their table first shipped (create_all doesn't alter existing tables). All nullable.
 _ADDED_COLUMNS = {
+    "budgets": [("effective_from", "VARCHAR(7)")],  # UI revamp Phase 5d: budgets per month
     "chat_messages": [("thread_id", "VARCHAR(36)")],  # UI revamp Phase 4: chat threads
     "review_items": [("extracted", "TEXT")],  # SPEC-FEEDBACK
     "transactions": [("currency", "VARCHAR(3)"), ("po_number", "VARCHAR"), ("file_key", "VARCHAR"),
@@ -58,6 +60,35 @@ def _add_missing_columns() -> None:
             logger.warning("Column migration for %s skipped: %s", table, e)
 
 
+def _migrate_budget_versions(engine) -> None:
+    """Phase 5d: budgets went from one row per (user, category) to one per month. Dates older rows by their last
+    update and swaps the unique constraint. SQLite can't drop a table constraint, so there the table is rebuilt.
+    Idempotent."""
+    insp = inspect(engine)
+    if "budgets" not in insp.get_table_names():
+        return
+    old = any(set(u["column_names"]) == {"user_id", "category"} for u in insp.get_unique_constraints("budgets"))
+    with engine.begin() as conn:
+        if old and engine.dialect.name == "sqlite":
+            conn.execute(text("ALTER TABLE budgets RENAME TO budgets_pre_5d"))
+            conn.execute(text("CREATE TABLE budgets (id VARCHAR(36) PRIMARY KEY, user_id VARCHAR(36) NOT NULL, "
+                              "category VARCHAR NOT NULL, monthly_limit FLOAT NOT NULL, effective_from VARCHAR(7), "
+                              "updated_at DATETIME)"))
+            have = {c["name"] for c in insp.get_columns("budgets_pre_5d")}
+            eff = "effective_from" if "effective_from" in have else "NULL"
+            conn.execute(text(f"INSERT INTO budgets (id, user_id, category, monthly_limit, effective_from, updated_at) "
+                              f"SELECT id, user_id, category, monthly_limit, {eff}, updated_at FROM budgets_pre_5d"))
+            conn.execute(text("DROP TABLE budgets_pre_5d"))
+            conn.execute(text("CREATE INDEX IF NOT EXISTS ix_budgets_user_id ON budgets (user_id)"))
+        elif old:
+            conn.execute(text("ALTER TABLE budgets DROP CONSTRAINT IF EXISTS u_budget_user_category"))
+        for row_id, updated in conn.execute(text("SELECT id, updated_at FROM budgets WHERE effective_from IS NULL")).all():
+            conn.execute(text("UPDATE budgets SET effective_from = :m WHERE id = :i"),
+                         {"m": str(updated or datetime.utcnow())[:7], "i": row_id})
+        conn.execute(text("CREATE UNIQUE INDEX IF NOT EXISTS u_budget_user_category_month "
+                          "ON budgets (user_id, category, effective_from)"))
+
+
 RAG_TABLES = {"documents", "document_chunks"}
 
 
@@ -90,6 +121,10 @@ def init_db(app: Flask):
             db.metadata.create_all(db.engine, tables=tables)
             _migrate_transaction_unique_index(app)
             _add_missing_columns()
+            try:
+                _migrate_budget_versions(db.engine)
+            except Exception as e:
+                logger.warning("Budget versions migration skipped: %s", e)
             logger.info("Database connected and tables initialized")
     except Exception as e:
         logger.warning("Could not connect to database: %s", e)

backend/tests/test_api_budget.py

@@ -195,3 +195,68 @@ def test_decide_claims_with_a_conditional_update(setup):
     claim = next(i for i, s in enumerate(statements) if s.startswith("update proposals") and "status = ?" in s.split("where", 1)[1])
     writes = [i for i, s in enumerate(statements) if s.startswith("insert into budgets") or s.startswith("insert into audit_events")]
     assert writes and claim < min(writes)
+
+
+def ym(n):
+    return months_back(n)[:7]
+
+
+def put_row(engine, user, category, limit, month):
+    from models import Budget
+
+    with Session(engine) as s, s.begin():
+        s.add(Budget(user_id=user, category=category, monthly_limit=limit, effective_from=month))
+
+
+def test_editing_this_month_keeps_last_months_limits(setup):
+    c, engine = setup
+    put_row(engine, A, "Food", 3000, ym(1))
+    seed(engine, A, [(months_back(1), 2500, "Food")])
+    c.put("/api/budget", headers=bearer(A), json={"categories": [{"category": "Food", "limit": 5000}]})
+    last = c.get(f"/api/budget?month={ym(1)}", headers=bearer(A)).json()
+    assert last["month"] == ym(1) and last["spent"] == 2500
+    assert last["categories"] == [{"category": "Food", "limit": 3000, "spent": 2500}]
+    assert c.get("/api/budget", headers=bearer(A)).json()["categories"][0]["limit"] == 5000
+
+
+def test_removing_a_category_only_applies_from_this_month(setup):
+    c, engine = setup
+    put_row(engine, A, "Food", 3000, ym(1))
+    body = c.put("/api/budget", headers=bearer(A), json={"categories": [{"category": "food", "limit": None}]}).json()
+    assert body["categories"] == []
+    assert c.get(f"/api/budget?month={ym(1)}", headers=bearer(A)).json()["categories"][0]["limit"] == 3000
+
+
+def test_a_bad_month_is_rejected(setup):
+    c, _ = setup
+    for bad in ("2026-13", "26-01", "2026-1", "x"):
+        assert c.get(f"/api/budget?month={bad}", headers=bearer(A)).status_code == 422
+    assert c.get("/api/budget/history?months=0", headers=bearer(A)).status_code == 422
+
+
+def test_history_compares_each_month_with_the_limits_then(setup):
+    c, engine = setup
+    put_row(engine, A, "Food", 3000, ym(2))
+    put_row(engine, A, "Food", 4000, ym(1))
+    seed(engine, A, [(months_back(3), 900, "Food"), (months_back(2), 3500, "Food"), (months_back(1), 1000, "Food")])
+    months = c.get("/api/budget/history?months=4", headers=bearer(A)).json()["months"]
+    assert months == [
+        {"month": ym(3), "budget": None, "spent": 900},  # before any budget: spend only
+        {"month": ym(2), "budget": 3000, "spent": 3500},
+        {"month": ym(1), "budget": 4000, "spent": 1000},
+        {"month": ym(0), "budget": 4000, "spent": 0},
+    ]
+    c.put("/api/budget", headers=bearer(A), json={"monthly_budget": 50000})
+    assert c.get("/api/budget/history?months=1", headers=bearer(A)).json()["months"][0]["budget"] == 50000
+
+
+def test_approving_a_draft_writes_it_from_the_drafts_month(setup):
+    from models import Budget
+
+    c, engine = setup
+    history(engine, A)
+    food = next(p for p in c.post("/api/budget/draft", headers=bearer(A)).json()["proposals"] if p["target"] == "Food")
+    assert food["payload"]["month"] == ym(0)
+    c.post(f"/api/agent/proposals/{food['id']}/approve", headers=bearer(A))
+    with Session(engine) as s:
+        assert [(b.category, b.effective_from) for b in s.scalars(select(Budget))] == [("Food", ym(0))]