feat(org): model grants, org memory, MCP gateway (v2 phases A1-3..5)
Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This commit is contained in:
mjalmoudy94
co-authored by
Claude Sonnet 5.5
Claude Sonnet 5.5
parent
c93b25741e
commit
dee2f72c46
@@ -0,0 +1,337 @@
|
||||
"""MCP tokens: issue, list, revoke, verify.
|
||||
|
||||
Owned by Agent 1. One table, one service, three rules that matter:
|
||||
|
||||
1. **A token is shown once.** Only ``sha256(token)`` and a six-character prefix
|
||||
are stored, so a database leak cannot authenticate anyone; the prefix is for
|
||||
the audit trail and the settings screen, and is never enough to log in.
|
||||
2. **Verification is constant-time and total.** ``hmac.compare_digest`` on the
|
||||
hash, and every failure mode — unknown, revoked, expired — returns the same
|
||||
``None``. A caller cannot tell "wrong token" from "revoked token", because a
|
||||
revoked token is not a fact about the token, it is a decision.
|
||||
3. **Scopes are a floor, not a ceiling.** A token carries the scopes it was
|
||||
minted with and nothing else; an empty scope list is "no scopes", never "all".
|
||||
The per-organization switch lives in ``org_policies.mcp_enabled`` (Phase 3).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import hmac
|
||||
import secrets
|
||||
from collections import deque
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import timedelta
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.core.config import settings
|
||||
from app.core.errors import NotFoundError, ValidationError
|
||||
from app.domain.base import new_id, utcnow
|
||||
from app.infrastructure.db.models_org import McpTokenRow
|
||||
|
||||
#: Every scope the gateway understands. A scope nobody granted a token is a
|
||||
#: scope the gateway refuses, whatever the asks say.
|
||||
MCP_SCOPES: tuple[str, ...] = (
|
||||
"app:read",
|
||||
"tasks:read",
|
||||
"tasks:write",
|
||||
"memory:read",
|
||||
"memory:write",
|
||||
"delivery:read",
|
||||
)
|
||||
|
||||
#: Enough entropy that a token is unguessable, short enough to paste into a
|
||||
#: client config file.
|
||||
TOKEN_ENTROPY_BYTES = 24
|
||||
|
||||
#: How the raw token looks to a human: the prefix is what the UI shows back.
|
||||
TOKEN_PREFIX = "smo_"
|
||||
|
||||
#: Displayed characters; ``token_hash`` is what authenticates.
|
||||
PREFIX_DISPLAY_CHARS = 6
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class McpToken:
|
||||
"""A token as the service hands it back — never the raw secret."""
|
||||
|
||||
id: str
|
||||
organization_id: str
|
||||
name: str
|
||||
prefix: str
|
||||
scopes: list[str]
|
||||
created_by_user_id: str | None = None
|
||||
expires_at: object | None = None
|
||||
last_used_at: object | None = None
|
||||
revoked_at: object | None = None
|
||||
created_at: object | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class IssuedToken:
|
||||
"""The one and only response that carries the raw token."""
|
||||
|
||||
token: str
|
||||
record: McpToken
|
||||
|
||||
|
||||
@dataclass(slots=True, kw_only=True)
|
||||
class TokenIdentity:
|
||||
"""What a verified token proves: which organization, and what it may do."""
|
||||
|
||||
organization_id: str
|
||||
token_id: str
|
||||
prefix: str
|
||||
scopes: frozenset[str] = field(default_factory=frozenset)
|
||||
|
||||
def allows(self, scope: str) -> bool:
|
||||
return scope in self.scopes
|
||||
|
||||
|
||||
def as_utc(value):
|
||||
"""Normalise a timestamp read from the database for comparison.
|
||||
|
||||
SQLite (the development and test store) hands back naive datetimes; every
|
||||
expiry and revocation check compares against ``utcnow()``, so a naive value
|
||||
is assumed to be UTC rather than raising in the middle of a refusal.
|
||||
"""
|
||||
if value is None:
|
||||
return None
|
||||
from datetime import UTC
|
||||
|
||||
return value.replace(tzinfo=UTC) if value.tzinfo is None else value
|
||||
|
||||
|
||||
def hash_token(raw: str) -> str:
|
||||
"""The only form of a token that is ever stored."""
|
||||
return hashlib.sha256(raw.encode("utf-8")).hexdigest()
|
||||
|
||||
|
||||
def _mint() -> str:
|
||||
return f"{TOKEN_PREFIX}{secrets.token_urlsafe(TOKEN_ENTROPY_BYTES)}"
|
||||
|
||||
|
||||
def _scopes_of(raw: str) -> list[str]:
|
||||
"""Parse a CSV scope field, keeping only known scopes and their order."""
|
||||
seen: list[str] = []
|
||||
for piece in (raw or "").split(","):
|
||||
scope = piece.strip()
|
||||
if scope and scope in MCP_SCOPES and scope not in seen:
|
||||
seen.append(scope)
|
||||
return seen
|
||||
|
||||
|
||||
def validate_scopes(scopes: list[str] | None) -> list[str]:
|
||||
"""Refuse an unknown scope rather than silently dropping it."""
|
||||
if not scopes:
|
||||
return []
|
||||
unknown = [scope for scope in scopes if scope not in MCP_SCOPES]
|
||||
if unknown:
|
||||
raise ValidationError(f"محدودهٔ ناشناخته: {', '.join(sorted(unknown))}")
|
||||
# De-duplicated, order preserved, so the UI's checkbox set round-trips.
|
||||
return list(dict.fromkeys(scopes))
|
||||
|
||||
|
||||
class TokenRateLimiter:
|
||||
"""In-memory, per-token call budget for the current minute.
|
||||
|
||||
Deliberately process-local: the gateway is one process, and a distributed
|
||||
limiter would need the Redis the deployment does not have. A restart forgives
|
||||
a burst — which is the right trade for a spare process only the owner's own
|
||||
agents reach.
|
||||
"""
|
||||
|
||||
def __init__(self, *, limit_per_minute: int | None = None) -> None:
|
||||
self.limit = int(limit_per_minute or settings.mcp_rate_limit_per_minute)
|
||||
self._calls: dict[str, deque[float]] = {}
|
||||
|
||||
def allow(self, key: str) -> bool:
|
||||
import time
|
||||
|
||||
now = time.monotonic()
|
||||
window_start = now - 60.0
|
||||
bucket = self._calls.setdefault(key, deque())
|
||||
while bucket and bucket[0] < window_start:
|
||||
bucket.popleft()
|
||||
if len(bucket) >= self.limit:
|
||||
return False
|
||||
bucket.append(now)
|
||||
return True
|
||||
|
||||
def reset(self) -> None:
|
||||
self._calls.clear()
|
||||
|
||||
|
||||
class McpService:
|
||||
"""Token lifecycle and verification for the MCP gateway."""
|
||||
|
||||
def __init__(self, session: AsyncSession) -> None:
|
||||
self.session = session
|
||||
|
||||
# -- lifecycle -----------------------------------------------------------
|
||||
|
||||
async def issue(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
name: str,
|
||||
scopes: list[str],
|
||||
created_by_user_id: str | None = None,
|
||||
ttl_days: int | None = None,
|
||||
) -> IssuedToken:
|
||||
clean_name = name.strip()
|
||||
if not clean_name:
|
||||
raise ValidationError("نام توکن نمیتواند خالی باشد.")
|
||||
wanted = validate_scopes(scopes)
|
||||
if not wanted:
|
||||
raise ValidationError("حداقل یک محدودهٔ دسترسی انتخاب کنید.")
|
||||
days = int(ttl_days or settings.mcp_token_ttl_days)
|
||||
if days < 1:
|
||||
raise ValidationError("مدت اعتبار توکن نامعتبر است.")
|
||||
|
||||
raw = _mint()
|
||||
row = McpTokenRow(
|
||||
id=new_id(),
|
||||
organization_id=organization_id,
|
||||
name=clean_name[:120],
|
||||
token_hash=hash_token(raw),
|
||||
prefix=raw[: len(TOKEN_PREFIX) + PREFIX_DISPLAY_CHARS],
|
||||
scopes=",".join(wanted),
|
||||
created_by_user_id=created_by_user_id,
|
||||
expires_at=utcnow() + timedelta(days=days),
|
||||
)
|
||||
self.session.add(row)
|
||||
await self.session.flush()
|
||||
return IssuedToken(token=raw, record=_to_domain(row))
|
||||
|
||||
async def list(self, organization_id: str) -> list[McpToken]:
|
||||
rows = (
|
||||
(
|
||||
await self.session.execute(
|
||||
select(McpTokenRow)
|
||||
.where(McpTokenRow.organization_id == organization_id)
|
||||
.order_by(McpTokenRow.created_at.desc())
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
return [_to_domain(row) for row in rows]
|
||||
|
||||
async def revoke(self, *, organization_id: str, token_id: str) -> None:
|
||||
"""Revoke a token. Revoking twice is a 404, not a silent success — the
|
||||
same idempotence rule the grant endpoints follow, so a caller can tell
|
||||
"already gone" from "never existed" only by looking, not by guessing."""
|
||||
row = await self._row(organization_id, token_id)
|
||||
if row is None or row.revoked_at is not None:
|
||||
raise NotFoundError("توکن یافت نشد.")
|
||||
row.revoked_at = utcnow()
|
||||
await self.session.flush()
|
||||
|
||||
# -- verification --------------------------------------------------------
|
||||
|
||||
async def verify(self, raw_token: str) -> TokenIdentity | None:
|
||||
"""Resolve a raw token, or ``None`` for every kind of refusal.
|
||||
|
||||
The lookup is by hash (the token is never stored), and the comparison is
|
||||
constant-time so a token that shares a prefix with a real one cannot be
|
||||
discovered by timing. An expired or revoked token is simply not an
|
||||
identity any more.
|
||||
"""
|
||||
if not raw_token or not raw_token.startswith(TOKEN_PREFIX):
|
||||
return None
|
||||
digest = hash_token(raw_token)
|
||||
rows = (
|
||||
(
|
||||
await self.session.execute(
|
||||
select(McpTokenRow).where(
|
||||
McpTokenRow.token_hash == digest,
|
||||
McpTokenRow.revoked_at.is_(None),
|
||||
)
|
||||
)
|
||||
)
|
||||
.scalars()
|
||||
.all()
|
||||
)
|
||||
if not rows:
|
||||
return None
|
||||
row = rows[0]
|
||||
if not hmac.compare_digest(row.token_hash, digest):
|
||||
return None
|
||||
if as_utc(row.expires_at) is not None and as_utc(row.expires_at) <= utcnow():
|
||||
return None
|
||||
return TokenIdentity(
|
||||
organization_id=row.organization_id,
|
||||
token_id=row.id,
|
||||
prefix=row.prefix,
|
||||
scopes=frozenset(_scopes_of(row.scopes)),
|
||||
)
|
||||
|
||||
async def creator_of(self, token_id: str) -> str | None:
|
||||
"""The user who minted this token — the accountable actor for its writes.
|
||||
|
||||
A task created through the gateway is attributed to the administrator who
|
||||
created the token, not to a ghost user: the token is a delegated key, and
|
||||
delegation never erases who holds it.
|
||||
"""
|
||||
row = await self.session.get(McpTokenRow, token_id)
|
||||
return row.created_by_user_id if row else None
|
||||
|
||||
async def touch(self, token_id: str) -> None:
|
||||
"""Record last use. Never allowed to fail a call."""
|
||||
row = await self.session.get(McpTokenRow, token_id)
|
||||
if row is not None:
|
||||
row.last_used_at = utcnow()
|
||||
await self.session.flush()
|
||||
|
||||
async def count_per_org(self) -> dict[str, int]:
|
||||
"""Active tokens per organization, for the owner's /admin/mcp screen."""
|
||||
rows = (
|
||||
await self.session.execute(
|
||||
select(McpTokenRow.organization_id).where(
|
||||
McpTokenRow.revoked_at.is_(None)
|
||||
)
|
||||
)
|
||||
).scalars()
|
||||
counts: dict[str, int] = {}
|
||||
for org_id in rows:
|
||||
counts[org_id] = counts.get(org_id, 0) + 1
|
||||
return counts
|
||||
|
||||
async def _row(self, organization_id: str, token_id: str) -> McpTokenRow | None:
|
||||
return await self.session.scalar(
|
||||
select(McpTokenRow).where(
|
||||
McpTokenRow.organization_id == organization_id,
|
||||
McpTokenRow.id == token_id,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _to_domain(row: McpTokenRow) -> McpToken:
|
||||
return McpToken(
|
||||
id=row.id,
|
||||
organization_id=row.organization_id,
|
||||
name=row.name,
|
||||
prefix=row.prefix,
|
||||
scopes=_scopes_of(row.scopes),
|
||||
created_by_user_id=row.created_by_user_id,
|
||||
expires_at=row.expires_at,
|
||||
last_used_at=row.last_used_at,
|
||||
revoked_at=row.revoked_at,
|
||||
created_at=row.created_at,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"MCP_SCOPES",
|
||||
"TOKEN_PREFIX",
|
||||
"IssuedToken",
|
||||
"McpService",
|
||||
"McpToken",
|
||||
"TokenIdentity",
|
||||
"TokenRateLimiter",
|
||||
"hash_token",
|
||||
"validate_scopes",
|
||||
]
|
||||
@@ -0,0 +1,535 @@
|
||||
"""Org memory use cases: what an organization remembers, and how it learns.
|
||||
|
||||
This module is the real ``MemoryPort`` implementation (see
|
||||
``container_org.build_memory``). The engine calls it three ways — build a
|
||||
context for a task, record something a human taught, and react to a finished
|
||||
task — and this service owns three invariants around those calls:
|
||||
|
||||
1. **Memory is capped.** One organization may spend at most
|
||||
``MEMORY_ORG_CHAR_CAP`` characters on live memory. When a write would cross
|
||||
it, the oldest unpinned *lesson* is evicted first — a lesson is disposable by
|
||||
definition, a rule or a preference is not.
|
||||
2. **Teaching the same thing twice teaches nothing.** A write whose content (or
|
||||
whose key) already exists updates that entry instead of adding a duplicate,
|
||||
which is also what lets "the port is 8001" correct "the port is 8000".
|
||||
3. **A reflection can never break the engine.** ``on_task_finished`` asks one
|
||||
cheap model for what the task taught. If the call fails, the error is
|
||||
recorded and the function returns: the task is finished, and the org simply
|
||||
learned nothing this time.
|
||||
|
||||
Nothing here returns a model name, a provider or an infrastructure detail — the
|
||||
engine's contract is a block of text.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.application.ports import ChatMessage, ModelProviderPort, UsageRecord
|
||||
from app.application.ports.memory import TaskOutcome
|
||||
from app.application.ports.model_provider import ModelResponse
|
||||
from app.core.config import settings
|
||||
from app.core.errors import NotFoundError, PermissionDeniedError, ValidationError
|
||||
from app.core.redact import redact_text
|
||||
from app.domain.ai_control.models import RateUnit, RoutingMode
|
||||
from app.domain.base import utcnow
|
||||
from app.domain.identity.models import ROLE_RANK, Role
|
||||
from app.domain.memory.models import (
|
||||
CONFIRMED_CONFIDENCE,
|
||||
KIND_LABELS,
|
||||
MAX_WRITES_PER_TASK,
|
||||
REFLECTION_CONFIDENCE,
|
||||
Memory,
|
||||
MemoryKind,
|
||||
MemorySource,
|
||||
coerce_kind,
|
||||
coerce_source,
|
||||
content_hash,
|
||||
lesson_expiry,
|
||||
normalise,
|
||||
relevance,
|
||||
tokenise,
|
||||
validate_confidence,
|
||||
validate_content,
|
||||
validate_key,
|
||||
visible_kinds,
|
||||
)
|
||||
from app.infrastructure.error_log import record_error
|
||||
from app.infrastructure.repositories.ai_control import (
|
||||
AiModelRepository,
|
||||
ModelCredentialRepository,
|
||||
RateCardRepository,
|
||||
)
|
||||
from app.infrastructure.repositories.audit import UsageLedgerRepository
|
||||
from app.infrastructure.repositories.memory import MemoryRepository
|
||||
|
||||
#: What a reflection is asked to produce. Deliberately narrow: two lessons and a
|
||||
#: code-map paragraph, in a shape the parser can read without guessing.
|
||||
REFLECTION_SYSTEM = (
|
||||
"تو بازبین کارهای یک سازمان هستی. حداکثر دو درس آموختنی و یک پاراگراف نقشهٔ "
|
||||
"کد بنویس. هر درس یک خط است و فقط یک واقعیت اجرایی را میگوید — نه تعریف. "
|
||||
"اگر درسی نیست، فقط کلمهٔ «هیچ» را بنویس."
|
||||
)
|
||||
|
||||
REFLECTION_TEMPLATE = """کار: {title}
|
||||
وضعیت: {status}
|
||||
خلاصه: {summary}
|
||||
فایلهای تغییرکرده: {files}
|
||||
خطاهای دروازه: {failures}
|
||||
|
||||
پاسخ را دقیقاً در این قالب بنویس:
|
||||
درس ۱: ...
|
||||
درس ۲: ...
|
||||
نقشهٔ کد: ..."""
|
||||
|
||||
#: The block header. Short on purpose: the block is a prompt fragment, and a
|
||||
#: prompt fragment pays for every character.
|
||||
CONTEXT_HEADER = "یادداشتهای سازمان:"
|
||||
|
||||
|
||||
class MemoryService:
|
||||
"""Recalls, teaches, corrects and prunes one organization's memory."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session: AsyncSession,
|
||||
memories: MemoryRepository,
|
||||
models: AiModelRepository,
|
||||
credentials: ModelCredentialRepository,
|
||||
rate_cards: RateCardRepository,
|
||||
ledger: UsageLedgerRepository,
|
||||
provider: ModelProviderPort,
|
||||
) -> None:
|
||||
self.session = session
|
||||
self.memories = memories
|
||||
self.models = models
|
||||
self.credentials = credentials
|
||||
self.rate_cards = rate_cards
|
||||
self.ledger = ledger
|
||||
self.provider = provider
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# MemoryPort #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def record(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
kind: str,
|
||||
content: str,
|
||||
source: str,
|
||||
evidence_task_id: str | None = None,
|
||||
key: str | None = None,
|
||||
created_by_user_id: str | None = None,
|
||||
) -> str | None:
|
||||
"""Store — or correct — one memory. Returns its id, or ``None``.
|
||||
|
||||
``None`` is an answer, not a failure: the write was refused because the
|
||||
task had already taught its three memories. The caller (usually the
|
||||
reflection path) must not treat that as an error.
|
||||
"""
|
||||
cleaned = validate_content(content)
|
||||
# A memory lands in a future prompt. A credential pasted into one would
|
||||
# be broadcast to every employee from then on, so nothing is stored
|
||||
# unredacted — not even what "looks safe".
|
||||
cleaned = redact_text(cleaned)
|
||||
entry_kind = coerce_kind(kind)
|
||||
entry_source = coerce_source(source)
|
||||
entry_key = validate_key(key)
|
||||
digest = content_hash(cleaned)
|
||||
|
||||
if evidence_task_id:
|
||||
writes = await self.memories.count_for_task(organization_id, evidence_task_id)
|
||||
if writes >= MAX_WRITES_PER_TASK:
|
||||
return None
|
||||
|
||||
duplicate = await self.memories.find_duplicate(
|
||||
organization_id, kind=entry_kind, key=entry_key, digest=digest
|
||||
)
|
||||
if duplicate is not None:
|
||||
# Same thought, same organization: correct the entry in place.
|
||||
# Confidence and pinned state survive, because re-reading a rule is
|
||||
# not a retraction of it.
|
||||
duplicate.content = cleaned
|
||||
duplicate.source = entry_source
|
||||
duplicate.confidence = _confidence_for(entry_source, duplicate.confidence)
|
||||
duplicate.expires_at = _expiry_for(entry_kind, duplicate)
|
||||
if created_by_user_id:
|
||||
duplicate.created_by_user_id = duplicate.created_by_user_id or created_by_user_id
|
||||
await self.memories.save(duplicate)
|
||||
return duplicate.id
|
||||
|
||||
memory = Memory(
|
||||
organization_id=organization_id,
|
||||
kind=entry_kind,
|
||||
key=entry_key,
|
||||
content=cleaned,
|
||||
source=entry_source,
|
||||
evidence_task_id=evidence_task_id,
|
||||
created_by_user_id=created_by_user_id,
|
||||
confidence=_confidence_for(entry_source, 1.0),
|
||||
)
|
||||
memory.expires_at = _expiry_for(entry_kind, memory)
|
||||
await self.memories.create(memory)
|
||||
await self._enforce_cap(organization_id)
|
||||
return memory.id
|
||||
|
||||
async def build_context(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
role: str,
|
||||
query: str,
|
||||
budget_chars: int = 2500,
|
||||
) -> str:
|
||||
"""The memory block for one task, within ``budget_chars``.
|
||||
|
||||
Pinned rules and preferences lead: they are standing decisions, not
|
||||
search results. Everything else follows by relevance to the query.
|
||||
Whatever does not fit is dropped — including a pinned entry, because a
|
||||
prompt over budget is worse than an incomplete one.
|
||||
"""
|
||||
kinds = visible_kinds(role)
|
||||
if budget_chars <= 0 or not kinds:
|
||||
return ""
|
||||
memories = await self.memories.list_for_org(organization_id, limit=200)
|
||||
live = [m for m in memories if m.kind in kinds and m.is_live()]
|
||||
if not live:
|
||||
return ""
|
||||
|
||||
query_tokens = tokenise(query)
|
||||
leading = [m for m in live if m.pinned and m.kind in _LEADING_KINDS]
|
||||
leading_ids = {m.id for m in leading}
|
||||
rest = [m for m in live if m.id not in leading_ids]
|
||||
|
||||
leading.sort(key=lambda m: (_kind_order(m.kind), -(m.updated_at or utcnow()).timestamp()))
|
||||
rest.sort(
|
||||
key=lambda m: (
|
||||
-relevance(m, query_tokens),
|
||||
-(m.updated_at or utcnow()).timestamp(),
|
||||
)
|
||||
)
|
||||
|
||||
budget = int(budget_chars)
|
||||
lines: list[str] = []
|
||||
used = len(CONTEXT_HEADER) + 1
|
||||
included: list[Memory] = []
|
||||
for memory in [*leading, *rest]:
|
||||
line = _format_entry(memory)
|
||||
if used + len(line) + 1 > budget:
|
||||
continue
|
||||
used += len(line) + 1
|
||||
lines.append(line)
|
||||
included.append(memory)
|
||||
if not lines:
|
||||
return ""
|
||||
await self.memories.mark_used([m.id for m in included])
|
||||
return "\n".join([CONTEXT_HEADER, *lines])
|
||||
|
||||
async def on_task_finished(self, *, organization_id: str, task: TaskOutcome) -> None:
|
||||
"""Let the organization learn from a finished task. Never raises.
|
||||
|
||||
One cheap completion, at most ``MAX_WRITES_PER_TASK`` memories, all
|
||||
written with ``source=reflection`` and half confidence. A failure is
|
||||
recorded and swallowed: the task is already finished, and the engine must
|
||||
not have to know whether an LLM happened to answer.
|
||||
"""
|
||||
try:
|
||||
answer = await self._reflect(organization_id=organization_id, task=task)
|
||||
except Exception as exc: # noqa: BLE001 - reflection must never raise
|
||||
await record_error(
|
||||
kind="memory.reflection_failed",
|
||||
message=str(exc),
|
||||
exc=exc,
|
||||
organization_id=organization_id,
|
||||
metadata={"task_id": task.task_id, "status": task.status},
|
||||
)
|
||||
return
|
||||
if not answer:
|
||||
return
|
||||
|
||||
for kind, content in _parse_reflection(answer)[:MAX_WRITES_PER_TASK]:
|
||||
await self.record(
|
||||
organization_id=organization_id,
|
||||
kind=str(kind),
|
||||
content=content,
|
||||
source=str(MemorySource.REFLECTION),
|
||||
evidence_task_id=task.task_id,
|
||||
)
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Writes the UI and the MCP gateway expose #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def list(
|
||||
self,
|
||||
organization_id: str,
|
||||
*,
|
||||
kind: str | None = None,
|
||||
search: str | None = None,
|
||||
) -> list[Memory]:
|
||||
parsed = coerce_kind(kind) if kind else None
|
||||
return await self.memories.list_for_org(organization_id, kind=parsed, search=search)
|
||||
|
||||
async def teach(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
user_id: str,
|
||||
role: Role,
|
||||
kind: str,
|
||||
content: str,
|
||||
key: str | None = None,
|
||||
) -> Memory:
|
||||
"""A human teaching the organization — a rule or a preference.
|
||||
|
||||
The role check lives here and not only in the router: this service is
|
||||
also reached by the MCP gateway (Phase 5), and "someone with a token"
|
||||
must hit the same wall as someone with a session.
|
||||
"""
|
||||
entry_kind = coerce_kind(kind)
|
||||
if entry_kind not in _TEACHABLE_KINDS:
|
||||
raise PermissionDeniedError("فقط «قاعده» یا «ترجیح» قابل آموزش است.")
|
||||
if ROLE_RANK[role] < ROLE_RANK[Role.MEMBER]:
|
||||
raise PermissionDeniedError("برای آموزش سازمان باید عضو آن باشید.")
|
||||
memory_id = await self.record(
|
||||
organization_id=organization_id,
|
||||
kind=str(entry_kind),
|
||||
content=content,
|
||||
source=str(MemorySource.USER),
|
||||
key=key,
|
||||
created_by_user_id=user_id,
|
||||
)
|
||||
if memory_id is None: # pragma: no cover - only the task cap returns None
|
||||
raise ValidationError("یادداشت ذخیره نشد.")
|
||||
memory = await self.memories.get(organization_id, memory_id)
|
||||
if memory is None: # pragma: no cover - created one call above
|
||||
raise NotFoundError("یادداشت یافت نشد.")
|
||||
return memory
|
||||
|
||||
async def get(self, organization_id: str, memory_id: str) -> Memory:
|
||||
memory = await self.memories.get(organization_id, memory_id)
|
||||
if memory is None or memory.deleted_at is not None:
|
||||
raise NotFoundError("یادداشت یافت نشد.")
|
||||
return memory
|
||||
|
||||
async def edit(
|
||||
self, *, organization_id: str, memory_id: str, content: str, key: str | None = None
|
||||
) -> Memory:
|
||||
memory = await self.get(organization_id, memory_id)
|
||||
memory.content = validate_content(content)
|
||||
if key is not None:
|
||||
memory.key = validate_key(key)
|
||||
memory.confidence = validate_confidence(memory.confidence)
|
||||
memory.expires_at = _expiry_for(memory.kind, memory)
|
||||
await self.memories.save(memory)
|
||||
return memory
|
||||
|
||||
async def set_pinned(
|
||||
self, *, organization_id: str, memory_id: str, pinned: bool
|
||||
) -> Memory:
|
||||
memory = await self.get(organization_id, memory_id)
|
||||
memory.pinned = bool(pinned)
|
||||
# Pinning is a promise the organization keeps, so a pinned lesson stops
|
||||
# expiring and becomes part of every context.
|
||||
memory.expires_at = _expiry_for(memory.kind, memory)
|
||||
await self.memories.save(memory)
|
||||
return memory
|
||||
|
||||
async def confirm(self, *, organization_id: str, memory_id: str) -> Memory:
|
||||
"""A reviewed lesson becomes settled knowledge: full confidence, no expiry."""
|
||||
memory = await self.get(organization_id, memory_id)
|
||||
if memory.kind is not MemoryKind.LESSON:
|
||||
raise ValidationError("فقط یک «درس» قابل تأیید است.")
|
||||
memory.confidence = CONFIRMED_CONFIDENCE
|
||||
memory.pinned = True
|
||||
memory.expires_at = None
|
||||
await self.memories.save(memory)
|
||||
return memory
|
||||
|
||||
async def delete(self, *, organization_id: str, memory_id: str) -> None:
|
||||
if not await self.memories.soft_delete(organization_id, memory_id):
|
||||
raise NotFoundError("یادداشت یافت نشد.")
|
||||
|
||||
async def expire_memories(self, organization_id: str | None = None) -> int:
|
||||
"""Soft-delete expired lessons. Exposed for the worker hook (see report)."""
|
||||
return await self.memories.expire(organization_id)
|
||||
|
||||
async def stats(self, organization_id: str) -> dict[str, int]:
|
||||
memories = await self.memories.list_for_org(organization_id, limit=200)
|
||||
return {
|
||||
"total": len(memories),
|
||||
"pinned": sum(1 for m in memories if m.pinned),
|
||||
"chars": await self.memories.total_chars(organization_id),
|
||||
}
|
||||
|
||||
# ------------------------------------------------------------------ #
|
||||
# Internals #
|
||||
# ------------------------------------------------------------------ #
|
||||
|
||||
async def _reflect(self, *, organization_id: str, task: TaskOutcome) -> str:
|
||||
"""One cheap completion about the finished task, priced in the ledger."""
|
||||
model = await self.models.pick_for_mode(RoutingMode.CHEAP)
|
||||
if model is None:
|
||||
# No catalogue means nothing to route to. Record it once rather than
|
||||
# letting the engine wonder why nobody learned anything.
|
||||
await record_error(
|
||||
kind="memory.reflection_no_model",
|
||||
message="no active model for the cheap routing mode",
|
||||
organization_id=organization_id,
|
||||
metadata={"task_id": task.task_id},
|
||||
)
|
||||
return ""
|
||||
pool = await self.credentials.list_active_for_provider(model.provider_id)
|
||||
response = await self.provider.complete(
|
||||
model=model,
|
||||
credential=pool[0] if pool else None,
|
||||
credentials=pool or None,
|
||||
messages=[
|
||||
ChatMessage(
|
||||
role="user",
|
||||
content=REFLECTION_TEMPLATE.format(
|
||||
title=task.title,
|
||||
status=task.status,
|
||||
summary=task.summary or "—",
|
||||
files="، ".join(task.changed_files[:20]) or "—",
|
||||
failures="، ".join(task.gate_failures) or "—",
|
||||
),
|
||||
)
|
||||
],
|
||||
system=REFLECTION_SYSTEM,
|
||||
temperature=0.2,
|
||||
max_tokens=600,
|
||||
)
|
||||
await self._record_usage(organization_id=organization_id, response=response)
|
||||
return response.content or ""
|
||||
|
||||
async def _record_usage(self, *, organization_id: str, response: ModelResponse) -> None:
|
||||
"""Price a reflection exactly like any other model call."""
|
||||
for unit, quantity in (
|
||||
(RateUnit.MODEL_INPUT_TOKENS, float(response.input_tokens or 0)),
|
||||
(RateUnit.MODEL_OUTPUT_TOKENS, float(response.output_tokens or 0)),
|
||||
):
|
||||
card = await self.rate_cards.get_by_unit(unit)
|
||||
unit_price = card.customer_price if card else 0.0
|
||||
internal_cost = card.internal_cost_estimate if card else 0.0
|
||||
await self.ledger.record(
|
||||
usage=UsageRecord(
|
||||
organization_id=organization_id,
|
||||
unit=unit,
|
||||
quantity=quantity,
|
||||
unit_price=unit_price,
|
||||
internal_cost_estimate=quantity * internal_cost,
|
||||
total_price=quantity * unit_price,
|
||||
metadata_json={"memory_reflection": True},
|
||||
)
|
||||
)
|
||||
|
||||
async def _enforce_cap(self, organization_id: str) -> int:
|
||||
"""Trim the organization to its character budget. Returns evicted count."""
|
||||
cap = int(settings.memory_org_char_cap)
|
||||
evicted = 0
|
||||
while await self.memories.total_chars(organization_id) > cap:
|
||||
victims = await self.memories.evictable_lessons(organization_id, limit=1)
|
||||
if not victims:
|
||||
# Only rules and preferences are left: the organization chose to
|
||||
# spend its whole budget on standing decisions, and silently
|
||||
# deleting one of those is worse than being one character over.
|
||||
return evicted
|
||||
await self.memories.soft_delete(organization_id, victims[0].id)
|
||||
evicted += 1
|
||||
return evicted
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Pure helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
#: Kinds a human may teach directly; everything else is the system's voice.
|
||||
_TEACHABLE_KINDS = frozenset({MemoryKind.RULE, MemoryKind.PREFERENCE})
|
||||
|
||||
#: A pinned rule outranks a pinned preference: one says what to do, the other
|
||||
#: says how to like it.
|
||||
_LEADING_KINDS = frozenset({MemoryKind.RULE, MemoryKind.PREFERENCE})
|
||||
|
||||
_KIND_ORDER: dict[MemoryKind, int] = {
|
||||
MemoryKind.RULE: 0,
|
||||
MemoryKind.PREFERENCE: 1,
|
||||
}
|
||||
|
||||
|
||||
def _kind_order(kind: MemoryKind) -> int:
|
||||
return _KIND_ORDER.get(kind, 9)
|
||||
|
||||
|
||||
def _confidence_for(source: MemorySource, fallback: float) -> float:
|
||||
"""Reflections get half confidence; a human's entry keeps its own."""
|
||||
if source is MemorySource.REFLECTION:
|
||||
return REFLECTION_CONFIDENCE
|
||||
return validate_confidence(fallback)
|
||||
|
||||
|
||||
def _expiry_for(kind: MemoryKind, memory: Memory):
|
||||
"""Lessons expire; a rule, a preference, a fact and the code map do not."""
|
||||
if kind is not MemoryKind.LESSON:
|
||||
return None
|
||||
return lesson_expiry(
|
||||
pinned=memory.pinned, confirmed=memory.confidence >= CONFIRMED_CONFIDENCE
|
||||
)
|
||||
|
||||
|
||||
def _format_entry(memory: Memory) -> str:
|
||||
label = KIND_LABELS.get(str(memory.kind), "یادداشت")
|
||||
suffix = ""
|
||||
if memory.kind is MemoryKind.LESSON:
|
||||
suffix = f" (اعتماد {int(memory.confidence * 100)}٪)"
|
||||
elif memory.source is MemorySource.REFLECTION:
|
||||
suffix = " (بازتاب)"
|
||||
return f"- [{label}]{suffix} {memory.content.strip()}"
|
||||
|
||||
|
||||
def _parse_reflection(text: str) -> list[tuple[MemoryKind, str]]:
|
||||
"""Read a reflection's answer into ``(kind, content)`` pairs.
|
||||
|
||||
Tolerant by design: the model is asked for a rigid template, but a model is
|
||||
a model. A line that does not start with a known marker is ignored rather
|
||||
than stored, so a chatty answer cannot smuggle in a third lesson.
|
||||
"""
|
||||
entries: list[tuple[MemoryKind, str]] = []
|
||||
for raw_line in (text or "").splitlines():
|
||||
line = normalise(raw_line).lstrip("-*• \t").strip()
|
||||
if not line:
|
||||
continue
|
||||
lowered = line.casefold()
|
||||
for marker, kind in _REFLECTION_MARKERS:
|
||||
if lowered.startswith(marker):
|
||||
body = line[len(marker):].strip(" :،-–")
|
||||
if body:
|
||||
entries.append((kind, body))
|
||||
break
|
||||
return entries
|
||||
|
||||
|
||||
#: Marker → kind. Longest marker first, so «نقشهٔ کد:» cannot be swallowed by a
|
||||
#: shorter prefix, and both the Persian and the English spelling is accepted.
|
||||
_REFLECTION_MARKERS: tuple[tuple[str, MemoryKind], ...] = (
|
||||
("نقشهٔ کد:", MemoryKind.CODE_MAP),
|
||||
("نقشه کد:", MemoryKind.CODE_MAP),
|
||||
("code map:", MemoryKind.CODE_MAP),
|
||||
("درس ۱:", MemoryKind.LESSON),
|
||||
("درس ۱", MemoryKind.LESSON),
|
||||
("درس ۲:", MemoryKind.LESSON),
|
||||
("درس ۲", MemoryKind.LESSON),
|
||||
("درس:", MemoryKind.LESSON),
|
||||
("درس", MemoryKind.LESSON),
|
||||
("lesson 1:", MemoryKind.LESSON),
|
||||
("lesson 1", MemoryKind.LESSON),
|
||||
("lesson 2:", MemoryKind.LESSON),
|
||||
("lesson 2", MemoryKind.LESSON),
|
||||
("lesson:", MemoryKind.LESSON),
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["CONTEXT_HEADER", "MemoryService"]
|
||||
@@ -0,0 +1,456 @@
|
||||
"""Owner model grants and per-organization policy.
|
||||
|
||||
Owned by Agent 1. This is the real ``ModelGrantPort`` implementation (see
|
||||
``container_org.build_model_grants``): the engine asks "which models may this
|
||||
organization use?" and gets a *neutral* answer — a tier (``fast`` /
|
||||
``balanced`` / ``strong`` ≈ «سریع» / «متعادل» / «قوی») and the capability list,
|
||||
never a model name. Model names stay an owner-only concept, visible in the
|
||||
admin grant matrix and nowhere else.
|
||||
|
||||
Two independent switches have to agree before a model is usable:
|
||||
|
||||
* the model itself is active in the catalogue (``ai_models.is_active``), and
|
||||
* the owner granted it to *this* organization (``org_model_grants.is_active``).
|
||||
|
||||
Either one being false means "not granted". That is what makes revoking a model
|
||||
from one organization a real revocation while other organizations keep working.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass, field
|
||||
|
||||
from sqlalchemy import select
|
||||
from sqlalchemy.ext.asyncio import AsyncSession
|
||||
|
||||
from app.application.ports.model_grants import GrantedModel, tier_for
|
||||
from app.core.config import settings
|
||||
from app.core.errors import NotFoundError, ValidationError
|
||||
from app.infrastructure.db.models import AiModelRow
|
||||
from app.infrastructure.repositories.org_grants import (
|
||||
DefaultGrantRecord,
|
||||
DefaultModelGrantRepository,
|
||||
OrgGrantRecord,
|
||||
OrgModelGrantRepository,
|
||||
OrgPolicyRecord,
|
||||
OrgPolicyRepository,
|
||||
)
|
||||
|
||||
#: The only vocabulary an organization's users ever see for a model.
|
||||
TIER_LABELS: dict[str, str] = {
|
||||
"fast": "سریع",
|
||||
"balanced": "متعادل",
|
||||
"strong": "قوی",
|
||||
}
|
||||
|
||||
#: Policy defaults for an organization that has never been configured. Kept in
|
||||
#: code (not only in the row default) so ``get_policy`` can answer for any org
|
||||
#: without writing a row first.
|
||||
DEFAULT_MAX_EMPLOYEES = 5
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class ModelView:
|
||||
"""One catalogue model as the owner sees it, plus its grant state for an org."""
|
||||
|
||||
ai_model_id: str
|
||||
label: str
|
||||
provider_id: str
|
||||
capabilities: list[str] = field(default_factory=list)
|
||||
tier: str = "balanced"
|
||||
tier_label: str = "متعادل"
|
||||
output_price_per_1m: float = 0.0
|
||||
model_active: bool = True
|
||||
granted: bool = False
|
||||
daily_credit_cap: float | None = None
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class OrgGrantView:
|
||||
ai_model_id: str
|
||||
label: str
|
||||
tier: str
|
||||
tier_label: str
|
||||
capabilities: list[str]
|
||||
daily_credit_cap: float | None
|
||||
is_active: bool
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class PolicyView:
|
||||
organization_id: str
|
||||
max_employees: int
|
||||
monthly_credit_budget: float | None
|
||||
automation_enabled: bool
|
||||
mcp_enabled: bool
|
||||
max_employees_limit: int
|
||||
|
||||
|
||||
def _tier_of(row: AiModelRow) -> str:
|
||||
return tier_for(
|
||||
capabilities=list(row.capabilities or []),
|
||||
output_price_per_1m=float(row.output_price_per_1m or 0.0),
|
||||
)
|
||||
|
||||
|
||||
def _tier_label(tier: str) -> str:
|
||||
return TIER_LABELS.get(tier, TIER_LABELS["balanced"])
|
||||
|
||||
|
||||
class ModelGrantService:
|
||||
"""Grants, default grant set and org policy — all in one place."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session: AsyncSession,
|
||||
grants: OrgModelGrantRepository,
|
||||
defaults: DefaultModelGrantRepository,
|
||||
policies: OrgPolicyRepository,
|
||||
) -> None:
|
||||
self.session = session
|
||||
self.grants = grants
|
||||
self.defaults = defaults
|
||||
self.policies = policies
|
||||
|
||||
# -- catalogue ----------------------------------------------------------
|
||||
|
||||
async def _model(self, ai_model_id: str) -> AiModelRow:
|
||||
"""The catalogue row, or a 404. Never returns an inactive model's row.
|
||||
|
||||
Inactive means "the owner retired it": granting it would be meaningless,
|
||||
so it is reported as missing rather than as a validation error.
|
||||
"""
|
||||
row = await self.session.scalar(
|
||||
select(AiModelRow).where(
|
||||
AiModelRow.id == ai_model_id, AiModelRow.is_active
|
||||
)
|
||||
)
|
||||
if row is None:
|
||||
raise NotFoundError("مدل مورد نظر یافت نشد.")
|
||||
return row
|
||||
|
||||
async def list_models(self, *, organization_id: str) -> list[ModelView]:
|
||||
"""The full grant matrix for one organization (owner view)."""
|
||||
models = (
|
||||
await self.session.execute(
|
||||
select(AiModelRow).where(AiModelRow.is_active).order_by(AiModelRow.id)
|
||||
)
|
||||
).scalars().all()
|
||||
active_grants = {
|
||||
g.ai_model_id: g
|
||||
for g in await self.grants.list_for_org(organization_id, active_only=True)
|
||||
}
|
||||
return [
|
||||
self._view(row, active_grants.get(row.id)) for row in models
|
||||
]
|
||||
|
||||
def _view(self, row: AiModelRow, grant: OrgGrantRecord | None) -> ModelView:
|
||||
tier = _tier_of(row)
|
||||
return ModelView(
|
||||
ai_model_id=row.id,
|
||||
label=row.display_name_for_admin or row.internal_name,
|
||||
provider_id=row.provider_id,
|
||||
capabilities=list(row.capabilities or []),
|
||||
tier=tier,
|
||||
tier_label=_tier_label(tier),
|
||||
output_price_per_1m=float(row.output_price_per_1m or 0.0),
|
||||
model_active=bool(row.is_active),
|
||||
granted=grant is not None,
|
||||
daily_credit_cap=grant.daily_credit_cap if grant else None,
|
||||
)
|
||||
|
||||
# -- ModelGrantPort -----------------------------------------------------
|
||||
|
||||
async def granted_models(self, *, organization_id: str) -> list[GrantedModel]:
|
||||
"""Active models with an active grant for this organization."""
|
||||
grants = await self.grants.list_for_org(organization_id, active_only=True)
|
||||
if not grants:
|
||||
return []
|
||||
by_id = {g.ai_model_id: g for g in grants}
|
||||
rows = (
|
||||
await self.session.execute(
|
||||
select(AiModelRow).where(
|
||||
AiModelRow.is_active, AiModelRow.id.in_(list(by_id))
|
||||
)
|
||||
)
|
||||
).scalars().all()
|
||||
granted = [
|
||||
GrantedModel(
|
||||
ai_model_id=row.id,
|
||||
tier=_tier_of(row),
|
||||
capabilities=tuple(row.capabilities or []),
|
||||
daily_credit_cap=by_id[row.id].daily_credit_cap,
|
||||
)
|
||||
for row in rows
|
||||
]
|
||||
# Cheapest first: a caller that picks the head of the list gets the
|
||||
# economical model, which is the sane default for bulk work.
|
||||
granted.sort(key=lambda m: m.daily_credit_cap if m.daily_credit_cap is not None else -1)
|
||||
return granted
|
||||
|
||||
async def is_granted(self, *, organization_id: str, ai_model_id: str) -> bool:
|
||||
grant = await self.grants.get(organization_id, ai_model_id)
|
||||
if grant is None or not grant.is_active:
|
||||
return False
|
||||
exists = await self.session.scalar(
|
||||
select(AiModelRow.id).where(
|
||||
AiModelRow.id == ai_model_id, AiModelRow.is_active
|
||||
)
|
||||
)
|
||||
return exists is not None
|
||||
|
||||
# -- grants -------------------------------------------------------------
|
||||
|
||||
async def list_grants(
|
||||
self, *, organization_id: str, active_only: bool = False
|
||||
) -> list[OrgGrantView]:
|
||||
records = await self.grants.list_for_org(organization_id, active_only=active_only)
|
||||
if not records:
|
||||
return []
|
||||
rows = (
|
||||
await self.session.execute(
|
||||
select(AiModelRow).where(
|
||||
AiModelRow.id.in_([r.ai_model_id for r in records])
|
||||
)
|
||||
)
|
||||
).scalars().all()
|
||||
by_id = {row.id: row for row in rows}
|
||||
views: list[OrgGrantView] = []
|
||||
for record in records:
|
||||
row = by_id.get(record.ai_model_id)
|
||||
if row is None: # deleted from the catalogue: keep it out of the UI
|
||||
continue
|
||||
tier = _tier_of(row)
|
||||
views.append(
|
||||
OrgGrantView(
|
||||
ai_model_id=row.id,
|
||||
label=row.display_name_for_admin or row.internal_name,
|
||||
tier=tier,
|
||||
tier_label=_tier_label(tier),
|
||||
capabilities=list(row.capabilities or []),
|
||||
daily_credit_cap=record.daily_credit_cap,
|
||||
is_active=record.is_active,
|
||||
)
|
||||
)
|
||||
return views
|
||||
|
||||
@staticmethod
|
||||
def _validate_cap(daily_credit_cap: float | None) -> float | None:
|
||||
if daily_credit_cap is None:
|
||||
return None
|
||||
cap = float(daily_credit_cap)
|
||||
if cap < 0:
|
||||
raise ValidationError("سقف اعتبار روزانه نمیتواند منفی باشد.")
|
||||
return cap
|
||||
|
||||
async def grant_model(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
ai_model_id: str,
|
||||
daily_credit_cap: float | None = None,
|
||||
actor_user_id: str | None = None,
|
||||
) -> OrgGrantView:
|
||||
await self._model(ai_model_id)
|
||||
cap = self._validate_cap(daily_credit_cap)
|
||||
await self.grants.upsert(
|
||||
organization_id=organization_id,
|
||||
ai_model_id=ai_model_id,
|
||||
daily_credit_cap=cap,
|
||||
is_active=True,
|
||||
created_by_user_id=actor_user_id,
|
||||
)
|
||||
views = await self.list_grants(organization_id=organization_id, active_only=True)
|
||||
return next(v for v in views if v.ai_model_id == ai_model_id)
|
||||
|
||||
async def revoke_model(
|
||||
self, *, organization_id: str, ai_model_id: str, actor_user_id: str | None = None
|
||||
) -> None:
|
||||
"""Revoke by deactivating, never by deleting.
|
||||
|
||||
The row is kept so the history ("was granted, then revoked") survives
|
||||
and re-granting is idempotent instead of racing a unique constraint.
|
||||
"""
|
||||
existing = await self.grants.get(organization_id, ai_model_id)
|
||||
if existing is None or not existing.is_active:
|
||||
raise NotFoundError("این مدل برای سازمان فعال نیست.")
|
||||
await self.grants.upsert(
|
||||
organization_id=organization_id,
|
||||
ai_model_id=ai_model_id,
|
||||
daily_credit_cap=existing.daily_credit_cap,
|
||||
is_active=False,
|
||||
created_by_user_id=actor_user_id or existing.created_by_user_id,
|
||||
)
|
||||
|
||||
async def set_grant_cap(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
ai_model_id: str,
|
||||
daily_credit_cap: float | None,
|
||||
actor_user_id: str | None = None,
|
||||
) -> OrgGrantView:
|
||||
existing = await self.grants.get(organization_id, ai_model_id)
|
||||
if existing is None or not existing.is_active:
|
||||
raise NotFoundError("این مدل برای سازمان فعال نیست.")
|
||||
return await self.grant_model(
|
||||
organization_id=organization_id,
|
||||
ai_model_id=ai_model_id,
|
||||
daily_credit_cap=daily_credit_cap,
|
||||
actor_user_id=actor_user_id,
|
||||
)
|
||||
|
||||
# -- default grant set --------------------------------------------------
|
||||
|
||||
async def list_defaults(self) -> list[ModelView]:
|
||||
models = (
|
||||
await self.session.execute(
|
||||
select(AiModelRow).where(AiModelRow.is_active).order_by(AiModelRow.id)
|
||||
)
|
||||
).scalars().all()
|
||||
active = {d.ai_model_id: d for d in await self.defaults.list_active()}
|
||||
return [self._view(row, _as_grant(active.get(row.id))) for row in models]
|
||||
|
||||
async def set_default(
|
||||
self,
|
||||
*,
|
||||
ai_model_id: str,
|
||||
daily_credit_cap: float | None = None,
|
||||
is_active: bool = True,
|
||||
actor_user_id: str | None = None,
|
||||
) -> None:
|
||||
await self._model(ai_model_id)
|
||||
await self.defaults.upsert(
|
||||
ai_model_id=ai_model_id,
|
||||
daily_credit_cap=self._validate_cap(daily_credit_cap),
|
||||
is_active=is_active,
|
||||
created_by_user_id=actor_user_id,
|
||||
)
|
||||
|
||||
async def remove_default(self, *, ai_model_id: str) -> None:
|
||||
if not await self.defaults.delete(ai_model_id):
|
||||
raise NotFoundError("این مدل در مجموعهٔ پیشفرض نیست.")
|
||||
|
||||
async def copy_defaults_into(
|
||||
self, *, organization_id: str, actor_user_id: str | None = None
|
||||
) -> int:
|
||||
"""Materialise the default set into a new organization. Returns the count.
|
||||
|
||||
Only the *active* defaults are copied: a default the owner switched off
|
||||
must not silently reappear in the next organization created.
|
||||
"""
|
||||
copied = 0
|
||||
for record in await self.defaults.list_active():
|
||||
await self.grants.upsert(
|
||||
organization_id=organization_id,
|
||||
ai_model_id=record.ai_model_id,
|
||||
daily_credit_cap=record.daily_credit_cap,
|
||||
is_active=True,
|
||||
created_by_user_id=actor_user_id,
|
||||
)
|
||||
copied += 1
|
||||
return copied
|
||||
|
||||
# -- org policy ---------------------------------------------------------
|
||||
|
||||
@staticmethod
|
||||
def _policy_view(record: OrgPolicyRecord | None, organization_id: str) -> PolicyView:
|
||||
limit = max(1, int(settings.engine_max_employees))
|
||||
if record is None:
|
||||
return PolicyView(
|
||||
organization_id=organization_id,
|
||||
max_employees=min(DEFAULT_MAX_EMPLOYEES, limit),
|
||||
monthly_credit_budget=None,
|
||||
automation_enabled=True,
|
||||
mcp_enabled=True,
|
||||
max_employees_limit=limit,
|
||||
)
|
||||
return PolicyView(
|
||||
organization_id=organization_id,
|
||||
max_employees=min(int(record.max_employees), limit),
|
||||
monthly_credit_budget=record.monthly_credit_budget,
|
||||
automation_enabled=bool(record.automation_enabled),
|
||||
mcp_enabled=bool(record.mcp_enabled),
|
||||
max_employees_limit=limit,
|
||||
)
|
||||
|
||||
async def get_org_policy(self, organization_id: str) -> PolicyView:
|
||||
"""Readable by anyone who may manage the org (Agent 2 reads this later)."""
|
||||
return self._policy_view(await self.policies.get(organization_id), organization_id)
|
||||
|
||||
async def set_org_policy(
|
||||
self,
|
||||
*,
|
||||
organization_id: str,
|
||||
max_employees: int | None = None,
|
||||
monthly_credit_budget: float | None = None,
|
||||
automation_enabled: bool | None = None,
|
||||
mcp_enabled: bool | None = None,
|
||||
actor_user_id: str | None = None,
|
||||
) -> PolicyView:
|
||||
current = await self.get_org_policy(organization_id)
|
||||
employees = current.max_employees if max_employees is None else int(max_employees)
|
||||
if employees < 1:
|
||||
raise ValidationError("حداقل تعداد پرسنل یک است.")
|
||||
limit = current.max_employees_limit
|
||||
if employees > limit:
|
||||
raise ValidationError(
|
||||
f"حداکثر تعداد پرسنل نمیتواند از سقف سامانه ({limit}) بیشتر باشد."
|
||||
)
|
||||
budget = (
|
||||
current.monthly_credit_budget
|
||||
if monthly_credit_budget is None
|
||||
else float(monthly_credit_budget)
|
||||
)
|
||||
if budget is not None and budget < 0:
|
||||
raise ValidationError("بودجهٔ ماهانه نمیتواند منفی باشد.")
|
||||
automation = (
|
||||
current.automation_enabled
|
||||
if automation_enabled is None
|
||||
else bool(automation_enabled)
|
||||
)
|
||||
mcp = current.mcp_enabled if mcp_enabled is None else bool(mcp_enabled)
|
||||
record = await self.policies.upsert(
|
||||
organization_id=organization_id,
|
||||
max_employees=employees,
|
||||
monthly_credit_budget=budget,
|
||||
automation_enabled=automation,
|
||||
mcp_enabled=mcp,
|
||||
)
|
||||
return self._policy_view(record, organization_id)
|
||||
|
||||
async def set_mcp_enabled(
|
||||
self, *, organization_id: str, enabled: bool, actor_user_id: str | None = None
|
||||
) -> PolicyView:
|
||||
"""Owner switch for the MCP gateway (Phase 5)."""
|
||||
return await self.set_org_policy(
|
||||
organization_id=organization_id, mcp_enabled=enabled, actor_user_id=actor_user_id
|
||||
)
|
||||
|
||||
async def mcp_enabled(self, organization_id: str) -> bool:
|
||||
return (await self.get_org_policy(organization_id)).mcp_enabled
|
||||
|
||||
|
||||
def _as_grant(record: DefaultGrantRecord | None) -> OrgGrantRecord | None:
|
||||
"""Default rows reuse the org-grant view; only the shape matters here."""
|
||||
if record is None:
|
||||
return None
|
||||
return OrgGrantRecord(
|
||||
id=record.id,
|
||||
organization_id="",
|
||||
ai_model_id=record.ai_model_id,
|
||||
daily_credit_cap=record.daily_credit_cap,
|
||||
is_active=record.is_active,
|
||||
created_by_user_id=record.created_by_user_id,
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_MAX_EMPLOYEES",
|
||||
"ModelGrantService",
|
||||
"ModelView",
|
||||
"OrgGrantView",
|
||||
"PolicyView",
|
||||
"TIER_LABELS",
|
||||
]
|
||||
@@ -243,6 +243,9 @@ class OrgSetupService:
|
||||
if payload.deploy_target_kind:
|
||||
deploy_target_id = await self._setup_deploy_target(org.id, project_row.id, payload)
|
||||
|
||||
# -- default model grants (owner-provided, Phase 3) -------------------
|
||||
await self._copy_default_grants(org.id, user.id)
|
||||
|
||||
await self.session.flush()
|
||||
|
||||
# -- commit so hooks see a consistent world --------------------------
|
||||
@@ -377,6 +380,30 @@ class OrgSetupService:
|
||||
row = await self.deploys.upsert(values=values)
|
||||
return str(row.id)
|
||||
|
||||
async def _copy_default_grants(self, organization_id: str, user_id: str) -> int:
|
||||
"""Copy the platform owner's default grant set into the new organization.
|
||||
|
||||
Runs inside the creation transaction on purpose: an organization that
|
||||
exists without its grants would be an organization whose AI staff cannot
|
||||
think. An empty default set is legitimate (the owner granted nothing).
|
||||
"""
|
||||
from app.application.services.model_grant_service import ModelGrantService
|
||||
from app.infrastructure.repositories.org_grants import (
|
||||
DefaultModelGrantRepository,
|
||||
OrgModelGrantRepository,
|
||||
OrgPolicyRepository,
|
||||
)
|
||||
|
||||
service = ModelGrantService(
|
||||
session=self.session,
|
||||
grants=OrgModelGrantRepository(self.session),
|
||||
defaults=DefaultModelGrantRepository(self.session),
|
||||
policies=OrgPolicyRepository(self.session),
|
||||
)
|
||||
return await service.copy_defaults_into(
|
||||
organization_id=organization_id, actor_user_id=user_id
|
||||
)
|
||||
|
||||
async def _call_hooks(self, *, organization_id: str, project_id: str) -> None:
|
||||
if self.workforce_bootstrap:
|
||||
try:
|
||||
|
||||
Reference in New Issue
Block a user