Coverage for stackone_ai / toolset.py: 99%
513 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-10-09 08:03 +0000
1"""The StackOne toolset: fetches the served tool catalog and exposes it to agent frameworks."""
3from __future__ import annotations
5import concurrent.futures
6import contextlib
7import copy
8import fnmatch
9import logging
10import os
11import threading
12import time
13from collections.abc import Callable, Collection, Sequence
14from dataclasses import dataclass, field
15from typing import Any, TypeVar
17import httpx
19from stackone_ai.tools import (
20 McpToolDefinition,
21 RateLimitRetryingClient,
22 RateLimitTimeout,
23 StackOneMcpTool,
24 StackOneTool,
25 Tools,
26 _js_number,
27 _json_text,
28 _json_type,
29 build_request_headers,
30 fetch_mcp_tools,
31 is_end_user_refusal,
32 is_rate_limited,
33 is_sdk_owned_header,
34)
35from stackone_ai.types import (
36 DEFAULT_BASE_URL,
37 SUBMIT_FEEDBACK_TOOL_NAME,
38 ExecuteToolsConfig,
39 FeedbackCategory,
40 FeedbackRating,
41 FeedbackSource,
42 Headers,
43 JsonDict,
44 StackOneAPIError,
45 StackOneError,
46 ToolMode,
47 ToolParameters,
48 ToolsetConfigError,
49 ToolsetError,
50 ToolsetLoadError,
51)
53logger = logging.getLogger("stackone.tools")
55_UNSET = object()
56"""Sentinel: `mode=None` is a real mode (individual), so it cannot mean "use the default"."""
58_Listing = tuple[McpToolDefinition, str | None, str]
59"""A served tool, the account that listed it, and the endpoint it was listed from."""
61_RetryFailed = Callable[["_Catalog"], Collection[str | None]]
62"""Given a cached catalog, the accounts marked failed in it to list again now."""
64# The search_actions meta tool's served schema caps top_k at 50.
65_MAX_TOP_K = 50
67FAILED_ACCOUNT_RETRY_SECONDS = 30.0
68"""How long an account that failed to list tools is left out of a cached catalog.
70Its healthy siblings' listings are cached all the same. A later call within this window
71serves them without the failed account, and without warning about it again; the first
72call after it lists the failed account again. ``execute()`` does not wait: it lists one
73that could serve its action again straight away, unless a provider lookup failed to name
74it within the window.
75"""
77# Looked up at call time, so tests can move it on rather than wait.
78_clock = time.monotonic
81@dataclass(frozen=True)
82class _Catalog:
83 """A cached catalog: each healthy account's listing, and when and how each failed one failed."""
85 listings: dict[str | None, list[_Listing]]
86 failed_at: dict[str | None, float] = field(default_factory=dict)
87 errors: dict[str | None, Exception] = field(default_factory=dict)
89 def in_order(self, account_scope: list[str | None]) -> list[_Listing]:
90 return [listing for account in account_scope for listing in self.listings.get(account, [])]
92 def failures(self) -> dict[str | None, Exception]:
93 """Each account left out of this catalog, with the error it failed with."""
94 return {account: self.errors[account] for account in self.failed_at if account in self.errors}
97_Item = TypeVar("_Item")
98_Result = TypeVar("_Result")
101def _fan_out(
102 fn: Callable[[_Item], _Result], items: Sequence[_Item]
103) -> list[tuple[_Item, _Result | Exception]]:
104 """Run ``fn`` over ``items`` concurrently; each item's result or failure, in item order.
106 A 429 that outlasted its retries is raised at once instead of being returned: the
107 caller would skip it like any other failure and hand back a partial result. Work not
108 yet started is cancelled; calls already running finish in the background, unread.
109 """
110 pool = concurrent.futures.ThreadPoolExecutor(max_workers=min(len(items), 10))
111 try:
112 futures = {pool.submit(fn, item): index for index, item in enumerate(items)}
113 outcomes: dict[int, _Result | Exception] = {}
114 for future in concurrent.futures.as_completed(futures):
115 try:
116 outcomes[futures[future]] = future.result()
117 except Exception as exc:
118 if is_rate_limited(exc):
119 raise
120 outcomes[futures[future]] = exc
121 finally:
122 pool.shutdown(wait=False, cancel_futures=True)
123 return [(item, outcomes[index]) for index, item in enumerate(items)]
126def _all_accounts_failed(failures: list[tuple[str | None, Exception]]) -> Exception:
127 """The error for a listing in which every account failed, in account order.
129 When every failure is a StackOneAPIError with the same status, the first of them: two
130 accounts answering 401 are a revoked key, and a caller must be able to tell that from a
131 429 as it can with one account. Otherwise a ToolsetLoadError, caused by an
132 ExceptionGroup of the failures and keeping them in ``failures`` too.
133 """
134 errors = [failure for _, failure in failures]
135 first = errors[0]
136 if isinstance(first, StackOneAPIError) and all(
137 isinstance(error, StackOneAPIError) and error.status_code == first.status_code for error in errors
138 ):
139 return first
140 error = ToolsetLoadError(
141 "Every account failed to list tools: "
142 + "; ".join(f"{account}: {failure}" for account, failure in failures),
143 failures=errors,
144 )
145 error.__cause__ = ExceptionGroup("Every account failed to list tools", errors)
146 return error
149class StackOneToolSet:
150 """Main class for accessing StackOne tools.
152 The toolset is a thin client over the served catalog: it lists tools from the
153 MCP endpoint and executes each one over MCP ``tools/call`` on the endpoint, and
154 with the account, that listed it. It does not rewrite, filter or invent schemas.
155 """
157 def __init__(
158 self,
159 api_key: str | None = None,
160 account_id: str | None = None,
161 base_url: str | None = None,
162 execute: ExecuteToolsConfig | None = None,
163 timeout: float | None = None,
164 tool_mode: ToolMode | None = None,
165 headers: Headers | None = None,
166 include_non_shared: bool = False,
167 ) -> None:
168 """Initialize StackOne tools with authentication
170 Args:
171 api_key: Optional API key. If not provided, will try to get from STACKONE_API_KEY env var
172 account_id: Optional account ID
173 base_url: Optional base URL override for API requests. Falls back to
174 STACKONE_BASE_URL, then to the default, https://api.stackone.com. An
175 empty string, whether passed here or read from the environment, is
176 treated as not set rather than as a literal empty host.
177 execute: Execution configuration. Controls default account scoping
178 for tool execution. Pass ``{"account_ids": ["acc-1"]}`` to scope
179 tools to specific accounts.
180 timeout: Request timeout in seconds for tool listing and execution.
181 Default: 60. Takes precedence over ``execute.timeout`` if set.
182 Increase for slow providers (e.g. Workday).
183 tool_mode: How the endpoint lists tools. ``"search_execute"`` returns
184 two meta tools per connector instead of one tool per action,
185 keeping the catalog small enough for a model's context.
186 headers: Extra HTTP headers sent with every request this toolset makes
187 (``GET /accounts`` and every MCP request). ``Authorization``,
188 ``x-account-id`` and ``User-Agent`` are always the SDK's own and
189 cannot be set here. ``x-end-user-id`` can: it is passed through as
190 given, unless ``GET /accounts`` reported the account's end user,
191 which then replaces it.
192 include_non_shared: Whether account discovery uses non-shared accounts
193 (``shared`` false in ``GET /accounts``). Each belongs to a single end
194 user, so by default discovery skips them, with a warning naming them;
195 pass their ids to use them without this.
197 Raises:
198 ToolsetConfigError: If no API key is provided or found in environment, or
199 ``account_id`` is an empty string
200 """
201 # An empty account_id is usually an unset variable, and treating it as unset would
202 # silently widen every call to all active shared accounts.
203 if account_id == "":
204 raise ToolsetConfigError("account_id must not be an empty string")
205 api_key_value = api_key or os.getenv("STACKONE_API_KEY")
206 if not api_key_value:
207 raise ToolsetConfigError(
208 "An API key must be provided, either to the toolset or in the STACKONE_API_KEY "
209 "environment variable"
210 )
211 ignored_headers = [name for name in (headers or {}) if is_sdk_owned_header(name)]
212 if ignored_headers:
213 logger.warning(
214 "Ignoring headers %s: the SDK sets them itself. Pass the API key and account ids "
215 "through their own options instead.",
216 ", ".join(_json_text(name) for name in ignored_headers),
217 )
219 self.api_key: str = api_key_value
220 self.account_id = account_id
221 self.base_url = base_url or os.getenv("STACKONE_BASE_URL") or DEFAULT_BASE_URL
222 self._headers: Headers = dict(headers or {})
223 self._account_ids: list[str] = self._validate_account_ids(
224 execute.get("account_ids") if execute else None
225 )
226 self._execute_config: ExecuteToolsConfig | None = execute
227 execute_timeout = execute.get("timeout") if execute else None
228 self._timeout: float = (
229 timeout if timeout is not None else (execute_timeout if execute_timeout is not None else 60.0)
230 )
231 # Cache the listing, not the Tools wrapper. StackOneTool objects are mutable
232 # (Tools.set_account_id rebinds them), so handing the same instances back on a
233 # cache hit let one caller silently rescope every later caller's tools.
234 self._catalog_cache: dict[tuple[Any, ...], _Catalog] = {}
235 self._discovered_account_ids: list[str] | None = None
236 # Shared while in flight, so concurrent fetch_tools() and search() calls on a
237 # fresh toolset make one GET /accounts between them rather than one each.
238 self._discovering: concurrent.futures.Future[list[str]] | None = None
239 # Bumped by clear_catalog_cache(). A listing already in flight when the cache is
240 # cleared captured the generation it started under, and refuses to write back if
241 # it has moved — otherwise the stale catalog lands *after* the clear and is
242 # served for the life of the process, which is the one thing the clear exists
243 # to prevent.
244 self._cache_generation = 0
245 self._cache_lock = threading.Lock()
246 # Each non-shared account's end-user id, from the last successful GET /accounts.
247 # Replaced whole by the next one, and kept by clear_catalog_cache(): it describes
248 # the accounts, not the catalog.
249 self._end_user_ids: dict[str, str] = {}
250 # Each account's provider, from the same GET /accounts and replaced with it: what
251 # execute() reads to tell whether an account that failed to list could have served
252 # an action.
253 self._providers: dict[str, str] = {}
254 # When a provider lookup last failed to name each account, which execute() then
255 # neither looks up nor lists again early for FAILED_ACCOUNT_RETRY_SECONDS.
256 self._provider_misses: dict[str, float] = {}
257 # The GET /accounts in flight, if any: an account's end user is looked up by joining
258 # it rather than making another.
259 self._fetching_accounts: concurrent.futures.Future[list[JsonDict]] | None = None
260 self._include_non_shared = include_non_shared
261 # Bumped at the start of every fetch_accounts() call, so two overlapping calls —
262 # fetch_accounts() racing discovery, say — are ordered by when they started
263 # rather than when they returned: a slower, older-started response must not
264 # overwrite a newer one that already landed.
265 self._accounts_sequence = 0
266 self._end_user_ids_sequence = 0
267 self._tool_mode: ToolMode | None = tool_mode
269 # 2.x read STACKONE_ACCOUNT_ID; 3.x never does, and without an account it uses every
270 # active shared account on the key. Anyone upgrading with only the variable set would
271 # otherwise widen to other end users' accounts with no sign of it. An empty
272 # account_ids list is unset too, the same as _resolve_account_ids() treats it.
273 scoped = account_id is not None or bool(self._account_ids)
274 if os.getenv("STACKONE_ACCOUNT_ID") and not scoped:
275 logger.warning(
276 "STACKONE_ACCOUNT_ID is set, but the SDK does not read it: with no account id passed, "
277 "every active shared account on this API key is used. Pass an account id to scope the "
278 "toolset."
279 )
281 def set_accounts(self, account_ids: list[str] | None) -> StackOneToolSet:
282 """Set account IDs for filtering tools
284 Returns:
285 This toolset instance for chaining
286 """
287 self._account_ids = self._validate_account_ids(account_ids)
288 self.clear_catalog_cache()
289 return self
291 def clear_catalog_cache(self) -> None:
292 """Invalidate the cached tool catalog.
294 Call when linked accounts change outside of ``set_accounts`` or when
295 you need to force a fresh fetch from the StackOne MCP endpoint. The end-user ids
296 recorded from GET /accounts are kept until the next one replaces them; the accounts
297 a provider lookup failed to name are forgotten.
298 """
299 with self._cache_lock:
300 self._cache_generation += 1
301 self._catalog_cache.clear()
302 self._provider_misses.clear()
303 self._discovered_account_ids = None
304 self._discovering = None
306 def fetch_tools(
307 self,
308 *,
309 account_ids: list[str] | None = None,
310 providers: list[str] | None = None,
311 actions: list[str] | None = None,
312 mode: Any = _UNSET,
313 ) -> Tools:
314 """Fetch tools with optional filtering by account IDs, providers, and actions
316 Args:
317 account_ids: Optional list of account IDs to filter by.
318 If not provided, uses accounts set via set_accounts()
319 providers: Optional list of provider names (e.g., ['hibob', 'bamboohr']).
320 Case-insensitive matching.
321 actions: Optional list of action patterns with glob support
322 (e.g., ['*_list_employees', 'hibob_create_employees'])
324 Returns:
325 Collection of tools matching the filter criteria
327 Raises:
328 ToolsetLoadError: If there is an error loading the tools. When every account
329 fails for differing reasons, its ``failures`` holds each account's error,
330 and its ``__cause__`` is an ExceptionGroup of them.
331 StackOneAPIError: With ``status_code`` 429 if the API is still rate limiting
332 after retries, even when only one of several accounts is; or with the API's
333 status when every account fails with that same status, so a caller can tell
334 a revoked key's 401 from a 429. Other per-account failures skip that account
335 with a warning, and leave it out for ``FAILED_ACCOUNT_RETRY_SECONDS``.
337 A 429 is retried up to three times, after the server's ``Retry-After`` (capped at
338 30 seconds) or else a 1s, 2s, 4s backoff with jitter. A wait that would not end
339 within ``timeout`` is not started, and the 429 is raised instead.
341 Examples:
342 tools = toolset.fetch_tools(account_ids=['123', '456'])
343 tools = toolset.fetch_tools(providers=['hibob', 'bamboohr'])
344 tools = toolset.fetch_tools(actions=['*_list_employees'])
346 Note:
347 For organizations with multiple connected accounts, calling `fetch_tools()`
348 without `account_ids` discovers and fetches the catalog for every active shared
349 account (non-shared ones too with ``include_non_shared=True``).
350 If your organization has many accounts, pass explicit `account_ids` to avoid
351 excessive round trips and blowing model context limits.
353 Discovery skips non-shared accounts unless the toolset was built with
354 ``include_non_shared=True``.
355 """
356 mode = self._tool_mode if mode is _UNSET else mode
357 tools, _ = self._fetch_tools(account_ids, providers, actions, mode, retry_failed=None)
358 return tools
360 def _fetch_tools(
361 self,
362 account_ids: list[str] | None,
363 providers: list[str] | None,
364 actions: list[str] | None,
365 mode: ToolMode | None,
366 *,
367 retry_failed: _RetryFailed | None,
368 ) -> tuple[Tools, dict[str | None, Exception]]:
369 """:meth:`fetch_tools`, and each account in scope it left out with the error it failed
370 with. ``retry_failed`` picks, from a cached catalog, the accounts marked failed to list
371 again now rather than once they are due.
372 """
373 try:
374 # Captured before account resolution, which may discover accounts: a
375 # clear_catalog_cache() landing mid-discovery must stop the catalog it
376 # resolves to from being cached too, not just the discovered list itself.
377 generation = self._cache_generation
378 account_scope: list[str | None] = sorted(
379 dict.fromkeys(self._resolve_account_ids(account_ids)), key=lambda a: (a is None, a)
380 )
382 # Keyed on what was fetched, not on how it is filtered: providers and
383 # actions narrow the list in memory, so they must not force a refetch.
384 # base_url and api_key belong here — leaving them out meant reassigning
385 # either one kept serving the old catalog, still pointed at the old host.
386 listings, failures = self._catalog(account_scope, mode, generation, retry_failed=retry_failed)
388 all_tools = [
389 self._create_tool(tool_def, account, endpoint)
390 for tool_def, account, endpoint in self._dedupe_global_tools(listings)
391 ]
393 if providers:
394 all_tools = [tool for tool in all_tools if self._filter_by_provider(tool.name, providers)]
396 if actions:
397 all_tools = [tool for tool in all_tools if self._filter_by_action(tool.name, actions)]
399 return Tools(all_tools), failures
401 except (ToolsetError, StackOneError):
402 # StackOneAPIError carries the HTTP status. Re-wrapping it below would throw
403 # that away, so a caller could not tell a 401 from a 429.
404 raise
405 except Exception as exc: # pragma: no cover - unexpected runtime errors
406 raise ToolsetLoadError(f"Error fetching tools: {exc}") from exc
408 @staticmethod
409 def _validate_account_ids(account_ids: Any) -> list[str]:
410 """A copy of the account ids, refusing a bare string, a non-list, a non-string id or
411 an empty id. ``None`` means not given: no ids.
413 An empty or ``None`` id would be sent with no ``x-account-id``, so it is rejected
414 rather than letting the server answer for an account nobody chose, as in Node. A
415 dict is refused rather than iterated, which read ``{"id": "acc-1"}`` as the account
416 ``"id"``.
417 """
418 if account_ids is None:
419 return []
420 if isinstance(account_ids, str):
421 raise ToolsetConfigError(
422 "account_ids must be a list of account ids, not a string. "
423 f"Did you mean [{_json_text(account_ids)}]?"
424 )
425 if not isinstance(account_ids, list) or not all(isinstance(i, str) for i in account_ids):
426 raise ToolsetConfigError("account_ids must be a list of account id strings")
427 ids = list(account_ids)
428 if "" in ids:
429 raise ToolsetConfigError("account_ids must not contain an empty account id")
430 return ids
432 def _resolve_account_ids(self, account_ids: list[str] | None) -> list[str]:
433 """The accounts a call is scoped to, in the order they were given or discovered.
435 The argument, then ``set_accounts()``, then the constructor's ``account_id``, then
436 every active shared account ``GET /accounts`` lists (non-shared ones too with
437 ``include_non_shared=True``).
438 """
439 if account_ids is not None:
440 account_ids = self._validate_account_ids(account_ids)
441 resolved = account_ids or self._account_ids
442 if not resolved and self.account_id:
443 resolved = [self.account_id]
444 if not resolved:
445 resolved = self._discover_account_ids()
446 return list(resolved)
448 def _catalog(
449 self,
450 account_scope: list[str | None],
451 mode: ToolMode | None,
452 generation: int,
453 *,
454 retry_failed: _RetryFailed | None = None,
455 ) -> tuple[list[_Listing], dict[str | None, Exception]]:
456 """Every scoped account's catalog, from the cache where it can be.
458 One unusable account must not cost the caller every other account's tools, so a
459 failing account is skipped with a warning. The healthy accounts' listings are cached
460 all the same, and the failed account is left out, silently, until
461 ``FAILED_ACCOUNT_RETRY_SECONDS`` have passed, when the next call lists it again. Not
462 caching at all made every call re-list every account, and wait out the full timeout
463 on one that hangs.
465 A 429 that outlasts its retries is raised, not skipped: it is the key's, not the
466 account's, and skipping it would hand back a catalog missing whichever accounts
467 happened to be throttled. When every account fails, see :func:`_all_accounts_failed`.
469 ``generation`` is the cache generation as of the start of the call that resolved
470 ``account_scope``: a clear_catalog_cache() during account discovery must stop the
471 catalog this scope resolves to from being cached, same as the discovered list.
473 Returned with each account left out and the error it failed with. ``retry_failed`` is
474 given the cached catalog and picks accounts marked failed in it to list again now,
475 without waiting for them to be due.
476 """
477 key = self._cache_key(account_scope, mode)
478 with self._cache_lock:
479 cached = self._catalog_cache.get(key)
480 now = _clock()
481 if cached is None:
482 due = account_scope
483 else:
484 retried = set(retry_failed(cached)) if retry_failed else set()
485 due = [
486 account
487 for account, failed_at in cached.failed_at.items()
488 if account in retried or now - failed_at >= FAILED_ACCOUNT_RETRY_SECONDS
489 ]
490 if not due:
491 return cached.in_order(account_scope), cached.failures()
493 # No param-style pin: arguments are sent verbatim and the server maps them with its
494 # own reverse map, so the model sees whatever style the server serves.
495 endpoint = f"{self.base_url.rstrip('/')}/mcp"
496 if mode:
497 endpoint = f"{endpoint}?tool-mode={mode}"
499 def _fetch_for_account(
500 account: str | None,
501 ) -> list[_Listing]:
502 without_end_user = account is not None and self._end_user_id(account) is None
503 try:
504 listed = fetch_mcp_tools(endpoint, self._build_mcp_headers(account), timeout=self._timeout)
505 except StackOneAPIError as exc:
506 # Refused for want of an end user it was sent without: look it up, and retry
507 # once with it.
508 if not (without_end_user and account is not None and self._look_up_end_user(account, exc)):
509 raise
510 listed = fetch_mcp_tools(endpoint, self._build_mcp_headers(account), timeout=self._timeout)
511 return [(tool_def, account, endpoint) for tool_def in listed]
513 def _store(catalog: _Catalog) -> None:
514 # Merged with whatever is cached now, not overwritten: a concurrent call for the
515 # same key may have stored a healthier catalog since this call read it, and a
516 # wholesale overwrite would hide its successes behind this call's failures for
517 # FAILED_ACCOUNT_RETRY_SECONDS. A listing wins over a failure for the same account.
518 with self._cache_lock:
519 if generation != self._cache_generation:
520 return
521 existing = self._catalog_cache.get(key)
522 if existing is not None:
523 merged_listings = {**existing.listings, **catalog.listings}
524 merged_failed_at = {**existing.failed_at, **catalog.failed_at}
525 merged_errors = {**existing.errors, **catalog.errors}
526 for account in catalog.listings:
527 merged_failed_at.pop(account, None)
528 merged_errors.pop(account, None)
529 catalog = _Catalog(merged_listings, merged_failed_at, merged_errors)
530 self._catalog_cache[key] = catalog
532 if len(account_scope) == 1:
533 catalog = _Catalog({account_scope[0]: _fetch_for_account(account_scope[0])})
534 _store(catalog)
535 return catalog.in_order(account_scope), {}
537 listings = dict(cached.listings) if cached else {}
538 failed_at = dict(cached.failed_at) if cached else {}
539 errors = dict(cached.errors) if cached else {}
540 failures: list[tuple[str | None, Exception]] = []
541 for account, outcome in _fan_out(_fetch_for_account, due):
542 if isinstance(outcome, Exception):
543 failures.append((account, outcome))
544 else:
545 listings[account] = outcome
546 failed_at.pop(account, None)
547 errors.pop(account, None)
548 if not listings:
549 raise _all_accounts_failed(failures)
550 failed_now = _clock()
551 for account, failure in failures:
552 logger.warning("Skipping account that failed to list tools — %s: %s", account, failure)
553 failed_at[account] = failed_now
554 errors[account] = failure
556 catalog = _Catalog(listings, failed_at, errors)
557 _store(catalog)
558 return catalog.in_order(account_scope), catalog.failures()
560 @staticmethod
561 def _dedupe_global_tools(
562 listings: list[_Listing],
563 ) -> list[_Listing]:
564 """Keep only the first listing of the feedback tool.
566 It is global rather than account-scoped, so every account's listing carries an
567 identical copy. Left in, N accounts meant N same-named tools and a duplicate-name
568 warning about a clash that cannot misroute anything — noise that teaches callers
569 to ignore the warning that does matter.
570 """
571 seen_feedback = False
572 kept: list[_Listing] = []
573 for listing in listings:
574 if listing[0].name == SUBMIT_FEEDBACK_TOOL_NAME:
575 if seen_feedback:
576 continue
577 seen_feedback = True
578 kept.append(listing)
579 return kept
581 def _cache_key(self, account_scope: list[str | None], mode: ToolMode | None) -> tuple[Any, ...]:
582 return (
583 tuple(sorted(account_scope, key=lambda a: (a is None, a))),
584 mode,
585 self.base_url,
586 self.api_key,
587 )
589 def _meta_tools(self, suffix: str, account_ids: list[str] | None = None) -> list[StackOneTool]:
590 """The server's per-connector meta tools, regardless of this toolset's mode.
592 The mode is passed down rather than assigned to ``self``. Flipping instance
593 state here raced with any concurrent ``fetch_tools()``: that call could read
594 the flipped mode partway through and cache search_execute meta tools under the
595 individual-mode key, permanently, for every later caller.
596 """
597 tools, _ = self._meta_tools_and_failures(suffix, account_ids, retry_failed=None)
598 return tools
600 def _meta_tools_and_failures(
601 self, suffix: str, account_ids: list[str] | None, *, retry_failed: _RetryFailed | None
602 ) -> tuple[list[StackOneTool], dict[str | None, Exception]]:
603 tools, failures = self._fetch_tools(
604 account_ids, None, None, "search_execute", retry_failed=retry_failed
605 )
606 return [tool for tool in tools if tool.name.endswith(suffix)], failures
608 def search(
609 self,
610 query: str,
611 *,
612 top_k: int = 10,
613 account_ids: list[str] | None = None,
614 ) -> list[JsonDict]:
615 """Find actions matching a natural-language query.
617 Searches every linked connector and merges the results, so the catalog
618 never has to be loaded into a model's context.
620 Args:
621 query: What you want to do, e.g. "list recent comments".
622 top_k: Maximum results, across every connector searched.
623 account_ids: Restrict to these accounts. Defaults to every active shared one
624 (non-shared ones too with ``include_non_shared=True``).
626 Returns:
627 At most ``top_k`` action dicts, best first across every connector, each carrying
628 at least ``action_id`` and ``description``, the ``account_id`` of the account
629 whose connector found it, and the ``session_id`` of the search when the server
630 issued one. The same action linked on two accounts is two hits. Pass
631 ``session_id`` to :meth:`execute` and :meth:`submit_feedback` to link the calls,
632 and ``account_id`` in ``account_ids`` to run the action on that account.
634 Raises:
635 StackOneAPIError: With ``status_code`` 429 if the API is still rate limiting
636 after retries, for any account or connector. A 429 is retried as in
637 :meth:`fetch_tools`; other per-connector failures are skipped with a
638 warning.
639 """
640 # The server rejects anything outside 1..50, but only after a round trip per
641 # connector — and reports it as a load failure, which reads like an outage
642 # rather than a typo. Fail here instead, where the caller can see why.
643 if not isinstance(top_k, int) or isinstance(top_k, bool) or not 1 <= top_k <= _MAX_TOP_K:
644 raise ToolsetConfigError(
645 f"top_k must be an integer between 1 and {_MAX_TOP_K}, got {_json_text(top_k)}"
646 )
648 tools = self._meta_tools("_search_actions", account_ids)
649 if not tools: 649 ↛ 650line 649 didn't jump to line 650 because the condition on line 649 was never true
650 return []
652 def _search_one(tool: StackOneTool) -> list[JsonDict]:
653 found = tool.execute({"query": query, "top_k": top_k})
654 actions = list(found.get("actions", []))
655 # The server returns session_id once per search, beside the actions. Results
656 # from every connector are merged and re-ranked below, so this is the last
657 # point at which a hit can still be traced to the search, and the account,
658 # that produced it.
659 session_id = found.get("session_id")
660 traced: JsonDict = {}
661 if isinstance(session_id, str) and session_id:
662 traced["session_id"] = session_id
663 account_id = tool.get_account_id()
664 if account_id: 664 ↛ 666line 664 didn't jump to line 666 because the condition on line 664 was always true
665 traced["account_id"] = account_id
666 return [{**action, **traced} for action in actions]
668 results: list[JsonDict] = []
669 failures: list[str] = []
670 # Fan out the way fetch_tools() does. Serially, a customer with a dozen
671 # connectors pays the sum of every connector's latency on the headline call.
672 for tool, outcome in _fan_out(_search_one, tools):
673 if isinstance(outcome, Exception):
674 # Skip everything but a rate limit, as fetch_tools() does. Catching only
675 # StackOneError let a transport failure — which _describe_mcp_failure
676 # reports as a ToolsetLoadError — abort the whole search, which is
677 # precisely the flakiness this guard exists to absorb.
678 failures.append(f"{tool.name}: {outcome}")
679 else:
680 results.extend(outcome)
682 if failures and not results:
683 raise ToolsetLoadError("No connector returned results. " + " | ".join(failures))
684 for failure in failures:
685 logger.warning("Skipping connector that failed to search — %s", failure)
687 # Concatenating per-connector results leaves the list grouped by connector, so
688 # results[0] would be the best hit of whichever connector answered first rather
689 # than the best hit overall. Rank globally; the server scores every action on the
690 # same scale. Actions without a score sort last rather than raising.
691 def _score(action: JsonDict) -> float:
692 raw = action.get("similarity_score")
693 if isinstance(raw, bool) or not isinstance(raw, int | float):
694 return 0.0
695 return float(raw)
697 # Then cut to top_k: each connector was asked for top_k, so the merged list can
698 # hold many more.
699 results.sort(key=_score, reverse=True)
700 return results[:top_k]
702 def execute(
703 self,
704 action_id: str,
705 arguments: JsonDict | None = None,
706 *,
707 account_ids: list[str] | None = None,
708 session_id: str | None = None,
709 ) -> JsonDict:
710 """Execute an action by id, as returned by :meth:`search`.
712 Always runs through the connector's ``_execute_action`` meta tool, so
713 ``arguments`` is the nested form every action's ``example_request``
714 shows — ``{"query": {...}, "path": {...}, "body": {...}}``. A
715 ``fetch_tools()`` tool takes the keys its own served schema names instead;
716 routing by whether an id happened to be in the catalog would make the
717 argument shape depend on something the caller cannot see.
719 ``arguments["headers"]`` is forwarded to the action: ``*_execute_action`` serves an
720 open ``headers`` object, so any header name is accepted except Authorization,
721 x-account-id, User-Agent and x-end-user-id, which the SDK sets itself and which are
722 dropped.
724 ``session_id`` is the value a :meth:`search` hit carries. Passing it links
725 this call to that search server-side.
727 ``account_ids`` restricts routing to these accounts. Pass a search hit's
728 ``account_id`` to run the action on the account that found it.
730 Raises:
731 ToolsetConfigError: If ``action_id``, ``arguments`` or ``session_id`` is malformed,
732 or the action's connector is linked on more than one account in scope.
733 ToolArgumentsError: If the arguments cannot be encoded as JSON.
734 ToolsetLoadError: If no connector matches, or an account in scope that could serve
735 the action failed to list. An account left out after failing whose provider in
736 ``GET /accounts`` is the action's connector, or whose provider no
737 ``GET /accounts`` has named, is listed again first, without waiting out
738 ``FAILED_ACCOUNT_RETRY_SECONDS``; if it still fails, the action is not run on
739 another account in its place. When a failed account's provider is not known,
740 one ``GET /accounts`` is made first to learn it; if that fails, it stays
741 unknown, and is neither looked up nor listed again early for
742 ``FAILED_ACCOUNT_RETRY_SECONDS``.
743 StackOneAPIError: If the action fails, including when the server rejects the
744 arguments.
745 """
746 if not isinstance(action_id, str) or not action_id:
747 raise ToolsetConfigError(f"action_id must be a non-empty string, got {_json_text(action_id)}")
748 if arguments is not None and not isinstance(arguments, dict):
749 raise ToolsetConfigError(f"arguments must be a JSON object, got {_json_type(arguments)}")
750 if session_id is not None and (not isinstance(session_id, str) or not session_id):
751 raise ToolsetConfigError(f"session_id must be a non-empty string, got {_json_text(session_id)}")
753 suffix = "_execute_action"
754 # At most one GET /accounts per call, to learn the providers of failed accounts: with
755 # explicit ids none was made, and an unknown provider could serve any action.
756 learnt = False
758 def _learn_providers(accounts: list[str]) -> set[str]:
759 nonlocal learnt
760 if learnt:
761 return set()
762 missed = self._learn_providers(accounts)
763 learnt = missed is not None
764 return missed or set()
766 def _retry_failed(cached: _Catalog) -> list[str | None]:
767 # An account marked failed that could serve the action is listed again now: left
768 # out, it would leave the action to whichever account did list, which may be
769 # another end user's. One whose provider cannot serve it is not, so it does not
770 # cost every other connector's actions its timeout. Nor is one a lookup recently
771 # failed to name: it would pay its timeout on every call in the window.
772 failed = [account for account in cached.failed_at if account is not None]
773 missed = self._recently_missed(failed) | _learn_providers(failed)
774 listed = [
775 self._connector_of_name(tool_def.name, account, suffix)
776 for account, listings in cached.listings.items()
777 for tool_def, _, _ in listings
778 if tool_def.name.endswith(suffix)
779 ]
780 return [
781 account
782 for account in failed
783 if account not in missed and self._may_serve(account, action_id, listed)
784 ]
786 meta_tools, failed = self._meta_tools_and_failures(suffix, account_ids, retry_failed=_retry_failed)
787 _learn_providers([account for account in failed if account is not None])
788 connectors = [self._connector_of(tool, suffix) for tool in meta_tools]
789 blocking = sorted(
790 (
791 (account, failure)
792 for account, failure in failed.items()
793 if account is not None and self._may_serve(account, action_id, connectors)
794 ),
795 key=lambda pair: pair[0],
796 )
797 if blocking:
798 raise ToolsetLoadError(
799 f"{_json_text(action_id)} may be served by an account that failed to list ("
800 + "; ".join(f"{account}: {failure}" for account, failure in blocking)
801 + "). Pass the account id to use, such as a search hit's account_id."
802 )
803 matches = [
804 tool
805 for tool in meta_tools
806 if action_id.lower().startswith(self._connector_of(tool, "_execute_action") + "_")
807 ]
808 if matches:
809 # Longest connector wins: with both `browser` and `browser_linkedin` linked,
810 # the first token alone would route every browser_linkedin action to browser.
811 best = max(len(self._connector_of(t, "_execute_action")) for t in matches)
812 finalists = [t for t in matches if len(self._connector_of(t, "_execute_action")) == best]
813 if len(finalists) > 1:
814 # The same provider linked twice, which discovery makes common. Picking one
815 # would run the action against an account the caller never chose — another
816 # end user's, possibly.
817 ordered = sorted(finalists, key=lambda t: t.get_account_id() or "")
818 raise ToolsetConfigError(
819 f"{_json_text(action_id)} matches {len(finalists)} connectors on different accounts "
820 f"({', '.join(f'{t.name} on {t.get_account_id()}' for t in ordered)}). "
821 "Pass the account id to use, such as a search hit's account_id."
822 )
823 tool = finalists[0]
824 # action_id LAST, removed first so it is last in key order too, as in Node.
825 # Spreading arguments over it let a model-supplied "action_id" silently replace
826 # the action the caller pinned — the exact thing a host app pins it for.
827 # session_id only when given: the served schema makes it an optional string, so
828 # an absent key is valid and a null is not.
829 call_arguments: JsonDict = dict(arguments or {})
830 call_arguments.pop("action_id", None)
831 if session_id is not None:
832 call_arguments.pop("session_id", None)
833 call_arguments["session_id"] = session_id
834 call_arguments["action_id"] = action_id
835 # Returned as the server wrote it, the same shape tool.execute() returns.
836 return tool.execute(call_arguments)
838 raise ToolsetLoadError(
839 f"No connector found for {_json_text(action_id)}. Use search() to discover valid action ids."
840 )
842 def submit_feedback(
843 self,
844 rating: FeedbackRating,
845 tool_names: Sequence[str],
846 *,
847 feedback: str | None = None,
848 category: FeedbackCategory | None = None,
849 session_id: str | None = None,
850 source: FeedbackSource = "model",
851 account_ids: list[str] | None = None,
852 action_run_id: str | None = None,
853 ) -> JsonDict:
854 """Record a verdict on how well the tools served this session.
856 Calls the server's ``stackone_submit_feedback`` tool once, through the account with
857 the lowest id among those the call is scoped to: ``account_ids``, or else the
858 toolset's, or else every active shared one (non-shared ones too with
859 ``include_non_shared=True``). Pass the ``session_id`` from a :meth:`search`
860 hit to attach the feedback to that session.
862 Args:
863 rating: ``"positive"``, ``"negative"`` or ``"neutral"``.
864 tool_names: The tools or action ids the feedback is about.
865 feedback: An optional one-line reason.
866 category: What the feedback is about, e.g. ``"search"`` or ``"execute"``.
867 session_id: The session to link this feedback to.
868 source: Who produced the feedback.
869 account_ids: Accounts to choose from; only the lowest id is used. Defaults to
870 the toolset's own.
871 action_run_id: The action run the feedback is about. Sent only when given.
873 Raises:
874 ToolsetConfigError: If ``tool_names`` is a string or ``session_id`` is empty or
875 not a string.
876 ToolsetLoadError: If feedback is not enabled for this project.
877 """
878 if isinstance(tool_names, str):
879 raise ToolsetConfigError(
880 "tool_names must be a list of tool names, not a string. "
881 f"Did you mean [{_json_text(tool_names)}]?"
882 )
883 if session_id is not None and (not isinstance(session_id, str) or not session_id):
884 raise ToolsetConfigError(f"session_id must be a non-empty string, got {_json_text(session_id)}")
886 # Found in the served catalog, never built here: the server only serves the tool
887 # when the org flag and the project setting are both on, and a client-side stand-in
888 # would report success for feedback that went nowhere. It is global and served in
889 # every mode, so one account's search_execute listing — two tools per connector,
890 # not one per action — is enough to find it.
891 # The lowest id is the one a caller can predict, whatever order GET /accounts lists
892 # them in.
893 scope = self._resolve_account_ids(account_ids)
894 tool = self.fetch_tools(account_ids=[min(scope)], mode="search_execute").get_tool(
895 SUBMIT_FEEDBACK_TOOL_NAME
896 )
897 if tool is None:
898 raise ToolsetLoadError(
899 f"The server did not serve {SUBMIT_FEEDBACK_TOOL_NAME}: feedback is not enabled "
900 "for this project."
901 )
903 # Unset optionals are omitted, never sent as null: the served schema declares them
904 # as optional strings, so a null is a present key of the wrong type and fails
905 # validation where an absent one would not.
906 arguments: JsonDict = {"rating": rating, "tool_names": list(tool_names)}
907 optional = {
908 "feedback": feedback,
909 "category": category,
910 "session_id": session_id,
911 "action_run_id": action_run_id,
912 "source": source,
913 }
914 arguments.update({key: value for key, value in optional.items() if value is not None})
915 return tool.execute(arguments)
917 def openai(self, *, account_ids: list[str] | None = None) -> list[JsonDict]:
918 """Get tools in OpenAI function calling format."""
919 return self.fetch_tools(account_ids=account_ids).to_openai()
921 def langchain(self, *, account_ids: list[str] | None = None) -> Any:
922 """Get tools in LangChain format."""
923 return self.fetch_tools(account_ids=account_ids).to_langchain()
925 def pydantic_ai(self, *, account_ids: list[str] | None = None) -> list[Any]:
926 """Get tools as Pydantic AI ``Tool`` instances.
928 Requires ``stackone-ai[pydantic-ai]`` (installs ``pydantic-ai-slim``).
929 """
930 return self.fetch_tools(account_ids=account_ids).to_pydantic_ai()
932 @staticmethod
933 def _connector_of(tool: StackOneTool, suffix: str) -> str:
934 """The connector a meta tool belongs to: its name minus the account id and suffix.
936 The account id is stripped by identity, not by splitting on the last underscore.
937 Account ids are nanoid-shaped and nanoid's default alphabet includes ``_``, so
938 splitting turned ``mock_acc_1_execute_action`` into connector ``mock_acc`` and
939 made every action on that account unroutable — with an error blaming the
940 caller's action id.
941 """
942 return StackOneToolSet._connector_of_name(tool.name, tool.get_account_id(), suffix)
944 @staticmethod
945 def _connector_of_name(name: str, account: str | None, suffix: str) -> str:
946 """:meth:`_connector_of` for a meta tool's served name and the account that listed it."""
947 stem = name[: -len(suffix)] if name.endswith(suffix) else name
948 if account and stem.endswith(f"_{account}"): 948 ↛ 950line 948 didn't jump to line 950 because the condition on line 948 was always true
949 stem = stem[: -(len(account) + 1)]
950 return stem.lower()
952 def _may_serve(self, account_id: str, action_id: str, listed_connectors: list[str]) -> bool:
953 """Whether an account could serve an action: its provider is the action's connector,
954 or no GET /accounts has named its provider.
956 The connector is the longest prefix of the action among the connectors listed and the
957 providers GET /accounts named, as execute() routes: a failed ``browser`` account cannot
958 serve ``browser_linkedin_search_people``.
959 """
960 with self._cache_lock:
961 provider = self._providers.get(account_id)
962 known = [p.lower() for p in self._providers.values()]
963 if provider is None:
964 return True
965 lowered = action_id.lower()
966 matching = [c for c in [*listed_connectors, *known] if lowered.startswith(c + "_")]
967 connector = max(matching, key=len, default=None)
968 return provider.lower() == connector
970 def _filter_by_provider(self, tool_name: str, providers: list[str]) -> bool:
971 """Whether a tool belongs to one of the given providers (case-insensitive).
973 Matched as a full prefix rather than on the first underscore-separated token:
974 splitting on "_" reads `browser_linkedin_search_people` as provider `browser`,
975 so asking for `browser_linkedin` returned nothing at all — silently, since an
976 empty result is indistinguishable from a provider with no tools.
977 """
978 lowered = tool_name.lower()
979 return any(lowered.startswith(provider.lower() + "_") for provider in providers)
981 def _filter_by_action(self, tool_name: str, actions: list[str]) -> bool:
982 """Whether a tool name matches any of the given glob patterns."""
983 return any(fnmatch.fnmatch(tool_name, pattern) for pattern in actions)
985 def fetch_accounts(self) -> list[JsonDict]:
986 """List the accounts linked to this API key.
988 Each entry carries at least ``id``, ``provider`` and ``status``. Only
989 accounts with ``status == "active"`` can serve tools.
991 For each account with ``shared`` false and an ``origin_username``, that username is
992 recorded as the account's end-user id, and each account's ``provider`` as its
993 provider. The record replaces that of any call that started before this one; a call
994 that started earlier but finishes after this one records nothing. Every MCP request
995 this toolset's tools make for that account then carries it in ``x-end-user-id``,
996 which the API requires for a non-shared account. An account passed by id is listed
997 here only when the API refuses one of its requests for want of an end user: the
998 toolset then calls this once, and retries the request with the end user it recorded.
1000 Raises:
1001 StackOneAPIError: If the API answers with an error, including a 429 that
1002 outlasted its retries.
1003 """
1004 with self._cache_lock:
1005 self._accounts_sequence += 1
1006 sequence = self._accounts_sequence
1007 fetching: concurrent.futures.Future[list[JsonDict]] = concurrent.futures.Future()
1008 self._fetching_accounts = fetching
1009 return self._settle_accounts(fetching, sequence)
1011 def _settle_accounts(
1012 self, fetching: concurrent.futures.Future[list[JsonDict]], sequence: int
1013 ) -> list[JsonDict]:
1014 """Make the GET /accounts started as ``sequence``, settling ``fetching`` with it."""
1015 try:
1016 accounts = self._request_accounts(sequence)
1017 except BaseException as exc:
1018 fetching.set_exception(exc)
1019 raise
1020 else:
1021 fetching.set_result(accounts)
1022 return accounts
1023 finally:
1024 with self._cache_lock:
1025 if self._fetching_accounts is fetching:
1026 self._fetching_accounts = None
1028 def _request_accounts(self, sequence: int) -> list[JsonDict]:
1029 url = f"{self.base_url.rstrip('/')}/accounts"
1030 try:
1031 with RateLimitRetryingClient(timeout=self._timeout, retry_within=self._timeout) as client:
1032 response = client.get(
1033 url,
1034 headers=build_request_headers(api_key=self.api_key, extra_headers=self._headers),
1035 )
1036 except RateLimitTimeout as exc:
1037 # Timing out while a 429 is retried is still the rate limit, as it is over MCP.
1038 raise StackOneAPIError(
1039 f"Listing accounts at {url} was rate limited (429) and timed out after "
1040 f"{_js_number(self._timeout)}s while retrying",
1041 429,
1042 None,
1043 ) from exc
1044 except httpx.HTTPError as exc:
1045 # The only public method with no error handling at all: a dead host, a bad
1046 # scheme or a timeout leaked httpx's own exception type straight out of the
1047 # SDK, outside the documented hierarchy.
1048 raise ToolsetLoadError(f"Could not reach {url}: {exc}") from exc
1050 if response.is_error:
1051 # Without this the catch-all in fetch_tools flattens it to a message and
1052 # the status is lost, so a caller cannot tell 401 from 429.
1053 error = StackOneAPIError(
1054 f"Listing accounts at {url} failed with "
1055 f"{response.status_code} {response.reason_phrase}: {response.text.strip()}",
1056 response.status_code,
1057 response.text,
1058 )
1059 if response.status_code == 429:
1060 # Chained to the HTTP 429 so is_rate_limited() reads it as the rate limit it
1061 # is: fatal to a fan-out, as a 429 over MCP is.
1062 raise error from httpx.HTTPStatusError("429", request=response.request, response=response)
1063 raise error
1064 try:
1065 body = response.json()
1066 except (ValueError, UnicodeDecodeError) as exc:
1067 raise ToolsetLoadError(f"Invalid JSON returned by {url}: {exc}") from exc
1068 accounts = body.get("data", body) if isinstance(body, dict) else body
1069 if not isinstance(accounts, list):
1070 # list(dict) yields the KEYS, so coercing here turned an unexpected wrapper
1071 # into a list of strings that blew up much later as an AttributeError.
1072 raise ToolsetLoadError(
1073 f"Unexpected /accounts response shape: expected a list, got {_json_type(accounts)}"
1074 )
1075 end_user_ids = {
1076 account["id"]: account["origin_username"]
1077 for account in accounts
1078 if isinstance(account, dict)
1079 and isinstance(account.get("id"), str)
1080 and account["id"]
1081 and account.get("shared") is False
1082 and isinstance(account.get("origin_username"), str)
1083 and account["origin_username"]
1084 }
1085 providers = {
1086 account["id"]: account["provider"]
1087 for account in accounts
1088 if isinstance(account, dict)
1089 and isinstance(account.get("id"), str)
1090 and account["id"]
1091 and isinstance(account.get("provider"), str)
1092 and account["provider"]
1093 }
1094 with self._cache_lock:
1095 # Only if newer, by start order, than the result last written: a call that
1096 # started first but finished last is answering a question a newer call has
1097 # since answered better, and must not clobber it. Compared with the last
1098 # write rather than the last start, so a newer call that fails leaves this
1099 # one to write.
1100 if sequence > self._end_user_ids_sequence:
1101 self._end_user_ids = end_user_ids
1102 self._providers = providers
1103 self._end_user_ids_sequence = sequence
1104 return accounts
1106 def _look_up_end_user(self, account_id: str, refusal: Exception) -> bool:
1107 """Given the API's refusal of an account's request for want of its end user, list the
1108 accounts and say whether that one now has an end user recorded to retry with.
1110 Only the API's 400 naming this account counts. A GET /accounts already in flight is
1111 joined rather than another made. If it fails, the refusal is raised with that failure
1112 as its cause: a key without the scope to list accounts would otherwise be told only
1113 that the header does not match.
1114 """
1115 if not is_end_user_refusal(refusal, account_id):
1116 return False
1117 try:
1118 self._join_accounts_listing()
1119 except Exception as exc:
1120 # A rate limit is the key's, not the account's, and stays fatal to a fan-out.
1121 if is_rate_limited(exc):
1122 raise
1123 raise refusal from exc
1124 return self._end_user_id(account_id) is not None
1126 def _join_accounts_listing(self) -> None:
1127 """Wait for the GET /accounts in flight, or make one: what it records is then current.
1129 Raises:
1130 Exception: Whatever that GET /accounts failed with.
1131 """
1132 with self._cache_lock:
1133 fetching = self._fetching_accounts
1134 sequence = 0
1135 if fetching is None:
1136 self._accounts_sequence += 1
1137 sequence = self._accounts_sequence
1138 fetching = self._fetching_accounts = concurrent.futures.Future()
1139 if sequence:
1140 with contextlib.suppress(Exception):
1141 self._settle_accounts(fetching, sequence)
1142 fetching.result()
1144 def _learn_providers(self, account_ids: list[str]) -> set[str] | None:
1145 """If any of these accounts' providers is unknown, and no lookup has missed it within
1146 ``FAILED_ACCOUNT_RETRY_SECONDS``, join or make one GET /accounts to learn it, and return
1147 those it still could not name; ``None`` if it made no lookup. If that fails the providers
1148 stay unknown, except that a rate limit is raised. Each one still unknown is recorded as
1149 missed, unless clear_catalog_cache() ran meanwhile, which promises to forget misses.
1150 """
1151 missed = self._recently_missed(account_ids)
1152 with self._cache_lock:
1153 unknown = [a for a in account_ids if a not in self._providers and a not in missed]
1154 generation = self._cache_generation
1155 if not unknown:
1156 return None
1157 try:
1158 self._join_accounts_listing()
1159 except Exception as exc:
1160 if is_rate_limited(exc):
1161 raise
1162 missed_at = _clock()
1163 with self._cache_lock:
1164 still_unknown = {account for account in unknown if account not in self._providers}
1165 if generation == self._cache_generation:
1166 for account in still_unknown:
1167 self._provider_misses[account] = missed_at
1168 return still_unknown
1170 def _recently_missed(self, account_ids: list[str]) -> set[str]:
1171 """Those of these accounts whose provider is unknown and a lookup failed to name
1172 within ``FAILED_ACCOUNT_RETRY_SECONDS``.
1173 """
1174 now = _clock()
1175 with self._cache_lock:
1176 return {
1177 account
1178 for account in account_ids
1179 if account not in self._providers
1180 and (missed_at := self._provider_misses.get(account)) is not None
1181 and now - missed_at < FAILED_ACCOUNT_RETRY_SECONDS
1182 }
1184 def _end_user_id(self, account_id: str) -> str | None:
1185 """The end-user id GET /accounts recorded for an account, if it is not shared."""
1186 with self._cache_lock:
1187 return self._end_user_ids.get(account_id)
1189 def _discover_account_ids(self) -> list[str]:
1190 """List the linked accounts this API key can use.
1192 The MCP endpoint requires an ``x-account-id`` on every request, so an API
1193 key on its own is not enough to list tools. Rather than make every caller
1194 supply one, ask the API which accounts the key has.
1196 Warning: For organizations with many connected accounts, relying on discovery
1197 fetches the tool catalog for every account. If you have a large number of
1198 accounts, it is highly recommended to supply specific ``account_id`` or
1199 ``account_ids`` to avoid excessive API round trips and huge model contexts.
1201 Non-shared accounts are skipped, with one warning naming them, unless the toolset was
1202 built with ``include_non_shared=True``: each belongs to a single end user.
1204 Raises:
1205 ToolsetConfigError: If the key has no accounts, none are active, or every active
1206 one is non-shared and skipped.
1207 """
1208 if self._discovered_account_ids is not None:
1209 return self._discovered_account_ids
1211 with self._cache_lock:
1212 if self._discovered_account_ids is not None: 1212 ↛ 1213line 1212 didn't jump to line 1213 because the condition on line 1212 was never true
1213 return self._discovered_account_ids
1214 future = self._discovering
1215 owner = future is None
1216 if owner:
1217 future = self._discovering = concurrent.futures.Future()
1219 if future is None: 1219 ↛ 1220line 1219 didn't jump to line 1220 because the condition on line 1219 was never true
1220 raise RuntimeError("account discovery future was not initialized")
1221 if not owner:
1222 # A discovery already in flight: wait for it rather than issue a second
1223 # GET /accounts. Its result, or its exception, is shared with every waiter.
1224 return future.result()
1226 try:
1227 active = self._fetch_active_account_ids()
1228 except BaseException as exc:
1229 with self._cache_lock:
1230 if self._discovering is future:
1231 self._discovering = None
1232 future.set_exception(exc)
1233 raise
1234 with self._cache_lock:
1235 if self._discovering is future:
1236 self._discovering = None
1237 future.set_result(active)
1238 return active
1240 def _fetch_active_account_ids(self) -> list[str]:
1241 generation = self._cache_generation
1242 accounts = self.fetch_accounts()
1244 active = [a["id"] for a in accounts if a.get("status") == "active" and a.get("id")]
1245 if not active:
1246 if not accounts:
1247 raise ToolsetConfigError(
1248 "This API key has no linked accounts. Link one in the StackOne "
1249 "dashboard, or pass an account id explicitly."
1250 )
1251 listed = ", ".join(f"{a.get('provider')} ({a.get('status')})" for a in accounts)
1252 raise ToolsetConfigError(
1253 f"None of this API key's {len(accounts)} linked accounts are active: {listed}. "
1254 "Re-link them in the StackOne dashboard, or pass an account id explicitly."
1255 )
1257 # Each non-shared account belongs to a single end user: discovering it would put one
1258 # end user's tools in front of whoever uses this toolset.
1259 if not self._include_non_shared:
1260 skipped = sorted(
1261 {a["id"] for a in accounts if a.get("id") in active and a.get("shared") is False}
1262 )
1263 if skipped and len(skipped) == len(set(active)):
1264 raise ToolsetConfigError(
1265 f"None of this API key's {len(skipped)} active account(s) are shared: each "
1266 "belongs to a single end user. Pass their account ids, or opt in to non-shared "
1267 "accounts, to use them."
1268 )
1269 if skipped:
1270 logger.warning(
1271 "Discovery skipped %d non-shared account(s) (%s): each belongs to a single end "
1272 "user. Pass their account ids, or opt in to non-shared accounts, to use them.",
1273 len(skipped),
1274 ", ".join(skipped),
1275 )
1276 active = [account for account in active if account not in skipped]
1278 # Not kept if clear_catalog_cache() ran while the list was in flight: it may name
1279 # accounts the clear was meant to forget.
1280 with self._cache_lock:
1281 if generation == self._cache_generation:
1282 self._discovered_account_ids = active
1283 return active
1285 def _build_mcp_headers(self, account_id: str | None) -> dict[str, str]:
1286 end_user_id = self._end_user_id(account_id) if account_id else None
1287 return build_request_headers(
1288 api_key=self.api_key,
1289 account_id=account_id,
1290 end_user_id=end_user_id,
1291 extra_headers=self._headers,
1292 )
1294 def _create_tool(
1295 self,
1296 tool_def: McpToolDefinition,
1297 account_id: str | None,
1298 endpoint: str,
1299 ) -> StackOneTool:
1300 """Build an executable tool from a served catalog entry.
1302 Every tool, in every mode, executes over MCP ``tools/call`` on the endpoint that
1303 listed it, with the ``x-account-id`` of the account that listed it, and the
1304 ``x-end-user-id`` recorded for that account when the request is made.
1305 """
1306 schema = tool_def.input_schema or {}
1307 # Pop keys we explicitly override to avoid "multiple values for keyword argument"
1308 rest = dict(schema)
1309 # A string type is passed through; anything else, such as ["object", "null"], is
1310 # "object", as in Node. str() handed the model "['object', 'null']".
1311 served_type = rest.pop("type", "object")
1312 schema_type = served_type if isinstance(served_type, str) else "object"
1313 rest.pop("properties", None)
1314 schema_properties = self._normalize_schema_properties(schema)
1316 parameters = ToolParameters(
1317 **rest,
1318 type=schema_type,
1319 properties=schema_properties,
1320 )
1321 tool = StackOneMcpTool(
1322 name=tool_def.name,
1323 description=tool_def.description or "",
1324 parameters=parameters,
1325 api_key=self.api_key,
1326 endpoint=endpoint,
1327 account_id=account_id,
1328 headers=self._headers,
1329 timeout=self._timeout,
1330 )
1331 tool._end_user_id_of = self._end_user_id
1332 tool._look_up_end_user = self._look_up_end_user
1333 return tool
1335 def _normalize_schema_properties(self, schema: dict[str, Any]) -> dict[str, Any]:
1336 """Mirror the served schema's properties, recording requiredness per property.
1338 The served schema expresses requiredness as a top-level ``required`` list.
1339 The SDK carries it per-property as ``nullable`` so each property is
1340 self-describing; everything else in the property is preserved verbatim.
1341 """
1342 properties = schema.get("properties")
1343 if not isinstance(properties, dict):
1344 return {}
1346 raw_required = schema.get("required")
1347 # A string `required` would iterate as characters and mark every real property
1348 # optional; a null would raise and fail the whole catalog over one bad tool. A
1349 # non-string entry is dropped, as to_openai_function() drops it, so the marker and
1350 # the `required` a model is shown cannot disagree.
1351 served_required = raw_required if isinstance(raw_required, list) else []
1352 required_fields = {name for name in served_required if isinstance(name, str)}
1353 normalized: dict[str, Any] = {}
1355 for name, details in properties.items():
1356 if isinstance(details, dict):
1357 prop = copy.deepcopy(details)
1358 else:
1359 prop = {"description": str(details)}
1361 # Assign, never setdefault: a served `nullable` (OpenAPI 3.0 style) means
1362 # "accepts null", not "optional". Letting it stand would drop the field
1363 # from `required` and the model would omit it.
1364 prop["nullable"] = name not in required_fields
1365 normalized[str(name)] = prop
1367 return normalized