Skip to content

Commit c2c0811

Browse files
Ilanlidoclaude
andcommitted
CM-64735 - Reduce ai-guardrails scan latency
Cluster of CLI changes that take a 0-detection ai-guardrails scan from ~5.5s to ~1.55s measured via the packaged onedir with warm tokens: - Reuse a process-wide requests.Session so TCP+TLS handshakes amortize across the multiple API calls per scan (was one-shot per request on macOS/Linux, paying ~300-450ms of handshake on each call). - Argv-peek lazy subapp registration: import only the invoked subapp at startup instead of the full set, skipping ~300ms of unrelated imports on hot paths. - Skip get_detection_rules when the scan returned zero detections (the common case for ai-guardrails hooks). - Skip POST /cli-scan/{id}/status on sync flows where the /sync response already returned the full result inline. - Share access tokens across CycodeClientBase instances by re-reading the on-disk cache before doing the HTTP refresh, so the ai-security client doesn't duplicate work the scan client just did. - Skip version-checker on the ai-guardrails scan hot path (it emits JSON to stdout — an upgrade notice would corrupt the response, plus the PyPI round-trip is a ~500ms cache-miss cost). - Send cli_start_time in /sync so the server can compute honest end-to-end execution_time independent of CLI-side wall clock. - Add coverage for the argv-peek invariant: root options that take a value must be registered so argv-peek skips past their values. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent d53991b commit c2c0811

9 files changed

Lines changed: 266 additions & 62 deletions

File tree

cycode/__init__.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1 +1,8 @@
1+
import time as _time
2+
3+
# Unix-epoch wall clock captured at the earliest possible moment of CLI
4+
# startup. Sent as `scan_parameters.cli_start_time` so the server can compute
5+
# end-to-end scan duration from the moment the user actually triggered it.
6+
_BOOT_WALL: float = _time.time()
7+
18
__version__ = '0.0.0' # DON'T TOUCH. Placeholder. Will be filled automatically on poetry build from Git Tag

cycode/cli/app.py

Lines changed: 88 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import importlib
12
import logging
23
import sys
34
from typing import Annotated, Optional
@@ -10,12 +11,7 @@
1011
from typer.completion import install_callback, show_callback
1112

1213
from cycode import __version__
13-
from cycode.cli.apps import ai_guardrails, ai_remediation, auth, configure, ignore, report, report_import, scan, status
1414
from cycode.cli.apps.api import get_platform_group
15-
16-
if sys.version_info >= (3, 10):
17-
from cycode.cli.apps import mcp
18-
1915
from cycode.cli.cli_types import OutputTypeOption
2016
from cycode.cli.consts import CLI_CONTEXT_SETTINGS
2117
from cycode.cli.printers import ConsolePrinter
@@ -46,17 +42,88 @@
4642
add_completion=False, # we add it manually to control the rich help panel
4743
)
4844

