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.
203 lines
8.5 KiB
Python
203 lines
8.5 KiB
Python
"""Secrets, auth resolver (csrf/bearer), REST tool + projection, and builtin /test validation."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import httpx
|
|
from sqlalchemy import select
|
|
|
|
from forge.auth_providers.resolver import AuthResolver
|
|
from forge.db.base import SessionLocal
|
|
from forge.models import AuditLog, AuthProvider
|
|
from forge.secrets.store import SecretStore
|
|
from forge.services.tools import ToolService
|
|
from forge.tools.rest import build_args_schema, execute_rest
|
|
|
|
GET_ORDER = {
|
|
"name": "get_order",
|
|
"description": "Fetch an order.",
|
|
"kind": "rest_api",
|
|
"request": {
|
|
"method": "GET",
|
|
"url_template": "https://api.acme.dev/v2/orders/{order_id}",
|
|
"fields": [
|
|
{"path": "order_id", "type": "string", "in": "path", "required": True, "llm_visible": True},
|
|
{"path": "include", "type": "string", "in": "query", "required": False, "llm_visible": False, "default": "totals"},
|
|
],
|
|
"headers": [{"name": "Accept", "value": "application/json"}],
|
|
},
|
|
"response": {"projection_jmespath": "data.{subtotal: totals.subtotal, total: totals.grand_total, status: status}"},
|
|
}
|
|
|
|
|
|
# --- secrets ---
|
|
async def test_secret_roundtrip_encrypts_and_decrypts():
|
|
store = SecretStore()
|
|
async with SessionLocal() as s:
|
|
await store.write(s, tenant_id="t_sec", project_id="p_sec", name="creds", value={"u": "a", "p": "b"}, kind="generic")
|
|
got = await store.read_ref(tenant_id="t_sec", project_id="p_sec", ref="secret://proj/creds")
|
|
assert got == {"u": "a", "p": "b"}
|
|
async with SessionLocal() as s:
|
|
audit = (
|
|
await s.execute(
|
|
select(AuditLog).where(
|
|
AuditLog.tenant_id == "t_sec",
|
|
AuditLog.project_id == "p_sec",
|
|
AuditLog.action == "secret.read",
|
|
AuditLog.resource_type == "secret",
|
|
AuditLog.resource_id == "creds",
|
|
)
|
|
)
|
|
).scalar_one()
|
|
assert audit.meta == {"scheme": "secret"}
|
|
|
|
|
|
# --- REST tool ---
|
|
def test_args_schema_excludes_non_llm_visible_fields():
|
|
Args = build_args_schema(GET_ORDER)
|
|
assert set(Args.model_fields) == {"order_id"} # `include` is hidden from the model
|
|
|
|
|
|
async def test_rest_execute_projects_payload():
|
|
seen = {}
|
|
|
|
def handler(req: httpx.Request) -> httpx.Response:
|
|
seen["url"] = str(req.url)
|
|
return httpx.Response(200, json={
|
|
"data": {"totals": {"subtotal": 90, "grand_total": 99}, "line_items": [1, 2, 3, 4], "status": "open"},
|
|
})
|
|
|
|
client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
|
|
res = await execute_rest(GET_ORDER, {"order_id": "A-1"}, tenant_id="t", project_id="p", context={}, auth_resolver=None, client=client)
|
|
await client.aclose()
|
|
|
|
assert "/orders/A-1" in seen["url"] and "include=totals" in seen["url"]
|
|
assert res["projected"] == {"subtotal": 90, "total": 99, "status": "open"}
|
|
assert res["raw"]["data"]["line_items"] == [1, 2, 3, 4] # raw keeps everything
|
|
|
|
|
|
# --- auth resolver ---
|
|
async def test_csrf_session_extract_and_inject():
|
|
ap = AuthProvider(
|
|
id="ap1", tenant_id="t", project_id="p", name="orders", kind="csrf_session",
|
|
config={
|
|
"kind": "csrf_session",
|
|
"token_fetch": {"method": "POST", "url": "https://api.acme.dev/auth/login", "body": {}},
|
|
"extract": [
|
|
{"name": "csrf", "from": "header", "header": "X-CSRF-Token"},
|
|
{"name": "session", "from": "cookie", "cookie": "SESSIONID"},
|
|
],
|
|
"inject": [
|
|
{"to": "header", "name": "X-CSRF-Token", "value": "{{extracted.csrf}}"},
|
|
{"to": "cookie", "name": "SESSIONID", "value": "{{extracted.session}}"},
|
|
],
|
|
"cache_ttl_seconds": 1800,
|
|
},
|
|
)
|
|
|
|
def handler(req: httpx.Request) -> httpx.Response:
|
|
return httpx.Response(200, headers=[("X-CSRF-Token", "abc123"), ("set-cookie", "SESSIONID=zzz; Path=/")])
|
|
|
|
client = httpx.AsyncClient(transport=httpx.MockTransport(handler))
|
|
resolved = await AuthResolver().resolve(tenant_id="t", project_id="p", provider_id="ap1", provider=ap, client=client, force=True)
|
|
await client.aclose()
|
|
assert resolved.headers["X-CSRF-Token"] == "abc123"
|
|
assert resolved.cookies["SESSIONID"] == "zzz"
|
|
|
|
|
|
async def test_bearer_static_auth():
|
|
store = SecretStore()
|
|
async with SessionLocal() as s:
|
|
await store.write(s, tenant_id="t", project_id="p", name="tok", value="T0KEN", kind="bearer")
|
|
ap = AuthProvider(id="b1", tenant_id="t", project_id="p", name="b", kind="bearer", config={"kind": "bearer", "token_ref": "secret://proj/tok"})
|
|
resolved = await AuthResolver().resolve(tenant_id="t", project_id="p", provider_id="b1", provider=ap, force=True)
|
|
assert resolved.headers["Authorization"] == "Bearer T0KEN"
|
|
|
|
|
|
# --- builtin tool /test ---
|
|
def test_resolve_model_injects_project_provider_key(monkeypatch):
|
|
"""A per-project key in ctx.provider_credentials is passed to init_chat_model."""
|
|
import langchain.chat_models as cm
|
|
|
|
captured: dict = {}
|
|
|
|
def fake_init(model, **kwargs):
|
|
captured["model"] = model
|
|
captured.update(kwargs)
|
|
return object()
|
|
|
|
monkeypatch.setattr(cm, "init_chat_model", fake_init)
|
|
|
|
from forge.engine.context import CompileContext
|
|
from forge.engine.models import resolve_model
|
|
|
|
ctx = CompileContext(tenant_id="t", project_id="p", provider_credentials={"openai": "sk-proj-test"})
|
|
resolve_model("openai:gpt-5.4-mini", ctx)
|
|
assert captured["model"] == "openai:gpt-5.4-mini"
|
|
assert captured.get("api_key") == "sk-proj-test"
|
|
|
|
# Google uses google_api_key, not api_key.
|
|
captured.clear()
|
|
ctx2 = CompileContext(tenant_id="t", project_id="p", provider_credentials={"google_genai": "g-key"})
|
|
resolve_model("google_genai:gemini-3.5-flash", ctx2)
|
|
assert captured.get("google_api_key") == "g-key"
|
|
|
|
|
|
async def test_calculator_builtin_test_endpoint_logic():
|
|
res = await ToolService.test("t", "p", {"name": "calc", "kind": "builtin", "builtin": "calculator", "description": "d"}, {"expression": "2*(3+4)"})
|
|
assert res["ok"] is True
|
|
assert res["projected"] == "14"
|
|
|
|
|
|
# --- tools wired into an agent ---
|
|
_CALC = {"name": "calculator", "kind": "builtin", "builtin": "calculator", "description": "Evaluate arithmetic."}
|
|
|
|
|
|
async def test_agent_compiles_and_runs_with_bound_tool():
|
|
from langchain_core.messages import HumanMessage
|
|
from langgraph.checkpoint.memory import InMemorySaver
|
|
|
|
from forge.engine.compiler import compile_workflow
|
|
from forge.services.runtime import make_runtime_ctx
|
|
from forge.tools.materialize import materialize_tool
|
|
|
|
ctx = make_runtime_ctx("t", "p")
|
|
ctx.checkpointer = InMemorySaver()
|
|
ctx.tool_registry = {"tool_calc": materialize_tool(_CALC, ctx)}
|
|
wf = {
|
|
"id": "wf_tool", "version": 1,
|
|
"state": {"messages": {"type": "list[message]", "reducer": "add_messages"}},
|
|
"entry_node": "agent",
|
|
"nodes": [
|
|
{"id": "agent", "type": "agent", "config": {"flavor": "agent", "model": "fake:Done.", "tools": ["tool_calc"]}},
|
|
{"id": "end", "type": "end", "config": {}},
|
|
],
|
|
"edges": [{"source": "agent", "target": "end"}],
|
|
}
|
|
graph = compile_workflow(wf, ctx)
|
|
out = await graph.ainvoke({"messages": [HumanMessage(content="hi")]}, {"configurable": {"thread_id": "t1"}})
|
|
assert out["messages"][-1].content == "Done."
|
|
|
|
|
|
async def test_agent_actually_invokes_tool_full_loop():
|
|
"""Scripted fake model emits a tool call → calculator runs → 6*7 = 42 appears."""
|
|
from langchain.agents import create_agent
|
|
from langchain_core.language_models.fake_chat_models import GenericFakeChatModel
|
|
from langchain_core.messages import AIMessage, HumanMessage
|
|
|
|
from forge.services.runtime import make_runtime_ctx
|
|
from forge.tools.materialize import materialize_tool
|
|
|
|
class Fake(GenericFakeChatModel):
|
|
def bind_tools(self, tools=None, **kwargs):
|
|
return self
|
|
|
|
scripted = iter([
|
|
AIMessage(content="", tool_calls=[{"name": "calculator", "args": {"expression": "6*7"}, "id": "c1", "type": "tool_call"}]),
|
|
AIMessage(content="It is 42."),
|
|
])
|
|
tool = materialize_tool(_CALC, make_runtime_ctx("t", "p"))
|
|
agent = create_agent(model=Fake(messages=scripted), tools=[tool], system_prompt="Use the calculator.")
|
|
out = await agent.ainvoke({"messages": [HumanMessage(content="what is 6*7")]})
|
|
tool_msgs = [m for m in out["messages"] if getattr(m, "type", None) == "tool"]
|
|
assert any("42" in str(m.content) for m in tool_msgs), [m.content for m in out["messages"]]
|