mirror of
https://github.com/tiennm99/DocsGPT.git
synced 2026-10-15 04:17:56 +00:00
Keep outside callers' write limits in scheduled and webhook runs
A scheduled or webhook run acts as the agent's owner with no one to approve, so a public-link user or API-key caller could have the agent schedule a write and have it run on the owner's accounts. Runs now keep the caller's rules: a schedule set by someone who reaches the agent only by its public link runs as a public-link caller, one set through the API (recorded as created_via 'api', migration 0042) and every webhook run as an external caller, each with the agent's API write allowlist.
This commit is contained in:
1 parent
bfa67d3138
commit
daf6af57ae
13 files changed
+294
-4
No files matched your search
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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'));"
|
||||
)
|
||||
@@ -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
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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"]
|
||||
@@ -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}
|
||||
@@ -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"
|
||||
@@ -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")
|
||||
@@ -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
|
||||
|
||||
Reference in new issue
Block a user