diff --git a/docs/content/Guides/Connectors.mdx b/docs/content/Guides/Connectors.mdx index 20c00417..940d550a 100644 --- a/docs/content/Guides/Connectors.mdx +++ b/docs/content/Guides/Connectors.mdx @@ -49,9 +49,9 @@ Tools from OAuth MCP servers default to each member's own account. An admin can ### Agents used through an API key -An agent called with its API key (the website widget, the API) runs with its owner's connections, and nobody can approve an action there. It can read through those connections, but it can't take write actions on them unless the owner allows each one under **Access details > Actions anyone with this key can take as you**. Only the owner can change that list; team editors can't. The owner previewing their own agent in DocsGPT is not limited. +An agent called with its API key (the website widget, the API) or fired by a webhook runs with its owner's connections, and nobody can approve an action there. It can read through those connections, but it can't take write actions on them unless the owner allows each one under **Access details > Actions API, widget and public-link users can take as you**. The same goes for tools that hold the owner's own credentials without a connection: an API tool, or an MCP server the owner signed in to. Only the owner can change that list; team editors can't. The owner previewing their own agent in DocsGPT is not limited. -The same list covers people who open the agent from its public link without being on a team it's shared with. They can't approve write actions on the owner's account, so those run only when the owner allows them there. On a tool where each member connects their own account, they approve writes on their own account as usual. +The same list covers people who open the agent from its public link without being on a team it's shared with, and anything they schedule from that chat. They can't approve write actions on the owner's accounts or credentials, so those run only when the owner allows them there. On a tool where each member connects their own account, they approve writes on their own account as usual. ### What answers show diff --git a/docsgpt/agents/headless_runner.py b/docsgpt/agents/headless_runner.py index 4c1098c1..31388977 100644 --- a/docsgpt/agents/headless_runner.py +++ b/docsgpt/agents/headless_runner.py @@ -20,6 +20,7 @@ from docsgpt.api.answer.services.stream_processor import ( get_prompt, ) from docsgpt.core.settings import settings +from docsgpt.guardrails.config import AgentConfig from docsgpt.quotas.service import QuotaExceededError, QuotaService from docsgpt.retriever.dispatcher import build_dispatcher from docsgpt.retriever.retriever_creator import RetrieverCreator @@ -75,6 +76,8 @@ def run_agent_headless( endpoint: str = "headless", chat_history: Optional[List[Dict[str, Any]]] = None, conversation_id: Optional[str] = None, + external_caller: bool = False, + public_link_caller: bool = False, request_id: Optional[str] = None, trace_user_id: Optional[str] = None, ) -> Dict[str, Any]: @@ -87,6 +90,11 @@ def run_agent_headless( ``trace_user_id`` owns the trace when the run belongs to someone other than the agent's owner (a schedule a user set on a shared agent), so the trace is visible wherever that user sees the run; it defaults to the owner. + ``external_caller`` (a webhook, a schedule set through the API) and + ``public_link_caller`` (a schedule a public-link user set) mark a run for + someone who can't approve for the owner: writes on the owner's accounts + and credentials then run only when the agent's API write allowlist has + them. Raises: QuotaExceededError: If the agent owner's usage quota is exhausted. @@ -109,6 +117,8 @@ def run_agent_headless( endpoint=endpoint, chat_history=chat_history, conversation_id=conversation_id, + external_caller=external_caller, + public_link_caller=public_link_caller, ) if outcome.get("error"): status = tracing.STATUS_ERROR @@ -129,6 +139,8 @@ def _run_agent_headless( endpoint: str = "headless", chat_history: Optional[List[Dict[str, Any]]] = None, conversation_id: Optional[str] = None, + external_caller: bool = False, + public_link_caller: bool = False, ) -> Dict[str, Any]: from docsgpt.core.model_utils import ( get_api_key_for_provider, @@ -232,6 +244,9 @@ def _run_agent_headless( agent_id=agent_id, headless=True, tool_allowlist=list(tool_allowlist or []), + external_caller=external_caller, + public_link_caller=public_link_caller, + api_write_allowlist=AgentConfig.parse(agent_config.get("config")).api_write_allowlist, ) if conversation_id: tool_executor.conversation_id = str(conversation_id) diff --git a/docsgpt/agents/tool_executor.py b/docsgpt/agents/tool_executor.py index ee68c7d2..a6f321d7 100644 --- a/docsgpt/agents/tool_executor.py +++ b/docsgpt/agents/tool_executor.py @@ -1908,6 +1908,9 @@ class ToolExecutor: # falls back to ``origin_conversation_id`` as the schedule's # conversation home. tool_config["agent_id"] = str(self.agent_id) if self.agent_id else None + if self.external_caller: + # Its runs act as the owner for an API-key caller. + tool_config["created_via"] = "api" if tool_data["name"] == "mcp_tool": tool_config["query_mode"] = True diff --git a/docsgpt/agents/tools/scheduler.py b/docsgpt/agents/tools/scheduler.py index cc0b64ed..bb30d675 100644 --- a/docsgpt/agents/tools/scheduler.py +++ b/docsgpt/agents/tools/scheduler.py @@ -44,6 +44,9 @@ class SchedulerTool(Tool): self.user_id: Optional[str] = user_id self.agent_id: Optional[str] = cfg.get("agent_id") self.conversation_id: Optional[str] = cfg.get("conversation_id") + # ``api`` when an API-key caller asked (set by the executor, never the + # model): the schedule's runs then keep that caller's write limits. + self.created_via: str = "api" if cfg.get("created_via") == "api" else "chat" def execute_action(self, action_name: str, **kwargs: Any) -> str: """Dispatch on the LLM-supplied action name.""" @@ -191,7 +194,7 @@ class SchedulerTool(Tool): timezone=tz or "UTC", tool_allowlist=allowlist, origin_conversation_id=self.conversation_id, - created_via="chat", + created_via=self.created_via, ) except Exception as exc: logger.exception("schedule_task create failed: %s", exc) diff --git a/docsgpt/alembic/versions/0042_schedule_created_via_api.py b/docsgpt/alembic/versions/0042_schedule_created_via_api.py new file mode 100644 index 00000000..975c0f08 --- /dev/null +++ b/docsgpt/alembic/versions/0042_schedule_created_via_api.py @@ -0,0 +1,40 @@ +"""0042 schedule created_via api — schedules set by an API-key caller. + +A schedule the agent sets from a widget or API chat is stored under the +agent's owner, like the run itself. ``created_via = 'api'`` records that +it came from outside the app, so its runs keep that caller's limits: no +writes on the owner's accounts or credentials unless the owner allowed +them in the agent's API write allowlist. + +Idempotent both ways. Downgrade folds ``api`` back into ``chat``. + +Revision ID: 0042_schedule_created_via_api +Revises: 0041_connection_account_name +""" + +from typing import Sequence, Union + +from alembic import op + + +revision: str = "0042_schedule_created_via_api" +down_revision: Union[str, None] = "0041_connection_account_name" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + op.execute("ALTER TABLE schedules DROP CONSTRAINT IF EXISTS schedules_created_via_chk;") + op.execute( + "ALTER TABLE schedules ADD CONSTRAINT schedules_created_via_chk " + "CHECK (created_via IN ('chat', 'ui', 'api'));" + ) + + +def downgrade() -> None: + op.execute("ALTER TABLE schedules DROP CONSTRAINT IF EXISTS schedules_created_via_chk;") + op.execute("UPDATE schedules SET created_via = 'chat' WHERE created_via = 'api';") + op.execute( + "ALTER TABLE schedules ADD CONSTRAINT schedules_created_via_chk " + "CHECK (created_via IN ('chat', 'ui'));" + ) diff --git a/docsgpt/api/user/scheduler_worker.py b/docsgpt/api/user/scheduler_worker.py index fe581759..0183c633 100644 --- a/docsgpt/api/user/scheduler_worker.py +++ b/docsgpt/api/user/scheduler_worker.py @@ -233,6 +233,33 @@ def _scheduler_still_allowed(schedule: Dict[str, Any], agent_config: Dict[str, A return resolve(conn, "agent", str(agent_id), user_id) is not None +def _schedule_caller_rules(schedule: Dict[str, Any], agent_config: Dict[str, Any]) -> Dict[str, bool]: + """The caller rules a scheduled run keeps, though it acts as the owner. + + A schedule set through the agent's API key is an external caller's; one + a user without a team grant set (they reach the agent by its public + link) is a public-link caller's. Neither can approve for the owner, so + the run writes on the owner's accounts and credentials only as far as + the agent's API write allowlist allows. + + Args: + schedule: The schedule row. + agent_config: The agent row (or the agentless ephemeral config). + + Returns: + ``external_caller`` and ``public_link_caller`` flags for the run. + """ + public = False + user_id = schedule.get("user_id") + agent_id = agent_config.get("id") + if agent_id and user_id and agent_config.get("user_id") != user_id: + from docsgpt.api.user.resource_access import resolve + + with get_engine().connect() as conn: + public = resolve(conn, "agent", str(agent_id), user_id) is None + return {"external_caller": schedule.get("created_via") == "api", "public_link_caller": public} + + def execute_scheduled_run_body(run_id: str, celery_task_id: Optional[str]) -> Dict[str, Any]: """Execute one scheduled run by id; returns a result dict for tracing.""" if not settings.POSTGRES_URI: @@ -324,6 +351,7 @@ def execute_scheduled_run_body(run_id: str, celery_task_id: Optional[str]) -> Di # to the user who scheduled it (not always the agent's owner). request_id=str(run_id), trace_user_id=run.get("user_id") or schedule.get("user_id"), + **_schedule_caller_rules(schedule, agent_config), ) except SoftTimeLimitExceeded: timed_out = True diff --git a/docsgpt/worker.py b/docsgpt/worker.py index 2c3fe9ca..6bd59f4c 100755 --- a/docsgpt/worker.py +++ b/docsgpt/worker.py @@ -2189,6 +2189,9 @@ def agent_webhook_worker(self, agent_id, payload): input_data, tool_allowlist=_webhook_tool_allowlist(agent_config), endpoint="webhook", + # The payload comes from outside: nobody there can approve for + # the owner, so owner-account writes need the API allowlist. + external_caller=True, request_id=getattr(getattr(self, "request", None), "id", None), ) result = { diff --git a/frontend/src/agents/types/schedule.ts b/frontend/src/agents/types/schedule.ts index 02241354..f277d97c 100644 --- a/frontend/src/agents/types/schedule.ts +++ b/frontend/src/agents/types/schedule.ts @@ -37,7 +37,7 @@ export type Schedule = { model_id?: string | null; token_budget?: number | null; origin_conversation_id?: string | null; - created_via: 'chat' | 'ui'; + created_via: 'chat' | 'ui' | 'api'; consecutive_failure_count: number; created_at: string; updated_at: string; diff --git a/tests/agents/test_headless_runner_callers.py b/tests/agents/test_headless_runner_callers.py new file mode 100644 index 00000000..7c67f665 --- /dev/null +++ b/tests/agents/test_headless_runner_callers.py @@ -0,0 +1,53 @@ +"""``run_agent_headless`` holds a run for an outside caller to the write allowlist. + +Webhooks, schedules set through the API, and schedules a public-link user +set all run as the owner, but nobody behind them can approve for the owner. +""" + +from __future__ import annotations + +from unittest.mock import MagicMock, patch + +import pytest + + +def _executor_kwargs(monkeypatch, config_extra=None, **run_kwargs): + from docsgpt.agents import headless_runner as hr + + agent = MagicMock(name="agent") + agent.gen.return_value = iter([{"answer": "ok"}]) + agent.llm.token_usage = {} + retriever = MagicMock(search=MagicMock(return_value=[])) + captured = {} + + def _executor(*_args, **kwargs): + captured.update(kwargs) + return MagicMock(headless_denials=[]) + + monkeypatch.setattr(hr, "get_prompt", lambda _pid: "system prompt") + monkeypatch.setattr(hr.RetrieverCreator, "create_retriever", classmethod(lambda cls, *a, **kw: retriever)) + monkeypatch.setattr(hr, "ToolExecutor", _executor) + monkeypatch.setattr(hr.AgentCreator, "create_agent", classmethod(lambda cls, *a, **kw: agent)) + monkeypatch.setattr(hr.QuotaService, "check", lambda *a, **kw: None) + config = {"user_id": "u1", "id": "agent-1", "default_model_id": "m", **(config_extra or {})} + with patch("docsgpt.core.model_utils.validate_model_id", return_value=True), \ + patch("docsgpt.core.model_utils.get_provider_from_model_id", return_value="openai"), \ + patch("docsgpt.core.model_utils.get_api_key_for_provider", return_value="k"), \ + patch("docsgpt.utils.calculate_doc_token_budget", return_value=1000): + hr.run_agent_headless(config, "do the thing", **run_kwargs) + return captured + + +@pytest.mark.unit +class TestHeadlessCallerRules: + def test_owner_run_has_no_outside_caller(self, monkeypatch): + kwargs = _executor_kwargs(monkeypatch) + assert kwargs.get("external_caller") is False + assert kwargs.get("public_link_caller") is False + + @pytest.mark.parametrize("flag", ["external_caller", "public_link_caller"]) + def test_outside_caller_carries_the_agents_allowlist(self, monkeypatch, flag): + config = {"config": {"api_write_allowlist": ["tool-1:send"]}} + kwargs = _executor_kwargs(monkeypatch, config, **{flag: True}) + assert kwargs[flag] is True + assert kwargs["api_write_allowlist"] == ["tool-1:send"] diff --git a/tests/api/user/test_scheduler_worker.py b/tests/api/user/test_scheduler_worker.py index 71afc851..1f9f3176 100644 --- a/tests/api/user/test_scheduler_worker.py +++ b/tests/api/user/test_scheduler_worker.py @@ -680,3 +680,63 @@ class TestRunAsOwnerAccessRecheck: assert result["status"] == "success" # Runs with the owner's agent row (owner's context). assert headless.call_args.args[0]["user_id"] == "owner-x" + + +_OK = { + "answer": "ok", "tool_calls": [], "sources": [], "thought": "", + "prompt_tokens": 1, "generated_tokens": 1, "denied": [], + "error_type": None, "model_id": "m", +} + + +class TestCallerRules: + """A scheduled run acts as the owner but keeps the rules of whoever set it.""" + + def _run(self, conn, *, owner="owner-x", user="owner-x", shared=False, created_via="ui", grant=False): + import uuid as _uuid + + from docsgpt.storage.db.repositories.team_members import TeamMembersRepository + from docsgpt.storage.db.repositories.team_resource_grants import ( + TeamResourceGrantsRepository, + ) + from docsgpt.storage.db.repositories.teams import TeamsRepository + + agent_id = _make_agent(conn, owner) + conn.execute(text("UPDATE agents SET shared = :s WHERE id = CAST(:id AS uuid)"), + {"s": shared, "id": agent_id}) + if grant: + team = TeamsRepository(conn).create("T", f"t-{_uuid.uuid4().hex[:8]}", owner) + TeamMembersRepository(conn).add_member(str(team["id"]), user) + TeamResourceGrantsRepository(conn).grant( + str(team["id"]), "agent", agent_id, owner, owner, access_level="viewer", + ) + schedule = SchedulesRepository(conn).create( + user_id=user, agent_id=agent_id, trigger_type="once", instruction="hello", + run_at=_now(), next_run_at=_now(), created_via=created_via, + ) + return ScheduleRunsRepository(conn).record_pending(str(schedule["id"]), user, agent_id, _now()) + + def _caller(self, pg_engine, **kwargs): + with pg_engine.begin() as conn: + run = self._run(conn, **kwargs) + with patch("docsgpt.api.user.scheduler_worker.run_agent_headless", return_value=_OK) as headless: + assert execute_scheduled_run_body(str(run["id"]), "celery-c")["status"] == "success" + return {k: headless.call_args.kwargs.get(k) for k in ("external_caller", "public_link_caller")} + + def test_public_link_users_schedule_runs_as_a_public_link_caller(self, pg_engine, patched_engine, stub_events): + assert self._caller(pg_engine, user="stranger", shared=True) == { + "external_caller": False, "public_link_caller": True, + } + + def test_team_members_schedule_runs_as_a_teammate(self, pg_engine, patched_engine, stub_events): + assert self._caller(pg_engine, user="member-y", shared=True, grant=True) == { + "external_caller": False, "public_link_caller": False, + } + + def test_schedule_set_through_the_api_runs_as_an_external_caller(self, pg_engine, patched_engine, stub_events): + assert self._caller(pg_engine, created_via="api") == { + "external_caller": True, "public_link_caller": False, + } + + def test_owners_own_schedule_runs_as_the_owner(self, pg_engine, patched_engine, stub_events): + assert self._caller(pg_engine) == {"external_caller": False, "public_link_caller": False} diff --git a/tests/connectors/test_public_link_callers.py b/tests/connectors/test_public_link_callers.py index f9060df8..7d6bfc85 100644 --- a/tests/connectors/test_public_link_callers.py +++ b/tests/connectors/test_public_link_callers.py @@ -8,6 +8,7 @@ account (member mode) stays theirs to approve, and teammates keep the card. from __future__ import annotations +import json import uuid from contextlib import ExitStack, contextmanager from types import SimpleNamespace @@ -306,3 +307,30 @@ class TestScheduledRunsForOthers: cid = _connection(pg_conn, "alice") with _db(pg_conn): assert _pause(ToolExecutor(user="alice", headless=True), _tool(cid)) is None + + +class TestSchedulesSetThroughTheApi: + def test_scheduler_loaded_for_an_api_caller_records_it(self): + from docsgpt.agents.tool_executor import ToolExecutor + + executor = ToolExecutor(user="alice", agent_id="11111111-1111-1111-1111-111111111111", external_caller=True) + row = {"id": "sched", "name": "scheduler", "config": {}, "actions": []} + with patch("docsgpt.agents.tool_executor.ToolManager") as manager: + executor._get_or_load_tool(row, "sched", "schedule_task") + assert manager.return_value.load_tool.call_args.kwargs["tool_config"]["created_via"] == "api" + + def test_api_schedule_is_stored_as_such(self, pg_conn): + from docsgpt.agents.tools.scheduler import SchedulerTool + from docsgpt.storage.db.repositories.agents import AgentsRepository + from docsgpt.storage.db.repositories.schedules import SchedulesRepository + + agent_id = str(AgentsRepository(pg_conn).create("alice", "A", "published", key=f"k-{uuid.uuid4().hex}")["id"]) + tool = SchedulerTool({"agent_id": agent_id, "created_via": "api"}, user_id="alice") + + @contextmanager + def _yield(): + yield pg_conn + + with patch.multiple("docsgpt.agents.tools.scheduler", db_session=_yield, db_readonly=_yield): + created = json.loads(tool.execute_action("schedule_task", instruction="ping", delay="1h")) + assert SchedulesRepository(pg_conn).get_internal(created["task_id"])["created_via"] == "api" diff --git a/tests/storage/db/test_migration_0042.py b/tests/storage/db/test_migration_0042.py new file mode 100644 index 00000000..95f297c8 --- /dev/null +++ b/tests/storage/db/test_migration_0042.py @@ -0,0 +1,55 @@ +"""Migration round-trip test for 0042_schedule_created_via_api.""" + +from __future__ import annotations + +import os +import subprocess +import sys +from pathlib import Path + +import pytest +from sqlalchemy import text +from sqlalchemy.exc import IntegrityError + + +pytestmark = pytest.mark.integration + +_0041 = "0041_connection_account_name" + + +def _run_alembic(url: str, *args: str) -> None: + ini = Path(__file__).resolve().parents[3] / "docsgpt" / "alembic.ini" + subprocess.check_call( + [sys.executable, "-m", "alembic", "-c", str(ini), *args], + timeout=120, + env={**os.environ, "POSTGRES_URI": url}, + ) + + +def _insert(conn, created_via: str) -> None: + conn.execute( + text( + "INSERT INTO schedules (user_id, trigger_type, instruction, run_at, next_run_at, created_via) " + "VALUES ('u', 'once', 'x', now(), now(), :v)" + ), + {"v": created_via}, + ) + + +class TestMigration0042RoundTrip: + def test_head_accepts_api(self, pg_engine): + with pg_engine.begin() as conn: + _insert(conn, "api") + + def test_downgrade_folds_api_into_chat_then_upgrade_restores(self, pg_engine): + with pg_engine.begin() as conn: + _insert(conn, "api") + url = pg_engine.url.render_as_string(hide_password=False) + _run_alembic(url, "downgrade", _0041) + with pg_engine.connect() as conn: + assert conn.execute(text("SELECT count(*) FROM schedules WHERE created_via = 'api'")).scalar() == 0 + with pytest.raises(IntegrityError), pg_engine.begin() as conn: + _insert(conn, "api") + _run_alembic(url, "upgrade", "head") + with pg_engine.begin() as conn: + _insert(conn, "api") diff --git a/tests/worker/test_agent_workers.py b/tests/worker/test_agent_workers.py index 24c27129..bfda6631 100644 --- a/tests/worker/test_agent_workers.py +++ b/tests/worker/test_agent_workers.py @@ -75,6 +75,8 @@ class TestAgentWebhookWorker: assert captured["input"] == '{"event": "ping"}' # Webhook caller should pass endpoint='webhook'. assert captured["kwargs"].get("endpoint") == "webhook" + # The payload comes from outside: nobody there can approve for the owner. + assert captured["kwargs"].get("external_caller") is True def test_missing_agent_raises( self, pg_conn, patch_worker_db, task_self, monkeypatch