"""Python 3.10+ REST example. Standard library only; no credential output.""" from __future__ import annotations import base64 import json import os import re import sys import time from collections.abc import Mapping from typing import Any from urllib.error import HTTPError, URLError from urllib.parse import urlencode, urlsplit from urllib.request import HTTPRedirectHandler, ProxyHandler, Request, build_opener TIMEOUT_SECONDS = 15 MAX_BYTES = 1024 * 1024 class ClientError(Exception): """A safe, credential-free error suitable for command-line output.""" def __init__(self, message: str, status: int | None = None) -> None: super().__init__(message) self.status = status class NoRedirect(HTTPRedirectHandler): def redirect_request( self, req: Request, fp: Any, code: int, msg: str, headers: Any, newurl: str, ) -> None: return None def status_error(status: int) -> ClientError: reasons = { 401: "Credentials or token are expired, revoked, or invalid.", 403: "This agent does not have the exact tool permission.", 424: "A required workspace tool is not ready.", 429: "A usage limit was reached; do not retry automatically.", } return ClientError( f"TaskVegas request failed (HTTP {status}). " + reasons.get(status, "Check the service and request before trying again."), status=status, ) class TaskVegasClient: def __init__(self, environment: Mapping[str, str] | None = None) -> None: env = os.environ if environment is None else environment try: url = urlsplit(env.get("TASKVEGAS_URL", "https://taskvegas.si")) port = url.port except ValueError as error: raise ClientError("TASKVEGAS_URL must be a bare HTTPS origin.") from error loopback = url.hostname in {"localhost", "127.0.0.1", "::1"} if ( not url.hostname or port == 0 or (url.scheme != "https" and not (url.scheme == "http" and loopback)) or url.username is not None or url.password is not None or url.query or url.fragment or url.path not in {"", "/"} or re.search(r"[\s\\\x00-\x1f\x7f]", env.get("TASKVEGAS_URL", "")) ): raise ClientError( "Use a bare HTTPS origin, or HTTP on localhost for development." ) hostname = f"[{url.hostname}]" if ":" in url.hostname else url.hostname default_port = 443 if url.scheme == "https" else 80 suffix = f":{port}" if port and port != default_port else "" self.origin = f"{url.scheme}://{hostname}{suffix}" self.resource = self.origin + "/api/v1" client_id = env.get("TASKVEGAS_CLIENT_ID", "") client_secret = env.get("TASKVEGAS_CLIENT_SECRET", "") if not all( isinstance(value, str) and 0 < len(value) <= 4096 and not re.search(r"[\s\x00-\x1f\x7f]", value) for value in (client_id, client_secret) ) or ":" in client_id: raise ClientError( "Set valid TASKVEGAS_CLIENT_ID and TASKVEGAS_CLIENT_SECRET " "environment variables." ) pair = f"{client_id}:{client_secret}".encode() self.basic = base64.b64encode(pair).decode("ascii") self.access: dict[str, Any] | None = None # Avoid environment proxies for local development and never follow redirects. self.opener = build_opener(ProxyHandler({}), NoRedirect()) def _json(self, request: Request, maximum: int = MAX_BYTES) -> dict[str, Any]: deadline = time.monotonic() + TIMEOUT_SECONDS try: with self.opener.open(request, timeout=TIMEOUT_SECONDS) as response: if response.headers.get_content_type() != "application/json": raise ClientError("TaskVegas returned an unexpected response type.") parts: list[bytes] = [] size = 0 while True: if time.monotonic() > deadline: raise ClientError( "TaskVegas response exceeded the client deadline." ) chunk = response.read1(min(65536, maximum + 1 - size)) # urllib applies a socket timeout, not hard wall-clock cancellation. # Reject a late read, including EOF, before accepting any response. if time.monotonic() > deadline: raise ClientError( "TaskVegas response exceeded the client elapsed-time limit." ) if not chunk: break parts.append(chunk) size += len(chunk) if size > maximum: raise ClientError( "TaskVegas response exceeded the client size limit." ) value = json.loads(b"".join(parts).decode("utf-8")) if not isinstance(value, dict): raise ClientError("TaskVegas returned an unexpected JSON shape.") return value except HTTPError as error: error.close() raise status_error(error.code) from None except (URLError, TimeoutError, OSError) as error: raise ClientError( "TaskVegas could not be reached within its network timeout." ) from error except (ValueError, UnicodeError) as error: raise ClientError("TaskVegas returned invalid JSON.") from error def token(self, force: bool = False) -> str: if ( not force and self.access and self.access["expires_at"] > time.monotonic() + 30 ): return self.access["token"] self.access = None request = Request( self.origin + "/token", data=urlencode( {"grant_type": "client_credentials", "resource": self.resource} ).encode(), headers={ "Authorization": "Basic " + self.basic, "Content-Type": "application/x-www-form-urlencoded", "Accept": "application/json", }, method="POST", ) value = self._json(request, 16384) token = value.get("access_token") expiry = value.get("expires_in") if ( not isinstance(token, str) or re.fullmatch(r"[A-Za-z0-9_.~-]{1,1024}", token) is None or str(value.get("token_type", "")).lower() != "bearer" or type(expiry) is not int or not 1 <= expiry <= 3600 ): raise ClientError("TaskVegas returned an invalid token response.") self.access = {"token": token, "expires_at": time.monotonic() + expiry} return token def request(self, path: str, body: dict[str, Any] | None = None) -> dict[str, Any]: if path not in {"/api/v1/tools", "/api/v1/tools/taskvegas_search_docs/execute"}: raise ClientError("This example supports only its fixed tool endpoints.") method = "GET" if body is None else "POST" def send(token: str) -> dict[str, Any]: headers = {"Authorization": "Bearer " + token, "Accept": "application/json"} data = None if body is not None: data = json.dumps(body).encode() headers["Content-Type"] = "application/json" return self._json( Request(self.origin + path, data=data, headers=headers, method=method) ) token = self.token() try: return send(token) except ClientError as error: if error.status != 401: raise self.access = None if method != "GET": raise # Only the discovery GET may be retried after renewing a rejected token. renewed = self.token(force=True) try: return send(renewed) except ClientError as final: if final.status == 401: self.access = None raise def search_documentation(self, query: str = "diagrams") -> dict[str, Any]: catalog = self.request("/api/v1/tools") tools = catalog.get("tools") if catalog.get("resource") != self.resource or not isinstance(tools, list): raise ClientError("Tool discovery returned the wrong resource or shape.") tool = next( ( item for item in tools if isinstance(item, dict) and item.get("name") == "taskvegas_search_docs" ), None, ) if ( not tool or tool.get("readOnly") is not True or tool.get("paid") is not False ): raise ClientError( "Grant the free taskvegas_search_docs tool before running this example." ) if not isinstance(query, str) or not query.strip() or len(query) > 300: raise ClientError("Use a documentation query of 1–300 characters.") value = self.request( "/api/v1/tools/taskvegas_search_docs/execute", {"input": {"query": query, "limit": 3}}, ) if not isinstance(value.get("result"), dict): raise ClientError("Tool execution returned an incomplete response.") return value["result"] def main() -> int: try: result = TaskVegasClient().search_documentation() items = result.get("items", []) count = len(items) if isinstance(items, list) else 0 print(f"Documentation search completed: {count} result(s).") return 0 except ClientError as error: print(str(error), file=sys.stderr) except Exception: print( "Agent client failed. Check configuration, connectivity " "and service status. " "The tool request was not retried.", file=sys.stderr, ) return 1 if __name__ == "__main__": sys.exit(main())