49-
app.add_typer(ai_guardrails.app)
50-
app.add_typer(ai_remediation.app)
51-
app.add_typer(auth.app)
52-
app.add_typer(configure.app)
53-
app.add_typer(ignore.app)
54-
app.add_typer(report.app)
55-
app.add_typer(report_import.app)
56-
app.add_typer(scan.app)
57-
app.add_typer(status.app)
45+
# Top-level subcommand → module providing its Typer app. Peeking at sys.argv
46+
# lets us import only the invoked subapp on the hot path (e.g.
47+
# `cycode ai-guardrails scan`), skipping ~300ms of unrelated imports.
48+
_SUBAPP_MODULES: dict[str, str] = {
49+
'ai-guardrails': 'cycode.cli.apps.ai_guardrails',
50+
'ai-remediation': 'cycode.cli.apps.ai_remediation',
51+
'auth': 'cycode.cli.apps.auth',
52+
'configure': 'cycode.cli.apps.configure',
53+
'ignore': 'cycode.cli.apps.ignore',
54+
'report': 'cycode.cli.apps.report',
55+
'import': 'cycode.cli.apps.report_import',
56+
'scan': 'cycode.cli.apps.scan',
57+
'status': 'cycode.cli.apps.status',
58+
}
5859
if sys.version_info >= (3, 10):
59-
app.add_typer(mcp.app)
60+
_SUBAPP_MODULES['mcp'] = 'cycode.cli.apps.mcp'
61+
62+
# Aliases: alternate spellings that resolve to a primary subcommand key.
63+
_SUBAPP_ALIASES: dict[str, str] = {
64+
'ai_remediation': 'ai-remediation', # backward-compat underscore form
65+
'version': 'status',
66+
}
67+
68+
# Root-level options that consume a following value; argv-peek must skip past
69+
# both the option and its value when scanning for the first positional arg.
70+
_ROOT_OPTS_WITH_VALUE = frozenset(
71+
{
72+
'--output',
73+
'-o',
74+
'--user-agent',
75+
'--client-secret',
76+
'--client-id',
77+
'--id-token',
78+
'--show-completion',
79+
}
80+
)
81+
82+
83+
def _detect_invocation() -> tuple[Optional[str], Optional[str]]:
84+
"""Return (top-level-subapp, second-level-subcommand) parsed from sys.argv.
85+
86+
Both values may be None: when no positional arg matches a known subapp,
87+
or when the user only provided a top-level subcommand.
88+
"""
89+
positionals = []
90+
args = sys.argv[1:]
91+
i = 0
92+
while i < len(args):
93+
arg = args[i]
94+
if arg in _ROOT_OPTS_WITH_VALUE:
95+
i += 2
96+
elif arg.startswith('-'):
97+
# Any flag form: short, long, --key=value, or '--' marker. Skip the token only.
98+
i += 1
99+
else:
100+
positionals.append(arg)
101+
if len(positionals) >= 2:
102+
break
103+
i += 1
104+
subapp = positionals[0] if positionals else None
105+
subapp = _SUBAPP_ALIASES.get(subapp, subapp)
106+
if subapp not in _SUBAPP_MODULES:
107+
return None, None
108+
subcommand = positionals[1] if len(positionals) >= 2 else None
109+
return subapp, subcommand
110+
111+
112+
# Computed once at import; reused by lazy registration and the version-checker skip.
113+
_INVOKED_SUBAPP, _INVOKED_SUBCOMMAND = _detect_invocation()
114+
115+
116+
def _register_subapps(only: Optional[str]) -> None:
117+
if only is not None:
118+
app.add_typer(importlib.import_module(_SUBAPP_MODULES[only]).app)
119+
return
120+
# Cold path (--help, completion, unknown subcommand): load all modules so
121+
# root help lists everything. Deduplicate since aliases share modules.
122+
for module_path in dict.fromkeys(_SUBAPP_MODULES.values()):
123+
app.add_typer(importlib.import_module(module_path).app)
124+
125+
126+
_register_subapps(_INVOKED_SUBAPP)
60127

61128
# Register the `platform` command group (dynamically built from the OpenAPI spec).
62129
# The group itself is constructed cheaply at import time; the spec is only fetched
@@ -81,6 +148,12 @@ def _get_group_with_platform(app_typer: typer.Typer) -> click.Group:
81148

82149

83150
def check_latest_version_on_close(ctx: typer.Context) -> None:
151+
# Skip on `cycode ai-guardrails scan` — it emits JSON to stdout, so an
152+
# upgrade notice would corrupt the response. Human-driven sibling commands
153+
# (install, uninstall, status, session-start) still get the notice.
154+
if (_INVOKED_SUBAPP, _INVOKED_SUBCOMMAND) == ('ai-guardrails', 'scan'):
155+
return
156+
84157
output = ctx.obj.get('output')
85158
# don't print anything if the output is JSON
86159
if output == OutputTypeOption.JSON:

cycode/cli/apps/scan/code_scanner.py

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -204,18 +204,21 @@ def _scan_batch_thread_func(batch: list[Document]) -> tuple[str, CliError, Local
204204
'zip_file_size': zip_file_size,
205205
},
206206
)
207-
report_scan_status(
208-
cycode_client,
209-
scan_type,
210-
scan_id,
211-
scan_completed,
212-
relevant_detections_count,
213-
detections_count,
214-
len(batch),
215-
zip_file_size,
216-
command_scan_type,
217-
error_message,
218-
)
207+
# Sync flows already received the full result inline; only async flows
208+
# need a separate status report to signal polling completion.
209+
if not should_use_sync_flow:
210+
report_scan_status(
211+
cycode_client,
212+
scan_type,
213+
scan_id,
214+
scan_completed,
215+
relevant_detections_count,
216+
detections_count,
217+
len(batch),
218+
zip_file_size,
219+
command_scan_type,
220+
error_message,
221+
)
219222

