From 89eb3f68fc0e323e54047b3e0fbd63dcc4fdbfc9 Mon Sep 17 00:00:00 2001 From: Alex Date: Wed, 8 Jul 2026 21:05:40 +0100 Subject: [PATCH] fix: remove mcp artefacts, handle non fatal issues on workflow gracefully --- application/agents/workflow_agent.py | 9 +- application/api/answer/routes/base.py | 5 + application/core/settings.py | 2 - application/mcp_server.py | 73 +---- .../services/artifact_resource_service.py | 267 ------------------ .../storage/db/repositories/artifacts.py | 21 +- .../src/conversation/ConversationMessages.tsx | 54 ++-- .../src/conversation/conversationModels.ts | 3 + .../conversation/conversationSlice.test.ts | 31 ++ .../src/conversation/conversationSlice.ts | 43 +++ tests/agents/test_workflow_input_documents.py | 11 +- tests/api/answer/routes/test_base.py | 34 +++ .../test_artifact_resource_service.py | 251 ---------------- tests/services/test_mcp_server.py | 133 --------- 14 files changed, 166 insertions(+), 771 deletions(-) delete mode 100644 application/services/artifact_resource_service.py delete mode 100644 tests/services/test_artifact_resource_service.py diff --git a/application/agents/workflow_agent.py b/application/agents/workflow_agent.py index 331cd263..ec13d507 100644 --- a/application/agents/workflow_agent.py +++ b/application/agents/workflow_agent.py @@ -90,10 +90,13 @@ class WorkflowAgent(BaseAgent): self._finalize_workflow_run(workflow_owner_id, run_user_id, pg_workflow_id, query) return - # Non-fatal: some attachments were dropped (oversize / unreadable). Tell the - # user which, then still run with the documents that did bridge. + # Non-fatal: some attachments were dropped (oversize / unreadable). Emit a + # ``notice`` -- NOT an ``error`` -- so the client surfaces which were dropped + # without marking the turn failed or ending the stream (an ``error`` event is + # terminal client-side and disables reconnect). The run then still executes + # with the documents that did bridge. if dropped: - yield {"type": "error", "user_facing": True, "error": " ".join(dropped)} + yield {"type": "notice", "notice": " ".join(dropped)} self._engine.run_persisted = self._run_persisted interrupted = True diff --git a/application/api/answer/routes/base.py b/application/api/answer/routes/base.py index 6fa3574b..0902e34b 100644 --- a/application/api/answer/routes/base.py +++ b/application/api/answer/routes/base.py @@ -552,6 +552,11 @@ class BaseAnswerResource: if not line.get("user_facing"): error_text = sanitize_api_error(error_text) yield _emit({"type": "error", "error": error_text}) + elif line.get("type") == "notice": + # Non-fatal, non-terminal notice (e.g. some workflow input + # documents were dropped). Forwarded verbatim so the client can + # surface it without failing the turn; never sanitized as an error. + yield _emit({"type": "notice", "notice": line.get("notice", "")}) else: yield _emit(line) if is_structured and structured_chunks: diff --git a/application/core/settings.py b/application/core/settings.py index dcfd3946..7ea4e4b6 100644 --- a/application/core/settings.py +++ b/application/core/settings.py @@ -392,8 +392,6 @@ class Settings(BaseSettings): ARTIFACT_MAX_BYTES: int = 50 * 1024 * 1024 # cap on a single stored artifact version's bytes ARTIFACT_MAX_COUNT_PER_USER: int = 5000 # cap on artifacts a user may own ARTIFACT_MAX_TOTAL_BYTES_PER_USER: int = 5 * 1024 * 1024 * 1024 # cap on a user's total stored bytes - # Cap on bytes served per MCP ``resources/read`` so a giant artifact never streams into LLM context. - ARTIFACT_RESOURCE_READ_MAX_BYTES: int = 1 * 1024 * 1024 @field_validator("POSTGRES_URI", mode="before") @classmethod diff --git a/application/mcp_server.py b/application/mcp_server.py index 8b96fd20..347fd2b4 100644 --- a/application/mcp_server.py +++ b/application/mcp_server.py @@ -1,11 +1,8 @@ -"""FastMCP server exposing DocsGPT retrieval + artifacts over streamable HTTP. +"""FastMCP server exposing DocsGPT retrieval over streamable HTTP. Mounted at ``/mcp`` by ``application/asgi.py``. Bearer tokens are the existing DocsGPT agent API keys — no new credential surface. The -``search_docs`` tool searches the caller's knowledge base; the artifact -resources middleware exposes the caller's own artifacts as MCP resources -(``resources/list`` / ``resources/read`` over ``artifact://`` URIs), -scoped strictly to the Bearer key's owning principal. +``search_docs`` tool searches the caller's knowledge base. The tool reads the ``Authorization`` header directly via ``get_http_headers(include={"authorization"})``. The ``include`` kwarg @@ -18,23 +15,11 @@ token, we opt it back in. from __future__ import annotations import asyncio -import base64 import logging -from typing import Sequence from fastmcp import FastMCP -from fastmcp.resources import ResourceContent, ResourceResult -from fastmcp.resources.base import Resource from fastmcp.server.dependencies import get_http_headers -from fastmcp.server.middleware import CallNext, Middleware, MiddlewareContext -from application.services.artifact_resource_service import ( - ArtifactReadResult, - ResourceDenied, - ResourceNotFound, - list_artifact_resources, - read_artifact_resource, -) from application.services.search_service import ( InvalidAPIKey, SearchFailed, @@ -52,61 +37,7 @@ def _extract_bearer_token() -> str | None: return parts[1] -def _read_result_to_mcp(result: ArtifactReadResult) -> ResourceResult: - """Wrap a service read result as a FastMCP ``ResourceResult`` (text or blob).""" - if result.blob_b64 is not None: - raw = base64.b64decode(result.blob_b64) - return ResourceResult([ResourceContent(raw, mime_type=result.mime_type)]) - return ResourceResult([ResourceContent(result.text or "", mime_type=result.mime_type)]) - - -class ArtifactResourcesMiddleware(Middleware): - """Expose the calling principal's artifacts as MCP resources (read/list). - - The principal is the Bearer api_key's owner; resources are scoped to that - owner and a foreign/unauthenticated read is denied. Static resources - registered on the server (if any) are preserved by chaining ``call_next``. - """ - - async def on_list_resources( - self, - context: MiddlewareContext, - call_next: CallNext, - ) -> Sequence[Resource]: - """Append the principal's artifact resources to the static resource list.""" - existing = list(await call_next(context)) - try: - artifacts = await asyncio.to_thread( - list_artifact_resources, _extract_bearer_token() - ) - except Exception: - logger.exception("on_list_resources: artifact listing failed") - return existing - existing.extend(artifacts) - return existing - - async def on_read_resource( - self, - context: MiddlewareContext, - call_next: CallNext, - ) -> ResourceResult: - """Serve ``artifact://`` reads from the artifact store; defer others.""" - uri = str(getattr(context.message, "uri", "")) - if not uri.startswith("artifact://"): - return await call_next(context) - try: - result = await asyncio.to_thread( - read_artifact_resource, _extract_bearer_token(), uri - ) - except ResourceDenied as exc: - raise PermissionError(str(exc) or "forbidden") from exc - except ResourceNotFound as exc: - raise ValueError(str(exc) or "resource not found") from exc - return _read_result_to_mcp(result) - - mcp = FastMCP("docsgpt") -mcp.add_middleware(ArtifactResourcesMiddleware()) @mcp.tool diff --git a/application/services/artifact_resource_service.py b/application/services/artifact_resource_service.py deleted file mode 100644 index 7e401f2e..00000000 --- a/application/services/artifact_resource_service.py +++ /dev/null @@ -1,267 +0,0 @@ -"""Flask-free service exposing a principal's artifacts as MCP Resources. - -The MCP server (``application/mcp_server.py``) authenticates a request with -``Authorization: Bearer ``; that key resolves to the owning -``user_id`` via ``AgentsRepository.find_by_key`` (the same api_key->owner path -as the HTTP artifact routes). Resources are scoped strictly to that principal -AND to the owner's own conversations: because the api_key is embedded publicly -in the widget and MCP carries no per-visitor conversation context, shared-usage -(widget/visitor) conversations are excluded. ``resources/list`` returns only the -owner's own-chat artifacts, and ``resources/read`` re-checks ownership + owner -scope before serving any bytes. An unresolvable principal yields an empty list / -a denied read -- never another principal's or another visitor's artifact. - -``resources/list`` returns FastMCP ``Resource`` objects so they pass straight -through the server's list/dedupe/wire pipeline; ``resources/read`` is served by -the MCP middleware, which streams the real bytes for each ``artifact://`` uri. -""" - -from __future__ import annotations - -import base64 -import logging -import re -from dataclasses import dataclass -from typing import List, Optional - -from fastmcp.resources.base import Resource -from fastmcp.resources.types import TextResource -from sqlalchemy.exc import DataError, DBAPIError - -from application.core.settings import settings -from application.storage.db.base_repository import looks_like_uuid -from application.storage.db.repositories.agents import AgentsRepository -from application.storage.db.repositories.artifacts import ArtifactsRepository -from application.storage.db.session import db_readonly -from application.storage.storage_creator import StorageCreator - -logger = logging.getLogger(__name__) - -# artifact://{artifact_id}/v{version} -_URI_RE = re.compile(r"^artifact://(?P[^/]+)/v(?P\d+)$") - -# Max resources advertised by ``resources/list`` so the model is not flooded -# with thousands of rows even when a principal owns far more artifacts. -_RESOURCE_LIST_LIMIT = 500 - -# mime types served as inline ``text`` rather than base64 ``blob``. -_TEXT_MIME_PREFIXES = ("text/",) -_TEXT_MIME_EXACT = { - "application/json", - "application/xml", - "application/javascript", - "application/x-ndjson", - "image/svg+xml", -} -_TEXT_MIME_SUFFIXES = ("+json", "+xml") - - -class ResourceDenied(Exception): - """The principal may not access the requested resource (unauth or foreign).""" - - -class ResourceNotFound(Exception): - """The requested ``artifact://`` uri does not resolve to a stored version.""" - - -@dataclass(frozen=True) -class ArtifactReadResult: - """Materialized artifact contents for a ``resources/read`` response.""" - - uri: str - mime_type: str - text: Optional[str] = None - blob_b64: Optional[str] = None - - -def _read_cap() -> int: - """Return the per-read byte cap, or 0 (no cap) when the setting disables it.""" - return int(getattr(settings, "ARTIFACT_RESOURCE_READ_MAX_BYTES", 0) or 0) - - -def _cap_text(text: str) -> str: - """Truncate ``text`` to the read cap; unchanged when the cap is disabled.""" - cap = _read_cap() - return text[:cap] if cap > 0 else text - - -def _is_texty(mime_type: str) -> bool: - """Return True when ``mime_type`` should be served as inline UTF-8 text.""" - mime = (mime_type or "").split(";", 1)[0].strip().lower() - if mime in _TEXT_MIME_EXACT: - return True - if any(mime.startswith(p) for p in _TEXT_MIME_PREFIXES): - return True - return any(mime.endswith(s) for s in _TEXT_MIME_SUFFIXES) - - -def _resolve_agent(api_key: Optional[str]) -> Optional[dict]: - """Resolve a Bearer api_key to its owning agent row (carries ``id`` + ``user_id``). - - Returns the whole agent so callers can scope artifact visibility to that agent's - conversations, not the owner's entire corpus. None when unresolvable. - """ - if not api_key: - return None - try: - with db_readonly() as conn: - agent = AgentsRepository(conn).find_by_key(api_key) - except Exception: - logger.exception("artifact resource: principal resolution failed") - return None - if not agent or not agent.get("user_id") or not agent.get("id"): - return None - return agent - - -def _resource_uri(artifact_id: str, version: int) -> str: - """Build the stable ``artifact://{id}/v{version}`` resource uri.""" - return f"artifact://{artifact_id}/v{version}" - - -def list_artifact_resources(api_key: Optional[str]) -> List[Resource]: - """List the calling principal's artifacts as FastMCP resources (empty if unresolved). - - Returns FastMCP ``Resource`` objects (not raw ``mcp.types.Resource``): the - server's ``resources/list`` pipeline reads FastMCP-only attributes - (``.version``, ``.auth``, ``is_enabled``) and calls ``to_mcp_resource()`` for - the wire encoding, so a raw ``mcp.types.Resource`` mixed into the list would - raise ``AttributeError``. The placeholder ``text`` is never served -- the - middleware's ``on_read_resource`` intercepts every ``artifact://`` read and - streams the real bytes via :func:`read_artifact_resource`. - """ - agent = _resolve_agent(api_key) - if not agent: - return [] - try: - with db_readonly() as conn: - # MCP-over-key has no per-visitor conversation context, and the api_key - # is embedded publicly in the widget, so scope to the owner's OWN - # (non-shared-usage) conversations only. Without this, the key would - # enumerate artifacts from every visitor's conversation with the agent. - rows = ArtifactsRepository(conn).list_artifacts_for_agent( - str(agent["id"]), str(agent["user_id"]), owner_only=True - ) - except Exception: - logger.exception("artifact resource: list failed") - return [] - - resources: List[Resource] = [] - for row in rows[:_RESOURCE_LIST_LIMIT]: - artifact_id = str(row.get("id")) - version = row.get("current_version") or 1 - title = row.get("title") or f"artifact-{artifact_id}" - resources.append( - TextResource( - uri=_resource_uri(artifact_id, version), - name=title, - title=title, - description=f"DocsGPT {row.get('kind') or 'file'} artifact", - mime_type=_kind_mime_hint(row.get("kind")), - text="", - ) - ) - return resources - - -def _kind_mime_hint(kind: Optional[str]) -> str: - """Concrete mime hint for a list row; ambiguous kinds fall back to octet-stream.""" - # Only kinds with a single unambiguous mime get a concrete type; everything - # else stays octet-stream so a list row never mismatches the read's bytes. - return { - "html": "text/html", - "data": "application/json", - }.get((kind or "").lower(), "application/octet-stream") - - -def read_artifact_resource(api_key: Optional[str], uri: str) -> ArtifactReadResult: - """Authorize and materialize an ``artifact://`` resource for the principal. - - Raises: - ResourceDenied: principal unresolved, or the artifact is not theirs. - ResourceNotFound: uri malformed, or the version/file does not exist. - """ - match = _URI_RE.match(uri or "") - if not match: - raise ResourceNotFound(f"unsupported resource uri: {uri!r}") - artifact_id = match.group("id") - version = int(match.group("version")) - - # A non-UUID id would reach a ``CAST(:id AS uuid)`` and raise a DB DataError; - # gate it up front like the HTTP artifact routes do. - if not looks_like_uuid(artifact_id): - raise ResourceNotFound(f"artifact {artifact_id} not found") - - agent = _resolve_agent(api_key) - if not agent: - raise ResourceDenied("unauthenticated") - user_id = str(agent["user_id"]) - agent_id = str(agent["id"]) - - try: - with db_readonly() as conn: - repo = ArtifactsRepository(conn) - artifact = repo.get_artifact(artifact_id) - if artifact is None: - raise ResourceNotFound(f"artifact {artifact_id} not found") - # Ownership is the first authz point: never serve another principal's - # artifact over MCP, regardless of conversation/workflow parent sharing. - if str(artifact.get("user_id")) != user_id: - raise ResourceDenied("forbidden") - # Agent scope is the second: a per-agent key only reads artifacts from - # its own conversations, not the owner's other agents / workflow runs. - # ``owner_only`` further excludes shared-usage (widget/visitor) - # conversations: the public key reads the owner's own chats only, never - # another visitor's artifact (parity with the HTTP route's confinement). - if not repo.artifact_in_agent_scope(artifact_id, agent_id, owner_only=True): - raise ResourceDenied("forbidden") - version_row = repo.get_version(artifact_id, version) - except (DataError, DBAPIError) as exc: - raise ResourceNotFound(f"artifact {artifact_id} not found") from exc - - if version_row is None: - raise ResourceNotFound(f"version {version} of {artifact_id} not found") - - mime_type = version_row.get("mime_type") or "application/octet-stream" - uri = _resource_uri(artifact_id, version) - - # Prefer the stored preview/extracted text for texty kinds: it is already - # bounded and avoids a storage round-trip. - preview = version_row.get("preview_text") - if _is_texty(mime_type) and preview: - return ArtifactReadResult(uri=uri, mime_type=mime_type, text=_cap_text(preview)) - - storage_path = version_row.get("storage_path") - if not storage_path: - if _is_texty(mime_type) and preview is not None: - return ArtifactReadResult(uri=uri, mime_type=mime_type, text=_cap_text(preview)) - raise ResourceNotFound(f"version {version} of {artifact_id} has no stored bytes") - - try: - data = _read_capped_bytes(storage_path) - except FileNotFoundError as exc: - raise ResourceNotFound(f"version {version} of {artifact_id} has no stored bytes") from exc - - if _is_texty(mime_type): - # ``errors="ignore"`` keeps valid text texty even when the cap splits a - # multibyte char at the boundary, instead of demoting it to a blob. - return ArtifactReadResult(uri=uri, mime_type=mime_type, text=data.decode("utf-8", errors="ignore")) - return ArtifactReadResult( - uri=uri, mime_type=mime_type, blob_b64=base64.b64encode(data).decode("ascii") - ) - - -def _read_capped_bytes(storage_path: str) -> bytes: - """Read at most ``_read_cap()`` bytes of a stored artifact version (0 == all).""" - cap = _read_cap() - storage = StorageCreator.get_storage() - file_obj = storage.get_file(storage_path) - try: - return file_obj.read(cap) if cap > 0 else file_obj.read() - finally: - close = getattr(file_obj, "close", None) - if callable(close): - try: - close() - except Exception: - logger.debug("artifact resource: file close failed", exc_info=True) diff --git a/application/storage/db/repositories/artifacts.py b/application/storage/db/repositories/artifacts.py index 2c4b44f6..ba8f6700 100644 --- a/application/storage/db/repositories/artifacts.py +++ b/application/storage/db/repositories/artifacts.py @@ -206,25 +206,19 @@ class ArtifactsRepository: agent_id: str, user_id: str, conversation_id: Optional[str] = None, - owner_only: bool = False, ) -> list[dict]: """List an agent's artifacts (owner-scoped), optionally narrowed to one conversation. Scopes a per-agent api-key's artifact visibility to the conversations that agent produced, so the key cannot enumerate the owner's whole corpus. When ``conversation_id`` is given, the SQL narrows to that single conversation - (the per-visitor bearer capability for a public widget key). When - ``owner_only`` is set, shared-usage (widget/visitor) conversations are - excluded so a credential with no per-visitor conversation context (the MCP - Bearer key) sees only the owner's own chats, never other visitors' artifacts. + (the per-visitor bearer capability for a public widget key). """ clauses = ["c.agent_id = CAST(:agent_id AS uuid)", "a.user_id = :user_id"] params: dict[str, Any] = {"agent_id": str(agent_id), "user_id": user_id} if conversation_id is not None: clauses.append("a.conversation_id = CAST(:conversation_id AS uuid)") params["conversation_id"] = conversation_id - if owner_only: - clauses.append("c.is_shared_usage = false") result = self._conn.execute( text( "SELECT a.* FROM artifacts a " @@ -236,21 +230,12 @@ class ArtifactsRepository: ) return [_artifact_to_dict(r) for r in result.fetchall()] - def artifact_in_agent_scope( - self, artifact_id: str, agent_id: str, owner_only: bool = False - ) -> bool: - """True if ``artifact_id``'s parent conversation belongs to ``agent_id``. - - With ``owner_only`` set, shared-usage (widget/visitor) conversations are - excluded so a per-visitor-less credential (the MCP Bearer key) can only - confirm scope for the owner's own chats, never another visitor's artifact. - """ + def artifact_in_agent_scope(self, artifact_id: str, agent_id: str) -> bool: + """True if ``artifact_id``'s parent conversation belongs to ``agent_id``.""" clauses = [ "a.id = CAST(:id AS uuid)", "c.agent_id = CAST(:agent_id AS uuid)", ] - if owner_only: - clauses.append("c.is_shared_usage = false") result = self._conn.execute( text( "SELECT 1 FROM artifacts a " diff --git a/frontend/src/conversation/ConversationMessages.tsx b/frontend/src/conversation/ConversationMessages.tsx index 66eaf064..2356d729 100644 --- a/frontend/src/conversation/ConversationMessages.tsx +++ b/frontend/src/conversation/ConversationMessages.tsx @@ -280,35 +280,47 @@ export default function ConversationMessages({ } // tool_calls.length, not tool_calls — empty arrays are JS-truthy. + // ``notice`` is included so a non-fatal notice still surfaces even when the + // run produced no textual answer (e.g. an artifact-only workflow). const hasContent = query.thought || query.response || (query.tool_calls && query.tool_calls.length > 0) || - query.research; + query.research || + query.notice; if (hasContent) { const isCurrentlyStreaming = status === 'loading' && index === queries.length - 1; return ( - handleFeedback(query, feedback, index) - : undefined - } - /> + + {query.notice ? ( +
+ {query.notice} +
+ ) : null} + handleFeedback(query, feedback, index) + : undefined + } + /> + ); } diff --git a/frontend/src/conversation/conversationModels.ts b/frontend/src/conversation/conversationModels.ts index 26b6bb9c..01c87a58 100644 --- a/frontend/src/conversation/conversationModels.ts +++ b/frontend/src/conversation/conversationModels.ts @@ -63,6 +63,9 @@ export interface Query { sources?: { title: string; text: string; link: string }[]; tool_calls?: ToolCallsType[]; error?: string; + // Non-fatal notice (e.g. some workflow input documents were dropped). Shown + // alongside the answer; unlike ``error`` it does not fail the turn or end the stream. + notice?: string; attachments?: { id: string; fileName: string }[]; structured?: boolean; schema?: object; diff --git a/frontend/src/conversation/conversationSlice.test.ts b/frontend/src/conversation/conversationSlice.test.ts index db942434..fa012cb8 100644 --- a/frontend/src/conversation/conversationSlice.test.ts +++ b/frontend/src/conversation/conversationSlice.test.ts @@ -2,6 +2,7 @@ import { describe, expect, it } from 'vitest'; import reducer, { applyMessageTail, + raiseNotice, setConversation, } from './conversationSlice'; @@ -151,3 +152,33 @@ describe('applyMessageTail — streaming partial', () => { expect(next.queries[0].response).toBeUndefined(); }); }); + +describe('raiseNotice — non-fatal notice', () => { + it('records a notice without setting an error (turn is not failed)', () => { + // seedSlice leaves conversationId at its initial null; match that. + const state = seedSlice(); + const next = reducer( + state, + raiseNotice({ + conversationId: null, + index: 0, + message: 'big.txt was dropped (too large)', + }), + ); + expect(next.queries[0].notice).toBe('big.txt was dropped (too large)'); + expect(next.queries[0].error).toBeUndefined(); + }); + + it('is a no-op when the conversationId does not match', () => { + const state = seedSlice(); + const next = reducer( + state, + raiseNotice({ + conversationId: 'some-other-conversation', + index: 0, + message: 'ignored', + }), + ); + expect(next.queries[0].notice).toBeUndefined(); + }); +}); diff --git a/frontend/src/conversation/conversationSlice.ts b/frontend/src/conversation/conversationSlice.ts index 06b48934..cc2c2af8 100644 --- a/frontend/src/conversation/conversationSlice.ts +++ b/frontend/src/conversation/conversationSlice.ts @@ -297,6 +297,14 @@ export const fetchAnswer = createAsyncThunk< dispatch( conversationSlice.actions.setStatus('awaiting_tool_actions'), ); + } else if (data.type === 'notice') { + dispatch( + conversationSlice.actions.raiseNotice({ + conversationId: currentConversationId, + index: targetIndex, + message: data.notice ?? '', + }), + ); } else if (data.type === 'error') { dispatch(conversationSlice.actions.setStatus('failed')); dispatch( @@ -442,6 +450,14 @@ export const fetchAnswer = createAsyncThunk< progress: data.data, }), ); + } else if (data.type === 'notice') { + dispatch( + conversationSlice.actions.raiseNotice({ + conversationId: currentConversationId, + index: targetIndex, + message: data.notice ?? '', + }), + ); } else if (data.type === 'error') { // set status to 'failed' dispatch(conversationSlice.actions.setStatus('failed')); @@ -703,6 +719,14 @@ export const submitToolActions = createAsyncThunk< ); } else if (data.type === 'tool_calls_pending') { dispatch(conversationSlice.actions.setStatus('awaiting_tool_actions')); + } else if (data.type === 'notice') { + dispatch( + conversationSlice.actions.raiseNotice({ + conversationId, + index: targetIndex, + message: data.notice ?? '', + }), + ); } else if (data.type === 'error') { dispatch(conversationSlice.actions.setStatus('failed')); dispatch( @@ -1009,6 +1033,24 @@ export const conversationSlice = createSlice({ state.queries[index].error = message; }, + // Non-fatal counterpart to ``raiseError``: records a notice on the query + // (e.g. some workflow input documents were dropped) WITHOUT setting the + // 'failed' status, so the turn keeps streaming and can still complete. + raiseNotice( + state, + action: PayloadAction<{ + conversationId: string | null; + index: number; + message: string; + }>, + ) { + const { conversationId, index, message } = action.payload; + if (state.conversationId !== conversationId) return; + if (!state.queries[index]) return; + + state.queries[index].notice = message; + }, + resetConversation: (state) => { state.queries = initialState.queries; state.status = initialState.status; @@ -1054,6 +1096,7 @@ export const { setConversationId, setStatus, raiseError, + raiseNotice, resetConversation, applyMessageTail, updateMessageMeta, diff --git a/tests/agents/test_workflow_input_documents.py b/tests/agents/test_workflow_input_documents.py index f7e53730..0c1f5918 100644 --- a/tests/agents/test_workflow_input_documents.py +++ b/tests/agents/test_workflow_input_documents.py @@ -415,11 +415,12 @@ def test_oversize_declared_attachment_skipped_with_notice(pg_engine, tmp_path, m # The oversize doc was dropped -> no input documents bridged, but the run ran. assert engine.captured_inputs is not None assert engine.captured_inputs["input_documents"] == [] - # A non-fatal notice naming the dropped document was surfaced, flagged user_facing - # so the route emits the filename verbatim instead of a generic sanitized message. - errors = [e for e in events if e.get("type") == "error"] - assert errors and "big.txt" in errors[0]["error"] - assert errors[0].get("user_facing") is True + # A non-fatal notice naming the dropped document was surfaced as a ``notice`` + # (NOT an ``error``, which is terminal client-side) so the run still completes. + notices = [e for e in events if e.get("type") == "notice"] + assert notices and "big.txt" in notices[0]["notice"] + # It must not be an error event (that would fail the turn and disable reconnect). + assert not [e for e in events if e.get("type") == "error"] # Nothing was persisted for the oversize doc. with pg_engine.connect() as conn: n = conn.execute( diff --git a/tests/api/answer/routes/test_base.py b/tests/api/answer/routes/test_base.py index 65da96d0..9707f436 100644 --- a/tests/api/answer/routes/test_base.py +++ b/tests/api/answer/routes/test_base.py @@ -241,6 +241,40 @@ class TestCompleteStreamMethod: assert "artifact storage quota" in error_chunks[0] assert "Rate limit exceeded" not in error_chunks[0] + def test_notice_is_forwarded_verbatim_and_not_an_error(self, mock_mongo_db, flask_app): + """A non-fatal ``notice`` streams through as a notice, never as an error. + + A ``notice`` (e.g. some workflow input documents were dropped) must not be + emitted as ``type: error`` -- the client treats an error event as terminal and + disables reconnect -- and its text must not be run through sanitize_api_error. + """ + from application.api.answer.routes.base import BaseAnswerResource + + with flask_app.app_context(): + resource = BaseAnswerResource() + + mock_agent = MagicMock() + mock_agent.gen.return_value = iter( + [{"type": "notice", "notice": "big.txt exceeds the per-file size limit"}] + ) + + stream = list( + resource.complete_stream( + question="Test?", + agent=mock_agent, + conversation_id=None, + user_api_key=None, + decoded_token={"sub": "user123"}, + should_persist=False, + ) + ) + + notice_chunks = [s for s in stream if '"type": "notice"' in s] + assert notice_chunks + assert "big.txt exceeds the per-file size limit" in notice_chunks[0] + # Crucially, it is not surfaced as an error event. + assert not [s for s in stream if '"type": "error"' in s] + def test_non_user_facing_error_is_sanitized(self, mock_mongo_db, flask_app): """A raw error without the flag is still routed through sanitize_api_error.""" from application.api.answer.routes.base import BaseAnswerResource diff --git a/tests/services/test_artifact_resource_service.py b/tests/services/test_artifact_resource_service.py deleted file mode 100644 index 2bd26903..00000000 --- a/tests/services/test_artifact_resource_service.py +++ /dev/null @@ -1,251 +0,0 @@ -"""Tests for application/services/artifact_resource_service.py. - -The service exposes a Bearer-key principal's own artifacts as MCP resources. -These tests patch the DB/storage seams (``db_readonly``, the repositories, and -``StorageCreator``) so they run without Postgres, mirroring the light mocking -in ``tests/services/test_mcp_server.py``. They assert that: - -- ``resources/list`` returns only the principal's artifacts; -- ``resources/read`` returns ``text`` vs ``blob`` by mime with the right type; -- a foreign-owner artifact is denied (no cross-principal exposure); -- an unauthenticated / unresolved principal gets an empty list / denied read; -- a non-UUID id is rejected as not-found (no leaked DB error); -- the read is byte-capped. -""" - -from __future__ import annotations - -import base64 -import io -from contextlib import contextmanager - -import pytest - -from application.services import artifact_resource_service as svc - -OWNER = "owner-1" -STRANGER = "stranger-2" - -# Real UUIDs: the read path gates non-UUID ids before they reach the DB. -ART_TEXT = "11111111-1111-4111-8111-111111111111" -ART_BIN = "22222222-2222-4222-8222-222222222222" -ART_FOREIGN = "33333333-3333-4333-8333-333333333333" -# Owned by OWNER but produced by a DIFFERENT agent — must be invisible to owner-key. -ART_OTHER_AGENT = "44444444-4444-4444-8444-444444444444" -# In the owner's agent scope but from a shared-usage (widget/visitor) conversation -# — must be invisible to the public MCP key (owner-only scope). -ART_SHARED = "55555555-5555-4555-8555-555555555555" - - -@contextmanager -def _fake_conn(): - yield object() - - -class _FakeAgents: - """Stub AgentsRepository: maps api_key -> agent row (id + user_id) or None.""" - - _MAP = { - "owner-key": {"id": "agent-owner", "user_id": OWNER}, - "stranger-key": {"id": "agent-stranger", "user_id": STRANGER}, - } - - def __init__(self, conn): - pass - - def find_by_key(self, key): - return self._MAP.get(key) - - -class _FakeArtifacts: - """Stub ArtifactsRepository backed by in-memory artifact/version dicts.""" - - artifacts: dict = {} - versions: dict = {} - - def __init__(self, conn): - pass - - def list_artifacts_for_agent(self, agent_id, user_id, conversation_id=None, owner_only=False): - return [ - a - for a in self.artifacts.values() - if a.get("agent_id") == agent_id - and a["user_id"] == user_id - and (conversation_id is None or a.get("conversation_id") == conversation_id) - and not (owner_only and a.get("is_shared_usage")) - ] - - def artifact_in_agent_scope(self, artifact_id, agent_id, owner_only=False): - art = self.artifacts.get(artifact_id) - return ( - art is not None - and art.get("agent_id") == agent_id - and not (owner_only and art.get("is_shared_usage")) - ) - - def get_artifact(self, artifact_id): - return self.artifacts.get(artifact_id) - - def get_version(self, artifact_id, version): - return self.versions.get((artifact_id, version)) - - -class _FakeStorage: - """Stub BaseStorage.get_file returning a capped BytesIO of fixed bytes.""" - - blob = b"x" * 10 - - def get_file(self, path): - return io.BytesIO(self.blob) - - -@pytest.fixture(autouse=True) -def _wire(monkeypatch): - """Point the service's DB/storage seams at the in-memory fakes.""" - _FakeArtifacts.artifacts = { - ART_TEXT: {"id": ART_TEXT, "user_id": OWNER, "agent_id": "agent-owner", "kind": "data", "title": "notes", "current_version": 2}, - ART_BIN: {"id": ART_BIN, "user_id": OWNER, "agent_id": "agent-owner", "kind": "image", "title": "chart", "current_version": 1}, - ART_FOREIGN: { - "id": ART_FOREIGN, - "user_id": STRANGER, - "agent_id": "agent-stranger", - "kind": "data", - "title": "secret", - "current_version": 1, - }, - ART_OTHER_AGENT: { - "id": ART_OTHER_AGENT, - "user_id": OWNER, - "agent_id": "agent-other", - "kind": "data", - "title": "other-agent", - "current_version": 1, - }, - ART_SHARED: { - "id": ART_SHARED, - "user_id": OWNER, - "agent_id": "agent-owner", - "kind": "data", - "title": "visitor-doc", - "current_version": 1, - "is_shared_usage": True, - }, - } - _FakeArtifacts.versions = { - (ART_TEXT, 2): {"mime_type": "text/csv", "storage_path": "k/text.csv", "preview_text": None}, - (ART_BIN, 1): {"mime_type": "image/png", "storage_path": "k/chart.png", "preview_text": None}, - (ART_FOREIGN, 1): {"mime_type": "text/plain", "storage_path": "k/secret.txt", "preview_text": None}, - (ART_OTHER_AGENT, 1): {"mime_type": "text/plain", "storage_path": "k/other.txt", "preview_text": None}, - (ART_SHARED, 1): {"mime_type": "text/plain", "storage_path": "k/visitor.txt", "preview_text": None}, - } - monkeypatch.setattr(svc, "db_readonly", _fake_conn) - monkeypatch.setattr(svc, "AgentsRepository", _FakeAgents) - monkeypatch.setattr(svc, "ArtifactsRepository", _FakeArtifacts) - monkeypatch.setattr(svc.StorageCreator, "get_storage", staticmethod(lambda: _FakeStorage())) - - -@pytest.mark.unit -class TestListArtifactResources: - def test_lists_only_principal_artifacts(self): - out = svc.list_artifact_resources("owner-key") - uris = {str(r.uri) for r in out} - assert uris == {f"artifact://{ART_TEXT}/v2", f"artifact://{ART_BIN}/v1"} - assert f"artifact://{ART_FOREIGN}/v1" not in uris - - def test_excludes_shared_usage_visitor_artifacts(self): - # The MCP Bearer key is embedded publicly in the widget and has no - # per-visitor conversation context, so shared-usage (visitor) artifacts - # must never be enumerated — only the owner's own-chat artifacts. - uris = {str(r.uri) for r in svc.list_artifact_resources("owner-key")} - assert f"artifact://{ART_SHARED}/v1" not in uris - - def test_unresolved_principal_is_empty(self): - assert svc.list_artifact_resources("bogus-key") == [] - - def test_missing_token_is_empty(self): - assert svc.list_artifact_resources(None) == [] - - def test_resource_carries_name_and_mime(self): - out = {str(r.uri): r for r in svc.list_artifact_resources("owner-key")} - assert out[f"artifact://{ART_TEXT}/v2"].name == "notes" - # The list row never advertises a wildcard/wrong type; the image kind - # falls back to the generic octet-stream hint. FastMCP Resource uses - # ``mime_type`` (and must survive ``to_mcp_resource()`` for the wire). - assert out[f"artifact://{ART_BIN}/v1"].mime_type == "application/octet-stream" - - def test_resources_are_fastmcp_and_render_to_wire(self): - # Regression: ``resources/list`` must yield FastMCP Resource objects so - # the server's list pipeline (reads ``.version``/``.auth``, calls - # ``to_mcp_resource``) does not raise AttributeError on a raw mcp.types. - from fastmcp.resources.base import Resource as FastMCPResource - - out = svc.list_artifact_resources("owner-key") - assert out and all(isinstance(r, FastMCPResource) for r in out) - for r in out: - wire = r.to_mcp_resource() - assert str(wire.uri).startswith("artifact://") - - -@pytest.mark.unit -class TestReadArtifactResource: - def test_text_mime_returns_text(self): - res = svc.read_artifact_resource("owner-key", f"artifact://{ART_TEXT}/v2") - assert res.text == _FakeStorage.blob.decode("utf-8") - assert res.blob_b64 is None - assert res.mime_type == "text/csv" - - def test_binary_mime_returns_blob(self): - res = svc.read_artifact_resource("owner-key", f"artifact://{ART_BIN}/v1") - assert res.blob_b64 == base64.b64encode(_FakeStorage.blob).decode("ascii") - assert res.text is None - assert res.mime_type == "image/png" - - def test_prefers_preview_text_when_present(self, monkeypatch): - _FakeArtifacts.versions[(ART_TEXT, 2)]["preview_text"] = "cached preview" - res = svc.read_artifact_resource("owner-key", f"artifact://{ART_TEXT}/v2") - assert res.text == "cached preview" - - def test_read_denies_owner_artifact_from_another_agent(self): - # Owned by the same user but produced by a different agent: a per-agent key - # is scoped like its search and must NOT read the owner's other-agent corpus. - with pytest.raises(svc.ResourceDenied): - svc.read_artifact_resource("owner-key", f"artifact://{ART_OTHER_AGENT}/v1") - - def test_read_denies_shared_usage_visitor_artifact(self): - # In the owner's agent scope but produced in a shared-usage (widget/visitor) - # conversation: the public MCP key must not read another visitor's artifact. - with pytest.raises(svc.ResourceDenied): - svc.read_artifact_resource("owner-key", f"artifact://{ART_SHARED}/v1") - - def test_foreign_owner_is_denied(self): - with pytest.raises(svc.ResourceDenied): - svc.read_artifact_resource("owner-key", f"artifact://{ART_FOREIGN}/v1") - - def test_unauthenticated_is_denied(self): - with pytest.raises(svc.ResourceDenied): - svc.read_artifact_resource(None, f"artifact://{ART_TEXT}/v2") - with pytest.raises(svc.ResourceDenied): - svc.read_artifact_resource("bogus-key", f"artifact://{ART_TEXT}/v2") - - def test_unknown_uri_scheme_not_found(self): - with pytest.raises(svc.ResourceNotFound): - svc.read_artifact_resource("owner-key", "https://example.com/x") - - def test_non_uuid_id_is_not_found(self): - # A non-UUID id must be rejected before the DB cast (no leaked DataError). - with pytest.raises(svc.ResourceNotFound): - svc.read_artifact_resource("owner-key", "artifact://not-a-uuid/v1") - - def test_missing_version_not_found(self): - with pytest.raises(svc.ResourceNotFound): - svc.read_artifact_resource("owner-key", f"artifact://{ART_TEXT}/v99") - - def test_read_is_byte_capped(self, monkeypatch): - monkeypatch.setattr(svc.settings, "ARTIFACT_RESOURCE_READ_MAX_BYTES", 3) - _FakeStorage.blob = b"abcdefghij" - try: - res = svc.read_artifact_resource("owner-key", f"artifact://{ART_BIN}/v1") - assert base64.b64decode(res.blob_b64) == b"abc" - finally: - _FakeStorage.blob = b"x" * 10 diff --git a/tests/services/test_mcp_server.py b/tests/services/test_mcp_server.py index ede019aa..c1da1cbd 100644 --- a/tests/services/test_mcp_server.py +++ b/tests/services/test_mcp_server.py @@ -9,7 +9,6 @@ full HTTP-layer plumbing (mount, lifespan, session handshake) is covered by ``tests/test_asgi.py``. """ -from types import SimpleNamespace from unittest.mock import patch import pytest @@ -133,135 +132,3 @@ class TestSearchDocsTool: ): await search_docs(query="q") mock_search.assert_called_once_with("lowercase-scheme", "q", 5) - - -async def _call_next_empty(context): - """Stand-in downstream handler returning no static resources.""" - return [] - - -@pytest.mark.unit -class TestArtifactResourcesMiddleware: - @pytest.mark.asyncio - async def test_list_appends_principal_artifacts(self): - from fastmcp.resources.types import TextResource - - from application.mcp_server import ArtifactResourcesMiddleware - - mw = ArtifactResourcesMiddleware() - res = TextResource( - uri="artifact://a/v1", name="x", mime_type="text/plain", text="" - ) - with ( - patch("application.mcp_server._extract_bearer_token", return_value="k"), - patch( - "application.mcp_server.list_artifact_resources", return_value=[res] - ), - ): - out = await mw.on_list_resources(SimpleNamespace(), _call_next_empty) - assert [str(r.uri) for r in out] == ["artifact://a/v1"] - - @pytest.mark.asyncio - async def test_list_resources_wire_pipeline_accepts_artifact(self): - # Regression for the resources/list crash: drive the real server pipeline - # (is_enabled / auth filter / dedupe / to_mcp_resource), which raises - # AttributeError on a raw mcp.types.Resource. A FastMCP resource passes. - from fastmcp.resources.types import TextResource - - from application.mcp_server import mcp - - res = TextResource( - uri="artifact://a/v1", - name="x", - title="x", - description="d", - mime_type="text/plain", - text="", - ) - with ( - patch("application.mcp_server._extract_bearer_token", return_value="k"), - patch( - "application.mcp_server.list_artifact_resources", return_value=[res] - ), - ): - resources = await mcp.list_resources() - assert "artifact://a/v1" in [str(r.uri) for r in resources] - # Every listed resource must encode to the MCP wire shape. - for r in resources: - assert str(r.to_mcp_resource().uri) - - @pytest.mark.asyncio - async def test_read_text_wraps_as_text_contents(self): - from application.mcp_server import ArtifactResourcesMiddleware - from application.services.artifact_resource_service import ArtifactReadResult - - mw = ArtifactResourcesMiddleware() - ctx = SimpleNamespace(message=SimpleNamespace(uri="artifact://a/v2")) - rr = ArtifactReadResult(uri="artifact://a/v2", mime_type="text/csv", text="a,b") - with ( - patch("application.mcp_server._extract_bearer_token", return_value="k"), - patch( - "application.mcp_server.read_artifact_resource", return_value=rr - ), - ): - result = await mw.on_read_resource(ctx, _call_next_empty) - out = result.to_mcp_result("artifact://a/v2").contents[0] - assert out.text == "a,b" - assert out.mimeType == "text/csv" - - @pytest.mark.asyncio - async def test_read_blob_wraps_as_blob_contents(self): - import base64 - - from application.mcp_server import ArtifactResourcesMiddleware - from application.services.artifact_resource_service import ArtifactReadResult - - mw = ArtifactResourcesMiddleware() - ctx = SimpleNamespace(message=SimpleNamespace(uri="artifact://a/v1")) - blob = base64.b64encode(b"\x89PNG").decode("ascii") - rr = ArtifactReadResult(uri="artifact://a/v1", mime_type="image/png", blob_b64=blob) - with ( - patch("application.mcp_server._extract_bearer_token", return_value="k"), - patch( - "application.mcp_server.read_artifact_resource", return_value=rr - ), - ): - result = await mw.on_read_resource(ctx, _call_next_empty) - out = result.to_mcp_result("artifact://a/v1").contents[0] - assert out.mimeType == "image/png" - assert base64.b64decode(out.blob) == b"\x89PNG" - - @pytest.mark.asyncio - async def test_non_artifact_uri_is_deferred(self): - from application.mcp_server import ArtifactResourcesMiddleware - - mw = ArtifactResourcesMiddleware() - ctx = SimpleNamespace(message=SimpleNamespace(uri="https://example.com/x")) - sentinel = object() - - async def _call_next(context): - return sentinel - - with patch( - "application.mcp_server.read_artifact_resource" - ) as mock_read: - out = await mw.on_read_resource(ctx, _call_next) - assert out is sentinel - mock_read.assert_not_called() - - @pytest.mark.asyncio - async def test_denied_read_raises_permission_error(self): - from application.mcp_server import ArtifactResourcesMiddleware - from application.services.artifact_resource_service import ResourceDenied - - mw = ArtifactResourcesMiddleware() - ctx = SimpleNamespace(message=SimpleNamespace(uri="artifact://foreign/v1")) - with ( - patch("application.mcp_server._extract_bearer_token", return_value="k"), - patch( - "application.mcp_server.read_artifact_resource", - side_effect=ResourceDenied("forbidden"), - ), - ): - with pytest.raises(PermissionError): - await mw.on_read_resource(ctx, _call_next_empty)