feat(integration): prepaid billing, engine UI, credits console, direct tasks, v2 final
- Prepaid credit system: org wallet, debit on usage, chat gate, admin credit console - Engine UI: /manager, /tasks, /tasks/[id], /workforce, /dashboard, /traces pages - Direct task creation (HTTP POST /tasks, MCP tasks_create) at pending_manager_review - Pause/resume/cancel/answer actions in engine API and UI - Employee cap (7 per org default), review pipeline guard (active reviewer required) - Scheduler: wide window prevents org starvation; archived employees excluded from cap - Tenant isolation: project-id cross-org 404, engine never uses another tenant's sandbox - Order route fix: manager built with pipeline states + employee profiles - Contract extension: ModelGrantPort.get_org_policy with DefaultOrgPolicy - Engine event types in @agy/contracts (task.step, task.gate, run.log, run.tool, team.changed) - api-engine.ts typed client; api-credit.ts client - Alembic migration c4d5e6f7a8b9: prepaid credit + engine UI tables - 859 tests passed, 1 skipped; ruff/mypy/tsc/eslint/vitest all green Refs: PLAN-v2.md §9 (T1–T43), todo-v2-final.md
This commit is contained in:
@@ -10,6 +10,7 @@ from sqlalchemy.ext.asyncio import AsyncSession
|
||||
from app.application.ports.audit import UsageAggregate, UsageBreakdown, UsageRecord
|
||||
from app.domain.ai_control.models import RateUnit
|
||||
from app.domain.audit.models import AuditEvent, ErrorEvent, QualityMetric, Span, Trace
|
||||
from app.domain.base import new_id
|
||||
from app.infrastructure.db.models import (
|
||||
AuditEventRow,
|
||||
ErrorEventRow,
|
||||
@@ -18,6 +19,7 @@ from app.infrastructure.db.models import (
|
||||
TraceRow,
|
||||
UsageEventRow,
|
||||
)
|
||||
from app.infrastructure.repositories.credit import CreditLedgerRepository
|
||||
|
||||
|
||||
def _audit_to_domain(row: AuditEventRow) -> AuditEvent:
|
||||
@@ -107,14 +109,26 @@ class AuditRepository:
|
||||
|
||||
|
||||
class UsageLedgerRepository:
|
||||
"""Concrete ``UsageLedgerPort``. The ledger is the billing source of truth."""
|
||||
"""Concrete ``UsageLedgerPort``. The ledger is the billing source of truth.
|
||||
|
||||
Prepaid: every usage row also debits its organization's wallet, once, in the
|
||||
same transaction (the debit is keyed on the usage row's id). The ledger stays
|
||||
the source of truth; the wallet is the balance it leaves behind.
|
||||
"""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def record(self, *, usage: UsageRecord) -> None:
|
||||
async def record(self, *, usage: UsageRecord) -> float | None:
|
||||
"""Append one usage row, debit the wallet for it, return the new balance.
|
||||
|
||||
Both writes share the caller's transaction, so a rolled-back request leaves
|
||||
neither behind. Returns ``None`` when the row cost nothing and nothing was
|
||||
debited.
|
||||
"""
|
||||
total = usage.total_price or (usage.quantity * usage.unit_price)
|
||||
row = UsageEventRow(
|
||||
id=new_id(),
|
||||
organization_id=usage.organization_id,
|
||||
project_id=usage.project_id,
|
||||
user_id=usage.user_id,
|
||||
@@ -129,6 +143,12 @@ class UsageLedgerRepository:
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
debit = await CreditLedgerRepository(self.session).debit_for_usage(
|
||||
organization_id=usage.organization_id,
|
||||
usage_event_id=row.id,
|
||||
amount=total,
|
||||
)
|
||||
return debit.balance_after if debit is not None else None
|
||||
|
||||
async def breakdown(
|
||||
self,
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
"""The prepaid credit ledger: every balance change is one atomic update and one history row.
|
||||
|
||||
``credit_wallets.balance`` is the number the engine reads; ``credit_transactions``
|
||||
is the history that explains it. Both change in the caller's transaction, so a
|
||||
rolled-back request leaves neither behind.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from datetime import datetime
|
||||
|
||||
from sqlalchemy import func, select, update
|
||||
from sqlalchemy.exc import IntegrityError
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.domain.base import utcnow
|
||||
from app.domain.billing.models import CreditKind, CreditTransaction
|
||||
from app.infrastructure.db.models import CreditWalletRow
|
||||
from app.infrastructure.db.models_billing import CreditTransactionRow
|
||||
from app.infrastructure.repositories.billing import CreditWalletRepository
|
||||
|
||||
#: Longest history page one call may return; the UI never needs more at once.
|
||||
_MAX_HISTORY = 200
|
||||
|
||||
|
||||
def _transaction(row: CreditTransactionRow) -> CreditTransaction:
|
||||
return CreditTransaction(
|
||||
id=row.id,
|
||||
organization_id=row.organization_id,
|
||||
kind=CreditKind(row.kind),
|
||||
amount=float(row.amount),
|
||||
balance_after=float(row.balance_after),
|
||||
reason=row.reason,
|
||||
actor_user_id=row.actor_user_id,
|
||||
usage_event_id=row.usage_event_id,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
class CreditLedgerRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
async def apply(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
amount: float,
|
||||
kind: CreditKind,
|
||||
reason: str = "",
|
||||
actor_user_id: str | None = None,
|
||||
usage_event_id: str | None = None,
|
||||
) -> CreditTransaction:
|
||||
"""Move the wallet by a signed ``amount`` and record the move.
|
||||
|
||||
The balance changes in one SQL ``UPDATE`` (``balance = balance + amount``),
|
||||
never as a Python read-modify-write, so two writers cannot lose each
|
||||
other's change. The history row records the balance that update left.
|
||||
"""
|
||||
await CreditWalletRepository(self.session).ensure(organization_id)
|
||||
await self.session.execute(
|
||||
update(CreditWalletRow)
|
||||
.where(CreditWalletRow.organization_id == organization_id)
|
||||
.values(balance=CreditWalletRow.balance + amount, updated_at=utcnow())
|
||||
.execution_options(synchronize_session="fetch")
|
||||
)
|
||||
balance_after = (
|
||||
await self.session.execute(
|
||||
select(CreditWalletRow.balance).where(
|
||||
CreditWalletRow.organization_id == organization_id
|
||||
)
|
||||
)
|
||||
).scalar_one()
|
||||
row = CreditTransactionRow(
|
||||
organization_id=organization_id,
|
||||
kind=str(kind),
|
||||
amount=amount,
|
||||
balance_after=float(balance_after),
|
||||
reason=reason[:500],
|
||||
actor_user_id=actor_user_id,
|
||||
usage_event_id=usage_event_id,
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return _transaction(row)
|
||||
|
||||
async def debit_for_usage(
|
||||
self, *, organization_id: str, usage_event_id: str, amount: float
|
||||
) -> CreditTransaction | None:
|
||||
"""Debit the wallet for one usage row, at most once.
|
||||
|
||||
Returns ``None`` when the usage costs nothing or was already debited. The
|
||||
debit runs in a savepoint: when two writers race on the same usage row, the
|
||||
loser's insert is rolled back and it returns ``None`` instead of poisoning
|
||||
the caller's transaction. Any other integrity failure is re-raised.
|
||||
"""
|
||||
if amount <= 0 or await self._debited(organization_id, usage_event_id):
|
||||
return None
|
||||
try:
|
||||
async with self.session.begin_nested():
|
||||
return await self.apply(
|
||||
organization_id=organization_id,
|
||||
amount=-amount,
|
||||
kind=CreditKind.USAGE,
|
||||
usage_event_id=usage_event_id,
|
||||
)
|
||||
except IntegrityError:
|
||||
if await self._debited(organization_id, usage_event_id):
|
||||
return None
|
||||
raise
|
||||
|
||||
async def history(
|
||||
self,
|
||||
organization_id: str,
|
||||
*,
|
||||
kinds: tuple[CreditKind, ...] | None = None,
|
||||
limit: int = 50,
|
||||
) -> list[CreditTransaction]:
|
||||
"""One organization's history, newest first, optionally one kind of entry."""
|
||||
stmt = select(CreditTransactionRow).where(
|
||||
CreditTransactionRow.organization_id == organization_id
|
||||
)
|
||||
if kinds:
|
||||
stmt = stmt.where(CreditTransactionRow.kind.in_([str(k) for k in kinds]))
|
||||
stmt = stmt.order_by(CreditTransactionRow.created_at.desc()).limit(
|
||||
min(limit, _MAX_HISTORY)
|
||||
)
|
||||
return [_transaction(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def last_credit_at_by_org(self) -> dict[str, datetime]:
|
||||
"""When each organization last received credit (usage debits do not count).
|
||||
|
||||
Cross-organization on purpose: only the platform owner's views may use it.
|
||||
"""
|
||||
stmt = (
|
||||
select(CreditTransactionRow.organization_id, func.max(CreditTransactionRow.created_at))
|
||||
.where(CreditTransactionRow.kind != str(CreditKind.USAGE))
|
||||
.group_by(CreditTransactionRow.organization_id)
|
||||
)
|
||||
result = await self.session.execute(stmt)
|
||||
return dict(result.tuples().all())
|
||||
|
||||
async def _debited(self, organization_id: str, usage_event_id: str) -> bool:
|
||||
stmt = select(CreditTransactionRow.id).where(
|
||||
CreditTransactionRow.organization_id == organization_id,
|
||||
CreditTransactionRow.usage_event_id == usage_event_id,
|
||||
)
|
||||
return (await self.session.execute(stmt)).first() is not None
|
||||
@@ -36,7 +36,7 @@ from app.domain.workforce.pipeline import (
|
||||
TaskKind,
|
||||
TaskPipelineState,
|
||||
)
|
||||
from app.infrastructure.db.models import RunRow, UsageEventRow
|
||||
from app.infrastructure.db.models import RunRow, TaskRow, UsageEventRow
|
||||
from app.infrastructure.db.models_engine import (
|
||||
EmployeeProfileRow,
|
||||
EngineTurnRow,
|
||||
@@ -75,6 +75,8 @@ _PATCHABLE = frozenset(
|
||||
"same_diff_count",
|
||||
"retry_count",
|
||||
"waiver_reason",
|
||||
"last_diff",
|
||||
"merged_sha",
|
||||
"parent_task_id",
|
||||
"depends_on",
|
||||
}
|
||||
@@ -97,6 +99,8 @@ def _state(row: TaskPipelineStateRow) -> TaskPipelineState:
|
||||
same_diff_count=row.same_diff_count,
|
||||
retry_count=row.retry_count,
|
||||
waiver_reason=row.waiver_reason,
|
||||
last_diff=row.last_diff,
|
||||
merged_sha=row.merged_sha,
|
||||
parent_task_id=row.parent_task_id,
|
||||
depends_on=list(row.depends_on or []),
|
||||
version=row.version,
|
||||
@@ -185,6 +189,8 @@ class PipelineStateRepository:
|
||||
same_diff_count=state.same_diff_count,
|
||||
retry_count=state.retry_count,
|
||||
waiver_reason=state.waiver_reason,
|
||||
last_diff=state.last_diff,
|
||||
merged_sha=state.merged_sha,
|
||||
parent_task_id=state.parent_task_id,
|
||||
depends_on=list(state.depends_on),
|
||||
version=state.version,
|
||||
@@ -268,6 +274,27 @@ class PipelineStateRepository:
|
||||
)
|
||||
return int((await self.session.execute(stmt)).scalar_one())
|
||||
|
||||
async def count_at_by_org(self, stage: Stage) -> dict[str, int]:
|
||||
"""How many tasks sit at one stage, per organization.
|
||||
|
||||
Cross-organization on purpose: only the platform owner's views may use it.
|
||||
"""
|
||||
stmt = (
|
||||
select(TaskPipelineStateRow.organization_id, func.count())
|
||||
.where(TaskPipelineStateRow.stage == str(stage))
|
||||
.group_by(TaskPipelineStateRow.organization_id)
|
||||
)
|
||||
return {org: int(n) for org, n in (await self.session.execute(stmt)).all()}
|
||||
|
||||
async def stage_counts(self, organization_id: str) -> dict[str, int]:
|
||||
"""How many of one organization's tasks sit at each stage."""
|
||||
stmt = (
|
||||
select(TaskPipelineStateRow.stage, func.count())
|
||||
.where(TaskPipelineStateRow.organization_id == organization_id)
|
||||
.group_by(TaskPipelineStateRow.stage)
|
||||
)
|
||||
return {stage: int(n) for stage, n in (await self.session.execute(stmt)).all()}
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# employee_profiles
|
||||
@@ -410,6 +437,21 @@ class EngineTurnRepository:
|
||||
)
|
||||
return [_turn(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def list_for_org(self, organization_id: str, *, limit: int = 50) -> list[EngineTurn]:
|
||||
"""The newest turns across one organization's tasks.
|
||||
|
||||
Turns carry no organization column, so the tenant filter joins through
|
||||
the task: another organization's turn can never match.
|
||||
"""
|
||||
stmt = (
|
||||
select(EngineTurnRow)
|
||||
.join(TaskRow, EngineTurnRow.task_id == TaskRow.id)
|
||||
.where(TaskRow.organization_id == organization_id)
|
||||
.order_by(EngineTurnRow.started_at.desc())
|
||||
.limit(min(limit, 200))
|
||||
)
|
||||
return [_turn(r) for r in (await self.session.execute(stmt)).scalars()]
|
||||
|
||||
async def finish(
|
||||
self,
|
||||
turn_id: str,
|
||||
|
||||
@@ -39,7 +39,7 @@ class DefaultGrantRecord:
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class OrgPolicyRecord:
|
||||
organization_id: str
|
||||
max_employees: int = 5
|
||||
max_employees: int = 7
|
||||
monthly_credit_budget: float | None = None
|
||||
automation_enabled: bool = True
|
||||
mcp_enabled: bool = True
|
||||
|
||||
@@ -510,6 +510,15 @@ class TaskRepository:
|
||||
await self.session.flush()
|
||||
return _task(row)
|
||||
|
||||
async def append_description(self, task_id: str, text: str) -> Task | None:
|
||||
"""Add text to the end of a task's description (a person's answer, for triage)."""
|
||||
row = await self.session.get(TaskRow, task_id)
|
||||
if row is None:
|
||||
return None
|
||||
row.description = f"{row.description or ''}{text}"
|
||||
await self.session.flush()
|
||||
return _task(row)
|
||||
|
||||
|
||||
class RunRepository:
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
|
||||
Reference in New Issue
Block a user