220223
return scan_id, error, local_scan_result
221224

cycode/cli/apps/scan/scan_parameters.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import typer
44

5+
from cycode import _BOOT_WALL
56
from cycode.cli.apps.scan.remote_url_resolver import get_remote_url_scan_parameter
67
from cycode.cli.utils.scan_utils import generate_unique_scan_id
78
from cycode.logger import get_logger
@@ -17,6 +18,7 @@ def _get_default_scan_parameters(ctx: typer.Context) -> dict:
1718
'license_compliance': ctx.obj.get('license-compliance'),
1819
'command_type': ctx.info_name.replace('-', '_'), # save backward compatibility
1920
'aggregation_id': str(generate_unique_scan_id()),
21+
'cli_start_time': _BOOT_WALL,
2022
}
2123

2224

cycode/cli/apps/scan/scan_result.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -189,6 +189,10 @@ def enrich_scan_result_with_data_from_detection_rules(
189189
for detection in detections_per_file.detections:
190190
detection_rule_ids.add(detection.detection_rule_id)
191191

192+
if not detection_rule_ids:
193+
logger.debug('No detections to enrich, skipping detection_rules fetch')
194+
return
195+
192196
detection_rules = cycode_client.get_detection_rules(detection_rule_ids)
193197
detection_rules_by_id = {detection_rule.detection_rule_id: detection_rule for detection_rule in detection_rules}
194198

cycode/cyclient/base_token_auth_client.py

Lines changed: 27 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -24,19 +24,10 @@ def __init__(self, client_id: str) -> None:
2424
self.client_id = client_id
2525

2626
self._credentials_manager = CredentialsManager()
27-
# load cached access token
28-
access_token, expires_in, creator = self._credentials_manager.get_access_token()
29-
30-
self._access_token = self._expires_in = None
31-
expected_creator = self._create_jwt_creator()
32-
if creator == expected_creator:
33-
# we must be sure that cached access token is created using the same client id and client secret.
34-
# because client id and client secret could be passed via command, via env vars or via config file.
35-
# we must not use cached access token if client id or client secret was changed.
36-
self._access_token = access_token
37-
self._expires_in = arrow.get(expires_in) if expires_in else None
38-
27+
self._access_token = None
28+
self._expires_in = None
3929
self._lock = Lock()
30+
self._load_token_from_disk()
4031

4132
def get_access_token(self) -> str:
4233
with self._lock:
@@ -51,8 +42,30 @@ def invalidate_access_token(self, in_storage: bool = False) -> None:
5142
self._credentials_manager.update_access_token(None, None, None)
5243

5344
def refresh_access_token_if_needed(self) -> None:
54-
if self._access_token is None or self._expires_in is None or arrow.utcnow() >= self._expires_in:
55-
self.refresh_access_token()
45+
if self._has_valid_token():
46+
return
47+
# Re-check disk before doing the network refresh: another client instance
48+
# in this process may have already refreshed and persisted a fresh token.
49+
self._load_token_from_disk()
50+
if self._has_valid_token():
51+
return
52+
self.refresh_access_token()
53+
54+
def _has_valid_token(self) -> bool:
55+
return self._access_token is not None and self._expires_in is not None and arrow.utcnow() < self._expires_in
56+
57+
def _load_token_from_disk(self) -> None:
58+
access_token, expires_in, creator = self._credentials_manager.get_access_token()
59+
expected_creator = self._create_jwt_creator()
60+
# We must be sure that cached access token is created using the same client id and client secret.
61+
# Because client id and client secret could be passed via command, via env vars or via config file.
62+
# We must not use cached access token if client id or client secret was changed.
63+
if creator == expected_creator and access_token:
64+
self._access_token = access_token
65+
self._expires_in = arrow.get(expires_in) if expires_in else None
66+
else:
67+
self._access_token = None
68+
self._expires_in = None
5669

5770
def refresh_access_token(self) -> None:
5871
auth_response = self._request_new_access_token()

cycode/cyclient/cycode_client_base.py

Lines changed: 24 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1+
import functools
12
import os
23
import platform
34
import ssl
@@ -39,16 +40,29 @@ def cert_verify(self, *args, **kwargs) -> None:
3940
conn.ca_certs = None
4041

4142

43+
@functools.cache
44+
def _get_session() -> requests.Session:
45+
"""Process-wide Session so TCP+TLS connections are reused across all API calls."""
46+
session = requests.Session()
47+
# On Windows without an explicit CA bundle env var, fall back to the system
48+
# trust store via a custom SSL context.
49+
if platform.system() == 'Windows' and not (
50+
os.environ.get('REQUESTS_CA_BUNDLE') or os.environ.get('CURL_CA_BUNDLE')
51+
):
52+
session.mount('https://', SystemStorageSslContext())
53+
return session
54+
55+
4256
def _get_request_function() -> Callable:
43-
if os.environ.get('REQUESTS_CA_BUNDLE') or os.environ.get('CURL_CA_BUNDLE'):
44-
return requests.request
57+
return _get_session().request
4558

46-
if platform.system() != 'Windows':
47-
return requests.request
4859

49-
session = requests.Session()
50-
session.mount('https://', SystemStorageSslContext())
51-
return session.request
60+
def _log_response(response: Response, url: str, hide_response_content_log: bool) -> None:
61+
content = 'HIDDEN' if hide_response_content_log else response.text
62+
logger.debug(
63+
'Receiving response, %s',
64+
{'status_code': response.status_code, 'url': url, 'content': content},
65+
)
5266

5367

5468
_REQUEST_ERRORS_TO_RETRY = (
@@ -182,12 +196,7 @@ def _send_multipart(
182196
response = _get_request_function()(
183197
method='post', url=url, data=tracker, headers=headers, timeout=self.timeout
184198
)
185-
186-
content = 'HIDDEN' if hide_response_content_log else response.text
187-
logger.debug(
188-
'Receiving response, %s',
189-
{'status_code': response.status_code, 'url': url, 'content': content},
190-
)
199+
_log_response(response, url, hide_response_content_log)
191200

192201
response.raise_for_status()
193202
return response
@@ -231,14 +240,8 @@ def _execute(
231240

232241
try:
233242
headers = self.get_request_headers(headers, without_auth=without_auth)
234-
request = _get_request_function()
235-
response = request(method=method, url=url, timeout=timeout, headers=headers, **kwargs)
236-
237-
content = 'HIDDEN' if hide_response_content_log else response.text
238-
logger.debug(
239-
'Receiving response, %s',
240-
{'status_code': response.status_code, 'url': url, 'content': content},
241-
)
243+
response = _get_request_function()(method=method, url=url, timeout=timeout, headers=headers, **kwargs)
244+
_log_response(response, url, hide_response_content_log)
242245

243246
response.raise_for_status()
244247
return response

pyinstaller.spec

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,26 @@ CLI_VERSION = _dunamai.get_version('cycode', first_choice=_dunamai.Version.from_
2121
with open(_INIT_FILE_PATH, 'w', encoding='UTF-8') as file:
2222
file.write(prev_content.replace(VERSION_PLACEHOLDER, CLI_VERSION))
2323

24+
# Top-level subapp modules are loaded lazily via importlib.import_module() in
25+
# cycode/cli/app.py to keep startup fast on hot paths (e.g. ai-guardrails scan).
26+
# PyInstaller's static analyzer can't see those imports, so list them explicitly.
27+
_hiddenimports = [
28+
'cycode.cli.apps.ai_guardrails',
29+
'cycode.cli.apps.ai_remediation',
30+
'cycode.cli.apps.auth',
31+
'cycode.cli.apps.configure',
32+
'cycode.cli.apps.ignore',
33+
'cycode.cli.apps.report',
34+
'cycode.cli.apps.report_import',
35+
'cycode.cli.apps.scan',
36+
'cycode.cli.apps.status',
37+
'cycode.cli.apps.mcp',
38+
]
39+
2440
a = Analysis(
2541
scripts=['cycode/cli/main.py'],
2642
excludes=['tests', 'setuptools', 'pkg_resources'],
43+
hiddenimports=_hiddenimports,
2744
)
2845

2946
exe_args = [PYZ(a.pure), a.scripts, a.binaries, a.datas]

0 commit comments

Comments
 (0)