aggrete 0.1.0__py3-none-any.whl

aggrete/accumulator.py

@@ -0,0 +1,111 @@
+"""Layer 4 state: what each user has already pulled, across every connector.
+
+Keyed to the *user*, not the session. A new chat must not reset the budget.
+MemoryStore is for tests and single-process runs; RedisStore is what you
+deploy, because the whole point is state shared across clients and gateways.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Iterable, Protocol
+
+
+def parse_window(w: str | int) -> int:
+    if isinstance(w, int):
+        return w
+    units = {"s": 1, "m": 60, "h": 3600, "d": 86400}
+    return int(w[:-1]) * units[w[-1]]
+
+
+class Store(Protocol):
+    def record(self, user: str, domain: str, entities: Iterable[str], ttl: int) -> None: ...
+    def entities(self, user: str, domain: str) -> set[str]: ...
+    def domains(self, user: str) -> set[str]: ...
+    def grant(self, user: str, rule_id: str, ttl: int, purpose: str) -> None: ...
+    def granted(self, user: str, rule_id: str) -> str | None: ...
+    def reset(self, user: str) -> None: ...
+
+
+class MemoryStore:
+    def __init__(self) -> None:
+        self._ents: dict[tuple[str, str], dict[str, float]] = {}
+        self._grants: dict[tuple[str, str], tuple[float, str]] = {}
+
+    def _live(self, key: tuple[str, str]) -> dict[str, float]:
+        now = time.time()
+        bucket = {e: exp for e, exp in self._ents.get(key, {}).items() if exp > now}
+        self._ents[key] = bucket
+        return bucket
+
+    def record(self, user, domain, entities, ttl):
+        key = (user, domain)
+        bucket = self._live(key)
+        exp = time.time() + ttl
+        # touching a domain with no entities still marks the domain as seen
+        bucket.setdefault("__touched__", exp)
+        for e in entities:
+            bucket[e] = exp
+        self._ents[key] = bucket
+
+    def entities(self, user, domain):
+        return {e for e in self._live((user, domain)) if e != "__touched__"}
+
+    def domains(self, user):
+        return {d for (u, d) in list(self._ents) if u == user and self._live((u, d))}
+
+    def grant(self, user, rule_id, ttl, purpose):
+        self._grants[(user, rule_id)] = (time.time() + ttl, purpose)
+
+    def granted(self, user, rule_id):
+        exp, purpose = self._grants.get((user, rule_id), (0, ""))
+        return purpose if exp > time.time() else None
+
+    def reset(self, user):
+        for k in [k for k in self._ents if k[0] == user]:
+            del self._ents[k]
+        for k in [k for k in self._grants if k[0] == user]:
+            del self._grants[k]
+
+
+class RedisStore:
+    """Same contract, backed by Redis sets with TTLs."""
+
+    def __init__(self, client, prefix: str = "coc"):
+        self.r = client
+        self.p = prefix
+
+    def _key(self, user, domain):
+        return f"{self.p}:ents:{user}:{domain}"
+
+    def record(self, user, domain, entities, ttl):
+        key = self._key(user, domain)
+        pipe = self.r.pipeline()
+        pipe.sadd(key, "__touched__", *entities)
+        pipe.expire(key, ttl)
+        pipe.sadd(f"{self.p}:domains:{user}", domain)
+        pipe.expire(f"{self.p}:domains:{user}", ttl)
+        pipe.execute()
+
+    def entities(self, user, domain):
+        raw = self.r.smembers(self._key(user, domain))
+        return {e.decode() if isinstance(e, bytes) else e for e in raw} - {"__touched__"}
+
+    def domains(self, user):
+        live = set()
+        for d in self.r.smembers(f"{self.p}:domains:{user}"):
+            d = d.decode() if isinstance(d, bytes) else d
+            if self.r.exists(self._key(user, d)):
+                live.add(d)
+        return live
+
+    def grant(self, user, rule_id, ttl, purpose):
+        self.r.setex(f"{self.p}:grant:{user}:{rule_id}", ttl, purpose)
+
+    def granted(self, user, rule_id):
+        v = self.r.get(f"{self.p}:grant:{user}:{rule_id}")
+        return (v.decode() if isinstance(v, bytes) else v) if v else None
+
+    def reset(self, user):
+        for k in self.r.scan_iter(f"{self.p}:*:{user}*"):
+            self.r.delete(k)

aggrete/auth.py

@@ -0,0 +1,96 @@
+"""Who is calling. Over HTTP the answer comes from the bearer token, never from config.
+
+Two verifiers:
+
+- JWTVerifier. Production. Validates RS256/ES256 JWTs from your IdP (Okta,
+                  Entra, Google, Auth0, Keycloak…) against its JWKS, checks
+                  issuer/audience/expiry, and derives the user from a claim.
+- StaticTokens. Development and tests. A fixed token → subject map.
+
+The identity the policy engine sees is `identity_for(token)`: the configured
+claim (default `email`, falling back to `sub`). Per-user state in the
+accumulator is keyed on it, so pick a claim that is stable across sessions.
+"""
+
+from __future__ import annotations
+
+import time
+from typing import Any
+
+import jwt
+from jwt import PyJWKClient
+from mcp.server.auth.provider import AccessToken, TokenVerifier
+
+DEFAULT_IDENTITY_CLAIMS = ("email", "preferred_username", "sub")
+
+
+class StaticTokens(TokenVerifier):
+    """`auth: {mode: static, tokens: {<token>: {subject: ..., scopes: [...]}}}`."""
+
+    def __init__(self, tokens: dict[str, dict[str, Any]]):
+        self.tokens = tokens
+
+    async def verify_token(self, token: str) -> AccessToken | None:
+        spec = self.tokens.get(token)
+        if not spec:
+            return None
+        return AccessToken(token=token, client_id=spec.get("client_id", "static"),
+                           scopes=list(spec.get("scopes", [])), subject=spec["subject"],
+                           claims={"email": spec.get("email", spec["subject"])})
+
+
+class JWTVerifier(TokenVerifier):
+    """`auth: {mode: jwt, issuer: ..., audience: ..., jwks_url: ... | public_key: ...}`."""
+
+    def __init__(self, issuer: str, audience: str | list[str], *, jwks_url: str | None = None,
+                 public_key: str | None = None, algorithms: list[str] | None = None,
+                 required_scopes: list[str] | None = None, leeway: int = 30):
+        if not (jwks_url or public_key):
+            jwks_url = issuer.rstrip("/") + "/.well-known/jwks.json"
+        self.issuer, self.audience = issuer, audience
+        self.jwks = PyJWKClient(jwks_url, cache_keys=True) if jwks_url else None
+        self.public_key = public_key
+        self.algorithms = algorithms or ["RS256", "ES256", "RS384", "ES384", "RS512", "ES512"]
+        self.required_scopes = required_scopes or []
+        self.leeway = leeway
+
+    async def verify_token(self, token: str) -> AccessToken | None:
+        try:
+            key = self.public_key or self.jwks.get_signing_key_from_jwt(token).key
+            claims = jwt.decode(token, key, algorithms=self.algorithms, issuer=self.issuer,
+                                audience=self.audience, leeway=self.leeway,
+                                options={"require": ["exp", "iss", "sub"]})
+        except jwt.PyJWTError:
+            return None
+        scopes = claims.get("scope", "")
+        scopes = scopes.split() if isinstance(scopes, str) else list(scopes or [])
+        if any(s not in scopes for s in self.required_scopes):
+            return None
+        return AccessToken(token=token, client_id=str(claims.get("client_id") or claims.get("azp") or claims.get("aud")),
+                           scopes=scopes, expires_at=claims.get("exp"), subject=claims["sub"], claims=claims)
+
+
+def build_verifier(auth_cfg: dict[str, Any]) -> TokenVerifier:
+    mode = auth_cfg.get("mode", "jwt")
+    if mode == "static":
+        return StaticTokens(auth_cfg["tokens"])
+    if mode == "jwt":
+        return JWTVerifier(auth_cfg["issuer"], auth_cfg["audience"], jwks_url=auth_cfg.get("jwks_url"),
+                           public_key=auth_cfg.get("public_key"), algorithms=auth_cfg.get("algorithms"),
+                           required_scopes=auth_cfg.get("required_scopes"))
+    raise ValueError(f"unknown auth mode {mode!r}")
+
+
+def identity_for(token: AccessToken, claim: str | None = None) -> str:
+    claims = token.claims or {}
+    order = (claim,) if claim else DEFAULT_IDENTITY_CLAIMS
+    for c in order:
+        if c and claims.get(c):
+            return str(claims[c])
+    if token.subject:
+        return token.subject
+    raise ValueError("token carries no usable identity claim")
+
+
+def unexpired(token: AccessToken) -> bool:
+    return token.expires_at is None or token.expires_at > time.time()

aggrete/connectors/drive.py

@@ -0,0 +1,170 @@
+"""Google Drive as an Aggrete upstream, one tool pair per folder.
+
+    python -m aggrete.connectors.drive --credentials sa.json --root "Northwind"
+
+The proxy runs this over stdio and holds the service-account key; people never
+do. Each subfolder of the root becomes two tools, `search_<folder>` and
+`read_<folder>`, so the policy can name folders as domains:
+
+    domains:
+      "drive__*_restructuring_plan": restructuring-plan
+      "drive__*_legal_hold": legal-hold
+      "drive__*": drive-general
+
+Results carry the owner's and last editor's email, which is what the policy
+counts as "people". Read only: the service account needs Viewer on the root.
+"""
+
+from __future__ import annotations
+
+import argparse
+import asyncio
+import json
+import re
+import time
+from pathlib import Path
+
+import httpx2 as httpx
+import jwt
+from mcp.server.mcpserver import MCPServer
+
+API = "https://www.googleapis.com/drive/v3"
+SCOPE = "https://www.googleapis.com/auth/drive.readonly"
+EXPORT = {"application/vnd.google-apps.document": "text/plain",
+          "application/vnd.google-apps.spreadsheet": "text/csv",
+          "application/vnd.google-apps.presentation": "text/plain"}
+FIELDS = "files(id,name,mimeType,modifiedTime,owners(emailAddress,displayName),lastModifyingUser(emailAddress),parents,webViewLink)"
+
+
+def slug(name: str) -> str:
+    return re.sub(r"[^a-z0-9]+", "_", name.lower()).strip("_")
+
+
+class Drive:
+    """Minimal Drive v3 client on a service account. No Google SDK: PyJWT + httpx."""
+
+    def __init__(self, credentials: str | Path):
+        self.sa = json.loads(Path(credentials).read_text())
+        self._tok, self._exp = None, 0.0
+        self.http = httpx.Client(timeout=30)
+
+    def token(self) -> str:
+        if self._tok and time.time() < self._exp - 60:
+            return self._tok
+        now = int(time.time())
+        assertion = jwt.encode({"iss": self.sa["client_email"], "scope": SCOPE, "aud": self.sa["token_uri"],
+                                "iat": now, "exp": now + 3600}, self.sa["private_key"], algorithm="RS256")
+        r = self.http.post(self.sa["token_uri"], data={"grant_type": "urn:ietf:params:oauth:grant-type:jwt-bearer", "assertion": assertion})
+        r.raise_for_status()
+        self._tok, self._exp = r.json()["access_token"], time.time() + r.json().get("expires_in", 3600)
+        return self._tok
+
+    def get(self, path: str, **params):
+        r = self.http.get(API + path, params=params, headers={"Authorization": f"Bearer {self.token()}"})
+        r.raise_for_status()
+        return r
+
+    def list(self, q: str, page_size: int = 20) -> list[dict]:
+        return self.get("/files", q=q, fields=FIELDS, pageSize=page_size, supportsAllDrives="true",
+                        includeItemsFromAllDrives="true").json().get("files", [])
+
+    def folder_by_name(self, name: str) -> dict | None:
+        fs = self.list(f"mimeType='application/vnd.google-apps.folder' and name='{name}' and trashed=false", 5)
+        return fs[0] if fs else None
+
+    def subfolders(self, folder_id: str) -> list[dict]:
+        return self.list(f"'{folder_id}' in parents and mimeType='application/vnd.google-apps.folder' and trashed=false", 50)
+
+    def descendants(self, folder_id: str, depth: int = 3) -> set[str]:
+        ids, frontier = {folder_id}, [folder_id]
+        for _ in range(depth):
+            nxt = []
+            for f in frontier:
+                for sub in self.subfolders(f):
+                    if sub["id"] not in ids:
+                        ids.add(sub["id"]); nxt.append(sub["id"])
+            frontier = nxt
+        return ids
+
+    def search(self, folder_id: str, query: str) -> list[dict]:
+        out = []
+        for fid in self.descendants(folder_id):
+            q = f"'{fid}' in parents and trashed=false and mimeType!='application/vnd.google-apps.folder'"
+            if query.strip():
+                safe = query.replace("'", "\\'")
+                q += f" and (fullText contains '{safe}' or name contains '{safe}')"
+            out += self.list(q, 20)
+        return out
+
+    def read(self, file_id: str, folder_id: str) -> tuple[dict, str]:
+        meta = self.get(f"/files/{file_id}", fields="id,name,mimeType,parents,owners(emailAddress,displayName),lastModifyingUser(emailAddress),webViewLink", supportsAllDrives="true").json()
+        if not (set(meta.get("parents") or []) & self.descendants(folder_id)):
+            raise PermissionError("that file is not in this folder")
+        mime = meta.get("mimeType", "")
+        if mime in EXPORT:
+            text = self.get(f"/files/{file_id}/export", mimeType=EXPORT[mime]).text
+        elif mime.startswith("text/") or mime in ("application/json",):
+            text = self.get(f"/files/{file_id}", alt="media", supportsAllDrives="true").text
+        else:
+            text = f"[{mime}: binary file, {meta.get('name')}; not rendered]"
+        return meta, text[:20000]
+
+
+def build(drive: Drive, root_name: str) -> MCPServer:
+    server = MCPServer("drive")
+    root = drive.folder_by_name(root_name)
+    if not root:
+        # Stay up so the proxy starts; tell whoever asks what is missing.
+        @server.tool(name="status", description="Why no Drive folders are available yet.")
+        def status() -> str:
+            return json.dumps({"error": f"No folder named {root_name!r} is shared with {drive.sa['client_email']}. "
+                                        "Share it (Viewer) and restart the proxy."})
+        return server
+    folders = drive.subfolders(root["id"]) or [root]
+
+    @server.tool(name="folders", description="List the Google Drive folders you can search here. Call this first when asked about documents, files or anything in Google Drive.")
+    def folders_tool() -> str:
+        return json.dumps({"drive_folders": [{"name": f["name"], "search_tool": f"search_{slug(f['name'])}", "read_tool": f"read_{slug(f['name'])}"} for f in folders]})
+
+    for f in folders:
+        s = slug(f["name"]); fid = f["id"]; label = f["name"]
+
+        def make(fid=fid, label=label):
+            def search(query: str = "") -> str:
+                return json.dumps({"folder": label, "files": [
+                    {"id": x["id"], "name": x["name"], "type": x.get("mimeType"), "modified": x.get("modifiedTime"),
+                     "owner_email": (x.get("owners") or [{}])[0].get("emailAddress"),
+                     "editor_email": (x.get("lastModifyingUser") or {}).get("emailAddress"), "link": x.get("webViewLink")}
+                    for x in drive.search(fid, query)]})
+            def read(file_id: str) -> str:
+                meta, text = drive.read(file_id, fid)
+                return json.dumps({"folder": label, "name": meta["name"], "owner_email": (meta.get("owners") or [{}])[0].get("emailAddress"),
+                                   "editor_email": (meta.get("lastModifyingUser") or {}).get("emailAddress"), "text": text})
+            return search, read
+
+        search, read = make()
+        sdesc = f"Search Google Drive for documents in the '{label}' folder, by words in the title or full text. Use this to look for {label.lower()} in Drive; leave the query empty to list everything in the folder."
+        rdesc = f"Read a Google Drive document from the '{label}' folder (Docs, Sheets and Slides come back as text)."
+        server.tool(name=f"search_{s}", description=sdesc)(search)
+        server.tool(name=f"read_{s}", description=rdesc)(read)
+    return server
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--credentials", required=True, help="service-account JSON key")
+    ap.add_argument("--root", default="Aggrete", help="name of the shared root folder; its subfolders become tools")
+    ap.add_argument("--list", action="store_true", help="print the tools that would be exposed and exit")
+    a = ap.parse_args()
+    drive = Drive(a.credentials)
+    if a.list:
+        root = drive.folder_by_name(a.root)
+        print("root:", root["name"] if root else None, root["id"] if root else "")
+        for f in (drive.subfolders(root["id"]) if root else []):
+            print(f"  {f['name']!r:32} -> drive__search_{slug(f['name'])}, drive__read_{slug(f['name'])}")
+        return
+    asyncio.run(build(drive, a.root).run_stdio_async())
+
+
+if __name__ == "__main__":
+    main()

aggrete/entities.py

@@ -0,0 +1,69 @@
+"""Pull stable person identifiers out of whatever a connector hands back.
+
+This function is where this whole design succeeds or fails. Names are not
+identifiers. Prefer source-system IDs, then email, and treat free text as a
+last resort. Tune `IDENTIFIER_KEYS` against your own connectors before trusting
+any threshold you set in coc.yaml.
+
+Linking: one JSON object that carries several identifier fields (e.g. both
+`email` and `employee_id`) is ONE person and yields ONE canonical key. Email is
+preferred as the canonical form because it is the identifier most likely to be
+shared across connectors. Cross-domain overlap (COC-HR-004) only works when
+both sides produce the same key. When a record has no email the first ID field
+found (in `ID_KEYS` order) is used instead.
+"""
+
+from __future__ import annotations
+
+import json
+import re
+
+EMAIL_KEYS = ("email", "user_email", "primary_email", "mail", "owner_email", "editor_email", "emailaddress")
+ID_KEYS = ("employee_id", "person_id", "worker_id", "user_id", "assignee_id",
+           "owner_id", "sfid", "slack_user_id")
+IDENTIFIER_KEYS = set(EMAIL_KEYS) | set(ID_KEYS)
+EMAIL_RE = re.compile(r"[\w.+-]+@[\w-]+\.[\w.-]+")
+
+
+def normalize(value: str) -> str:
+    v = str(value).strip().lower()
+    return f"p:{v}"
+
+
+def _canonical(record: dict) -> str | None:
+    """Return the single canonical key for a record, or None if it has none."""
+    lowered = {k.lower(): v for k, v in record.items() if isinstance(v, (str, int))}
+    for k in EMAIL_KEYS:
+        if k in lowered and str(lowered[k]).strip():
+            return normalize(lowered[k])
+    for k in ID_KEYS:
+        if k in lowered and str(lowered[k]).strip():
+            return normalize(lowered[k])
+    return None
+
+
+def from_json(obj, out: set[str]) -> set[str]:
+    if isinstance(obj, dict):
+        key = _canonical(obj)
+        if key:
+            out.add(key)
+        # Nested objects may describe other people (e.g. "manager": {...}).
+        for k, v in obj.items():
+            if k.lower() not in IDENTIFIER_KEYS:
+                from_json(v, out)
+    elif isinstance(obj, list):
+        for item in obj:
+            from_json(item, out)
+    return out
+
+
+def extract(text: str) -> list[str]:
+    """Best-effort extraction from one tool result payload."""
+    found: set[str] = set()
+    try:
+        from_json(json.loads(text), found)
+    except (ValueError, TypeError):
+        pass
+    if not found:  # fall back to emails in prose
+        found |= {normalize(m) for m in EMAIL_RE.findall(text or "")}
+    return sorted(found)