nihalashetty
ae67bff5a3
Self-hosted platform for building, testing, and shipping LangChain/LangGraph agents. Deep-agent sub-agents on the canvas, a live tracing/observability timeline, auto-provisioned built-in tools with import/export, per-environment tool variables, streamed evaluations, and per-user auth token forwarding.
334 lines
18 KiB
Python
334 lines
18 KiB
Python
"""Tool Sets: service CRUD + membership, agent toolset->tools resolution, and the REST API."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import uuid
|
|
|
|
import httpx
|
|
|
|
from forge.db.base import SessionLocal
|
|
from forge.main import create_app
|
|
from forge.models import Project, Tool, User
|
|
from forge.services.runtime import build_compile_context
|
|
from forge.services.tool_sets import ToolSetService
|
|
from forge.services.tools import ToolService
|
|
|
|
|
|
async def _seed(tenant: str, slug: str) -> tuple[str, str, str]:
|
|
async with SessionLocal() as s:
|
|
proj = Project(tenant_id=tenant, name="TS Proj", slug=slug, config={})
|
|
s.add(proj)
|
|
await s.flush()
|
|
t1 = Tool(tenant_id=tenant, project_id=proj.id, name="alpha", kind="builtin",
|
|
config={"builtin": "calculator", "description": "a"})
|
|
t2 = Tool(tenant_id=tenant, project_id=proj.id, name="beta", kind="builtin",
|
|
config={"builtin": "current_time", "description": "b"})
|
|
s.add_all([t1, t2])
|
|
await s.commit()
|
|
for obj in (proj, t1, t2):
|
|
await s.refresh(obj)
|
|
return proj.id, t1.id, t2.id
|
|
|
|
|
|
async def test_tool_set_service_crud_and_membership():
|
|
tenant = "t_ts_svc"
|
|
pid, t1, t2 = await _seed(tenant, "ts-svc")
|
|
async with SessionLocal() as s:
|
|
ts = await ToolSetService.create(s, tenant, pid, name="Billing Tools", description="billing", tool_ids=[t1, t2])
|
|
assert ts.slug == "billing-tools"
|
|
assert set(await ToolSetService.member_ids(s, tenant, ts.id)) == {t1, t2}
|
|
assert set((await ToolSetService.members_map(s, tenant, pid))[ts.id]) == {t1, t2}
|
|
assert set(await ToolSetService.tool_ids_for_sets(s, tenant, pid, [ts.id])) == {t1, t2}
|
|
|
|
# unknown / cross-project ids are filtered out of membership
|
|
ts2 = await ToolSetService.create(s, tenant, pid, name="X", tool_ids=[t1, "does-not-exist"])
|
|
assert await ToolSetService.member_ids(s, tenant, ts2.id) == [t1]
|
|
|
|
# add / remove membership
|
|
await ToolSetService.remove_member(s, ts, t1)
|
|
assert await ToolSetService.member_ids(s, tenant, ts.id) == [t2]
|
|
await ToolSetService.add_member(s, ts, t1)
|
|
assert set(await ToolSetService.member_ids(s, tenant, ts.id)) == {t1, t2}
|
|
await ToolSetService.add_member(s, ts, t1) # idempotent (no duplicate row)
|
|
assert len(await ToolSetService.member_ids(s, tenant, ts.id)) == 2
|
|
|
|
# rename regenerates a unique slug (collides with ts2's "x")
|
|
ts = await ToolSetService.update(s, ts, name="X")
|
|
assert ts.slug == "x-2"
|
|
|
|
# update can replace membership wholesale
|
|
ts = await ToolSetService.update(s, ts, tool_ids=[t2])
|
|
assert await ToolSetService.member_ids(s, tenant, ts.id) == [t2]
|
|
|
|
# delete removes the set and its membership rows
|
|
set_id = ts.id
|
|
await ToolSetService.delete(s, ts)
|
|
assert await ToolSetService.get(s, tenant, set_id) is None
|
|
assert set_id not in await ToolSetService.members_map(s, tenant, pid)
|
|
|
|
|
|
async def test_tool_deletion_removes_membership():
|
|
tenant = "t_ts_del"
|
|
pid, t1, t2 = await _seed(tenant, "ts-del")
|
|
async with SessionLocal() as s:
|
|
ts = await ToolSetService.create(s, tenant, pid, name="S", tool_ids=[t1, t2])
|
|
tool = await ToolService.get(s, tenant, t1)
|
|
await ToolService.delete(s, tool) # deleting a tool must drop its membership rows
|
|
assert await ToolSetService.member_ids(s, tenant, ts.id) == [t2]
|
|
|
|
|
|
async def test_build_compile_context_resolves_toolset_to_member_tools():
|
|
tenant = "t_ts_ctx"
|
|
pid, t1, t2 = await _seed(tenant, "ts-ctx")
|
|
async with SessionLocal() as s:
|
|
set_id = (await ToolSetService.create(s, tenant, pid, name="Set A", tool_ids=[t1, t2])).id
|
|
async with SessionLocal() as s:
|
|
ctx = await build_compile_context(s, tenant_id=tenant, project_id=pid)
|
|
# membership is loaded onto the compile context
|
|
assert set(ctx.toolset_members.get(set_id, [])) == {t1, t2}
|
|
# an agent granted only the set resolves to the set's member tool ids...
|
|
assert set(ctx.resolve_tool_ids([], [set_id])) == {t1, t2}
|
|
# ...and to the materialized tools (both builtins compiled into the registry)
|
|
assert len(ctx.tools_for(ctx.resolve_tool_ids([], [set_id]))) == 2
|
|
|
|
|
|
async def test_tool_sets_api_end_to_end():
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
# register a real user (mutations require a real principal, not the dev fallback)
|
|
reg = await c.post("/v1/auth/register", json={"email": f"u{uuid.uuid4().hex[:10]}@example.com", "password": "supersecret1"})
|
|
assert reg.status_code == 201, reg.text
|
|
c.headers["Authorization"] = f"Bearer {reg.json()['access_token']}"
|
|
# project + two tools, all via the API (one consistent tenant = the registered workspace)
|
|
pid = (await c.post("/v1/projects", json={"name": "API TS", "slug": "api-ts"})).json()["id"]
|
|
|
|
def _mk(name: str, builtin: str) -> dict:
|
|
return {"name": name, "kind": "builtin", "config": {"builtin": builtin, "description": name}}
|
|
|
|
t1 = (await c.post(f"/v1/projects/{pid}/tools", json=_mk("aa", "calculator"))).json()["id"]
|
|
t2 = (await c.post(f"/v1/projects/{pid}/tools", json=_mk("bb", "current_time"))).json()["id"]
|
|
|
|
# create a set with one member
|
|
r = await c.post(f"/v1/projects/{pid}/tool-sets", json={"name": "Group One", "description": "g1", "tool_ids": [t1]})
|
|
assert r.status_code == 201, r.text
|
|
st = r.json()
|
|
assert st["slug"] == "group-one" and st["tool_ids"] == [t1] and st["description"] == "g1"
|
|
sid = st["id"]
|
|
|
|
# list
|
|
r = await c.get(f"/v1/projects/{pid}/tool-sets")
|
|
assert r.status_code == 200 and any(x["id"] == sid for x in r.json())
|
|
|
|
# add + remove via the membership endpoints
|
|
assert (await c.post(f"/v1/projects/{pid}/tool-sets/{sid}/tools/{t2}")).status_code == 204
|
|
assert set((await c.get(f"/v1/projects/{pid}/tool-sets/{sid}")).json()["tool_ids"]) == {t1, t2}
|
|
assert (await c.delete(f"/v1/projects/{pid}/tool-sets/{sid}/tools/{t1}")).status_code == 204
|
|
|
|
# patch: rename + replace membership
|
|
r = await c.patch(f"/v1/projects/{pid}/tool-sets/{sid}", json={"name": "Renamed", "tool_ids": [t1, t2]})
|
|
assert r.json()["slug"] == "renamed" and set(r.json()["tool_ids"]) == {t1, t2}
|
|
|
|
# delete
|
|
assert (await c.delete(f"/v1/projects/{pid}/tool-sets/{sid}")).status_code == 204
|
|
assert (await c.get(f"/v1/projects/{pid}/tool-sets/{sid}")).status_code == 404
|
|
|
|
|
|
async def _seed_mcp_project(tenant: str, slug: str) -> tuple[str, str, str]:
|
|
"""Project + two builtin tools; returns (project_id, calc_tool_id, clock_tool_id)."""
|
|
async with SessionLocal() as s:
|
|
proj = Project(tenant_id=tenant, name="MCP TS", slug=slug, config={})
|
|
s.add(proj)
|
|
await s.flush()
|
|
ta = Tool(tenant_id=tenant, project_id=proj.id, name="calc", kind="builtin",
|
|
config={"builtin": "calculator", "description": "c"})
|
|
tb = Tool(tenant_id=tenant, project_id=proj.id, name="clock", kind="builtin",
|
|
config={"builtin": "current_time", "description": "t"})
|
|
s.add_all([ta, tb])
|
|
await s.commit()
|
|
for obj in (proj, ta, tb):
|
|
await s.refresh(obj)
|
|
return proj.id, ta.id, tb.id
|
|
|
|
|
|
async def _list_names(c: httpx.AsyncClient, path: str) -> set[str]:
|
|
r = await c.post(path, json={"jsonrpc": "2.0", "id": 1, "method": "tools/list"})
|
|
return {t["name"] for t in r.json()["result"]["tools"]}
|
|
|
|
|
|
async def test_mcp_toolset_scoped_exposure():
|
|
tenant = "t_mcp_ts"
|
|
pid, a_id, b_id = await _seed_mcp_project(tenant, "mcp-ts")
|
|
async with SessionLocal() as s:
|
|
await ToolSetService.create(s, tenant, pid, name="Set A", tool_ids=[a_id])
|
|
set_b = await ToolSetService.create(s, tenant, pid, name="Set B", tool_ids=[b_id])
|
|
b_slug, b_set_id = set_b.slug, set_b.id
|
|
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
# base endpoint = flat union of every EXPOSED set's enabled tools
|
|
assert await _list_names(c, f"/v1/mcp/{pid}") == {"calc", "clock"}
|
|
# per-set endpoint = just that set's flat list
|
|
assert await _list_names(c, f"/v1/mcp/{pid}/toolset/{b_slug}") == {"clock"}
|
|
assert await _list_names(c, f"/v1/mcp/{pid}/toolset/nope") == set() # unknown slug => empty
|
|
r = await c.post(f"/v1/mcp/{pid}/toolset/{b_slug}", json={"jsonrpc": "2.0", "id": 3, "method": "tools/call",
|
|
"params": {"name": "clock", "arguments": {}}})
|
|
assert r.json()["result"]["isError"] is False
|
|
|
|
# un-expose Set B -> it drops off both the base surface and its own endpoint
|
|
async with SessionLocal() as s:
|
|
await ToolSetService.update(s, await ToolSetService.get(s, tenant, b_set_id), exposed=False)
|
|
assert await _list_names(c, f"/v1/mcp/{pid}") == {"calc"}
|
|
assert await _list_names(c, f"/v1/mcp/{pid}/toolset/{b_slug}") == set()
|
|
r = await c.post(f"/v1/mcp/{pid}", json={"jsonrpc": "2.0", "id": 2, "method": "tools/call",
|
|
"params": {"name": "clock", "arguments": {}}})
|
|
assert "not exposed" in r.json()["error"]["message"]
|
|
|
|
|
|
async def test_mcp_tool_level_exclusion():
|
|
"""Everything in an exposed set is published by default; an operator can untick individual
|
|
tools via project.config.mcp_excluded_tools."""
|
|
tenant = "t_mcp_excl"
|
|
pid, a_id, b_id = await _seed_mcp_project(tenant, "mcp-excl")
|
|
async with SessionLocal() as s:
|
|
await ToolSetService.create(s, tenant, pid, name="Set", tool_ids=[a_id, b_id])
|
|
proj = await s.get(Project, pid)
|
|
proj.config = {"mcp_excluded_tools": [b_id]} # untick 'clock'
|
|
await s.commit()
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
assert await _list_names(c, f"/v1/mcp/{pid}") == {"calc"}
|
|
|
|
|
|
async def test_mcp_no_toolsets_exposes_nothing():
|
|
tenant = "t_mcp_none"
|
|
pid, a_id, _b = await _seed_mcp_project(tenant, "mcp-none")
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
# No tool sets => nothing published (there are no loose / "direct" tools).
|
|
assert await _list_names(c, f"/v1/mcp/{pid}") == set()
|
|
# once a tool is placed in an exposed set, it appears
|
|
async with SessionLocal() as s:
|
|
await ToolSetService.create(s, tenant, pid, name="General", tool_ids=[a_id])
|
|
assert await _list_names(c, f"/v1/mcp/{pid}") == {"calc"}
|
|
|
|
|
|
async def test_mcp_session_token_authorizes_as_end_user():
|
|
"""A project-scoped Forge session token authenticates an MCP caller AS its end_user
|
|
(the portable per-user identity channel), alongside the shared project key."""
|
|
from forge.security import create_session_token
|
|
|
|
tenant = "t_mcp_sess"
|
|
async with SessionLocal() as s:
|
|
proj = Project(tenant_id=tenant, name="Sess", slug="mcp-sess", config={"mcp_api_key": "shared-key"})
|
|
s.add(proj)
|
|
await s.flush()
|
|
s.add(Tool(tenant_id=tenant, project_id=proj.id, name="calc", kind="builtin",
|
|
config={"builtin": "calculator", "description": "c"}))
|
|
await s.commit()
|
|
await s.refresh(proj)
|
|
pid = proj.id
|
|
|
|
good = create_session_token(tenant_id=tenant, project_id=pid, end_user={"id": "u1", "entitlements": ["billing"]})
|
|
wrong_project = create_session_token(tenant_id=tenant, project_id="another", end_user={"id": "u2"})
|
|
body = {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
# shared key -> authorized (no per-user identity)
|
|
assert (await c.post(f"/v1/mcp/{pid}", headers={"Authorization": "Bearer shared-key"}, json=body)).status_code == 200
|
|
# project-scoped session token -> authorized as that end user
|
|
assert (await c.post(f"/v1/mcp/{pid}", headers={"Authorization": f"Bearer {good}"}, json=body)).status_code == 200
|
|
# session token scoped to a different project -> rejected
|
|
assert (await c.post(f"/v1/mcp/{pid}", headers={"Authorization": f"Bearer {wrong_project}"}, json=body)).status_code == 401
|
|
# no credential -> rejected (a key is configured)
|
|
assert (await c.post(f"/v1/mcp/{pid}", json=body)).status_code == 401
|
|
|
|
|
|
async def test_mcp_personal_access_token_authorizes_as_user():
|
|
"""A per-user Personal Access Token (forge_pat_) authenticates an MCP client as that user."""
|
|
from forge.services.apikeys import ApiKeyService
|
|
|
|
tenant = "t_mcp_pat"
|
|
async with SessionLocal() as s:
|
|
proj = Project(tenant_id=tenant, name="PAT", slug="mcp-pat", config={"mcp_api_key": "shared-key"})
|
|
s.add(proj)
|
|
await s.flush()
|
|
s.add(Tool(tenant_id=tenant, project_id=proj.id, name="calc", kind="builtin",
|
|
config={"builtin": "calculator", "description": "c"}))
|
|
user = User(tenant_id=tenant, email="pat-user@example.com", role="editor", status="active")
|
|
s.add(user)
|
|
await s.commit()
|
|
await s.refresh(proj)
|
|
await s.refresh(user)
|
|
pid = proj.id
|
|
_k1, pat = await ApiKeyService.create_personal(s, tenant_id=tenant, user_id=user.id, name="t", project_id=pid)
|
|
key_id = _k1.id
|
|
_k2, pat_other = await ApiKeyService.create_personal(s, tenant_id=tenant, user_id=user.id, name="t2", project_id="another-project")
|
|
|
|
body = {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
# a project-scoped PAT authorizes
|
|
assert (await c.post(f"/v1/mcp/{pid}", headers={"Authorization": f"Bearer {pat}"}, json=body)).status_code == 200
|
|
# a PAT scoped to a different project is rejected here
|
|
assert (await c.post(f"/v1/mcp/{pid}", headers={"Authorization": f"Bearer {pat_other}"}, json=body)).status_code == 401
|
|
# once revoked, the PAT no longer authorizes
|
|
async with SessionLocal() as s:
|
|
await ApiKeyService.revoke_personal(s, tenant_id=tenant, user_id=user.id, key_id=key_id)
|
|
assert (await c.post(f"/v1/mcp/{pid}", headers={"Authorization": f"Bearer {pat}"}, json=body)).status_code == 401
|
|
|
|
|
|
async def test_mcp_token_api_crud():
|
|
"""The user-facing PAT endpoints mint / list / revoke tokens, and a minted token works on MCP."""
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
reg = await c.post("/v1/auth/register", json={"email": f"u{uuid.uuid4().hex[:10]}@example.com", "password": "supersecret1"})
|
|
assert reg.status_code == 201, reg.text
|
|
c.headers["Authorization"] = f"Bearer {reg.json()['access_token']}"
|
|
pid = (await c.post("/v1/projects", json={"name": "PAT API", "slug": "pat-api"})).json()["id"]
|
|
# lock the MCP surface behind a key so credential checks are meaningful
|
|
await c.patch(f"/v1/projects/{pid}", json={"config": {"mcp_api_key": "k"}})
|
|
|
|
r = await c.post(f"/v1/projects/{pid}/mcp-tokens", json={"name": "my token"})
|
|
assert r.status_code == 201, r.text
|
|
tok = r.json()
|
|
assert tok["token"].startswith("forge_pat_") and tok["status"] == "active"
|
|
|
|
body = {"jsonrpc": "2.0", "id": 1, "method": "tools/list"}
|
|
# the freshly minted PAT authenticates against the MCP endpoint (auth is enforced by the key)
|
|
assert (await c.post(f"/v1/mcp/{pid}", json=body)).status_code == 401
|
|
pat_headers = {"Authorization": f"Bearer {tok['token']}"}
|
|
assert (await c.post(f"/v1/mcp/{pid}", headers=pat_headers, json=body)).status_code == 200
|
|
|
|
# listed without the plaintext, then revoked
|
|
lst = (await c.get(f"/v1/projects/{pid}/mcp-tokens")).json()
|
|
assert any(t["id"] == tok["id"] and t.get("token") is None for t in lst)
|
|
assert (await c.delete(f"/v1/projects/{pid}/mcp-tokens/{tok['id']}")).status_code == 204
|
|
assert (await c.post(f"/v1/mcp/{pid}", headers=pat_headers, json=body)).status_code == 401
|
|
|
|
|
|
async def test_connector_role_is_mcp_only():
|
|
"""A 'connector' user can manage their own MCP tokens but cannot mutate project resources."""
|
|
from forge.security import create_access_token
|
|
|
|
tenant = "t_conn_role"
|
|
async with SessionLocal() as s:
|
|
u = User(tenant_id=tenant, email="connector@example.com", role="connector", status="active")
|
|
proj = Project(tenant_id=tenant, name="Conn", slug="conn-p", config={})
|
|
s.add_all([u, proj])
|
|
await s.commit()
|
|
await s.refresh(u)
|
|
await s.refresh(proj)
|
|
uid, pid = u.id, proj.id
|
|
|
|
token = create_access_token(user_id=uid, tenant_id=tenant, role="connector")
|
|
app = create_app()
|
|
async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as c:
|
|
c.headers["Authorization"] = f"Bearer {token}"
|
|
# cannot create/mutate project resources (needs editor+)
|
|
assert (await c.post("/v1/projects", json={"name": "X", "slug": "x-conn"})).status_code == 403
|
|
assert (await c.post(f"/v1/projects/{pid}/tool-sets", json={"name": "S"})).status_code == 403
|
|
# but can mint their own MCP personal access token
|
|
r = await c.post(f"/v1/projects/{pid}/mcp-tokens", json={"name": "my token"})
|
|
assert r.status_code == 201, r.text
|
|
assert r.json()["token"].startswith("forge_pat_")
|