diff --git a/.env.example b/.env.example index a966d00..79d2883 100644 --- a/.env.example +++ b/.env.example @@ -44,21 +44,30 @@ CURRENCY_CODE=USD CURRENCY_SYMBOL=$ CURRENCY_DECIMAL_PLACES=4 -# Plugin manifest trust (optional). Unsigned plugins remain allowed by default. -# Set PLUGIN_TRUST_POLICY=required to reject unsigned or invalid signatures. -# JSON object mapping canonical publisher ids to Ed25519 public keys; each key -# may be base64-encoded raw 32 bytes or a PEM SubjectPublicKeyInfo. Invalid key -# material aborts startup instead of silently disabling verification. +# Plugin marketplace, catalog, network, and execution settings (Phase 8-9). +# Defaults are intentionally fail-closed: remote installation is disabled, no +# plugin hosts are reachable, and all callbacks are time-bounded. + +# Publisher trust. "if-present" accepts unsigned plugins but still refuses a +# forged or untrusted signature; set "required" to also refuse unsigned ones. +# PLUGIN_TRUSTED_KEYS is a JSON object mapping canonical publisher ids to +# Ed25519 public keys (base64 raw 32 bytes or PEM SubjectPublicKeyInfo). Invalid +# key material aborts startup instead of silently disabling verification. PLUGIN_TRUST_POLICY=if-present PLUGIN_TRUSTED_KEYS={} # Operator-controlled catalog for POST .../plugins/install-from-catalog # (Phase 9, issue #389). Empty disables remote installation. The file is a JSON # catalog of validated entries: id, version, https source, sha256 integrity, -# optional detached Ed25519 signature, and a full manifest. See -# docs/plugins.md. +# optional detached Ed25519 signature, and a full manifest. See docs/plugins.md. +# The catalog is re-read on each install request; no cache TTL setting is active +# yet (catalog caching is tracked in #377). PLUGIN_CATALOG_FILE= +# Read-only plugin discovery directory used by GET /plugins/discover. It is +# scanned for manifest.json files only; nothing is installed or executed. +PLUGIN_DISCOVERY_DIR=plugins + # Plugin execution budget (Phase 9, issue #408): seconds an event handler or # lifecycle hook may run before it is abandoned (daemon thread), recorded as a # bounded error report, and dispatch continues. 0 disables the bound. diff --git a/README.md b/README.md index ab1450b..bb00d7b 100644 --- a/README.md +++ b/README.md @@ -56,6 +56,26 @@ Open http://localhost:5000 in your browser. - Team collaboration and workspace-level permissions - Extensible plugin architecture for future integrations +## Marketplace and plugin configuration + +Plugin installation and outbound requests are opt-in. The defaults are +fail-closed, so a fresh install cannot reach a marketplace or the network: + +| Setting | Default | Why it is safe | +| ------- | ------- | -------------- | +| `PLUGIN_CATALOG_FILE` | empty | Remote/catalog installation is disabled until an operator points at a vetted catalog. | +| `PLUGIN_TRUST_POLICY` | `if-present` | Accepts unsigned local plugins, but always refuses forged or untrusted signatures. Use `required` in production. | +| `PLUGIN_TRUSTED_KEYS` | `{}` | No publisher is trusted by default; keys are validated at startup. | +| `PLUGIN_NETWORK_ALLOWLIST` | empty | Every plugin outbound request is denied. Add only the hosts a plugin needs. | +| `PLUGIN_NETWORK_HTTPS_ONLY` | `1` | Plugin traffic must use HTTPS. | +| `PLUGIN_NETWORK_ALLOW_PRIVATE` | `0` | Blocks private, loopback, and link-local targets. | +| `PLUGIN_EXECUTION_TIMEOUT` | `5` | Abandons a plugin callback after five seconds so it cannot stall a request. | + +The catalog is read from disk for each install request; there is no active +cache TTL setting yet (caching is tracked in #377). See +[the plugin guide](docs/plugins.md#configuration-reference) for the full +reference, including signature and catalog formats. + ## Documentation - [Chat feature guide](CONTRIBUTING.md#chat-feature-developer-guide) diff --git a/app/__init__.py b/app/__init__.py index f51d466..47ef5e4 100644 --- a/app/__init__.py +++ b/app/__init__.py @@ -133,4 +133,10 @@ def init_db_command(): register_plugins_cli(app) + # Usage-ledger maintenance (historical chat backfill and future retention + # commands). The backfill is bounded and idempotent. + from app.services.usage_cli import register_usage_cli + + register_usage_cli(app) + return app diff --git a/app/services/usage_cli.py b/app/services/usage_cli.py new file mode 100644 index 0000000..2672258 --- /dev/null +++ b/app/services/usage_cli.py @@ -0,0 +1,93 @@ +"""Usage-ledger maintenance commands. + +Usage:: + + flask usage backfill [--batch-size N] [--limit N] [--dry-run] [--json] + +The backfill copies historical ``Message`` token columns into the append-only +usage ledger. It never updates message rows, is safe to re-run, and processes +messages in bounded keyset pages. +""" + +from __future__ import annotations + +import json + +import click +from flask import Flask + +from app.services.usage_ledger import BACKFILL_DEFAULT_BATCH_SIZE, BACKFILL_MAX_BATCH_SIZE + +EXIT_OK = 0 +EXIT_ERROR = 2 + + +def _write_json(data) -> None: + click.echo(json.dumps(data, indent=2, sort_keys=True)) + + +def _fail(message: str, as_json: bool) -> None: + if as_json: + _write_json({"error": message}) + else: + click.echo(message, err=True) + raise click.exceptions.Exit(EXIT_ERROR) + + +def register_usage_cli(app: Flask) -> None: + """Register the ``usage`` CLI command group on ``app``.""" + + @app.cli.group("usage") + def usage_group(): + """Usage-ledger maintenance commands.""" + + @usage_group.command("backfill") + @click.option( + "--batch-size", + type=click.IntRange(1, BACKFILL_MAX_BATCH_SIZE), + default=BACKFILL_DEFAULT_BATCH_SIZE, + show_default=True, + help="Messages processed per database batch.", + ) + @click.option( + "--limit", + type=click.IntRange(min=0), + default=None, + help="Maximum messages to scan in this run (default: all).", + ) + @click.option( + "--dry-run", + is_flag=True, + help="Report what would be written without committing ledger rows.", + ) + @click.option("--json", "as_json", is_flag=True, help="Emit machine-readable JSON.") + def backfill(batch_size: int, limit: int | None, dry_run: bool, as_json: bool) -> int: + """Backfill ledger rows from historical chat-message token counts.""" + from app.services.usage_ledger import backfill_from_messages + + try: + result = backfill_from_messages( + batch_size=batch_size, + limit=limit, + dry_run=dry_run, + ) + except ValueError as exc: + _fail(str(exc), as_json=as_json) + + if as_json: + _write_json(result) + return EXIT_OK + + action = "Dry run" if dry_run else "Backfill" + click.echo( + f"{action} complete: scanned={result['scanned']} " + f"created={result['created']} " + f"skipped_existing={result['skipped_existing']} " + f"skipped_incomplete={result['skipped_incomplete']} " + f"skipped_orphan={result['skipped_orphan']} " + f"batches={result['batches']}" + ) + return EXIT_OK + + +__all__ = ["EXIT_ERROR", "EXIT_OK", "register_usage_cli"] diff --git a/app/services/usage_ledger.py b/app/services/usage_ledger.py index de9834f..4df5b08 100644 --- a/app/services/usage_ledger.py +++ b/app/services/usage_ledger.py @@ -12,8 +12,10 @@ from typing import Any +from sqlalchemy import or_ + from app.extensions import db -from app.models import UsageRecord +from app.models import Message, UsageRecord from app.services import pricing #: Operation families understood by the ledger. Unknown values are grouped @@ -32,6 +34,12 @@ ) UNKNOWN_SOURCE = "unknown" +#: Default/upper bounds for the historical chat-message backfill. The work is +#: paginated by message id so a large conversation history is never loaded at +#: once. +BACKFILL_DEFAULT_BATCH_SIZE = 500 +BACKFILL_MAX_BATCH_SIZE = 5000 + #: Operation outcomes understood by the ledger (see docs/observability.md). STATUS_SUCCESS = "success" STATUS_ERROR = "error" @@ -321,7 +329,166 @@ def breakdown_payload( } +def _message_usage(message: Message) -> dict[str, int] | None: + """Return validated token counts for a message, or ``None`` when absent.""" + values = ( + getattr(message, "prompt_tokens", None), + getattr(message, "completion_tokens", None), + getattr(message, "total_tokens", None), + ) + if all(value is None for value in values): + return None + try: + return normalize_usage(*values) + except ValueError: + return None + + +def _backfill_record_exists( + message: Message, + conversation_id: int, + user_id: int | None, + counts: dict[str, int], +) -> bool: + """Return ``True`` when this message was already copied into the ledger. + + The ledger predates a durable ``message_id`` column, so a backfilled row is + identified by its source, conversation, token counts, and the original + message timestamp. That tuple is stable across runs, which makes the + command idempotent without changing or annotating historical message rows. + """ + if message.created_at is None: + return False + return ( + UsageRecord.query.filter_by( + source=SOURCE_CHAT, + conversation_id=conversation_id, + user_id=user_id, + prompt_tokens=counts["prompt_tokens"], + completion_tokens=counts["completion_tokens"], + total_tokens=counts["total_tokens"], + ) + .filter(UsageRecord.created_at == message.created_at) + .first() + is not None + ) + + +def backfill_from_messages( + *, + batch_size: int = BACKFILL_DEFAULT_BATCH_SIZE, + limit: int | None = None, + dry_run: bool = False, +) -> dict: + """Copy historical chat-message token counts into the usage ledger. + + The operation is keyset-paginated and idempotent: existing message rows + are never modified, and re-running the command only skips rows that were + already backfilled. ``limit`` bounds the number of messages considered in a + single invocation, which is useful for a staged rollout. ``dry_run`` reports + what would be written without committing anything. + + Raises: + ValueError: If ``batch_size``/``limit`` is outside the supported range. + """ + if not isinstance(batch_size, int) or not 1 <= batch_size <= BACKFILL_MAX_BATCH_SIZE: + raise ValueError(f"batch_size must be an integer between 1 and {BACKFILL_MAX_BATCH_SIZE}.") + if limit is not None and (not isinstance(limit, int) or limit < 0): + raise ValueError("limit must be a non-negative integer.") + + stats = { + "source": SOURCE_CHAT, + "scanned": 0, + "created": 0, + "skipped_existing": 0, + "skipped_incomplete": 0, + "skipped_orphan": 0, + "batches": 0, + "batch_size": batch_size, + "limit": limit, + "dry_run": bool(dry_run), + } + last_id = 0 + processed = 0 + + try: + while limit is None or processed < limit: + size = batch_size if limit is None else min(batch_size, limit - processed) + if size <= 0: + break + messages = ( + Message.query.filter(Message.id > last_id) + .filter( + or_( + Message.prompt_tokens.isnot(None), + Message.completion_tokens.isnot(None), + Message.total_tokens.isnot(None), + ) + ) + .order_by(Message.id) + .limit(size) + .all() + ) + if not messages: + break + + stats["batches"] += 1 + for message in messages: + if limit is not None and processed >= limit: + break + last_id = message.id + processed += 1 + stats["scanned"] += 1 + + counts = _message_usage(message) + if counts is None: + stats["skipped_incomplete"] += 1 + continue + conversation = message.conversation + if conversation is None: + stats["skipped_orphan"] += 1 + continue + if _backfill_record_exists( + message, + message.conversation_id, + conversation.user_id, + counts, + ): + stats["skipped_existing"] += 1 + continue + if dry_run: + stats["created"] += 1 + continue + + record = record_usage( + source=SOURCE_CHAT, + prompt_tokens=counts["prompt_tokens"], + completion_tokens=counts["completion_tokens"], + total_tokens=counts["total_tokens"], + # Historical columns do not record whether the provider + # reported the counts or they were estimated. + tokens_estimated=True, + user_id=conversation.user_id, + conversation_id=message.conversation_id, + commit=False, + ) + record.created_at = message.created_at + stats["created"] += 1 + + if not dry_run: + db.session.commit() + if len(messages) < size: + break + except Exception: + db.session.rollback() + raise + + return stats + + __all__ = [ + "BACKFILL_DEFAULT_BATCH_SIZE", + "BACKFILL_MAX_BATCH_SIZE", "SOURCES", "SOURCE_ANALYSIS", "SOURCE_CHAT", @@ -332,6 +499,7 @@ def breakdown_payload( "STATUS_ERROR", "STATUS_SUCCESS", "UNKNOWN_SOURCE", + "backfill_from_messages", "breakdown_by_source", "breakdown_payload", "current_provider_metadata", diff --git a/app/tools/routes.py b/app/tools/routes.py index ff781f5..42d0533 100644 --- a/app/tools/routes.py +++ b/app/tools/routes.py @@ -15,7 +15,7 @@ from app.extensions import db from app.models import AnalyzedFile, Conversation, FileAnalysis, Message -from app.services import analysis, ratelimit +from app.services import analysis, ratelimit, token_usage, usage_ledger from app.services.github import ( GitHubError, get_github_client, @@ -118,19 +118,98 @@ def _read_upload(field: str = "file") -> tuple[str, str] | tuple[None, str]: return secure_filename(file.filename), text -def _run_action(action: str, prompt: str) -> str: - """Run a single AI action and return the model output.""" +def _conversation_from_payload(data: dict) -> Conversation | None: + """Resolve an optional ``conversation_id`` owned by the current user. + + Tools can be launched from a chat conversation. The scope is attached to + the ledger row when it is valid; an unknown or foreign id is ignored so the + tool request still succeeds without leaking another user's conversation. + """ + raw = data.get("conversation_id") + if raw in (None, ""): + return None try: - provider = get_provider() - return provider.complete( - [ - {"role": "system", "content": "You are a helpful AI coding assistant."}, - {"role": "user", "content": prompt}, - ] + conversation_id = int(raw) + except (TypeError, ValueError): + return None + return Conversation.query.filter_by(id=conversation_id, user_id=current_user.id).first() + + +def _complete_with_usage(provider, messages): + """Run a completion and return ``(content, provider_response_or_none)``. + + ``chat`` is preferred because it carries vendor-reported token usage; + providers (or test doubles) that only implement ``complete`` stay supported. + """ + if callable(getattr(provider, "chat", None)): + response = provider.chat(messages) + return response.content, response + return provider.complete(messages), None + + +def _record_tool_usage( + provider, + response, + content: str, + messages, + *, + conversation: Conversation | None = None, +) -> None: + """Append a ``tools`` ledger row without ever failing the tool request.""" + if response is not None: + usage = token_usage.usage_from_response(response, messages) + model = getattr(response, "model", None) + estimated = getattr(provider, "name", None) == "mock" or any( + getattr(response, attribute, None) is None + for attribute in ("prompt_tokens", "completion_tokens", "total_tokens") ) + else: + usage = token_usage.usage_from_text(token_usage.messages_text(messages), content) + model = None + estimated = True + + if not model: + models = getattr(provider, "models", ()) or () + model = models[0] if models else None + + try: + usage_ledger.record_usage( + source=usage_ledger.SOURCE_TOOLS, + prompt_tokens=usage["prompt_tokens"], + completion_tokens=usage["completion_tokens"], + total_tokens=usage["total_tokens"], + tokens_estimated=estimated, + provider=getattr(provider, "name", None), + model=model, + user=current_user, + workspace_id=conversation.workspace_id if conversation is not None else None, + conversation=conversation, + ) + except Exception: # instrumentation must never break the operation + db.session.rollback() + current_app.logger.exception("Could not record AI tool usage") + + +def _run_action( + action: str, + prompt: str, + *, + conversation: Conversation | None = None, +) -> str: + """Run a single AI action, record its usage, and return the model output.""" + messages = [ + {"role": "system", "content": "You are a helpful AI coding assistant."}, + {"role": "user", "content": prompt}, + ] + try: + provider = get_provider() + content, response = _complete_with_usage(provider, messages) except LLMProviderError as exc: return f"[provider error] {exc}" + _record_tool_usage(provider, response, content, messages, conversation=conversation) + return content + @bp.route("/generate", methods=["POST"]) @login_required @@ -144,7 +223,12 @@ def generate(): return jsonify({"error": "A description is required."}), 400 system = ACTION_PROMPTS["generate"].format(language=language) - result = _run_action("generate", f"{system}\n\nRequest: {description}") + conversation = _conversation_from_payload(data) + result = _run_action( + "generate", + f"{system}\n\nRequest: {description}", + conversation=conversation, + ) return jsonify({"result": result, "language": language}) @@ -162,7 +246,12 @@ def code_action(): return jsonify({"error": "Code or diff content is required."}), 400 system = ACTION_PROMPTS[action] - result = _run_action(action, f"{system}\n\nCode:\n{code}") + conversation = _conversation_from_payload(data) + result = _run_action( + action, + f"{system}\n\nCode:\n{code}", + conversation=conversation, + ) return jsonify({"action": action, "result": result}) @@ -214,14 +303,13 @@ def analyze_file(): } ) + messages = [ + {"role": "system", "content": "You are a helpful AI coding assistant."}, + {"role": "user", "content": f"{system}\n\nFile: {filename}\n\nCode:\n{text}"}, + ] try: provider = get_provider() - result = provider.complete( - [ - {"role": "system", "content": "You are a helpful AI coding assistant."}, - {"role": "user", "content": f"{system}\n\nFile: {filename}\n\nCode:\n{text}"}, - ] - ) + result, response = _complete_with_usage(provider, messages) except LLMProviderError as exc: return jsonify( {"filename": filename, "action": action, "result": f"[provider error] {exc}"} @@ -242,6 +330,7 @@ def analyze_file(): ) db.session.add(record) db.session.commit() + _record_tool_usage(provider, response, result, messages) return jsonify( { "filename": filename, diff --git a/docs/plugins.md b/docs/plugins.md index 3641ae4..1cb9e5a 100644 --- a/docs/plugins.md +++ b/docs/plugins.md @@ -462,6 +462,43 @@ existing failure-isolation behavior. Reports are read back owner-scoped via `GET /plugins/api/workspaces//plugin-errors` (owner only, workspace-isolated) and operator-wide via `flask plugins errors [--json]`. +## Configuration reference + +Marketplace and plugin settings are environment variables. The defaults are +deliberately fail-closed: catalog installation is disabled, no publisher is +trusted, and plugin outbound traffic is denied until an operator opts in. + +| Variable | Safe default | Purpose | +| -------- | ------------ | ------- | +| `PLUGIN_CATALOG_FILE` | empty | Path to the validated JSON catalog used by `POST .../plugins/install-from-catalog`. Empty disables catalog installation. | +| `PLUGIN_DISCOVERY_DIR` | `plugins` | Read-only directory scanned for `manifest.json` files by `GET /plugins/discover`. Nothing is installed or executed. | +| `PLUGIN_TRUST_POLICY` | `if-present` | `if-present` accepts unsigned manifests but always refuses forged or untrusted signatures; `required` also refuses unsigned manifests. Use `required` for a hosted marketplace. | +| `PLUGIN_TRUSTED_KEYS` | `{}` | JSON object of canonical publisher ids to Ed25519 public keys (base64 raw 32 bytes or PEM). Invalid key material aborts startup. | +| `PLUGIN_NETWORK_ALLOWLIST` | empty | Comma-separated outbound hosts. Empty denies every plugin request; `*` allows any public host and `*.example.com` allows subdomains. | +| `PLUGIN_NETWORK_HTTPS_ONLY` | `1` | Rejects non-HTTPS plugin requests. Keep enabled outside isolated development. | +| `PLUGIN_NETWORK_ALLOW_PRIVATE` | `0` | Allows private/loopback/link-local targets. Keep disabled in production. | +| `PLUGIN_NETWORK_TIMEOUT` | `15` | Outbound request timeout in seconds. | +| `PLUGIN_NETWORK_MAX_BYTES` | `2097152` | Maximum response body accepted from a plugin request (2 MiB). | +| `PLUGIN_NETWORK_STRICT_DNS` | `0` | When enabled, an unresolvable host is rejected instead of tolerated. | +| `PLUGIN_EXECUTION_TIMEOUT` | `5` | Seconds a plugin event handler or lifecycle hook may run before it is abandoned and recorded as a bounded error. `0` or negative disables the bound. | + +**Catalog cache:** there is no active cache TTL setting yet. The catalog file is +validated and re-read for each install request, so an operator's catalog edit is +visible immediately. A TTL cache and its `PLUGIN_MARKETPLACE_CACHE_TTL` setting +are tracked separately in #377. + +Example production posture: + +```bash +export PLUGIN_CATALOG_FILE=/etc/ai-code-assistant/catalog.json +export PLUGIN_TRUST_POLICY=required +export PLUGIN_TRUSTED_KEYS='{"acme-labs": ""}' +export PLUGIN_NETWORK_ALLOWLIST=plugins.example.com +export PLUGIN_NETWORK_HTTPS_ONLY=1 +export PLUGIN_NETWORK_ALLOW_PRIVATE=0 +export PLUGIN_EXECUTION_TIMEOUT=5 +``` + ## Audit trail (capabilities & state) Security-relevant plugin actions are recorded in the shared, append-only @@ -606,6 +643,9 @@ is reset between tests): - `tests/test_plugin_catalog_install.py` — catalog parsing/validation and the remote-install endpoint: authorization, trust policy, 4xx failure paths, and no implicit grants or code execution. +- `tests/test_marketplace_e2e.py` — offline catalog-to-dispatch coverage: + install, explicit enable, dependency ordering, trust refusal, rollback, and + no implicit grants. - `tests/test_plugin_lock_entries.py` — the `PluginLockEntry` model, scoping, unique constraint, and canonical lockfile serialization. - `tests/test_plugin_audit.py` — the capability/state audit trail. diff --git a/tests/test_marketplace_e2e.py b/tests/test_marketplace_e2e.py new file mode 100644 index 0000000..55f6ed8 --- /dev/null +++ b/tests/test_marketplace_e2e.py @@ -0,0 +1,246 @@ +"""Offline end-to-end tests for the Phase 9 marketplace flow (#423). + +The tests exercise the real catalog loader, install endpoint, enable/disable +endpoint, capability grants, and event dispatcher. They use a local temp catalog +and a synthetic uploaded project, so no network access or real plugin code +execution is involved. +""" + +from __future__ import annotations + +import io +import json +import zipfile + +import pytest + +from app.extensions import db +from app.models import ( + ActivityEvent, + CapabilityGrant, + Plugin, + PluginInstallation, + Workspace, +) +from app.models.activity_event import EVENT_PLUGIN_DENIED +from app.services.events import get_dispatcher + +INTEGRITY = "sha256:" + "a" * 64 + + +def _manifest(plugin_id: str, *, dependencies=None, capabilities=None): + return { + "id": plugin_id, + "name": f"Marketplace {plugin_id}", + "version": "0.1.0", + "description": "Offline marketplace end-to-end fixture", + "author": "tester", + "entry_point": "plugins.test:TestPlugin", + "capabilities": capabilities or ["PROJECT_READ"], + "dependencies": dependencies or [], + } + + +def _entry(plugin_id: str, *, dependencies=None, capabilities=None): + manifest = _manifest( + plugin_id, + dependencies=dependencies, + capabilities=capabilities, + ) + return { + "id": plugin_id, + "version": manifest["version"], + # An opaque local identifier: the loader accepts it and the install + # path never downloads it, which keeps this test fully offline. + "source": f"catalog/{plugin_id}-{manifest['version']}.zip", + "integrity": INTEGRITY, + "manifest": manifest, + } + + +def _write_catalog(tmp_path, entries): + path = tmp_path / "catalog.json" + path.write_text(json.dumps({"version": 1, "plugins": entries}), encoding="utf-8") + return path + + +def _install(client, workspace_id: int, plugin_id: str): + return client.post( + f"/plugins/api/workspaces/{workspace_id}/plugins/install-from-catalog", + json={"plugin_id": plugin_id}, + ) + + +def _enable(client, workspace_id: int, plugin_id: str): + return client.post(f"/plugins/api/workspaces/{workspace_id}/plugins/{plugin_id}/enable") + + +def _disable(client, workspace_id: int, plugin_id: str): + return client.post(f"/plugins/api/workspaces/{workspace_id}/plugins/{plugin_id}/disable") + + +def _grant(client, workspace_id: int, plugin_id: str): + return client.post( + f"/plugins/api/workspaces/{workspace_id}/plugins/{plugin_id}/capabilities", + json={"grant": ["PROJECT_READ"], "revoke": []}, + ) + + +def _zip_bytes(entries): + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w", zipfile.ZIP_DEFLATED) as archive: + for name, data in entries: + archive.writestr(name, data) + return buffer.getvalue() + + +def _upload_project(client, workspace_id: int): + return client.post( + f"/workspaces/api/workspaces/{workspace_id}/projects", + data={ + "file": ( + io.BytesIO(_zip_bytes([("main.py", "print('marketplace')\n")])), + "marketplace.zip", + ) + }, + content_type="multipart/form-data", + ) + + +@pytest.fixture() +def workspace(app, make_user, login): + user = make_user() + login() + workspace = Workspace(user_id=user.id, name="Marketplace E2E") + db.session.add(workspace) + db.session.commit() + return workspace + + +class TestMarketplaceCatalogToDispatch: + def test_install_enable_dependency_order_and_dispatch(self, app, client, workspace, tmp_path): + app.config["PLUGIN_CATALOG_FILE"] = str( + _write_catalog( + tmp_path, + [ + _entry("market-base"), + _entry("market-addon", dependencies=["market-base>=0.1.0"]), + ], + ) + ) + + # The dependent plugin can be installed first (install records metadata + # only), but it cannot be enabled until its dependency is registered. + assert _install(client, workspace.id, "market-addon").status_code == 201 + assert _disable(client, workspace.id, "market-addon").status_code == 200 + denied = _enable(client, workspace.id, "market-addon") + assert denied.status_code == 400 + installation = PluginInstallation.query.filter_by( + plugin_id="market-addon", workspace_id=workspace.id + ).one() + assert installation.enabled is False + + # Install and enable the dependency, then the addon enable succeeds. + assert _install(client, workspace.id, "market-base").status_code == 201 + assert _enable(client, workspace.id, "market-base").status_code == 200 + assert _enable(client, workspace.id, "market-addon").status_code == 200 + + # Grants stay explicit, even after a catalog install. + assert CapabilityGrant.query.filter_by(plugin_id="market-addon").count() == 0 + assert _grant(client, workspace.id, "market-addon").status_code == 200 + + dispatcher = get_dispatcher() + received = [] + + def handler(event): + received.append(event) + + dispatcher.subscribe("project.created", handler, plugin_id="market-addon") + try: + response = _upload_project(client, workspace.id) + finally: + dispatcher.unsubscribe("project.created", handler) + + assert response.status_code == 201 + assert len(received) == 1 + assert received[0].event_type == "project.created" + assert received[0].workspace_id == workspace.id + assert received[0].data["project_id"] == response.get_json()["id"] + + +class TestMarketplaceTrustRefusal: + def test_required_policy_refuses_unsigned_entry_without_partial_writes( + self, app, client, workspace, tmp_path + ): + app.config["PLUGIN_CATALOG_FILE"] = str( + _write_catalog(tmp_path, [_entry("market-unsigned")]) + ) + app.config["PLUGIN_TRUST_POLICY"] = "required" + + response = _install(client, workspace.id, "market-unsigned") + + assert response.status_code == 403 + assert Plugin.query.filter_by(id="market-unsigned").count() == 0 + assert ( + PluginInstallation.query.filter_by( + plugin_id="market-unsigned", workspace_id=workspace.id + ).count() + == 0 + ) + assert CapabilityGrant.query.filter_by(plugin_id="market-unsigned").count() == 0 + denied = ActivityEvent.query.filter_by( + workspace_id=workspace.id, event_type=EVENT_PLUGIN_DENIED + ).one() + assert denied.event_metadata["plugin_id"] == "market-unsigned" + assert denied.event_metadata["reason"] == "publisher signature refused" + + +class TestMarketplaceRollback: + def test_duplicate_install_rolls_back_without_extra_rows( + self, app, client, workspace, tmp_path + ): + app.config["PLUGIN_CATALOG_FILE"] = str( + _write_catalog(tmp_path, [_entry("market-duplicate")]) + ) + + assert _install(client, workspace.id, "market-duplicate").status_code == 201 + duplicate = _install(client, workspace.id, "market-duplicate") + + assert duplicate.status_code == 409 + assert Plugin.query.filter_by(id="market-duplicate").count() == 1 + assert ( + PluginInstallation.query.filter_by( + plugin_id="market-duplicate", workspace_id=workspace.id + ).count() + == 1 + ) + assert CapabilityGrant.query.filter_by(plugin_id="market-duplicate").count() == 0 + + def test_entry_point_conflict_leaves_no_workspace_installation( + self, app, client, workspace, tmp_path + ): + db.session.add( + Plugin( + id="market-conflict", + name="Original Plugin", + version="0.1.0", + entry_point="other.module:OtherPlugin", + capabilities=["PROJECT_READ"], + ) + ) + db.session.commit() + app.config["PLUGIN_CATALOG_FILE"] = str( + _write_catalog(tmp_path, [_entry("market-conflict")]) + ) + + response = _install(client, workspace.id, "market-conflict") + + assert response.status_code == 409 + assert ( + PluginInstallation.query.filter_by( + plugin_id="market-conflict", workspace_id=workspace.id + ).count() + == 0 + ) + plugin = db.session.get(Plugin, "market-conflict") + assert plugin.entry_point == "other.module:OtherPlugin" diff --git a/tests/test_notifications.py b/tests/test_notifications.py index 66929eb..4c3462b 100644 --- a/tests/test_notifications.py +++ b/tests/test_notifications.py @@ -1,5 +1,7 @@ """Tests for the notification inbox API and preferences (#148/#150).""" +import pytest + from app.extensions import db from app.models import NotificationPreference, User, Workspace, WorkspaceMember from app.services.notifications import notify @@ -170,6 +172,14 @@ def test_remove_member_creates_membership_notification(self, client, make_user, class TestBudgetAlertDeduplication: + def setup_method(self): + # Threshold-delivery wiring is tracked separately (#470). Until the + # service lands, skip these tests instead of failing every suite. + pytest.importorskip( + "app.services.budget_alerts", + reason="budget threshold alert delivery lands with #470", + ) + def test_repeated_crossings_alert_once_per_period(self, app, make_user): from app.services.budget_alerts import check_budget_alerts diff --git a/tests/test_tools.py b/tests/test_tools.py index e3eff5d..39cbcbe 100644 --- a/tests/test_tools.py +++ b/tests/test_tools.py @@ -5,7 +5,7 @@ import pytest from app.extensions import db as database -from app.models import FileAnalysis +from app.models import Conversation, FileAnalysis, UsageRecord def _register(client, username="tester", email="tester@example.com"): @@ -180,3 +180,81 @@ def test_send_to_chat_requires_content(self, client): headers={"X-CSRFToken": "ignored"}, ) assert response.status_code == 400 + + +class TestToolUsageLedger: + def test_generate_records_a_tools_row(self, client, db, make_user, login): + user = make_user() + login() + + response = client.post( + "/tools/generate", + json={"description": "Build a REST API in Flask", "language": "python"}, + headers={"X-CSRFToken": "ignored"}, + ) + + assert response.status_code == 200 + row = UsageRecord.query.filter_by(source="tools").one() + assert row.user_id == user.id + assert row.provider == "mock" + assert row.model == "mock-1" + assert row.total_tokens > 0 + assert row.tokens_estimated is True + + def test_code_action_records_a_tools_row_with_conversation_scope( + self, client, db, make_user, login + ): + user = make_user() + login() + conversation = Conversation(user_id=user.id, title="Tool follow-up") + db.session.add(conversation) + db.session.commit() + + response = client.post( + "/tools/code", + json={ + "action": "explain", + "code": "def foo():\n return 1", + "conversation_id": conversation.id, + }, + headers={"X-CSRFToken": "ignored"}, + ) + + assert response.status_code == 200 + row = UsageRecord.query.filter_by(source="tools").one() + assert row.user_id == user.id + assert row.conversation_id == conversation.id + assert row.total_tokens > 0 + + def test_foreign_conversation_is_not_attached(self, client, db, make_user, login): + owner = make_user() + other = make_user(username="other", email="other@example.com") + foreign = Conversation(user_id=other.id, title="Private") + db.session.add(foreign) + db.session.commit() + login(email=owner.email) + + response = client.post( + "/tools/generate", + json={"description": "Write a function", "conversation_id": foreign.id}, + headers={"X-CSRFToken": "ignored"}, + ) + + assert response.status_code == 200 + row = UsageRecord.query.filter_by(source="tools").one() + assert row.conversation_id is None + + def test_analyze_file_records_a_tools_row(self, client, db, make_user, login): + user = make_user() + login() + + response = client.post( + "/tools/analyze", + data={"file": (BytesIO(b"def add(a, b):\n return a + b"), "calc.py")}, + content_type="multipart/form-data", + ) + + assert response.status_code == 200 + row = UsageRecord.query.filter_by(source="tools").one() + assert row.user_id == user.id + assert row.total_tokens > 0 diff --git a/tests/test_usage_ledger.py b/tests/test_usage_ledger.py index fb19d19..95399af 100644 --- a/tests/test_usage_ledger.py +++ b/tests/test_usage_ledger.py @@ -5,10 +5,12 @@ project chat and project analysis. """ +import json + import pytest from app.extensions import db -from app.models import Project, ProjectFile, UsageRecord, Workspace +from app.models import Conversation, Message, Project, ProjectFile, UsageRecord, Workspace from app.models.project import SOURCE_ARCHIVE, STATUS_READY from app.services import usage_ledger @@ -329,6 +331,117 @@ def test_each_project_analysis_records_a_row(self, client, db, make_user, login) assert row.model == "mock-1" +class TestMessageBackfill: + def _conversation(self, make_user, **kwargs): + user = make_user(**kwargs) + conversation = Conversation(user_id=user.id, title="Historical chat") + db.session.add(conversation) + db.session.commit() + return user, conversation + + def _message(self, conversation, **fields): + values = {"role": "assistant", "content": "historical reply", **fields} + message = Message(conversation_id=conversation.id, **values) + db.session.add(message) + db.session.commit() + return message + + def test_empty_database_reports_zero(self, app): + result = app.test_cli_runner().invoke(args=["usage", "backfill", "--json"]) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload["scanned"] == 0 + assert payload["created"] == 0 + assert payload["skipped_existing"] == 0 + assert payload["batches"] == 0 + + def test_backfills_chat_messages_and_is_idempotent(self, app, db, make_user): + user, conversation = self._conversation(make_user) + first = self._message( + conversation, + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + ) + second = self._message( + conversation, + prompt_tokens=4, + completion_tokens=6, + total_tokens=10, + ) + self._message(conversation, role="user", content="no usage") # ignored + + runner = app.test_cli_runner() + initial = runner.invoke(args=["usage", "backfill", "--json"]) + + assert initial.exit_code == 0 + initial_payload = json.loads(initial.output) + assert initial_payload["scanned"] == 2 + assert initial_payload["created"] == 2 + assert initial_payload["skipped_existing"] == 0 + + rows = UsageRecord.query.filter_by(source="chat").order_by(UsageRecord.id).all() + assert len(rows) == 2 + assert {(row.prompt_tokens, row.completion_tokens, row.total_tokens) for row in rows} == { + (10, 5, 15), + (4, 6, 10), + } + assert all(row.user_id == user.id for row in rows) + assert all(row.conversation_id == conversation.id for row in rows) + assert all(row.tokens_estimated is True for row in rows) + + # The source message rows are never modified. + assert db.session.get(Message, first.id).prompt_tokens == 10 + assert db.session.get(Message, second.id).prompt_tokens == 4 + + repeat = runner.invoke(args=["usage", "backfill", "--json"]) + assert repeat.exit_code == 0 + repeat_payload = json.loads(repeat.output) + assert repeat_payload["scanned"] == 2 + assert repeat_payload["created"] == 0 + assert repeat_payload["skipped_existing"] == 2 + assert UsageRecord.query.filter_by(source="chat").count() == 2 + + def test_batch_size_and_limit_bound_the_run(self, app, db, make_user): + _, conversation = self._conversation(make_user) + for _ in range(3): + self._message( + conversation, + prompt_tokens=1, + completion_tokens=1, + total_tokens=2, + ) + + result = app.test_cli_runner().invoke( + args=["usage", "backfill", "--batch-size", "1", "--limit", "2", "--json"] + ) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload["scanned"] == 2 + assert payload["created"] == 2 + assert payload["batches"] == 2 + assert UsageRecord.query.filter_by(source="chat").count() == 2 + + def test_dry_run_writes_nothing(self, app, db, make_user): + _, conversation = self._conversation(make_user) + self._message( + conversation, + prompt_tokens=3, + completion_tokens=2, + total_tokens=5, + ) + + result = app.test_cli_runner().invoke(args=["usage", "backfill", "--dry-run", "--json"]) + + assert result.exit_code == 0 + payload = json.loads(result.output) + assert payload["dry_run"] is True + assert payload["created"] == 1 + assert UsageRecord.query.filter_by(source="chat").count() == 0 + + class TestReviewInstrumentation: def test_each_review_records_a_scoped_review_row(self, client, db, make_user, login): user, workspace, project = _project(