515 lines
22 KiB
Python
515 lines
22 KiB
Python
#!/usr/bin/env python3
|
|
"""Manage remote connector accounts served through the tool gateway.
|
|
|
|
``manage_connections`` is the never-deferred surface for connection
|
|
lifecycle:
|
|
|
|
- ``status`` — which connectors exist for this account and whether each is
|
|
connected (read-only).
|
|
- ``connect`` / ``reconnect`` — start (or restart) an authorization flow.
|
|
The gateway returns a connect link, passed through UN-redacted: the model
|
|
shows it to the user, who opens it in a browser. Each connector's
|
|
``instruction`` text is surfaced once per session, not on every call.
|
|
- ``wait`` — block inside the call until the named connectors report
|
|
connected, or the budget runs out. A model has no clock: told to wait it
|
|
says "I'll check back in a minute" and its next action lands immediately,
|
|
so guidance produced a burst of polls rather than a paced one. Waiting
|
|
inside the call cannot be skipped and works the same on every platform.
|
|
|
|
Scope: gateway connectors ONLY. Local MCP servers stay with ``setup_mcp``,
|
|
which still exists and still works. An earlier draft folded ``install`` /
|
|
``enable`` / ``authorize`` in here, but the desktop consent card arrives
|
|
through a per-tool interception branch keyed on the name ``setup_mcp``
|
|
(agent/tool_executor.py, agent/agent_runtime_helpers.py) and
|
|
``registry.dispatch`` never forwards a ``callback``. So the fold could only
|
|
ever return the "use the terminal" fallback while its schema advertised the
|
|
consent flow — a promise with no delivery path.
|
|
|
|
De-authentication is deliberately NOT exposed to the model: disconnecting
|
|
an account is a user decision, made in the portal dashboard.
|
|
|
|
Availability: gated by the portal sign-in the managed tools already use
|
|
(``check_fn``), so signed-out sessions see exactly today's behavior.
|
|
"""
|
|
|
|
import json
|
|
import logging
|
|
import threading
|
|
import time
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
|
|
from tools.registry import registry, tool_error
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_CONNECTOR_ACTIONS = ("status", "connect", "reconnect", "wait")
|
|
|
|
# (session_id, connector) pairs whose `instruction` text has already been
|
|
# shown. Keyed per session, not per process: the gateway multiplexes many
|
|
# sessions through one process, and guidance suppressed for session A must
|
|
# still reach session B. Module-level dict + lock is the house idiom
|
|
# (browser_use `_pending_create_keys` precedent); an unknown session keys
|
|
# on "" and degrades to per-process, never crashes.
|
|
_seen_instructions: set = set()
|
|
_seen_instructions_lock = threading.Lock()
|
|
|
|
# session_id -> {connector slug: monotonic time the connect call addressed
|
|
# it}. Same keying and same idiom as `_seen_instructions` above, for the same
|
|
# reason. Honest scope: this records the moment THIS TOOL handed the model a
|
|
# link (or confirmed the connector already active) — it cannot prove the model
|
|
# relayed the link to the user. Two guards ride on it: a wait for a connector
|
|
# this session never addressed at all is refused (nothing to wait for), and a
|
|
# wait arriving within seconds of the mint — the connect-and-wait-in-one-batch
|
|
# shape, where no message with the links can have reached the user yet — is
|
|
# bounced as a never-error nudge instead of a silent multi-minute block.
|
|
_rendered_links: Dict[str, Dict[str, float]] = {}
|
|
_rendered_links_lock = threading.Lock()
|
|
|
|
# The just-minted window: a wait that begins this soon after its links were
|
|
# minted can only come from the same tool batch (a real turn boundary costs a
|
|
# model round trip). Bounced, not served — the user has not seen the links.
|
|
_LINKS_JUST_MINTED_SECONDS = 2.0
|
|
|
|
# How long a `wait` runs by default, and the bounds it is clamped into. The
|
|
# floor keeps a wait long enough to be worth the round trip; the ceiling keeps
|
|
# one call inside a span a user reads as "a moment" — the model can always ask
|
|
# and wait again. Honest scope: the budget bounds the loop's own decisions
|
|
# (sleeps, and whether another poll starts); a poll that has already begun
|
|
# runs to the transport's own per-request timeout, so a hung gateway can
|
|
# overrun the ceiling by one poll's worth.
|
|
_WAIT_DEFAULT_SECONDS = 120.0
|
|
_WAIT_MIN_SECONDS = 5.0
|
|
_WAIT_MAX_SECONDS = 180.0
|
|
# The gap between polls, and so the notice delay on a connector coming live.
|
|
# Every poll is a live `v1/connectors` read — nothing is cached, because the
|
|
# whole point of the loop is to see a change made outside this process.
|
|
_POLL_GAP_SECONDS = 5.0
|
|
# The budget is counted as it is spent — the waits plus the time each poll
|
|
# actually takes — rather than read off a wall clock. A slow gateway therefore
|
|
# costs polls instead of overrunning the call, and the loop stays testable
|
|
# without a fake clock.
|
|
#
|
|
# Waits are taken in slices so they stay answerable. Nothing outside a tool can
|
|
# end a call that has already started — the executor only checks for an
|
|
# interrupt between tools — so a tool that blocks this long watches the flag
|
|
# itself, and touches the activity heartbeat so the gateway's inactivity
|
|
# timeout does not kill the session underneath it.
|
|
_POLL_WAIT_SLICE_SECONDS = 1.0
|
|
|
|
_WAIT_UNFINISHED_NOTE = (
|
|
"This is NOT an error: the user simply has not finished connecting yet. "
|
|
"ASK THE USER what they want to do — keep waiting (call wait again), "
|
|
"continue without the pending apps, or get fresh connect links (action "
|
|
"'connect'). Do not retry silently and do not treat the pending apps as "
|
|
"broken."
|
|
)
|
|
|
|
|
|
def _connectors_available() -> bool:
|
|
try:
|
|
from tools.tool_gateway.config import connectors_available
|
|
|
|
return connectors_available()
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _default_client():
|
|
from tools.tool_gateway.client import ConnectorClient
|
|
|
|
return ConnectorClient()
|
|
|
|
|
|
def _clamp_timeout(raw: Any) -> tuple:
|
|
"""Return (seconds, note). *note* is None unless the ask was clamped."""
|
|
try:
|
|
asked = _WAIT_DEFAULT_SECONDS if raw is None else float(raw)
|
|
except (TypeError, ValueError):
|
|
return _WAIT_DEFAULT_SECONDS, None
|
|
if asked > _WAIT_MAX_SECONDS:
|
|
return _WAIT_MAX_SECONDS, (
|
|
f"timeout_seconds was capped at {int(_WAIT_MAX_SECONDS)}s "
|
|
f"(asked for {asked:g}s). Call wait again to keep waiting."
|
|
)
|
|
if asked < _WAIT_MIN_SECONDS:
|
|
return _WAIT_MIN_SECONDS, (
|
|
f"timeout_seconds was raised to the {int(_WAIT_MIN_SECONDS)}s "
|
|
f"minimum (asked for {asked:g}s)."
|
|
)
|
|
return asked, None
|
|
|
|
|
|
def _rendered_for(
|
|
session_id: Optional[str], rendered: Dict[str, Dict[str, float]]
|
|
) -> Dict[str, float]:
|
|
with _rendered_links_lock:
|
|
return dict(rendered.get(str(session_id or ""), {}))
|
|
|
|
|
|
def _record_rendered(
|
|
session_id: Optional[str],
|
|
connector: str,
|
|
rendered: Dict[str, Dict[str, float]],
|
|
*,
|
|
never_fresh: bool = False,
|
|
) -> None:
|
|
# Lowercased on the way in: the wait matcher and the input path both
|
|
# lowercase, and a case drift here would turn into a permanent refusal.
|
|
# `never_fresh` records a connector with no link to show (already active):
|
|
# membership holds, but the just-minted bounce can never fire for it.
|
|
stamp = float("-inf") if never_fresh else time.monotonic()
|
|
with _rendered_links_lock:
|
|
rendered.setdefault(str(session_id or ""), {})[connector.lower()] = stamp
|
|
|
|
|
|
def _wait_between_polls(seconds: float, activity_state: Dict[str, Any]) -> bool:
|
|
"""Hold the call open until the next poll; False if the user interrupted."""
|
|
from tools.interrupt import is_interrupted
|
|
|
|
try:
|
|
from tools.environments.base import touch_activity_if_due
|
|
except Exception:
|
|
touch_activity_if_due = None
|
|
|
|
remaining = seconds
|
|
while remaining > 0:
|
|
if is_interrupted():
|
|
return False
|
|
if touch_activity_if_due is not None:
|
|
try:
|
|
touch_activity_if_due(activity_state, "waiting for connections")
|
|
except Exception:
|
|
pass
|
|
this_slice = min(_POLL_WAIT_SLICE_SECONDS, remaining)
|
|
time.sleep(this_slice)
|
|
remaining -= this_slice
|
|
return True
|
|
|
|
|
|
def _wait_for_connections(
|
|
client: Any,
|
|
connectors: List[str],
|
|
*,
|
|
timeout_seconds: float,
|
|
timeout_note: Optional[str],
|
|
) -> str:
|
|
"""Poll the gateway until the named connectors are live, or time runs out.
|
|
|
|
Never reports a wait outcome as an error. A connector the user has not
|
|
finished authorizing is an ordinary, expected state — the model's next move
|
|
is a question to the user, not a repair.
|
|
"""
|
|
wanted = set(connectors)
|
|
activity_state = {"last_touch": time.monotonic(), "start": time.monotonic()}
|
|
spent = 0.0
|
|
live_entries: List[Dict[str, Any]] = []
|
|
pending: List[str] = list(connectors)
|
|
|
|
def result(status: str, note: str) -> str:
|
|
payload: Dict[str, Any] = {
|
|
"status": status,
|
|
"connectors": live_entries,
|
|
"pending": pending,
|
|
"note": note,
|
|
}
|
|
if timeout_note:
|
|
payload["timeout_note"] = timeout_note
|
|
return json.dumps(payload, ensure_ascii=False)
|
|
|
|
consecutive_errors = 0
|
|
|
|
while True:
|
|
poll_started = time.monotonic()
|
|
try:
|
|
items = client.list_connectors()
|
|
except Exception:
|
|
# A transient gateway blip costs one poll, never the whole wait.
|
|
# Three in a row means the gateway is genuinely down mid-wait —
|
|
# still not the model's error: report what the last good poll saw
|
|
# and hand the decision back, same as a timeout.
|
|
spent += time.monotonic() - poll_started
|
|
consecutive_errors += 1
|
|
if consecutive_errors >= 3:
|
|
return result(
|
|
"timeout",
|
|
"The connector gateway stopped answering while waiting; "
|
|
"still not confirmed: " + ", ".join(pending) + ". "
|
|
+ _WAIT_UNFINISHED_NOTE,
|
|
)
|
|
if spent >= timeout_seconds:
|
|
return result(
|
|
"timeout",
|
|
f"Waited about {int(spent)}s; still not connected: "
|
|
f"{', '.join(pending)}. " + _WAIT_UNFINISHED_NOTE,
|
|
)
|
|
gap = min(_POLL_GAP_SECONDS, timeout_seconds - spent)
|
|
if not _wait_between_polls(gap, activity_state):
|
|
return result(
|
|
"interrupted",
|
|
"The user interrupted the wait; still not connected: "
|
|
f"{', '.join(pending)}. " + _WAIT_UNFINISHED_NOTE,
|
|
)
|
|
spent += gap
|
|
continue
|
|
consecutive_errors = 0
|
|
spent += time.monotonic() - poll_started
|
|
|
|
live_entries = []
|
|
live_slugs = set()
|
|
for item in items or ():
|
|
if not isinstance(item, dict):
|
|
continue
|
|
slug = str(item.get("connector", "")).lower()
|
|
if slug in wanted and item.get("connected"):
|
|
live_entries.append(item)
|
|
live_slugs.add(slug)
|
|
pending = [c for c in connectors if c not in live_slugs]
|
|
|
|
if not pending:
|
|
return result(
|
|
"connected",
|
|
"All requested apps are connected. Go ahead and use them.",
|
|
)
|
|
if spent >= timeout_seconds:
|
|
return result(
|
|
"timeout",
|
|
f"Waited about {int(spent)}s; still not connected: "
|
|
f"{', '.join(pending)}. " + _WAIT_UNFINISHED_NOTE,
|
|
)
|
|
# The last gap before the budget line is the REMAINDER, not a full
|
|
# gap: requiring room for a whole gap silently halved short asks (a 6s
|
|
# ask returned after one poll and zero sleep) and undershot every
|
|
# budget by one gap (120s asks exited near 115s).
|
|
gap = min(_POLL_GAP_SECONDS, timeout_seconds - spent)
|
|
if not _wait_between_polls(gap, activity_state):
|
|
# Interrupted mid-wait: answer with what the last poll saw rather
|
|
# than spending a round trip the user has just asked us to stop for.
|
|
return result(
|
|
"interrupted",
|
|
"The user interrupted the wait; still not connected: "
|
|
f"{', '.join(pending)}. " + _WAIT_UNFINISHED_NOTE,
|
|
)
|
|
spent += gap
|
|
|
|
|
|
def manage_connections(
|
|
args: Dict[str, Any],
|
|
*,
|
|
client_factory: Optional[Callable[[], Any]] = None,
|
|
seen_instructions: Optional[set] = None,
|
|
rendered_links: Optional[Dict[str, Dict[str, float]]] = None,
|
|
session_id: Optional[str] = None,
|
|
) -> str:
|
|
"""Dispatch one ``manage_connections`` action. Returns a JSON string."""
|
|
action = str(args.get("action") or "status").strip().lower()
|
|
|
|
if action not in _CONNECTOR_ACTIONS:
|
|
return tool_error(
|
|
f"action must be one of {', '.join(_CONNECTOR_ACTIONS)}. "
|
|
"Local MCP servers are set up with setup_mcp, not here. "
|
|
"Disconnecting an account is done by the user in the Nous Portal "
|
|
"dashboard, not through this tool."
|
|
)
|
|
|
|
raw_connectors = args.get("connectors")
|
|
if isinstance(raw_connectors, str):
|
|
raw_connectors = [raw_connectors]
|
|
connectors: List[str] = []
|
|
if isinstance(raw_connectors, list):
|
|
for c in raw_connectors:
|
|
c = str(c or "").strip().lower()
|
|
if c and c not in connectors:
|
|
connectors.append(c)
|
|
|
|
try:
|
|
client = (client_factory or _default_client)()
|
|
if action == "status":
|
|
items = client.list_connectors()
|
|
if connectors:
|
|
wanted = set(connectors)
|
|
items = [i for i in items if str(i.get("connector", "")).lower() in wanted]
|
|
return json.dumps(
|
|
{
|
|
"connectors": items,
|
|
"hint": (
|
|
"connected=false means calls to that connector will return "
|
|
"CONNECTION_REQUIRED. Use action 'connect' to get an "
|
|
"authorization link for the user."
|
|
),
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
|
|
if not connectors:
|
|
return tool_error(
|
|
f"'{action}' requires 'connectors': the connector slugs to authorize "
|
|
"(e.g. [\"gmail\"]). Use action 'status' to list them."
|
|
)
|
|
|
|
rendered = rendered_links if rendered_links is not None else _rendered_links
|
|
if action == "wait":
|
|
shown = _rendered_for(session_id, rendered)
|
|
never_shown = [c for c in connectors if c not in shown]
|
|
if never_shown:
|
|
return tool_error(
|
|
"wait refused: this session never obtained a connect link "
|
|
f"for {', '.join(never_shown)}, so there is nothing to "
|
|
"wait for. Call action 'connect' for those connectors first "
|
|
"and put the links in front of the user, then wait."
|
|
)
|
|
just_minted = [
|
|
c
|
|
for c in connectors
|
|
if time.monotonic() - shown[c] < _LINKS_JUST_MINTED_SECONDS
|
|
]
|
|
if just_minted:
|
|
# The connect that minted these links ran moments ago — same
|
|
# tool batch, so no message carrying them has reached the user
|
|
# yet. Bounce (never an error) instead of blocking a spinner.
|
|
return json.dumps(
|
|
{
|
|
"status": "pending",
|
|
"connectors": [],
|
|
"pending": list(connectors),
|
|
"note": (
|
|
"Not waiting yet: the connect links for "
|
|
f"{', '.join(just_minted)} were minted moments ago, "
|
|
"in this same turn — the user has not seen them. "
|
|
"Send your message showing the links FIRST, then "
|
|
"call wait on your next turn."
|
|
),
|
|
},
|
|
ensure_ascii=False,
|
|
)
|
|
timeout_seconds, timeout_note = _clamp_timeout(args.get("timeout_seconds"))
|
|
return _wait_for_connections(
|
|
client,
|
|
connectors,
|
|
timeout_seconds=timeout_seconds,
|
|
timeout_note=timeout_note,
|
|
)
|
|
|
|
response = client.connections(connectors, reinitiate=(action == "reconnect"))
|
|
seen = seen_instructions if seen_instructions is not None else _seen_instructions
|
|
results = []
|
|
for entry in response.get("results", []):
|
|
connector = str(entry.get("connector") or "")
|
|
out: Dict[str, Any] = {
|
|
"connector": connector,
|
|
"status": entry.get("status"),
|
|
}
|
|
if entry.get("connect_url"):
|
|
out["connect_url"] = entry["connect_url"]
|
|
out["note"] = (
|
|
"Show this link to the user; they open it in a browser to "
|
|
"authorize. Then use action 'wait' to hold for the "
|
|
"connection instead of guessing when they are done."
|
|
)
|
|
_record_rendered(session_id, connector, rendered)
|
|
elif entry.get("status") == "active":
|
|
# Already authorized: the gateway mints no link for a live
|
|
# connection. This connector is still ADDRESSED by this call —
|
|
# record it, or the documented connect-then-wait sequence
|
|
# refuses on the success case and loops the model through
|
|
# fresh mints that can never fill the record.
|
|
out["note"] = "Already connected — no link needed."
|
|
# never_fresh: there is no link the user must see before a
|
|
# wait, so the just-minted bounce must not fire for this one —
|
|
# an immediate wait legitimately returns connected on poll one.
|
|
_record_rendered(session_id, connector, rendered, never_fresh=True)
|
|
instruction = entry.get("instruction")
|
|
if instruction:
|
|
seen_key = (str(session_id or ""), connector)
|
|
with _seen_instructions_lock:
|
|
if seen_key not in seen:
|
|
seen.add(seen_key)
|
|
out["instruction"] = instruction
|
|
results.append(out)
|
|
return json.dumps(
|
|
{"results": results, "summary": response.get("summary", {})},
|
|
ensure_ascii=False,
|
|
)
|
|
except Exception as exc:
|
|
# Registered tools go through the registry's catch-wrap, but keep the
|
|
# message model-actionable rather than a raw traceback.
|
|
logger.debug("manage_connections %s failed: %s", action, exc)
|
|
return tool_error(
|
|
f"The connector gateway request failed: {exc}. "
|
|
"If this persists, the user can manage connections in the Nous Portal."
|
|
)
|
|
|
|
|
|
MANAGE_CONNECTIONS_SCHEMA = {
|
|
"name": "manage_connections",
|
|
"description": (
|
|
"Manage remote connector accounts (Gmail, Linear, Notion, ...) served "
|
|
"through the tool gateway. Actions: "
|
|
"'status' lists connectors and whether each is connected; 'connect' "
|
|
"starts an authorization for the given connectors and returns a link "
|
|
"for the USER to open in a browser (never open it yourself); "
|
|
"'reconnect' restarts a broken authorization; "
|
|
"'wait' blocks until the given connectors report connected. Pass "
|
|
"SEVERAL slugs in one call to get all authorization links at once. "
|
|
"When a connector tool "
|
|
"call returns CONNECTION_REQUIRED, use 'connect' and show the link. "
|
|
"Send the message that shows the user the links FIRST; on your NEXT "
|
|
"turn call 'wait' with those same slugs instead of guessing when the "
|
|
"user is done — it polls for you (a wait in the same turn as the "
|
|
"connect is bounced, because the user cannot have seen the links "
|
|
"yet). 'wait' requires 'connectors', and only accepts connectors this "
|
|
"session already addressed with 'connect' (already-connected apps "
|
|
"count). A 'timeout' or 'interrupted' result is NOT an "
|
|
"error: the user has not finished connecting, so ask them whether to "
|
|
"keep waiting, continue without those apps, or get fresh links. "
|
|
"Local MCP servers are configured separately. "
|
|
"This tool can NOT disconnect, delete, or revoke an account — that is "
|
|
"deliberately user-only. When asked, say so and direct the user to "
|
|
"the Nous Portal (their org's Connectors page) or the desktop app."
|
|
),
|
|
"parameters": {
|
|
"type": "object",
|
|
"properties": {
|
|
"action": {
|
|
"type": "string",
|
|
"enum": list(_CONNECTOR_ACTIONS),
|
|
"description": "Defaults to status.",
|
|
},
|
|
"connectors": {
|
|
"type": "array",
|
|
"items": {"type": "string"},
|
|
"description": (
|
|
"Connector slugs. REQUIRED for connect, reconnect and wait "
|
|
"(e.g. [\"gmail\", \"linear\"]); optional filter for status."
|
|
),
|
|
},
|
|
"timeout_seconds": {
|
|
"type": "integer",
|
|
"description": (
|
|
"For action 'wait' only: how long to hold the call open. "
|
|
f"Defaults to {int(_WAIT_DEFAULT_SECONDS)}, clamped to "
|
|
f"{int(_WAIT_MIN_SECONDS)}-{int(_WAIT_MAX_SECONDS)}. Ask for "
|
|
"more and the result carries a 'timeout_note' saying the cap "
|
|
"was applied; call wait again to keep waiting."
|
|
),
|
|
},
|
|
},
|
|
"required": [],
|
|
},
|
|
}
|
|
|
|
|
|
registry.register(
|
|
name="manage_connections",
|
|
toolset="connections",
|
|
schema=MANAGE_CONNECTIONS_SCHEMA,
|
|
# Registry dispatch does not re-run check_fn: enforce the off switch for
|
|
# stale schemas without rebuilding a conversation's cached tool list.
|
|
handler=lambda args, **kw: (
|
|
manage_connections(args, session_id=kw.get("session_id"))
|
|
if _connectors_available()
|
|
else tool_error("Connectors are not available in this session.")
|
|
),
|
|
check_fn=_connectors_available,
|
|
emoji="🔗",
|
|
)
|