diff --git a/src/xagent/web/api/chat.py b/src/xagent/web/api/chat.py index 9b9ac2ec35..50479fe2aa 100644 --- a/src/xagent/web/api/chat.py +++ b/src/xagent/web/api/chat.py @@ -9,6 +9,7 @@ from typing import Any, Callable, Dict, List, Mapping, Optional, TypeVar, Union, cast from fastapi import APIRouter, Depends, HTTPException, Request +from fastapi.responses import JSONResponse from sqlalchemy import func, or_, update from sqlalchemy.orm import Session @@ -75,7 +76,10 @@ parse_user_sandbox_key, ) from ..schemas.chat import TaskCreateRequest, TaskCreateResponse -from ..schemas.connector_runtime import ConnectorRuntimeRequirementsModel +from ..schemas.connector_runtime import ( + ConnectorRuntimeRequirementsModel, + ConnectorRuntimeValuesRequest, +) from ..services.agent_access import list_accessible_published_agents from ..services.agent_team_scope import ( get_agent_team_scope, @@ -90,6 +94,7 @@ ) from ..services.client_error_messages import ClientErrorCode, client_error_message from ..services.connector_runtime import ( + apply_task_connector_runtime_context_values, bind_connector_runtime_selection_snapshot, build_task_runtime_requirements, resolve_agent_runtime_requirements, @@ -5301,6 +5306,32 @@ async def get_task_runtime_extensions( } +def _connector_runtime_error_response(exc: ConnectorRuntimeError) -> JSONResponse: + """Render a connector-runtime failure in this endpoint's error envelope. + + Mirrors the ``{"error": {"code", "message", "details"}}`` shape + ``_raise_v1_connector_runtime_error`` uses for the /v1 surface + (``api/v1/tasks.py``) -- only the envelope is shared, not ``V1ErrorCode``, + which is a separate SDK-facing contract this endpoint does not + participate in. The status code always comes from ``exc.status_code``, + never recomputed from ``exc.code`` via ``_status_for_code``: that helper + cannot produce 503, and both this endpoint's own conditional-update + failure and a team-scope resolution failure construct their + ``ConnectorRuntimeError`` with an explicit ``status_code=503``. + """ + + return JSONResponse( + status_code=exc.status_code, + content={ + "error": { + "code": exc.code, + "message": exc.safe_message, + "details": exc.to_public_error()["details"], + } + }, + ) + + @chat_router.get( "/agent/{agent_id}/connector-runtime-requirements", response_model=ConnectorRuntimeRequirementsModel, @@ -5387,6 +5418,94 @@ async def get_task_connector_runtime_requirements( ) from exc +@chat_router.post( + "/task/{task_id}/connector-runtime-values", + response_model=ConnectorRuntimeRequirementsModel, +) +def post_task_connector_runtime_values( + task_id: int, + request: ConnectorRuntimeValuesRequest, + db: Session = Depends(get_db), + user: User = Depends(get_current_user), +) -> ConnectorRuntimeRequirementsModel | JSONResponse: + """Accept ``context`` values for a task's connectors, merged key by key. + + A key not yet stored is written; one already stored with the identical + value is a no-op; one already stored with a different value fails the + whole request with no partial write (``runtime_context_immutable``, + 409). There is no override switch: a stored value is never replaced. + ``secrets``/``auth_selector`` are out of scope this phase and rejected + at the request-shape level (``extra="forbid"``), not accepted and + ignored. + + Access is the same plain task ownership the task-keyed read endpoint + applies -- ``Task.user_id == current_user.id`` in the query that loads + the task, with no admin exception -- so a task that does not exist and + one that is not the caller's own answer the same 404. + + On success, the response is the same requirements report the read + endpoints return, reflecting exactly what was just written -- not the + request's own echo, and not whatever the read endpoints would have + said before this call. A 200 here means only that the submitted keys + were merged in; it says nothing about whether every required input is + now present, which is what the response's own ``satisfied`` fields + answer. + + Every ``ConnectorRuntimeError`` raised anywhere on this request path -- + validation, a stored-value conflict, a concurrent write racing this one + to the same row, and the agent/workforce resolution that runs before + the write -- is rendered through ``_connector_runtime_error_response`` + rather than the plain-``detail`` ``HTTPException`` the two read + endpoints above use, because a caller needs the structured + ``code``/``details.reason`` to decide how to recover (retry, refresh + and drop already-satisfied keys, or give up), not just a status code. + + Anything else -- a resolver raising on a malformed workforce snapshot, + a driver error mid-write, a failing ``db.commit()`` -- is rolled back + and re-raised unchanged, so it ends as a bare 500. Such a failure is + deliberately not dressed up in this envelope: the envelope's ``code`` + is a contract about a condition the caller can act on, and an + unclassified fault is not one. What the rollback guarantees either way + is that no partial batch survives the request. + """ + + task = db.query(Task).filter(Task.id == task_id, Task.user_id == user.id).first() + if task is None: + raise HTTPException(status_code=404, detail="Task not found") + # The connector scope this endpoint writes into, and reports on, must + # be the scope a turn would actually run under, so the agent is + # resolved by the same two calls the per-turn tool build makes for + # this task, in the same order -- the identical pair the task-keyed + # read endpoint above uses, so neither endpoint can answer with a + # connector set the other would not. Reading the agent row directly + # would key team-shared connectors on the raw row's team even where + # the runtime resolves the agent to None, and resolving with no + # workforce runtime would drop the team of a workforce manager agent + # whose run the runtime does find. Here the scope decides not only + # what the response lists but which connectors the caller may write + # to at all, so neither the over-report nor the under-report is + # acceptable. + try: + workforce_runtime = resolve_workforce_task_runtime(db, task) + agent = _load_agent_for_task_runtime(db, task, workforce_runtime) + requirements = apply_task_connector_runtime_context_values( + db=db, task=task, agent=agent, payload_items=request.items + ) + db.commit() + except ConnectorRuntimeError as exc: + db.rollback() + return _connector_runtime_error_response(exc) + except Exception: + # Not a fallback: the exception keeps travelling and this request + # still ends as a bare 500. The rollback is the whole point -- + # a commit that raises leaves the session's transaction open with + # this batch's rows already flushed into it, and every other + # failure past the flush leaves the same thing behind. + db.rollback() + raise + return requirements + + @chat_router.delete("/task/{task_id}") async def delete_task( task_id: int, diff --git a/src/xagent/web/schemas/connector_runtime.py b/src/xagent/web/schemas/connector_runtime.py index 20f5d6c6f7..f81301b95e 100644 --- a/src/xagent/web/schemas/connector_runtime.py +++ b/src/xagent/web/schemas/connector_runtime.py @@ -1,25 +1,25 @@ -"""Connector runtime requirements: the shared response shape both read -endpoints return. - -The agent-keyed and task-keyed read endpoints both return this same -requirements report -- which runtime inputs a task's (or a prospective -task's) connectors declare, and whether each one already has a value -- -never a stored value itself, and never a connector's transport or -authentication configuration -- the same shape from both, with one field, -``satisfied``, answering a different question on each; see +"""Connector runtime requirements: shared response shape and the values +request body. + +Four producers share the response shape: the agent-keyed and task-keyed +read endpoints, the task-create response, and the values endpoint's 200 +response. All four describe a requirements report -- which runtime inputs +a task's (or a prospective task's) connectors declare, and whether each +one already has a value -- never a stored value itself, and never a +connector's transport or authentication configuration. One field, +``satisfied``, answers a different question depending on the producer; see ``ConnectorRuntimeRequirementsModel``. -Placed in its own module rather than ``schemas/chat.py`` because a values- -submission endpoint lands on top of it shortly and will share this same -response shape as its own 200 body; putting it here now avoids a later -move that would touch every existing importer. +Placed in its own module rather than ``schemas/chat.py`` or ``schemas/v1.py`` +because it has more than one audience: folding it into either of those +modules would couple that module's own audience to the others. """ from __future__ import annotations from typing import Literal -from pydantic import BaseModel +from pydantic import BaseModel, ConfigDict class ConnectorRuntimeRefModel(BaseModel): @@ -107,3 +107,30 @@ class ConnectorRuntimeRequirementsModel(BaseModel): satisfied: bool secrets_expires_at: str | None connectors: list[ConnectorRuntimeConnectorModel] + + +class ConnectorRuntimeValueItem(BaseModel): + """One connector's caller-supplied context values. + + ``secrets`` and ``auth_selector`` are deliberately absent from this + phase's request shape: ``extra="forbid"`` turns either one into a 422 + rather than silently accepting and discarding it. + """ + + model_config = ConfigDict(extra="forbid") + + connector_ref: ConnectorRuntimeRefModel + context: dict[str, object] | None = None + + +class ConnectorRuntimeValuesRequest(BaseModel): + """Body of ``POST /api/chat/task/{task_id}/connector-runtime-values``. + + No override switch of any kind belongs here: a stored value is never + replaced, and ``extra="forbid"`` turns an attempt to add one (for + example ``if_absent`` or ``force``) into a 422. + """ + + model_config = ConfigDict(extra="forbid") + + items: list[ConnectorRuntimeValueItem] diff --git a/src/xagent/web/services/connector_runtime.py b/src/xagent/web/services/connector_runtime.py index 2ce93f781b..e25f7f0e6b 100644 --- a/src/xagent/web/services/connector_runtime.py +++ b/src/xagent/web/services/connector_runtime.py @@ -14,6 +14,12 @@ from threading import RLock from typing import Any, Iterable, cast +from sqlalchemy import Text +from sqlalchemy import cast as sa_cast +from sqlalchemy import select +from sqlalchemy import update as sa_update +from sqlalchemy.engine import CursorResult +from sqlalchemy.exc import IntegrityError from sqlalchemy.orm import Session from ...core.tools.adapters.vibe.connector_runtime import ( @@ -50,10 +56,17 @@ ConnectorRuntimeInputModel, ConnectorRuntimeRefModel, ConnectorRuntimeRequirementsModel, + ConnectorRuntimeValueItem, ) logger = logging.getLogger(__name__) +# Two independent size caps on the values endpoint's request: no single +# key's value may serialize past the first, and the whole batch may not +# pass the second even if every individual key is within the first. +_RUNTIME_CONTEXT_VALUE_MAX_BYTES = 64 * 1024 +_RUNTIME_CONTEXT_REQUEST_MAX_BYTES = 256 * 1024 + @dataclass(frozen=True) class ConnectorRuntimePayload: @@ -77,6 +90,21 @@ class ConnectorRuntimeAppendPlan: ephemeral_by_ref: dict[ConnectorRef, dict[str, dict[str, Any]]] +@dataclass(frozen=True) +class _ContextRowSnapshot: + """One ``task_connector_runtime_contexts`` row as read for a write attempt. + + ``context_text`` is ``CAST(context AS TEXT)`` rendered by the database + itself in the same query that read ``context`` -- never a Python-side + re-serialization. It is the only value a subsequent conditional UPDATE + may compare against; see ``_update_context_row``. + """ + + id: int + context: dict[str, Any] + context_text: str + + @dataclass(frozen=True) class ConnectorRuntimeValues: context: dict[str, Any] @@ -434,8 +462,9 @@ def resolve_agent_runtime_requirements( -- same filter, same canonical order -- because a caller creating a task persists them into ``Task.connector_runtime_selected_refs``, and every later reader of that column (the per-turn gate, - ``load_connector_runtime_view``) depends on it holding exactly that - set. Both the write and the read of that column sort by + ``load_connector_runtime_view``, and the values endpoint's selection + check in ``_validate_payload_refs``) depends on it holding exactly + that set. Both the write and the read of that column sort by ``(connector_type, connector_id)``, so the invariant the column carries is the set under that canonical order, not the order a caller happened to hand over. Do not derive the refs any other way, even one @@ -496,6 +525,198 @@ def build_task_runtime_requirements( return _build_task_requirements_model(selected_refs, visible, stored_context) +def apply_task_connector_runtime_context_values( + *, + db: Session, + task: Task, + agent: Agent | None, + payload_items: Iterable[ConnectorRuntimeValueItem], +) -> ConnectorRuntimeRequirementsModel: + """Merge caller-supplied ``context`` values into a task's connector rows. + + Task ownership is the caller's job -- this function trusts the ``task`` + it is handed and never queries for it. Everything else runs here, in + this order: + + 1. Request shape: a non-empty item list where every item carries at + least one ``context`` key (an empty mapping and an absent one are + the same thing). + 2. (caller's job -- task ownership.) + 3. A size cap on each key's serialized value and on the batch as a + whole. + 4. No ref repeated within the batch. + 5. Every ref is currently visible and was selected for this task + (``_validate_payload_refs``, shared with the create and append + paths). + 6. Every key is declared and syntactically valid + (``_validate_values_against_schema``, reused unchanged). + 7. Every value matches its declared, normalized type and carries + content -- an empty or whitespace-only string, and an empty + object, are rejected for optional keys as well as required ones. + 7b. (encryption-key-configuration gate -- a later phase's tier; this + phase's request shape has no ``secrets`` field, so there is + nothing for it to check yet. See the note below tier 8.) + 8. A key already stored with a different value fails the whole batch; + one stored with the same value is a no-op; one not yet stored is + added. + There is no tier 9: this endpoint never asserts that a required key is + present, only that what was sent is valid and conflict-free. + + A later phase that adds the ``secrets`` section must not let a stored + secret's decode failure (``EncryptionDecodeError``) be caught and + treated as "no value stored" -- that would resurrect a value that + failed to decrypt as though it had never been set, defeating tier 8's + conflict detection for that key. This phase never reads or decodes a + secret value, so the case cannot occur yet, but the failure mode is + reserved: do not add a broad ``except Exception`` around a future + secret read here. + + Tier 6 must run before 7 and 8: both of those format a caller-supplied + ``key`` into a ``reason`` string the caller gets back in the response, + and that is only safe once tier 6 has confirmed the key matches + ``validate_runtime_source_key`` and is declared by the connector -- + never an arbitrary string the caller wrote. + + Cross-ref ordering splits in two. Validation walks the batch in the + order ``payload_items`` arrives in, so which of several bad refs is + reported is the caller's own ordering; that is not a leak, because an + invisible ref answers a uniform 404 regardless of position and a + visible-but-unselected ref answers 400 for something the caller can + already see. The write, in contrast, is issued in canonical ref order + (``_sort_connector_refs``) no matter how the batch arrived, so two + concurrent requests naming the same rows in opposite order still take + those rows' locks in the same order and cannot deadlock each other. + + The write is a compare-and-swap on the text the database itself + rendered for a row's previous content, retried once after a full + reread-and-remerge if the swap misses (either an ``IntegrityError`` on + a first-ever row for a ref, or an updated-row-count of zero) -- a + second miss surfaces as 503, never a third attempt. A stored-value + conflict (tier 8) is a different thing entirely and is never retried: + it fails the request immediately, before anything is written. + + Never commits: it only ``add``s, ``flush``es and ``execute``s. The + caller commits once this returns, so every key in the batch lands in + the same transaction -- a mid-batch failure leaves nothing behind. + """ + + parsed = _parse_context_value_items(payload_items) + if not parsed: + raise ConnectorRuntimeError( + ERROR_INVALID_RUNTIME_CONTEXT, + _message_for_code(ERROR_INVALID_RUNTIME_CONTEXT), + details={"reason": "empty_items"}, + ) + for ref, context in parsed: + if not context: + _raise_runtime_error( + ERROR_INVALID_RUNTIME_CONTEXT, ref, reason="empty_item_payload" + ) + _reject_oversized_context_payload(parsed) + _reject_duplicate_context_refs(parsed) + payload_by_ref: dict[ConnectorRef, dict[str, Any]] = dict(parsed) + + connector_user_id = int(task.user_id) + agent_team_id = ( + int(agent.team_id) if agent is not None and agent.team_id is not None else None + ) + visible = _load_visible_runtime_connectors( + db, user_id=connector_user_id, agent_team_id=agent_team_id + ) + selected_refs = _load_task_selected_refs(task) + _validate_payload_refs(payload_by_ref, visible=visible, selected_refs=selected_refs) + for ref, context in payload_by_ref.items(): + connector = visible[ref] + _validate_values_against_schema(ref, connector, context, {}, {}) + _validate_context_value_types(ref, connector, context) + # Tier 7b, the encryption-key-configuration gate, belongs to the + # phase that adds the `secrets` section. This request shape has no + # `secrets` field (`extra="forbid"` rejects one), so this tier has + # no trigger condition here -- do not add a placeholder branch. + + task_id = int(task.id) + # Every statement this batch issues against a row goes out in canonical + # ref order, never the order the caller happened to list the refs in: + # two requests that touch the same two rows in opposite order would + # otherwise be able to take the same two row locks in opposite order + # and deadlock on PostgreSQL. The same ordering the task's selection + # snapshot is written in (`bind_connector_runtime_selection_snapshot`). + write_order = _sort_connector_refs(payload_by_ref) + rows: dict[ConnectorRef, _ContextRowSnapshot] = {} + merged_by_ref: dict[ConnectorRef, dict[str, Any]] = {} + written_keys_by_ref: dict[ConnectorRef, list[str]] = {} + for attempt in range(2): + rows = _load_task_context_row_snapshots(db, task_id=task_id) + merged_by_ref = {} + for ref, context in payload_by_ref.items(): + stored_row = rows.get(ref) + stored_context = stored_row.context if stored_row is not None else {} + merged_by_ref[ref] = _merge_context_values(ref, stored_context, context) + + conflict = False + written_keys_by_ref = {} + for ref in write_order: + merged = merged_by_ref[ref] + stored_row = rows.get(ref) + if stored_row is None: + db.add( + TaskConnectorRuntimeContext( + task_id=task_id, + connector_type=ref.connector_type, + connector_id=ref.connector_id, + context=merged, + ) + ) + written_keys_by_ref[ref] = sorted(merged) + continue + if _canonical_json_text(merged) == _canonical_json_text(stored_row.context): + # Same content already on the row: no write statement at + # all, not even a same-value UPDATE. + continue + if not _update_context_row( + db, + row_id=stored_row.id, + context_text_as_read=stored_row.context_text, + merged=merged, + ): + conflict = True + break + written_keys_by_ref[ref] = sorted(set(merged) - set(stored_row.context)) + + if conflict: + visible = _retry_or_503( + db, attempt=attempt, user_id=connector_user_id, team_id=agent_team_id + ) + continue + + try: + db.flush() + except IntegrityError: + visible = _retry_or_503( + db, attempt=attempt, user_id=connector_user_id, team_id=agent_team_id + ) + continue + break + + # The only record of a fill, per the design's disclosure: key names, + # the connector ref, and a count -- never a value. One line per ref + # that actually gained a key; a ref where every submitted key already + # matched what was stored logs nothing, because nothing was written. + for ref, keys in written_keys_by_ref.items(): + logger.info( + "Connector runtime context values written " + "(task_id=%d, connector_ref=%s, key_count=%d, keys=%s)", + task_id, + ref.storage_key, + len(keys), + keys, + ) + + stored_context_after = {ref: row.context for ref, row in rows.items()} + stored_context_after.update(merged_by_ref) + return _build_task_requirements_model(selected_refs, visible, stored_context_after) + + def bind_connector_runtime_selection_snapshot( *, task: Task, selected_refs: Iterable[ConnectorRef] ) -> None: @@ -817,13 +1038,18 @@ def _has_runtime_declaration(connector: Any) -> bool: def _validate_payload_refs( - payload_by_ref: dict[ConnectorRef, ConnectorRuntimePayload], + refs: Iterable[ConnectorRef], *, visible: dict[ConnectorRef, Any], selected_refs: tuple[ConnectorRef, ...], ) -> None: + """Every ref must be visible to the caller now and have been selected + by the task. Takes the refs alone: nothing here reads a submitted + value, so a caller holding only refs does not have to build payload + objects to call it. A mapping keyed by ref satisfies this as it is. + """ selected = set(selected_refs) - for ref in payload_by_ref: + for ref in refs: if ref not in visible: _raise_runtime_error(ERROR_CONNECTOR_NOT_FOUND, ref) if ref not in selected: @@ -1014,6 +1240,256 @@ def _build_task_requirements_model( ) +def _parse_context_value_items( + payload_items: Iterable[ConnectorRuntimeValueItem], +) -> list[tuple[ConnectorRef, dict[str, Any]]]: + """Parse the values endpoint's items into ``(ref, context)`` pairs, + without deduplicating or size-checking -- those are separate, later + tiers (4 and 3 respectively; see ``apply_task_connector_runtime_context_values``). + """ + parsed: list[tuple[ConnectorRef, dict[str, Any]]] = [] + for item in payload_items or (): + raw_ref = item.connector_ref.model_dump() + try: + ref = ConnectorRef.from_wire(raw_ref) + except ValueError as exc: + # An unknown connector_type, or a connector_id that fails + # ``ConnectorRef.from_wire``'s ``isinstance(connector_id, int)`` + # check, cannot identify any real connector, so this is the same + # "not found" the visibility check below reports for a + # well-formed ref outside the caller's visible set -- not a + # request-shape error. (A non-positive connector_id passes this + # check -- from_wire does not validate sign -- and is rejected + # by that same visibility check instead, one step later.) + raise ConnectorRuntimeError( + ERROR_CONNECTOR_NOT_FOUND, + _message_for_code(ERROR_CONNECTOR_NOT_FOUND), + status_code=_status_for_code(ERROR_CONNECTOR_NOT_FOUND), + ) from exc + context = item.context or {} + parsed.append((ref, context)) + return parsed + + +def _reject_oversized_context_payload( + items: list[tuple[ConnectorRef, dict[str, Any]]], +) -> None: + """Tier 3: a per-key cap and a whole-batch cap, checked before this + function touches the database -- the caller (the values-endpoint route) + has already looked up the task and agent by the time this runs, so this + is not the request's first database access. Which key was oversized + goes to the log only -- never into the response ``reason`` -- and the + log line never carries the value itself, only the key name, the + connector ref, and a byte count. + + Both caps measure the value in its stored form. ``json.dumps`` defaults + to ``ensure_ascii=True``, so a character outside ASCII counts as the + six-character escape sequence JSON writes it as -- which is exactly + what SQLAlchemy's JSON column puts in the row, and exactly the text the + conditional update in ``_update_context_row`` compares against. + Measuring the raw UTF-8 length instead would admit a value two to three + times the size of the column the caps exist to bound. A caller whose + text is CJK or emoji therefore reaches a cap at roughly half, or a + third, of the character count an ASCII caller does; that is the cap's + intended meaning, not an oversight. + """ + total_bytes = 0 + for ref, context in items: + for key, value in context.items(): + encoded_len = len(json.dumps(value).encode("utf-8")) + total_bytes += encoded_len + if encoded_len > _RUNTIME_CONTEXT_VALUE_MAX_BYTES: + logger.warning( + "Connector runtime context value exceeds the per-key " + "size cap (connector_ref=%s, key=%r, bytes=%d)", + ref.storage_key, + key, + encoded_len, + ) + _raise_runtime_error( + ERROR_INVALID_RUNTIME_CONTEXT, ref, reason="payload_too_large" + ) + if total_bytes > _RUNTIME_CONTEXT_REQUEST_MAX_BYTES: + logger.warning( + "Connector runtime values request exceeds the aggregate size " + "cap (bytes=%d)", + total_bytes, + ) + raise ConnectorRuntimeError( + ERROR_INVALID_RUNTIME_CONTEXT, + _message_for_code(ERROR_INVALID_RUNTIME_CONTEXT), + details={"reason": "payload_too_large"}, + ) + + +def _reject_duplicate_context_refs( + items: list[tuple[ConnectorRef, dict[str, Any]]], +) -> None: + """Tier 4: no ref may appear twice in the same batch.""" + seen: set[ConnectorRef] = set() + for ref, _context in items: + if ref in seen: + _raise_runtime_error( + ERROR_INVALID_RUNTIME_CONTEXT, ref, reason="duplicate_ref" + ) + seen.add(ref) + + +def _validate_context_value_types( + ref: ConnectorRef, connector: Any, context: dict[str, Any] +) -> None: + """Tier 7: every value matches its declaration's normalized type and + carries content. + + Two rules, in this order. The type first: a key declared ``object`` + takes a mapping, every other key takes a string. Then content: a + string that is empty or strips to nothing, and an empty mapping, are + both rejected. That holds whether or not the key is required -- a + stored value is never replaced, so a blank one that got in would mark + its key filled forever while carrying nothing a connector can use, and + the only way out would be a new task. + + Must run after tier 6 (``_validate_values_against_schema``): the ``key`` + formatted into ``reason`` here is only safe to reflect back to the + caller because tier 6 has already confirmed it is syntactically valid + and declared, never an arbitrary caller-written string. + """ + declarations = _schema_section( + _runtime_input_schema(connector), RUNTIME_INPUT_CONTEXT + ) + for key, value in context.items(): + declared_type = _normalize_runtime_input_type(declarations.get(key)) + matches = ( + isinstance(value, dict) + if declared_type == "object" + else isinstance(value, str) + ) + if not matches: + _raise_runtime_error( + ERROR_INVALID_RUNTIME_CONTEXT, + ref, + reason=f"type_mismatch.{RUNTIME_INPUT_CONTEXT}.{key}", + ) + blank = not value if declared_type == "object" else not value.strip() + if blank: + _raise_runtime_error( + ERROR_INVALID_RUNTIME_CONTEXT, + ref, + reason=f"empty_value.{RUNTIME_INPUT_CONTEXT}.{key}", + ) + + +def _merge_context_values( + ref: ConnectorRef, stored: dict[str, Any], incoming: dict[str, Any] +) -> dict[str, Any]: + """Tier 8: per-key merge. A key not yet stored is added; one stored + with the same value is a no-op; one stored with a different value + fails the whole batch immediately, before anything is written. + + "The same value" means the same canonical JSON text, not two objects + Python's ``==`` calls equal: a resubmission that differs only in JSON + form -- ``1`` against ``1.0``, or ``1`` against ``true`` -- would store + something different and so is a conflict, not a no-op. + """ + merged = dict(stored) + for key, value in incoming.items(): + if key in stored: + if _canonical_json_text(stored[key]) != _canonical_json_text(value): + _raise_runtime_error( + ERROR_RUNTIME_CONTEXT_IMMUTABLE, + ref, + reason=f"conflict.{RUNTIME_INPUT_CONTEXT}.{key}", + ) + continue + merged[key] = value + return merged + + +def _load_task_context_row_snapshots( + db: Session, *, task_id: int +) -> dict[ConnectorRef, _ContextRowSnapshot]: + """Read every context row for a task along with the database's own + ``CAST(context AS TEXT)`` rendering, in the same query -- the rendering + is the only value a later conditional UPDATE may compare against; see + ``_update_context_row``. + + Takes no row lock of any kind (plain ``SELECT``, no + ``with_for_update()``). Concurrent writers are serialized at that later + conditional UPDATE's ``WHERE`` clause, not here -- a row read by this + function may already be stale by the time ``_update_context_row`` runs, + which is exactly the case the conditional predicate exists to catch. + """ + stmt = select( + TaskConnectorRuntimeContext.id, + TaskConnectorRuntimeContext.connector_type, + TaskConnectorRuntimeContext.connector_id, + TaskConnectorRuntimeContext.context, + sa_cast(TaskConnectorRuntimeContext.context, Text).label("context_text"), + ).where(TaskConnectorRuntimeContext.task_id == task_id) + result: dict[ConnectorRef, _ContextRowSnapshot] = {} + for row in db.execute(stmt): + connector_type = cast(ConnectorType, str(row.connector_type)) + ref = ConnectorRef(connector_type, int(row.connector_id)) + result[ref] = _ContextRowSnapshot( + id=int(row.id), + context=row.context if isinstance(row.context, dict) else {}, + context_text=row.context_text, + ) + return result + + +def _update_context_row( + db: Session, *, row_id: int, context_text_as_read: str, merged: dict[str, Any] +) -> bool: + """Conditional update: only takes effect if the row's rendered text is + still exactly what was read. PostgreSQL's ``json`` column type has no + equality operator, so the comparison casts both sides to text rather + than comparing ``context`` directly; SQLite has no + ``IS NOT DISTINCT FROM``, so the comparison is a plain ``=`` and a row + with no matching id/text is simply not found, not an error. Returns + whether exactly one row was updated. + """ + stmt = ( + sa_update(TaskConnectorRuntimeContext) + .where( + TaskConnectorRuntimeContext.id == row_id, + sa_cast(TaskConnectorRuntimeContext.context, Text) == context_text_as_read, + ) + .values(context=merged) + ) + result = cast(CursorResult, db.execute(stmt)) + return int(result.rowcount) == 1 + + +def _retry_or_503( + db: Session, *, attempt: int, user_id: int, team_id: int | None +) -> dict[ConnectorRef, Any]: + """Recover a lost write inside the values endpoint's bounded retry + loop, or give up on it. + + A write is lost two ways -- a conditional update that matched no row, + and an ``IntegrityError`` on a first-ever row for a ref -- and both are + answered identically, so both branches come here. On the second + attempt there is no third: the request ends as a 503 the caller may + retry, with no named reason, which is what tells that caller this was a + collision rather than something underneath failing. Otherwise the + transaction is rolled back and a freshly loaded visible-connector set + is handed back, because ``db.rollback()`` expires every ORM object + loaded before it, including the connector rows in ``visible`` -- + refetching once here is cheaper than letting the retry (and the + response assembly after it) trigger a per-object lazy-refresh query + apiece. + """ + db.rollback() + if attempt == 1: + raise ConnectorRuntimeError( + ERROR_CONNECTOR_RUNTIME_UNAVAILABLE, + "Connector runtime context is unavailable.", + status_code=503, + ) + return _load_visible_runtime_connectors(db, user_id=user_id, agent_team_id=team_id) + + def _validate_values_against_schema( ref: ConnectorRef, connector: Any, @@ -1214,11 +1690,24 @@ def _is_required(declaration: Any) -> bool: return isinstance(declaration, dict) and bool(declaration.get("required")) +def _canonical_json_text(value: Any) -> str: + """The value's canonical JSON text: keys sorted, no incidental + whitespace. Two values render the same text only when they would store + the same JSON, which comparing the decoded objects does not answer -- + Python reads ``1``, ``1.0`` and ``True`` as equal to each other, while + the JSON they store (``1``, ``1.0``, ``true``) differs. Compare values + through this whenever the question is "would storing this change what + is already there". + """ + return json.dumps(value, sort_keys=True, separators=(",", ":")) + + +def _canonical_json(value: Any) -> Any: + return json.loads(_canonical_json_text(value)) + + def _canonical_json_value(value: dict[str, Any]) -> dict[str, Any]: - return cast( - dict[str, Any], - json.loads(json.dumps(value, sort_keys=True, separators=(",", ":"))), - ) + return cast(dict[str, Any], _canonical_json(value)) def _raise_runtime_error( diff --git a/tests/web/test_connector_runtime_entrypoints_e2e.py b/tests/web/test_connector_runtime_entrypoints_e2e.py index dda240a0e1..b0c963277e 100644 --- a/tests/web/test_connector_runtime_entrypoints_e2e.py +++ b/tests/web/test_connector_runtime_entrypoints_e2e.py @@ -12,8 +12,10 @@ import pytest from fastapi import FastAPI from fastapi.testclient import TestClient +from sqlalchemy import event from sqlalchemy.orm import Session +from xagent.core.tools.adapters.vibe.connector_runtime import ConnectorRuntimeError from xagent.web.api.auth import auth_router, create_access_token from xagent.web.api.chat import AgentServiceManager, chat_router from xagent.web.api.public_chat_access import create_public_chat_access_token @@ -30,6 +32,7 @@ Base, get_db, get_engine, + get_session_local, ) from xagent.web.models.deployment import Deployment, DeploymentOwnerType from xagent.web.models.mcp import MCPServer, UserMCPServer @@ -2864,3 +2867,2115 @@ def test_create_task_persists_same_selected_refs_as_legacy_snapshot( persisted_refs = _task(task_id).connector_runtime_selected_refs expected_wire = [ref.to_wire() for ref in expected_refs] assert persisted_refs == expected_wire + + +# --------------------------------------------------------------------------- +# POST /task/{task_id}/connector-runtime-values. +# --------------------------------------------------------------------------- + + +def _values_url(task_id: int) -> str: + return f"/api/chat/task/{task_id}/connector-runtime-values" + + +def _setup_context_task( + *, + required: bool = True, + key_type: str = "string", + server_name: str = "a4-server", +) -> tuple[dict[str, str], int, int]: + """Owner-only setup shared by most of the values-endpoint tests below: + one MCP connector with a single declared ``context`` key, one task that + selected it. Returns (owner headers, task_id, server_id). + """ + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + server = _mcp_server_with_context_schema( + db, + user, + name=server_name, + context_schema={"auth_token": {"type": key_type, "required": required}}, + ) + agent = _create_agent( + db, user, name=f"{server_name}-agent", tool_categories=["mcp"] + ) + db.commit() + db.refresh(agent) + db.refresh(server) + agent_id = int(agent.id) + server_id = int(server.id) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={"title": server_name, "description": "d", "agent_id": agent_id}, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + return headers, task_id, server_id + + +def test_values_endpoint_requires_task_ownership_by_caller(e2e_db: None) -> None: + """Three identities that must all see a uniform 404, matching the read + endpoints' ownership predicate -- including a non-owner whose request + body is itself oversized, so that the ownership check cannot be + inferred to run after the size check from a caller watching only the + status code (a non-owner would then get 400 for a huge body and 404 for + a small one, which leaks whether the task exists).""" + _, task_id, server_id = _setup_context_task() + db = _db_session() + try: + other = _create_user(db, "values-other-user") + db.commit() + other_headers = _auth_headers_for_user(other) + finally: + db.close() + + body = { + "items": [ + { + "connector_ref": {"connector_type": "mcp", "connector_id": server_id}, + "context": {"auth_token": "x"}, + } + ] + } + + other_response = client.post(_values_url(task_id), headers=other_headers, json=body) + assert other_response.status_code == 404, other_response.text + + oversized_body = { + "items": [ + { + "connector_ref": {"connector_type": "mcp", "connector_id": server_id}, + "context": {"auth_token": "x" * (64 * 1024 + 1)}, + } + ] + } + oversized_other_response = client.post( + _values_url(task_id), headers=other_headers, json=oversized_body + ) + assert oversized_other_response.status_code == 404, oversized_other_response.text + + anon_response = client.post(_values_url(task_id), json=body) + assert anon_response.status_code == 403, anon_response.text + + widget_token = create_public_chat_access_token( + {"guest_id": "values-anon-guest", "widget_agent_id": 1} + ) + widget_response = client.post( + _values_url(task_id), + headers={"Authorization": f"Bearer {widget_token}"}, + json=body, + ) + assert widget_response.status_code == 401, widget_response.text + assert _context_row_count(task_id) == 0 + + +def test_values_endpoint_gives_an_admin_no_write_to_another_task( + e2e_db: None, +) -> None: + """Being an admin buys no write to another user's task here either: + this endpoint filters on the caller's own user id with no admin + exception, and answers the same "Task not found" 404 the task-keyed + read endpoint answers for the same caller. The owner's own write of + the same body answers 200, so the admin's 404 is about the caller, not + about the id or the body. + """ + admin_headers = _setup_admin_headers() + db = _db_session() + try: + owner = _create_user(db, "values-admin-probe-owner") + db.flush() + server = _mcp_server_with_context_schema( + db, + owner, + name="values-admin-probe-server", + context_schema={"auth_token": {"type": "string", "required": True}}, + url="https://example.com/values-admin-probe/mcp", + ) + agent = _create_agent( + db, owner, name="Values Admin Probe Agent", tool_categories=["mcp"] + ) + db.commit() + db.refresh(agent) + db.refresh(server) + agent_id = int(agent.id) + server_id = int(server.id) + owner_headers = _auth_headers_for_user(owner) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=owner_headers, + json={ + "title": "values admin probe task", + "description": "d", + "agent_id": agent_id, + }, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + + body = { + "items": [ + { + "connector_ref": {"connector_type": "mcp", "connector_id": server_id}, + "context": {"auth_token": "x"}, + } + ] + } + + admin_response = client.post(_values_url(task_id), headers=admin_headers, json=body) + assert admin_response.status_code == 404, admin_response.text + assert admin_response.json()["detail"] == "Task not found" + assert _context_row_count(task_id) == 0 + + owner_response = client.post(_values_url(task_id), headers=owner_headers, json=body) + assert owner_response.status_code == 200, owner_response.text + assert _context_row_count(task_id) == 1 + + +def test_values_endpoint_surfaces_team_scope_failure_as_typed_503( + e2e_db: None, +) -> None: + """A team-scope resolution failure surfaces as 503 with + ``reason == "team_scope_resolution_failed"``, not a 400 -- setup must + install the team hook first (``_load_visible_runtime_connectors`` only + calls ``resolve_team_connector_ids_or_raise``, which is what raises + this, when the hook is installed; without it the branch that produces + this 503 is unreachable and the assertion below would be probing a + no-op setup instead of the endpoint). + """ + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + agent = _create_agent( + db, user, name="Team Scope Failure Agent", tool_categories=["mcp"] + ) + agent.team_id = 303 + db.commit() + db.refresh(agent) + agent_id = int(agent.id) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={"title": "i10b task", "description": "d", "agent_id": agent_id}, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + + def _raising_hook(db: Session, *, team_id: int) -> dict[str, set[int]]: + raise RuntimeError("team-scope-hook-failure-must-not-leak") + + connector_team_scope.set_connector_team_hooks(team_visibility=_raising_hook) + try: + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": {"connector_type": "mcp", "connector_id": 1}, + "context": {"auth_token": "x"}, + } + ] + }, + ) + finally: + connector_team_scope.set_connector_team_hooks() + + assert response.status_code == 503, response.text + assert "team-scope-hook-failure-must-not-leak" not in response.text + body = response.json() + assert body["error"]["details"]["reason"] == "team_scope_resolution_failed" + + +@pytest.mark.parametrize( + ("payload_builder", "expected_reason_prefix"), + [ + ( + lambda server_id: {"other_key": "x"}, + "undeclared_context_key", + ), + ( + lambda server_id: {"auth_token": {"nested": "object"}}, + "type_mismatch.context.auth_token", + ), + ( + lambda server_id: {"auth_token": 5}, + "type_mismatch.context.auth_token", + ), + ], +) +def test_values_endpoint_rejects_invalid_values_and_writes_nothing( + e2e_db: None, payload_builder: Any, expected_reason_prefix: str +) -> None: + """An undeclared key, a string-typed key given an object, and a + string-typed key given an int are each rejected, and each rejection + leaves zero rows behind.""" + headers, task_id, server_id = _setup_context_task() + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": payload_builder(server_id), + } + ] + }, + ) + assert response.status_code == 400, response.text + assert response.json()["error"]["details"]["reason"] == expected_reason_prefix + assert _context_row_count(task_id) == 0 + + +def test_values_endpoint_rejects_object_key_given_wrong_shapes(e2e_db: None) -> None: + """An object-typed key given a string, then given an array, is rejected + both times.""" + headers, task_id, server_id = _setup_context_task(key_type="object") + for bad_value in ("not-an-object", ["also", "not", "an", "object"]): + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": bad_value}, + } + ] + }, + ) + assert response.status_code == 400, response.text + assert ( + response.json()["error"]["details"]["reason"] + == "type_mismatch.context.auth_token" + ) + assert _context_row_count(task_id) == 0 + + +def test_values_endpoint_rejects_malformed_key_name(e2e_db: None) -> None: + """A key name that fails ``validate_runtime_source_key`` (e.g. + containing a dot) is rejected.""" + headers, task_id, server_id = _setup_context_task() + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth.token": "x"}, + } + ] + }, + ) + assert response.status_code == 400, response.text + assert _context_row_count(task_id) == 0 + + +def test_values_endpoint_response_stays_unsatisfied_for_a_malformed_declared_key( + e2e_db: None, +) -> None: + """A successful write does not make the response's top-level + ``satisfied`` true while the connector still declares a key whose name + the per-turn gate rejects -- even an optional one. The write itself + succeeds and the key it wrote reads satisfied, but the malformed key + stays listed and unsatisfied and holds the top-level flag at false, so + a caller cannot read a 200 here as "this task can now run". + """ + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + server = _mcp_server_with_context_schema( + db, + user, + name="values-malformed-declaration-server", + context_schema={ + "auth_token": {"type": "string", "required": True}, + "bad.key": {"type": "string", "required": False}, + }, + url="https://example.com/values-malformed-declaration/mcp", + ) + agent = _create_agent( + db, + user, + name="Values Malformed Declaration Agent", + tool_categories=["mcp"], + ) + db.commit() + db.refresh(agent) + db.refresh(server) + agent_id = int(agent.id) + server_id = int(server.id) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={ + "title": "values malformed declaration task", + "description": "d", + "agent_id": agent_id, + }, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": "x"}, + } + ] + }, + ) + assert response.status_code == 200, response.text + payload = response.json() + inputs_by_key = { + item["key"]: item + for connector in payload["connectors"] + for item in connector["inputs"] + } + assert set(inputs_by_key) == {"auth_token", "bad.key"} + assert inputs_by_key["auth_token"]["satisfied"] is True + assert inputs_by_key["bad.key"]["required"] is False + assert inputs_by_key["bad.key"]["satisfied"] is False + assert payload["satisfied"] is False + + +def test_values_endpoint_rejects_oversized_single_key_and_batch(e2e_db: None) -> None: + """The two size-cap cases: a single key over 64KB, and a batch over + 256KB in aggregate even though no single key trips the per-key cap. + Neither writes anything, and which key was oversized never + appears in the response ``reason`` -- only the fixed literal + ``payload_too_large``.""" + headers, task_id, server_id = _setup_context_task() + oversized_value = "x" * (64 * 1024 + 1) + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": oversized_value}, + } + ] + }, + ) + assert response.status_code == 400, response.text + assert response.json()["error"]["details"]["reason"] == "payload_too_large" + assert _context_row_count(task_id) == 0 + + # A batch that trips the aggregate cap without any single key tripping + # the per-key cap: one connector declaring five keys, each just under + # 64KB (5 * ~64KB > 256KB in total). + headers2 = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + key_names = [f"key_{i}" for i in range(5)] + server2 = _mcp_server_with_context_schema( + db, + user, + name="a4-batch-server", + context_schema={ + key: {"type": "string", "required": False} for key in key_names + }, + ) + agent2 = _create_agent(db, user, name="a4-batch-agent", tool_categories=["mcp"]) + db.commit() + db.refresh(agent2) + db.refresh(server2) + agent2_id = int(agent2.id) + server2_id = int(server2.id) + finally: + db.close() + + create_response2 = client.post( + "/api/chat/task/create", + headers=headers2, + json={"title": "a4-batch-task", "description": "d", "agent_id": agent2_id}, + ) + assert create_response2.status_code == 200, create_response2.text + task_id2 = int(create_response2.json()["task_id"]) + + just_under_per_key_cap = "x" * (64 * 1024 - 200) + response2 = client.post( + _values_url(task_id2), + headers=headers2, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server2_id, + }, + "context": {key: just_under_per_key_cap for key in key_names}, + }, + ] + }, + ) + assert response2.status_code == 400, response2.text + assert response2.json()["error"]["details"]["reason"] == "payload_too_large" + assert _context_row_count(task_id2) == 0 + + +def test_values_endpoint_rejects_duplicate_ref_in_batch(e2e_db: None) -> None: + """The same connector ref submitted twice in one batch is rejected as a + duplicate.""" + headers, task_id, server_id = _setup_context_task() + item = { + "connector_ref": {"connector_type": "mcp", "connector_id": server_id}, + "context": {"auth_token": "x"}, + } + response = client.post( + _values_url(task_id), headers=headers, json={"items": [item, item]} + ) + assert response.status_code == 400, response.text + assert response.json()["error"]["details"]["reason"] == "duplicate_ref" + assert _context_row_count(task_id) == 0 + + +def _stored_context(task_id: int, server_id: int) -> dict[str, Any] | None: + db = _db_session() + try: + row = ( + db.query(TaskConnectorRuntimeContext) + .filter( + TaskConnectorRuntimeContext.task_id == task_id, + TaskConnectorRuntimeContext.connector_type == "mcp", + TaskConnectorRuntimeContext.connector_id == server_id, + ) + .one_or_none() + ) + return dict(row.context) if row is not None else None + finally: + db.close() + + +def test_values_endpoint_merges_keys_and_never_replaces_a_stored_value( + e2e_db: None, +) -> None: + """Five requests against the same connector row, in sequence -- a + stored value is never replaced, a new key can be added, and the + response reflects the write it just made (not a stale pre-write read). + """ + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + server = _mcp_server_with_context_schema( + db, + user, + name="i12-server", + context_schema={ + "a": {"type": "string", "required": False}, + "b": {"type": "string", "required": False}, + "c": {"type": "string", "required": False}, + }, + ) + agent = _create_agent(db, user, name="i12-agent", tool_categories=["mcp"]) + db.commit() + db.refresh(agent) + db.refresh(server) + agent_id = int(agent.id) + server_id = int(server.id) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={"title": "i12 task", "description": "d", "agent_id": agent_id}, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + ref = {"connector_type": "mcp", "connector_id": server_id} + + def _post(context: dict[str, Any]) -> Any: + return client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": context}]}, + ) + + # (1) write "a", then resend the same value -- 200, one row, unchanged. + r1 = _post({"a": "1"}) + assert r1.status_code == 200, r1.text + r1b = _post({"a": "1"}) + assert r1b.status_code == 200, r1b.text + assert _stored_context(task_id, server_id) == {"a": "1"} + assert _context_row_count(task_id) == 1 + + # (2) resend "a" with a different value -- 409, unchanged. + r2 = _post({"a": "2"}) + assert r2.status_code == 409, r2.text + assert r2.json()["error"]["details"]["reason"] == "conflict.context.a" + assert _stored_context(task_id, server_id) == {"a": "1"} + + # (3) add "b" -- 200, merged, and this response's own inputs for both + # "a" and "b" already read satisfied=true: this response is assembled + # from the write it just made, not a stale pre-write read. + r3 = _post({"b": "2"}) + assert r3.status_code == 200, r3.text + assert _stored_context(task_id, server_id) == {"a": "1", "b": "2"} + r3_body = r3.json() + r3_inputs = {item["key"]: item for item in r3_body["connectors"][0]["inputs"]} + assert r3_inputs["a"]["satisfied"] is True + assert r3_inputs["b"]["satisfied"] is True + # The wire fields this phase emits as constants, pinned on the 200 body + # of the third producer of this model: `expired` is false for every + # input and `secrets_expires_at` is null, because no expiry-bearing + # value can be stored yet, and every input a `context`-only connector + # declares lands in the `context` section. + assert all(item["expired"] is False for item in r3_inputs.values()) + assert all(item["section"] == "context" for item in r3_inputs.values()) + assert r3_body["secrets_expires_at"] is None + + # (4) add "c" alongside the already-stored "a" -- 200, all three present. + r4 = _post({"a": "1", "c": "3"}) + assert r4.status_code == 200, r4.text + assert _stored_context(task_id, server_id) == {"a": "1", "b": "2", "c": "3"} + + # (5) conflict on "a" again, this time alongside a currently-valid "c" + # -- the whole request still fails, not a byte written. + r5 = _post({"a": "2", "c": "3"}) + assert r5.status_code == 409, r5.text + assert _stored_context(task_id, server_id) == {"a": "1", "b": "2", "c": "3"} + + +@pytest.mark.parametrize("override_key", ["if_absent", "overwrite", "force"]) +def test_values_endpoint_rejects_any_override_switch( + e2e_db: None, override_key: str +) -> None: + """No override switch is accepted at the request-shape level -- + ``extra="forbid"`` turns any of them into 422 before the merge logic + ever runs.""" + headers, task_id, server_id = _setup_context_task(required=False) + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": "x"}, + override_key: True, + } + ] + }, + ) + assert response.status_code == 422, response.text + assert _context_row_count(task_id) == 0 + + +@pytest.mark.parametrize( + ("section", "payload"), + [("secrets", {"api_key": "x"}), ("auth_selector", {"choice": "a"})], +) +def test_values_endpoint_rejects_secret_bearing_sections( + e2e_db: None, section: str, payload: dict[str, str] +) -> None: + """``secrets`` and ``auth_selector`` are deliberately absent from the + request shape at this phase, so an item carrying either section is + 422 before the merge logic runs and nothing is stored.""" + headers, task_id, server_id = _setup_context_task(required=False) + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": "x"}, + section: payload, + } + ] + }, + ) + assert response.status_code == 422, response.text + assert _context_row_count(task_id) == 0 + + +def test_values_endpoint_conflict_details_carry_reason_and_connector_ref( + e2e_db: None, +) -> None: + """A 409's ``details`` carries both ``reason`` and ``connector_ref`` -- + the latter is a documented part of the contract, not an accidental leak + of ``ConnectorRuntimeError``'s internals.""" + headers, task_id, server_id = _setup_context_task(required=False) + ref = {"connector_type": "mcp", "connector_id": server_id} + first = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"auth_token": "1"}}]}, + ) + assert first.status_code == 200, first.text + second = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"auth_token": "2"}}]}, + ) + assert second.status_code == 409, second.text + details = second.json()["error"]["details"] + assert details["reason"] == "conflict.context.auth_token" + assert details["connector_ref"] == ref + + +def test_values_endpoint_first_item_atomic_when_second_item_fails_validation( + e2e_db: None, +) -> None: + """Multi-item atomicity on the 400 path -- the first item alone would + have succeeded, but the whole batch fails together with the second, and + neither lands.""" + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + good_server = _mcp_server_with_context_schema( + db, + user, + name="i15-good-server", + context_schema={"auth_token": {"type": "string", "required": False}}, + ) + bad_server = _mcp_server_with_context_schema( + db, + user, + name="i15-bad-server", + context_schema={"auth_token": {"type": "string", "required": False}}, + ) + agent = _create_agent(db, user, name="i15-agent", tool_categories=["mcp"]) + db.commit() + db.refresh(agent) + db.refresh(good_server) + db.refresh(bad_server) + agent_id = int(agent.id) + good_id = int(good_server.id) + bad_id = int(bad_server.id) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={"title": "i15 task", "description": "d", "agent_id": agent_id}, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": {"connector_type": "mcp", "connector_id": good_id}, + "context": {"auth_token": "ok"}, + }, + { + "connector_ref": {"connector_type": "mcp", "connector_id": bad_id}, + "context": {"undeclared_key": "x"}, + }, + ] + }, + ) + assert response.status_code == 400, response.text + assert _context_row_count(task_id) == 0 + + +def test_values_endpoint_first_item_atomic_when_second_item_conflicts( + e2e_db: None, +) -> None: + """Multi-item atomicity on the 409 path. Ref A has no row yet; ref B + already has a different value stored under another session. One + request submitting A+B fails together (A gets zero residue too), then + resubmitting A alone (dropping B, the frontend rule after a 409) + succeeds.""" + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + server_a = _mcp_server_with_context_schema( + db, + user, + name="i11ba-server-a", + context_schema={"auth_token": {"type": "string", "required": False}}, + ) + server_b = _mcp_server_with_context_schema( + db, + user, + name="i11ba-server-b", + context_schema={"auth_token": {"type": "string", "required": False}}, + ) + agent = _create_agent(db, user, name="i11ba-agent", tool_categories=["mcp"]) + db.commit() + db.refresh(agent) + db.refresh(server_a) + db.refresh(server_b) + agent_id = int(agent.id) + a_id = int(server_a.id) + b_id = int(server_b.id) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={"title": "i11ba task", "description": "d", "agent_id": agent_id}, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + + ref_a = {"connector_type": "mcp", "connector_id": a_id} + ref_b = {"connector_type": "mcp", "connector_id": b_id} + pre_fill = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [{"connector_ref": ref_b, "context": {"auth_token": "original"}}] + }, + ) + assert pre_fill.status_code == 200, pre_fill.text + + combined = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + {"connector_ref": ref_a, "context": {"auth_token": "new-a"}}, + {"connector_ref": ref_b, "context": {"auth_token": "conflicting"}}, + ] + }, + ) + assert combined.status_code == 409, combined.text + assert ( + combined.json()["error"]["details"]["reason"] == "conflict.context.auth_token" + ) + assert _context_row_count(task_id) == 1 # only B's pre-fill row + assert _stored_context(task_id, a_id) is None + + only_a = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref_a, "context": {"auth_token": "new-a"}}]}, + ) + assert only_a.status_code == 200, only_a.text + assert _stored_context(task_id, a_id) == {"auth_token": "new-a"} + + +def test_values_endpoint_partial_fill_returns_200_and_persists(e2e_db: None) -> None: + """Partial fill is 200, not a rejection -- including when the + connector's declared schema grows a new required key between when the + caller last read the requirements and when it submits.""" + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + server = _mcp_server_with_context_schema( + db, + user, + name="i32b-server", + context_schema={ + "a": {"type": "string", "required": True}, + "b": {"type": "string", "required": True}, + }, + ) + agent = _create_agent(db, user, name="i32b-agent", tool_categories=["mcp"]) + db.commit() + db.refresh(agent) + db.refresh(server) + agent_id = int(agent.id) + server_id = int(server.id) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={"title": "i32b task", "description": "d", "agent_id": agent_id}, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + ref = {"connector_type": "mcp", "connector_id": server_id} + + response = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"a": "1"}}]}, + ) + assert response.status_code == 200, response.text + payload = response.json() + assert payload["satisfied"] is False + inputs = {item["key"]: item for item in payload["connectors"][0]["inputs"]} + assert inputs["a"]["satisfied"] is True + assert inputs["b"]["satisfied"] is False + assert _stored_context(task_id, server_id) == {"a": "1"} + + # Declaration grows a new required key after this read; the caller + # (having computed "only a was missing" from a stale read) submits + # exactly what it originally intended. Still 200; the new key is simply + # still unsatisfied afterward, not a rejection of this request. + db = _db_session() + try: + row = db.query(MCPServer).filter(MCPServer.id == server_id).one() + row.runtime_input_schema = { + "context": { + "a": {"type": "string", "required": True}, + "b": {"type": "string", "required": True}, + "d": {"type": "string", "required": True}, + } + } + db.commit() + finally: + db.close() + + response2 = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"b": "2"}}]}, + ) + assert response2.status_code == 200, response2.text + assert _stored_context(task_id, server_id) == {"a": "1", "b": "2"} + + +def _count_context_table_writes(fn: Any) -> int: + """Run ``fn`` while counting INSERT/UPDATE/DELETE statements issued + against ``task_connector_runtime_contexts`` on the app's engine.""" + from xagent.web.models.database import get_engine + + engine = get_engine() + counted = {"n": 0} + + def _listener( + conn: Any, cursor: Any, statement: str, *_args: Any, **_kwargs: Any + ) -> None: + upper = statement.upper() + if "TASK_CONNECTOR_RUNTIME_CONTEXTS" in upper and ( + "INSERT" in upper or "UPDATE" in upper or "DELETE" in upper + ): + counted["n"] += 1 + + event.listen(engine, "before_cursor_execute", _listener) + try: + fn() + finally: + event.remove(engine, "before_cursor_execute", _listener) + return counted["n"] + + +def test_values_endpoint_same_value_rewrite_issues_no_write_statement( + e2e_db: None, +) -> None: + """Writing the same value twice issues zero INSERT/UPDATE statements on + the second request -- the short circuit at "same content already on the + row" actually skips the write, not just the response change it would + otherwise have caused.""" + headers, task_id, server_id = _setup_context_task(required=False) + ref = {"connector_type": "mcp", "connector_id": server_id} + body = {"items": [{"connector_ref": ref, "context": {"auth_token": "same"}}]} + + first = client.post(_values_url(task_id), headers=headers, json=body) + assert first.status_code == 200, first.text + assert _context_row_count(task_id) == 1 + + second_writes = _count_context_table_writes( + lambda: client.post(_values_url(task_id), headers=headers, json=body) + ) + assert second_writes == 0 + assert _context_row_count(task_id) == 1 + + +@pytest.mark.parametrize( + ("key_type", "context_value"), + [ + ("string", "SENTINEL-CONTEXT-VALUE-DO-NOT-LOG"), + ("object", {"nested": "SENTINEL-CONTEXT-VALUE-DO-NOT-LOG"}), + ], + ids=["top_level_string", "nested_in_object"], +) +def test_values_endpoint_never_logs_submitted_values( + e2e_db: None, + caplog: pytest.LogCaptureFixture, + key_type: str, + context_value: Any, +) -> None: + """A submitted value never reaches a log line, checked at DEBUG across + the whole request, both in the rendered message and in any + lazily-formatted ``%s`` args -- for a plain top-level string value and + for one nested inside an object.""" + import logging + + headers, task_id, server_id = _setup_context_task(required=False, key_type=key_type) + sentinel = "SENTINEL-CONTEXT-VALUE-DO-NOT-LOG" + ref = {"connector_type": "mcp", "connector_id": server_id} + + with caplog.at_level(logging.DEBUG): + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": ref, + "context": {"auth_token": context_value}, + } + ] + }, + ) + assert response.status_code == 200, response.text + for record in caplog.records: + assert sentinel not in record.getMessage() + for arg in record.args or (): + assert sentinel not in str(arg) + # Positive half: a fill is still recorded, just without the value -- the + # key name alone shows up in the one log line the write path does emit. + assert any("auth_token" in record.getMessage() for record in caplog.records) + + +@pytest.mark.parametrize("blank_value", ["", " ", "\t\n"]) +def test_values_endpoint_rejects_a_blank_string_value( + e2e_db: None, blank_value: str +) -> None: + """A required key submitted empty, or with nothing but whitespace, is + rejected before anything is written. Accepting it would mark the key + filled forever -- a stored value is never replaced -- while carrying + nothing a connector can use.""" + headers, task_id, server_id = _setup_context_task() + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": blank_value}, + } + ] + }, + ) + assert response.status_code == 400, response.text + body = response.json() + assert body["error"]["code"] == "invalid_runtime_context" + assert body["error"]["details"]["reason"] == "empty_value.context.auth_token" + assert _context_row_count(task_id) == 0 + + +def test_values_endpoint_rejects_a_blank_value_for_an_optional_key( + e2e_db: None, +) -> None: + """The blank rule ignores ``required``: an optional key's value is just + as permanent once stored, so a blank one is refused there too.""" + headers, task_id, server_id = _setup_context_task(required=False) + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": ""}, + } + ] + }, + ) + assert response.status_code == 400, response.text + assert ( + response.json()["error"]["details"]["reason"] + == "empty_value.context.auth_token" + ) + assert _context_row_count(task_id) == 0 + + +def test_values_endpoint_rejects_an_empty_object_value(e2e_db: None) -> None: + """An object-typed key given ``{}`` passes the type check and is still + refused: an empty object is a value with no content, and ``{}`` is not + "a value" for the purpose of satisfying the key.""" + headers, task_id, server_id = _setup_context_task(key_type="object") + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": {}}, + } + ] + }, + ) + assert response.status_code == 400, response.text + assert ( + response.json()["error"]["details"]["reason"] + == "empty_value.context.auth_token" + ) + assert _context_row_count(task_id) == 0 + + +def test_values_endpoint_accepts_a_real_value_after_a_blank_one_was_refused( + e2e_db: None, +) -> None: + """The refusal writes nothing, so the caller can simply submit again -- + this is the whole point of refusing at the door rather than storing a + blank that could never be replaced.""" + headers, task_id, server_id = _setup_context_task() + ref = {"connector_type": "mcp", "connector_id": server_id} + + refused = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"auth_token": ""}}]}, + ) + assert refused.status_code == 400, refused.text + assert _stored_context(task_id, server_id) is None + + accepted = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [{"connector_ref": ref, "context": {"auth_token": "real-token"}}] + }, + ) + assert accepted.status_code == 200, accepted.text + assert accepted.json()["satisfied"] is True + assert _stored_context(task_id, server_id) == {"auth_token": "real-token"} + + +def test_values_endpoint_treats_a_json_reshaped_value_as_a_conflict( + e2e_db: None, +) -> None: + """``{"a": 1}`` already stored, then ``{"a": 1.0}`` or ``{"a": true}`` + submitted: Python calls those equal, but they store different JSON, so + the immutability rule must fire rather than answering 200 and writing + nothing. Resubmitting the identical value is still the no-op it was, + with no write statement at all.""" + headers, task_id, server_id = _setup_context_task(key_type="object") + ref = {"connector_type": "mcp", "connector_id": server_id} + + def _post(value: Any) -> Any: + return client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"auth_token": value}}]}, + ) + + first = _post({"a": 1}) + assert first.status_code == 200, first.text + assert _stored_context(task_id, server_id) == {"auth_token": {"a": 1}} + + for reshaped in ({"a": 1.0}, {"a": True}): + conflict = _post(reshaped) + assert conflict.status_code == 409, conflict.text + assert ( + conflict.json()["error"]["details"]["reason"] + == "conflict.context.auth_token" + ) + assert _stored_context(task_id, server_id) == {"auth_token": {"a": 1}} + + writes = _count_context_table_writes(lambda: _post({"a": 1})) + assert writes == 0 + assert _stored_context(task_id, server_id) == {"auth_token": {"a": 1}} + + +def test_values_endpoint_reports_satisfied_once_every_required_key_is_filled( + e2e_db: None, +) -> None: + """The endpoint exists to turn this flag true, so one test asserts it + does -- on the write response itself, and again on the task-keyed read + endpoint afterwards.""" + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + server = _mcp_server_with_context_schema( + db, + user, + name="satisfied-server", + context_schema={ + "a": {"type": "string", "required": True}, + "b": {"type": "string", "required": True}, + }, + ) + agent = _create_agent(db, user, name="satisfied-agent", tool_categories=["mcp"]) + db.commit() + db.refresh(agent) + db.refresh(server) + agent_id = int(agent.id) + server_id = int(server.id) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={"title": "satisfied task", "description": "d", "agent_id": agent_id}, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + ref = {"connector_type": "mcp", "connector_id": server_id} + + partial = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"a": "1"}}]}, + ) + assert partial.status_code == 200, partial.text + assert partial.json()["satisfied"] is False + + final = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"b": "2"}}]}, + ) + assert final.status_code == 200, final.text + final_body = final.json() + assert final_body["satisfied"] is True + final_inputs = {item["key"]: item for item in final_body["connectors"][0]["inputs"]} + assert final_inputs["a"]["satisfied"] is True + assert final_inputs["b"]["satisfied"] is True + + report = client.get( + f"/api/chat/task/{task_id}/connector-runtime-requirements", headers=headers + ) + assert report.status_code == 200, report.text + assert report.json()["satisfied"] is True + + +def test_values_endpoint_returns_503_when_the_conditional_update_keeps_missing( + e2e_db: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """Retry exhaustion on the conditional-update branch, asserted at the + HTTP layer: exactly two attempts, then the typed 503 whose ``details`` + carry no reason -- the shape that tells a caller this was a collision + with another writer and is safe to retry. Driven by forcing every + conditional update to match no row, which is what a lost race looks + like from inside this request.""" + task_id, server_id = _setup_concurrency_task(key_names=["a", "b"]) + headers = _setup_admin_headers() + ref = {"connector_type": "mcp", "connector_id": server_id} + + seed = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"a": "1"}}]}, + ) + assert seed.status_code == 200, seed.text + + from xagent.web.services import connector_runtime as connector_runtime_service + + calls = {"n": 0} + + def _always_missing_update( + db: Session, *, row_id: int, context_text_as_read: str, merged: dict[str, Any] + ) -> bool: + calls["n"] += 1 + return False + + monkeypatch.setattr( + connector_runtime_service, "_update_context_row", _always_missing_update + ) + + response = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"b": "2"}}]}, + ) + + assert response.status_code == 503, response.text + assert response.json() == { + "error": { + "code": "connector_runtime_unavailable", + "message": "Connector runtime context is unavailable.", + "details": {}, + } + } + assert calls["n"] == 2 + assert _stored_context(task_id, server_id) == {"a": "1"} + + +def test_values_endpoint_returns_503_when_the_insert_keeps_colliding( + e2e_db: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """The sibling branch, same assertion at the same layer: a snapshot + read that never sees the row already in the table sends this request + down the insert path twice, and the second ``IntegrityError`` ends it + as the same typed 503 with the same empty ``details``.""" + task_id, server_id = _setup_concurrency_task(key_names=["a", "b"]) + headers = _setup_admin_headers() + ref = {"connector_type": "mcp", "connector_id": server_id} + + seed = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"a": "1"}}]}, + ) + assert seed.status_code == 200, seed.text + + from xagent.web.services import connector_runtime as connector_runtime_service + + calls = {"n": 0} + + def _always_stale_read(db: Session, *, task_id: int) -> dict[str, Any]: + calls["n"] += 1 + return {} + + monkeypatch.setattr( + connector_runtime_service, + "_load_task_context_row_snapshots", + _always_stale_read, + ) + + response = client.post( + _values_url(task_id), + headers=headers, + json={"items": [{"connector_ref": ref, "context": {"b": "2"}}]}, + ) + + assert response.status_code == 503, response.text + assert response.json() == { + "error": { + "code": "connector_runtime_unavailable", + "message": "Connector runtime context is unavailable.", + "details": {}, + } + } + assert calls["n"] == 2 + assert _stored_context(task_id, server_id) == {"a": "1"} + + +def test_values_endpoint_rolls_back_when_the_commit_itself_fails( + e2e_db: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """A ``db.commit()`` that raises is rolled back explicitly and the + exception is left to travel: a bare 500, not the typed envelope, and + nothing of the batch in the table. The envelope is reserved for a + ``ConnectorRuntimeError``, whose ``code`` promises the caller a + condition they can act on; an unclassified commit fault is not one. + + ``Session.commit``/``Session.rollback`` are patched only for the + duration of the request, so the setup above and the table read below + both run against the real methods. + """ + headers, task_id, server_id = _setup_context_task(required=False) + ref = {"connector_type": "mcp", "connector_id": server_id} + body = {"items": [{"connector_ref": ref, "context": {"auth_token": "v"}}]} + + real_rollback = Session.rollback + rollbacks = {"n": 0} + + def _failing_commit(self: Session) -> None: + raise RuntimeError("commit-failure-must-not-leak") + + def _counting_rollback(self: Session) -> None: + rollbacks["n"] += 1 + real_rollback(self) + + monkeypatch.setattr(Session, "commit", _failing_commit) + monkeypatch.setattr(Session, "rollback", _counting_rollback) + try: + response = client.post(_values_url(task_id), headers=headers, json=body) + finally: + monkeypatch.undo() + + assert response.status_code == 500, response.text + assert "commit-failure-must-not-leak" not in response.text + assert rollbacks["n"] == 1 + assert _context_row_count(task_id) == 0 + + +def _setup_two_connector_context_task() -> tuple[dict[str, str], int, int, int]: + """One task whose agent selects two MCP connectors, each declaring the + same two optional context keys. Returns (owner headers, task_id, and + the two connector ids in canonical order -- ascending, which is what + ``_sort_connector_refs`` produces for two refs of the same type). + """ + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + schema = { + "a": {"type": "string", "required": False}, + "b": {"type": "string", "required": False}, + } + first = _mcp_server_with_context_schema( + db, user, name="order-server-1", context_schema=schema + ) + second = _mcp_server_with_context_schema( + db, user, name="order-server-2", context_schema=schema + ) + agent = _create_agent(db, user, name="order-agent", tool_categories=["mcp"]) + db.commit() + db.refresh(agent) + db.refresh(first) + db.refresh(second) + agent_id = int(agent.id) + server_ids = sorted([int(first.id), int(second.id)]) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={"title": "order task", "description": "d", "agent_id": agent_id}, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + return headers, task_id, server_ids[0], server_ids[1] + + +def test_values_endpoint_updates_rows_in_canonical_ref_order( + e2e_db: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """A batch naming two connectors in reverse canonical order still + updates their rows in canonical order. Two concurrent requests listing + the same rows in opposite order would otherwise be able to take those + rows' locks in opposite order and deadlock on PostgreSQL, so the + request's own ordering must not reach the write. + """ + headers, task_id, low_id, high_id = _setup_two_connector_context_task() + low_ref = {"connector_type": "mcp", "connector_id": low_id} + high_ref = {"connector_type": "mcp", "connector_id": high_id} + + # Both rows must already exist, so the second request drives every ref + # down the conditional-update branch rather than the insert branch. + seed = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + {"connector_ref": low_ref, "context": {"a": "1"}}, + {"connector_ref": high_ref, "context": {"a": "1"}}, + ] + }, + ) + assert seed.status_code == 200, seed.text + + db = _db_session() + try: + row_id_by_connector = { + int(row.connector_id): int(row.id) + for row in db.query(TaskConnectorRuntimeContext) + .filter(TaskConnectorRuntimeContext.task_id == task_id) + .all() + } + finally: + db.close() + assert set(row_id_by_connector) == {low_id, high_id} + + from xagent.web.services import connector_runtime as connector_runtime_service + + real_update = connector_runtime_service._update_context_row + updated_row_ids: list[int] = [] + + def _recording_update( + db: Session, *, row_id: int, context_text_as_read: str, merged: dict[str, Any] + ) -> bool: + updated_row_ids.append(row_id) + return bool( + real_update( + db, + row_id=row_id, + context_text_as_read=context_text_as_read, + merged=merged, + ) + ) + + monkeypatch.setattr( + connector_runtime_service, "_update_context_row", _recording_update + ) + + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + {"connector_ref": high_ref, "context": {"b": "2"}}, + {"connector_ref": low_ref, "context": {"b": "2"}}, + ] + }, + ) + + assert response.status_code == 200, response.text + assert updated_row_ids == [ + row_id_by_connector[low_id], + row_id_by_connector[high_id], + ] + assert _stored_context(task_id, low_id) == {"a": "1", "b": "2"} + assert _stored_context(task_id, high_id) == {"a": "1", "b": "2"} + + +# --------------------------------------------------------------------------- +# Values-endpoint concurrency: two sessions racing the write. Driven +# below the HTTP layer, directly against +# ``apply_task_connector_runtime_context_values``, because the race needs +# precise control over when each session reads versus when the other +# commits -- an ordinary sequential HTTP call cannot express "A read before +# B committed, then B committed, then A wrote". +# --------------------------------------------------------------------------- + + +def _setup_concurrency_task(*, key_names: list[str]) -> tuple[int, int]: + """One connector declaring the given (all optional) context keys, one + task that selected it. Returns (task_id, server_id).""" + headers = _setup_admin_headers() + db = _db_session() + try: + user = _admin_user(db) + server = _mcp_server_with_context_schema( + db, + user, + name=f"conc-server-{'-'.join(key_names)}-{secrets.token_hex(4)}", + context_schema={ + key: {"type": "string", "required": False} for key in key_names + }, + ) + agent = _create_agent( + db, + user, + name=f"conc-agent-{secrets.token_hex(4)}", + tool_categories=["mcp"], + ) + db.commit() + db.refresh(agent) + db.refresh(server) + agent_id = int(agent.id) + server_id = int(server.id) + finally: + db.close() + + create_response = client.post( + "/api/chat/task/create", + headers=headers, + json={"title": "conc task", "description": "d", "agent_id": agent_id}, + ) + assert create_response.status_code == 200, create_response.text + return int(create_response.json()["task_id"]), server_id + + +def _direct_apply( + session: Session, task_id: int, server_id: int, context: dict[str, Any] +) -> Any: + from xagent.web.models.agent import Agent as AgentModel + from xagent.web.schemas.connector_runtime import ConnectorRuntimeValueItem + from xagent.web.services.connector_runtime import ( + apply_task_connector_runtime_context_values, + ) + + task = session.query(Task).filter(Task.id == task_id).one() + agent = ( + session.query(AgentModel).filter(AgentModel.id == task.agent_id).one() + if task.agent_id is not None + else None + ) + item = ConnectorRuntimeValueItem( + connector_ref={"connector_type": "mcp", "connector_id": server_id}, + context=context, + ) + return apply_task_connector_runtime_context_values( + db=session, task=task, agent=agent, payload_items=[item] + ) + + +@pytest.mark.parametrize( + ("a_context", "b_context", "expect_conflict"), + [ + ({"a": "1"}, {"b": "2"}, False), + ({"a": "1"}, {"a": "1"}, False), + ({"a": "1"}, {"a": "2"}, True), + ], +) +def test_concurrent_first_write_has_one_winner( + e2e_db: None, + monkeypatch: pytest.MonkeyPatch, + a_context: dict[str, str], + b_context: dict[str, str], + expect_conflict: bool, +) -> None: + """Two sessions both find no row for a ref. B writes and commits + first; A -- still holding its stale "no row" read -- then attempts its + own insert, hits the real unique-constraint ``IntegrityError``, and + must recover: roll back, reread (now sees B's committed row), remerge, + and either succeed or conflict depending on whether the two writers' + keys collide. A's session must stay usable after the ``IntegrityError`` + -- proven here by querying through it once the call returns. + """ + task_id, server_id = _setup_concurrency_task(key_names=["a", "b"]) + + from xagent.web.services import connector_runtime as connector_runtime_service + + real_read = connector_runtime_service._load_task_context_row_snapshots + + session_a = get_session_local()() + try: + # B runs and commits first, through the real (unpatched) read -- + # this is genuinely the first write for this ref, no simulation + # needed here. + session_b = get_session_local()() + try: + _direct_apply(session_b, task_id, server_id, b_context) + session_b.commit() + finally: + session_b.close() + + # Now install the stale read for A's *first* attempt only: A is + # simulated as having looked before B committed, i.e. still empty, + # even though B's row already exists in the database by the time + # A's insert actually runs. The second attempt (the retry) uses + # the real reader, which does see B's row. + state = {"calls": 0} + + def _stale_first_read(db: Session, *, task_id: int) -> dict[str, Any]: + state["calls"] += 1 + if state["calls"] == 1: + return {} + return real_read(db, task_id=task_id) + + monkeypatch.setattr( + connector_runtime_service, + "_load_task_context_row_snapshots", + _stale_first_read, + ) + + if expect_conflict: + with pytest.raises(ConnectorRuntimeError) as exc_info: + _direct_apply(session_a, task_id, server_id, a_context) + assert exc_info.value.code == "runtime_context_immutable" + session_a.rollback() + else: + result = _direct_apply(session_a, task_id, server_id, a_context) + merged_keys = { + item.key: item.satisfied + for connector in result.connectors + for item in connector.inputs + } + for key in {**a_context, **b_context}: + assert merged_keys[key] is True + session_a.commit() + + # The session survived the IntegrityError-triggered rollback and + # retry -- prove it is not left in SQLAlchemy's + # ``PendingRollbackError`` state by issuing one more query on it. + assert session_a.query(Task).filter(Task.id == task_id).one() is not None + finally: + session_a.close() + + assert _context_row_count(task_id) == 1 + + +def test_concurrent_first_write_retry_is_bounded( + e2e_db: None, monkeypatch: pytest.MonkeyPatch +) -> None: + """A write that keeps colliding does not retry forever -- exactly two + attempts, then 503.""" + task_id, server_id = _setup_concurrency_task(key_names=["a"]) + + from xagent.web.services import connector_runtime as connector_runtime_service + + # B commits a row for this ref up front, unconditionally. + session_b = get_session_local()() + try: + _direct_apply(session_b, task_id, server_id, {"a": "occupied"}) + session_b.commit() + finally: + session_b.close() + + calls = {"n": 0} + + def _always_stale_read(db: Session, *, task_id: int) -> dict[str, Any]: + calls["n"] += 1 + return {} + + monkeypatch.setattr( + connector_runtime_service, + "_load_task_context_row_snapshots", + _always_stale_read, + ) + + session_a = get_session_local()() + try: + with pytest.raises(ConnectorRuntimeError) as exc_info: + _direct_apply(session_a, task_id, server_id, {"a": "occupied"}) + assert exc_info.value.status_code == 503 + session_a.rollback() + finally: + session_a.close() + + assert calls["n"] == 2 + + +@pytest.mark.parametrize( + ("a_context", "expect_conflict"), + [ + ({"b": "2"}, False), + ({"x": "0"}, False), + ({"x": "9"}, True), + ], +) +def test_concurrent_overwrite_of_an_existing_row( + e2e_db: None, + monkeypatch: pytest.MonkeyPatch, + a_context: dict[str, str], + expect_conflict: bool, +) -> None: + """SQLite half of this compare-and-swap proof; the PostgreSQL half, + which exercises the database's own text rendering of the JSON column, + is a separate change. A row already has + ``{"x": "0"}``. A read it before B committed a change to + the same row; B commits (adding "a"); A's conditional UPDATE, built on + the text it read before B's write, cannot match the row's current text + -- rowcount 0, not an error -- and must retry: reread, remerge, and + either succeed (A's own key doesn't collide with B's) or conflict + (A resubmits "x" with a different value than what's on the row now). + """ + task_id, server_id = _setup_concurrency_task(key_names=["x", "a", "b"]) + ref = {"connector_type": "mcp", "connector_id": server_id} + seed = client.post( + _values_url(task_id), + headers=_setup_admin_headers(), + json={"items": [{"connector_ref": ref, "context": {"x": "0"}}]}, + ) + assert seed.status_code == 200, seed.text + + from xagent.web.services import connector_runtime as connector_runtime_service + + real_read = connector_runtime_service._load_task_context_row_snapshots + + session_a = get_session_local()() + try: + stale_snapshot = real_read(session_a, task_id=task_id) + + session_b = get_session_local()() + try: + _direct_apply(session_b, task_id, server_id, {"a": "1"}) + session_b.commit() + finally: + session_b.close() + + state = {"calls": 0} + update_results: list[bool] = [] + real_update = connector_runtime_service._update_context_row + + def _stale_first_read(db: Session, *, task_id: int) -> dict[str, Any]: + state["calls"] += 1 + if state["calls"] == 1: + return stale_snapshot + return real_read(db, task_id=task_id) + + def _counting_update(*args: Any, **kwargs: Any) -> bool: + result = real_update(*args, **kwargs) + update_results.append(result) + return result + + monkeypatch.setattr( + connector_runtime_service, + "_load_task_context_row_snapshots", + _stale_first_read, + ) + monkeypatch.setattr( + connector_runtime_service, "_update_context_row", _counting_update + ) + + if expect_conflict: + with pytest.raises(ConnectorRuntimeError) as exc_info: + _direct_apply(session_a, task_id, server_id, a_context) + assert exc_info.value.details["reason"] == "conflict.context.x" + session_a.rollback() + else: + _direct_apply(session_a, task_id, server_id, a_context) + session_a.commit() + assert _stored_context(task_id, server_id)["x"] == "0" + for key, value in a_context.items(): + assert _stored_context(task_id, server_id)[key] == value + # The counting format: a genuinely new key (the "b" case) has + # something to write against both the stale and the fresh + # view, so its first (stale-text) UPDATE attempt hits 0 rows + # and only the retry, built on a fresh read, succeeds. + # Resubmitting a value the row already has (the "x" case) + # merges to the *same* content under the stale view too, so + # tier 8's same-content skip fires before any UPDATE is even + # attempted -- zero calls, not one that fails and one that + # succeeds. + if "b" in a_context: + assert update_results == [False, True] + else: + assert update_results == [] + finally: + session_a.close() + + if expect_conflict: + # No byte moved: the row still shows only what existed before A's + # failed request, plus B's unrelated write. + assert _stored_context(task_id, server_id) == {"x": "0", "a": "1"} + + +def _overwrite_context_row_text(task_id: int, server_id: int, raw_text: str) -> None: + """Write ``raw_text`` into a context row's ``context`` column via a bare + SQL UPDATE -- bypassing the JSON column's bind processor entirely, so + the row holds exactly ``raw_text``, byte for byte, whatever its + canonical form would have been. + """ + from sqlalchemy import text as sa_text + + db = _db_session() + try: + row = ( + db.query(TaskConnectorRuntimeContext) + .filter( + TaskConnectorRuntimeContext.task_id == task_id, + TaskConnectorRuntimeContext.connector_type == "mcp", + TaskConnectorRuntimeContext.connector_id == server_id, + ) + .one() + ) + db.execute( + sa_text( + "UPDATE task_connector_runtime_contexts SET context = :t WHERE id = :i" + ), + {"t": raw_text, "i": row.id}, + ) + db.commit() + finally: + db.close() + + +@pytest.mark.parametrize( + "raw_text", + [ + # Non-canonical rendering: different key order plus extra + # whitespace than SQLAlchemy's default `json.dumps` would produce. + '{"b": 1, "a": 2}', + # Non-ASCII, unescaped -- `ensure_ascii=True` (json.dumps's + # default) would have turned this into \\uXXXX escapes. + '{"name": "会议室预订"}', + ], +) +def test_values_endpoint_expected_old_value_is_the_database_rendering( + e2e_db: None, raw_text: str +) -> None: + """SQLite half of this compare-and-swap proof; the PostgreSQL half, + which exercises the database's own text rendering of the JSON column, + is a separate change. A row whose stored + text was never produced by this codebase's own + ``json.dumps`` call (a non-canonical rendering, or unescaped non-ASCII) + must still accept a new key. If the expected-old-value comparison ever + re-serializes the read value in Python instead of using the database's + own ``CAST(context AS TEXT)`` rendering, this row's CAS predicate can + never match -- every write to it fails forever, not just once. + """ + headers, task_id, server_id = _setup_context_task(required=False) + db = _db_session() + try: + db.add( + TaskConnectorRuntimeContext( + task_id=task_id, + connector_type="mcp", + connector_id=server_id, + context={"placeholder": "will be overwritten"}, + ) + ) + db.commit() + finally: + db.close() + _overwrite_context_row_text(task_id, server_id, raw_text) + + response = client.post( + _values_url(task_id), + headers=headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": "new-value"}, + } + ] + }, + ) + assert response.status_code == 200, response.text + assert _stored_context(task_id, server_id)["auth_token"] == "new-value" + + +def test_values_endpoint_accepts_team_shared_connector(e2e_db: None) -> None: + """The values-endpoint half of the team-visibility check: a connector + visible only through the agent's team is writable through the values + endpoint for a non-owning team member -- not a 404, which is what + happens if this + endpoint's own ``agent_team_id`` derivation is ever dropped back to + ``None``.""" + db = _db_session() + try: + owner = _create_user(db, "values-team-owner") + member = _create_user(db, "values-team-member") + db.flush() + member_id = int(member.id) + server = MCPServer( + name="values-team-shared-server", + description="values-team-shared-server description", + managed="external", + transport="streamable_http", + url="https://example.com/mcp", + runtime_input_schema={ + "context": {"auth_token": {"type": "string", "required": False}} + }, + runtime_bindings=[ + { + "source": {"input_type": "context", "key": "auth_token"}, + "target": {"target_type": "mcp_meta", "key": "auth_token"}, + } + ], + ) + db.add(server) + db.flush() + server_id = int(server.id) + agent = Agent( + user_id=owner.id, + name="Values Team Shared Agent", + instructions="i", + execution_mode="balanced", + status=AgentStatus.PUBLISHED, + tool_categories=["mcp"], + team_id=404, + visibility="team", + ) + db.add(agent) + db.commit() + db.refresh(agent) + agent_id = int(agent.id) + member_headers = _auth_headers_for_user(member) + finally: + db.close() + + set_agent_team_scope_hook( + lambda db, user_id: ( + AgentTeamScope(team_id=404, is_team_admin=False) + if user_id == member_id + else None + ) + ) + connector_team_scope.set_connector_team_hooks( + team_visibility=lambda db, *, team_id: ( + {"mcp": {server_id}, "custom_api": set()} + if team_id == 404 + else {"mcp": set(), "custom_api": set()} + ) + ) + try: + create_response = client.post( + "/api/chat/task/create", + headers=member_headers, + json={ + "title": "values team shared task", + "description": "d", + "agent_id": agent_id, + }, + ) + assert create_response.status_code == 200, create_response.text + task_id = int(create_response.json()["task_id"]) + + response = client.post( + _values_url(task_id), + headers=member_headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": server_id, + }, + "context": {"auth_token": "x"}, + } + ] + }, + ) + assert response.status_code == 200, response.text + finally: + set_agent_team_scope_hook(None) + connector_team_scope.set_connector_team_hooks() + + +def test_values_endpoint_uses_runtime_agent_resolution_for_team_scope( + e2e_db: None, +) -> None: + """The values-endpoint half of the runtime agent resolution: a task + whose agent is a workforce-generated manager agent, but for which no + matching ``WorkforceRun`` exists, gets its connector scope from + ``_load_agent_for_task_runtime`` returning ``None`` -- the same outcome + a turn would get -- so a connector reachable only through the agent's + team is neither writable here nor listed in the response, while a + personally linked connector is both. Resolving the agent from + ``task.agent_id`` directly would open this endpoint to a connector the + turn itself will never resolve, and make its report disagree with the + task-keyed read endpoint's for the same task. + + The task is built directly for the same reason the read endpoint's + counterpart is: ``POST /task/create`` 404s for a workforce-generated + manager agent, so a task with this agent can only exist as a row the + workforce side creates. + """ + db = _db_session() + try: + caller = _create_user(db, "wf-values-owner") + db.flush() + personal = _mcp_server_with_context_schema( + db, + caller, + name="wf-values-personal-server", + context_schema={"auth_token": {"type": "string", "required": True}}, + url="https://example.com/wf-values-personal/mcp", + ) + team_shared = MCPServer( + name="wf-values-team-shared-server", + description="wf-values-team-shared-server description", + managed="external", + transport="streamable_http", + url="https://example.com/wf-values-team/mcp", + runtime_input_schema={ + "context": {"auth_token": {"type": "string", "required": True}} + }, + runtime_bindings=[ + { + "source": {"input_type": "context", "key": "auth_token"}, + "target": {"target_type": "mcp_meta", "key": "auth_token"}, + } + ], + ) + db.add(team_shared) + db.flush() + # No UserMCPServer link for `caller` at all -- reachable only + # through the team hook below. + personal_id = int(personal.id) + team_shared_id = int(team_shared.id) + agent = Agent( + user_id=caller.id, + name="WF Values Manager Agent", + instructions="i", + execution_mode="balanced", + status=AgentStatus.PUBLISHED, + tool_categories=["mcp"], + team_id=505, + origin=AgentOrigin.WORKFORCE_GENERATED_MANAGER.value, + ) + db.add(agent) + db.flush() + agent_id = int(agent.id) + ordered_ids = sorted([personal_id, team_shared_id]) + task = Task( + user_id=caller.id, + agent_id=agent_id, + title="workforce manager values task", + description="d", + status=TaskStatus.PENDING, + connector_runtime_selected_refs=[ + {"connector_type": "mcp", "connector_id": ref_id} + for ref_id in ordered_ids + ], + ) + db.add(task) + db.commit() + db.refresh(task) + task_id = int(task.id) + caller_headers = _auth_headers_for_user(caller) + finally: + db.close() + + def _hook(db: Session, *, team_id: int) -> dict[str, set[int]]: + if team_id == 505: + return {"mcp": {team_shared_id}, "custom_api": set()} + return {"mcp": set(), "custom_api": set()} + + connector_team_scope.set_connector_team_hooks(team_visibility=_hook) + try: + personal_response = client.post( + _values_url(task_id), + headers=caller_headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": personal_id, + }, + "context": {"auth_token": "x"}, + } + ] + }, + ) + team_response = client.post( + _values_url(task_id), + headers=caller_headers, + json={ + "items": [ + { + "connector_ref": { + "connector_type": "mcp", + "connector_id": team_shared_id, + }, + "context": {"auth_token": "x"}, + } + ] + }, + ) + finally: + connector_team_scope.set_connector_team_hooks() + + assert personal_response.status_code == 200, personal_response.text + refs_seen = [ + item["connector_ref"] for item in personal_response.json()["connectors"] + ] + assert refs_seen == [{"connector_type": "mcp", "connector_id": personal_id}] + + # The team-shared connector is outside the scope a turn would resolve, + # so it is not writable here either -- the same uniform "not found" + # the endpoint answers for any ref the caller cannot see. + assert team_response.status_code == 404, team_response.text + assert _context_row_count(task_id) == 1